Update README.md

optimize controlnet loader

.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 }}

ComfyUI_IPAdapter_plus/CrossAttentionPatch.py

@@ -0,0 +1,209 @@
+import torch
+import math
+import torch.nn.functional as F
+from comfy.ldm.modules.attention import optimized_attention
+from .utils import tensor_to_size
+
+class Attn2Replace:
+    def __init__(self, callback=None, **kwargs):
+        self.callback = [callback]
+        self.kwargs = [kwargs]
+
+    def add(self, callback, **kwargs):
+        self.callback.append(callback)
+        self.kwargs.append(kwargs)
+
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def __call__(self, q, k, v, extra_options):
+        dtype = q.dtype
+        out = optimized_attention(q, k, v, extra_options["n_heads"])
+        sigma = extra_options["sigmas"].detach().cpu()[0].item() if 'sigmas' in extra_options else 999999999.9
+
+        for i, callback in enumerate(self.callback):
+            if sigma <= self.kwargs[i]["sigma_start"] and sigma >= self.kwargs[i]["sigma_end"]:
+                out = out + callback(out, q, k, v, extra_options, **self.kwargs[i])
+
+        return out.to(dtype=dtype)
+
+def ipadapter_attention(out, q, k, v, extra_options, module_key='', ipadapter=None, weight=1.0, cond=None, cond_alt=None, uncond=None, weight_type="linear", mask=None, sigma_start=0.0, sigma_end=1.0, unfold_batch=False, embeds_scaling='V only', **kwargs):
+    dtype = q.dtype
+    cond_or_uncond = extra_options["cond_or_uncond"]
+    block_type = extra_options["block"][0]
+    #block_id = extra_options["block"][1]
+    t_idx = extra_options["transformer_index"]
+    layers = 11 if '101_to_k_ip' in ipadapter.ip_layers.to_kvs else 16
+    k_key = module_key + "_to_k_ip"
+    v_key = module_key + "_to_v_ip"
+
+    # extra options for AnimateDiff
+    ad_params = extra_options['ad_params'] if "ad_params" in extra_options else None
+
+    b = q.shape[0]
+    seq_len = q.shape[1]
+    batch_prompt = b // len(cond_or_uncond)
+    _, _, oh, ow = extra_options["original_shape"]
+
+    if weight_type == 'ease in':
+        weight = weight * (0.05 + 0.95 * (1 - t_idx / layers))
+    elif weight_type == 'ease out':
+        weight = weight * (0.05 + 0.95 * (t_idx / layers))
+    elif weight_type == 'ease in-out':
+        weight = weight * (0.05 + 0.95 * (1 - abs(t_idx - (layers/2)) / (layers/2)))
+    elif weight_type == 'reverse in-out':
+        weight = weight * (0.05 + 0.95 * (abs(t_idx - (layers/2)) / (layers/2)))
+    elif weight_type == 'weak input' and block_type == 'input':
+        weight = weight * 0.2
+    elif weight_type == 'weak middle' and block_type == 'middle':
+        weight = weight * 0.2
+    elif weight_type == 'weak output' and block_type == 'output':
+        weight = weight * 0.2
+    elif weight_type == 'strong middle' and (block_type == 'input' or block_type == 'output'):
+        weight = weight * 0.2
+    elif isinstance(weight, dict):
+        if t_idx not in weight:
+            return 0
+
+        if weight_type == "style transfer precise":
+            if layers == 11 and t_idx == 3:
+                uncond = cond
+                cond = cond * 0
+            elif layers == 16 and (t_idx == 4 or t_idx == 5):
+                uncond = cond
+                cond = cond * 0
+        elif weight_type == "composition precise":
+            if layers == 11 and t_idx != 3:
+                uncond = cond
+                cond = cond * 0
+            elif layers == 16 and (t_idx != 4 and t_idx != 5):
+                uncond = cond
+                cond = cond * 0
+
+        weight = weight[t_idx]
+
+        if cond_alt is not None and t_idx in cond_alt:
+            cond = cond_alt[t_idx]
+            del cond_alt
+
+    if unfold_batch:
+        # Check AnimateDiff context window
+        if ad_params is not None and ad_params["sub_idxs"] is not None:
+            if isinstance(weight, torch.Tensor):
+                weight = tensor_to_size(weight, ad_params["full_length"])
+                weight = torch.Tensor(weight[ad_params["sub_idxs"]])
+                if torch.all(weight == 0):
+                    return 0
+                weight = weight.repeat(len(cond_or_uncond), 1, 1) # repeat for cond and uncond
+            elif weight == 0:
+                return 0
+
+            # if image length matches or exceeds full_length get sub_idx images
+            if cond.shape[0] >= ad_params["full_length"]:
+                cond = torch.Tensor(cond[ad_params["sub_idxs"]])
+                uncond = torch.Tensor(uncond[ad_params["sub_idxs"]])
+            # otherwise get sub_idxs images
+            else:
+                cond = tensor_to_size(cond, ad_params["full_length"])
+                uncond = tensor_to_size(uncond, ad_params["full_length"])
+                cond = cond[ad_params["sub_idxs"]]
+                uncond = uncond[ad_params["sub_idxs"]]
+        else:
+            if isinstance(weight, torch.Tensor):
+                weight = tensor_to_size(weight, batch_prompt)
+                if torch.all(weight == 0):
+                    return 0
+                weight = weight.repeat(len(cond_or_uncond), 1, 1) # repeat for cond and uncond
+            elif weight == 0:
+                return 0
+
+            cond = tensor_to_size(cond, batch_prompt)
+            uncond = tensor_to_size(uncond, batch_prompt)
+
+        k_cond = ipadapter.ip_layers.to_kvs[k_key](cond)
+        k_uncond = ipadapter.ip_layers.to_kvs[k_key](uncond)
+        v_cond = ipadapter.ip_layers.to_kvs[v_key](cond)
+        v_uncond = ipadapter.ip_layers.to_kvs[v_key](uncond)
+    else:
+        # TODO: should we always convert the weights to a tensor?
+        if isinstance(weight, torch.Tensor):
+            weight = tensor_to_size(weight, batch_prompt)
+            if torch.all(weight == 0):
+                return 0
+            weight = weight.repeat(len(cond_or_uncond), 1, 1) # repeat for cond and uncond
+        elif weight == 0:
+            return 0
+
+        k_cond = ipadapter.ip_layers.to_kvs[k_key](cond).repeat(batch_prompt, 1, 1)
+        k_uncond = ipadapter.ip_layers.to_kvs[k_key](uncond).repeat(batch_prompt, 1, 1)
+        v_cond = ipadapter.ip_layers.to_kvs[v_key](cond).repeat(batch_prompt, 1, 1)
+        v_uncond = ipadapter.ip_layers.to_kvs[v_key](uncond).repeat(batch_prompt, 1, 1)
+
+    if len(cond_or_uncond) == 3: # TODO: conxl, I need to check this
+        ip_k = torch.cat([(k_cond, k_uncond, k_cond)[i] for i in cond_or_uncond], dim=0)
+        ip_v = torch.cat([(v_cond, v_uncond, v_cond)[i] for i in cond_or_uncond], dim=0)
+    else:
+        ip_k = torch.cat([(k_cond, k_uncond)[i] for i in cond_or_uncond], dim=0)
+        ip_v = torch.cat([(v_cond, v_uncond)[i] for i in cond_or_uncond], dim=0)
+
+    if embeds_scaling == 'K+mean(V) w/ C penalty':
+        scaling = float(ip_k.shape[2]) / 1280.0
+        weight = weight * scaling
+        ip_k = ip_k * weight
+        ip_v_mean = torch.mean(ip_v, dim=1, keepdim=True)
+        ip_v = (ip_v - ip_v_mean) + ip_v_mean * weight
+        out_ip = optimized_attention(q, ip_k, ip_v, extra_options["n_heads"])
+        del ip_v_mean
+    elif embeds_scaling == 'K+V w/ C penalty':
+        scaling = float(ip_k.shape[2]) / 1280.0
+        weight = weight * scaling
+        ip_k = ip_k * weight
+        ip_v = ip_v * weight
+        out_ip = optimized_attention(q, ip_k, ip_v, extra_options["n_heads"])
+    elif embeds_scaling == 'K+V':
+        ip_k = ip_k * weight
+        ip_v = ip_v * weight
+        out_ip = optimized_attention(q, ip_k, ip_v, extra_options["n_heads"])
+    else:
+        #ip_v = ip_v * weight
+        out_ip = optimized_attention(q, ip_k, ip_v, extra_options["n_heads"])
+        out_ip = out_ip * weight # I'm doing this to get the same results as before
+
+    if mask is not None:
+        mask_h = oh / math.sqrt(oh * ow / seq_len)
+        mask_h = int(mask_h) + int((seq_len % int(mask_h)) != 0)
+        mask_w = seq_len // mask_h
+
+        # check if using AnimateDiff and sliding context window
+        if (mask.shape[0] > 1 and ad_params is not None and ad_params["sub_idxs"] is not None):
+            # if mask length matches or exceeds full_length, get sub_idx masks
+            if mask.shape[0] >= ad_params["full_length"]:
+                mask = torch.Tensor(mask[ad_params["sub_idxs"]])
+                mask = F.interpolate(mask.unsqueeze(1), size=(mask_h, mask_w), mode="bilinear").squeeze(1)
+            else:
+                mask = F.interpolate(mask.unsqueeze(1), size=(mask_h, mask_w), mode="bilinear").squeeze(1)
+                mask = tensor_to_size(mask, ad_params["full_length"])
+                mask = mask[ad_params["sub_idxs"]]
+        else:
+            mask = F.interpolate(mask.unsqueeze(1), size=(mask_h, mask_w), mode="bilinear").squeeze(1)
+            mask = tensor_to_size(mask, batch_prompt)
+
+        mask = mask.repeat(len(cond_or_uncond), 1, 1)
+        mask = mask.view(mask.shape[0], -1, 1).repeat(1, 1, out.shape[2])
+
+        # covers cases where extreme aspect ratios can cause the mask to have a wrong size
+        mask_len = mask_h * mask_w
+        if mask_len < seq_len:
+            pad_len = seq_len - mask_len
+            pad1 = pad_len // 2
+            pad2 = pad_len - pad1
+            mask = F.pad(mask, (0, 0, pad1, pad2), value=0.0)
+        elif mask_len > seq_len:
+            crop_start = (mask_len - seq_len) // 2
+            mask = mask[:, crop_start:crop_start+seq_len, :]
+
+        out_ip = out_ip * mask
+
+    #out = out + out_ip
+
+    return out_ip.to(dtype=dtype)

ComfyUI_IPAdapter_plus/image_proj_models.py

@@ -0,0 +1,275 @@
+import math
+import torch
+import torch.nn as nn
+from einops import rearrange
+from einops.layers.torch import Rearrange
+
+
+# FFN
+def FeedForward(dim, mult=4):
+    inner_dim = int(dim * mult)
+    return nn.Sequential(
+        nn.LayerNorm(dim),
+        nn.Linear(dim, inner_dim, bias=False),
+        nn.GELU(),
+        nn.Linear(inner_dim, dim, bias=False),
+    )
+
+
+def reshape_tensor(x, heads):
+    bs, length, width = x.shape
+    # (bs, length, width) --> (bs, length, n_heads, dim_per_head)
+    x = x.view(bs, length, heads, -1)
+    # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
+    x = x.transpose(1, 2)
+    # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
+    x = x.reshape(bs, heads, length, -1)
+    return x
+
+
+class PerceiverAttention(nn.Module):
+    def __init__(self, *, dim, dim_head=64, heads=8):
+        super().__init__()
+        self.scale = dim_head**-0.5
+        self.dim_head = dim_head
+        self.heads = heads
+        inner_dim = dim_head * heads
+
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
+        self.to_out = nn.Linear(inner_dim, dim, bias=False)
+
+    def forward(self, x, latents):
+        """
+        Args:
+            x (torch.Tensor): image features
+                shape (b, n1, D)
+            latent (torch.Tensor): latent features
+                shape (b, n2, D)
+        """
+        x = self.norm1(x)
+        latents = self.norm2(latents)
+
+        b, l, _ = latents.shape
+
+        q = self.to_q(latents)
+        kv_input = torch.cat((x, latents), dim=-2)
+        k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+
+        q = reshape_tensor(q, self.heads)
+        k = reshape_tensor(k, self.heads)
+        v = reshape_tensor(v, self.heads)
+
+        # attention
+        scale = 1 / math.sqrt(math.sqrt(self.dim_head))
+        weight = (q * scale) @ (k * scale).transpose(-2, -1)  # More stable with f16 than dividing afterwards
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        out = weight @ v
+
+        out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
+
+        return self.to_out(out)
+
+
+class Resampler(nn.Module):
+    def __init__(
+        self,
+        dim=1024,
+        depth=8,
+        dim_head=64,
+        heads=16,
+        num_queries=8,
+        embedding_dim=768,
+        output_dim=1024,
+        ff_mult=4,
+        max_seq_len: int = 257,  # CLIP tokens + CLS token
+        apply_pos_emb: bool = False,
+        num_latents_mean_pooled: int = 0,  # number of latents derived from mean pooled representation of the sequence
+    ):
+        super().__init__()
+        self.pos_emb = nn.Embedding(max_seq_len, embedding_dim) if apply_pos_emb else None
+
+        self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5)
+
+        self.proj_in = nn.Linear(embedding_dim, dim)
+
+        self.proj_out = nn.Linear(dim, output_dim)
+        self.norm_out = nn.LayerNorm(output_dim)
+
+        self.to_latents_from_mean_pooled_seq = (
+            nn.Sequential(
+                nn.LayerNorm(dim),
+                nn.Linear(dim, dim * num_latents_mean_pooled),
+                Rearrange("b (n d) -> b n d", n=num_latents_mean_pooled),
+            )
+            if num_latents_mean_pooled > 0
+            else None
+        )
+
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
+                        FeedForward(dim=dim, mult=ff_mult),
+                    ]
+                )
+            )
+
+    def forward(self, x):
+        if self.pos_emb is not None:
+            n, device = x.shape[1], x.device
+            pos_emb = self.pos_emb(torch.arange(n, device=device))
+            x = x + pos_emb
+
+        latents = self.latents.repeat(x.size(0), 1, 1)
+
+        x = self.proj_in(x)
+
+        if self.to_latents_from_mean_pooled_seq:
+            meanpooled_seq = masked_mean(x, dim=1, mask=torch.ones(x.shape[:2], device=x.device, dtype=torch.bool))
+            meanpooled_latents = self.to_latents_from_mean_pooled_seq(meanpooled_seq)
+            latents = torch.cat((meanpooled_latents, latents), dim=-2)
+
+        for attn, ff in self.layers:
+            latents = attn(x, latents) + latents
+            latents = ff(latents) + latents
+
+        latents = self.proj_out(latents)
+        return self.norm_out(latents)
+
+
+def masked_mean(t, *, dim, mask=None):
+    if mask is None:
+        return t.mean(dim=dim)
+
+    denom = mask.sum(dim=dim, keepdim=True)
+    mask = rearrange(mask, "b n -> b n 1")
+    masked_t = t.masked_fill(~mask, 0.0)
+
+    return masked_t.sum(dim=dim) / denom.clamp(min=1e-5)
+
+
+class FacePerceiverResampler(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim=768,
+        depth=4,
+        dim_head=64,
+        heads=16,
+        embedding_dim=1280,
+        output_dim=768,
+        ff_mult=4,
+    ):
+        super().__init__()
+
+        self.proj_in = nn.Linear(embedding_dim, dim)
+        self.proj_out = nn.Linear(dim, output_dim)
+        self.norm_out = nn.LayerNorm(output_dim)
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
+                        FeedForward(dim=dim, mult=ff_mult),
+                    ]
+                )
+            )
+
+    def forward(self, latents, x):
+        x = self.proj_in(x)
+        for attn, ff in self.layers:
+            latents = attn(x, latents) + latents
+            latents = ff(latents) + latents
+        latents = self.proj_out(latents)
+        return self.norm_out(latents)
+
+
+class MLPProjModel(nn.Module):
+    def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024):
+        super().__init__()
+
+        self.proj = nn.Sequential(
+            nn.Linear(clip_embeddings_dim, clip_embeddings_dim),
+            nn.GELU(),
+            nn.Linear(clip_embeddings_dim, cross_attention_dim),
+            nn.LayerNorm(cross_attention_dim)
+        )
+
+    def forward(self, image_embeds):
+        clip_extra_context_tokens = self.proj(image_embeds)
+        return clip_extra_context_tokens
+
+class MLPProjModelFaceId(nn.Module):
+    def __init__(self, cross_attention_dim=768, id_embeddings_dim=512, num_tokens=4):
+        super().__init__()
+
+        self.cross_attention_dim = cross_attention_dim
+        self.num_tokens = num_tokens
+
+        self.proj = nn.Sequential(
+            nn.Linear(id_embeddings_dim, id_embeddings_dim*2),
+            nn.GELU(),
+            nn.Linear(id_embeddings_dim*2, cross_attention_dim*num_tokens),
+        )
+        self.norm = nn.LayerNorm(cross_attention_dim)
+
+    def forward(self, id_embeds):
+        x = self.proj(id_embeds)
+        x = x.reshape(-1, self.num_tokens, self.cross_attention_dim)
+        x = self.norm(x)
+        return x
+
+class ProjModelFaceIdPlus(nn.Module):
+    def __init__(self, cross_attention_dim=768, id_embeddings_dim=512, clip_embeddings_dim=1280, num_tokens=4):
+        super().__init__()
+
+        self.cross_attention_dim = cross_attention_dim
+        self.num_tokens = num_tokens
+
+        self.proj = nn.Sequential(
+            nn.Linear(id_embeddings_dim, id_embeddings_dim*2),
+            nn.GELU(),
+            nn.Linear(id_embeddings_dim*2, cross_attention_dim*num_tokens),
+        )
+        self.norm = nn.LayerNorm(cross_attention_dim)
+
+        self.perceiver_resampler = FacePerceiverResampler(
+            dim=cross_attention_dim,
+            depth=4,
+            dim_head=64,
+            heads=cross_attention_dim // 64,
+            embedding_dim=clip_embeddings_dim,
+            output_dim=cross_attention_dim,
+            ff_mult=4,
+        )
+
+    def forward(self, id_embeds, clip_embeds, scale=1.0, shortcut=False):
+        x = self.proj(id_embeds)
+        x = x.reshape(-1, self.num_tokens, self.cross_attention_dim)
+        x = self.norm(x)
+        out = self.perceiver_resampler(x, clip_embeds)
+        if shortcut:
+            out = x + scale * out
+        return out
+
+class ImageProjModel(nn.Module):
+    def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
+        super().__init__()
+
+        self.cross_attention_dim = cross_attention_dim
+        self.clip_extra_context_tokens = clip_extra_context_tokens
+        self.proj = nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim)
+        self.norm = nn.LayerNorm(cross_attention_dim)
+
+    def forward(self, image_embeds):
+        embeds = image_embeds
+        x = self.proj(embeds).reshape(-1, self.clip_extra_context_tokens, self.cross_attention_dim)
+        x = self.norm(x)
+        return x

ComfyUI_IPAdapter_plus/utils.py

@@ -0,0 +1,270 @@
+import re
+import torch
+import os
+import folder_paths
+from comfy.clip_vision import clip_preprocess, Output
+import comfy.utils
+import comfy.model_management as model_management
+try:
+    import torchvision.transforms.v2 as T
+except ImportError:
+    import torchvision.transforms as T
+
+def get_clipvision_file(preset):
+    preset = preset.lower()
+    clipvision_list = folder_paths.get_filename_list("clip_vision")
+
+    if preset.startswith("vit-g"):
+        pattern = r'(ViT.bigG.14.*39B.b160k|ipadapter.*sdxl|sdxl.*model\.(bin|safetensors))'
+    else:
+        pattern = r'(ViT.H.14.*s32B.b79K|ipadapter.*sd15|sd1.?5.*model\.(bin|safetensors))'
+    clipvision_file = [e for e in clipvision_list if re.search(pattern, e, re.IGNORECASE)]
+
+    clipvision_file = folder_paths.get_full_path("clip_vision", clipvision_file[0]) if clipvision_file else None
+
+    return clipvision_file
+
+def get_ipadapter_file(preset, is_sdxl):
+    preset = preset.lower()
+    ipadapter_list = folder_paths.get_filename_list("ipadapter")
+    is_insightface = False
+    lora_pattern = None
+
+    if preset.startswith("light"):
+        if is_sdxl:
+            raise Exception("light model is not supported for SDXL")
+        pattern = r'sd15.light.v11\.(safetensors|bin)$'
+        # if v11 is not found, try with the old version
+        if not [e for e in ipadapter_list if re.search(pattern, e, re.IGNORECASE)]:
+            pattern = r'sd15.light\.(safetensors|bin)$'
+    elif preset.startswith("standard"):
+        if is_sdxl:
+            pattern = r'ip.adapter.sdxl.vit.h\.(safetensors|bin)$'
+        else:
+            pattern = r'ip.adapter.sd15\.(safetensors|bin)$'
+    elif preset.startswith("vit-g"):
+        if is_sdxl:
+            pattern = r'ip.adapter.sdxl\.(safetensors|bin)$'
+        else:
+            pattern = r'sd15.vit.g\.(safetensors|bin)$'
+    elif preset.startswith("plus ("):
+        if is_sdxl:
+            pattern = r'plus.sdxl.vit.h\.(safetensors|bin)$'
+        else:
+            pattern = r'ip.adapter.plus.sd15\.(safetensors|bin)$'
+    elif preset.startswith("plus face"):
+        if is_sdxl:
+            pattern = r'plus.face.sdxl.vit.h\.(safetensors|bin)$'
+        else:
+            pattern = r'plus.face.sd15\.(safetensors|bin)$'
+    elif preset.startswith("full"):
+        if is_sdxl:
+            raise Exception("full face model is not supported for SDXL")
+        pattern = r'full.face.sd15\.(safetensors|bin)$'
+    elif preset.startswith("faceid portrait ("):
+        if is_sdxl:
+            pattern = r'portrait.sdxl\.(safetensors|bin)$'
+        else:
+            pattern = r'portrait.v11.sd15\.(safetensors|bin)$'
+            # if v11 is not found, try with the old version
+            if not [e for e in ipadapter_list if re.search(pattern, e, re.IGNORECASE)]:
+                pattern = r'portrait.sd15\.(safetensors|bin)$'
+        is_insightface = True
+    elif preset.startswith("faceid portrait unnorm"):
+        if is_sdxl:
+            pattern = r'portrait.sdxl.unnorm\.(safetensors|bin)$'
+        else:
+            raise Exception("portrait unnorm model is not supported for SD1.5")
+        is_insightface = True
+    elif preset == "faceid":
+        if is_sdxl:
+            pattern = r'faceid.sdxl\.(safetensors|bin)$'
+            lora_pattern = r'faceid.sdxl.lora\.safetensors$'
+        else:
+            pattern = r'faceid.sd15\.(safetensors|bin)$'
+            lora_pattern = r'faceid.sd15.lora\.safetensors$'
+        is_insightface = True
+    elif preset.startswith("faceid plus -"):
+        if is_sdxl:
+            raise Exception("faceid plus model is not supported for SDXL")
+        pattern = r'faceid.plus.sd15\.(safetensors|bin)$'
+        lora_pattern = r'faceid.plus.sd15.lora\.safetensors$'
+        is_insightface = True
+    elif preset.startswith("faceid plus v2"):
+        if is_sdxl:
+            pattern = r'faceid.plusv2.sdxl\.(safetensors|bin)$'
+            lora_pattern = r'faceid.plusv2.sdxl.lora\.safetensors$'
+        else:
+            pattern = r'faceid.plusv2.sd15\.(safetensors|bin)$'
+            lora_pattern = r'faceid.plusv2.sd15.lora\.safetensors$'
+        is_insightface = True
+    # Community's models
+    elif preset.startswith("composition"):
+        if is_sdxl:
+            pattern = r'plus.composition.sdxl\.safetensors$'
+        else:
+            pattern = r'plus.composition.sd15\.safetensors$'
+    else:
+        raise Exception(f"invalid type '{preset}'")
+
+    ipadapter_file = [e for e in ipadapter_list if re.search(pattern, e, re.IGNORECASE)]
+    ipadapter_file = folder_paths.get_full_path("ipadapter", ipadapter_file[0]) if ipadapter_file else None
+
+    return ipadapter_file, is_insightface, lora_pattern
+
+def get_lora_file(pattern):
+    lora_list = folder_paths.get_filename_list("loras")
+    lora_file = [e for e in lora_list if re.search(pattern, e, re.IGNORECASE)]
+    lora_file = folder_paths.get_full_path("loras", lora_file[0]) if lora_file else None
+
+    return lora_file
+
+def ipadapter_model_loader(file):
+    model = comfy.utils.load_torch_file(file, safe_load=True)
+
+    if file.lower().endswith(".safetensors"):
+        st_model = {"image_proj": {}, "ip_adapter": {}}
+        for key in model.keys():
+            if key.startswith("image_proj."):
+                st_model["image_proj"][key.replace("image_proj.", "")] = model[key]
+            elif key.startswith("ip_adapter."):
+                st_model["ip_adapter"][key.replace("ip_adapter.", "")] = model[key]
+        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))
+
+    if 'plusv2' in file.lower():
+        model["faceidplusv2"] = True
+    
+    if 'unnorm' in file.lower():
+        model["portraitunnorm"] = True
+
+    return model
+
+def insightface_loader(provider):
+    try:
+        from insightface.app import FaceAnalysis
+    except ImportError as e:
+        raise Exception(e)
+
+    path = os.path.join(folder_paths.models_dir, "insightface")
+    model = FaceAnalysis(name="antelopev2", root=path, providers=[provider + 'ExecutionProvider',])
+    model.prepare(ctx_id=0, det_size=(640, 640))
+    return model
+
+def encode_image_masked(clip_vision, image, mask=None, batch_size=0, size=224):
+    model_management.load_model_gpu(clip_vision.patcher)
+    outputs = Output()
+
+    if batch_size == 0:
+        batch_size = image.shape[0]
+    elif batch_size > image.shape[0]:
+        batch_size = image.shape[0]
+
+    image_batch = torch.split(image, batch_size, dim=0)
+
+    for img in image_batch:
+        img = img.to(clip_vision.load_device)
+
+        pixel_values = clip_preprocess(img.to(clip_vision.load_device), size=size).float()
+
+        # TODO: support for multiple masks
+        if mask is not None:
+            pixel_values = pixel_values * mask.to(clip_vision.load_device)
+
+        out = clip_vision.model(pixel_values=pixel_values, intermediate_output=-2)
+
+        if not hasattr(outputs, "last_hidden_state"):
+            outputs["last_hidden_state"] = out[0].to(model_management.intermediate_device())
+            outputs["image_embeds"] = out[2].to(model_management.intermediate_device())
+            outputs["penultimate_hidden_states"] = out[1].to(model_management.intermediate_device())
+        else:
+            outputs["last_hidden_state"] = torch.cat((outputs["last_hidden_state"], out[0].to(model_management.intermediate_device())), dim=0)
+            outputs["image_embeds"] = torch.cat((outputs["image_embeds"], out[2].to(model_management.intermediate_device())), dim=0)
+            outputs["penultimate_hidden_states"] = torch.cat((outputs["penultimate_hidden_states"], out[1].to(model_management.intermediate_device())), dim=0)
+
+    del img, pixel_values, out
+    torch.cuda.empty_cache()
+
+    return outputs
+
+def tensor_to_size(source, dest_size):
+    if isinstance(dest_size, torch.Tensor):
+        dest_size = dest_size.shape[0]
+    source_size = source.shape[0]
+
+    if source_size < dest_size:
+        shape = [dest_size - source_size] + [1]*(source.dim()-1)
+        source = torch.cat((source, source[-1:].repeat(shape)), dim=0)
+    elif source_size > dest_size:
+        source = source[:dest_size]
+
+    return source
+
+def min_(tensor_list):
+    # return the element-wise min of the tensor list.
+    x = torch.stack(tensor_list)
+    mn = x.min(axis=0)[0]
+    return torch.clamp(mn, min=0)
+
+def max_(tensor_list):
+    # return the element-wise max of the tensor list.
+    x = torch.stack(tensor_list)
+    mx = x.max(axis=0)[0]
+    return torch.clamp(mx, max=1)
+
+# From https://github.com/Jamy-L/Pytorch-Contrast-Adaptive-Sharpening/
+def contrast_adaptive_sharpening(image, amount):
+    img = T.functional.pad(image, (1, 1, 1, 1)).cpu()
+
+    a = img[..., :-2, :-2]
+    b = img[..., :-2, 1:-1]
+    c = img[..., :-2, 2:]
+    d = img[..., 1:-1, :-2]
+    e = img[..., 1:-1, 1:-1]
+    f = img[..., 1:-1, 2:]
+    g = img[..., 2:, :-2]
+    h = img[..., 2:, 1:-1]
+    i = img[..., 2:, 2:]
+
+    # Computing contrast
+    cross = (b, d, e, f, h)
+    mn = min_(cross)
+    mx = max_(cross)
+
+    diag = (a, c, g, i)
+    mn2 = min_(diag)
+    mx2 = max_(diag)
+    mx = mx + mx2
+    mn = mn + mn2
+
+    # Computing local weight
+    inv_mx = torch.reciprocal(mx)
+    amp = inv_mx * torch.minimum(mn, (2 - mx))
+
+    # scaling
+    amp = torch.sqrt(amp)
+    w = - amp * (amount * (1/5 - 1/8) + 1/8)
+    div = torch.reciprocal(1 + 4*w)
+
+    output = ((b + d + f + h)*w + e) * div
+    output = torch.nan_to_num(output)
+    output = output.clamp(0, 1)
+
+    return output
+
+def tensor_to_image(tensor):
+    image = tensor.mul(255).clamp(0, 255).byte().cpu()
+    image = image[..., [2, 1, 0]].numpy()
+    return image
+
+def image_to_tensor(image):
+    tensor = torch.clamp(torch.from_numpy(image).float() / 255., 0, 1)
+    tensor = tensor[..., [2, 1, 0]]
+    return tensor

README.md

@@ -1,6 +1,20 @@
-![image](https://github.com/MinusZoneAI/ComfyUI-Kolors-MZ/blob/main/examples/workflow.png)
+![image](./examples/workflow_ipa.png)
 
-ComfyUI上Kolors的实现
+## Recent changes 
+* [2024-07-27] 支持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的实现
 
 参考自 https://github.com/kijai/ComfyUI-KwaiKolorsWrapper
 
@@ -8,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
 
@@ -28,26 +77,87 @@ 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
-Error occurred when executing MZ_ChatGLM3Loader: 'ChatGLMModel' object has no attribute 'transformer'
+加载模型时出现的错误
++ 目前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)
 
-检查ChatGLM3Loader节点选择的模型是否已经正确下载
+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本体到最新版本(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
 
-尝试使用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节点选择的模型是否已经正确下载
+
+
+## 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
 
@@ -61,7 +171,15 @@ RuntimeError: Only Tensors of floating point dtype can require gradients
 
 
 ## 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,42 +12,25 @@ NODE_CLASS_MAPPINGS = {
 NODE_DISPLAY_NAME_MAPPINGS = {
 }
 
+MAX_RESOLUTION = 16384
+
 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 = folder_paths.get_filename_list("LLM")
+
+        # 筛选safetensors结尾的文件
+        llm_models = [
+            model for model in llm_models if model.endswith("safetensors")]
 
         return {"required": {
             "chatglm3_checkpoint": (llm_models,),
@@ -68,14 +51,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 +69,135 @@ 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_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):
         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 +206,134 @@ 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"
+
+
+class MZ_KolorsCLIPVisionLoader:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"clip_name": (folder_paths.get_filename_list("clip_vision"), ),
+                             }}
+    RETURN_TYPES = ("CLIP_VISION",)
+    FUNCTION = "load_clip"
+
+    CATEGORY = CATEGORY_NAME + "/Legacy"
+
+    def load_clip(self, **kwargs):
+        from . import mz_kolors_core
+        importlib.reload(mz_kolors_core)
+        return mz_kolors_core.MZ_KolorsCLIPVisionLoader_call(kwargs)
+
+
+NODE_CLASS_MAPPINGS["MZ_KolorsCLIPVisionLoader"] = MZ_KolorsCLIPVisionLoader
+NODE_DISPLAY_NAME_MAPPINGS["MZ_KolorsCLIPVisionLoader"] = f"{AUTHOR_NAME} - KolorsCLIPVisionLoader - Legacy"
+
+
+class MZ_ApplySDXLSamplingSettings():
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL", ),
+            }
+        }
+
+    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"
+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
 

clip_vit_336/config.json

@@ -0,0 +1,19 @@
+{
+  "attention_dropout": 0.0,
+  "dropout": 0.0,
+  "hidden_act": "quick_gelu",
+  "hidden_size": 1024,
+  "image_size": 336,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-05,
+  "model_type": "clip_vision_model",
+  "num_attention_heads": 16,
+  "num_channels": 3,
+  "num_hidden_layers": 24,
+  "patch_size": 14,
+  "projection_dim": 768,
+  "torch_dtype": "float32"
+}
+

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,318 @@
+import os
+from types import MethodType
+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
+from comfy.cldm.cldm import ControlNet
+
+# 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):
+        with torch.cuda.amp.autocast(enabled=True):
+            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_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": 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)
+        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}
+
+    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
+        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
+
+
+def load_clipvision_336_from_sd(sd, prefix="", convert_keys=False):
+    from comfy.clip_vision import ClipVisionModel, convert_to_transformers
+
+    json_config = os.path.join(os.path.dirname(
+        os.path.realpath(__file__)), "clip_vit_336", "config.json")
+
+    clip = ClipVisionModel(json_config)
+
+    m, u = clip.load_sd(sd)
+    if len(m) > 0:
+        logging.warning("missing clip vision: {}".format(m))
+    u = set(u)
+    keys = list(sd.keys())
+    for k in keys:
+        if k not in u:
+            t = sd.pop(k)
+            del t
+
+    # def vis_forward(self, pixel_values, attention_mask=None, intermediate_output=None):
+    #     pixel_values = nn.functional.interpolate(
+    #         pixel_values, size=(336, 336), mode='bilinear', align_corners=False)
+    #     x = self.embeddings(pixel_values)
+    #     x = self.pre_layrnorm(x)
+    #     # TODO: attention_mask?
+    #     x, i = self.encoder(
+    #         x, mask=None, intermediate_output=intermediate_output)
+    #     pooled_output = self.post_layernorm(x[:, 0, :])
+    #     return x, i, pooled_output
+
+    # clip.model.vision_model.forward = MethodType(
+    #     vis_forward, clip.model.vision_model
+    # )
+
+    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.cuda.amp.autocast(enabled=True):
+            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
+
+        self.original_load_clipvision_from_sd = comfy.clip_vision.load_clipvision_from_sd
+        comfy.clip_vision.load_clipvision_from_sd = load_clipvision_336_from_sd
+
+        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]
+
+        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
+
+        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
+
+        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,13 +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
 
 
-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 +27,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 +64,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(
@@ -87,8 +80,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(
@@ -100,6 +91,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
@@ -113,6 +105,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)
@@ -122,7 +121,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")
@@ -131,107 +131,264 @@ 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)
-
+    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 load_unet_state_dict(sd):  # load unet in diffusers or regular format
-    from comfy import model_management, model_detection
-    import comfy.utils
+def MZ_ChatGLM3Embeds2Conditioning_call(args):
+    kolors_embeds = args.get("kolors_embeds")
 
-    # 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
+    # kolors_embeds = {
+    #     'prompt_embeds': prompt_embeds,
+    #     'negative_prompt_embeds': negative_prompt_embeds,
+    #     'pooled_prompt_embeds': text_proj,
+    #     'negative_pooled_prompt_embeds': negative_text_proj
+    # }
 
-    parameters = comfy.utils.calculate_parameters(sd)
-    unet_dtype = model_management.unet_dtype(model_params=parameters)
-    load_device = model_management.get_torch_device()
+    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")
+        }
+    ]]
 
-    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
+    negative = [[
+        kolors_embeds['negative_prompt_embeds'],
+        {
+            "pooled_output": kolors_embeds['negative_pooled_prompt_embeds'],
+        }
+    ]]
 
-        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
+    return (positive, negative)
 
 
-def MZ_KolorsUNETLoader_call(kwargs):
+def MZ_KolorsUNETLoaderV2_call(kwargs):
 
-    from . import hook_comfyui
-    with hook_comfyui.apply_kolors():
+    from . import hook_comfyui_kolors_v2
+    import comfy.sd
+
+    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)
-        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)
+def MZ_KolorsCheckpointLoaderSimple_call(kwargs):
+    checkpoint_name = kwargs.get("ckpt_name")
 
-    return ([[
-        cond,
-        {"pooled_output": pool},
-    ]],)
+    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_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 = []
+        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_forward = ControlNet.forward
+
+        def KolorsControlNet_forward(self, x, hint, timesteps, context, **kwargs):
+            with torch.cuda.amp.autocast(enabled=True):
+                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(
+                KolorsControlLora_pre_run, c)
+            return c
+
+        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 = 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 = self.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,)
+
+
+def MZ_KolorsCLIPVisionLoader_call(kwargs):
+    import comfy.clip_vision
+    from . import hook_comfyui_kolors_v2
+    clip_name = kwargs.get("clip_name")
+    clip_path = folder_paths.get_full_path("clip_vision", clip_name)
+    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,)

mz_kolors_legacy.py

@@ -0,0 +1,226 @@
+
+
+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)
+
+    dtype = kwargs.get("dtype")
+    if dtype == "fp16":
+        print("fp16")
+        cond = cond.half()
+        pool = pool.half()
+    elif dtype == "bf16":
+        print("bf16")
+        cond = cond.bfloat16()
+        pool = pool.bfloat16()
+    else:
+        print("fp32")
+        cond = cond.float()
+        pool = pool.float()
+
+    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}),
+                "dtype": ([
+                    "fp32",
+                    "fp16",
+                    "bf16",
+                ],),
+            }
+        }
+
+    RETURN_TYPES = ("CONDITIONING", )
+    RETURN_NAMES = ("prompt", )
+    FUNCTION = "encode"
+    CATEGORY = CATEGORY_NAME
+
+    def encode(self, **kwargs):
+        return MZ_FakeCond_call(kwargs)
+
+
+try:
+    if os.environ.get("MZ_DEV", None) is not None:
+        NODE_CLASS_MAPPINGS["MZ_FakeCond"] = MZ_FakeCond
+        NODE_DISPLAY_NAME_MAPPINGS[
+            "MZ_FakeCond"] = f"{AUTHOR_NAME} - FakeCond"
+except ImportError:
+    pass

mz_kolors_utils.py

@@ -927,63 +927,9 @@ class Utils:
                 m for m in source_model_zoo_json if tags_filter in m["tags"]]
 
         return source_model_zoo_json
+        
 
 
 modelscope_models_map = {
-    "llama3": {
-        "Meta-Llama-3-8B-Instruct.Q4_K_M.gguf": {
-            "url": "https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=Meta-Llama-3-8B-Instruct-GGUF%2FMeta-Llama-3-8B-Instruct.Q4_K_M.gguf",
-            "output": "modelscope_models/Meta-Llama-3-8B-Instruct-GGUF/Meta-Llama-3-8B-Instruct.Q4_K_M.gguf",
-            "SHA256": "647a2b64cbcdbe670432d0502ebb2592b36dd364d51a9ef7a1387b7a4365781f",
-        },
-    },
-    "llava": {
-        "ggml_bakllava-1/ggml-model-q5_k.gguf": {
-            "url": "https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=ggml_bakllava-1%2Fggml-model-q5_k.gguf",
-            "output": "modelscope_models/ggml_bakllava-1/ggml-model-q5_k.gguf",
-            "SHA256": "c93de1376be9b6977cc94d252a3d165d6059e07b528de0fa762534d9599b27d6",
-        },
-        "ggml_llava-v1.5-7b/ggml-model-q4_k.gguf": {
-            "url": "https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=ggml_llava-v1.5-7b%2Fggml-model-q4_k.gguf",
-            "output": "modelscope_models/ggml_llava-v1.5-7b/ggml-model-q4_k.gguf",
-            "SHA256": "7ac9c2f7b8d76cc7f3118cdf0953ebab7a7a9b12bad5dbe237219d2ab61765ea",
-        },
-        "ggml_bakllava-1/mmproj-model-f16.gguf": {
-            "url": "https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=ggml_bakllava-1%2Fmmproj-model-f16.gguf",
-            "output": "modelscope_models/ggml_bakllava-1/mmproj-model-f16.gguf",
-            "SHA256": "2e467eba710002839e0966d5e329942bb836eabd4e787bc713b07eff1d8ea13b"
-        },
-        "ggml_llava-v1.5-7b/mmproj-model-f16.gguf": {
-            "url": "https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=ggml_llava-v1.5-7b%2Fmmproj-model-f16.gguf",
-            "output": "modelscope_models/ggml_llava-v1.5-7b/mmproj-model-f16.gguf",
-            "SHA256": "b7c8ff0f58fca47d28ba92c4443adf8653f3349282cb8d9e6911f22d9b3814fe"
-        },
-        "llava-v1.6-vicuna-13b-gguf/llava-v1.6-vicuna-13b.Q5_K_M.gguf": {
-            "url": "https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=llava-v1.6-vicuna-13b-gguf%2Fllava-v1.6-vicuna-13b.Q5_K_M.gguf",
-            "output": "modelscope_models/llava-v1.6-vicuna-13b-gguf/llava-v1.6-vicuna-13b.Q5_K_M.gguf",
-            "SHA256": "fac8ecd642ac1c2dd8ef8fb6d6daac349293051b468f5fdc21921b94dfb2916a",
-        },
-        "llava-v1.6-vicuna-13b-gguf/mmproj-model-f16.gguf": {
-            "url": "https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=llava-v1.6-vicuna-13b-gguf%2Fmmproj-model-f16.gguf",
-            "output": "modelscope_models/llava-v1.6-vicuna-13b-gguf/mmproj-model-f16.gguf",
-            "SHA256": "9860fbc162d6df40d12f32211aed11ab4826e16e473e6fff595496b9943872e2"
-        },
-        "llava-1.6-mistral-7b-gguf/llava-v1.6-mistral-7b.Q5_K_M.gguf": {
-            "url": "https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=llava-1.6-mistral-7b-gguf%2Fllava-v1.6-mistral-7b.Q5_K_M.gguf",
-            "output": "modelscope_models/llava-1.6-mistral-7b-gguf/llava-v1.6-mistral-7b.Q5_K_M.gguf",
-            "SHA256": "b1d37fc65ecb80aa8f1ce185bf4d7605bc3c5cc5bcc77a160c3a1b0377631112",
-        },
-        "llava-1.6-mistral-7b-gguf/mmproj-model-f16.gguf": {
-            "url": "https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=llava-1.6-mistral-7b-gguf%2Fmmproj-model-f16.gguf",
-            "output": "modelscope_models/llava-1.6-mistral-7b-gguf/mmproj-model-f16.gguf",
-            "SHA256": "00205ee8a0d7a381900cd031e43105f86aa0d8c07bf329851e85c71a26632d16",
-        },
-    },
-    "phi3": {
-        "Phi-3-mini-4k-instruct-q4.gguf": {
-            "url": "https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=Phi-3-mini-4k-instruct-gguf%2FPhi-3-mini-4k-instruct-q4.gguf",
-            "output": "modelscope_models/Phi-3-mini-4k-instruct-gguf/Phi-3-mini-4k-instruct-q4.gguf",
-            "SHA256": "1cd9a9df07350196623f93bf4829cf228959e07ad32f787b8fdd7f5956f5b9de"
-        }
-    },
+   
 }

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 = ""