fix

hook_comfyui_kolors_v2.py

@@ -6,6 +6,7 @@ import comfy.model_detection as model_detection
 import comfy.supported_models
 import comfy.utils
 
+import comfy.model_management as model_management
 import torch
 from comfy import model_base
 from comfy.model_base import sdxl_pooled, CLIPEmbeddingNoiseAugmentation, Timestep, ModelType
@@ -13,6 +14,7 @@ from comfy.model_base import sdxl_pooled, CLIPEmbeddingNoiseAugmentation, Timest
 
 from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel
 from comfy.cldm.cldm import ControlNet
+import torch.amp
 
 # try:
 #     import comfy.samplers as samplers
@@ -50,26 +52,39 @@ from comfy.cldm.cldm import ControlNet
 #     print("CFGGuider not found, skipping patching")
 
 
+def get_torch_device(): 
+    if model_management.directml_enabled: 
+        return model_management.directml_device
+    if model_management.cpu_state == model_management.CPUState.MPS:
+        return "cpu"
+    if model_management.cpu_state == model_management.CPUState.CPU:
+        return "cpu"
+    else:
+        if model_management.is_intel_xpu():
+            return "xpu"
+        else:
+            return "cuda"
+ 
+def get_torch_dtype():
+    dtype = model_management.unet_dtype()
+    if dtype not in [torch.float32, torch.float16, torch.bfloat16]:
+        dtype = torch.float16 if model_management.should_use_fp16() else torch.float32
+
+    return dtype
+
 class KolorsUNetModel(UNetModel):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.encoder_hid_proj = nn.Linear(
             4096, 2048, bias=True)
 
-    def forward(self, *args, **kwargs):
-        with torch.cuda.amp.autocast(enabled=True):
+    def forward(self, *args, **kwargs): 
+        with torch.amp.autocast(enabled=True, device_type=get_torch_device(), dtype=get_torch_dtype()):
             if "context" in kwargs:
                 kwargs["context"] = self.encoder_hid_proj(
-                    kwargs["context"])
-
-                # if "y" in kwargs:
-                #     if kwargs["y"].shape[1] == 2816:
-                #         # 扩展至5632
-                #         kwargs["y"] = torch.cat(
-                #             torch.zeros(kwargs["y"].shape[0], 2816).to(kwargs["y"].device), kwargs["y"], dim=1)
-
+                    kwargs["context"]) 
             result = super().forward(*args, **kwargs)
-            return result
+        return result
 
 
 class KolorsSDXL(model_base.SDXL):
@@ -264,7 +279,7 @@ class KolorsControlNet(ControlNet):
             4096, 2048, bias=True)
 
     def forward(self, *args, **kwargs):
-        with torch.cuda.amp.autocast(enabled=True):
+        with torch.amp.autocast(enabled=True, device_type=get_torch_device(), dtype=get_torch_dtype()):
             if "context" in kwargs:
                 kwargs["context"] = self.encoder_hid_proj(
                     kwargs["context"])

mz_kolors_core.py

@@ -12,6 +12,9 @@ from types import MethodType
 import torch
 import folder_paths
 import comfy.model_management as mm
+import comfy.model_management as model_management
+from .hook_comfyui_kolors_v2 import get_torch_device, get_torch_dtype
+import torch.amp
 
 
 def chatglm3_text_encode(chatglm3_model, prompt):
@@ -272,7 +275,7 @@ def MZ_KolorsControlNetPatch_call(kwargs):
         super_forward = ControlNet.forward
 
         def KolorsControlNet_forward(self, x, hint, timesteps, context, **kwargs):
-            with torch.cuda.amp.autocast(enabled=True):
+            with torch.amp.autocast(enabled=True, device_type=get_torch_device(), dtype=get_torch_dtype()):
                 context = self.encoder_hid_proj(context)
                 return super_forward(self, x, hint, timesteps, context, **kwargs)
 
@@ -320,7 +323,7 @@ def MZ_KolorsControlNetPatch_call(kwargs):
         super_forward = ControlNet.forward
 
         def KolorsControlNet_forward(self, x, hint, timesteps, context, **kwargs):
-            with torch.cuda.amp.autocast(enabled=True):
+            with torch.amp.autocast(enabled=True, device_type=get_torch_device(), dtype=get_torch_dtype()):
                 context = self.encoder_hid_proj(context)
                 return super_forward(self, x, hint, timesteps, context, **kwargs)