Merge branch 'main' into 20240713

一些优化

.github/workflows/publish.yml

@@ -0,0 +1,24 @@
+name: Publish to Comfy registry
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+      - master
+    paths:
+      - "pyproject.toml"
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    # if this is a forked repository. Skipping the workflow.
+    if: github.event.repository.fork == false 
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+      - name: Publish Custom Node
+        uses: Comfy-Org/publish-node-action@main
+        with:
+          ## Add your own personal access token to your Github Repository secrets and reference it here.
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.gitignore

@@ -162,4 +162,3 @@ cython_debug/
 #.idea/
 exclude.txt
 copytoww.bat
-train_tools

README.md

@@ -1,4 +1,10 @@
-![image](https://github.com/MinusZoneAI/ComfyUI-Kolors-MZ/blob/main/examples/workflow.png)
+![image](./examples/workflow.png)
+
+## Recent changes 
+* [2024-07-14] 删除自动兼容ControlNet, 新增MZ_KolorsControlNetPatch节点
+  ![image](https://github.com/user-attachments/assets/73ae6447-c69d-4781-9c66-94e0029709ed)
+
+
 
 ComfyUI上Kolors的实现
 
@@ -41,13 +47,18 @@ https://huggingface.co/Kijai/ChatGLM3-safetensors/tree/main
 
 
 ## FAQ
-Error occurred when executing MZ_ChatGLM3Loader: 'ChatGLMModel' object has no attribute 'transformer'
+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
 
-检查ChatGLM3Loader节点选择的模型是否已经正确下载
+module 'comfy.model_detection' has no attribute 'unet_prefix_from_state_dict'
++ 更新ComfyUI本体到最新版本
 
 RuntimeError: Only Tensors of floating point dtype can require gradients
++ 尝试使用fp16版本的模型: https://huggingface.co/Kijai/ChatGLM3-safetensors/blob/main/chatglm3-fp16.safetensors
+
+Error occurred when executing MZ_ChatGLM3Loader: 'ChatGLMModel' object has no attribute 'transformer'
++ 检查ChatGLM3Loader节点选择的模型是否已经正确下载
 
-尝试使用fp16版本的模型: https://huggingface.co/Kijai/ChatGLM3-safetensors/blob/main/chatglm3-fp16.safetensors
 
 ## Star History
 

__init__.py

@@ -14,43 +14,20 @@ NODE_DISPLAY_NAME_MAPPINGS = {
 
 AUTHOR_NAME = "MinusZone"
 CATEGORY_NAME = f"{AUTHOR_NAME} - Kolors"
-
-
-class MZ_FakeCond:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "seed": ("INT", {"default": 0}),
-            }
-        }
-
-    RETURN_TYPES = ("CONDITIONING", )
-    RETURN_NAMES = ("prompt", )
-    FUNCTION = "encode"
-    CATEGORY = CATEGORY_NAME
-
-    def encode(self, **kwargs):
-        from . import mz_kolors_core
-        importlib.reload(mz_kolors_core)
-        return mz_kolors_core.MZ_FakeCond_call(kwargs)
-
-
-NODE_CLASS_MAPPINGS["MZ_FakeCond"] = MZ_FakeCond
-NODE_DISPLAY_NAME_MAPPINGS[
-    "MZ_FakeCond"] = f"{AUTHOR_NAME} - FakeCond"
+folder_paths.add_model_folder_path(
+    "LLM", os.path.join(folder_paths.models_dir, "LLM"))
 
 
 class MZ_ChatGLM3Loader:
     @classmethod
     def INPUT_TYPES(s):
-        from .mz_kolors_utils import Utils
-        llm_dir = os.path.join(Utils.get_models_path(), "LLM")
+        # 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 = Utils.listdir_models(llm_dir)
 
         return {"required": {
-            "chatglm3_checkpoint": (llm_models,),
+            "chatglm3_checkpoint": (folder_paths.get_filename_list("LLM"),),
         }}
 
     RETURN_TYPES = ("CHATGLM3MODEL",)
@@ -68,14 +45,13 @@ NODE_CLASS_MAPPINGS["MZ_ChatGLM3Loader"] = MZ_ChatGLM3Loader
 NODE_DISPLAY_NAME_MAPPINGS["MZ_ChatGLM3Loader"] = f"{AUTHOR_NAME} - ChatGLM3Loader"
 
 
-class MZ_ChatGLM3TextEncode:
+class MZ_ChatGLM3TextEncodeV2:
     @classmethod
     def INPUT_TYPES(s):
         return {
             "required": {
                 "chatglm3_model": ("CHATGLM3MODEL", ),
                 "text": ("STRING", {"multiline": True, "dynamicPrompts": True}),
-                "hid_proj": ("TorchLinear", ),
             }
         }
 
@@ -87,24 +63,24 @@ class MZ_ChatGLM3TextEncode:
     def encode(self, **kwargs):
         from . import mz_kolors_core
         importlib.reload(mz_kolors_core)
-        return mz_kolors_core.MZ_ChatGLM3TextEncode_call(kwargs)
+        return mz_kolors_core.MZ_ChatGLM3TextEncodeV2_call(kwargs)
 
 
-NODE_CLASS_MAPPINGS["MZ_ChatGLM3"] = MZ_ChatGLM3TextEncode
+NODE_CLASS_MAPPINGS["MZ_ChatGLM3_V2"] = MZ_ChatGLM3TextEncodeV2
 NODE_DISPLAY_NAME_MAPPINGS[
-    "MZ_ChatGLM3"] = f"{AUTHOR_NAME} - ChatGLM3TextEncode"
+    "MZ_ChatGLM3_V2"] = f"{AUTHOR_NAME} - ChatGLM3TextEncodeV2"
 
 
-class MZ_KolorsUNETLoader():
+class MZ_KolorsUNETLoaderV2():
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
                 "unet_name": (folder_paths.get_filename_list("unet"), ),
                 }}
 
-    RETURN_TYPES = ("MODEL", "TorchLinear")
+    RETURN_TYPES = ("MODEL",)
 
-    RETURN_NAMES = ("model", "hid_proj")
+    RETURN_NAMES = ("model",)
 
     FUNCTION = "load_unet"
 
@@ -113,9 +89,40 @@ class MZ_KolorsUNETLoader():
     def load_unet(self, **kwargs):
         from . import mz_kolors_core
         importlib.reload(mz_kolors_core)
-        return mz_kolors_core.MZ_KolorsUNETLoader_call(kwargs)
+        return mz_kolors_core.MZ_KolorsUNETLoaderV2_call(kwargs)
 
 
-NODE_CLASS_MAPPINGS["MZ_KolorsUNETLoader"] = MZ_KolorsUNETLoader
+NODE_CLASS_MAPPINGS["MZ_KolorsUNETLoaderV2"] = MZ_KolorsUNETLoaderV2
 NODE_DISPLAY_NAME_MAPPINGS[
-    "MZ_KolorsUNETLoader"] = f"{AUTHOR_NAME} - Kolors UNET Loader"
+    "MZ_KolorsUNETLoaderV2"] = f"{AUTHOR_NAME} - KolorsUNETLoaderV2"
+
+
+class MZ_KolorsControlNetPatch:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "control_net": ("CONTROL_NET", ),
+                "model": ("MODEL", ),
+            }
+        }
+
+    RETURN_TYPES = ("CONTROL_NET",)
+
+    FUNCTION = "start"
+    CATEGORY = CATEGORY_NAME
+
+    def start(self, **kwargs):
+        from . import mz_kolors_core
+        importlib.reload(mz_kolors_core)
+        return mz_kolors_core.MZ_KolorsControlNetPatch_call(kwargs)
+
+
+NODE_CLASS_MAPPINGS["MZ_KolorsControlNetPatch"] = MZ_KolorsControlNetPatch
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_KolorsControlNetPatch"] = f"{AUTHOR_NAME} - KolorsControlNetPatch"
+
+
+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)

hook_comfyui_kolors_v1.py

@@ -12,25 +12,6 @@ class Kolors(comfy.supported_models.SDXL):
         "adm_in_channels": 5632,
         "use_temporal_attention": False,
     }
-    
-    def process_clip_state_dict(self, state_dict):
-        raise NotImplementedError("Kolors does not support clip state dict processing")
-        return state_dict
-
-    
-    def process_clip_state_dict_for_saving(self, state_dict):
-        raise NotImplementedError("Kolors does not support clip state dict processing")
-        return state_dict
-    
-    
-    def clip_target(self):
-        raise NotImplementedError("Kolors does not support clip target")
-        # return supported_models_base.ClipTarget(sdxl_clip.SDXLTokenizer, sdxl_clip.SDXLClipModel)
-
-
-
-if Kolors not in comfy.supported_models.models:
-    comfy.supported_models.models += [Kolors]
 
 
 def kolors_unet_config_from_diffusers_unet(state_dict, dtype=None):
@@ -90,8 +71,15 @@ def kolors_unet_config_from_diffusers_unet(state_dict, dtype=None):
 
 class apply_kolors:
     def __enter__(self):
+        import comfy.supported_models
+        self.old_supported_models = comfy.supported_models.models
+        comfy.supported_models.models = [Kolors]
+
         self.old_unet_config_from_diffusers_unet = model_detection.unet_config_from_diffusers_unet
         model_detection.unet_config_from_diffusers_unet = kolors_unet_config_from_diffusers_unet
 
     def __exit__(self, type, value, traceback):
         model_detection.unet_config_from_diffusers_unet = self.old_unet_config_from_diffusers_unet
+
+        import comfy.supported_models
+        comfy.supported_models.models = self.old_supported_models

hook_comfyui_kolors_v2.py

@@ -0,0 +1,205 @@
+import warnings
+from comfy.model_detection import *
+import comfy.model_detection as model_detection
+import comfy.supported_models
+import comfy.utils
+
+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
+
+
+# try:
+#     import comfy.samplers as samplers
+#     original_CFGGuider_inner_set_conds = samplers.CFGGuider.set_conds
+
+#     def patched_set_conds(self, positive, negative):
+#         if isinstance(self.model_patcher.model, KolorsSDXL):
+#             import copy
+#             if "control" in positive[0][1]:
+#                 if hasattr(positive[0][1]["control"], "control_model"):
+#                     if positive[0][1]["control"].control_model.label_emb.shape[1] == 5632:
+#                         return
+
+
+#             warnings.warn("该方法不再维护")
+#             positive = copy.deepcopy(positive)
+#             negative = copy.deepcopy(negative)
+#             hid_proj = self.model_patcher.model.encoder_hid_proj
+#             if hid_proj is not None:
+#                 positive[0][0] = hid_proj(positive[0][0])
+#                 negative[0][0] = hid_proj(negative[0][0])
+
+#                 if "control" in positive[0][1]:
+#                     if hasattr(positive[0][1]["control"], "control_model"):
+#                         positive[0][1]["control"].control_model.label_emb = self.model_patcher.model.diffusion_model.label_emb
+
+#                 if "control" in negative[0][1]:
+#                     if hasattr(negative[0][1]["control"], "control_model"):
+#                         negative[0][1]["control"].control_model.label_emb = self.model_patcher.model.diffusion_model.label_emb
+
+#         return original_CFGGuider_inner_set_conds(self, positive, negative)
+
+#     samplers.CFGGuider.set_conds = patched_set_conds
+# except ImportError:
+#     print("CFGGuider not found, skipping patching")
+
+
+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):
+        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)
+
+        result = super().forward(*args, **kwargs)
+        return result
+
+
+class KolorsSDXL(model_base.SDXL):
+    def __init__(self, model_config, model_type=ModelType.EPS, device=None):
+        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})
+
+    def encode_adm(self, **kwargs):
+        clip_pooled = sdxl_pooled(kwargs, self.noise_augmentor)
+        width = kwargs.get("width", 768)
+        height = kwargs.get("height", 768)
+        crop_w = kwargs.get("crop_w", 0)
+        crop_h = kwargs.get("crop_h", 0)
+        target_width = kwargs.get("target_width", width)
+        target_height = kwargs.get("target_height", height)
+
+        out = []
+        out.append(self.embedder(torch.Tensor([height])))
+        out.append(self.embedder(torch.Tensor([width])))
+        out.append(self.embedder(torch.Tensor([crop_h])))
+        out.append(self.embedder(torch.Tensor([crop_w])))
+        out.append(self.embedder(torch.Tensor([target_height])))
+        out.append(self.embedder(torch.Tensor([target_width])))
+        flat = torch.flatten(torch.cat(out)).unsqueeze(
+            dim=0).repeat(clip_pooled.shape[0], 1)
+        return torch.cat((clip_pooled.to(flat.device), flat), dim=1)
+
+
+class KolorsSupported(comfy.supported_models.SDXL):
+    unet_config = {
+        "model_channels": 320,
+        "use_linear_in_transformer": True,
+        "transformer_depth": [0, 0, 2, 2, 10, 10],
+        "context_dim": 2048,
+        "adm_in_channels": 5632,
+        "use_temporal_attention": False,
+    }
+
+    def get_model(self, state_dict, prefix="", device=None):
+        out = KolorsSDXL(self, model_type=self.model_type(
+            state_dict, prefix), device=device,)
+        out.__class__ = model_base.SDXL
+        if self.inpaint_model():
+            out.set_inpaint()
+        return out
+
+
+def kolors_unet_config_from_diffusers_unet(state_dict, dtype=None):
+    match = {}
+    transformer_depth = []
+
+    attn_res = 1
+    down_blocks = count_blocks(state_dict, "down_blocks.{}")
+    for i in range(down_blocks):
+        attn_blocks = count_blocks(
+            state_dict, "down_blocks.{}.attentions.".format(i) + '{}')
+        res_blocks = count_blocks(
+            state_dict, "down_blocks.{}.resnets.".format(i) + '{}')
+        for ab in range(attn_blocks):
+            transformer_count = count_blocks(
+                state_dict, "down_blocks.{}.attentions.{}.transformer_blocks.".format(i, ab) + '{}')
+            transformer_depth.append(transformer_count)
+            if transformer_count > 0:
+                match["context_dim"] = state_dict["down_blocks.{}.attentions.{}.transformer_blocks.0.attn2.to_k.weight".format(
+                    i, ab)].shape[1]
+
+        attn_res *= 2
+        if attn_blocks == 0:
+            for i in range(res_blocks):
+                transformer_depth.append(0)
+
+    match["transformer_depth"] = transformer_depth
+
+    match["model_channels"] = state_dict["conv_in.weight"].shape[0]
+    match["in_channels"] = state_dict["conv_in.weight"].shape[1]
+    match["adm_in_channels"] = None
+    if "class_embedding.linear_1.weight" in state_dict:
+        match["adm_in_channels"] = state_dict["class_embedding.linear_1.weight"].shape[1]
+    elif "add_embedding.linear_1.weight" in state_dict:
+        match["adm_in_channels"] = state_dict["add_embedding.linear_1.weight"].shape[1]
+
+    Kolors = {'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': 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}
+
+    supported_models = [Kolors]
+
+    for unet_config in supported_models:
+        matches = True
+        for k in match:
+            if match[k] != unet_config[k]:
+                print("key {} does not match".format(
+                    k), match[k], "||", unet_config[k])
+                matches = False
+                break
+        if matches:
+            return convert_config(unet_config)
+    return None
+
+
+import comfy.ldm.modules.diffusionmodules.openaimodel
+from torch import nn
+
+
+class apply_kolors:
+    def __enter__(self):
+        import comfy.ldm.modules.diffusionmodules.openaimodel
+        import comfy.utils
+
+        self.original_UNET_MAP_BASIC = comfy.utils.UNET_MAP_BASIC.copy()
+        comfy.utils.UNET_MAP_BASIC.add(
+            ("encoder_hid_proj.weight", "encoder_hid_proj.weight"),
+        )
+        comfy.utils.UNET_MAP_BASIC.add(
+            ("encoder_hid_proj.bias", "encoder_hid_proj.bias"),
+        )
+
+        self.original_unet_config_from_diffusers_unet = model_detection.unet_config_from_diffusers_unet
+        model_detection.unet_config_from_diffusers_unet = kolors_unet_config_from_diffusers_unet
+
+        import comfy.supported_models
+        self.original_supported_models = comfy.supported_models.models
+        comfy.supported_models.models = [KolorsSupported]
+
+    def __exit__(self, type, value, traceback):
+        import comfy.ldm.modules.diffusionmodules.openaimodel
+        import comfy.utils
+        comfy.utils.UNET_MAP_BASIC = self.original_UNET_MAP_BASIC
+
+        model_detection.unet_config_from_diffusers_unet = self.original_unet_config_from_diffusers_unet
+
+        import comfy.supported_models
+        comfy.supported_models.models = self.original_supported_models

mz_kolors_core.py

@@ -5,13 +5,14 @@ import json
 import os
 import random
 import re
+from types import MethodType
 
 import torch
 import folder_paths
 import comfy.model_management as mm
 
 
-def KolorsTextEncode(chatglm3_model, prompt):
+def chatglm3_text_encode(chatglm3_model, prompt):
     device = mm.get_torch_device()
     offload_device = mm.unet_offload_device()
     mm.unload_all_models()
@@ -24,14 +25,6 @@ def KolorsTextEncode(chatglm3_model, prompt):
 
     prompt = re.sub(r'\{([^{}]*)\}', choose_random_option, prompt)
 
-    if "|" in prompt:
-        prompt = prompt.split("|")
-
-    if prompt is not None and isinstance(prompt, str):
-        batch_size = 1
-    elif prompt is not None and isinstance(prompt, list):
-        batch_size = len(prompt)
-
     # Define tokenizers and text encoders
     tokenizer = chatglm3_model['tokenizer']
     text_encoder = chatglm3_model['text_encoder']
@@ -69,15 +62,13 @@ def KolorsTextEncode(chatglm3_model, prompt):
     return prompt_embeds, text_proj
 
 
-import importlib
-
-
 def MZ_ChatGLM3Loader_call(args):
-    from .mz_kolors_utils import Utils
-    llm_dir = os.path.join(Utils.get_models_path(), "LLM")
+    # from .mz_kolors_utils import Utils
+    # llm_dir = os.path.join(Utils.get_models_path(), "LLM")
     chatglm3_checkpoint = args.get("chatglm3_checkpoint")
 
-    chatglm3_checkpoint_path = os.path.join(llm_dir, chatglm3_checkpoint)
+    chatglm3_checkpoint_path = folder_paths.get_full_path(
+        'LLM', chatglm3_checkpoint)
 
     if not os.path.exists(chatglm3_checkpoint_path):
         raise RuntimeError(
@@ -131,115 +122,103 @@ def MZ_ChatGLM3Loader_call(args):
     return ({"text_encoder": text_encoder, "tokenizer": tokenizer},)
 
 
-def MZ_ChatGLM3TextEncode_call(args):
-
+def MZ_ChatGLM3TextEncodeV2_call(args):
     text = args.get("text")
     chatglm3_model = args.get("chatglm3_model")
-
-    prompt_embeds, pooled_output = KolorsTextEncode(
+    prompt_embeds, pooled_output = chatglm3_text_encode(
         chatglm3_model,
         text,
     )
-
-    from torch import nn
-    hid_proj: nn.Linear = args.get("hid_proj")
-
-    prompt_embeds = hid_proj(prompt_embeds)
-
     return ([[
         prompt_embeds,
         {"pooled_output": pooled_output},
     ]], )
 
 
-def load_unet_state_dict(sd):  # load unet in diffusers or regular format
-    from comfy import model_management, model_detection
-    import comfy.utils
+def MZ_KolorsUNETLoaderV2_call(kwargs):
 
-    # Allow loading unets from checkpoint files
-    checkpoint = False
-    diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
-    temp_sd = comfy.utils.state_dict_prefix_replace(
-        sd, {diffusion_model_prefix: ""}, filter_keys=True)
-    if len(temp_sd) > 0:
-        sd = temp_sd
-        checkpoint = True
+    from . import hook_comfyui_kolors_v2
+    import comfy.sd
 
-    parameters = comfy.utils.calculate_parameters(sd)
-    unet_dtype = model_management.unet_dtype(model_params=parameters)
-    load_device = model_management.get_torch_device()
-
-    from torch import nn
-    hid_proj: nn.Linear = None
-    if True:
-        model_config = model_detection.model_config_from_diffusers_unet(sd)
-        if model_config is None:
-            return None
-
-        diffusers_keys = comfy.utils.unet_to_diffusers(
-            model_config.unet_config)
-
-        new_sd = {}
-        for k in diffusers_keys:
-            if k in sd:
-                new_sd[diffusers_keys[k]] = sd.pop(k)
-            else:
-                print("{} {}".format(diffusers_keys[k], k))
-
-        encoder_hid_proj_weight = sd.pop("encoder_hid_proj.weight")
-        encoder_hid_proj_bias = sd.pop("encoder_hid_proj.bias")
-        hid_proj = nn.Linear(
-            encoder_hid_proj_weight.shape[1], encoder_hid_proj_weight.shape[0])
-        hid_proj.weight.data = encoder_hid_proj_weight
-        hid_proj.bias.data = encoder_hid_proj_bias
-        hid_proj = hid_proj.to(load_device)
-
-    offload_device = model_management.unet_offload_device()
-    unet_dtype = model_management.unet_dtype(
-        model_params=parameters, supported_dtypes=model_config.supported_inference_dtypes)
-    manual_cast_dtype = model_management.unet_manual_cast(
-        unet_dtype, load_device, model_config.supported_inference_dtypes)
-    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
-    model = model_config.get_model(new_sd, "")
-    model = model.to(offload_device)
-    model.load_model_weights(new_sd, "")
-    left_over = sd.keys()
-    if len(left_over) > 0:
-        print("left over keys in unet: {}".format(left_over))
-    return comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=offload_device), hid_proj
-
-
-def MZ_KolorsUNETLoader_call(kwargs):
-
-    from . import hook_comfyui
-    # with hook_comfyui.apply_kolors():
-    if True:
+    with hook_comfyui_kolors_v2.apply_kolors():
         unet_name = kwargs.get("unet_name")
         unet_path = folder_paths.get_full_path("unet", unet_name)
         import comfy.utils
         sd = comfy.utils.load_torch_file(unet_path)
-        comfy.mz_log("add_embedding.linear_1.weight", sd["add_embedding.linear_1.weight"])
-        # add_embedding_linear_1_weight : torch.Size([1280, 5632])
-        add_embedding_linear_1_weight = sd["add_embedding.linear_1.weight"]
-        # new_add_embedding_linear_1_weight : torch.Size([1280, 2816])
-        new_add_embedding_linear_1_weight = add_embedding_linear_1_weight[:, :2816]
-        sd["add_embedding.linear_1.weight"] = new_add_embedding_linear_1_weight
-         
-        model, hid_proj = load_unet_state_dict(sd)
+        model = comfy.sd.load_unet_state_dict(sd)
         if model is None:
             raise RuntimeError(
                 "ERROR: Could not detect model type of: {}".format(unet_path))
-        return (model, hid_proj)
+
+        return (model, )
 
 
-def MZ_FakeCond_call(kwargs):
-    import torch
-    # cond: torch.Size([1, 77, ])
-    cond = torch.zeros(2, 256, 4096)
-    # torch.Size([1, 1280])
-    pool = torch.zeros(2, 4096)
+from comfy.cldm.cldm import ControlNet
+from comfy.controlnet import ControlLora
 
-    return ([[
-        cond,
-        {"pooled_output": pool},
-    ]],)
+
+def MZ_KolorsControlNetPatch_call(kwargs):
+    from . import hook_comfyui_kolors_v2
+    model = kwargs.get("model")
+    control_net = kwargs.get("control_net")
+    import comfy.controlnet
+    if isinstance(control_net, ControlLora):
+        del_keys = []
+        for k in control_net.control_weights:
+            if k.startswith("label_emb.0.0."):
+                del_keys.append(k)
+
+        for k in del_keys:
+            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):
+            context = model.model.diffusion_model.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"):
+                self.control_model.forward = MethodType(
+                    KolorsControlNet_forward, self.control_model)
+            return result
+
+        control_net.pre_run = MethodType(
+            KolorsControlLora_pre_run, control_net)
+
+        def KolorsControlLora_copy(self, *args, **kwargs):
+            c = super_copy(self, *args, **kwargs)
+            c.pre_run = MethodType(
+                KolorsControlLora_pre_run, c)
+            return c
+
+        control_net.copy = MethodType(
+            KolorsControlLora_copy, 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
+
+        control_net.control_model_wrapped.model.label_emb = model_label_emb
+
+        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)
+                return super_forward(self, x, hint, timesteps, context, **kwargs)
+
+        control_net.control_model.forward = MethodType(
+            KolorsControlNet_forward, control_net.control_model)
+
+    else:
+        raise NotImplementedError(
+            f"Type {control_net} not supported for KolorsControlNetPatch")
+
+    return (control_net,)

mz_kolors_legacy.py

@@ -0,0 +1,203 @@
+
+
+import gc
+import json
+import os
+import random
+import re
+
+import torch
+import folder_paths
+import comfy.model_management as mm
+from . import mz_kolors_core
+
+
+def MZ_ChatGLM3TextEncode_call(args):
+
+    text = args.get("text")
+    chatglm3_model = args.get("chatglm3_model")
+
+    prompt_embeds, pooled_output = mz_kolors_core.chatglm3_text_encode(
+        chatglm3_model,
+        text,
+    )
+
+    from torch import nn
+    hid_proj: nn.Linear = args.get("hid_proj")
+
+    if hid_proj.weight.dtype != prompt_embeds.dtype:
+        with torch.cuda.amp.autocast(dtype=hid_proj.weight.dtype):
+            prompt_embeds = hid_proj(prompt_embeds)
+    else:
+        prompt_embeds = hid_proj(prompt_embeds)
+
+    return ([[
+        prompt_embeds,
+        {"pooled_output": pooled_output},
+    ]], )
+
+
+def load_unet_state_dict(sd):  # load unet in diffusers or regular format
+    from comfy import model_management, model_detection
+    import comfy.utils
+
+    # Allow loading unets from checkpoint files
+    checkpoint = False
+    diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
+    temp_sd = comfy.utils.state_dict_prefix_replace(
+        sd, {diffusion_model_prefix: ""}, filter_keys=True)
+    if len(temp_sd) > 0:
+        sd = temp_sd
+        checkpoint = True
+
+    parameters = comfy.utils.calculate_parameters(sd)
+    unet_dtype = model_management.unet_dtype(model_params=parameters)
+    load_device = model_management.get_torch_device()
+
+    from torch import nn
+    hid_proj: nn.Linear = None
+    if True:
+        model_config = model_detection.model_config_from_diffusers_unet(sd)
+        if model_config is None:
+            return None
+
+        diffusers_keys = comfy.utils.unet_to_diffusers(
+            model_config.unet_config)
+
+        new_sd = {}
+        for k in diffusers_keys:
+            if k in sd:
+                new_sd[diffusers_keys[k]] = sd.pop(k)
+            else:
+                print("{} {}".format(diffusers_keys[k], k))
+
+        encoder_hid_proj_weight = sd.pop("encoder_hid_proj.weight")
+        encoder_hid_proj_bias = sd.pop("encoder_hid_proj.bias")
+        hid_proj = nn.Linear(
+            encoder_hid_proj_weight.shape[1], encoder_hid_proj_weight.shape[0])
+        hid_proj.weight.data = encoder_hid_proj_weight
+        hid_proj.bias.data = encoder_hid_proj_bias
+        hid_proj = hid_proj.to(load_device)
+
+    offload_device = model_management.unet_offload_device()
+    unet_dtype = model_management.unet_dtype(
+        model_params=parameters, supported_dtypes=model_config.supported_inference_dtypes)
+    manual_cast_dtype = model_management.unet_manual_cast(
+        unet_dtype, load_device, model_config.supported_inference_dtypes)
+    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
+    model = model_config.get_model(new_sd, "")
+    model = model.to(offload_device)
+    model.load_model_weights(new_sd, "")
+    left_over = sd.keys()
+    if len(left_over) > 0:
+        print("left over keys in unet: {}".format(left_over))
+    return comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=offload_device), hid_proj
+
+
+def MZ_KolorsUNETLoader_call(kwargs):
+
+    from . import hook_comfyui_kolors_v1
+    with hook_comfyui_kolors_v1.apply_kolors():
+        unet_name = kwargs.get("unet_name")
+        unet_path = folder_paths.get_full_path("unet", unet_name)
+        import comfy.utils
+        sd = comfy.utils.load_torch_file(unet_path)
+        model, hid_proj = load_unet_state_dict(sd)
+        if model is None:
+            raise RuntimeError(
+                "ERROR: Could not detect model type of: {}".format(unet_path))
+        return (model, hid_proj)
+
+
+def MZ_FakeCond_call(kwargs):
+    import torch
+    cond = torch.zeros(2, 256, 4096)
+    pool = torch.zeros(2, 4096)
+
+    return ([[
+        cond,
+        {"pooled_output": pool},
+    ]],)
+
+
+NODE_CLASS_MAPPINGS = {
+}
+
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+}
+
+AUTHOR_NAME = "MinusZone"
+CATEGORY_NAME = f"{AUTHOR_NAME} - Kolors"
+
+
+class MZ_ChatGLM3TextEncode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "chatglm3_model": ("CHATGLM3MODEL", ),
+                "text": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+                "hid_proj": ("TorchLinear", ),
+            }
+        }
+
+    RETURN_TYPES = ("CONDITIONING",)
+
+    FUNCTION = "encode"
+    CATEGORY = CATEGORY_NAME + "/Legacy"
+
+    def encode(self, **kwargs):
+        return MZ_ChatGLM3TextEncode_call(kwargs)
+
+
+NODE_CLASS_MAPPINGS["MZ_ChatGLM3"] = MZ_ChatGLM3TextEncode
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_ChatGLM3"] = f"{AUTHOR_NAME} - ChatGLM3TextEncode"
+
+
+class MZ_KolorsUNETLoader():
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                "unet_name": (folder_paths.get_filename_list("unet"), ),
+                }}
+
+    RETURN_TYPES = ("MODEL", "TorchLinear")
+
+    RETURN_NAMES = ("model", "hid_proj")
+
+    FUNCTION = "load_unet"
+
+    CATEGORY = CATEGORY_NAME + "/Legacy"
+
+    def load_unet(self, **kwargs):
+        return MZ_KolorsUNETLoader_call(kwargs)
+
+
+NODE_CLASS_MAPPINGS["MZ_KolorsUNETLoader"] = MZ_KolorsUNETLoader
+NODE_DISPLAY_NAME_MAPPINGS[
+    "MZ_KolorsUNETLoader"] = f"{AUTHOR_NAME} - Kolors UNET Loader"
+
+
+class MZ_FakeCond:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "seed": ("INT", {"default": 0}),
+            }
+        }
+
+    RETURN_TYPES = ("CONDITIONING", )
+    RETURN_NAMES = ("prompt", )
+    FUNCTION = "encode"
+    CATEGORY = CATEGORY_NAME
+
+    def encode(self, **kwargs):
+        return MZ_FakeCond_call(kwargs)
+
+
+# NODE_CLASS_MAPPINGS["MZ_FakeCond"] = MZ_FakeCond
+# NODE_DISPLAY_NAME_MAPPINGS[
+#     "MZ_FakeCond"] = f"{AUTHOR_NAME} - FakeCond"

pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+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"
+
+[project.urls]
+Repository = "https://github.com/MinusZoneAI/ComfyUI-Kolors-MZ"
+#  Used by Comfy Registry https://comfyregistry.org
+
+[tool.comfy]
+PublisherId = "wailovet"
+DisplayName = "ComfyUI-Kolors-MZ"
+Icon = ""

test.py

@@ -0,0 +1,76 @@
+from typing import Any
+
+
+class A:
+    def __init__(self):
+        self.a = 1
+
+    def get_a(self):
+        return self.a
+
+
+class B(A):
+    def __init__(self):
+        super().__init__()
+        self.__name__ = A.__name__
+        self.b = 2
+
+    def get_b(self):
+        return self.b
+
+    def __getattribute__(self, name: str) -> Any:
+        print("getattribute")
+        return super().__getattribute__(name)
+
+    def __call__(self, *args: Any, **kwds: Any) -> Any:
+        print("call")
+        return super().__call__(*args, **kwds)
+
+    def __reduce__(self):
+        print("reduce")
+        return super().__reduce__()
+
+    def __reduce_ex__(self, protocol: int):
+        print("reduce_ex")
+        return super().__reduce_ex__(protocol)
+
+    def __repr__(self) -> str:
+        print("repr")
+        return super().__repr__()
+
+    def __str__(self) -> str:
+        print("str")
+        return super().__str__()
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        print("setattr:", name, value)
+        super().__setattr__(name, value)
+
+    def __delattr__(self, name: str) -> None:
+        print("delattr")
+        super().__delattr__(name)
+
+    def __dir__(self):
+        print("dir")
+        return super().__dir__()
+
+    def __eq__(self, other: Any) -> bool:
+        print("eq")
+        return super().__eq__(other)
+
+    def __ne__(self, other: Any) -> bool:
+        print("ne")
+        return super().__ne__(other)
+
+    def __hash__(self) -> int:
+        print("hash")
+        return super().__hash__()
+
+    def __bool__(self) -> bool:
+        print("bool")
+        return super().__bool__()
+
+
+b = B()
+print(b.get_b())
+print(type(b))