nodes/city96-ComfyUI_ExtraModels
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

HunYuanDiT: Add 1.2 support

Browse Source
This commit is contained in:
City
2024-07-04 18:01:53 +02:00
parent 62b84531f4
commit 40d2a76310
3 changed files with 54 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
HunYuanDiT/models/models.py
+24 -13
View File
@@ -169,6 +169,8 @@ class HunYuanDiT(nn.Module):
num_heads=16,
mlp_ratio=4.0,
log_fn=print,
cond_style=True,
cond_res=True,
**kwargs,
):
super().__init__()
@@ -187,6 +189,8 @@ class HunYuanDiT(nn.Module):
self.text_len = args.text_len
self.text_len_t5 = args.text_len_t5
self.norm = args.norm
self.cond_res = cond_res
self.cond_style = cond_style
use_flash_attn = args.infer_mode == 'fa'
if use_flash_attn:
@@ -205,11 +209,15 @@ class HunYuanDiT(nn.Module):
# Attention pooling
self.pooler = AttentionPool(self.text_len_t5, self.text_states_dim_t5, num_heads=8, output_dim=1024)
# Here we use a default learned embedder layer for future extension.
self.style_embedder = nn.Embedding(1, hidden_size)
# Image size and crop size conditions
self.extra_in_dim = 256 * 6 + hidden_size
self.extra_in_dim = 0
if self.cond_res:
# Image size and crop size conditions
self.extra_in_dim += 256 * 6
if self.cond_style:
# Here we use a default learned embedder layer for future extension.
self.style_embedder = nn.Embedding(1, hidden_size)
self.extra_in_dim += hidden_size
# Text embedding for `add`
self.last_size = input_size
@@ -310,16 +318,19 @@ class HunYuanDiT(nn.Module):
# Build text tokens with pooling
extra_vec = self.pooler(encoder_hidden_states_t5)
# Build image meta size tokens
image_meta_size = timestep_embedding(image_meta_size.view(-1), 256) # [B * 6, 256]
# if self.args.use_fp16:
# image_meta_size = image_meta_size.half()
image_meta_size = image_meta_size.view(-1, 6 * 256)
extra_vec = torch.cat([extra_vec, image_meta_size], dim=1) # [B, D + 6 * 256]
if self.cond_res:
# Build image meta size tokens
image_meta_size = timestep_embedding(image_meta_size.view(-1), 256) # [B * 6, 256]
# if self.args.use_fp16:
# image_meta_size = image_meta_size.half()
image_meta_size = image_meta_size.view(-1, 6 * 256)
extra_vec = torch.cat([extra_vec, image_meta_size], dim=1) # [B, D + 6 * 256]
# Build style tokens
style_embedding = self.style_embedder(style)
extra_vec = torch.cat([extra_vec, style_embedding], dim=1)
if self.cond_style:
# Build style tokens
style_embedding = self.style_embedder(style)
extra_vec = torch.cat([extra_vec, style_embedding], dim=1)
# Concatenate all extra vectors
c = t + self.extra_embedder(extra_vec.to(self.dtype)) # [B, D]