diff --git a/configs/text_encoder_config.json b/configs/text_encoder_config.json new file mode 100644 index 0000000..c6e1930 --- /dev/null +++ b/configs/text_encoder_config.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 +} \ No newline at end of file diff --git a/configs/tokenizer/tokenizer.model b/configs/tokenizer/tokenizer.model new file mode 100644 index 0000000..c8336ad Binary files /dev/null and b/configs/tokenizer/tokenizer.model differ diff --git a/configs/tokenizer/tokenizer_config.json b/configs/tokenizer/tokenizer_config.json new file mode 100644 index 0000000..f6f13c8 --- /dev/null +++ b/configs/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 + ] + } +} diff --git a/configs/tokenizer/vocab.txt b/configs/tokenizer/vocab.txt new file mode 100644 index 0000000..c8336ad Binary files /dev/null and b/configs/tokenizer/vocab.txt differ diff --git a/nodes.py b/nodes.py index bf0ed78..1722f62 100755 --- a/nodes.py +++ b/nodes.py @@ -3,17 +3,18 @@ import os import random import re import gc -import sys +import json import comfy.model_management as mm from comfy.utils import ProgressBar, load_torch_file import folder_paths script_directory = os.path.dirname(os.path.abspath(__file__)) -sys.path.append(script_directory) + +folder_paths.add_model_folder_path("LLM", os.path.join(folder_paths.models_dir, "LLM", "checkpoints")) from .kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256 import StableDiffusionXLPipeline -from .kolors.models.modeling_chatglm import ChatGLMModel +from .kolors.models.modeling_chatglm import ChatGLMModel, ChatGLMConfig from .kolors.models.tokenization_chatglm import ChatGLMTokenizer from diffusers import UNet2DConditionModel from diffusers import (DPMSolverMultistepScheduler, @@ -83,6 +84,52 @@ class DownloadAndLoadKolorsModel: return (kolors_model,) +class LoadChatGLM3: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "chatglm3_checkpoint": (folder_paths.get_filename_list("LLM"),), + "precision": ([ 'fp16', 'quant4', 'quant8'], + { + "default": 'fp16' + }), + }, + } + + RETURN_TYPES = ("CHATGLM3MODEL",) + RETURN_NAMES = ("chatglm3_model",) + FUNCTION = "loadmodel" + CATEGORY = "KwaiKolorsWrapper" + + def loadmodel(self, chatglm3_checkpoint, precision): + + pbar = ProgressBar(2) + chatglm3_path = folder_paths.get_full_path("LLM", chatglm3_checkpoint) + print("Load TEXT_ENCODER...") + text_encoder_config = os.path.join(script_directory, 'configs', 'text_encoder_config.json') + with open(text_encoder_config, 'r') as file: + config = json.load(file) + + text_encoder_config = ChatGLMConfig(**config) + text_encoder = ChatGLMModel(text_encoder_config) + text_encoder.load_state_dict(load_torch_file(chatglm3_path)) + + if precision == 'quant8': + text_encoder.quantize(8) + elif precision == 'quant4': + text_encoder.quantize(4) + + tokenizer_path = os.path.join(script_directory,'configs',"tokenizer") + tokenizer = ChatGLMTokenizer.from_pretrained(tokenizer_path) + pbar.update(1) + + chatglm3_model = { + 'text_encoder': text_encoder, + 'tokenizer': tokenizer + } + + return (chatglm3_model,) + class DownloadAndLoadChatGLM3: @classmethod def INPUT_TYPES(s): @@ -396,11 +443,13 @@ NODE_CLASS_MAPPINGS = { "DownloadAndLoadKolorsModel": DownloadAndLoadKolorsModel, "DownloadAndLoadChatGLM3": DownloadAndLoadChatGLM3, "KolorsSampler": KolorsSampler, - "KolorsTextEncode": KolorsTextEncode + "KolorsTextEncode": KolorsTextEncode, + "LoadChatGLM3": LoadChatGLM3 } NODE_DISPLAY_NAME_MAPPINGS = { "DownloadAndLoadKolorsModel": "(Down)load Kolors Model", "DownloadAndLoadChatGLM3": "(Down)load ChatGLM3 Model", "KolorsSampler": "Kolors Sampler", - "KolorsTextEncode": "Kolors Text Encode" + "KolorsTextEncode": "Kolors Text Encode", + "LoadChatGLM3": "Load ChatGLM3 Model" } \ No newline at end of file