mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-11-12 12:13:29 +08:00
116 lines
4.1 KiB
Python
116 lines
4.1 KiB
Python
import comfy.supported_models_base
|
|
import comfy.latent_formats
|
|
import comfy.model_patcher
|
|
import comfy.model_base
|
|
import comfy.utils
|
|
import torch
|
|
from collections import namedtuple
|
|
from .hydit.modules.models import HunYuanDiT as HYDiT
|
|
import pdb
|
|
from .hydit.modules.posemb_layers import get_2d_rotary_pos_embed, get_fill_resize_and_crop
|
|
from .hydit.modules.models import HUNYUAN_DIT_CONFIG
|
|
|
|
def calc_rope(height, width):
|
|
th = height // 8 // HUNYUAN_DIT_CONFIG['DiT-g/2']['patch_size']
|
|
tw = width // 8 // HUNYUAN_DIT_CONFIG['DiT-g/2']['patch_size']
|
|
base_size = 512 // 8 // HUNYUAN_DIT_CONFIG['DiT-g/2']['patch_size']
|
|
start, stop = get_fill_resize_and_crop((th, tw), base_size)
|
|
sub_args = [start, stop, (th, tw)]
|
|
head_size = HUNYUAN_DIT_CONFIG['DiT-g/2']['hidden_size'] // HUNYUAN_DIT_CONFIG['DiT-g/2']['num_heads']
|
|
rope = get_2d_rotary_pos_embed(head_size, *sub_args)
|
|
return rope
|
|
|
|
def batch_embeddings(embeds, batch_size):
|
|
bs_embed, seq_len, _ = embeds.shape
|
|
embeds = embeds.repeat(1, batch_size, 1)
|
|
embeds = embeds.view(bs_embed * batch_size, seq_len, -1)
|
|
return embeds
|
|
|
|
class HunYuan_DiT(comfy.supported_models_base.BASE):
|
|
Conf = namedtuple('DiT', ['learn_sigma', 'text_states_dim', 'text_states_dim_t5', 'text_len', 'text_len_t5', 'norm', 'infer_mode',
|
|
'use_fp16', 'use_flash_attn', 'qk_norm', 'size_cond', 'use_style_cond'])
|
|
conf = {
|
|
'learn_sigma': True,
|
|
'text_states_dim': 1024,
|
|
'text_states_dim_t5': 2048,
|
|
'text_len': 77,
|
|
'text_len_t5': 256,
|
|
'norm': 'layer',
|
|
'infer_mode': 'torch',
|
|
'use_fp16': True,
|
|
'use_flash_attn': False,
|
|
'qk_norm': True,
|
|
'size_cond': None,
|
|
'use_style_cond': False,
|
|
|
|
}
|
|
|
|
unet_config = {}
|
|
unet_extra_config = {
|
|
"num_heads": 16
|
|
}
|
|
latent_format = comfy.latent_formats.SDXL
|
|
|
|
dit_conf = Conf(**conf)
|
|
|
|
def __init__(self, model_conf):
|
|
self.unet_config = model_conf.get("unet_config", {})
|
|
#print(model_conf)
|
|
print(self.unet_config)
|
|
self.sampling_settings = model_conf.get("sampling_settings", {})
|
|
self.latent_format = self.latent_format()
|
|
self.unet_config["disable_unet_model_creation"] = True
|
|
#self.unet_config["disable_unet_model_creation"] = self.unet_config.get("disable_unet_model_creation", True)
|
|
|
|
def model_type(self, state_dict, prefix=""):
|
|
return comfy.model_base.ModelType.V_PREDICTION
|
|
|
|
class HYDiT_Model(comfy.model_base.BaseModel):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def extra_conds(self, **kwargs):
|
|
out = super().extra_conds(**kwargs)
|
|
addit_embeds = kwargs['cross_attn'].addit_embeds
|
|
for name in addit_embeds:
|
|
out[name] = comfy.conds.CONDRegular(addit_embeds[name])
|
|
|
|
return out
|
|
|
|
class ModifiedHunYuanDiT(HYDiT):
|
|
def forward_core(self, *args, **kwargs):
|
|
return super().forward(*args, **kwargs)
|
|
|
|
def forward(self, x, timesteps, context, t5_embeds=None, attention_mask=None, t5_attention_mask=None, image_meta_size=None, **kwargs):
|
|
batch_size, _, height, width = x.shape
|
|
ratio = 8
|
|
#assert(0)
|
|
#pdb.set_trace()
|
|
|
|
|
|
style = None
|
|
size_cond = None
|
|
image_meta_size = None
|
|
rope = calc_rope(int(ratio*height), int(ratio*width))
|
|
#import pdb
|
|
#pdb.set_trace()
|
|
|
|
|
|
noise_pred = self.forward_core(
|
|
x = x.to(self.dtype),
|
|
t = timesteps.to(self.dtype),
|
|
encoder_hidden_states = context.to(self.dtype),
|
|
text_embedding_mask = attention_mask.to(self.dtype),
|
|
encoder_hidden_states_t5 = t5_embeds.to(self.dtype),
|
|
text_embedding_mask_t5 = t5_attention_mask.to(self.dtype),
|
|
image_meta_size = image_meta_size,
|
|
style = style,
|
|
cos_cis_img = rope[0],
|
|
sin_cis_img = rope[1],
|
|
return_dict=False
|
|
)
|
|
noise_pred = noise_pred.to(torch.float)
|
|
eps, _ = noise_pred[:, :self.in_channels], noise_pred[:, self.in_channels:]
|
|
return eps
|
|
|