xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-11-12 12:13:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix code format style

Browse Source
This commit is contained in:
CrazyBoyM 2024-08-07 18:38:14 +08:00
parent 2b68ab0a42
commit bda7c82134
1 changed files with 153 additions and 82 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

235
comfy/ldm/hydit/controlnet.py
View File

@ -6,7 +6,12 @@ import torch.nn.functional as F
from torch.utils import checkpoint from torch.utils import checkpoint
from comfy.ldm.modules.diffusionmodules.mmdit import Mlp, TimestepEmbedder, PatchEmbed, RMSNorm from comfy.ldm.modules.diffusionmodules.mmdit import (
Mlp,
TimestepEmbedder,
PatchEmbed,
RMSNorm,
)
from comfy.ldm.modules.diffusionmodules.util import timestep_embedding from comfy.ldm.modules.diffusionmodules.util import timestep_embedding
from .poolers import AttentionPool from .poolers import AttentionPool
@ -15,11 +20,13 @@ from .models import HunYuanDiTBlock
from .posemb_layers import get_2d_rotary_pos_embed, get_fill_resize_and_crop from .posemb_layers import get_2d_rotary_pos_embed, get_fill_resize_and_crop
def zero_module(module): def zero_module(module):
for p in module.parameters(): for p in module.parameters():
nn.init.zeros_(p) nn.init.zeros_(p)
return module return module
def calc_rope(x, patch_size, head_size): def calc_rope(x, patch_size, head_size):
th = (x.shape[2] + (patch_size // 2)) // patch_size th = (x.shape[2] + (patch_size // 2)) // patch_size
tw = (x.shape[3] + (patch_size // 2)) // patch_size tw = (x.shape[3] + (patch_size // 2)) // patch_size
@ -60,29 +67,31 @@ class HunYuanControlNet(nn.Module):
log_fn: callable log_fn: callable
The logging function. The logging function.
""" """
def __init__(self,
input_size: tuple = 128, def __init__(
patch_size: int = 2, self,
in_channels: int = 4, input_size: tuple = 128,
hidden_size: int = 1408, patch_size: int = 2,
depth: int = 40, in_channels: int = 4,
num_heads: int = 16, hidden_size: int = 1408,
mlp_ratio: float = 4.3637, depth: int = 40,
text_states_dim = 1024, num_heads: int = 16,
text_states_dim_t5 = 2048, mlp_ratio: float = 4.3637,
text_len = 77, text_states_dim=1024,
text_len_t5 = 256, text_states_dim_t5=2048,
qk_norm = True,# See http://arxiv.org/abs/2302.05442 for details. text_len=77,
size_cond = False, text_len_t5=256,
use_style_cond = False, qk_norm=True, # See http://arxiv.org/abs/2302.05442 for details.
learn_sigma = True, size_cond=False,
norm = "layer", use_style_cond=False,
log_fn: callable = print, learn_sigma=True,
attn_precision=None, norm="layer",
dtype=None, log_fn: callable = print,
device=None, attn_precision=None,
operations=None, dtype=None,
**kwargs, device=None,
operations=None,
**kwargs,
): ):
super().__init__() super().__init__()
self.log_fn = log_fn self.log_fn = log_fn
@ -103,19 +112,44 @@ class HunYuanControlNet(nn.Module):
self.dtype = dtype self.dtype = dtype
self.latent_format = comfy.latent_formats.SDXL self.latent_format = comfy.latent_formats.SDXL
self.mlp_t5 = nn.Sequential( self.mlp_t5 = nn.Sequential(
nn.Linear(self.text_states_dim_t5, self.text_states_dim_t5 * 4, bias=True, dtype=dtype, device=device), nn.Linear(
self.text_states_dim_t5,
self.text_states_dim_t5 * 4,
bias=True,
dtype=dtype,
device=device,
),
nn.SiLU(), nn.SiLU(),
nn.Linear(self.text_states_dim_t5 * 4, self.text_states_dim, bias=True, dtype=dtype, device=device), nn.Linear(
self.text_states_dim_t5 * 4,
self.text_states_dim,
bias=True,
dtype=dtype,
device=device,
),
) )
# learnable replace # learnable replace
self.text_embedding_padding = nn.Parameter( self.text_embedding_padding = nn.Parameter(
torch.randn(self.text_len + self.text_len_t5, self.text_states_dim, dtype=dtype, device=device)) torch.randn(
self.text_len + self.text_len_t5,
self.text_states_dim,
dtype=dtype,
device=device,
)
)
# Attention pooling # Attention pooling
pooler_out_dim = 1024 pooler_out_dim = 1024
self.pooler = AttentionPool(self.text_len_t5, self.text_states_dim_t5, num_heads=8, output_dim=pooler_out_dim, dtype=dtype, device=device, operations=operations) self.pooler = AttentionPool(
self.text_len_t5,
self.text_states_dim_t5,
num_heads=8,
output_dim=pooler_out_dim,
dtype=dtype,
device=device,
operations=operations,
)
# Dimension of the extra input vectors # Dimension of the extra input vectors
self.extra_in_dim = pooler_out_dim self.extra_in_dim = pooler_out_dim
@ -126,60 +160,89 @@ class HunYuanControlNet(nn.Module):
if self.use_style_cond: if self.use_style_cond:
# Here we use a default learned embedder layer for future extension. # Here we use a default learned embedder layer for future extension.
self.style_embedder = nn.Embedding(1, hidden_size, dtype=dtype, device=device) self.style_embedder = nn.Embedding(
1, hidden_size, dtype=dtype, device=device
)
self.extra_in_dim += hidden_size self.extra_in_dim += hidden_size
# Text embedding for `add` # Text embedding for `add`
self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, dtype=dtype, device=device, operations=operations) self.x_embedder = PatchEmbed(
self.t_embedder = TimestepEmbedder(hidden_size, dtype=dtype, device=device, operations=operations) input_size,
patch_size,
in_channels,
hidden_size,
dtype=dtype,
device=device,
operations=operations,
)
self.t_embedder = TimestepEmbedder(
hidden_size, dtype=dtype, device=device, operations=operations
)
self.extra_embedder = nn.Sequential( self.extra_embedder = nn.Sequential(
operations.Linear(self.extra_in_dim, hidden_size * 4, dtype=dtype, device=device), operations.Linear(
self.extra_in_dim, hidden_size * 4, dtype=dtype, device=device
),
nn.SiLU(), nn.SiLU(),
operations.Linear(hidden_size * 4, hidden_size, bias=True, dtype=dtype, device=device), operations.Linear(
hidden_size * 4, hidden_size, bias=True, dtype=dtype, device=device
),
) )
# Image embedding # Image embedding
num_patches = self.x_embedder.num_patches num_patches = self.x_embedder.num_patches
# HUnYuanDiT Blocks # HUnYuanDiT Blocks
self.blocks = nn.ModuleList([ self.blocks = nn.ModuleList(
HunYuanDiTBlock(hidden_size=hidden_size, [
c_emb_size=hidden_size, HunYuanDiTBlock(
num_heads=num_heads, hidden_size=hidden_size,
mlp_ratio=mlp_ratio, c_emb_size=hidden_size,
text_states_dim=self.text_states_dim, num_heads=num_heads,
qk_norm=qk_norm, mlp_ratio=mlp_ratio,
norm_type=self.norm, text_states_dim=self.text_states_dim,
skip=False, qk_norm=qk_norm,
attn_precision=attn_precision, norm_type=self.norm,
dtype=dtype, skip=False,
device=device, attn_precision=attn_precision,
operations=operations, dtype=dtype,
) device=device,
for _ in range(19) operations=operations,
]) )
for _ in range(19)
]
)
# Input zero linear for the first block # Input zero linear for the first block
self.before_proj = zero_module(nn.Linear(self.hidden_size, self.hidden_size, dtype=dtype, device=device)) self.before_proj = zero_module(
nn.Linear(self.hidden_size, self.hidden_size, dtype=dtype, device=device)
)
# Output zero linear for the every block # Output zero linear for the every block
self.after_proj_list = nn.ModuleList( self.after_proj_list = nn.ModuleList(
[zero_module(nn.Linear(self.hidden_size, self.hidden_size, dtype=dtype, device=device)) for _ in range(len(self.blocks))] [
zero_module(
nn.Linear(
self.hidden_size, self.hidden_size, dtype=dtype, device=device
)
)
for _ in range(len(self.blocks))
]
) )
def forward(self, def forward(
x: torch.Tensor, self,
t: torch.Tensor = None, x: torch.Tensor,
condition = None, t: torch.Tensor = None,
encoder_hidden_states: Optional[torch.Tensor] = None, condition=None,
text_embedding_mask=None, encoder_hidden_states: Optional[torch.Tensor] = None,
encoder_hidden_states_t5=None, text_embedding_mask=None,
text_embedding_mask_t5=None, encoder_hidden_states_t5=None,
image_meta_size=None, text_embedding_mask_t5=None,
style=None, image_meta_size=None,
transformer_options=None, style=None,
**kwarg transformer_options=None,
): **kwarg,
):
""" """
Forward pass of the encoder. Forward pass of the encoder.
@ -208,32 +271,41 @@ class HunYuanControlNet(nn.Module):
""" """
if condition.shape[0] == 1: if condition.shape[0] == 1:
condition = torch.repeat_interleave(condition, x.shape[0], dim=0) condition = torch.repeat_interleave(condition, x.shape[0], dim=0)
text_states = encoder_hidden_states # 2,77,1024 text_states = encoder_hidden_states # 2,77,1024
text_states_t5 = encoder_hidden_states_t5 # 2,256,2048 text_states_t5 = encoder_hidden_states_t5 # 2,256,2048
text_states_mask = text_embedding_mask.bool() # 2,77 text_states_mask = text_embedding_mask.bool() # 2,77
text_states_t5_mask = text_embedding_mask_t5.bool() # 2,256 text_states_t5_mask = text_embedding_mask_t5.bool() # 2,256
b_t5, l_t5, c_t5 = text_states_t5.shape b_t5, l_t5, c_t5 = text_states_t5.shape
text_states_t5 = self.mlp_t5(text_states_t5.view(-1, c_t5)).view(b_t5, l_t5, -1) text_states_t5 = self.mlp_t5(text_states_t5.view(-1, c_t5)).view(b_t5, l_t5, -1)
padding = comfy.ops.cast_to_input(self.text_embedding_padding, text_states) padding = comfy.ops.cast_to_input(self.text_embedding_padding, text_states)
text_states[:,-self.text_len:] = torch.where(text_states_mask[:,-self.text_len:].unsqueeze(2), text_states[:,-self.text_len:], padding[:self.text_len]) text_states[:, -self.text_len :] = torch.where(
text_states_t5[:,-self.text_len_t5:] = torch.where(text_states_t5_mask[:,-self.text_len_t5:].unsqueeze(2), text_states_t5[:,-self.text_len_t5:], padding[self.text_len:]) text_states_mask[:, -self.text_len :].unsqueeze(2),
text_states[:, -self.text_len :],
padding[: self.text_len],
)
text_states_t5[:, -self.text_len_t5 :] = torch.where(
text_states_t5_mask[:, -self.text_len_t5 :].unsqueeze(2),
text_states_t5[:, -self.text_len_t5 :],
padding[self.text_len :],
)
text_states = torch.cat([text_states, text_states_t5], dim=1) # 2,205，1024 text_states = torch.cat([text_states, text_states_t5], dim=1) # 2,205，1024
_, _, oh, ow = x.shape # _, _, oh, ow = x.shape
th, tw = oh // self.patch_size, ow // self.patch_size # th, tw = oh // self.patch_size, ow // self.patch_size
# Get image RoPE embedding according to `reso`lution. # Get image RoPE embedding according to `reso`lution.
freqs_cis_img = calc_rope(x, self.patch_size, self.hidden_size // self.num_heads) # (cos_cis_img, sin_cis_img) freqs_cis_img = calc_rope(
x, self.patch_size, self.hidden_size // self.num_heads
) # (cos_cis_img, sin_cis_img)
# ========================= Build time and image embedding ========================= # ========================= Build time and image embedding =========================
t = self.t_embedder(t, dtype=self.dtype) t = self.t_embedder(t, dtype=self.dtype)
x = self.x_embedder(x) x = self.x_embedder(x)
# ========================= Concatenate all extra vectors ========================= # ========================= Concatenate all extra vectors =========================
# Build text tokens with pooling # Build text tokens with pooling
extra_vec = self.pooler(encoder_hidden_states_t5) extra_vec = self.pooler(encoder_hidden_states_t5)
@ -259,10 +331,9 @@ class HunYuanControlNet(nn.Module):
# ========================= Forward pass through HunYuanDiT blocks ========================= # ========================= Forward pass through HunYuanDiT blocks =========================
controls = [] controls = []
x = x + self.before_proj(condition) # add condition x = x + self.before_proj(condition) # add condition
for layer, block in enumerate(self.blocks): for layer, block in enumerate(self.blocks):
x = block(x, c, text_states, freqs_cis_img) x = block(x, c, text_states, freqs_cis_img)
controls.append(self.after_proj_list[layer](x)) # zero linear for output controls.append(self.after_proj_list[layer](x)) # zero linear for output
return {'output': controls} return {"output": controls}