123 lines
3.6 KiB
Python
123 lines
3.6 KiB
Python
from enum import Enum
|
|
|
|
import torch
|
|
|
|
from comfy import supported_models_base
|
|
from comfy import latent_formats
|
|
from comfy.model_detection import convert_config
|
|
|
|
|
|
class ModelVersion(Enum):
|
|
SD15 = "sd15"
|
|
SDXL = "sdxl"
|
|
SDXL_REFINER = "sdxl_refiner"
|
|
LCM = "lcm"
|
|
|
|
|
|
config_map = {
|
|
ModelVersion.SD15: {
|
|
"use_checkpoint": False,
|
|
"image_size": 32,
|
|
"out_channels": 4,
|
|
"use_spatial_transformer": True,
|
|
"legacy": False,
|
|
"adm_in_channels": None,
|
|
"dtype": torch.float16,
|
|
"in_channels": 4,
|
|
"model_channels": 320,
|
|
"num_res_blocks": 2,
|
|
"attention_resolutions": [1, 2, 4],
|
|
"transformer_depth": [1, 1, 1, 0],
|
|
"channel_mult": [1, 2, 4, 4],
|
|
"transformer_depth_middle": 1,
|
|
"use_linear_in_transformer": False,
|
|
"context_dim": 768,
|
|
"num_heads": 8,
|
|
"disable_unet_model_creation": True,
|
|
},
|
|
ModelVersion.SDXL: {
|
|
"use_checkpoint": False,
|
|
"image_size": 32,
|
|
"out_channels": 4,
|
|
"use_spatial_transformer": True,
|
|
"legacy": False,
|
|
"num_classes": "sequential",
|
|
"adm_in_channels": 2816,
|
|
"dtype": torch.float16,
|
|
"in_channels": 4,
|
|
"model_channels": 320,
|
|
"num_res_blocks": 2,
|
|
"attention_resolutions": [2, 4],
|
|
"transformer_depth": [0, 2, 10],
|
|
"channel_mult": [1, 2, 4],
|
|
"transformer_depth_middle": 10,
|
|
"use_linear_in_transformer": True,
|
|
"context_dim": 2048,
|
|
"num_head_channels": 64,
|
|
"disable_unet_model_creation": True,
|
|
},
|
|
ModelVersion.SDXL_REFINER: {
|
|
"use_checkpoint": False,
|
|
"image_size": 32,
|
|
"out_channels": 4,
|
|
"use_spatial_transformer": True,
|
|
"legacy": False,
|
|
"num_classes": "sequential",
|
|
"adm_in_channels": 2560,
|
|
"dtype": torch.float16,
|
|
"in_channels": 4,
|
|
"model_channels": 384,
|
|
"num_res_blocks": 2,
|
|
"attention_resolutions": [2, 4],
|
|
"transformer_depth": [0, 4, 4, 0],
|
|
"channel_mult": [1, 2, 4, 4],
|
|
"transformer_depth_middle": 4,
|
|
"use_linear_in_transformer": True,
|
|
"context_dim": 1280,
|
|
"num_head_channels": 64,
|
|
"disable_unet_model_creation": True,
|
|
},
|
|
}
|
|
|
|
latent_format_map = {
|
|
ModelVersion.SD15: latent_formats.SD15,
|
|
ModelVersion.SDXL: latent_formats.SDXL,
|
|
ModelVersion.SDXL_REFINER: latent_formats.SDXL,
|
|
}
|
|
|
|
|
|
def get_model_config(model_version: ModelVersion):
|
|
unet_config = convert_config(config_map[model_version])
|
|
config = supported_models_base.BASE(unet_config)
|
|
config.latent_format = latent_format_map[model_version]()
|
|
return config
|
|
|
|
|
|
def unet_config_from_diffusers_unet(state_dict):
|
|
match = {}
|
|
attention_resolutions = []
|
|
|
|
attn_res = 1
|
|
for i in range(5):
|
|
k = "down_blocks.{}.attentions.1.transformer_blocks.0.attn2.to_k.weight".format(
|
|
i
|
|
)
|
|
if k in state_dict:
|
|
match["context_dim"] = state_dict[k].shape[1]
|
|
attention_resolutions.append(attn_res)
|
|
attn_res *= 2
|
|
|
|
match["attention_resolutions"] = attention_resolutions
|
|
|
|
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]
|
|
|
|
print(match)
|