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

Mixed precision diffusion models with scaled fp8.

Browse Source 
This change allows supports for diffusion models where all the linears are
scaled fp8 while the other weights are the original precision.
This commit is contained in:
comfyanonymous
2024-10-21 18:12:51 -04:00
parent 83ca891118
commit 0075c6d096
5 changed files with 14 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
comfy/model_detection.py
View File

@@ -288,8 +288,11 @@ def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=Fal
if model_config is None and use_base_if_no_match:
model_config = comfy.supported_models_base.BASE(unet_config)
if "{}scaled_fp8".format(unet_key_prefix) in state_dict:
model_config.scaled_fp8 = True
scaled_fp8_weight = state_dict.get("{}scaled_fp8".format(unet_key_prefix), None)
if scaled_fp8_weight is not None:
model_config.scaled_fp8 = scaled_fp8_weight.dtype
if model_config.scaled_fp8 == torch.float32:
model_config.scaled_fp8 = torch.float8_e4m3fn
return model_config