Use fp16 if checkpoint weights are fp16 and the model supports it. · code/app-python-comfyui@1804397 · GitHub
Skip to content

Commit 1804397

Browse files
Use fp16 if checkpoint weights are fp16 and the model supports it.
1 parent f4dac8a commit 1804397

3 files changed

Lines changed: 12 additions & 20 deletions

File tree

‎comfy/controlnet.py‎

Lines changed: 2 additions & 8 deletions

‎comfy/model_management.py‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -674,7 +674,7 @@ def unet_inital_load_device(parameters, dtype):
674674
def maximum_vram_for_weights(device=None):
675675
return (get_total_memory(device) * 0.88 - minimum_inference_memory())
676676

677-
def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]):
677+
def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32], weight_dtype=None):
678678
if model_params < 0:
679679
model_params = 1000000000000000000000
680680
if args.fp32_unet:
@@ -692,10 +692,8 @@ def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, tor
692692

693693
fp8_dtype = None
694694
try:
695-
for dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
696-
if dtype in supported_dtypes:
697-
fp8_dtype = dtype
698-
break
695+
if weight_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
696+
fp8_dtype = weight_dtype
699697
except:
700698
pass
701699

@@ -707,7 +705,7 @@ def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, tor
707705
if model_params * 2 > free_model_memory:
708706
return fp8_dtype
709707

710-
if PRIORITIZE_FP16:
708+
if PRIORITIZE_FP16 or weight_dtype == torch.float16:
711709
if torch.float16 in supported_dtypes and should_use_fp16(device=device, model_params=model_params):
712710
return torch.float16
713711

‎comfy/sd.py‎

Lines changed: 6 additions & 6 deletions

0 commit comments

Comments
 (0)