98 lines
3.9 KiB
Python
98 lines
3.9 KiB
Python
from comfy.model_detection import *
|
|
import comfy.model_detection as model_detection
|
|
import comfy.supported_models
|
|
|
|
|
|
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")
|
|
# 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):
|
|
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
|
|
|
|
|
|
class apply_kolors:
|
|
def __enter__(self):
|
|
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
|