147 lines
5.1 KiB
Python
147 lines
5.1 KiB
Python
import math
|
|
import logging
|
|
|
|
import comfy.utils
|
|
import comfy.model_base
|
|
import comfy.model_detection
|
|
|
|
import comfy.supported_models_base
|
|
import comfy.supported_models
|
|
import comfy.latent_formats
|
|
|
|
from .models.sana import Sana
|
|
from .models.sana_multi_scale import SanaMS
|
|
from .diffusers_convert import convert_state_dict
|
|
from ..utils.loader import load_state_dict_from_config
|
|
|
|
class SanaLatent(comfy.latent_formats.LatentFormat):
|
|
latent_channels = 32
|
|
def __init__(self):
|
|
self.scale_factor = 0.41407
|
|
self.latent_rgb_factors = latent_rgb_factors.copy()
|
|
self.latent_rgb_factors_bias = latent_rgb_factors_bias.copy()
|
|
|
|
class SanaConfig(comfy.supported_models_base.BASE):
|
|
unet_class = SanaMS
|
|
unet_config = {}
|
|
unet_extra_config = {}
|
|
|
|
latent_format = SanaLatent
|
|
sampling_settings = {
|
|
"shift": 3.0,
|
|
}
|
|
|
|
def model_type(self, state_dict, prefix=""):
|
|
return comfy.model_base.ModelType.FLOW
|
|
|
|
def get_model(self, state_dict, prefix="", device=None):
|
|
return SanaModel(model_config=self, unet_model=self.unet_class, device=device)
|
|
|
|
class SanaModel(comfy.model_base.BaseModel):
|
|
def __init__(self, *args, model_type=comfy.model_base.ModelType.FLOW, unet_model=SanaMS, **kwargs):
|
|
super().__init__(*args, model_type=model_type, unet_model=unet_model, **kwargs)
|
|
|
|
def load_sana_state_dict(sd, model_options={}):
|
|
# prefix / format
|
|
sd = sd.get("model", sd) # ref ckpt
|
|
diffusion_model_prefix = comfy.model_detection.unet_prefix_from_state_dict(sd)
|
|
temp_sd = comfy.utils.state_dict_prefix_replace(sd, {diffusion_model_prefix: ""}, filter_keys=True)
|
|
if len(temp_sd) > 0:
|
|
sd = temp_sd
|
|
|
|
# diffusers convert
|
|
if "adaln_single.linear.weight" in sd:
|
|
sd = convert_state_dict(sd)
|
|
|
|
# model config
|
|
model_config = model_config_from_unet(sd)
|
|
return load_state_dict_from_config(model_config, sd, model_options)
|
|
|
|
def model_config_from_unet(sd):
|
|
"""
|
|
Guess config based on (converted) state dict.
|
|
"""
|
|
# shared settings that match between all models
|
|
# TODO: some can (should) be enumerated
|
|
config = {
|
|
"in_channels": 32,
|
|
"linear_head_dim": 32,
|
|
"model_max_length": 300,
|
|
"y_norm": True,
|
|
"attn_type": "linear",
|
|
"ffn_type": "glumbconv",
|
|
"mlp_ratio": 2.5,
|
|
"mlp_acts": ["silu", "silu", None],
|
|
"use_pe": False,
|
|
"pred_sigma": False,
|
|
"learn_sigma": False,
|
|
"fp32_attention": True,
|
|
"patch_size": 1,
|
|
}
|
|
config["depth"] = sum([key.endswith(".point_conv.conv.weight") for key in sd.keys()]) or 28
|
|
|
|
if "pos_embed" in sd:
|
|
config["input_size"] = int(math.sqrt(sd["pos_embed"].shape[1])) * config["patch_size"]
|
|
else:
|
|
# TODO: this isn't optimal though most models don't use it
|
|
config["use_pe"] = False
|
|
|
|
if "x_embedder.proj.bias" in sd:
|
|
config["hidden_size"] = sd["x_embedder.proj.bias"].shape[0]
|
|
|
|
if config["hidden_size"] == 1152:
|
|
config["num_heads"] = 16
|
|
elif config["hidden_size"] == 2240:
|
|
config["num_heads"] = 20
|
|
else:
|
|
raise RuntimeError(f"Unknown model config.")
|
|
|
|
model_config = SanaConfig(config)
|
|
logging.debug(f"Sana config:\n{config}")
|
|
return model_config
|
|
|
|
# for fast latent preview
|
|
latent_rgb_factors = [
|
|
[-2.0022e-03, -6.0736e-03, -1.7096e-03],
|
|
[ 1.4221e-03, 3.6703e-03, 4.1083e-03],
|
|
[ 1.0081e-02, 2.6456e-04, -1.4333e-02],
|
|
[-2.4253e-03, 3.0967e-03, -1.0301e-03],
|
|
[ 2.2158e-03, 7.7274e-03, -1.3151e-02],
|
|
[ 1.1235e-02, 5.7630e-03, 3.6146e-03],
|
|
[-7.2899e-02, 1.1062e-02, 3.6103e-02],
|
|
[ 3.2346e-02, 2.8678e-02, 2.5014e-02],
|
|
[ 1.6469e-03, -1.1364e-03, 2.8366e-03],
|
|
[-3.5597e-02, -2.3447e-02, -3.1172e-03],
|
|
[-1.9985e-04, -2.0647e-03, -1.2702e-02],
|
|
[ 2.1318e-04, 1.2196e-03, -8.3461e-04],
|
|
[ 1.3766e-02, 2.7559e-03, 7.3567e-03],
|
|
[ 1.3027e-02, 2.6365e-03, 3.0405e-03],
|
|
[ 1.5335e-02, 9.4682e-03, 6.7312e-03],
|
|
[ 5.1827e-03, -9.4865e-03, 8.5080e-03],
|
|
[ 1.4365e-02, -3.2867e-03, 9.5108e-03],
|
|
[-4.1216e-03, -1.9177e-03, -3.3726e-03],
|
|
[-2.4757e-03, 5.1739e-04, 2.0280e-03],
|
|
[-3.5950e-03, 1.0720e-03, 5.3043e-03],
|
|
[-5.1758e-03, 8.1040e-03, -3.7564e-02],
|
|
[-3.8555e-03, -1.5529e-03, -3.5799e-03],
|
|
[-6.6175e-03, -6.8484e-03, -9.9609e-03],
|
|
[-2.1656e-03, 5.5770e-05, 1.4936e-03],
|
|
[-9.2857e-02, -1.1379e-01, -1.0919e-01],
|
|
[ 7.7044e-04, 5.5594e-03, 3.4755e-02],
|
|
[ 1.2714e-02, 2.9729e-02, 3.1989e-03],
|
|
[-1.1805e-03, 9.0548e-03, -4.1063e-04],
|
|
[ 8.3309e-04, 4.9694e-03, 2.3087e-03],
|
|
[ 7.8456e-03, 3.9750e-03, 3.5655e-03],
|
|
[-1.7552e-03, 4.9306e-03, 1.4210e-02],
|
|
[-1.4790e-03, 2.8837e-03, -4.5687e-03]
|
|
]
|
|
latent_rgb_factors_bias = [0.4358, 0.3814, 0.3388]
|
|
|
|
# 512/1024/2K match, TODO: 4K is new, add on release
|
|
from ..PixArt.loader import resolutions as pixart_res
|
|
resolutions = {
|
|
"Sana 512": pixart_res["PixArt 512"],
|
|
"Sana 1024": pixart_res["PixArt 1024"],
|
|
"Sana 2K": pixart_res["PixArt 2K"],
|
|
}
|