Spaces:
Configuration error
Configuration error
File size: 3,632 Bytes
8866644 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
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 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):
import comfy.supported_models
self.old_supported_models = comfy.supported_models.models
comfy.supported_models.models = [Kolors]
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
import comfy.supported_models
comfy.supported_models.models = self.old_supported_models
|