fred-dev/comfy_ui_ali
0
1import torch2from . import model_base3from . import utils4 5from . import sd1_clip6from . import sdxl_clip7import comfy.text_encoders.sd2_clip8import comfy.text_encoders.sd3_clip9import comfy.text_encoders.sa_t510import comfy.text_encoders.aura_t511import comfy.text_encoders.pixart_t512import comfy.text_encoders.hydit13import comfy.text_encoders.flux14import comfy.text_encoders.genmo15import comfy.text_encoders.lt16import comfy.text_encoders.hunyuan_video17import comfy.text_encoders.cosmos18import comfy.text_encoders.lumina219import comfy.text_encoders.wan20 21from . import supported_models_base22from . import latent_formats23 24from . import diffusers_convert25 26class SD15(supported_models_base.BASE):27 unet_config = {28 "context_dim": 768,29 "model_channels": 320,30 "use_linear_in_transformer": False,31 "adm_in_channels": None,32 "use_temporal_attention": False,33 }34 35 unet_extra_config = {36 "num_heads": 8,37 "num_head_channels": -1,38 }39 40 latent_format = latent_formats.SD1541 memory_usage_factor = 1.042 43 def process_clip_state_dict(self, state_dict):44 k = list(state_dict.keys())45 for x in k:46 if x.startswith("cond_stage_model.transformer.") and not x.startswith("cond_stage_model.transformer.text_model."):47 y = x.replace("cond_stage_model.transformer.", "cond_stage_model.transformer.text_model.")48 state_dict[y] = state_dict.pop(x)49 50 if 'cond_stage_model.transformer.text_model.embeddings.position_ids' in state_dict:51 ids = state_dict['cond_stage_model.transformer.text_model.embeddings.position_ids']52 if ids.dtype == torch.float32:53 state_dict['cond_stage_model.transformer.text_model.embeddings.position_ids'] = ids.round()54 55 replace_prefix = {}56 replace_prefix["cond_stage_model."] = "clip_l."57 state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=True)58 return state_dict59 60 def process_clip_state_dict_for_saving(self, state_dict):61 pop_keys = ["clip_l.transformer.text_projection.weight", "clip_l.logit_scale"]62 for p in pop_keys:63 if p in state_dict:64 state_dict.pop(p)65 66 replace_prefix = {"clip_l.": "cond_stage_model."}67 return utils.state_dict_prefix_replace(state_dict, replace_prefix)68 69 def clip_target(self, state_dict={}):70 return supported_models_base.ClipTarget(sd1_clip.SD1Tokenizer, sd1_clip.SD1ClipModel)71 72class SD20(supported_models_base.BASE):73 unet_config = {74 "context_dim": 1024,75 "model_channels": 320,76 "use_linear_in_transformer": True,77 "adm_in_channels": None,78 "use_temporal_attention": False,79 }80 81 unet_extra_config = {82 "num_heads": -1,83 "num_head_channels": 64,84 "attn_precision": torch.float32,85 }86 87 latent_format = latent_formats.SD1588 memory_usage_factor = 1.089 90 def model_type(self, state_dict, prefix=""):91 if self.unet_config["in_channels"] == 4: #SD2.0 inpainting models are not v prediction92 k = "{}output_blocks.11.1.transformer_blocks.0.norm1.bias".format(prefix)93 out = state_dict.get(k, None)94 if out is not None and torch.std(out, unbiased=False) > 0.09: # not sure how well this will actually work. I guess we will find out.95 return model_base.ModelType.V_PREDICTION96 return model_base.ModelType.EPS97 98 def process_clip_state_dict(self, state_dict):99 replace_prefix = {}100 replace_prefix["conditioner.embedders.0.model."] = "clip_h." #SD2 in sgm format101 replace_prefix["cond_stage_model.model."] = "clip_h."102 state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=True)103 state_dict = utils.clip_text_transformers_convert(state_dict, "clip_h.", "clip_h.transformer.")104 return state_dict105 106 def process_clip_state_dict_for_saving(self, state_dict):107 replace_prefix = {}108 replace_prefix["clip_h"] = "cond_stage_model.model"109 state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix)110 state_dict = diffusers_convert.convert_text_enc_state_dict_v20(state_dict)111 return state_dict112 113 def clip_target(self, state_dict={}):114 return supported_models_base.ClipTarget(comfy.text_encoders.sd2_clip.SD2Tokenizer, comfy.text_encoders.sd2_clip.SD2ClipModel)115 116class SD21UnclipL(SD20):117 unet_config = {118 "context_dim": 1024,119 "model_channels": 320,120 "use_linear_in_transformer": True,121 "adm_in_channels": 1536,122 "use_temporal_attention": False,123 }124 125 clip_vision_prefix = "embedder.model.visual."126 noise_aug_config = {"noise_schedule_config": {"timesteps": 1000, "beta_schedule": "squaredcos_cap_v2"}, "timestep_dim": 768}127 128 129class SD21UnclipH(SD20):130 unet_config = {131 "context_dim": 1024,132 "model_channels": 320,133 "use_linear_in_transformer": True,134 "adm_in_channels": 2048,135 "use_temporal_attention": False,136 }137 138 clip_vision_prefix = "embedder.model.visual."139 noise_aug_config = {"noise_schedule_config": {"timesteps": 1000, "beta_schedule": "squaredcos_cap_v2"}, "timestep_dim": 1024}140 141class SDXLRefiner(supported_models_base.BASE):142 unet_config = {143 "model_channels": 384,144 "use_linear_in_transformer": True,145 "context_dim": 1280,146 "adm_in_channels": 2560,147 "transformer_depth": [0, 0, 4, 4, 4, 4, 0, 0],148 "use_temporal_attention": False,149 }150 151 latent_format = latent_formats.SDXL152 memory_usage_factor = 1.0153 154 def get_model(self, state_dict, prefix="", device=None):155 return model_base.SDXLRefiner(self, device=device)156 157 def process_clip_state_dict(self, state_dict):158 keys_to_replace = {}159 replace_prefix = {}160 replace_prefix["conditioner.embedders.0.model."] = "clip_g."161 state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=True)162 163 state_dict = utils.clip_text_transformers_convert(state_dict, "clip_g.", "clip_g.transformer.")164 state_dict = utils.state_dict_key_replace(state_dict, keys_to_replace)165 return state_dict166 167 def process_clip_state_dict_for_saving(self, state_dict):168 replace_prefix = {}169 state_dict_g = diffusers_convert.convert_text_enc_state_dict_v20(state_dict, "clip_g")170 if "clip_g.transformer.text_model.embeddings.position_ids" in state_dict_g:171 state_dict_g.pop("clip_g.transformer.text_model.embeddings.position_ids")172 replace_prefix["clip_g"] = "conditioner.embedders.0.model"173 state_dict_g = utils.state_dict_prefix_replace(state_dict_g, replace_prefix)174 return state_dict_g175 176 def clip_target(self, state_dict={}):177 return supported_models_base.ClipTarget(sdxl_clip.SDXLTokenizer, sdxl_clip.SDXLRefinerClipModel)178 179class SDXL(supported_models_base.BASE):180 unet_config = {181 "model_channels": 320,182 "use_linear_in_transformer": True,183 "transformer_depth": [0, 0, 2, 2, 10, 10],184 "context_dim": 2048,185 "adm_in_channels": 2816,186 "use_temporal_attention": False,187 }188 189 latent_format = latent_formats.SDXL190 191 memory_usage_factor = 0.8192 193 def model_type(self, state_dict, prefix=""):194 if 'edm_mean' in state_dict and 'edm_std' in state_dict: #Playground V2.5195 self.latent_format = latent_formats.SDXL_Playground_2_5()196 self.sampling_settings["sigma_data"] = 0.5197 self.sampling_settings["sigma_max"] = 80.0198 self.sampling_settings["sigma_min"] = 0.002199 return model_base.ModelType.EDM200 elif "edm_vpred.sigma_max" in state_dict:201 self.sampling_settings["sigma_max"] = float(state_dict["edm_vpred.sigma_max"].item())202 if "edm_vpred.sigma_min" in state_dict:203 self.sampling_settings["sigma_min"] = float(state_dict["edm_vpred.sigma_min"].item())204 return model_base.ModelType.V_PREDICTION_EDM205 elif "v_pred" in state_dict:206 if "ztsnr" in state_dict: #Some zsnr anime checkpoints207 self.sampling_settings["zsnr"] = True208 return model_base.ModelType.V_PREDICTION209 else:210 return model_base.ModelType.EPS211 212 def get_model(self, state_dict, prefix="", device=None):213 out = model_base.SDXL(self, model_type=self.model_type(state_dict, prefix), device=device)214 if self.inpaint_model():215 out.set_inpaint()216 return out217 218 def process_clip_state_dict(self, state_dict):219 keys_to_replace = {}220 replace_prefix = {}221 222 replace_prefix["conditioner.embedders.0.transformer.text_model"] = "clip_l.transformer.text_model"223 replace_prefix["conditioner.embedders.1.model."] = "clip_g."224 state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=True)225 226 state_dict = utils.state_dict_key_replace(state_dict, keys_to_replace)227 state_dict = utils.clip_text_transformers_convert(state_dict, "clip_g.", "clip_g.transformer.")228 return state_dict229 230 def process_clip_state_dict_for_saving(self, state_dict):231 replace_prefix = {}232 state_dict_g = diffusers_convert.convert_text_enc_state_dict_v20(state_dict, "clip_g")233 for k in state_dict:234 if k.startswith("clip_l"):235 state_dict_g[k] = state_dict[k]236 237 state_dict_g["clip_l.transformer.text_model.embeddings.position_ids"] = torch.arange(77).expand((1, -1))238 pop_keys = ["clip_l.transformer.text_projection.weight", "clip_l.logit_scale"]239 for p in pop_keys:240 if p in state_dict_g:241 state_dict_g.pop(p)242 243 replace_prefix["clip_g"] = "conditioner.embedders.1.model"244 replace_prefix["clip_l"] = "conditioner.embedders.0"245 state_dict_g = utils.state_dict_prefix_replace(state_dict_g, replace_prefix)246 return state_dict_g247 248 def clip_target(self, state_dict={}):249 return supported_models_base.ClipTarget(sdxl_clip.SDXLTokenizer, sdxl_clip.SDXLClipModel)250 251class SSD1B(SDXL):252 unet_config = {253 "model_channels": 320,254 "use_linear_in_transformer": True,255 "transformer_depth": [0, 0, 2, 2, 4, 4],256 "context_dim": 2048,257 "adm_in_channels": 2816,258 "use_temporal_attention": False,259 }260 261class Segmind_Vega(SDXL):262 unet_config = {263 "model_channels": 320,264 "use_linear_in_transformer": True,265 "transformer_depth": [0, 0, 1, 1, 2, 2],266 "context_dim": 2048,267 "adm_in_channels": 2816,268 "use_temporal_attention": False,269 }270 271class KOALA_700M(SDXL):272 unet_config = {273 "model_channels": 320,274 "use_linear_in_transformer": True,275 "transformer_depth": [0, 2, 5],276 "context_dim": 2048,277 "adm_in_channels": 2816,278 "use_temporal_attention": False,279 }280 281class KOALA_1B(SDXL):282 unet_config = {283 "model_channels": 320,284 "use_linear_in_transformer": True,285 "transformer_depth": [0, 2, 6],286 "context_dim": 2048,287 "adm_in_channels": 2816,288 "use_temporal_attention": False,289 }290 291class SVD_img2vid(supported_models_base.BASE):292 unet_config = {293 "model_channels": 320,294 "in_channels": 8,295 "use_linear_in_transformer": True,296 "transformer_depth": [1, 1, 1, 1, 1, 1, 0, 0],297 "context_dim": 1024,298 "adm_in_channels": 768,299 "use_temporal_attention": True,300 "use_temporal_resblock": True301 }302 303 unet_extra_config = {304 "num_heads": -1,305 "num_head_channels": 64,306 "attn_precision": torch.float32,307 }308 309 clip_vision_prefix = "conditioner.embedders.0.open_clip.model.visual."310 311 latent_format = latent_formats.SD15312 313 sampling_settings = {"sigma_max": 700.0, "sigma_min": 0.002}314 315 def get_model(self, state_dict, prefix="", device=None):316 out = model_base.SVD_img2vid(self, device=device)317 return out318 319 def clip_target(self, state_dict={}):320 return None321 322class SV3D_u(SVD_img2vid):323 unet_config = {324 "model_channels": 320,325 "in_channels": 8,326 "use_linear_in_transformer": True,327 "transformer_depth": [1, 1, 1, 1, 1, 1, 0, 0],328 "context_dim": 1024,329 "adm_in_channels": 256,330 "use_temporal_attention": True,331 "use_temporal_resblock": True332 }333 334 vae_key_prefix = ["conditioner.embedders.1.encoder."]335 336 def get_model(self, state_dict, prefix="", device=None):337 out = model_base.SV3D_u(self, device=device)338 return out339 340class SV3D_p(SV3D_u):341 unet_config = {342 "model_channels": 320,343 "in_channels": 8,344 "use_linear_in_transformer": True,345 "transformer_depth": [1, 1, 1, 1, 1, 1, 0, 0],346 "context_dim": 1024,347 "adm_in_channels": 1280,348 "use_temporal_attention": True,349 "use_temporal_resblock": True350 }351 352 353 def get_model(self, state_dict, prefix="", device=None):354 out = model_base.SV3D_p(self, device=device)355 return out356 357class Stable_Zero123(supported_models_base.BASE):358 unet_config = {359 "context_dim": 768,360 "model_channels": 320,361 "use_linear_in_transformer": False,362 "adm_in_channels": None,363 "use_temporal_attention": False,364 "in_channels": 8,365 }366 367 unet_extra_config = {368 "num_heads": 8,369 "num_head_channels": -1,370 }371 372 required_keys = {373 "cc_projection.weight": None,374 "cc_projection.bias": None,375 }376 377 clip_vision_prefix = "cond_stage_model.model.visual."378 379 latent_format = latent_formats.SD15380 381 def get_model(self, state_dict, prefix="", device=None):382 out = model_base.Stable_Zero123(self, device=device, cc_projection_weight=state_dict["cc_projection.weight"], cc_projection_bias=state_dict["cc_projection.bias"])383 return out384 385 def clip_target(self, state_dict={}):386 return None387 388class SD_X4Upscaler(SD20):389 unet_config = {390 "context_dim": 1024,391 "model_channels": 256,392 'in_channels': 7,393 "use_linear_in_transformer": True,394 "adm_in_channels": None,395 "use_temporal_attention": False,396 }397 398 unet_extra_config = {399 "disable_self_attentions": [True, True, True, False],400 "num_classes": 1000,401 "num_heads": 8,402 "num_head_channels": -1,403 }404 405 latent_format = latent_formats.SD_X4406 407 sampling_settings = {408 "linear_start": 0.0001,409 "linear_end": 0.02,410 }411 412 def get_model(self, state_dict, prefix="", device=None):413 out = model_base.SD_X4Upscaler(self, device=device)414 return out415 416class Stable_Cascade_C(supported_models_base.BASE):417 unet_config = {418 "stable_cascade_stage": 'c',419 }420 421 unet_extra_config = {}422 423 latent_format = latent_formats.SC_Prior424 supported_inference_dtypes = [torch.bfloat16, torch.float32]425 426 sampling_settings = {427 "shift": 2.0,428 }429 430 vae_key_prefix = ["vae."]431 text_encoder_key_prefix = ["text_encoder."]432 clip_vision_prefix = "clip_l_vision."433 434 def process_unet_state_dict(self, state_dict):435 key_list = list(state_dict.keys())436 for y in ["weight", "bias"]:437 suffix = "in_proj_{}".format(y)438 keys = filter(lambda a: a.endswith(suffix), key_list)439 for k_from in keys:440 weights = state_dict.pop(k_from)441 prefix = k_from[:-(len(suffix) + 1)]442 shape_from = weights.shape[0] // 3443 for x in range(3):444 p = ["to_q", "to_k", "to_v"]445 k_to = "{}.{}.{}".format(prefix, p[x], y)446 state_dict[k_to] = weights[shape_from*x:shape_from*(x + 1)]447 return state_dict448 449 def process_clip_state_dict(self, state_dict):450 state_dict = utils.state_dict_prefix_replace(state_dict, {k: "" for k in self.text_encoder_key_prefix}, filter_keys=True)451 if "clip_g.text_projection" in state_dict:452 state_dict["clip_g.transformer.text_projection.weight"] = state_dict.pop("clip_g.text_projection").transpose(0, 1)453 return state_dict454 455 def get_model(self, state_dict, prefix="", device=None):456 out = model_base.StableCascade_C(self, device=device)457 return out458 459 def clip_target(self, state_dict={}):460 return supported_models_base.ClipTarget(sdxl_clip.StableCascadeTokenizer, sdxl_clip.StableCascadeClipModel)461 462class Stable_Cascade_B(Stable_Cascade_C):463 unet_config = {464 "stable_cascade_stage": 'b',465 }466 467 unet_extra_config = {}468 469 latent_format = latent_formats.SC_B470 supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]471 472 sampling_settings = {473 "shift": 1.0,474 }475 476 clip_vision_prefix = None477 478 def get_model(self, state_dict, prefix="", device=None):479 out = model_base.StableCascade_B(self, device=device)480 return out481 482class SD15_instructpix2pix(SD15):483 unet_config = {484 "context_dim": 768,485 "model_channels": 320,486 "use_linear_in_transformer": False,487 "adm_in_channels": None,488 "use_temporal_attention": False,489 "in_channels": 8,490 }491 492 def get_model(self, state_dict, prefix="", device=None):493 return model_base.SD15_instructpix2pix(self, device=device)494 495class SDXL_instructpix2pix(SDXL):496 unet_config = {497 "model_channels": 320,498 "use_linear_in_transformer": True,499 "transformer_depth": [0, 0, 2, 2, 10, 10],500 "context_dim": 2048,501 "adm_in_channels": 2816,502 "use_temporal_attention": False,503 "in_channels": 8,504 }505 506 def get_model(self, state_dict, prefix="", device=None):507 return model_base.SDXL_instructpix2pix(self, model_type=self.model_type(state_dict, prefix), device=device)508 509class SD3(supported_models_base.BASE):510 unet_config = {511 "in_channels": 16,512 "pos_embed_scaling_factor": None,513 }514 515 sampling_settings = {516 "shift": 3.0,517 }518 519 unet_extra_config = {}520 latent_format = latent_formats.SD3521 522 memory_usage_factor = 1.2523 524 text_encoder_key_prefix = ["text_encoders."]525 526 def get_model(self, state_dict, prefix="", device=None):527 out = model_base.SD3(self, device=device)528 return out529 530 def clip_target(self, state_dict={}):531 clip_l = False532 clip_g = False533 t5 = False534 pref = self.text_encoder_key_prefix[0]535 if "{}clip_l.transformer.text_model.final_layer_norm.weight".format(pref) in state_dict:536 clip_l = True537 if "{}clip_g.transformer.text_model.final_layer_norm.weight".format(pref) in state_dict:538 clip_g = True539 t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref))540 if "dtype_t5" in t5_detect:541 t5 = True542 543 return supported_models_base.ClipTarget(comfy.text_encoders.sd3_clip.SD3Tokenizer, comfy.text_encoders.sd3_clip.sd3_clip(clip_l=clip_l, clip_g=clip_g, t5=t5, **t5_detect))544 545class StableAudio(supported_models_base.BASE):546 unet_config = {547 "audio_model": "dit1.0",548 }549 550 sampling_settings = {"sigma_max": 500.0, "sigma_min": 0.03}551 552 unet_extra_config = {}553 latent_format = latent_formats.StableAudio1554 555 text_encoder_key_prefix = ["text_encoders."]556 vae_key_prefix = ["pretransform.model."]557 558 def get_model(self, state_dict, prefix="", device=None):559 seconds_start_sd = utils.state_dict_prefix_replace(state_dict, {"conditioner.conditioners.seconds_start.": ""}, filter_keys=True)560 seconds_total_sd = utils.state_dict_prefix_replace(state_dict, {"conditioner.conditioners.seconds_total.": ""}, filter_keys=True)561 return model_base.StableAudio1(self, seconds_start_embedder_weights=seconds_start_sd, seconds_total_embedder_weights=seconds_total_sd, device=device)562 563 def process_unet_state_dict(self, state_dict):564 for k in list(state_dict.keys()):565 if k.endswith(".cross_attend_norm.beta") or k.endswith(".ff_norm.beta") or k.endswith(".pre_norm.beta"): #These weights are all zero566 state_dict.pop(k)567 return state_dict568 569 def process_unet_state_dict_for_saving(self, state_dict):570 replace_prefix = {"": "model.model."}571 return utils.state_dict_prefix_replace(state_dict, replace_prefix)572 573 def clip_target(self, state_dict={}):574 return supported_models_base.ClipTarget(comfy.text_encoders.sa_t5.SAT5Tokenizer, comfy.text_encoders.sa_t5.SAT5Model)575 576class AuraFlow(supported_models_base.BASE):577 unet_config = {578 "cond_seq_dim": 2048,579 }580 581 sampling_settings = {582 "multiplier": 1.0,583 "shift": 1.73,584 }585 586 unet_extra_config = {}587 latent_format = latent_formats.SDXL588 589 vae_key_prefix = ["vae."]590 text_encoder_key_prefix = ["text_encoders."]591 592 def get_model(self, state_dict, prefix="", device=None):593 out = model_base.AuraFlow(self, device=device)594 return out595 596 def clip_target(self, state_dict={}):597 return supported_models_base.ClipTarget(comfy.text_encoders.aura_t5.AuraT5Tokenizer, comfy.text_encoders.aura_t5.AuraT5Model)598 599class PixArtAlpha(supported_models_base.BASE):600 unet_config = {601 "image_model": "pixart_alpha",602 }603 604 sampling_settings = {605 "beta_schedule" : "sqrt_linear",606 "linear_start" : 0.0001,607 "linear_end" : 0.02,608 "timesteps" : 1000,609 }610 611 unet_extra_config = {}612 latent_format = latent_formats.SD15613 614 memory_usage_factor = 0.5615 616 vae_key_prefix = ["vae."]617 text_encoder_key_prefix = ["text_encoders."]618 619 def get_model(self, state_dict, prefix="", device=None):620 out = model_base.PixArt(self, device=device)621 return out.eval()622 623 def clip_target(self, state_dict={}):624 return supported_models_base.ClipTarget(comfy.text_encoders.pixart_t5.PixArtTokenizer, comfy.text_encoders.pixart_t5.PixArtT5XXL)625 626class PixArtSigma(PixArtAlpha):627 unet_config = {628 "image_model": "pixart_sigma",629 }630 latent_format = latent_formats.SDXL631 632class HunyuanDiT(supported_models_base.BASE):633 unet_config = {634 "image_model": "hydit",635 }636 637 unet_extra_config = {638 "attn_precision": torch.float32,639 }640 641 sampling_settings = {642 "linear_start": 0.00085,643 "linear_end": 0.018,644 }645 646 latent_format = latent_formats.SDXL647 648 memory_usage_factor = 1.3649 650 vae_key_prefix = ["vae."]651 text_encoder_key_prefix = ["text_encoders."]652 653 def get_model(self, state_dict, prefix="", device=None):654 out = model_base.HunyuanDiT(self, device=device)655 return out656 657 def clip_target(self, state_dict={}):658 return supported_models_base.ClipTarget(comfy.text_encoders.hydit.HyditTokenizer, comfy.text_encoders.hydit.HyditModel)659 660class HunyuanDiT1(HunyuanDiT):661 unet_config = {662 "image_model": "hydit1",663 }664 665 unet_extra_config = {}666 667 sampling_settings = {668 "linear_start" : 0.00085,669 "linear_end" : 0.03,670 }671 672class Flux(supported_models_base.BASE):673 unet_config = {674 "image_model": "flux",675 "guidance_embed": True,676 }677 678 sampling_settings = {679 }680 681 unet_extra_config = {}682 latent_format = latent_formats.Flux683 684 memory_usage_factor = 2.8685 686 supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]687 688 vae_key_prefix = ["vae."]689 text_encoder_key_prefix = ["text_encoders."]690 691 def get_model(self, state_dict, prefix="", device=None):692 out = model_base.Flux(self, device=device)693 return out694 695 def clip_target(self, state_dict={}):696 pref = self.text_encoder_key_prefix[0]697 t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref))698 return supported_models_base.ClipTarget(comfy.text_encoders.flux.FluxTokenizer, comfy.text_encoders.flux.flux_clip(**t5_detect))699 700class FluxInpaint(Flux):701 unet_config = {702 "image_model": "flux",703 "guidance_embed": True,704 "in_channels": 96,705 }706 707 supported_inference_dtypes = [torch.bfloat16, torch.float32]708 709class FluxSchnell(Flux):710 unet_config = {711 "image_model": "flux",712 "guidance_embed": False,713 }714 715 sampling_settings = {716 "multiplier": 1.0,717 "shift": 1.0,718 }719 720 def get_model(self, state_dict, prefix="", device=None):721 out = model_base.Flux(self, model_type=model_base.ModelType.FLOW, device=device)722 return out723 724class GenmoMochi(supported_models_base.BASE):725 unet_config = {726 "image_model": "mochi_preview",727 }728 729 sampling_settings = {730 "multiplier": 1.0,731 "shift": 6.0,732 }733 734 unet_extra_config = {}735 latent_format = latent_formats.Mochi736 737 memory_usage_factor = 2.0 #TODO738 739 supported_inference_dtypes = [torch.bfloat16, torch.float32]740 741 vae_key_prefix = ["vae."]742 text_encoder_key_prefix = ["text_encoders."]743 744 def get_model(self, state_dict, prefix="", device=None):745 out = model_base.GenmoMochi(self, device=device)746 return out747 748 def clip_target(self, state_dict={}):749 pref = self.text_encoder_key_prefix[0]750 t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref))751 return supported_models_base.ClipTarget(comfy.text_encoders.genmo.MochiT5Tokenizer, comfy.text_encoders.genmo.mochi_te(**t5_detect))752 753class LTXV(supported_models_base.BASE):754 unet_config = {755 "image_model": "ltxv",756 }757 758 sampling_settings = {759 "shift": 2.37,760 }761 762 unet_extra_config = {}763 latent_format = latent_formats.LTXV764 765 memory_usage_factor = 5.5 # TODO: img2vid is about 2x vs txt2vid766 767 supported_inference_dtypes = [torch.bfloat16, torch.float32]768 769 vae_key_prefix = ["vae."]770 text_encoder_key_prefix = ["text_encoders."]771 772 def get_model(self, state_dict, prefix="", device=None):773 out = model_base.LTXV(self, device=device)774 return out775 776 def clip_target(self, state_dict={}):777 pref = self.text_encoder_key_prefix[0]778 t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref))779 return supported_models_base.ClipTarget(comfy.text_encoders.lt.LTXVT5Tokenizer, comfy.text_encoders.lt.ltxv_te(**t5_detect))780 781class HunyuanVideo(supported_models_base.BASE):782 unet_config = {783 "image_model": "hunyuan_video",784 }785 786 sampling_settings = {787 "shift": 7.0,788 }789 790 unet_extra_config = {}791 latent_format = latent_formats.HunyuanVideo792 793 memory_usage_factor = 1.8 #TODO794 795 supported_inference_dtypes = [torch.bfloat16, torch.float32]796 797 vae_key_prefix = ["vae."]798 text_encoder_key_prefix = ["text_encoders."]799 800 def get_model(self, state_dict, prefix="", device=None):801 out = model_base.HunyuanVideo(self, device=device)802 return out803 804 def process_unet_state_dict(self, state_dict):805 out_sd = {}806 for k in list(state_dict.keys()):807 key_out = k808 key_out = key_out.replace("txt_in.t_embedder.mlp.0.", "txt_in.t_embedder.in_layer.").replace("txt_in.t_embedder.mlp.2.", "txt_in.t_embedder.out_layer.")809 key_out = key_out.replace("txt_in.c_embedder.linear_1.", "txt_in.c_embedder.in_layer.").replace("txt_in.c_embedder.linear_2.", "txt_in.c_embedder.out_layer.")810 key_out = key_out.replace("_mod.linear.", "_mod.lin.").replace("_attn_qkv.", "_attn.qkv.")811 key_out = key_out.replace("mlp.fc1.", "mlp.0.").replace("mlp.fc2.", "mlp.2.")812 key_out = key_out.replace("_attn_q_norm.weight", "_attn.norm.query_norm.scale").replace("_attn_k_norm.weight", "_attn.norm.key_norm.scale")813 key_out = key_out.replace(".q_norm.weight", ".norm.query_norm.scale").replace(".k_norm.weight", ".norm.key_norm.scale")814 key_out = key_out.replace("_attn_proj.", "_attn.proj.")815 key_out = key_out.replace(".modulation.linear.", ".modulation.lin.")816 key_out = key_out.replace("_in.mlp.2.", "_in.out_layer.").replace("_in.mlp.0.", "_in.in_layer.")817 out_sd[key_out] = state_dict[k]818 return out_sd819 820 def process_unet_state_dict_for_saving(self, state_dict):821 replace_prefix = {"": "model.model."}822 return utils.state_dict_prefix_replace(state_dict, replace_prefix)823 824 def clip_target(self, state_dict={}):825 pref = self.text_encoder_key_prefix[0]826 hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}llama.transformer.".format(pref))827 return supported_models_base.ClipTarget(comfy.text_encoders.hunyuan_video.HunyuanVideoTokenizer, comfy.text_encoders.hunyuan_video.hunyuan_video_clip(**hunyuan_detect))828 829class HunyuanVideoI2V(HunyuanVideo):830 unet_config = {831 "image_model": "hunyuan_video",832 "in_channels": 33,833 }834 835 def get_model(self, state_dict, prefix="", device=None):836 out = model_base.HunyuanVideoI2V(self, device=device)837 return out838 839class HunyuanVideoSkyreelsI2V(HunyuanVideo):840 unet_config = {841 "image_model": "hunyuan_video",842 "in_channels": 32,843 }844 845 def get_model(self, state_dict, prefix="", device=None):846 out = model_base.HunyuanVideoSkyreelsI2V(self, device=device)847 return out848 849class CosmosT2V(supported_models_base.BASE):850 unet_config = {851 "image_model": "cosmos",852 "in_channels": 16,853 }854 855 sampling_settings = {856 "sigma_data": 0.5,857 "sigma_max": 80.0,858 "sigma_min": 0.002,859 }860 861 unet_extra_config = {}862 latent_format = latent_formats.Cosmos1CV8x8x8863 864 memory_usage_factor = 1.6 #TODO865 866 supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32] #TODO867 868 vae_key_prefix = ["vae."]869 text_encoder_key_prefix = ["text_encoders."]870 871 def get_model(self, state_dict, prefix="", device=None):872 out = model_base.CosmosVideo(self, device=device)873 return out874 875 def clip_target(self, state_dict={}):876 pref = self.text_encoder_key_prefix[0]877 t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref))878 return supported_models_base.ClipTarget(comfy.text_encoders.cosmos.CosmosT5Tokenizer, comfy.text_encoders.cosmos.te(**t5_detect))879 880class CosmosI2V(CosmosT2V):881 unet_config = {882 "image_model": "cosmos",883 "in_channels": 17,884 }885 886 def get_model(self, state_dict, prefix="", device=None):887 out = model_base.CosmosVideo(self, image_to_video=True, device=device)888 return out889 890class Lumina2(supported_models_base.BASE):891 unet_config = {892 "image_model": "lumina2",893 }894 895 sampling_settings = {896 "multiplier": 1.0,897 "shift": 6.0,898 }899 900 memory_usage_factor = 1.2901 902 unet_extra_config = {}903 latent_format = latent_formats.Flux904 905 supported_inference_dtypes = [torch.bfloat16, torch.float32]906 907 vae_key_prefix = ["vae."]908 text_encoder_key_prefix = ["text_encoders."]909 910 def get_model(self, state_dict, prefix="", device=None):911 out = model_base.Lumina2(self, device=device)912 return out913 914 def clip_target(self, state_dict={}):915 pref = self.text_encoder_key_prefix[0]916 hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}gemma2_2b.transformer.".format(pref))917 return supported_models_base.ClipTarget(comfy.text_encoders.lumina2.LuminaTokenizer, comfy.text_encoders.lumina2.te(**hunyuan_detect))918 919class WAN21_T2V(supported_models_base.BASE):920 unet_config = {921 "image_model": "wan2.1",922 "model_type": "t2v",923 }924 925 sampling_settings = {926 "shift": 8.0,927 }928 929 unet_extra_config = {}930 latent_format = latent_formats.Wan21931 932 memory_usage_factor = 1.0933 934 supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]935 936 vae_key_prefix = ["vae."]937 text_encoder_key_prefix = ["text_encoders."]938 939 def __init__(self, unet_config):940 super().__init__(unet_config)941 self.memory_usage_factor = self.unet_config.get("dim", 2000) / 2000942 943 def get_model(self, state_dict, prefix="", device=None):944 out = model_base.WAN21(self, device=device)945 return out946 947 def clip_target(self, state_dict={}):948 pref = self.text_encoder_key_prefix[0]949 t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}umt5xxl.transformer.".format(pref))950 return supported_models_base.ClipTarget(comfy.text_encoders.wan.WanT5Tokenizer, comfy.text_encoders.wan.te(**t5_detect))951 952class WAN21_I2V(WAN21_T2V):953 unet_config = {954 "image_model": "wan2.1",955 "model_type": "i2v",956 }957 958 def get_model(self, state_dict, prefix="", device=None):959 out = model_base.WAN21(self, image_to_video=True, device=device)960 return out961 962models = [Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, Lumina2, WAN21_T2V, WAN21_I2V]963 964models += [SVD_img2vid]965