fix

Update PyProject Toml - License

ComfyUI_IPAdapter_plus/IPAdapterPlus.py

@@ -46,7 +46,7 @@ WEIGHT_TYPES = ["linear", "ease in", "ease out", 'ease in-out', 'reverse in-out'
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 """
 class IPAdapter(nn.Module):
-    def __init__(self, ipadapter_model, cross_attention_dim=1024, output_cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4, is_sdxl=False, is_plus=False, is_full=False, is_faceid=False, is_portrait_unnorm=False, is_kwai_kolors=False):
+    def __init__(self, ipadapter_model, cross_attention_dim=1024, output_cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4, is_sdxl=False, is_plus=False, is_full=False, is_faceid=False, is_portrait_unnorm=False, is_kwai_kolors=False, encoder_hid_proj=None):
         super().__init__()
 
         self.clip_embeddings_dim = clip_embeddings_dim
@@ -69,7 +69,8 @@ class IPAdapter(nn.Module):
             self.image_proj_model = self.init_proj()
 
         self.image_proj_model.load_state_dict(ipadapter_model["image_proj"])
-        self.ip_layers = To_KV(ipadapter_model["ip_adapter"])
+
+        self.ip_layers = To_KV(ipadapter_model["ip_adapter"],  encoder_hid_proj_state_dict=encoder_hid_proj)
 
     def init_proj(self):
         image_proj_model = ImageProjModel(
@@ -173,13 +174,24 @@ class IPAdapter(nn.Module):
         return embeds
 
 class To_KV(nn.Module):
-    def __init__(self, state_dict):
+    def __init__(self, state_dict, encoder_hid_proj_state_dict=None):
         super().__init__()
 
+        if encoder_hid_proj_state_dict is not None:
+            encoder_hid_proj_linear = nn.Linear(encoder_hid_proj_state_dict["weight"].shape[1], encoder_hid_proj_state_dict["weight"].shape[0], bias=True)
+            encoder_hid_proj_linear.weight.data = encoder_hid_proj_state_dict["weight"]
+            encoder_hid_proj_linear.bias.data = encoder_hid_proj_state_dict["bias"]
+
         self.to_kvs = nn.ModuleDict()
         for key, value in state_dict.items():
-            self.to_kvs[key.replace(".weight", "").replace(".", "_")] = nn.Linear(value.shape[1], value.shape[0], bias=False)
-            self.to_kvs[key.replace(".weight", "").replace(".", "_")].weight.data = value
+            if encoder_hid_proj_state_dict is None:
+                self.to_kvs[key.replace(".weight", "").replace(".", "_")] = nn.Linear(value.shape[1], value.shape[0], bias=False)
+                self.to_kvs[key.replace(".weight", "").replace(".", "_")].weight.data = value
+            else:
+                linear = nn.Linear(value.shape[1], value.shape[0], bias=False)
+                linear.weight.data = value
+                self.to_kvs[key.replace(".weight", "").replace(".", "_")] = nn.Sequential(encoder_hid_proj_linear, linear)
+                
 
 def set_model_patch_replace(model, patch_kwargs, key):
     to = model.model_options["transformer_options"].copy()
@@ -237,6 +249,14 @@ def ipadapter_execute(model,
     is_sdxl = output_cross_attention_dim == 2048
     is_kwai_kolors = is_sdxl and "layers.0.0.to_out.weight" in ipadapter["image_proj"] and ipadapter["image_proj"]["layers.0.0.to_out.weight"].shape[0] == 2048
 
+    
+    # kwai-kolors faceid
+    is_kwai_kolors_faceid = "perceiver_resampler.layers.0.0.to_out.weight" in ipadapter["image_proj"] and ipadapter["image_proj"]["perceiver_resampler.layers.0.0.to_out.weight"].shape[0] == 4096
+    if is_kwai_kolors_faceid:
+        is_faceid = True 
+        is_kwai_kolors = True
+        is_faceidv2 = True
+
     if is_faceid and not insightface:
         raise Exception("insightface model is required for FaceID models")
 
@@ -246,6 +266,10 @@ def ipadapter_execute(model,
     cross_attention_dim = 1280 if (is_plus and is_sdxl and not is_faceid and not is_kwai_kolors) or is_portrait_unnorm else output_cross_attention_dim
     clip_extra_context_tokens = 16 if (is_plus and not is_faceid) or is_portrait or is_portrait_unnorm else 4
 
+    if is_kwai_kolors_faceid: 
+        clip_extra_context_tokens = 6
+        cross_attention_dim = 4096
+
     if image is not None and image.shape[1] != image.shape[2]:
         print("\033[33mINFO: the IPAdapter reference image is not a square, CLIPImageProcessor will resize and crop it at the center. If the main focus of the picture is not in the middle the result might not be what you are expecting.\033[0m")
 
@@ -309,6 +333,9 @@ def ipadapter_execute(model,
         face_cond_embeds = []
         image = []
 
+        face_image_size = 256 if is_sdxl else 224
+        if is_kwai_kolors_faceid:
+            face_image_size = 336
         for i in range(image_iface.shape[0]):
             for size in [(size, size) for size in range(640, 256, -64)]:
                 insightface.det_model.input_size = size # TODO: hacky but seems to be working
@@ -318,7 +345,7 @@ def ipadapter_execute(model,
                         face_cond_embeds.append(torch.from_numpy(face[0].normed_embedding).unsqueeze(0))
                     else:
                         face_cond_embeds.append(torch.from_numpy(face[0].embedding).unsqueeze(0))
-                    image.append(image_to_tensor(face_align.norm_crop(image_iface[i], landmark=face[0].kps, image_size=256 if is_sdxl else 224)))
+                    image.append(image_to_tensor(face_align.norm_crop(image_iface[i], landmark=face[0].kps, image_size=face_image_size)))
 
                     if 640 not in size:
                         print(f"\033[33mINFO: InsightFace detection resolution lowered to {size}.\033[0m")
@@ -408,6 +435,11 @@ def ipadapter_execute(model,
     if attn_mask is not None:
         attn_mask = attn_mask.to(device, dtype=dtype)
 
+    encoder_hid_proj = None
+    if is_kwai_kolors_faceid:
+        if hasattr(model.model, "diffusion_model") and hasattr(model.model.diffusion_model, "encoder_hid_proj"):
+            encoder_hid_proj = model.model.diffusion_model.encoder_hid_proj.state_dict()  
+     
     ipa = IPAdapter(
         ipadapter,
         cross_attention_dim=cross_attention_dim,
@@ -420,6 +452,7 @@ def ipadapter_execute(model,
         is_faceid=is_faceid,
         is_portrait_unnorm=is_portrait_unnorm,
         is_kwai_kolors=is_kwai_kolors,
+        encoder_hid_proj=encoder_hid_proj
     ).to(device, dtype=dtype)
 
     if is_faceid and is_plus:
@@ -443,6 +476,7 @@ def ipadapter_execute(model,
     sigma_start = model.get_model_object("model_sampling").percent_to_sigma(start_at)
     sigma_end = model.get_model_object("model_sampling").percent_to_sigma(end_at)
 
+
     patch_kwargs = {
         "ipadapter": ipa,
         "weight": weight,
@@ -1266,11 +1300,11 @@ class IPAdapterEncoder:
             mask = transforms(mask).squeeze(1)
             #mask = T.Resize((image.shape[1], image.shape[2]), interpolation=T.InterpolationMode.BICUBIC, antialias=True)(mask.unsqueeze(1)).squeeze(1)
 
-        img_cond_embeds = encode_image_masked(clip_vision, image, mask)
+        img_cond_embeds = encode_image_masked(clip_vision, image, mask, size=image_size)
 
         if is_plus:
             img_cond_embeds = img_cond_embeds.penultimate_hidden_states
-            img_uncond_embeds = encode_image_masked(clip_vision, torch.zeros([1, image_size, image_size, 3])).penultimate_hidden_states
+            img_uncond_embeds = encode_image_masked(clip_vision, torch.zeros([1, image_size, image_size, 3]), size=image_size).penultimate_hidden_states
         else:
             img_cond_embeds = img_cond_embeds.image_embeds
             img_uncond_embeds = torch.zeros_like(img_cond_embeds)
@@ -1798,92 +1832,4 @@ class IPAdapterCombineParams:
             ipadapter_params["end_at"] += params_5["end_at"]
 
         return (ipadapter_params, )
-
-"""
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
- Register
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-"""
-NODE_CLASS_MAPPINGS = {
-    # Main Apply Nodes
-    "IPAdapter": IPAdapterSimple,
-    "IPAdapterAdvanced": IPAdapterAdvanced,
-    "IPAdapterBatch": IPAdapterBatch,
-    "IPAdapterFaceID": IPAdapterFaceID,
-    "IPAAdapterFaceIDBatch": IPAAdapterFaceIDBatch,
-    "IPAdapterTiled": IPAdapterTiled,
-    "IPAdapterTiledBatch": IPAdapterTiledBatch,
-    "IPAdapterEmbeds": IPAdapterEmbeds,
-    "IPAdapterEmbedsBatch": IPAdapterEmbedsBatch,
-    "IPAdapterStyleComposition": IPAdapterStyleComposition,
-    "IPAdapterStyleCompositionBatch": IPAdapterStyleCompositionBatch,
-    "IPAdapterMS": IPAdapterMS,
-    "IPAdapterFromParams": IPAdapterFromParams,
-    "IPAdapterPreciseStyleTransfer": IPAdapterPreciseStyleTransfer,
-    "IPAdapterPreciseStyleTransferBatch": IPAdapterPreciseStyleTransferBatch,
-    "IPAdapterPreciseComposition": IPAdapterPreciseComposition,
-    "IPAdapterPreciseCompositionBatch": IPAdapterPreciseCompositionBatch,
-
-    # Loaders
-    "IPAdapterUnifiedLoader": IPAdapterUnifiedLoader,
-    "IPAdapterUnifiedLoaderFaceID": IPAdapterUnifiedLoaderFaceID,
-    "IPAdapterModelLoader": IPAdapterModelLoader,
-    "IPAdapterInsightFaceLoader": IPAdapterInsightFaceLoader,
-    "IPAdapterUnifiedLoaderCommunity": IPAdapterUnifiedLoaderCommunity,
-
-    # Helpers
-    "IPAdapterEncoder": IPAdapterEncoder,
-    "IPAdapterCombineEmbeds": IPAdapterCombineEmbeds,
-    "IPAdapterNoise": IPAdapterNoise,
-    "PrepImageForClipVision": PrepImageForClipVision,
-    "IPAdapterSaveEmbeds": IPAdapterSaveEmbeds,
-    "IPAdapterLoadEmbeds": IPAdapterLoadEmbeds,
-    "IPAdapterWeights": IPAdapterWeights,
-    "IPAdapterCombineWeights": IPAdapterCombineWeights,
-    "IPAdapterWeightsFromStrategy": IPAdapterWeightsFromStrategy,
-    "IPAdapterPromptScheduleFromWeightsStrategy": IPAdapterPromptScheduleFromWeightsStrategy,
-    "IPAdapterRegionalConditioning": IPAdapterRegionalConditioning,
-    "IPAdapterCombineParams": IPAdapterCombineParams,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    # Main Apply Nodes
-    "IPAdapter": "IPAdapter",
-    "IPAdapterAdvanced": "IPAdapter Advanced",
-    "IPAdapterBatch": "IPAdapter Batch (Adv.)",
-    "IPAdapterFaceID": "IPAdapter FaceID",
-    "IPAAdapterFaceIDBatch": "IPAdapter FaceID Batch",
-    "IPAdapterTiled": "IPAdapter Tiled",
-    "IPAdapterTiledBatch": "IPAdapter Tiled Batch",
-    "IPAdapterEmbeds": "IPAdapter Embeds",
-    "IPAdapterEmbedsBatch": "IPAdapter Embeds Batch",
-    "IPAdapterStyleComposition": "IPAdapter Style & Composition SDXL",
-    "IPAdapterStyleCompositionBatch": "IPAdapter Style & Composition Batch SDXL",
-    "IPAdapterMS": "IPAdapter Mad Scientist",
-    "IPAdapterFromParams": "IPAdapter from Params",
-    "IPAdapterPreciseStyleTransfer": "IPAdapter Precise Style Transfer",
-    "IPAdapterPreciseStyleTransferBatch": "IPAdapter Precise Style Transfer Batch",
-    "IPAdapterPreciseComposition": "IPAdapter Precise Composition",
-    "IPAdapterPreciseCompositionBatch": "IPAdapter Precise Composition Batch",
-
-    # Loaders
-    "IPAdapterUnifiedLoader": "IPAdapter Unified Loader",
-    "IPAdapterUnifiedLoaderFaceID": "IPAdapter Unified Loader FaceID",
-    "IPAdapterModelLoader": "IPAdapter Model Loader",
-    "IPAdapterInsightFaceLoader": "IPAdapter InsightFace Loader",
-    "IPAdapterUnifiedLoaderCommunity": "IPAdapter Unified Loader Community",
-
-    # Helpers
-    "IPAdapterEncoder": "IPAdapter Encoder",
-    "IPAdapterCombineEmbeds": "IPAdapter Combine Embeds",
-    "IPAdapterNoise": "IPAdapter Noise",
-    "PrepImageForClipVision": "Prep Image For ClipVision",
-    "IPAdapterSaveEmbeds": "IPAdapter Save Embeds",
-    "IPAdapterLoadEmbeds": "IPAdapter Load Embeds",
-    "IPAdapterWeights": "IPAdapter Weights",
-    "IPAdapterWeightsFromStrategy": "IPAdapter Weights From Strategy",
-    "IPAdapterPromptScheduleFromWeightsStrategy": "Prompt Schedule From Weights Strategy",
-    "IPAdapterCombineWeights": "IPAdapter Combine Weights",
-    "IPAdapterRegionalConditioning": "IPAdapter Regional Conditioning",
-    "IPAdapterCombineParams": "IPAdapter Combine Params",
-}
+ 

ComfyUI_IPAdapter_plus/utils.py

@@ -132,6 +132,10 @@ def ipadapter_model_loader(file):
         model = st_model
         del st_model
 
+    if "adapter_modules" in model.keys():
+        model["ip_adapter"] = model["adapter_modules"]
+        del model["adapter_modules"]
+
     if not "ip_adapter" in model.keys() or not model["ip_adapter"]:
         raise Exception("invalid IPAdapter model {}".format(file))
 
@@ -150,7 +154,7 @@ def insightface_loader(provider):
         raise Exception(e)
 
     path = os.path.join(folder_paths.models_dir, "insightface")
-    model = FaceAnalysis(name="buffalo_l", root=path, providers=[provider + 'ExecutionProvider',])
+    model = FaceAnalysis(name="antelopev2", root=path, providers=[provider + 'ExecutionProvider',])
     model.prepare(ctx_id=0, det_size=(640, 640))
     return model
 

README.md

@@ -1,12 +1,20 @@
-![image](./examples/workflow.png)
+![image](./examples/workflow_ipa.png)
 
 ## Recent changes 
+* [2024-08-02] 支持faceid,新增一些相关节点,工作流见examples/workflow_ipa_faceid.png
+* [2024-07-27] 新增MZ_KolorsControlNetLoader节点,用于加载可图ControlNet官方模型
+* [2024-07-26] 新增MZ_ApplySDXLSamplingSettings节点,用于V2版本重新回到SDXL的scheduler配置.
+  ![image](https://github.com/user-attachments/assets/8c3be6bf-4744-478f-8660-4842a4558a1f)
+
+* [2024-07-25] 修正sampling_settings,参数来自 [scheduler_config.json](https://huggingface.co/Kwai-Kolors/Kolors/blob/main/scheduler/scheduler_config.json),仅V2生效
+* [2024-07-21] 感谢来自yiwangsimple对Mac修复和测试的分支  https://github.com/yiwangsimple/ComfyUI-Kolors-MZ
+* [2024-07-21] 新增MZ_ChatGLM3TextEncodeAdvanceV2节点
+* [2024-07-18] IPA相关节点已在ComfyUI_IPAdapter_plus中支持
+* [2024-07-17] 新增支持IPAdapter_plus的加载器和高级应用节点 MZ_KolorsCLIPVisionLoader,MZ_IPAdapterModelLoaderKolors,MZ_IPAdapterAdvancedKolors
 * [2024-07-14] 删除自动兼容ControlNet, 新增MZ_KolorsControlNetPatch节点
   ![image](https://github.com/user-attachments/assets/73ae6447-c69d-4781-9c66-94e0029709ed)
 
-
-
-ComfyUI上Kolors的实现
+## ComfyUI上Kolors的实现
 
 参考自 https://github.com/kijai/ComfyUI-KwaiKolorsWrapper
 
@@ -14,18 +22,53 @@ ComfyUI上Kolors的实现
 
 工作流在examples/workflow.png中获取
 
-
+ 
+### UNET模型下载
 unet模型放置在 models/unet/ 文件夹下
-下载地址:
-https://huggingface.co/Kwai-Kolors/Kolors/blob/main/unet/diffusion_pytorch_model.fp16.safetensors
 
+模型主页: https://huggingface.co/Kwai-Kolors/Kolors
 
+下载地址: https://huggingface.co/Kwai-Kolors/Kolors/resolve/main/unet/diffusion_pytorch_model.fp16.safetensors
+ 
+
+### ChatGLM3模型下载
 chatglm3放置在 models/LLM/ 文件夹下
-下载地址:
-https://huggingface.co/Kijai/ChatGLM3-safetensors/blob/main/chatglm3-fp16.safetensors
+
+模型主页: https://huggingface.co/Kijai/ChatGLM3-safetensors
+
+下载地址: https://huggingface.co/Kijai/ChatGLM3-safetensors/resolve/main/chatglm3-fp16.safetensors
 
 
-Implementation of Kolors on ComfyUI
+### 官方IP-Adapter-Plus模型下载地址
+模型主页: https://huggingface.co/Kwai-Kolors/Kolors-IP-Adapter-Plus
+
+https://huggingface.co/Kwai-Kolors/Kolors-IP-Adapter-Plus/resolve/main/ip_adapter_plus_general.bin 下载至 models/ipadapter/
+
+https://huggingface.co/Kwai-Kolors/Kolors-IP-Adapter-Plus/resolve/main/image_encoder/pytorch_model.bin 下载至 models/clip_vision/
+
+### 官方ControlNet模型下载地址
+模型主页(Depth): https://huggingface.co/Kwai-Kolors/Kolors-ControlNet-Depth
+
+模型主页(Canny): https://huggingface.co/Kwai-Kolors/Kolors-ControlNet-Canny
+
+### Kolors-Inpainting模型下载地址
+模型主页: https://huggingface.co/Kwai-Kolors/Kolors-Inpainting
+
+https://huggingface.co/Kwai-Kolors/Kolors-Inpainting/resolve/main/unet/diffusion_pytorch_model.safetensors 下载至 models/unet/
+
+### 官方Kolors-IP-Adapter-FaceID-Plus模型下载地址
+模型主页: https://huggingface.co/Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus
+
+https://huggingface.co/Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus/resolve/main/ipa-faceid-plus.bin 下载至 models/ipadapter/
+
+https://huggingface.co/Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus/resolve/main/clip-vit-large-patch14-336/pytorch_model.bin 下载至 models/clip_vision/
+
+https://huggingface.co/MonsterMMORPG/tools/resolve/main/antelopev2.zip 下载并解压至 models/insightface/models/antelopev2/*.onnx
+
+
+
+
+## Implementation of Kolors on ComfyUI
 
 Reference from https://github.com/kijai/ComfyUI-KwaiKolorsWrapper
 
@@ -34,24 +77,74 @@ Using ComfyUI Native Sampling
 The workflow is obtained in examples/workflow.png
 
 
+### UNET model download
 The unet model is placed in the models/unet/ folder
+
+Model homepage: https://huggingface.co/Kwai-Kolors/Kolors
+
 Download link:
-https://huggingface.co/Kwai-Kolors/Kolors/blob/main/unet/diffusion_pytorch_model.fp16.safetensors
+https://huggingface.co/Kwai-Kolors/Kolors/resolve/main/unet/diffusion_pytorch_model.fp16.safetensors
 
 
+### ChatGLM3 model download
 The chatglm3 is placed in the models/LLM/ folder
+
+Model homepage: https://huggingface.co/Kijai/ChatGLM3-safetensors
+
 Download link:
-https://huggingface.co/Kijai/ChatGLM3-safetensors/tree/main
+https://huggingface.co/Kijai/ChatGLM3-safetensors/resolve/main/chatglm3-fp16.safetensors
 
 
+### Official IP-Adapter-Plus model download link
 
+Model homepage: https://huggingface.co/Kwai-Kolors/Kolors-IP-Adapter-Plus
+
+https://huggingface.co/Kwai-Kolors/Kolors-IP-Adapter-Plus/resolve/main/ip_adapter_plus_general.bin Download to models/ipadapter/
+
+https://huggingface.co/Kwai-Kolors/Kolors-IP-Adapter-Plus/resolve/main/image_encoder/pytorch_model.bin Download to models/clip_vision/
+
+### Official ControlNet model download link
+Model homepage(Depth): https://huggingface.co/Kwai-Kolors/Kolors-ControlNet-Depth
+
+Model homepage(Canny): https://huggingface.co/Kwai-Kolors/Kolors-ControlNet-Canny
+
+### Kolors-Inpainting model download link
+Model homepage: https://huggingface.co/Kwai-Kolors/Kolors-Inpainting
+
+https://huggingface.co/Kwai-Kolors/Kolors-Inpainting/resolve/main/unet/diffusion_pytorch_model.safetensors Download to models/unet/
+
+### Official Kolors-IP-Adapter-FaceID-Plus model download link
+Model homepage: https://huggingface.co/Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus
+
+https://huggingface.co/Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus/resolve/main/ipa-faceid-plus.bin Download to models/ipadapter/
+
+https://huggingface.co/Kwai-Kolors/Kolors-IP-Adapter-FaceID-Plus/resolve/main/clip-vit-large-patch14-336/pytorch_model.bin Download to models/clip_vision/
+
+https://huggingface.co/MonsterMMORPG/tools/resolve/main/antelopev2.zip Download and unzip to models/insightface/models/antelopev2/*.onnx
+
+## 使用ComfyUI-KwaiKolorsWrapper在相同种子下测试结果 (Testing results with the same seed using ComfyUI-KwaiKolorsWrapper)
+测试工作流来自examples/workflow_same_seed_test.png 
+
+The test workflow comes from examples/workflow_same_seed_test.png
+
+![image](./examples/workflow_same_seed_test.png)
 
 ## FAQ
+加载模型时出现的错误
++ 目前kolors有两个版本, 一种是unet类型采用unet加载器, 一种是放checkpoints文件夹用KolorsCheckpointLoaderSimple加载器, 如果你的模型来自huggingface的unet文件夹, 优先尝试使用unet加载器 (Currently, there are two versions of kolors, one is unet type using unet loader, and the other is placed in the checkpoints folder using KolorsCheckpointLoaderSimple loader. If your model comes from the huggingface unet folder, try to use the unet loader first)
++ 如果你无法确定模型类型, 那就都试一下 (If you are not sure about the model type, try both)
+
+Mac用户无法使用(Mac users cannot use)
++ Mac用户可移步至[ComfyUI-Kolors-MZ](https://github.com/yiwangsimple/ComfyUI-Kolors-MZ)  (Mac users can go to [ComfyUI-Kolors-MZ](https://github.com/yiwangsimple/ComfyUI-Kolors-MZ) )
+
+和IPAdapter有关的错误(Errors related to IPAdapter)
++ 确保ComfyUI本体和ComfyUI_IPAdapter_plus已经更新到最新版本(Make sure ComfyUI ontology and ComfyUI_IPAdapter_plus are updated to the latest version)
+
 name 'round_up' is not defined
 + 参考:https://github.com/THUDM/ChatGLM2-6B/issues/272#issuecomment-1632164243 , 使用 pip install cpm_kernels 或者 pip install -U cpm_kernels 更新 cpm_kernels
 
 module 'comfy.model_detection' has no attribute 'unet_prefix_from_state_dict'
-+ 更新ComfyUI本体到最新版本
++ 更新ComfyUI本体到最新版本(Update ComfyUI ontology to the latest version)
 
 RuntimeError: Only Tensors of floating point dtype can require gradients
 + 尝试使用fp16版本的模型: https://huggingface.co/Kijai/ChatGLM3-safetensors/blob/main/chatglm3-fp16.safetensors
@@ -60,6 +153,12 @@ Error occurred when executing MZ_ChatGLM3Loader: 'ChatGLMModel' object has no at
 + 检查ChatGLM3Loader节点选择的模型是否已经正确下载
 
 
+## Credits
+
+- [Kolors](https://github.com/Kwai-Kolors/Kolors)
+- [ComfyUI](https://github.com/comfyanonymous/ComfyUI) 
+- [ComfyUI_IPAdapter_plus](https://github.com/cubiq/ComfyUI_IPAdapter_plus)
+
 ## Star History
 
 <a href="https://star-history.com/#MinusZoneAI/ComfyUI-Kolors-MZ&Date">
@@ -72,7 +171,15 @@ Error occurred when executing MZ_ChatGLM3Loader: 'ChatGLMModel' object has no at
 
 
 ## Contact
-- 绿泡泡: minrszone
+- 微信Wechat: minrszone
 - Bilibili: [minus_zone](https://space.bilibili.com/5950992)
 - 小红书: [MinusZoneAI](https://www.xiaohongshu.com/user/profile/5f072e990000000001005472)
-- 爱发电: [MinusZoneAI](https://afdian.net/@MinusZoneAI)
+
+## Stargazers
+[![Stargazers repo roster for @MinusZoneAI/ComfyUI-Kolors-MZ](https://reporoster.com/stars/MinusZoneAI/ComfyUI-Kolors-MZ)](https://github.com/MinusZoneAI/ComfyUI-Kolors-MZ/stargazers)
+
+
+## Sponsorship
+<img src="https://github.com/user-attachments/assets/a7ef9684-4911-45b6-8071-a9b433dca6af"  width="200"/>
+
+

__init__.py

@@ -12,6 +12,8 @@ NODE_CLASS_MAPPINGS = {
 NODE_DISPLAY_NAME_MAPPINGS = {
 }
 
+MAX_RESOLUTION = 16384
+
 AUTHOR_NAME = "MinusZone"
 CATEGORY_NAME = f"{AUTHOR_NAME} - Kolors"
 folder_paths.add_model_folder_path(
@@ -24,10 +26,14 @@ class MZ_ChatGLM3Loader:
         # from .mz_kolors_utils import Utils
         # llm_dir = os.path.join(Utils.get_models_path(), "LLM")
         # print("llm_dir:", llm_dir)
-        # llm_models = Utils.listdir_models(llm_dir)
+        llm_models = folder_paths.get_filename_list("LLM")
+
+        # 筛选safetensors结尾的文件
+        llm_models = [
+            model for model in llm_models if model.endswith("safetensors")]
 
         return {"required": {
-            "chatglm3_checkpoint": (folder_paths.get_filename_list("LLM"),),
+            "chatglm3_checkpoint": (llm_models,),
         }}
 
     RETURN_TYPES = ("CHATGLM3MODEL",)
@@ -71,6 +77,117 @@ NODE_DISPLAY_NAME_MAPPINGS[
     "MZ_ChatGLM3_V2"] = f"{AUTHOR_NAME} - ChatGLM3TextEncodeV2"
 
 
+class MZ_ChatGLM3Embeds2Conditioning:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "kolors_embeds": ("KOLORS_EMBEDS", ),
+                "width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+                "height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+                "crop_w": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}),
+                "crop_h": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}),
+                "target_width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+                "target_height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+            }
+        }
+
+    RETURN_TYPES = ("CONDITIONING", "CONDITIONING",)
+    RETURN_NAMES = ("positive", "negative",)
+
+    FUNCTION = "embeds2conditioning"
+    CATEGORY = CATEGORY_NAME
+
+    def embeds2conditioning(self, **kwargs):
+        from . import mz_kolors_core
+        importlib.reload(mz_kolors_core)
+        return mz_kolors_core.MZ_ChatGLM3Embeds2Conditioning_call(kwargs)
+
+
+NODE_CLASS_MAPPINGS["MZ_ChatGLM3Embeds2Conditioning"] = MZ_ChatGLM3Embeds2Conditioning
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_ChatGLM3Embeds2Conditioning"] = f"{AUTHOR_NAME} - ChatGLM3Embeds2Conditioning"
+
+
+# for 2048 resolution
+class MZ_ChatGLM3TextEncodeAdvanceV2:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "chatglm3_model": ("CHATGLM3MODEL", ),
+                "text": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+                "width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+                "height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+                "crop_w": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}),
+                "crop_h": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}),
+                "target_width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+                "target_height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
+            }
+        }
+
+    RETURN_TYPES = ("CONDITIONING",)
+
+    FUNCTION = "encode"
+    CATEGORY = CATEGORY_NAME
+
+    def encode(self, **kwargs):
+        from . import mz_kolors_core
+        importlib.reload(mz_kolors_core)
+        return mz_kolors_core.MZ_ChatGLM3TextEncodeV2_call(kwargs)
+
+
+NODE_CLASS_MAPPINGS["MZ_ChatGLM3_Advance_V2"] = MZ_ChatGLM3TextEncodeAdvanceV2
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_ChatGLM3_Advance_V2"] = f"{AUTHOR_NAME} - ChatGLM3TextEncodeAdvanceV2"
+
+
+class MZ_KolorsCheckpointLoaderSimple():
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
+                             }}
+    RETURN_TYPES = ("MODEL", "VAE")
+    FUNCTION = "load_checkpoint"
+
+    CATEGORY = CATEGORY_NAME
+
+    def load_checkpoint(self, **kwargs):
+        from . import mz_kolors_core
+        importlib.reload(mz_kolors_core)
+        return mz_kolors_core.MZ_KolorsCheckpointLoaderSimple_call(kwargs)
+
+
+NODE_CLASS_MAPPINGS["MZ_KolorsCheckpointLoaderSimple"] = MZ_KolorsCheckpointLoaderSimple
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_KolorsCheckpointLoaderSimple"] = f"{AUTHOR_NAME} - KolorsCheckpointLoaderSimple"
+
+
+class MZ_KolorsControlNetLoader:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "control_net_name": (folder_paths.get_filename_list("controlnet"), ),
+            # "seed": ("INT", {"default": 0, "min": 0, "max": 1000000}),
+        }}
+
+    RETURN_TYPES = ("CONTROL_NET",)
+    RETURN_NAMES = ("ControlNet",)
+    FUNCTION = "load_controlnet"
+
+    CATEGORY = CATEGORY_NAME
+
+    def load_controlnet(self, **kwargs):
+        from . import mz_kolors_core
+        importlib.reload(mz_kolors_core)
+        return mz_kolors_core.MZ_KolorsControlNetLoader_call(kwargs)
+
+
+NODE_CLASS_MAPPINGS["MZ_KolorsControlNetLoader"] = MZ_KolorsControlNetLoader
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_KolorsControlNetLoader"] = f"{AUTHOR_NAME} - KolorsControlNetLoader"
+
+
 class MZ_KolorsUNETLoaderV2():
     @classmethod
     def INPUT_TYPES(s):
@@ -131,7 +248,7 @@ class MZ_KolorsCLIPVisionLoader:
     RETURN_TYPES = ("CLIP_VISION",)
     FUNCTION = "load_clip"
 
-    CATEGORY = CATEGORY_NAME
+    CATEGORY = CATEGORY_NAME + "/Legacy"
 
     def load_clip(self, **kwargs):
         from . import mz_kolors_core
@@ -140,21 +257,83 @@ class MZ_KolorsCLIPVisionLoader:
 
 
 NODE_CLASS_MAPPINGS["MZ_KolorsCLIPVisionLoader"] = MZ_KolorsCLIPVisionLoader
-NODE_DISPLAY_NAME_MAPPINGS["MZ_KolorsCLIPVisionLoader"] = f"{AUTHOR_NAME} - KolorsCLIPVisionLoader"
+NODE_DISPLAY_NAME_MAPPINGS["MZ_KolorsCLIPVisionLoader"] = f"{AUTHOR_NAME} - KolorsCLIPVisionLoader - Legacy"
 
 
-from .ComfyUI_IPAdapter_plus.IPAdapterPlus import IPAdapterAdvanced, IPAdapterModelLoader
+class MZ_ApplySDXLSamplingSettings():
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL", ),
+            }
+        }
 
-IPAdapterModelLoader.CATEGORY = CATEGORY_NAME + "/IPAdapter_plus"
+    RETURN_TYPES = ("MODEL", )
+
+    FUNCTION = "apply_sampling_settings"
+    CATEGORY = CATEGORY_NAME
+
+    def apply_sampling_settings(self, **kwargs):
+        from . import mz_kolors_core
+        importlib.reload(mz_kolors_core)
+        return mz_kolors_core.MZ_ApplySDXLSamplingSettings_call(kwargs)
+
+
+NODE_CLASS_MAPPINGS["MZ_ApplySDXLSamplingSettings"] = MZ_ApplySDXLSamplingSettings
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_ApplySDXLSamplingSettings"] = f"{AUTHOR_NAME} - ApplySDXLSamplingSettings"
+
+
+class MZ_ApplyCUDAGenerator():
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL", ),
+            }
+        }
+
+    RETURN_TYPES = ("MODEL", )
+
+    FUNCTION = "apply_cuda_generator"
+    CATEGORY = CATEGORY_NAME
+
+    def apply_cuda_generator(self, **kwargs):
+        from . import mz_kolors_core
+        importlib.reload(mz_kolors_core)
+        return mz_kolors_core.MZ_ApplyCUDAGenerator_call(kwargs)
+
+
+NODE_CLASS_MAPPINGS["MZ_ApplyCUDAGenerator"] = MZ_ApplyCUDAGenerator
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_ApplyCUDAGenerator"] = f"{AUTHOR_NAME} - ApplyCUDAGenerator"
+
+
+from .ComfyUI_IPAdapter_plus.IPAdapterPlus import IPAdapterAdvanced, IPAdapterModelLoader, IPAdapterInsightFaceLoader, IPAdapterFaceID
+
+IPAdapterModelLoader.CATEGORY = CATEGORY_NAME + "/IPAdapter"
 NODE_CLASS_MAPPINGS["MZ_IPAdapterModelLoaderKolors"] = IPAdapterModelLoader
 NODE_DISPLAY_NAME_MAPPINGS[
     "MZ_IPAdapterModelLoaderKolors"] = f"IPAdapterModelLoader(kolors)"
 
-IPAdapterAdvanced.CATEGORY = CATEGORY_NAME + "/IPAdapter_plus"
+IPAdapterAdvanced.CATEGORY = CATEGORY_NAME + "/IPAdapter"
 NODE_CLASS_MAPPINGS["MZ_IPAdapterAdvancedKolors"] = IPAdapterAdvanced
 NODE_DISPLAY_NAME_MAPPINGS[
     "MZ_IPAdapterAdvancedKolors"] = f"IPAdapterAdvanced(kolors)"
 
+IPAdapterInsightFaceLoader.CATEGORY = CATEGORY_NAME + "/IPAdapter"
+NODE_CLASS_MAPPINGS["MZ_IPAdapterInsightFaceLoader"] = IPAdapterInsightFaceLoader
+
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_IPAdapterInsightFaceLoader"] = f"IPAdapterInsightFaceLoader(kolors)"
+
+IPAdapterFaceID.CATEGORY = CATEGORY_NAME + "/IPAdapter"
+NODE_CLASS_MAPPINGS["MZ_IPAdapterFaceID"] = IPAdapterFaceID
+
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_IPAdapterFaceID"] = f"IPAdapterFaceID(kolors)"
+
 from . import mz_kolors_legacy
 NODE_CLASS_MAPPINGS.update(mz_kolors_legacy.NODE_CLASS_MAPPINGS)
 NODE_DISPLAY_NAME_MAPPINGS.update(mz_kolors_legacy.NODE_DISPLAY_NAME_MAPPINGS)

chatglm3/quantization.py

@@ -15,6 +15,26 @@ logger = logging.get_logger(__name__)
 try:
     from cpm_kernels.kernels.base import LazyKernelCModule, KernelFunction, round_up
 
+    import cpm_kernels.library.base
+
+    original_windows_find_lib = cpm_kernels.library.base.windows_find_lib
+    
+    def windows_find_lib(name):
+        result = original_windows_find_lib(name)
+        if result is not None:
+            return result
+        import torch
+        import os
+        torch_dir = os.path.dirname(torch.__file__)
+        torch_lib_dir = os.path.join(torch_dir, "lib")
+
+        for name in os.listdir(torch_lib_dir):
+            if name.startswith(name) and name.lower().endswith(".dll"):
+                return os.path.join(torch_lib_dir, name)
+            
+        return None
+    cpm_kernels.library.base.windows_find_lib = windows_find_lib
+
     class Kernel:
         def __init__(self, code: bytes, function_names: List[str]):
             self.code = code
@@ -137,8 +157,11 @@ class QuantizedLinear(torch.nn.Module):
             self.weight = torch.round(weight / self.weight_scale[:, None]).to(torch.int8)
             if weight_bit_width == 4:
                 self.weight = compress_int4_weight(self.weight)
-
-        self.weight = Parameter(self.weight.to(device), requires_grad=False)
+ 
+        try:
+            self.weight = Parameter(self.weight.to(device), requires_grad=False)
+        except:
+            self.weight.to(device, dtype=self.weight.dtype)
         self.weight_scale = Parameter(self.weight_scale.to(device), requires_grad=False)
         self.bias = Parameter(bias.to(device), requires_grad=False) if bias is not None else None
 

hook_comfyui_kolors_v2.py

@@ -6,13 +6,15 @@ 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
 
 
 from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel
-
+from comfy.cldm.cldm import ControlNet
+import torch.amp
 
 # try:
 #     import comfy.samplers as samplers
@@ -50,35 +52,53 @@ from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel
 #     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):
     def __init__(self, model_config, model_type=ModelType.EPS, device=None):
+        model_config.sampling_settings["beta_schedule"] = "linear"
+        model_config.sampling_settings["linear_start"] = 0.00085
+        model_config.sampling_settings["linear_end"] = 0.014
+        model_config.sampling_settings["timesteps"] = 1100
+        model_type = ModelType.EPS
         model_base.BaseModel.__init__(
             self, model_config, model_type, device=device, unet_model=KolorsUNetModel)
         self.embedder = Timestep(256)
         self.noise_augmentor = CLIPEmbeddingNoiseAugmentation(
-            **{"noise_schedule_config": {"timesteps": 1000, "beta_schedule": "squaredcos_cap_v2"}, "timestep_dim": 1280})
+            **{"noise_schedule_config": {"timesteps": 1100, "beta_schedule": "linear", "linear_start": 0.00085, "linear_end": 0.014}, "timestep_dim": 1280})
 
     def encode_adm(self, **kwargs):
         clip_pooled = sdxl_pooled(kwargs, self.noise_augmentor)
@@ -160,7 +180,38 @@ def kolors_unet_config_from_diffusers_unet(state_dict, dtype=None):
               'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10],
               'use_temporal_attention': False, 'use_temporal_resblock': False}
 
-    supported_models = [Kolors]
+    Kolors_inpaint = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
+                      'num_classes': 'sequential', 'adm_in_channels': 5632, 'dtype': dtype, 'in_channels': 9, 'model_channels': 320,
+                      'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 10, 10], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 10,
+                      'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10],
+                      'use_temporal_attention': False, 'use_temporal_resblock': False}
+
+    Kolors_ip2p = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
+                   'num_classes': 'sequential', 'adm_in_channels': 5632, 'dtype': dtype, 'in_channels': 8, 'model_channels': 320,
+                   'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 10, 10], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 10,
+                   'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10],
+                   'use_temporal_attention': False, 'use_temporal_resblock': False}
+
+    SDXL = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
+            'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,
+            'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 10, 10], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 10,
+            'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10],
+            'use_temporal_attention': False, 'use_temporal_resblock': False}
+
+    SDXL_mid_cnet = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
+                     'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,
+                     'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 0, 0, 1, 1], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 1,
+                     'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 0, 0, 0, 1, 1, 1],
+                     'use_temporal_attention': False, 'use_temporal_resblock': False}
+
+    SDXL_small_cnet = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
+                       'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,
+                       'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 0, 0, 0, 0], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 0,
+                       'use_linear_in_transformer': True, 'num_head_channels': 64, 'context_dim': 1, 'transformer_depth_output': [0, 0, 0, 0, 0, 0, 0, 0, 0],
+                       'use_temporal_attention': False, 'use_temporal_resblock': False}
+
+    supported_models = [Kolors, Kolors_inpaint,
+                        Kolors_ip2p, SDXL, SDXL_mid_cnet, SDXL_small_cnet]
 
     for unet_config in supported_models:
         matches = True
@@ -215,9 +266,32 @@ def load_clipvision_336_from_sd(sd, prefix="", convert_keys=False):
     return clip
 
 
+class KolorsControlNet(ControlNet):
+    def __init__(self, *args, **kwargs):
+        adm_in_channels = kwargs["adm_in_channels"]
+        if adm_in_channels == 2816:
+            # 异常: 该加载器不支持SDXL类型, 请使用ControlNet加载器+KolorsControlNetPatch节点
+            raise Exception(
+                "This loader does not support SDXL type, please use ControlNet loader + KolorsControlNetPatch node")
+
+        super().__init__(*args, **kwargs)
+        self.encoder_hid_proj = nn.Linear(
+            4096, 2048, bias=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"])
+
+            result = super().forward(*args, **kwargs)
+            return result
+
+
 class apply_kolors:
     def __enter__(self):
         import comfy.ldm.modules.diffusionmodules.openaimodel
+        import comfy.cldm.cldm
         import comfy.utils
         import comfy.clip_vision
 
@@ -239,8 +313,12 @@ class apply_kolors:
         self.original_supported_models = comfy.supported_models.models
         comfy.supported_models.models = [KolorsSupported]
 
+        self.original_controlnet = comfy.cldm.cldm.ControlNet
+        comfy.cldm.cldm.ControlNet = KolorsControlNet
+
     def __exit__(self, type, value, traceback):
         import comfy.ldm.modules.diffusionmodules.openaimodel
+        import comfy.cldm.cldm
         import comfy.utils
         comfy.utils.UNET_MAP_BASIC = self.original_UNET_MAP_BASIC
 
@@ -251,3 +329,5 @@ class apply_kolors:
 
         import comfy.clip_vision
         comfy.clip_vision.load_clipvision_from_sd = self.original_load_clipvision_from_sd
+
+        comfy.cldm.cldm.ControlNet = self.original_controlnet

mz_kolors_core.py

@@ -5,11 +5,16 @@ import json
 import os
 import random
 import re
+import subprocess
+import sys
 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):
@@ -78,8 +83,6 @@ def MZ_ChatGLM3Loader_call(args):
     from .chatglm3.modeling_chatglm import ChatGLMModel
     from .chatglm3.tokenization_chatglm import ChatGLMTokenizer
 
-    from .mz_kolors_utils import Utils
-
     offload_device = mm.unet_offload_device()
 
     text_encoder_config = os.path.join(
@@ -91,6 +94,7 @@ def MZ_ChatGLM3Loader_call(args):
 
     from comfy.utils import load_torch_file
     from contextlib import nullcontext
+    is_accelerate_available = False
     try:
         from accelerate import init_empty_weights
         from accelerate.utils import set_module_tensor_to_device
@@ -104,6 +108,13 @@ def MZ_ChatGLM3Loader_call(args):
             print("torch version:", torch.__version__)
             text_encoder = ChatGLMModel(text_encoder_config).eval()
             if '4bit' in chatglm3_checkpoint:
+                try:
+                    import cpm_kernels
+                except ImportError:
+                    print("Installing cpm_kernels...")
+                    subprocess.run(
+                        [sys.executable, "-m", "pip", "install", "cpm_kernels"], check=True)
+                    pass
                 text_encoder.quantize(4)
             elif '8bit' in chatglm3_checkpoint:
                 text_encoder.quantize(8)
@@ -113,7 +124,8 @@ def MZ_ChatGLM3Loader_call(args):
             set_module_tensor_to_device(
                 text_encoder, key, device=offload_device, value=text_encoder_sd[key])
     else:
-        text_encoder.load_state_dict()
+        print("WARNING: Accelerate not available, use load_state_dict load model")
+        text_encoder.load_state_dict(text_encoder_sd)
 
     tokenizer_path = os.path.join(
         os.path.dirname(__file__), 'configs', "tokenizer")
@@ -129,12 +141,60 @@ def MZ_ChatGLM3TextEncodeV2_call(args):
         chatglm3_model,
         text,
     )
+    extra_kwargs = {
+        "pooled_output": pooled_output,
+    }
+    extra_cond_keys = [
+        "width",
+        "height",
+        "crop_w",
+        "crop_h",
+        "target_width",
+        "target_height"
+    ]
+    for key, value in args.items():
+        if key in extra_cond_keys:
+            extra_kwargs[key] = value
     return ([[
         prompt_embeds,
-        {"pooled_output": pooled_output},
+        # {"pooled_output": pooled_output},
+        extra_kwargs
     ]], )
 
 
+def MZ_ChatGLM3Embeds2Conditioning_call(args):
+    kolors_embeds = args.get("kolors_embeds")
+
+    # kolors_embeds = {
+    #     'prompt_embeds': prompt_embeds,
+    #     'negative_prompt_embeds': negative_prompt_embeds,
+    #     'pooled_prompt_embeds': text_proj,
+    #     'negative_pooled_prompt_embeds': negative_text_proj
+    # }
+
+    positive = [[
+        kolors_embeds['prompt_embeds'],
+        {
+            "pooled_output": kolors_embeds['pooled_prompt_embeds'],
+            "width": args.get("width"),
+            "height": args.get("height"),
+            "crop_w": args.get("crop_w"),
+            "crop_h": args.get("crop_h"),
+            "target_width": args.get("target_width"),
+            "target_height": args.get("target_height")
+        }
+    ]]
+
+    negative = [[
+        kolors_embeds['negative_prompt_embeds'],
+        {
+            "pooled_output": kolors_embeds['negative_pooled_prompt_embeds'],
+        }
+    ]]
+
+    return (positive, negative)
+
+
 def MZ_KolorsUNETLoaderV2_call(kwargs):
 
     from . import hook_comfyui_kolors_v2
@@ -153,14 +213,54 @@ def MZ_KolorsUNETLoaderV2_call(kwargs):
         return (model, )
 
 
+def MZ_KolorsCheckpointLoaderSimple_call(kwargs):
+    checkpoint_name = kwargs.get("ckpt_name")
+
+    ckpt_path = folder_paths.get_full_path("checkpoints", checkpoint_name)
+
+    from . import hook_comfyui_kolors_v2
+    import comfy.sd
+
+    with hook_comfyui_kolors_v2.apply_kolors():
+        out = comfy.sd.load_checkpoint_guess_config(
+            ckpt_path, output_vae=True, output_clip=False, embedding_directory=folder_paths.get_folder_paths("embeddings"))
+
+        unet, _, vae = out[:3]
+        return (unet, vae)
+
+
 from comfy.cldm.cldm import ControlNet
 from comfy.controlnet import ControlLora
 
 
-def MZ_KolorsControlNetPatch_call(kwargs):
+def MZ_KolorsControlNetLoader_call(kwargs):
+    control_net_name = kwargs.get("control_net_name")
+    controlnet_path = folder_paths.get_full_path(
+        "controlnet", control_net_name)
+
+    from torch import nn
     from . import hook_comfyui_kolors_v2
+    import comfy.controlnet
+
+    with hook_comfyui_kolors_v2.apply_kolors():
+        control_net = comfy.controlnet.load_controlnet(controlnet_path)
+        return (control_net, )
+
+
+def MZ_KolorsControlNetPatch_call(kwargs):
+    import copy
+    from . import hook_comfyui_kolors_v2
+    import comfy.model_management
+    import comfy.model_patcher
+
     model = kwargs.get("model")
     control_net = kwargs.get("control_net")
+
+    if hasattr(control_net, "control_model") and hasattr(control_net.control_model, "encoder_hid_proj"):
+        return (control_net,)
+
+    control_net = copy.deepcopy(control_net)
+
     import comfy.controlnet
     if isinstance(control_net, ControlLora):
         del_keys = []
@@ -172,26 +272,33 @@ def MZ_KolorsControlNetPatch_call(kwargs):
             control_net.control_weights.pop(k)
 
         super_pre_run = ControlLora.pre_run
-        super_copy = ControlLora.copy
-
         super_forward = ControlNet.forward
 
         def KolorsControlNet_forward(self, x, hint, timesteps, context, **kwargs):
-            with torch.cuda.amp.autocast(enabled=True):
-                context = model.model.diffusion_model.encoder_hid_proj(context)
+            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)
 
         def KolorsControlLora_pre_run(self, *args, **kwargs):
             result = super_pre_run(self, *args, **kwargs)
 
             if hasattr(self, "control_model"):
+                if hasattr(self.control_model, "encoder_hid_proj"):
+                    return result
+
+                setattr(self.control_model, "encoder_hid_proj",
+                        model.model.diffusion_model.encoder_hid_proj)
+
                 self.control_model.forward = MethodType(
                     KolorsControlNet_forward, self.control_model)
+
             return result
 
         control_net.pre_run = MethodType(
             KolorsControlLora_pre_run, control_net)
 
+        super_copy = ControlLora.copy
+
         def KolorsControlLora_copy(self, *args, **kwargs):
             c = super_copy(self, *args, **kwargs)
             c.pre_run = MethodType(
@@ -201,18 +308,23 @@ def MZ_KolorsControlNetPatch_call(kwargs):
         control_net.copy = MethodType(
             KolorsControlLora_copy, control_net)
 
+        control_net = copy.deepcopy(control_net)
+
     elif isinstance(control_net, comfy.controlnet.ControlNet):
         model_label_emb = model.model.diffusion_model.label_emb
 
         control_net.control_model.label_emb = model_label_emb
+        setattr(control_net.control_model, "encoder_hid_proj",
+                model.model.diffusion_model.encoder_hid_proj)
 
-        control_net.control_model_wrapped.model.label_emb = model_label_emb
+        control_net.control_model_wrapped = comfy.model_patcher.ModelPatcher(
+            control_net.control_model, load_device=control_net.load_device, offload_device=comfy.model_management.unet_offload_device())
 
         super_forward = ControlNet.forward
 
         def KolorsControlNet_forward(self, x, hint, timesteps, context, **kwargs):
-            with torch.cuda.amp.autocast(enabled=True):
-                context = model.model.diffusion_model.encoder_hid_proj(context)
+            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)
 
         control_net.control_model.forward = MethodType(
@@ -233,3 +345,53 @@ def MZ_KolorsCLIPVisionLoader_call(kwargs):
     with hook_comfyui_kolors_v2.apply_kolors():
         clip_vision = comfy.clip_vision.load(clip_path)
         return (clip_vision,)
+
+
+def MZ_ApplySDXLSamplingSettings_call(kwargs):
+    model = kwargs.get("model").clone()
+
+    import comfy.model_sampling
+    sampling_base = comfy.model_sampling.ModelSamplingDiscrete
+    sampling_type = comfy.model_sampling.EPS
+
+    class SDXLSampling(sampling_base, sampling_type):
+        pass
+
+    model.model.model_config.sampling_settings["beta_schedule"] = "linear"
+    model.model.model_config.sampling_settings["linear_start"] = 0.00085
+    model.model.model_config.sampling_settings["linear_end"] = 0.012
+    model.model.model_config.sampling_settings["timesteps"] = 1000
+
+    model_sampling = SDXLSampling(model.model.model_config)
+
+    model.add_object_patch("model_sampling", model_sampling)
+
+    return (model,)
+
+
+def MZ_ApplyCUDAGenerator_call(kwargs):
+    model = kwargs.get("model")
+
+    def prepare_noise(latent_image, seed, noise_inds=None):
+        """
+        creates random noise given a latent image and a seed.
+        optional arg skip can be used to skip and discard x number of noise generations for a given seed
+        """
+        generator = torch.Generator(device="cuda").manual_seed(seed)
+        if noise_inds is None:
+            return torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cuda")
+
+        unique_inds, inverse = np.unique(noise_inds, return_inverse=True)
+        noises = []
+        for i in range(unique_inds[-1] + 1):
+            noise = torch.randn([1] + list(latent_image.size())[1:], dtype=latent_image.dtype,
+                                layout=latent_image.layout, generator=generator, device="cuda")
+            if i in unique_inds:
+                noises.append(noise)
+        noises = [noises[i] for i in inverse]
+        noises = torch.cat(noises, axis=0)
+        return noises
+
+    import comfy.sample
+    comfy.sample.prepare_noise = prepare_noise
+    return (model,)

pyproject.toml

@@ -2,7 +2,7 @@
 name = "comfyui-kolors-mz"
 description = "Implementation of Kolors on ComfyUI\nReference from [a/https://github.com/kijai/ComfyUI-KwaiKolorsWrapper](https://github.com/kijai/ComfyUI-KwaiKolorsWrapper)\nUsing ComfyUI Native Sampling"
 version = "1.0.0"
-license = "LICENSE"
+license = { file = "GPL-3.0 license" }
 
 [project.urls]
 Repository = "https://github.com/MinusZoneAI/ComfyUI-Kolors-MZ"