Fix no longer working on old pytorch.

guill · Aug 2, 2024 · 369f459 · 369f459
1 parent ce9ac2f
commit 369f459
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/nodes.py b/nodes.py
@@ -826,9 +826,14 @@ def INPUT_TYPES(s):
     CATEGORY = "advanced/loaders"
 
     def load_unet(self, unet_name, weight_dtype):
-        weight_dtype = {"default":None, "fp8_e4m3fn":torch.float8_e4m3fn, "fp8_e5m2":torch.float8_e4m3fn}[weight_dtype]
+        dtype = None
+        if weight_dtype == "fp8_e4m3fn":
+            dtype = torch.float8_e4m3fn
+        elif weight_dtype == "fp8_e5m2":
+            dtype = torch.float8_e5m2
+
         unet_path = folder_paths.get_full_path("unet", unet_name)
-        model = comfy.sd.load_unet(unet_path, dtype=weight_dtype)
+        model = comfy.sd.load_unet(unet_path, dtype=dtype)
         return (model,)
 
 class CLIPLoader: