1
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-08-02 15:04:50 +08:00

load_unet -> load_diffusion_model with a model_options argument.

This commit is contained in:
comfyanonymous
2024-08-12 23:18:54 -04:00
parent 5942c17d55
commit a562c17e8a
3 changed files with 19 additions and 8 deletions

View File

@@ -826,14 +826,14 @@ class UNETLoader:
CATEGORY = "advanced/loaders"
def load_unet(self, unet_name, weight_dtype):
dtype = None
model_options = {}
if weight_dtype == "fp8_e4m3fn":
dtype = torch.float8_e4m3fn
model_options["dtype"] = torch.float8_e4m3fn
elif weight_dtype == "fp8_e5m2":
dtype = torch.float8_e5m2
model_options["dtype"] = torch.float8_e5m2
unet_path = folder_paths.get_full_path("unet", unet_name)
model = comfy.sd.load_unet(unet_path, dtype=dtype)
model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options)
return (model,)
class CLIPLoader: