303 lines
14 KiB
Python
303 lines
14 KiB
Python
import json
|
|
import os
|
|
import torch
|
|
import subprocess
|
|
import sys
|
|
import comfy.supported_models
|
|
import comfy.model_patcher
|
|
import comfy.model_management
|
|
import comfy.model_detection as model_detection
|
|
import comfy.model_base as model_base
|
|
from comfy.model_base import sdxl_pooled, CLIPEmbeddingNoiseAugmentation, Timestep, ModelType
|
|
from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel
|
|
from comfy.clip_vision import ClipVisionModel, Output
|
|
from comfy.utils import load_torch_file
|
|
from .chatglm.modeling_chatglm import ChatGLMModel, ChatGLMConfig
|
|
from .chatglm.tokenization_chatglm import ChatGLMTokenizer
|
|
|
|
class KolorsUNetModel(UNetModel):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.encoder_hid_proj = torch.nn.Linear(4096, 2048, bias=True)
|
|
|
|
def forward(self, *args, **kwargs):
|
|
with torch.cuda.amp.autocast(enabled=True):
|
|
if "context" in kwargs:
|
|
kwargs["context"] = self.encoder_hid_proj(kwargs["context"])
|
|
result = super().forward(*args, **kwargs)
|
|
return result
|
|
|
|
class KolorsSDXL(model_base.SDXL):
|
|
def __init__(self, model_config, model_type=ModelType.EPS, device=None):
|
|
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": 1000, "beta_schedule": "squaredcos_cap_v2"}, "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 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 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
|
|
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}
|
|
|
|
Kolors_inpaint = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True,
|
|
'legacy': False,
|
|
'num_classes': 'sequential', 'adm_in_channels': 5632, 'dtype': dtype, 'in_channels': 9,
|
|
'model_channels': 320,
|
|
'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 10, 10], 'channel_mult': [1, 2, 4],
|
|
'transformer_depth_middle': 10,
|
|
'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64,
|
|
'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10],
|
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
|
|
|
Kolors_ip2p = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True,
|
|
'legacy': False,
|
|
'num_classes': 'sequential', 'adm_in_channels': 5632, 'dtype': dtype, 'in_channels': 8,
|
|
'model_channels': 320,
|
|
'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 10, 10], 'channel_mult': [1, 2, 4],
|
|
'transformer_depth_middle': 10,
|
|
'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64,
|
|
'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10],
|
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
|
|
|
SDXL = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True,
|
|
'legacy': False,
|
|
'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4,
|
|
'model_channels': 320,
|
|
'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 10, 10], 'channel_mult': [1, 2, 4],
|
|
'transformer_depth_middle': 10,
|
|
'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64,
|
|
'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10],
|
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
|
|
|
SDXL_mid_cnet = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True,
|
|
'legacy': False,
|
|
'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4,
|
|
'model_channels': 320,
|
|
'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 0, 0, 1, 1], 'channel_mult': [1, 2, 4],
|
|
'transformer_depth_middle': 1,
|
|
'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64,
|
|
'transformer_depth_output': [0, 0, 0, 0, 0, 0, 1, 1, 1],
|
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
|
|
|
SDXL_small_cnet = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True,
|
|
'legacy': False,
|
|
'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4,
|
|
'model_channels': 320,
|
|
'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 0, 0, 0, 0], 'channel_mult': [1, 2, 4],
|
|
'transformer_depth_middle': 0,
|
|
'use_linear_in_transformer': True, 'num_head_channels': 64, 'context_dim': 1,
|
|
'transformer_depth_output': [0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
|
|
|
supported_models = [Kolors, Kolors_inpaint,
|
|
Kolors_ip2p, SDXL, SDXL_mid_cnet, SDXL_small_cnet]
|
|
|
|
|
|
for unet_config in supported_models:
|
|
matches = True
|
|
for k in match:
|
|
if match[k] != unet_config[k]:
|
|
# print("key {} does not match".format(k), match[k], "||", unet_config[k])
|
|
matches = False
|
|
break
|
|
if matches:
|
|
return model_detection.convert_config(unet_config)
|
|
return None
|
|
|
|
# chatglm3 model
|
|
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()):
|
|
with torch.no_grad():
|
|
print('torch version:', torch.__version__)
|
|
self.text_encoder = ChatGLMModel(textmodel_json_config).eval()
|
|
if '4bit' in model_path:
|
|
try:
|
|
import cpm_kernels
|
|
except ImportError:
|
|
print("Installing cpm_kernels...")
|
|
subprocess.run([sys.executable, "-m", "pip", "install", "cpm_kernels"], check=True)
|
|
pass
|
|
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:
|
|
print("WARNING: Accelerate not available, use load_state_dict load model")
|
|
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}
|
|
|
|
|
|
# clipvision model
|
|
def load_clipvision_vitl_336(path):
|
|
sd = load_torch_file(path)
|
|
if "vision_model.encoder.layers.22.layer_norm1.weight" in sd:
|
|
json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl_336.json")
|
|
else:
|
|
raise Exception("Unsupported clip vision model")
|
|
clip = ClipVisionModel(json_config)
|
|
m, u = clip.load_sd(sd)
|
|
if len(m) > 0:
|
|
print("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
|
|
return clip
|
|
|
|
class applyKolorsUnet:
|
|
def __enter__(self):
|
|
import comfy.ldm.modules.diffusionmodules.openaimodel
|
|
import comfy.utils
|
|
import comfy.clip_vision
|
|
|
|
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 = [Kolors]
|
|
|
|
self.original_load_clipvision_from_sd = comfy.clip_vision.load_clipvision_from_sd
|
|
comfy.clip_vision.load_clipvision_from_sd = load_clipvision_vitl_336
|
|
|
|
def __exit__(self, type, value, traceback):
|
|
import comfy.ldm.modules.diffusionmodules.openaimodel
|
|
import comfy.utils
|
|
import comfy.supported_models
|
|
import comfy.clip_vision
|
|
|
|
comfy.utils.UNET_MAP_BASIC = self.original_UNET_MAP_BASIC
|
|
|
|
model_detection.unet_config_from_diffusers_unet = self.original_unet_config_from_diffusers_unet
|
|
comfy.supported_models.models = self.original_supported_models
|
|
|
|
comfy.clip_vision.load_clipvision_from_sd = self.original_load_clipvision_from_sd
|
|
|
|
|
|
def is_kolors_model(model):
|
|
unet_config = model.model.model_config.unet_config
|
|
if unet_config and "adm_in_channels" in unet_config and unet_config["adm_in_channels"] == 5632:
|
|
return True
|
|
else:
|
|
return False |