'update'

.github/workflows/publish.yml

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

.gitignore

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

ComfyUI_IPAdapter_plus/CrossAttentionPatch.py

@@ -1,209 +0,0 @@
-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

@@ -1,275 +0,0 @@
-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

@@ -1,266 +0,0 @@
-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 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="buffalo_l", 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,46 +1,25 @@
-![image](./examples/workflow_ipa.png)
+![image](https://github.com/MinusZoneAI/ComfyUI-Kolors-MZ/blob/main/examples/workflow.png)
 
-## Recent changes 
-* [2024-07-21] 感谢来自yiwangsimple对Mac修复和测试的分支  https://github.com/yiwangsimple/ComfyUI-Kolors-MZ
-* [2024-07-21] 新增MZ_ChatGLM3TextEncodeAdvanceV2节点
-* [2024-07-18] IPA相关节点已在ComfyUI_IPAdapter_plus中支持
-* [2024-07-17] 新增支持IPAdapter_plus的加载器和高级应用节点 MZ_KolorsCLIPVisionLoader,MZ_IPAdapterModelLoaderKolors,MZ_IPAdapterAdvancedKolors
-* [2024-07-14] 删除自动兼容ControlNet, 新增MZ_KolorsControlNetPatch节点
-  ![image](https://github.com/user-attachments/assets/73ae6447-c69d-4781-9c66-94e0029709ed)
-
-## ComfyUI上Kolors的实现
+ComfyUI上Kolors的实现
 
 参考自 https://github.com/kijai/ComfyUI-KwaiKolorsWrapper
 
 使用ComfyUI原生采样
 
 工作流在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
-
-下载地址: https://huggingface.co/Kijai/ChatGLM3-safetensors/resolve/main/chatglm3-fp16.safetensors
+下载地址:
+https://huggingface.co/Kijai/ChatGLM3-safetensors/blob/main/chatglm3-fp16.safetensors
 
 
-### 官方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/
- 
-
-## Implementation of Kolors on ComfyUI
+Implementation of Kolors on ComfyUI
 
 Reference from https://github.com/kijai/ComfyUI-KwaiKolorsWrapper
 
@@ -49,60 +28,26 @@ 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/resolve/main/unet/diffusion_pytorch_model.fp16.safetensors
+https://huggingface.co/Kwai-Kolors/Kolors/blob/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/resolve/main/chatglm3-fp16.safetensors
+https://huggingface.co/Kijai/ChatGLM3-safetensors/tree/main
 
 
-### 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/
 
 
 ## FAQ
+Error occurred when executing MZ_ChatGLM3Loader: 'ChatGLMModel' object has no attribute 'transformer'
 
-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)
+检查ChatGLM3Loader节点选择的模型是否已经正确下载
 
 RuntimeError: Only Tensors of floating point dtype can require gradients
-+ 尝试使用fp16版本的模型: https://huggingface.co/Kijai/ChatGLM3-safetensors/blob/main/chatglm3-fp16.safetensors
 
-Error occurred when executing MZ_ChatGLM3Loader: 'ChatGLMModel' object has no attribute 'transformer'
-+ 检查ChatGLM3Loader节点选择的模型是否已经正确下载
-
-
-## Credits
-
-- [Kolors](https://github.com/Kwai-Kolors/Kolors)
-- [ComfyUI](https://github.com/comfyanonymous/ComfyUI) 
-- [ComfyUI_IPAdapter_plus](https://github.com/cubiq/ComfyUI_IPAdapter_plus)
+尝试使用fp16版本的模型: https://huggingface.co/Kijai/ChatGLM3-safetensors/blob/main/chatglm3-fp16.safetensors
 
 ## Star History
 
@@ -116,15 +61,7 @@ Error occurred when executing MZ_ChatGLM3Loader: 'ChatGLMModel' object has no at
 
 
 ## Contact
-- 微信Wechat: minrszone
+- 绿泡泡: minrszone
 - Bilibili: [minus_zone](https://space.bilibili.com/5950992)
 - 小红书: [MinusZoneAI](https://www.xiaohongshu.com/user/profile/5f072e990000000001005472)
-
-## 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"/>
-
-
+- 爱发电: [MinusZoneAI](https://afdian.net/@MinusZoneAI)

__init__.py

@@ -12,25 +12,42 @@ NODE_CLASS_MAPPINGS = {
 NODE_DISPLAY_NAME_MAPPINGS = {
 }
 
-MAX_RESOLUTION = 16384
-
 AUTHOR_NAME = "MinusZone"
 CATEGORY_NAME = f"{AUTHOR_NAME} - Kolors"
-folder_paths.add_model_folder_path(
-    "LLM", os.path.join(folder_paths.models_dir, "LLM"))
+
+
+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"
 
 
 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 = folder_paths.get_filename_list("LLM")
-
-        # 筛选safetensors结尾的文件
-        llm_models = [
-            model for model in llm_models if model.endswith("safetensors")]
+        llm_models = Utils.listdir_models(llm_dir)
 
         return {"required": {
             "chatglm3_checkpoint": (llm_models,),
@@ -51,13 +68,14 @@ NODE_CLASS_MAPPINGS["MZ_ChatGLM3Loader"] = MZ_ChatGLM3Loader
 NODE_DISPLAY_NAME_MAPPINGS["MZ_ChatGLM3Loader"] = f"{AUTHOR_NAME} - ChatGLM3Loader"
 
 
-class MZ_ChatGLM3TextEncodeV2:
+class MZ_ChatGLM3TextEncode:
     @classmethod
     def INPUT_TYPES(s):
         return {
             "required": {
                 "chatglm3_model": ("CHATGLM3MODEL", ),
                 "text": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+                "hid_proj": ("TorchLinear", ),
             }
         }
 
@@ -69,89 +87,24 @@ class MZ_ChatGLM3TextEncodeV2:
     def encode(self, **kwargs):
         from . import mz_kolors_core
         importlib.reload(mz_kolors_core)
-        return mz_kolors_core.MZ_ChatGLM3TextEncodeV2_call(kwargs)
+        return mz_kolors_core.MZ_ChatGLM3TextEncode_call(kwargs)
 
 
-NODE_CLASS_MAPPINGS["MZ_ChatGLM3_V2"] = MZ_ChatGLM3TextEncodeV2
+NODE_CLASS_MAPPINGS["MZ_ChatGLM3"] = MZ_ChatGLM3TextEncode
 NODE_DISPLAY_NAME_MAPPINGS[
-    "MZ_ChatGLM3_V2"] = f"{AUTHOR_NAME} - ChatGLM3TextEncodeV2"
+    "MZ_ChatGLM3"] = f"{AUTHOR_NAME} - ChatGLM3TextEncode"
 
 
-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_KolorsUNETLoaderV2():
+class MZ_KolorsUNETLoader():
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
                 "unet_name": (folder_paths.get_filename_list("unet"), ),
                 }}
 
-    RETURN_TYPES = ("MODEL",)
+    RETURN_TYPES = ("MODEL", "TorchLinear")
 
-    RETURN_NAMES = ("model",)
+    RETURN_NAMES = ("model", "hid_proj")
 
     FUNCTION = "load_unet"
 
@@ -160,72 +113,9 @@ class MZ_KolorsUNETLoaderV2():
     def load_unet(self, **kwargs):
         from . import mz_kolors_core
         importlib.reload(mz_kolors_core)
-        return mz_kolors_core.MZ_KolorsUNETLoaderV2_call(kwargs)
+        return mz_kolors_core.MZ_KolorsUNETLoader_call(kwargs)
 
 
-NODE_CLASS_MAPPINGS["MZ_KolorsUNETLoaderV2"] = MZ_KolorsUNETLoaderV2
+NODE_CLASS_MAPPINGS["MZ_KolorsUNETLoader"] = MZ_KolorsUNETLoader
 NODE_DISPLAY_NAME_MAPPINGS[
-    "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"
-
-
-from .ComfyUI_IPAdapter_plus.IPAdapterPlus import IPAdapterAdvanced, IPAdapterModelLoader
-
-IPAdapterModelLoader.CATEGORY = CATEGORY_NAME + "/Legacy"
-NODE_CLASS_MAPPINGS["MZ_IPAdapterModelLoaderKolors"] = IPAdapterModelLoader
-NODE_DISPLAY_NAME_MAPPINGS[
-    "MZ_IPAdapterModelLoaderKolors"] = f"IPAdapterModelLoader(kolors) - Legacy"
-
-IPAdapterAdvanced.CATEGORY = CATEGORY_NAME + "/Legacy"
-NODE_CLASS_MAPPINGS["MZ_IPAdapterAdvancedKolors"] = IPAdapterAdvanced
-NODE_DISPLAY_NAME_MAPPINGS[
-    "MZ_IPAdapterAdvancedKolors"] = f"IPAdapterAdvanced(kolors) - Legacy"
-
-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)
+    "MZ_KolorsUNETLoader"] = f"{AUTHOR_NAME} - Kolors UNET Loader"

chatglm3/quantization.py

@@ -15,26 +15,6 @@ 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
@@ -157,11 +137,8 @@ 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)
- 
-        try:
-            self.weight = Parameter(self.weight.to(device), requires_grad=False)
-        except:
-            self.weight.to(device, dtype=self.weight.dtype)
+
+        self.weight = Parameter(self.weight.to(device), requires_grad=False)
         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

@@ -1,19 +0,0 @@
-{
-  "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.py

@@ -12,6 +12,25 @@ 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):
@@ -71,15 +90,8 @@ 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

@@ -1,258 +0,0 @@
-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
-
-
-# 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}
-
-    supported_models = [Kolors]
-
-    for unet_config in supported_models:
-        matches = True
-        for k in match:
-            if match[k] != unet_config[k]:
-                print("key {} does not match".format(
-                    k), match[k], "||", unet_config[k])
-                matches = False
-                break
-        if matches:
-            return convert_config(unet_config)
-    return None
-
-
-import comfy.ldm.modules.diffusionmodules.openaimodel
-from torch import nn
-
-
-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 apply_kolors:
-    def __enter__(self):
-        import comfy.ldm.modules.diffusionmodules.openaimodel
-        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]
-
-    def __exit__(self, type, value, traceback):
-        import comfy.ldm.modules.diffusionmodules.openaimodel
-        import comfy.utils
-        comfy.utils.UNET_MAP_BASIC = self.original_UNET_MAP_BASIC
-
-        model_detection.unet_config_from_diffusers_unet = self.original_unet_config_from_diffusers_unet
-
-        import comfy.supported_models
-        comfy.supported_models.models = self.original_supported_models
-
-        import comfy.clip_vision
-        comfy.clip_vision.load_clipvision_from_sd = self.original_load_clipvision_from_sd

mz_kolors_core.py

@@ -5,16 +5,13 @@ 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 chatglm3_text_encode(chatglm3_model, prompt):
+def KolorsTextEncode(chatglm3_model, prompt):
     device = mm.get_torch_device()
     offload_device = mm.unet_offload_device()
     mm.unload_all_models()
@@ -27,6 +24,14 @@ def chatglm3_text_encode(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']
@@ -64,13 +69,15 @@ def chatglm3_text_encode(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 = folder_paths.get_full_path(
-        'LLM', chatglm3_checkpoint)
+    chatglm3_checkpoint_path = os.path.join(llm_dir, chatglm3_checkpoint)
 
     if not os.path.exists(chatglm3_checkpoint_path):
         raise RuntimeError(
@@ -80,6 +87,8 @@ def MZ_ChatGLM3Loader_call(args):
     from .chatglm3.modeling_chatglm import ChatGLMModel
     from .chatglm3.tokenization_chatglm import ChatGLMTokenizer
 
+    from .mz_kolors_utils import Utils
+
     offload_device = mm.unet_offload_device()
 
     text_encoder_config = os.path.join(
@@ -91,7 +100,6 @@ 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
@@ -105,13 +113,6 @@ 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)
@@ -121,8 +122,7 @@ def MZ_ChatGLM3Loader_call(args):
             set_module_tensor_to_device(
                 text_encoder, key, device=offload_device, value=text_encoder_sd[key])
     else:
-        print("WARNING: Accelerate not available, use load_state_dict load model")
-        text_encoder.load_state_dict(text_encoder_sd)
+        text_encoder.load_state_dict()
 
     tokenizer_path = os.path.join(
         os.path.dirname(__file__), 'configs', "tokenizer")
@@ -131,162 +131,115 @@ def MZ_ChatGLM3Loader_call(args):
     return ({"text_encoder": text_encoder, "tokenizer": tokenizer},)
 
 
-def MZ_ChatGLM3TextEncodeV2_call(args):
+def MZ_ChatGLM3TextEncode_call(args):
+
     text = args.get("text")
     chatglm3_model = args.get("chatglm3_model")
-    prompt_embeds, pooled_output = chatglm3_text_encode(
+
+    prompt_embeds, pooled_output = KolorsTextEncode(
         chatglm3_model,
         text,
     )
-    extra_kwargs = {
-        "pooled_output": pooled_output,
-    }
-    extra_cond_keys = [
-        "width",
-        "height",
-        "crop_w",
-        "crop_h",
-        "target_width",
-        "target_height"
-    ]
-    for key, value in args.items():
-        if key in extra_cond_keys:
-            extra_kwargs[key] = value
+
+    from torch import nn
+    hid_proj: nn.Linear = args.get("hid_proj")
+
+    prompt_embeds = hid_proj(prompt_embeds)
+
     return ([[
         prompt_embeds,
-        # {"pooled_output": pooled_output},
-        extra_kwargs
+        {"pooled_output": pooled_output},
     ]], )
 
 
-def MZ_ChatGLM3Embeds2Conditioning_call(args):
-    kolors_embeds = args.get("kolors_embeds")
+def load_unet_state_dict(sd):  # load unet in diffusers or regular format
+    from comfy import model_management, model_detection
+    import comfy.utils
 
-    # kolors_embeds = {
-    #     'prompt_embeds': prompt_embeds,
-    #     'negative_prompt_embeds': negative_prompt_embeds,
-    #     'pooled_prompt_embeds': text_proj,
-    #     'negative_pooled_prompt_embeds': negative_text_proj
-    # }
+    # 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
 
-    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")
-        }
-    ]]
+    parameters = comfy.utils.calculate_parameters(sd)
+    unet_dtype = model_management.unet_dtype(model_params=parameters)
+    load_device = model_management.get_torch_device()
 
-    negative = [[
-        kolors_embeds['negative_prompt_embeds'],
-        {
-            "pooled_output": kolors_embeds['negative_pooled_prompt_embeds'],
-        }
-    ]]
+    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
 
-    return (positive, negative)
+        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_KolorsUNETLoaderV2_call(kwargs):
+def MZ_KolorsUNETLoader_call(kwargs):
 
-    from . import hook_comfyui_kolors_v2
-    import comfy.sd
-
-    with hook_comfyui_kolors_v2.apply_kolors():
+    from . import hook_comfyui
+    # with hook_comfyui.apply_kolors():
+    if True:
         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 = comfy.sd.load_unet_state_dict(sd)
+        comfy.mz_log("add_embedding.linear_1.weight", sd["add_embedding.linear_1.weight"])
+        # add_embedding_linear_1_weight : torch.Size([1280, 5632])
+        add_embedding_linear_1_weight = sd["add_embedding.linear_1.weight"]
+        # new_add_embedding_linear_1_weight : torch.Size([1280, 2816])
+        new_add_embedding_linear_1_weight = add_embedding_linear_1_weight[:, :2816]
+        sd["add_embedding.linear_1.weight"] = new_add_embedding_linear_1_weight
+         
+        model, hid_proj = load_unet_state_dict(sd)
         if model is None:
             raise RuntimeError(
                 "ERROR: Could not detect model type of: {}".format(unet_path))
-
-        return (model, )
+        return (model, hid_proj)
 
 
-from comfy.cldm.cldm import ControlNet
-from comfy.controlnet import ControlLora
+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_KolorsControlNetPatch_call(kwargs):
-    from . import hook_comfyui_kolors_v2
-    model = kwargs.get("model")
-    control_net = kwargs.get("control_net")
-    import comfy.controlnet
-    if isinstance(control_net, ControlLora):
-        del_keys = []
-        for k in control_net.control_weights:
-            if k.startswith("label_emb.0.0."):
-                del_keys.append(k)
-
-        for k in del_keys:
-            control_net.control_weights.pop(k)
-
-        super_pre_run = ControlLora.pre_run
-        super_copy = ControlLora.copy
-
-        super_forward = ControlNet.forward
-
-        def KolorsControlNet_forward(self, x, hint, timesteps, context, **kwargs):
-            with torch.cuda.amp.autocast(enabled=True):
-                context = model.model.diffusion_model.encoder_hid_proj(context)
-                return super_forward(self, x, hint, timesteps, context, **kwargs)
-
-        def KolorsControlLora_pre_run(self, *args, **kwargs):
-            result = super_pre_run(self, *args, **kwargs)
-
-            if hasattr(self, "control_model"):
-                self.control_model.forward = MethodType(
-                    KolorsControlNet_forward, self.control_model)
-            return result
-
-        control_net.pre_run = MethodType(
-            KolorsControlLora_pre_run, control_net)
-
-        def KolorsControlLora_copy(self, *args, **kwargs):
-            c = super_copy(self, *args, **kwargs)
-            c.pre_run = MethodType(
-                KolorsControlLora_pre_run, c)
-            return c
-
-        control_net.copy = MethodType(
-            KolorsControlLora_copy, control_net)
-
-    elif isinstance(control_net, comfy.controlnet.ControlNet):
-        model_label_emb = model.model.diffusion_model.label_emb
-
-        control_net.control_model.label_emb = model_label_emb
-
-        control_net.control_model_wrapped.model.label_emb = model_label_emb
-
-        super_forward = ControlNet.forward
-
-        def KolorsControlNet_forward(self, x, hint, timesteps, context, **kwargs):
-            with torch.cuda.amp.autocast(enabled=True):
-                context = model.model.diffusion_model.encoder_hid_proj(context)
-                return super_forward(self, x, hint, timesteps, context, **kwargs)
-
-        control_net.control_model.forward = MethodType(
-            KolorsControlNet_forward, control_net.control_model)
-
-    else:
-        raise NotImplementedError(
-            f"Type {control_net} not supported for KolorsControlNetPatch")
-
-    return (control_net,)
-
-
-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,)
+    return ([[
+        cond,
+        {"pooled_output": pool},
+    ]],)

mz_kolors_legacy.py

@@ -1,226 +0,0 @@
-
-
-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,9 +927,63 @@ 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

@@ -1,14 +0,0 @@
-[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 = ""