jinx/ComfyUI
1
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-08-02 23:14:49 +08:00
Code Issues Activity

Automatically use fp8 for diffusion model weights if:

Browse Source 
Checkpoint contains weights in fp8.

There isn't enough memory to load the diffusion model in GPU vram.
This commit is contained in:
comfyanonymous
2024-08-03 13:45:19 -04:00
parent f123328b82
commit ba9095e5bd
4 changed files with 34 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
comfy/utils.py
View File

@@ -40,9 +40,19 @@ def calculate_parameters(sd, prefix=""):
params = 0
for k in sd.keys():
if k.startswith(prefix):
params += sd[k].nelement()
w = sd[k]
params += w.nelement()
return params
def weight_dtype(sd, prefix=""):
dtypes = {}
for k in sd.keys():
if k.startswith(prefix):
w = sd[k]
dtypes[w.dtype] = dtypes.get(w.dtype, 0) + 1
return max(dtypes, key=dtypes.get)
def state_dict_key_replace(state_dict, keys_to_replace):
for x in keys_to_replace:
if x in state_dict: