nodes/kijai-ComfyUI-KwaiKolorsWrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

text encoder quantization

Browse Source
This commit is contained in:
kijai
2024-07-07 02:25:26 +03:00
parent 51b89de9fc
commit 39df37422f
5 changed files with 506 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/kolors_example.json
+244 -188
View File
@@ -1,90 +1,7 @@
{ {
"last_node_id": 11, "last_node_id": 15,
"last_link_id": 13, "last_link_id": 18,
"nodes": [ "nodes": [
{
"id": 6,
"type": "DownloadAndLoadKolorsModel",
"pos": [
547,
372
],
"size": {
"0": 315,
"1": 82
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "kolors_model",
"type": "KOLORSMODEL",
"links": [
7,
9
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "DownloadAndLoadKolorsModel"
},
"widgets_values": [
"Kwai-Kolors/Kolors",
"fp16"
]
},
{
"id": 9,
"type": "KolorsSampler",
"pos": [
1011,
371
],
"size": {
"0": 315,
"1": 198
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "kolors_model",
"type": "KOLORSMODEL",
"link": 9
},
{
"name": "kolors_embeds",
"type": "KOLORS_EMBEDS",
"link": 10
}
],
"outputs": [
{
"name": "latent",
"type": "LATENT",
"links": [
11
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "KolorsSampler"
},
"widgets_values": [
1024,
1024,
243590846571465,
"randomize",
25,
5
]
},
{ {
"id": 11, "id": 11,
"type": "VAELoader", "type": "VAELoader",
@@ -97,7 +14,7 @@
"1": 58 "1": 58
}, },
"flags": {}, "flags": {},
"order": 1, "order": 0,
"mode": 0, "mode": 0,
"outputs": [ "outputs": [
{ {
@@ -116,71 +33,6 @@
"sdxl.vae.safetensors" "sdxl.vae.safetensors"
] ]
}, },
{
"id": 8,
"type": "KolorsTextEncode",
"pos": [
549,
522
],
"size": {
"0": 400,
"1": 200
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "kolors_model",
"type": "KOLORSMODEL",
"link": 7
}
],
"outputs": [
{
"name": "kolors_embeds",
"type": "KOLORS_EMBEDS",
"links": [
10
],
"shape": 3
}
],
"properties": {
"Node name for S&R": "KolorsTextEncode"
},
"widgets_values": [
"cinematic photograph monkey giving two thumbs up",
"nsfw, naked",
4
]
},
{
"id": 3,
"type": "PreviewImage",
"pos": [
1367,
467
],
"size": [
670,
646.6666259765625
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 13
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{ {
"id": 10, "id": 10,
"type": "VAEDecode", "type": "VAEDecode",
@@ -193,13 +45,13 @@
"1": 46 "1": 46
}, },
"flags": {}, "flags": {},
"order": 4, "order": 6,
"mode": 0, "mode": 0,
"inputs": [ "inputs": [
{ {
"name": "samples", "name": "samples",
"type": "LATENT", "type": "LATENT",
"link": 11 "link": 18
}, },
{ {
"name": "vae", "name": "vae",
@@ -222,41 +74,213 @@
"properties": { "properties": {
"Node name for S&R": "VAEDecode" "Node name for S&R": "VAEDecode"
} }
},
{
"id": 14,
"type": "KolorsSampler",
"pos": [
1011,
371
],
"size": {
"0": 315,
"1": 222
},
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "kolors_model",
"type": "KOLORSMODEL",
"link": 16
},
{
"name": "kolors_embeds",
"type": "KOLORS_EMBEDS",
"link": 17
}
],
"outputs": [
{
"name": "latent",
"type": "LATENT",
"links": [
18
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "KolorsSampler"
},
"widgets_values": [
1024,
1024,
1000102404233412,
"fixed",
25,
5,
"EulerDiscreteScheduler"
]
},
{
"id": 6,
"type": "DownloadAndLoadKolorsModel",
"pos": [
201,
368
],
"size": {
"0": 315,
"1": 82
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "kolors_model",
"type": "KOLORSMODEL",
"links": [
16
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "DownloadAndLoadKolorsModel"
},
"widgets_values": [
"Kwai-Kolors/Kolors",
"fp16"
]
},
{
"id": 3,
"type": "PreviewImage",
"pos": [
1366,
468
],
"size": [
535.4001724243165,
562.2001106262207
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 13
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 12,
"type": "KolorsTextEncode",
"pos": [
519,
529
],
"size": [
457.2893696934723,
225.28656056301645
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "chatglm3_model",
"type": "CHATGLM3MODEL",
"link": 14,
"slot_index": 0
}
],
"outputs": [
{
"name": "kolors_embeds",
"type": "KOLORS_EMBEDS",
"links": [
17
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "KolorsTextEncode"
},
"widgets_values": [
"cinematic photograph of an astronaut riding a horse in space |\nillustration of a cat wearing a top hat and a scarf |\nphotograph of a goldfish in a bowl |\nanime screencap of a red haired girl",
"",
1
]
},
{
"id": 15,
"type": "Note",
"pos": [
200,
636
],
"size": [
273.5273818969726,
149.55464588512064
],
"flags": {},
"order": 2,
"mode": 0,
"properties": {
"text": ""
},
"widgets_values": [
"Text encoding takes the most VRAM, quantization can reduce that a lot.\n\nApproximate values I have observed:\nfp16 - 12 GB\nquant8 - 8-9 GB\nquant4 - 4-5 GB\n\nquant4 reduces the quality quite a bit, 8 seems fine"
],
"color": "#432",
"bgcolor": "#653"
},
{
"id": 13,
"type": "DownloadAndLoadChatGLM3",
"pos": [
206,
522
],
"size": [
274.5334274291992,
58
],
"flags": {},
"order": 3,
"mode": 0,
"outputs": [
{
"name": "chatglm3_model",
"type": "CHATGLM3MODEL",
"links": [
14
],
"shape": 3
}
],
"properties": {
"Node name for S&R": "DownloadAndLoadChatGLM3"
},
"widgets_values": [
"fp16"
]
} }
], ],
"links": [ "links": [
[
7,
6,
0,
8,
0,
"KOLORSMODEL"
],
[
9,
6,
0,
9,
0,
"KOLORSMODEL"
],
[
10,
8,
0,
9,
1,
"KOLORS_EMBEDS"
],
[
11,
9,
0,
10,
0,
"LATENT"
],
[ [
12, 12,
11, 11,
@@ -272,16 +296,48 @@
3, 3,
0, 0,
"IMAGE" "IMAGE"
],
[
14,
13,
0,
12,
0,
"CHATGLM3MODEL"
],
[
16,
6,
0,
14,
0,
"KOLORSMODEL"
],
[
17,
12,
0,
14,
1,
"KOLORS_EMBEDS"
],
[
18,
14,
0,
10,
0,
"LATENT"
] ]
], ],
"groups": [], "groups": [],
"config": {}, "config": {},
"extra": { "extra": {
"ds": { "ds": {
"scale": 1, "scale": 1.1,
"offset": { "offset": {
"0": -404.6999816894531, "0": -114.73954010009766,
"1": -111.86663818359375 "1": -139.79705810546875
} }
} }
}, },
kolors/models/quantization.py
+188
View File
File diff suppressed because one or more lines are too long
kolors/pipelines/pipeline_stable_diffusion_xl_chatglm_256.py
+4 -4
View File
@@ -107,8 +107,8 @@ class StableDiffusionXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoad
def __init__( def __init__(
self, self,
text_encoder: ChatGLMModel, # text_encoder: ChatGLMModel,
tokenizer: ChatGLMTokenizer, # tokenizer: ChatGLMTokenizer,
unet: UNet2DConditionModel, unet: UNet2DConditionModel,
scheduler: KarrasDiffusionSchedulers, scheduler: KarrasDiffusionSchedulers,
force_zeros_for_empty_prompt: bool = True, force_zeros_for_empty_prompt: bool = True,
@@ -117,8 +117,8 @@ class StableDiffusionXLPipeline(DiffusionPipeline, FromSingleFileMixin, LoraLoad
self.register_modules( self.register_modules(
#vae=vae, #vae=vae,
text_encoder=text_encoder, #text_encoder=text_encoder,
tokenizer=tokenizer, #tokenizer=tokenizer,
unet=unet, unet=unet,
scheduler=scheduler, scheduler=scheduler,
) )
nodes.py
+68 -19
View File
@@ -3,13 +3,14 @@ import os
import random import random
import re import re
import gc import gc
import sys
import comfy.model_management as mm import comfy.model_management as mm
from comfy.utils import ProgressBar, load_torch_file from comfy.utils import ProgressBar, load_torch_file
import folder_paths import folder_paths
script_directory = os.path.dirname(os.path.abspath(__file__)) script_directory = os.path.dirname(os.path.abspath(__file__))
sys.path.append(script_directory)
from .kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256 import StableDiffusionXLPipeline 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
@@ -59,7 +60,8 @@ class DownloadAndLoadKolorsModel:
print(f"Downloading Kolor model to: {model_path}") print(f"Downloading Kolor model to: {model_path}")
from huggingface_hub import snapshot_download from huggingface_hub import snapshot_download
snapshot_download(repo_id=model, snapshot_download(repo_id=model,
allow_patterns=['*fp16.safetensors*', '*.json', 'text_encoder/*', 'tokenizer/*'], allow_patterns=['*fp16.safetensors*', '*.json'],
ignore_patthers=['text_encoder/*', 'tokenizer/*'],
local_dir=model_path, local_dir=model_path,
local_dir_use_symlinks=False) local_dir_use_symlinks=False)
pbar.update(1) pbar.update(1)
@@ -67,28 +69,19 @@ class DownloadAndLoadKolorsModel:
scheduler = EulerDiscreteScheduler.from_pretrained(model_path, subfolder= 'scheduler') scheduler = EulerDiscreteScheduler.from_pretrained(model_path, subfolder= 'scheduler')
print("Load UNET...") print("Load UNET...")
unet = UNet2DConditionModel.from_pretrained(model_path, subfolder= 'unet', variant="fp16", revision=None, low_cpu_mem_usage=True).to(dtype).eval() unet = UNet2DConditionModel.from_pretrained(model_path, subfolder= 'unet', variant="fp16", revision=None, low_cpu_mem_usage=True).to(dtype).eval()
print("Load TEXT_ENCODER...")
pbar.update(1)
text_encoder_path = os.path.join(model_path, "text_encoder")
text_encoder = ChatGLMModel.from_pretrained(
text_encoder_path,
torch_dtype=dtype,
)
tokenizer = ChatGLMTokenizer.from_pretrained(text_encoder_path)
pbar.update(1)
pipeline = StableDiffusionXLPipeline( pipeline = StableDiffusionXLPipeline(
#vae=None, #vae=None,
text_encoder=text_encoder, #text_encoder=None,
tokenizer=tokenizer, #tokenizer=None,
unet=unet, unet=unet,
scheduler=scheduler, scheduler=scheduler,
force_zeros_for_empty_prompt=False force_zeros_for_empty_prompt=False
) )
#pipeline = pipeline.to(device) #pipeline = pipeline.to(device)
pipeline.enable_model_cpu_offload() #pipeline.enable_model_cpu_offload()
kolors_model = { kolors_model = {
'pipeline': pipeline, 'pipeline': pipeline,
@@ -97,13 +90,67 @@ class DownloadAndLoadKolorsModel:
return (kolors_model,) return (kolors_model,)
class DownloadAndLoadChatGLM3:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"precision": ([ 'fp16', 'quant4', 'quant8'],
{
"default": 'fp16'
}),
},
}
RETURN_TYPES = ("CHATGLM3MODEL",)
RETURN_NAMES = ("chatglm3_model",)
FUNCTION = "loadmodel"
CATEGORY = "KwaiKolorsWrapper"
def loadmodel(self, precision):
pbar = ProgressBar(2)
model = "Kwai-Kolors/Kolors"
model_name = model.rsplit('/', 1)[-1]
model_path = os.path.join(folder_paths.models_dir, "diffusers", model_name)
text_encoder_path = os.path.join(model_path, "text_encoder")
if not os.path.exists(text_encoder_path):
print(f"Downloading Kolor model to: {model_path}")
from huggingface_hub import snapshot_download
snapshot_download(repo_id=model,
allow_patterns=['text_encoder/*', 'tokenizer/*'],
local_dir=model_path,
local_dir_use_symlinks=False)
pbar.update(1)
print("Load TEXT_ENCODER...")
text_encoder_path = os.path.join(model_path, "text_encoder")
text_encoder = ChatGLMModel.from_pretrained(
text_encoder_path,
torch_dtype=torch.float16,
)
if precision == 'quant8':
text_encoder.quantize(8)
elif precision == 'quant4':
text_encoder.quantize(4)
tokenizer = ChatGLMTokenizer.from_pretrained(text_encoder_path)
pbar.update(1)
chatglm3_model = {
'text_encoder': text_encoder,
'tokenizer': tokenizer
}
return (chatglm3_model,)
class KolorsTextEncode: class KolorsTextEncode:
@classmethod @classmethod
def INPUT_TYPES(s): def INPUT_TYPES(s):
return { return {
"required": { "required": {
"kolors_model": ("KOLORSMODEL", ), "chatglm3_model": ("CHATGLM3MODEL", ),
"prompt": ("STRING", {"multiline": True, "default": "",}), "prompt": ("STRING", {"multiline": True, "default": "",}),
"negative_prompt": ("STRING", {"multiline": True, "default": "",}), "negative_prompt": ("STRING", {"multiline": True, "default": "",}),
"num_images_per_prompt": ("INT", {"default": 1, "min": 1, "max": 128, "step": 1}), "num_images_per_prompt": ("INT", {"default": 1, "min": 1, "max": 128, "step": 1}),
@@ -115,7 +162,7 @@ class KolorsTextEncode:
FUNCTION = "encode" FUNCTION = "encode"
CATEGORY = "KwaiKolorsWrapper" CATEGORY = "KwaiKolorsWrapper"
def encode(self, kolors_model, prompt, negative_prompt, num_images_per_prompt): def encode(self, chatglm3_model, prompt, negative_prompt, num_images_per_prompt):
device = mm.get_torch_device() device = mm.get_torch_device()
offload_device = mm.unet_offload_device() offload_device = mm.unet_offload_device()
mm.unload_all_models() mm.unload_all_models()
@@ -143,8 +190,8 @@ class KolorsTextEncode:
batch_size = len(prompt) batch_size = len(prompt)
# Define tokenizers and text encoders # Define tokenizers and text encoders
tokenizer = kolors_model['pipeline'].tokenizer tokenizer = chatglm3_model['tokenizer']
text_encoder = kolors_model['pipeline'].text_encoder text_encoder = chatglm3_model['text_encoder']
text_encoder.to(device) text_encoder.to(device)
@@ -341,11 +388,13 @@ class KolorsSampler:
NODE_CLASS_MAPPINGS = { NODE_CLASS_MAPPINGS = {
"DownloadAndLoadKolorsModel": DownloadAndLoadKolorsModel, "DownloadAndLoadKolorsModel": DownloadAndLoadKolorsModel,
"DownloadAndLoadChatGLM3": DownloadAndLoadChatGLM3,
"KolorsSampler": KolorsSampler, "KolorsSampler": KolorsSampler,
"KolorsTextEncode": KolorsTextEncode "KolorsTextEncode": KolorsTextEncode
} }
NODE_DISPLAY_NAME_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = {
"DownloadAndLoadKolorsModel": "(Down)load Kolors Model", "DownloadAndLoadKolorsModel": "(Down)load Kolors Model",
"DownloadAndLoadChatGLM3": "(Down)load ChatGLM3 Model",
"KolorsSampler": "Kolors Sampler", "KolorsSampler": "Kolors Sampler",
"KolorsTextEncode": "Kolors Text Encode" "KolorsTextEncode": "Kolors Text Encode"
} }
requirements.txt
+2 -1
View File
@@ -1,4 +1,5 @@
diffusers>=0.28.2 diffusers>=0.28.2
transformers>=4.26.1 transformers>=4.26.1
sentencepiece sentencepiece
accelerate accelerate
cpm-kernels