Add easy kolorsLoader

.gitignore

@@ -7,6 +7,7 @@ wildcards/**
 styles/**
 workflow/**
 autocomplete/**
+web_beta/**
 docs/**
 .vscode/
 .idea/

README.en.md

@@ -43,6 +43,10 @@ Double-click install.bat to install the required dependencies
 
 ## Changelog
 
+**v1.2.1**
+
+- Added `easy kolorsLoader` - Code based on [MinusZoneAI](https://github.com/MinusZoneAI/ComfyUI-Kolors-MZ)'s and [kijai](https://github.com/kijai/ComfyUI-KwaiKolorsWrapper)'s repo, thanks for their contribution.
+
 **v1.2.0**
 
 - Added `easy pulIDApply` and `easy pulIDApplyADV`

README.md

@@ -50,6 +50,10 @@ git clone https://github.com/yolain/ComfyUI-Easy-Use
 
 ## 更新日志
 
+**v1.2.1**
+
+- 增加 `easy kolorsLoader` - 可灵加载器，参考了 [MinusZoneAI](https://github.com/MinusZoneAI/ComfyUI-Kolors-MZ) 和 [kijai](https://github.com/kijai/ComfyUI-KwaiKolorsWrapper) 的代码。
+
 **v1.2.0**
 
 - 增加 `easy pulIDApply` 和 `easy pulIDApplyADV`
@@ -377,6 +381,7 @@ git clone https://github.com/yolain/ComfyUI-Easy-Use
 | easy imageChooser              | [cg-image-picker](https://github.com/chrisgoringe/cg-image-picker) | Preview Chooser         | 
 | easy styleAlignedBatchAlign    | [style_aligned_comfy](https://github.com/chrisgoringe/cg-image-picker) | styleAlignedBatchAlign  | 
 | easy icLightApply              | [ComfyUI-IC-Light](https://github.com/huchenlei/ComfyUI-IC-Light) | ICLightApply等           |
+| easy kolorsLoader              | [ComfyUI-Kolors-MZ](https://github.com/MinusZoneAI/ComfyUI-Kolors-MZ) | kolorsLoader            |
 
 ## Credits
 

__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.2.0"
+__version__ = "1.2.1"
 
 import os
 import folder_paths

prestartup_script.py

@@ -31,6 +31,7 @@ add_folder_path_and_extensions("mediapipe", [os.path.join(model_path, "mediapipe
 add_folder_path_and_extensions("inpaint", [os.path.join(model_path, "inpaint")], folder_paths.supported_pt_extensions)
 add_folder_path_and_extensions("prompt_generator", [os.path.join(model_path, "prompt_generator")], folder_paths.supported_pt_extensions)
 add_folder_path_and_extensions("t5", [os.path.join(model_path, "t5")], folder_paths.supported_pt_extensions)
+add_folder_path_and_extensions("llm", [os.path.join(model_path, "LLM")], folder_paths.supported_pt_extensions)
 
 add_folder_path_and_extensions("checkpoints_thumb", [os.path.join(model_path, "checkpoints")], image_suffixs)
 add_folder_path_and_extensions("loras_thumb", [os.path.join(model_path, "loras")], image_suffixs)

py/easyNodes.py

@@ -1739,8 +1739,6 @@ class dynamiCrafterLoader(DynamiCrafter):
         model_path = get_local_filepath(DYNAMICRAFTER_MODELS[model_name]['model_url'], DYNAMICRAFTER_DIR)
         model_patcher, image_proj_model = self.load_dynamicrafter(model_path)
 
-        # rescale cfg
-
         # apply
         model, empty_latent, image_latent = self.process_image_conditioning(model_patcher, clip_vision, vae, image_proj_model, init_image, use_interpolate, fps, frames, scale_latents)
 
@@ -1785,6 +1783,117 @@ class dynamiCrafterLoader(DynamiCrafter):
 
         return (pipe, model, vae)
 
+# kolors Loader
+from .kolors.text_encode import chatglm3_adv_text_encode
+hid_proj = None
+class kolorsLoader:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required":{
+                "unet_name": (folder_paths.get_filename_list("unet"),),
+                "vae_name": (folder_paths.get_filename_list("vae"),),
+                "chatglm3_name": (folder_paths.get_filename_list("llm"),),
+                "lora_name": (["None"] + folder_paths.get_filename_list("loras"),),
+                "lora_model_strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
+                "lora_clip_strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
+                "resolution": (resolution_strings, {"default": "1024 x 576"}),
+                "empty_latent_width": ("INT", {"default": 1024, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
+                "empty_latent_height": ("INT", {"default": 1024, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
+
+                "positive": ("STRING", {"default": "", "placeholder": "Positive", "multiline": True}),
+                "negative": ("STRING", {"default": "", "placeholder": "Negative", "multiline": True}),
+
+                "batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
+            },
+            "optional": {
+                "model_override": ("MODEL",),
+                "optional_lora_stack": ("LORA_STACK",),
+            },
+            "hidden": {"prompt": "PROMPT", "my_unique_id": "UNIQUE_ID"}
+        }
+
+    RETURN_TYPES = ("PIPE_LINE", "MODEL", "VAE")
+    RETURN_NAMES = ("pipe", "model", "vae")
+
+    FUNCTION = "adv_pipeloader"
+    CATEGORY = "EasyUse/Loaders"
+
+    def adv_pipeloader(self, unet_name, vae_name, chatglm3_name, lora_name, lora_model_strength, lora_clip_strength, resolution, empty_latent_width, empty_latent_height, positive, negative, batch_size, model_override=None, optional_lora_stack=None, prompt=None, my_unique_id=None):
+        # load unet
+        global hid_proj
+        if model_override:
+           model = model_override
+        else:
+           model, _hid_proj = easyCache.load_kolors_unet(unet_name)
+           if hid_proj is None:
+             hid_proj = _hid_proj
+        # load vae
+        vae = easyCache.load_vae(vae_name)
+        # load chatglm3
+        chatglm3_model = easyCache.load_chatglm3(chatglm3_name)
+        # load lora
+        lora_stack = []
+        if optional_lora_stack is not None:
+            for lora in optional_lora_stack:
+                lora = {"lora_name": lora[0], "model": model, "clip": None, "model_strength": lora[1],
+                        "clip_strength": lora[2]}
+                model, _ = easyCache.load_lora(lora)
+                lora['model'] = model
+                lora['clip'] = None
+                lora_stack.append(lora)
+
+        if lora_name != "None":
+            lora = {"lora_name": lora_name, "model": model, "clip": None, "model_strength": lora_model_strength,
+                    "clip_strength": lora_clip_strength}
+            model, _ = easyCache.load_lora(lora)
+            lora_stack.append(lora)
+
+
+        # text encode
+        positive_embeddings_final = chatglm3_adv_text_encode(chatglm3_model, positive, hid_proj)
+        negative_embeddings_final = chatglm3_adv_text_encode(chatglm3_model, negative, hid_proj)
+
+        # empty latent
+        samples = sampler.emptyLatent(resolution, empty_latent_width, empty_latent_height, batch_size)
+
+        log_node_warn("加载完毕...")
+        pipe = {
+            "model": model,
+            "positive": positive_embeddings_final,
+            "negative": negative_embeddings_final,
+            "vae": vae,
+            "clip": None,
+
+            "samples": samples,
+            "images": None,
+
+            "loader_settings": {
+                "unet_name": unet_name,
+                "vae_name": vae_name,
+                "chatglm3_name": chatglm3_name,
+
+                "lora_name": lora_name,
+                "lora_model_strength": lora_model_strength,
+                "lora_clip_strength": lora_clip_strength,
+
+                "positive": positive,
+                "negative": negative,
+                "resolution": resolution,
+                "empty_latent_width": empty_latent_width,
+                "empty_latent_height": empty_latent_height,
+                "batch_size": batch_size,
+            }
+        }
+
+        return {"ui": {},
+                "result": (pipe, model, vae, chatglm3_model, positive_embeddings_final, negative_embeddings_final, samples)}
+
+
+        return (chatglm3_model, None, None)
+
+
 # Dit Loader
 from .dit.utils import string_to_dtype
 from .dit.hunyuanDiT.config import hydit_conf, dtypes, devices
@@ -7326,6 +7435,7 @@ NODE_CLASS_MAPPINGS = {
     "easy zero123Loader": zero123Loader,
     "easy dynamiCrafterLoader": dynamiCrafterLoader,
     "easy cascadeLoader": cascadeLoader,
+    "easy kolorsLoader": kolorsLoader,
     "easy hunyuanDiTLoader": hunyuanDiTLoader,
     "easy pixArtLoader": pixArtLoader,
     "easy loraStack": loraStack,
@@ -7443,6 +7553,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "easy zero123Loader": "EasyLoader (Zero123)",
     "easy dynamiCrafterLoader": "EasyLoader (DynamiCrafter)",
     "easy cascadeLoader": "EasyCascadeLoader",
+    "easy kolorsLoader": "EasyLoader (Kolors)",
     "easy hunyuanDiTLoader": "EasyLoader (HunYuanDiT)",
     "easy pixArtLoader": "EasyLoader (PixArt)",
     "easy loraStack": "EasyLoraStack",

py/kolors/chatglm/config_chatglm.json

@@ -0,0 +1,42 @@
+{
+  "_name_or_path": "THUDM/chatglm3-6b-base",
+  "model_type": "chatglm",
+  "architectures": [
+    "ChatGLMModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_chatglm.ChatGLMConfig",
+    "AutoModel": "modeling_chatglm.ChatGLMForConditionalGeneration",
+    "AutoModelForCausalLM": "modeling_chatglm.ChatGLMForConditionalGeneration",
+    "AutoModelForSeq2SeqLM": "modeling_chatglm.ChatGLMForConditionalGeneration",
+    "AutoModelForSequenceClassification": "modeling_chatglm.ChatGLMForSequenceClassification"
+  },
+  "add_bias_linear": false,
+  "add_qkv_bias": true,
+  "apply_query_key_layer_scaling": true,
+  "apply_residual_connection_post_layernorm": false,
+  "attention_dropout": 0.0,
+  "attention_softmax_in_fp32": true,
+  "bias_dropout_fusion": true,
+  "ffn_hidden_size": 13696,
+  "fp32_residual_connection": false,
+  "hidden_dropout": 0.0,
+  "hidden_size": 4096,
+  "kv_channels": 128,
+  "layernorm_epsilon": 1e-05,
+  "multi_query_attention": true,
+  "multi_query_group_num": 2,
+  "num_attention_heads": 32,
+  "num_layers": 28,
+  "original_rope": true,
+  "padded_vocab_size": 65024,
+  "post_layer_norm": true,
+  "rmsnorm": true,
+  "seq_length": 32768,
+  "use_cache": true,
+  "torch_dtype": "float16",
+  "transformers_version": "4.30.2",
+  "tie_word_embeddings": false,
+  "eos_token_id": 2,
+  "pad_token_id": 0
+}

py/kolors/chatglm/configuration_chatglm.py

@@ -0,0 +1,60 @@
+from transformers import PretrainedConfig
+
+class ChatGLMConfig(PretrainedConfig):
+    model_type = "chatglm"
+    def __init__(
+        self,
+        num_layers=28,
+        padded_vocab_size=65024,
+        hidden_size=4096,
+        ffn_hidden_size=13696,
+        kv_channels=128,
+        num_attention_heads=32,
+        seq_length=2048,
+        hidden_dropout=0.0,
+        classifier_dropout=None,
+        attention_dropout=0.0,
+        layernorm_epsilon=1e-5,
+        rmsnorm=True,
+        apply_residual_connection_post_layernorm=False,
+        post_layer_norm=True,
+        add_bias_linear=False,
+        add_qkv_bias=False,
+        bias_dropout_fusion=True,
+        multi_query_attention=False,
+        multi_query_group_num=1,
+        apply_query_key_layer_scaling=True,
+        attention_softmax_in_fp32=True,
+        fp32_residual_connection=False,
+        quantization_bit=0,
+        pre_seq_len=None,
+        prefix_projection=False,
+        **kwargs
+    ):
+        self.num_layers = num_layers
+        self.vocab_size = padded_vocab_size
+        self.padded_vocab_size = padded_vocab_size
+        self.hidden_size = hidden_size
+        self.ffn_hidden_size = ffn_hidden_size
+        self.kv_channels = kv_channels
+        self.num_attention_heads = num_attention_heads
+        self.seq_length = seq_length
+        self.hidden_dropout = hidden_dropout
+        self.classifier_dropout = classifier_dropout
+        self.attention_dropout = attention_dropout
+        self.layernorm_epsilon = layernorm_epsilon
+        self.rmsnorm = rmsnorm
+        self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm
+        self.post_layer_norm = post_layer_norm
+        self.add_bias_linear = add_bias_linear
+        self.add_qkv_bias = add_qkv_bias
+        self.bias_dropout_fusion = bias_dropout_fusion
+        self.multi_query_attention = multi_query_attention
+        self.multi_query_group_num = multi_query_group_num
+        self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
+        self.attention_softmax_in_fp32 = attention_softmax_in_fp32
+        self.fp32_residual_connection = fp32_residual_connection
+        self.quantization_bit = quantization_bit
+        self.pre_seq_len = pre_seq_len
+        self.prefix_projection = prefix_projection
+        super().__init__(**kwargs)

py/kolors/chatglm/tokenization_chatglm.py

@@ -0,0 +1,299 @@
+import json
+import os
+import re
+from typing import List, Optional, Union, Dict
+from sentencepiece import SentencePieceProcessor
+from transformers import PreTrainedTokenizer
+from transformers.utils import logging, PaddingStrategy
+from transformers.tokenization_utils_base import EncodedInput, BatchEncoding
+
+class SPTokenizer:
+    def __init__(self, model_path: str):
+        # reload tokenizer
+        assert os.path.isfile(model_path), model_path
+        self.sp_model = SentencePieceProcessor(model_file=model_path)
+
+        # BOS / EOS token IDs
+        self.n_words: int = self.sp_model.vocab_size()
+        self.bos_id: int = self.sp_model.bos_id()
+        self.eos_id: int = self.sp_model.eos_id()
+        self.pad_id: int = self.sp_model.unk_id()
+        assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
+
+        role_special_tokens = ["<|system|>", "<|user|>", "<|assistant|>", "<|observation|>"]
+        special_tokens = ["[MASK]", "[gMASK]", "[sMASK]", "sop", "eop"] + role_special_tokens
+        self.special_tokens = {}
+        self.index_special_tokens = {}
+        for token in special_tokens:
+            self.special_tokens[token] = self.n_words
+            self.index_special_tokens[self.n_words] = token
+            self.n_words += 1
+        self.role_special_token_expression = "|".join([re.escape(token) for token in role_special_tokens])
+
+    def tokenize(self, s: str, encode_special_tokens=False):
+        if encode_special_tokens:
+            last_index = 0
+            t = []
+            for match in re.finditer(self.role_special_token_expression, s):
+                if last_index < match.start():
+                    t.extend(self.sp_model.EncodeAsPieces(s[last_index:match.start()]))
+                t.append(s[match.start():match.end()])
+                last_index = match.end()
+            if last_index < len(s):
+                t.extend(self.sp_model.EncodeAsPieces(s[last_index:]))
+            return t
+        else:
+            return self.sp_model.EncodeAsPieces(s)
+
+    def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]:
+        assert type(s) is str
+        t = self.sp_model.encode(s)
+        if bos:
+            t = [self.bos_id] + t
+        if eos:
+            t = t + [self.eos_id]
+        return t
+
+    def decode(self, t: List[int]) -> str:
+        text, buffer = "", []
+        for token in t:
+            if token in self.index_special_tokens:
+                if buffer:
+                    text += self.sp_model.decode(buffer)
+                    buffer = []
+                text += self.index_special_tokens[token]
+            else:
+                buffer.append(token)
+        if buffer:
+            text += self.sp_model.decode(buffer)
+        return text
+
+    def decode_tokens(self, tokens: List[str]) -> str:
+        text = self.sp_model.DecodePieces(tokens)
+        return text
+
+    def convert_token_to_id(self, token):
+        """ Converts a token (str) in an id using the vocab. """
+        if token in self.special_tokens:
+            return self.special_tokens[token]
+        return self.sp_model.PieceToId(token)
+
+    def convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        if index in self.index_special_tokens:
+            return self.index_special_tokens[index]
+        if index in [self.eos_id, self.bos_id, self.pad_id] or index < 0:
+            return ""
+        return self.sp_model.IdToPiece(index)
+
+
+class ChatGLMTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "tokenizer.model"}
+
+    model_input_names = ["input_ids", "attention_mask", "position_ids"]
+
+    def __init__(self, vocab_file, padding_side="left", clean_up_tokenization_spaces=False, encode_special_tokens=False,
+                 **kwargs):
+        self.name = "GLMTokenizer"
+
+        self.vocab_file = vocab_file
+        self.tokenizer = SPTokenizer(vocab_file)
+        self.special_tokens = {
+            "<bos>": self.tokenizer.bos_id,
+            "<eos>": self.tokenizer.eos_id,
+            "<pad>": self.tokenizer.pad_id
+        }
+        self.encode_special_tokens = encode_special_tokens
+        super().__init__(padding_side=padding_side, clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+                         encode_special_tokens=encode_special_tokens,
+                         **kwargs)
+
+    def get_command(self, token):
+        if token in self.special_tokens:
+            return self.special_tokens[token]
+        assert token in self.tokenizer.special_tokens, f"{token} is not a special token for {self.name}"
+        return self.tokenizer.special_tokens[token]
+
+    @property
+    def unk_token(self) -> str:
+        return "<unk>"
+
+    @property
+    def pad_token(self) -> str:
+        return "<unk>"
+
+    @property
+    def pad_token_id(self):
+        return self.get_command("<pad>")
+
+    @property
+    def eos_token(self) -> str:
+        return "</s>"
+
+    @property
+    def eos_token_id(self):
+        return self.get_command("<eos>")
+
+    @property
+    def vocab_size(self):
+        return self.tokenizer.n_words
+
+    def get_vocab(self):
+        """ Returns vocab as a dict """
+        vocab = {self._convert_id_to_token(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def _tokenize(self, text, **kwargs):
+        return self.tokenizer.tokenize(text, encode_special_tokens=self.encode_special_tokens)
+
+    def _convert_token_to_id(self, token):
+        """ Converts a token (str) in an id using the vocab. """
+        return self.tokenizer.convert_token_to_id(token)
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        return self.tokenizer.convert_id_to_token(index)
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        return self.tokenizer.decode_tokens(tokens)
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        """
+        Save the vocabulary and special tokens file to a directory.
+
+        Args:
+            save_directory (`str`):
+                The directory in which to save the vocabulary.
+            filename_prefix (`str`, *optional*):
+                An optional prefix to add to the named of the saved files.
+
+        Returns:
+            `Tuple(str)`: Paths to the files saved.
+        """
+        if os.path.isdir(save_directory):
+            vocab_file = os.path.join(
+                save_directory, self.vocab_files_names["vocab_file"]
+            )
+        else:
+            vocab_file = save_directory
+
+        with open(self.vocab_file, 'rb') as fin:
+            proto_str = fin.read()
+
+        with open(vocab_file, "wb") as writer:
+            writer.write(proto_str)
+
+        return (vocab_file,)
+
+    def get_prefix_tokens(self):
+        prefix_tokens = [self.get_command("[gMASK]"), self.get_command("sop")]
+        return prefix_tokens
+
+    def build_single_message(self, role, metadata, message):
+        assert role in ["system", "user", "assistant", "observation"], role
+        role_tokens = [self.get_command(f"<|{role}|>")] + self.tokenizer.encode(f"{metadata}\n")
+        message_tokens = self.tokenizer.encode(message)
+        tokens = role_tokens + message_tokens
+        return tokens
+
+    def build_chat_input(self, query, history=None, role="user"):
+        if history is None:
+            history = []
+        input_ids = []
+        for item in history:
+            content = item["content"]
+            if item["role"] == "system" and "tools" in item:
+                content = content + "\n" + json.dumps(item["tools"], indent=4, ensure_ascii=False)
+            input_ids.extend(self.build_single_message(item["role"], item.get("metadata", ""), content))
+        input_ids.extend(self.build_single_message(role, "", query))
+        input_ids.extend([self.get_command("<|assistant|>")])
+        return self.batch_encode_plus([input_ids], return_tensors="pt", is_split_into_words=True)
+
+    def build_inputs_with_special_tokens(
+            self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+        adding special tokens. A BERT sequence has the following format:
+
+        - single sequence: `[CLS] X [SEP]`
+        - pair of sequences: `[CLS] A [SEP] B [SEP]`
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs to which the special tokens will be added.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+        """
+        prefix_tokens = self.get_prefix_tokens()
+        token_ids_0 = prefix_tokens + token_ids_0
+        if token_ids_1 is not None:
+            token_ids_0 = token_ids_0 + token_ids_1 + [self.get_command("<eos>")]
+        return token_ids_0
+
+    def _pad(
+            self,
+            encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
+            max_length: Optional[int] = None,
+            padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
+            pad_to_multiple_of: Optional[int] = None,
+            return_attention_mask: Optional[bool] = None,
+    ) -> dict:
+        """
+        Pad encoded inputs (on left/right and up to predefined length or max length in the batch)
+
+        Args:
+            encoded_inputs:
+                Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`).
+            max_length: maximum length of the returned list and optionally padding length (see below).
+                Will truncate by taking into account the special tokens.
+            padding_strategy: PaddingStrategy to use for padding.
+
+                - PaddingStrategy.LONGEST Pad to the longest sequence in the batch
+                - PaddingStrategy.MAX_LENGTH: Pad to the max length (default)
+                - PaddingStrategy.DO_NOT_PAD: Do not pad
+                The tokenizer padding sides are defined in self.padding_side:
+
+                    - 'left': pads on the left of the sequences
+                    - 'right': pads on the right of the sequences
+            pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value.
+                This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability
+                `>= 7.5` (Volta).
+            return_attention_mask:
+                (optional) Set to False to avoid returning attention mask (default: set to model specifics)
+        """
+        # Load from model defaults
+        assert self.padding_side == "left"
+
+        required_input = encoded_inputs[self.model_input_names[0]]
+        seq_length = len(required_input)
+
+        if padding_strategy == PaddingStrategy.LONGEST:
+            max_length = len(required_input)
+
+        if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
+            max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of
+
+        needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length
+
+        # Initialize attention mask if not present.
+        if "attention_mask" not in encoded_inputs:
+            encoded_inputs["attention_mask"] = [1] * seq_length
+
+        if "position_ids" not in encoded_inputs:
+            encoded_inputs["position_ids"] = list(range(seq_length))
+
+        if needs_to_be_padded:
+            difference = max_length - len(required_input)
+
+            if "attention_mask" in encoded_inputs:
+                encoded_inputs["attention_mask"] = [0] * difference + encoded_inputs["attention_mask"]
+            if "position_ids" in encoded_inputs:
+                encoded_inputs["position_ids"] = [0] * difference + encoded_inputs["position_ids"]
+            encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input
+
+        return encoded_inputs

py/kolors/chatglm/tokenizer/tokenizer_config.json

@@ -0,0 +1,12 @@
+{
+  "name_or_path": "THUDM/chatglm3-6b-base",
+  "remove_space": false,
+  "do_lower_case": false,
+  "tokenizer_class": "ChatGLMTokenizer",
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_chatglm.ChatGLMTokenizer",
+      null
+      ]
+  }
+}

py/kolors/loader.py

@@ -0,0 +1,203 @@
+import json
+import os
+import torch
+import comfy.model_management
+import comfy.model_detection as model_detection
+import comfy.supported_models
+from comfy.utils import load_torch_file
+from .chatglm.modeling_chatglm import ChatGLMModel, ChatGLMConfig
+from .chatglm.tokenization_chatglm import ChatGLMTokenizer
+
+class Kolors(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 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")
+
+
+if Kolors not in comfy.supported_models.models:
+    comfy.supported_models.models += [Kolors]
+
+class applyKolorsUnet:
+    def __enter__(self):
+        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
+
+    def __exit__(self, type, value, traceback):
+        model_detection.unet_config_from_diffusers_unet = self.original_unet_config_from_diffusers_unet
+
+def kolors_unet_config_from_diffusers_unet(state_dict, dtype=None):
+    match = {}
+    transformer_depth = []
+
+    attn_res = 1
+    count_blocks = model_detection.count_blocks
+    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 model_detection.convert_config(unet_config)
+    return None
+
+def load_kolors_unet_state_dict(sd):
+    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 = comfy.model_management.unet_dtype(model_params=parameters)
+    load_device = comfy.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))
+                pass
+        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 = comfy.model_management.unet_offload_device()
+    unet_dtype = comfy.model_management.unet_dtype(
+        model_params=parameters, supported_dtypes=model_config.supported_inference_dtypes)
+    manual_cast_dtype = comfy.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
+
+
+class chatGLM3Model(torch.nn.Module):
+    def __init__(self, textmodel_json_config=None, device='cpu', offload_device='cpu', model_path=None):
+        super().__init__()
+        if model_path is None:
+            raise ValueError("model_path is required")
+        self.device = device
+        if textmodel_json_config is None:
+            textmodel_json_config = os.path.join(
+                os.path.dirname(os.path.realpath(__file__)),
+                "chatglm",
+                "config_chatglm.json"
+            )
+        with open(textmodel_json_config, 'r') as file:
+            config = json.load(file)
+        textmodel_json_config = ChatGLMConfig(**config)
+        is_accelerate_available = False
+        try:
+            from accelerate import init_empty_weights
+            from accelerate.utils import set_module_tensor_to_device
+            is_accelerate_available = True
+        except:
+            pass
+
+        from contextlib import nullcontext
+        with (init_empty_weights() if is_accelerate_available else nullcontext()):
+            self.text_encoder = ChatGLMModel(textmodel_json_config)
+            if '4bit' in model_path:
+                self.text_encoder.quantize(4)
+            elif '8bit' in model_path:
+                self.text_encoder.quantize(8)
+
+        sd = load_torch_file(model_path)
+        if is_accelerate_available:
+            for key in sd:
+                set_module_tensor_to_device(self.text_encoder, key, device=offload_device, value=sd[key])
+        else:
+            self.text_encoder.load_state_dict()
+
+def load_chatglm3(model_path=None):
+    if model_path is None:
+        return
+
+    load_device = comfy.model_management.text_encoder_device()
+    offload_device = comfy.model_management.text_encoder_offload_device()
+
+    glm3model = chatGLM3Model(
+        device=load_device,
+        offload_device=offload_device,
+        model_path=model_path
+    )
+    tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'chatglm', "tokenizer")
+    tokenizer = ChatGLMTokenizer.from_pretrained(tokenizer_path)
+    text_encoder = glm3model.text_encoder
+    return {"text_encoder":text_encoder, "tokenizer":tokenizer}

py/kolors/text_encode.py

@@ -0,0 +1,92 @@
+import re
+import gc
+import folder_paths
+import comfy.model_management as mm
+from nodes import ConditioningConcat, ConditioningZeroOut, ConditioningSetTimestepRange, ConditioningCombine
+
+def chatglm3_text_encode(chatglm3_model, prompt, hid_proj):
+    device = mm.get_torch_device()
+    offload_device = mm.unet_offload_device()
+    mm.unload_all_models()
+    mm.soft_empty_cache()
+    # Function to randomly select an option from the brackets
+
+    def choose_random_option(match):
+        options = match.group(1).split('|')
+        return random.choice(options)
+
+    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']
+    text_encoder.to(device)
+    text_inputs = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=256,
+        truncation=True,
+        return_tensors="pt",
+    ).to(device)
+
+    output = text_encoder(
+        input_ids=text_inputs['input_ids'],
+        attention_mask=text_inputs['attention_mask'],
+        position_ids=text_inputs['position_ids'],
+        output_hidden_states=True)
+
+    # [batch_size, 77, 4096]
+    prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone()
+    text_proj = output.hidden_states[-1][-1, :, :].clone()  # [batch_size, 4096]
+    bs_embed, seq_len, _ = prompt_embeds.shape
+    prompt_embeds = prompt_embeds.repeat(1, 1, 1)
+    prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1)
+
+    bs_embed = text_proj.shape[0]
+    text_proj = text_proj.repeat(1, 1).view(bs_embed, -1)
+    text_encoder.to(offload_device)
+    mm.soft_empty_cache()
+    gc.collect()
+
+    if hid_proj:
+        prompt_embeds = hid_proj(prompt_embeds)
+    return [[prompt_embeds, {"pooled_output": text_proj},]]
+
+def chatglm3_adv_text_encode(chatglm3_model, text, hid_proj):
+    time_start = 0
+    time_end = 1
+    match = re.search(r'TIMESTEP.*$', text)
+    if match:
+        timestep = match.group()
+        timestep = timestep.split(' ')
+        timestep = timestep[0]
+        text = text.replace(timestep, '')
+        value = timestep.split(':')
+        if len(value) >= 3:
+            time_start = float(value[1])
+            time_end = float(value[2])
+        elif len(value) == 2:
+            time_start = float(value[1])
+            time_end = 1
+        elif len(value) == 1:
+            time_start = 0.1
+            time_end = 1
+
+    conditioning = chatglm3_text_encode(chatglm3_model, text, hid_proj)
+
+    # setTimeStepRange
+    if time_start > 0 or time_end < 1:
+        conditioning_2, = ConditioningSetTimestepRange().set_range(conditioning, 0, time_start)
+        conditioning_1, = ConditioningZeroOut().zero_out(conditioning)
+        conditioning_1, = ConditioningSetTimestepRange().set_range(conditioning_1, time_start, time_end)
+        conditioning, = ConditioningCombine().combine(conditioning_1, conditioning_2)
+
+    return conditioning

py/libs/loader.py

@@ -8,6 +8,7 @@ from comfy.model_patcher import ModelPatcher
 from nodes import NODE_CLASS_MAPPINGS
 from collections import defaultdict
 from .log import log_node_info, log_node_error
+from ..kolors.loader import load_chatglm3, applyKolorsUnet, load_kolors_unet_state_dict
 from ..dit.hunyuanDiT.loader import EXM_HyDiT_Tenc_Temp, load_hydit
 from ..dit.pixArt.loader import load_pixart
 
@@ -32,6 +33,7 @@ class easyLoader:
             "lora": defaultdict(dict),  # {lora_name: {UID: (model_lora, clip_lora)}}
             "controlnet": defaultdict(dict),
             "t5": defaultdict(tuple),
+            "chatglm3": defaultdict(tuple),
         }
         self.memory_threshold = self.determine_memory_threshold(0.7)
         self.lora_name_cache = []
@@ -91,6 +93,7 @@ class easyLoader:
         desired_lora_settings = set()
         desired_controlnet_names = set()
         desired_t5_names = set()
+        desired_glm3_names = set()
 
         for entry in prompt.values():
             class_type = entry["class_type"]
@@ -104,6 +107,11 @@ class easyLoader:
                 desired_ckpt_names.add(self.get_input_value(entry, "ckpt_name", prompt))
                 desired_vae_names.add(self.get_input_value(entry, "vae_name"))
 
+            elif class_type in ['easy kolorsLoader']:
+                desired_unet_names.add(self.get_input_value(entry, "unet_name"))
+                desired_vae_names.add(self.get_input_value(entry, "vae_name"))
+                desired_glm3_names.add(self.get_input_value(entry, "chatglm3_name"))
+
             elif class_type in dit_loaders:
                 t5_name = self.get_input_value(entry, "mt5_name") if "mt5_name" in entry["inputs"] else  None
                 clip_name = self.get_input_value(entry, "clip_name") if "clip_name" in entry["inputs"] else None
@@ -161,6 +169,8 @@ class easyLoader:
                 desired_names = desired_controlnet_names
             elif object_type == "t5":
                 desired_names = desired_t5_names
+            elif object_type == "chatglm3":
+                desired_names = desired_glm3_names
             else:
                 desired_names = desired_lora_names
             self.clear_unused_objects(desired_names, object_type)
@@ -199,7 +209,7 @@ class easyLoader:
         current_memory = self.get_memory_usage()
         if current_memory < self.memory_threshold:
             return
-        eviction_order = ["vae", "lora", "bvae", "clip", "ckpt", "controlnet"]
+        eviction_order = ["vae", "lora", "bvae", "clip", "ckpt", "controlnet", "unet", "t5", "chatglm3"]
         for obj_type in eviction_order:
             if current_memory < self.memory_threshold:
                 break
@@ -319,11 +329,14 @@ class easyLoader:
         lbw_b = lora["lbw_b"] if "lbw_b" in lora else None
 
         model_hash = str(model)[44:-1]
-        clip_hash = str(clip)[25:-1]
+        clip_hash = str(clip)[25:-1] if clip else ''
 
         unique_id = f'{model_hash};{clip_hash};{lora_name};{model_strength};{clip_strength}'
 
-        if unique_id in self.loaded_objects["lora"] and unique_id in self.loaded_objects["lora"][lora_name]:
+        print(unique_id)
+        print(self.loaded_objects["lora"])
+        if unique_id in self.loaded_objects["lora"]:
+            log_node_info("Load LORA",f"{lora_name} cached")
             return self.loaded_objects["lora"][unique_id][0]
 
         orig_lora_name = lora_name
@@ -466,6 +479,32 @@ class easyLoader:
 
         return model, clip, vae, clip_vision, lora_stack
 
+    # Kolors
+    def load_kolors_unet(self, unet_name):
+        if unet_name in self.loaded_objects["unet"]:
+            return (self.loaded_objects["unet"][unet_name][0], None)
+        else:
+            with applyKolorsUnet():
+
+                unet_path = folder_paths.get_full_path("unet", unet_name)
+                sd = comfy.utils.load_torch_file(unet_path)
+                model, hid_proj = load_kolors_unet_state_dict(sd)
+                self.add_to_cache("unet", unet_name, model)
+                self.eviction_based_on_memory()
+
+                return model, hid_proj
+
+    def load_chatglm3(self, chatglm3_name):
+        if chatglm3_name in self.loaded_objects["chatglm3"]:
+            return self.loaded_objects["chatglm3"][chatglm3_name][0]
+
+        chatglm_model = load_chatglm3(model_path=folder_paths.get_full_path("llm", chatglm3_name))
+        self.add_to_cache("chatglm3", chatglm3_name, chatglm_model)
+        self.eviction_based_on_memory()
+
+        return chatglm_model
+
+
     # DiT
     def load_dit_ckpt(self, ckpt_name, model_name, **kwargs):
         if (ckpt_name+'_'+model_name) in self.loaded_objects["ckpt"]:

web/js/easy/easyDynamicWidgets.js

@@ -694,6 +694,7 @@ app.registerExtension({
 			case "easy dynamiCrafterLoader":
 			case "easy hunyuanDiTLoader":
 			case "easy pixArtLoader":
+			case "easy kolorsLoader":
 			case "easy loraStack":
 			case "easy controlnetStack":
 			case "easy latentNoisy":

web/js/easy/easyInterface.js

@@ -113,13 +113,13 @@ try{
     }
     let theme_name = localStorage.getItem('Comfy.Settings.Comfy.ColorPalette')
     control_mode = localStorage.getItem('Comfy.Settings.Comfy.WidgetControlMode')
-    if(control_mode) {
-        control_mode = JSON.parse(control_mode)
-        if(control_mode == 'before'){
-            localStorage['Comfy.Settings.AE.mouseover'] = false
-            localStorage['Comfy.Settings.AE.highlight'] = false
-        }
-    }
+    // if(control_mode) {
+    //     control_mode = JSON.parse(control_mode)
+    //     if(control_mode == 'before'){
+    //         localStorage['Comfy.Settings.AE.mouseover'] = false
+    //         localStorage['Comfy.Settings.AE.highlight'] = false
+    //     }
+    // }
     // 兼容 ComfyUI Revision: 1887 [235727fe] 以上版本
     if(api.storeSettings){
         const _settings = await api.getSettings()