Spaces:
Configuration error
Configuration error
| 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 | |