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

Add a way to load the diffusion model in fp8 with UNETLoader node.

Browse Source
This commit is contained in:
comfyanonymous
2024-08-01 13:28:41 -04:00
parent f2b80f95d2
commit d7430a1651
2 changed files with 12 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
nodes.py
View File

@@ -818,15 +818,17 @@ class UNETLoader:
@classmethod
def INPUT_TYPES(s):
return {"required": { "unet_name": (folder_paths.get_filename_list("unet"), ),
"weight_dtype": (["default", "fp8_e4m3fn", "fp8_e5m2"],)
}}
RETURN_TYPES = ("MODEL",)
FUNCTION = "load_unet"
CATEGORY = "advanced/loaders"
def load_unet(self, unet_name):
def load_unet(self, unet_name, weight_dtype):
weight_dtype = {"default":None, "fp8_e4m3fn":torch.float8_e4m3fn, "fp8_e5m2":torch.float8_e4m3fn}[weight_dtype]
unet_path = folder_paths.get_full_path("unet", unet_name)
model = comfy.sd.load_unet(unet_path)
model = comfy.sd.load_unet(unet_path, dtype=weight_dtype)
return (model,)
class CLIPLoader: