Merge branch 'fix_unet_wrapper_function_name' of https://github.com/g…

…ameltb/ComfyUI
comfyanonymous · Nov 5, 2023 · a6c83b3 · a6c83b3
2 parents 02f062b + 7e455ad
commit a6c83b3
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py
@@ -107,10 +107,10 @@ def model_patches_to(self, device):
                 for k in patch_list:
                     if hasattr(patch_list[k], "to"):
                         patch_list[k] = patch_list[k].to(device)
-        if "unet_wrapper_function" in self.model_options:
-            wrap_func = self.model_options["unet_wrapper_function"]
+        if "model_function_wrapper" in self.model_options:
+            wrap_func = self.model_options["model_function_wrapper"]
             if hasattr(wrap_func, "to"):
-                self.model_options["unet_wrapper_function"] = wrap_func.to(device)
+                self.model_options["model_function_wrapper"] = wrap_func.to(device)
 
     def model_dtype(self):
         if hasattr(self.model, "get_dtype"):