RomanNum3ral
/
ComfyUI
mirror of https://github.com/Comfy-Org/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge c5f077451a into 71b73e3b2b

This commit is contained in:
Barish Ozbay 2026-07-17 14:22:17 +09:00 committed by GitHub
parent 71b73e3b2b c5f077451a
commit 1065961b6a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 390 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

50
comfy/ldm/wan/model.py
View File

@ -552,6 +552,7 @@ class WanModel(torch.nn.Module):
List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
"""
# embeddings
x_input = x
x = self.patch_embedding(x.float()).to(x.dtype)
grid_sizes = x.shape[2:]
transformer_options["grid_sizes"] = grid_sizes
@ -564,11 +565,13 @@ class WanModel(torch.nn.Module):
e0 = self.time_projection(e).unflatten(2, (6, self.dim))
full_ref = None
img_offset = 0
if self.ref_conv is not None:
full_ref = kwargs.get("reference_latent", None)
if full_ref is not None:
full_ref = self.ref_conv(full_ref).flatten(2).transpose(1, 2)
x = torch.concat((full_ref, x), dim=1)
img_offset = full_ref.shape[1]
# In-context reference (Bernini)
context_latents = kwargs.get("context_latents", None)
@ -589,6 +592,7 @@ class WanModel(torch.nn.Module):
context_img_len = clip_fea.shape[-2]
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@ -604,6 +608,11 @@ class WanModel(torch.nn.Module):
else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": img_offset, "transformer_options": transformer_options})
x = out["img"]
# head
x = self.head(x, e)
@ -777,6 +786,7 @@ class VaceWanModel(WanModel):
**kwargs,
):
# embeddings
x_input = x
x = self.patch_embedding(x.float()).to(x.dtype)
grid_sizes = x.shape[2:]
transformer_options["grid_sizes"] = grid_sizes
@ -807,6 +817,7 @@ class VaceWanModel(WanModel):
x_orig = x
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@ -822,6 +833,11 @@ class VaceWanModel(WanModel):
else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": 0, "transformer_options": transformer_options})
x = out["img"]
ii = self.vace_layers_mapping.get(i, None)
if ii is not None:
for iii in range(len(c)):
@ -887,6 +903,7 @@ class CameraWanModel(WanModel):
**kwargs,
):
# embeddings
x_input = x
x = self.patch_embedding(x.float()).to(x.dtype)
if self.control_adapter is not None and camera_conditions is not None:
x = x + self.control_adapter(camera_conditions).to(x.dtype)
@ -909,6 +926,7 @@ class CameraWanModel(WanModel):
context_img_len = clip_fea.shape[-2]
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@ -924,6 +942,11 @@ class CameraWanModel(WanModel):
else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": 0, "transformer_options": transformer_options})
x = out["img"]
# head
x = self.head(x, e)
@ -1335,6 +1358,7 @@ class WanModel_S2V(WanModel):
# embeddings
bs, _, time, height, width = x.shape
x_input = x
x = self.patch_embedding(x.float()).to(x.dtype)
if control_video is not None:
x = x + self.cond_encoder(control_video)
@ -1379,6 +1403,7 @@ class WanModel_S2V(WanModel):
context = self.text_embedding(context)
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@ -1393,6 +1418,12 @@ class WanModel_S2V(WanModel):
x = out["img"]
else:
x = block(x, e=e0, freqs=freqs, context=context, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": 0, "transformer_options": transformer_options})
x = out["img"]
if audio_emb is not None:
x = self.audio_injector(x, i, audio_emb, audio_emb_global, seq_len)
# head
@ -1599,6 +1630,7 @@ class HumoWanModel(WanModel):
bs, _, time, height, width = x.shape
# embeddings
x_input = x
x = self.patch_embedding(x.float()).to(x.dtype)
grid_sizes = x.shape[2:]
x = x.flatten(2).transpose(1, 2)
@ -1630,6 +1662,7 @@ class HumoWanModel(WanModel):
audio = None
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@ -1645,6 +1678,11 @@ class HumoWanModel(WanModel):
else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, audio=audio, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": 0, "transformer_options": transformer_options})
x = out["img"]
# head
x = self.head(x, e)
@ -1660,8 +1698,14 @@ class SCAILWanModel(WanModel):
def forward_orig(self, x, t, context, clip_fea=None, freqs=None, transformer_options={}, pose_latents=None, reference_latent=None, ref_mask_latents=None, sam_latents=None, **kwargs):
x_input = x
img_offset = 0
if reference_latent is not None:
x = torch.cat((reference_latent, x), dim=2)
img_offset = (reference_latent.shape[2] // self.patch_size[0]) * \
(reference_latent.shape[3] // self.patch_size[1]) * \
(reference_latent.shape[4] // self.patch_size[2])
# embeddings
x = self.patch_embedding(x.float()).to(x.dtype)
@ -1697,6 +1741,7 @@ class SCAILWanModel(WanModel):
context_img_len = clip_fea.shape[-2]
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@ -1712,6 +1757,11 @@ class SCAILWanModel(WanModel):
else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": img_offset, "transformer_options": transformer_options})
x = out["img"]
# head
x = self.head(x, e)

9
comfy/ldm/wan/model_animate.py
View File

@ -493,6 +493,7 @@ class AnimateWanModel(WanModel):
**kwargs,
):
# embeddings
x_input = x
x = self.patch_embedding(x.float()).to(x.dtype)
x, motion_vec = self.after_patch_embedding(x, pose_latents, face_pixel_values)
grid_sizes = x.shape[2:]
@ -505,11 +506,13 @@ class AnimateWanModel(WanModel):
e0 = self.time_projection(e).unflatten(2, (6, self.dim))
full_ref = None
img_offset = 0
if self.ref_conv is not None:
full_ref = kwargs.get("reference_latent", None)
if full_ref is not None:
full_ref = self.ref_conv(full_ref).flatten(2).transpose(1, 2)
x = torch.concat((full_ref, x), dim=1)
img_offset = full_ref.shape[1]
# context
context = self.text_embedding(context)
@ -522,6 +525,7 @@ class AnimateWanModel(WanModel):
context_img_len = clip_fea.shape[-2]
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@ -537,6 +541,11 @@ class AnimateWanModel(WanModel):
else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": img_offset, "transformer_options": transformer_options})
x = out["img"]
if i % 5 == 0 and motion_vec is not None:
x = x + self.face_adapter.fuser_blocks[i // 5](x, motion_vec)

10
comfy/ldm/wan/model_wandancer.py
View File

@ -111,6 +111,7 @@ class WanDancerModel(WanModel):
def forward_orig(self, x, t, context, clip_fea=None, clip_fea_ref=None, freqs=None, audio_embed=None, fps=30, audio_inject_scale=1.0, transformer_options={}, **kwargs):
# embeddings
x_input = x
if int(fps + 0.5) != 30:
x = self.patch_embedding_global(x.float()).to(x.dtype)
else:
@ -128,11 +129,13 @@ class WanDancerModel(WanModel):
e0 = self.time_projection(e).unflatten(2, (6, self.dim))
full_ref = None
img_offset = 0
if self.ref_conv is not None: # model has the weight, but this wasn't used in the original pipeline
full_ref = kwargs.get("reference_latent", None)
if full_ref is not None:
full_ref = self.ref_conv(full_ref).flatten(2).transpose(1, 2)
x = torch.concat((full_ref, x), dim=1)
img_offset = full_ref.shape[1]
# context
context = self.text_embedding(context)
@ -163,6 +166,7 @@ class WanDancerModel(WanModel):
context_img_len += clip_fea_ref.shape[-2]
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@ -177,6 +181,12 @@ class WanDancerModel(WanModel):
x = out["img"]
else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": img_offset, "transformer_options": transformer_options})
x = out["img"]
if audio_emb is not None:
x = self.music_injector(x, i, audio_emb, audio_emb_global=None, seq_len=seq_len, scale=audio_inject_scale)

149
comfy/ldm/wan/uni3c.py Normal file
View File

@ -0,0 +1,149 @@
# Uni3C controlnet for Wan 2.1: https://github.com/ewrfcas/Uni3C
# Converted from the original diffusers based implementation.
import torch
import torch.nn as nn
from comfy.ldm.flux.layers import EmbedND
from .model import WanSelfAttention
class Uni3CLayerNormZero(nn.Module):
def __init__(
self,
conditioning_dim,
embedding_dim,
eps=1e-5,
device=None, dtype=None, operations=None
):
super().__init__()
self.silu = nn.SiLU()
self.linear = operations.Linear(conditioning_dim, 3 * embedding_dim, device=device, dtype=dtype)
self.norm = operations.LayerNorm(embedding_dim, eps=eps, elementwise_affine=True, device=device, dtype=dtype)
def forward(self, x, temb):
shift, scale, gate = self.linear(self.silu(temb)).chunk(3, dim=1)
x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
return x, gate[:, None, :]
class Uni3CAttentionBlock(nn.Module):
def __init__(
self,
dim,
ffn_dim,
num_heads,
time_embed_dim=5120,
eps=1e-6,
device=None, dtype=None, operations=None
):
super().__init__()
operation_settings = {"operations": operations, "device": device, "dtype": dtype}
self.norm1 = Uni3CLayerNormZero(time_embed_dim, dim, device=device, dtype=dtype, operations=operations)
self.self_attn = WanSelfAttention(dim, num_heads, qk_norm=True, eps=eps, operation_settings=operation_settings)
self.norm2 = Uni3CLayerNormZero(time_embed_dim, dim, device=device, dtype=dtype, operations=operations)
self.ffn = nn.Sequential(
operations.Linear(dim, ffn_dim, device=device, dtype=dtype), nn.GELU(approximate='tanh'),
operations.Linear(ffn_dim, dim, device=device, dtype=dtype))
def forward(self, x, temb, freqs):
norm_x, gate_msa = self.norm1(x, temb)
x = x + gate_msa * self.self_attn(norm_x, freqs)
norm_x, gate_ff = self.norm2(x, temb)
x = x + gate_ff * self.ffn(norm_x)
return x
class MaskCamEmbed(nn.Module):
def __init__(
self,
add_channels=7,
mid_channels=256,
conv_out_dim=5120,
device=None, dtype=None, operations=None
):
super().__init__()
self.mask_padding = [0, 0, 0, 0, 3, 0] # first frame conditioning
self.mask_proj = nn.Sequential(
operations.Conv3d(add_channels, mid_channels, kernel_size=(4, 8, 8), stride=(4, 8, 8), device=device, dtype=dtype),
operations.GroupNorm(mid_channels // 8, mid_channels, device=device, dtype=dtype),
nn.SiLU())
self.mask_zero_proj = operations.Conv3d(mid_channels, conv_out_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2), device=device, dtype=dtype)
def forward(self, add_inputs):
add_padded = torch.nn.functional.pad(add_inputs, self.mask_padding, mode="constant", value=0)
add_embeds = self.mask_proj(add_padded)
add_embeds = self.mask_zero_proj(add_embeds)
add_embeds = add_embeds.flatten(2).transpose(1, 2)
return add_embeds
class WanUni3CControlnet(nn.Module):
def __init__(
self,
in_channels=36,
conv_out_dim=5120,
dim=1024,
ffn_dim=8192,
num_heads=16,
num_layers=20,
time_embed_dim=5120,
out_proj_dim=5120,
add_channels=7,
mid_channels=256,
device=None, dtype=None, operations=None
):
super().__init__()
patch_size = (1, 2, 2)
self.num_layers = num_layers
self.controlnet_patch_embedding = operations.Conv3d(
in_channels, conv_out_dim, kernel_size=patch_size, stride=patch_size, device=device, dtype=torch.float32)
self.controlnet_mask_embedding = MaskCamEmbed(add_channels, mid_channels, conv_out_dim, device=device, dtype=dtype, operations=operations)
if conv_out_dim != dim:
self.proj_in = operations.Linear(conv_out_dim, dim, device=device, dtype=dtype)
else:
self.proj_in = nn.Identity()
self.controlnet_blocks = nn.ModuleList([
Uni3CAttentionBlock(dim, ffn_dim, num_heads, time_embed_dim, device=device, dtype=dtype, operations=operations)
for _ in range(num_layers)])
self.proj_out = nn.ModuleList([
operations.Linear(dim, out_proj_dim, device=device, dtype=dtype)
for _ in range(num_layers)])
head_dim = dim // num_heads
self.rope_embedder = EmbedND(dim=head_dim, theta=10000.0, axes_dim=[head_dim - 4 * (head_dim // 6), 2 * (head_dim // 6), 2 * (head_dim // 6)])
def rope_encode(self, t_len, h_len, w_len, device=None, dtype=None):
img_ids = torch.zeros((t_len, h_len, w_len, 3), device=device, dtype=dtype)
img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + torch.arange(t_len, device=device, dtype=dtype).reshape(-1, 1, 1)
img_ids[:, :, :, 1] = img_ids[:, :, :, 1] + torch.arange(h_len, device=device, dtype=dtype).reshape(1, -1, 1)
img_ids[:, :, :, 2] = img_ids[:, :, :, 2] + torch.arange(w_len, device=device, dtype=dtype).reshape(1, 1, -1)
img_ids = img_ids.reshape(1, -1, img_ids.shape[-1])
freqs = self.rope_embedder(img_ids).movedim(1, 2)
return freqs
def process_input(self, control_input, render_mask=None, camera_embedding=None):
# render_mask/camera_embedding are the checkpoint's extra conditioning path, not wired up yet
hidden = self.controlnet_patch_embedding(control_input.float()).to(control_input.dtype)
t_len, h_len, w_len = hidden.shape[2:]
freqs = self.rope_encode(t_len, h_len, w_len, device=hidden.device, dtype=hidden.dtype)
hidden = hidden.flatten(2).transpose(1, 2)
add_inputs = None
if camera_embedding is not None and render_mask is not None:
add_inputs = torch.cat([render_mask, camera_embedding], dim=1)
elif render_mask is not None:
add_inputs = render_mask
if add_inputs is not None:
hidden = hidden + self.controlnet_mask_embedding(add_inputs.to(hidden.dtype))
hidden = self.proj_in(hidden)
return hidden, freqs
def forward_block(self, block_index, hidden, temb, freqs):
hidden = self.controlnet_blocks[block_index](hidden, temb, freqs)
residual = self.proj_out[block_index](hidden)
return hidden, residual

172
comfy_extras/nodes_model_patch.py
View File

@ -8,6 +8,7 @@ import comfy.ldm.common_dit
import comfy.latent_formats
import comfy.ldm.lumina.controlnet
import comfy.ldm.supir.supir_modules
import comfy.ldm.wan.uni3c
from comfy.ldm.wan.model_multitalk import WanMultiTalkAttentionBlock, MultiTalkAudioProjModel
from comfy_api.latest import io
from comfy.ldm.supir.supir_patch import SUPIRPatch
@ -261,6 +262,37 @@ class ModelPatchLoader:
if torch.count_nonzero(ref_weight) == 0:
config['broken'] = True
model = comfy.ldm.lumina.controlnet.ZImage_Control(device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast, **config)
elif 'controlnet_patch_embedding.weight' in sd: # Uni3C controlnet for Wan
attn_key_replace = {".self_attn.to_q.": ".self_attn.q.",
".self_attn.to_k.": ".self_attn.k.",
".self_attn.to_v.": ".self_attn.v.",
".self_attn.to_out.0.": ".self_attn.o."}
converted_sd = {}
for k, w in sd.items():
for r, rr in attn_key_replace.items():
k = k.replace(r, rr)
converted_sd[k] = w
sd = converted_sd
num_layers = sum(1 for k in sd if k.startswith("proj_out.") and k.endswith(".weight"))
conv_out_dim = sd["controlnet_patch_embedding.weight"].shape[0]
if "proj_in.weight" in sd:
dim = sd["proj_in.weight"].shape[0]
else:
dim = conv_out_dim
model = comfy.ldm.wan.uni3c.WanUni3CControlnet(
in_channels=sd["controlnet_patch_embedding.weight"].shape[1],
conv_out_dim=conv_out_dim,
dim=dim,
ffn_dim=sd["controlnet_blocks.0.ffn.0.bias"].shape[0],
num_layers=num_layers,
time_embed_dim=conv_out_dim,
out_proj_dim=sd["proj_out.0.weight"].shape[0],
add_channels=sd["controlnet_mask_embedding.mask_proj.0.weight"].shape[1],
mid_channels=sd["controlnet_mask_embedding.mask_proj.0.weight"].shape[0],
device=comfy.model_management.unet_offload_device(),
dtype=dtype,
operations=comfy.ops.manual_cast)
elif "audio_proj.proj1.weight" in sd:
model = MultiTalkModelPatch(
audio_window=5, context_tokens=32, vae_scale=4,
@ -514,6 +546,144 @@ class ZImageFunControlnet(QwenImageDiffsynthControlnet):
CATEGORY = "model/patch/z-image"
class WanUni3CCnetPatch:
def __init__(self, model_patch, render_video, vae, latent_format, strength, sigma_start, sigma_end):
self.model_patch = model_patch
self.render_video = render_video
self.vae = vae
self.latent_format = latent_format
self.strength = strength
self.sigma_start = sigma_start
self.sigma_end = sigma_end
self.prepared_render = None
self.temp_data = None
def encode_render_video(self, target_latent_shape):
t_len, h_len, w_len = target_latent_shape
temporal_compression = self.vae.temporal_compression_decode() or 1
spatial_compression = self.vae.spacial_compression_encode()
target_frames = (t_len - 1) * temporal_compression + 1
target_height = h_len * spatial_compression
target_width = w_len * spatial_compression
frames = self.render_video
if frames.shape[0] > target_frames:
frames = frames[:target_frames]
elif frames.shape[0] < target_frames:
last_frame = frames[-1:].expand(target_frames - frames.shape[0], -1, -1, -1)
frames = torch.cat([frames, last_frame], dim=0)
if frames.shape[1] != target_height or frames.shape[2] != target_width:
frames = comfy.utils.common_upscale(frames.movedim(-1, 1), target_width, target_height, "bilinear", "center").movedim(1, -1)
loaded_models = comfy.model_management.loaded_models(only_currently_used=True)
render_latent = self.vae.encode(frames)
comfy.model_management.load_models_gpu(loaded_models)
return self.latent_format.process_in(render_latent)
def build_controlnet_input(self, x, dtype, samples_per_cond):
# first 20 channels of the model input: noise latent + I2V mask (zero padded for T2V)
hidden = x[:samples_per_cond, :20].to(dtype)
if hidden.shape[1] < 20:
pad_shape = list(hidden.shape)
pad_shape[1] = 20 - hidden.shape[1]
hidden = torch.cat([hidden, torch.zeros(pad_shape, dtype=hidden.dtype, device=hidden.device)], dim=1)
render = self.prepared_render
if render is None or render.shape[2:] != hidden.shape[2:]:
render = self.encode_render_video(hidden.shape[2:])
render = render.to(device=hidden.device, dtype=dtype)
self.prepared_render = render
if render.shape[0] != hidden.shape[0]:
render = render.expand(hidden.shape[0], -1, -1, -1, -1)
return torch.cat([hidden, render], dim=1)
def __call__(self, kwargs):
img = kwargs.get("img")
block_index = kwargs.get("block_index")
transformer_options = kwargs.get("transformer_options", {})
if block_index == 0:
self.temp_data = None
active = True
sigmas = transformer_options.get("sigmas", None)
if sigmas is not None:
sigma = sigmas[0].item()
if sigma > self.sigma_start or sigma < self.sigma_end:
active = False
if active:
x = kwargs.get("x")
# cond and uncond chunks share latents, so we can reuse residuals
num_conds = len(transformer_options.get("cond_or_uncond", [0]))
samples_per_cond = x.shape[0]
if num_conds > 0 and x.shape[0] % num_conds == 0:
samples_per_cond = x.shape[0] // num_conds
temb = kwargs.get("vec")[:samples_per_cond]
if temb.ndim == 3:
temb = temb[:, 0]
model = self.model_patch.model
cnet_dim = model.controlnet_blocks[0].norm1.linear.in_features
if temb.shape[-1] != cnet_dim:
raise RuntimeError("This Uni3C controlnet expects a Wan model with dim {}, the loaded model has dim {}".format(cnet_dim, temb.shape[-1]))
controlnet_input = self.build_controlnet_input(x, img.dtype, samples_per_cond)
hidden, freqs = model.process_input(controlnet_input)
self.temp_data = (hidden, temb.to(img.dtype), freqs)
num_layers = self.model_patch.model.num_layers
if self.temp_data is not None and block_index < num_layers:
hidden, temb, freqs = self.temp_data
hidden, residual = self.model_patch.model.forward_block(block_index, hidden, temb, freqs)
residual = residual.to(img.dtype) * self.strength
if residual.shape[0] != img.shape[0]:
residual = residual.repeat(img.shape[0] // residual.shape[0], 1, 1)
img_offset = kwargs.get("img_offset", 0)
img[:, img_offset:img_offset + residual.shape[1]] += residual
if block_index >= num_layers - 1:
self.temp_data = None
else:
self.temp_data = (hidden, temb, freqs)
return kwargs
def to(self, device_or_dtype):
if isinstance(device_or_dtype, torch.device):
if self.prepared_render is not None:
self.prepared_render = self.prepared_render.to(device_or_dtype)
self.temp_data = None
return self
def models(self):
return [self.model_patch]
class WanUni3CControlnetApply:
@classmethod
def INPUT_TYPES(s):
return {"required": { "model": ("MODEL",),
"model_patch": ("MODEL_PATCH",),
"vae": ("VAE",),
"render_video": ("IMAGE", {"tooltip": "The guidance video rendered from the camera trajectory, most commonly warped point cloud renders of the input image."}),
"strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
}}
RETURN_TYPES = ("MODEL",)
FUNCTION = "apply_patch"
EXPERIMENTAL = True
CATEGORY = "model/patch/wan"
def apply_patch(self, model, model_patch, vae, render_video, strength, start_percent, end_percent):
model_patched = model.clone()
model_sampling = model.get_model_object("model_sampling")
sigma_start = model_sampling.percent_to_sigma(start_percent)
sigma_end = model_sampling.percent_to_sigma(end_percent)
latent_format = model.get_model_object("latent_format")
patch = WanUni3CCnetPatch(model_patch, render_video[:, :, :, :3], vae, latent_format, strength, sigma_start, sigma_end)
model_patched.set_model_double_block_patch(patch)
return (model_patched,)
class UsoStyleProjectorPatch:
def __init__(self, model_patch, encoded_image):
self.model_patch = model_patch
@ -672,6 +842,7 @@ NODE_CLASS_MAPPINGS = {
"ModelPatchLoader": ModelPatchLoader,
"QwenImageDiffsynthControlnet": QwenImageDiffsynthControlnet,
"ZImageFunControlnet": ZImageFunControlnet,
"WanUni3CControlnetApply": WanUni3CControlnetApply,
"USOStyleReference": USOStyleReference,
"SUPIRApply": SUPIRApply,
}
@ -680,6 +851,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"ModelPatchLoader": "Load Model Patch",
"QwenImageDiffsynthControlnet": "Apply Qwen Image DiffSynth ControlNet",
"ZImageFunControlnet": "Apply Z-Image Fun ControlNet",
"WanUni3CControlnetApply": "Apply Wan Uni3C ControlNet",
"USOStyleReference": "Apply USO Style Reference",
"SUPIRApply": "Apply SUPIR Patch",
}