# [linear, gelu, linear]

def __init__(self, dim: int = 3072, device=None, dtype=None, operations=None):

super().__init__()

self.x_rms = operations.RMSNorm(dim, eps=1e-6)

self.y_rms = operations.RMSNorm(dim, eps=1e-6)

self.input_proj = operations.Linear(dim, dim)

self.act = torch.nn.GELU()

self.output_proj = operations.Linear(dim, dim)

def forward(self, x, y):

x, y = self.x_rms(x), self.y_rms(y)

x = self.input_proj(x + y)

x = self.act(x)

x = self.output_proj(x)

def __init__(

self,

num_layers: int = 60,

in_dim: int = 64,

additional_in_dim: int = 0,

dim: int = 3072,

device=None, dtype=None, operations=None

):

super().__init__()

self.additional_in_dim = additional_in_dim

self.img_in = operations.Linear(in_dim + additional_in_dim, dim, device=device, dtype=dtype)

self.controlnet_blocks = torch.nn.ModuleList(

[

BlockWiseControlBlock(dim, device=device, dtype=dtype, operations=operations)

for _ in range(num_layers)

]

)

def process_input_latent_image(self, latent_image):

latent_image[:, :16] = comfy.latent_formats.Wan21().process_in(latent_image[:, :16])

patch_size = 2

hidden_states = comfy.ldm.common_dit.pad_to_patch_size(latent_image, (1, patch_size, patch_size))

orig_shape = hidden_states.shape

hidden_states = hidden_states.view(orig_shape[0], orig_shape[1], orig_shape[-2] // 2, 2, orig_shape[-1] // 2, 2)

hidden_states = hidden_states.permute(0, 2, 4, 1, 3, 5)

hidden_states = hidden_states.reshape(orig_shape[0], (orig_shape[-2] // 2) * (orig_shape[-1] // 2), orig_shape[1] * 4)

return self.img_in(hidden_states)

def control_block(self, img, controlnet_conditioning, block_id):

"""

def __init__(

self,

siglip_token_nums: int = 729,

style_token_nums: int = 64,

siglip_token_dims: int = 1152,

hidden_size: int = 3072,

context_layer_norm: bool = True,

device=None, dtype=None, operations=None

):

super().__init__()

# High-level feature processing (layer -2)

self.high_embedding_linear = nn.Sequential(

operations.Linear(siglip_token_nums, style_token_nums),

nn.SiLU()

)

self.high_layer_norm = (

operations.LayerNorm(siglip_token_dims) if context_layer_norm else nn.Identity()

)

self.high_projection = operations.Linear(siglip_token_dims, hidden_size, bias=True)

# Mid-level feature processing (layer -11)

self.mid_embedding_linear = nn.Sequential(

operations.Linear(siglip_token_nums, style_token_nums),

nn.SiLU()

)

self.mid_layer_norm = (

operations.LayerNorm(siglip_token_dims) if context_layer_norm else nn.Identity()

)

self.mid_projection = operations.Linear(siglip_token_dims, hidden_size, bias=True)

# Low-level feature processing (layer -20)

self.low_embedding_linear = nn.Sequential(

operations.Linear(siglip_token_nums, style_token_nums),

nn.SiLU()

)

self.low_layer_norm = (

operations.LayerNorm(siglip_token_dims) if context_layer_norm else nn.Identity()

)

self.low_projection = operations.Linear(siglip_token_dims, hidden_size, bias=True)

def forward(self, siglip_outputs):

"""

dtype = next(self.high_embedding_linear.parameters()).dtype

# Process high-level features (layer -2)

high_embedding = self._process_layer_features(

siglip_outputs[2],

self.high_embedding_linear,

self.high_layer_norm,

self.high_projection,

dtype

)

# Process mid-level features (layer -11)

mid_embedding = self._process_layer_features(

siglip_outputs[1],

self.mid_embedding_linear,

self.mid_layer_norm,

self.mid_projection,

dtype

)

# Process low-level features (layer -20)

low_embedding = self._process_layer_features(

siglip_outputs[0],

self.low_embedding_linear,

self.low_layer_norm,

self.low_projection,

dtype

)

# Concatenate features from all layersmodel_patch

return torch.cat((high_embedding, mid_embedding, low_embedding), dim=1)

def _process_layer_features(

self,

hidden_states: torch.Tensor,

embedding_linear: nn.Module,

layer_norm: nn.Module,

projection: nn.Module,

dtype: torch.dtype

) -> torch.Tensor:

"""

# Transform dimensions: [bs, seq_len, dim] -> [bs, dim, seq_len] -> [bs, dim, style_token_nums] -> [bs, style_token_nums, dim]

embedding = embedding_linear(

hidden_states.to(dtype).transpose(1, 2)

).transpose(1, 2)

# Apply layer normalization

embedding = layer_norm(embedding)

# Project to target hidden space

embedding = projection(embedding)

replace_keys = {".attention.to_out.0.bias": ".attention.out.bias",

".attention.norm_k.weight": ".attention.k_norm.weight",

".attention.norm_q.weight": ".attention.q_norm.weight",

".attention.to_out.0.weight": ".attention.out.weight"

}

out_sd = {}

for k in sorted(sd.keys()):

w = sd[k]

k_out = k

if k_out.endswith(".attention.to_k.weight"):

cc = [w]

continue

if k_out.endswith(".attention.to_q.weight"):

cc = [w] + cc

continue

if k_out.endswith(".attention.to_v.weight"):

cc = cc + [w]

w = torch.cat(cc, dim=0)

k_out = k_out.replace(".attention.to_v.weight", ".attention.qkv.weight")

for r, rr in replace_keys.items():

k_out = k_out.replace(r, rr)

out_sd[k_out] = w

# quantized files keep their layers quantized with bf16 compute, others follow the usual bf16/fp32 dtype pick

quant = comfy.utils.detect_layer_quantization(sd, "")

if quant is not None:

return torch.bfloat16, comfy.ops.mixed_precision_ops(quant, torch.bfloat16)

load_device = comfy.model_management.get_torch_device()

dtype = comfy.model_management.unet_dtype(model_params=-1, supported_dtypes=[torch.bfloat16, torch.float32], weight_dtype=comfy.utils.weight_dtype(sd))

manual_cast_dtype = comfy.model_management.unet_manual_cast(dtype, load_device, supported_dtypes=[torch.bfloat16, torch.float32])

@classmethod

def INPUT_TYPES(s):

return {"required": { "name": (folder_paths.get_filename_list("model_patches"), ),

}}

RETURN_TYPES = ("MODEL_PATCH",)

FUNCTION = "load_model_patch"

EXPERIMENTAL = True

CATEGORY = "model/loaders"

def load_model_patch(self, name):

model_patch_path = folder_paths.get_full_path_or_raise("model_patches", name)

sd, metadata = comfy.utils.load_torch_file(model_patch_path, safe_load=True, return_metadata=True)

dtype = comfy.utils.weight_dtype(sd)

if 'lllite_conditioning1.conv1.weight' in sd:

model = comfy.ldm.anima.lllite.AnimaLLLite(sd, metadata, device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast)

elif 'controlnet_blocks.0.y_rms.weight' in sd:

additional_in_dim = sd["img_in.weight"].shape[1] - 64

model = QwenImageBlockWiseControlNet(additional_in_dim=additional_in_dim, device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast)

elif 'feature_embedder.mid_layer_norm.bias' in sd:

sd = comfy.utils.state_dict_prefix_replace(sd, {"feature_embedder.": ""}, filter_keys=True)

model = SigLIPMultiFeatProjModel(device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast)

elif 'control_all_x_embedder.2-1.weight' in sd: # alipai z image fun controlnet

sd = z_image_convert(sd)

config = {}

if 'control_layers.4.adaLN_modulation.0.weight' not in sd:

config['n_control_layers'] = 3

config['additional_in_dim'] = 17

config['refiner_control'] = True

if 'control_layers.14.adaLN_modulation.0.weight' in sd:

config['n_control_layers'] = 15

config['additional_in_dim'] = 17

config['refiner_control'] = True

ref_weight = sd.get("control_noise_refiner.0.after_proj.weight", None)

if ref_weight is not None:

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 'control_img_in.weight' in sd and 'control_blocks.0.img_mlp.out.weight' in sd: # Qwen Image 2.1 Fun ControlNet

dtype, operations = dit_patch_operations(sd)

inner_dim = sd["control_img_in.weight"].shape[0]

gate_up = sd.get("control_blocks.0.img_mlp.gate_up.weight", None)

hidden_dim = gate_up.shape[0] // 2 if gate_up is not None else sd["control_blocks.0.img_mlp.proj.weight"].shape[0]

num_blocks = 0

while "control_blocks.{}.after_proj.weight".format(num_blocks) in sd:

num_blocks += 1

model = comfy.ldm.qwen_image21.model.QwenImage21FunControl(

num_blocks=num_blocks,

control_in_dim=129,

inner_dim=inner_dim,

attention_head_dim=sd["control_blocks.0.attn.norm_q.weight"].shape[0],

mlp_ratio=hidden_dim // inner_dim,

fused_mlp=gate_up is not None,

operations=operations,

device=comfy.model_management.unet_offload_device(),

dtype=dtype,

)

elif comfy.ldm.minimax.controlnet.is_minimax_h3_fun_state_dict(sd):

dtype, operations = dit_patch_operations(sd)

num_blocks = 0

while "control_blocks.{}.after_proj.weight".format(num_blocks) in sd:

num_blocks += 1

# spread evenly over the 50 base blocks: v1 has 5 (every 10), v2 has 10 (every 5)

injection_layers = tuple(range(0, 50, 50 // num_blocks))

if metadata is not None and "control_blocks_places" in metadata:

injection_layers = tuple(json.loads(metadata["control_blocks_places"]))

if len(injection_layers) != num_blocks:

raise ValueError("MiniMax H3 Fun control_blocks_places metadata does not match the checkpoint")

qkv = sd["control_blocks.0.attn.qkv_proj.weight"]

head_dim = sd["control_blocks.0.attn.q_norm.weight"].shape[0]

use_adaln_curves = metadata is not None and metadata.get("minimax_h3_fun_controlnet") == "adaln_basis"

time_embed_dim = 8 if use_adaln_curves else 2688

model = comfy.ldm.minimax.controlnet.MiniMaxH3FunControl(

control_in_dim=49,

injection_layers=injection_layers,

inpaint_post_norm=metadata is not None and metadata.get("inpaint_masked_pixel_mode") == "post_norm",

hidden_size=sd["control_proj_in.weight"].shape[0],

num_attention_heads=qkv.shape[0] // (3 * head_dim),

attention_head_dim=head_dim,

ffn_hidden_size=sd["control_blocks.0.mlp.fc1.weight"].shape[0] // 2,

time_embed_dim=time_embed_dim,

use_adaln_curves=use_adaln_curves,

operations=operations,

device=comfy.model_management.unet_offload_device(),

dtype=dtype,

)

model.requires_grad_(False)

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=sd["controlnet_blocks.0.norm1.linear.weight"].shape[1],

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 any(k.endswith("duration_head.attention_pooler.query_tokens") for k in sd) or "attention_pooler.query_tokens" in sd:

sd = comfy.ldm.lightricks.duration_head.normalize_state_dict(sd)

sd = {k: v.float() for k, v in sd.items()} # tiny head, keep fp32

model = comfy.ldm.lightricks.duration_head.DurationHead()

elif "audio_proj.proj1.weight" in sd:

model = MultiTalkModelPatch(

audio_window=5, context_tokens=32, vae_scale=4,

in_dim=sd["blocks.0.audio_cross_attn.proj.weight"].shape[0],

intermediate_dim=sd["audio_proj.proj1.weight"].shape[0],

out_dim=sd["audio_proj.norm.weight"].shape[0],

device=comfy.model_management.unet_offload_device(),

operations=comfy.ops.manual_cast)

elif 'model.control_model.input_hint_block.0.weight' in sd or 'control_model.input_hint_block.0.weight' in sd:

prefix_replace = {}

if 'model.control_model.input_hint_block.0.weight' in sd:

prefix_replace["model.control_model."] = "control_model."

prefix_replace["model.diffusion_model.project_modules."] = "project_modules."

else:

prefix_replace["control_model."] = "control_model."

prefix_replace["project_modules."] = "project_modules."

# Extract denoise_encoder weights before filter_keys discards them

de_prefix = "first_stage_model.denoise_encoder."

denoise_encoder_sd = {}

for k in list(sd.keys()):

if k.startswith(de_prefix):

denoise_encoder_sd[k[len(de_prefix):]] = sd.pop(k)

sd = comfy.utils.state_dict_prefix_replace(sd, prefix_replace, filter_keys=True)

sd.pop("control_model.mask_LQ", None)

model = comfy.ldm.supir.supir_modules.SUPIR(device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast)

if denoise_encoder_sd:

model.denoise_encoder_sd = denoise_encoder_sd

model_patcher = comfy.model_patcher.CoreModelPatcher(model, load_device=comfy.model_management.get_torch_device(), offload_device=comfy.model_management.unet_offload_device(), fast_disk=comfy.storage.state_dict_fast_disk(sd))

model.load_state_dict(sd, assign=model_patcher.is_dynamic())

@classmethod

def INPUT_TYPES(s):

return {"required": {"model": ("MODEL",),

"model_patch": ("MODEL_PATCH",),

"image": ("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}),

},

"optional": {"mask": ("MASK",),

}}

RETURN_TYPES = ("MODEL",)

FUNCTION = "apply_patch"

EXPERIMENTAL = True

CATEGORY = "model/patch/anima"

def apply_patch(self, model, model_patch, image, strength, start_percent, end_percent, mask=None):

image = image[..., :3]

if model_patch.model.cond_in_channels == 4 and mask is None:

mask = torch.zeros_like(image[..., 0])

elif model_patch.model.cond_in_channels != 4:

mask = None

model_sampling = model.get_model_object("model_sampling")

sigma_start = float(model_sampling.percent_to_sigma(start_percent))

sigma_end = float(model_sampling.percent_to_sigma(end_percent))

patch = comfy.ldm.anima.lllite.AnimaLLLitePatch(model_patch, image, mask, strength, sigma_start, sigma_end)

model_patched = model.clone()

model_patched.set_model_post_input_patch(patch)

model_patched.set_model_attn1_patch(comfy.ldm.anima.lllite.AnimaLLLiteAttentionPatch(

patch,

{"q": "self_attn_q_proj", "k": "self_attn_k_proj", "v": "self_attn_v_proj"},

))

model_patched.set_model_attn2_patch(comfy.ldm.anima.lllite.AnimaLLLiteAttentionPatch(

patch,

{"q": "cross_attn_q_proj"},

))

model_patched.set_model_patch(comfy.ldm.anima.lllite.AnimaLLLiteMLPPatch(patch), "mlp_patch")

def __init__(self, model_patch, vae, image, strength, mask=None, sigma_range=(float("inf"), 0.0)):

self.model_patch = model_patch

self.vae = vae

self.image = image

self.strength = strength

self.mask = mask

self.sigma_range = sigma_range

self.encoded_image = model_patch.model.process_input_latent_image(self.encode_latent_cond(image))

self.encoded_image_size = (image.shape[1], image.shape[2])

def encode_latent_cond(self, image):

latent_image = self.vae.encode(image)

if self.model_patch.model.additional_in_dim > 0:

if self.mask is None:

mask_ = torch.ones_like(latent_image)[:, :self.model_patch.model.additional_in_dim // 4]

else:

mask_ = comfy.utils.common_upscale(self.mask.mean(dim=1, keepdim=True), latent_image.shape[-1], latent_image.shape[-2], "bilinear", "none")

return torch.cat([latent_image, mask_], dim=1)

else:

return latent_image

def __call__(self, kwargs):

if not in_sigma_range(kwargs["transformer_options"], self.sigma_range):

return kwargs

x = kwargs.get("x")

img = kwargs.get("img")

block_index = kwargs.get("block_index")

spacial_compression = self.vae.spacial_compression_encode()

if self.encoded_image is None or self.encoded_image_size != (x.shape[-2] * spacial_compression, x.shape[-1] * spacial_compression):

image_scaled = comfy.utils.common_upscale(self.image.movedim(-1, 1), x.shape[-1] * spacial_compression, x.shape[-2] * spacial_compression, "area", "center")

loaded_models = comfy.model_management.loaded_models(only_currently_used=True)

self.encoded_image = self.model_patch.model.process_input_latent_image(self.encode_latent_cond(image_scaled.movedim(1, -1)))

self.encoded_image_size = (image_scaled.shape[-2], image_scaled.shape[-1])

comfy.model_management.load_models_gpu(loaded_models)

img[:, :self.encoded_image.shape[1]] += (self.model_patch.model.control_block(img[:, :self.encoded_image.shape[1]], self.encoded_image.to(img.dtype), block_index) * self.strength)

kwargs['img'] = img

return kwargs

def to(self, device_or_dtype):

if isinstance(device_or_dtype, torch.device):

self.encoded_image = self.encoded_image.to(device_or_dtype)

return self

def models(self):

def __init__(self, model_patch, vae, image, strength, inpaint_image=None, mask=None, sigma_range=(float("inf"), 0.0)):

self.model_patch = model_patch

self.sigma_range = sigma_range

self.vae = vae

self.image = image

self.inpaint_image = inpaint_image

self.mask = mask

self.strength = strength

self.is_inpaint = self.model_patch.model.additional_in_dim > 0

skip_encoding = False

if self.image is not None and self.inpaint_image is not None:

if self.image.shape != self.inpaint_image.shape:

skip_encoding = True

if skip_encoding:

self.encoded_image = None

else:

self.encoded_image = self.encode_latent_cond(self.image, self.inpaint_image)

if self.image is None:

self.encoded_image_size = (self.inpaint_image.shape[1], self.inpaint_image.shape[2])

else:

self.encoded_image_size = (self.image.shape[1], self.image.shape[2])

self.temp_data = None

def encode_latent_cond(self, control_image=None, inpaint_image=None):

latent_image = None

if control_image is not None:

latent_image = comfy.latent_formats.Flux().process_in(self.vae.encode(control_image))

if self.is_inpaint:

if inpaint_image is None:

inpaint_image = torch.ones_like(control_image) * 0.5

if self.mask is not None:

mask_inpaint = comfy.utils.common_upscale(self.mask.view(self.mask.shape[0], -1, self.mask.shape[-2], self.mask.shape[-1]).mean(dim=1, keepdim=True), inpaint_image.shape[-2], inpaint_image.shape[-3], "bilinear", "center")

inpaint_image = ((inpaint_image - 0.5) * mask_inpaint.movedim(1, -1).round()) + 0.5

inpaint_image_latent = comfy.latent_formats.Flux().process_in(self.vae.encode(inpaint_image))

if self.mask is None:

mask_ = torch.zeros_like(inpaint_image_latent)[:, :1]

else:

mask_ = comfy.utils.common_upscale(self.mask.view(self.mask.shape[0], -1, self.mask.shape[-2], self.mask.shape[-1]).mean(dim=1, keepdim=True).to(device=inpaint_image_latent.device), inpaint_image_latent.shape[-1], inpaint_image_latent.shape[-2], "nearest", "center")

if latent_image is None:

latent_image = comfy.latent_formats.Flux().process_in(self.vae.encode(torch.ones_like(inpaint_image) * 0.5))

return torch.cat([latent_image, mask_, inpaint_image_latent], dim=1)

else:

return latent_image

def __call__(self, kwargs):

if not in_sigma_range(kwargs["transformer_options"], self.sigma_range):

return kwargs

x = kwargs.get("x")

img = kwargs.get("img")

img_input = kwargs.get("img_input")

txt = kwargs.get("txt")

pe = kwargs.get("pe")

vec = kwargs.get("vec")

block_index = kwargs.get("block_index")

block_type = kwargs.get("block_type", "")

spacial_compression = self.vae.spacial_compression_encode()

if self.encoded_image is None or self.encoded_image_size != (x.shape[-2] * spacial_compression, x.shape[-1] * spacial_compression):

image_scaled = None

if self.image is not None:

image_scaled = comfy.utils.common_upscale(self.image.movedim(-1, 1), x.shape[-1] * spacial_compression, x.shape[-2] * spacial_compression, "area", "center").movedim(1, -1)

self.encoded_image_size = (image_scaled.shape[-3], image_scaled.shape[-2])

inpaint_scaled = None

if self.inpaint_image is not None:

inpaint_scaled = comfy.utils.common_upscale(self.inpaint_image.movedim(-1, 1), x.shape[-1] * spacial_compression, x.shape[-2] * spacial_compression, "area", "center").movedim(1, -1)

self.encoded_image_size = (inpaint_scaled.shape[-3], inpaint_scaled.shape[-2])

loaded_models = comfy.model_management.loaded_models(only_currently_used=True)

self.encoded_image = self.encode_latent_cond(image_scaled, inpaint_scaled)

comfy.model_management.load_models_gpu(loaded_models)

cnet_blocks = self.model_patch.model.n_control_layers

div = round(30 / cnet_blocks)

cnet_index = (block_index // div)

cnet_index_float = (block_index / div)

kwargs.pop("img") # we do ops in place

kwargs.pop("txt")

if cnet_index_float > (cnet_blocks - 1):

self.temp_data = None

return kwargs

if self.temp_data is None or self.temp_data[0] > cnet_index:

if block_type == "noise_refiner":

self.temp_data = (-3, (None, self.model_patch.model(txt, self.encoded_image.to(img.dtype), pe, vec)))

else:

self.temp_data = (-1, (None, self.model_patch.model(txt, self.encoded_image.to(img.dtype), pe, vec)))

if block_type == "noise_refiner":

next_layer = self.temp_data[0] + 1

self.temp_data = (next_layer, self.model_patch.model.forward_noise_refiner_block(block_index, self.temp_data[1][1], img_input[:, :self.temp_data[1][1].shape[1]], None, pe, vec))

if self.temp_data[1][0] is not None:

img[:, :self.temp_data[1][0].shape[1]] += (self.temp_data[1][0] * self.strength)

else:

while self.temp_data[0] < cnet_index and (self.temp_data[0] + 1) < cnet_blocks:

next_layer = self.temp_data[0] + 1

self.temp_data = (next_layer, self.model_patch.model.forward_control_block(next_layer, self.temp_data[1][1], img_input[:, :self.temp_data[1][1].shape[1]], None, pe, vec))

if cnet_index_float == self.temp_data[0]:

img[:, :self.temp_data[1][0].shape[1]] += (self.temp_data[1][0] * self.strength)

if cnet_blocks == self.temp_data[0] + 1:

self.temp_data = None

return kwargs

def to(self, device_or_dtype):

if isinstance(device_or_dtype, torch.device):

if self.encoded_image is not None:

self.encoded_image = self.encoded_image.to(device_or_dtype)

self.temp_data = None

return self

def models(self):

def __init__(self, model_patch, vae, image, strength, inpaint_image=None, mask=None, sigma_range=(float("inf"), 0.0)):

self.model_patch = model_patch

self.vae = vae

self.image = image

self.inpaint_image = inpaint_image

self.mask = mask

self.strength = strength

self.sigma_range = sigma_range

self.active = False

self.injection_layers = None

self.control = None

self.stream = None

self.pristine = None

def prepare(self, h, w):

# 129 channels per target token: control latents | keep mask | masked-image latents, zeros where not given

if self.control is not None and self.control.shape[-2:] == (h, w):

return

width, height = w * self.vae.spacial_compression_encode(), h * self.vae.spacial_compression_encode()

latent_format = comfy.latent_formats.QwenImage21()

loaded_models = comfy.model_management.loaded_models(only_currently_used=True)

try:

control = torch.zeros(1, latent_format.latent_channels, h, w)

if self.image is not None:

image = comfy.utils.common_upscale(self.image[:1].movedim(-1, 1), width, height, "bicubic", "disabled")

control = latent_format.process_in(self.vae.encode(image.movedim(1, -1))).float().cpu()

regen = torch.ones(1, 1, height, width)

if self.mask is not None:

mask = self.mask.reshape(-1, 1, *self.mask.shape[-2:])[:1].float().cpu()

regen = (comfy.utils.common_upscale(mask, width, height, "bilinear", "disabled") >= 0.5).float()

inpaint = torch.zeros_like(control)

if self.inpaint_image is not None:

# regenerated pixels at mid-gray, the zero of the VAE's [-1, 1] input

image = comfy.utils.common_upscale(self.inpaint_image[:1].movedim(-1, 1), width, height, "bicubic", "disabled").float().cpu()

inpaint = latent_format.process_in(self.vae.encode((image * (1 - regen) + 0.5 * regen).movedim(1, -1))).float().cpu()

keep = 1 - torch.nn.functional.interpolate(regen, size=(h, w), mode="nearest")

finally:

comfy.model_management.load_models_gpu(loaded_models)

self.control = torch.cat([control, keep, inpaint], dim=1)

def diffusion_model_wrapper(self, executor, x, timestep, context, ref_latents, image_slots, transformer_options, **kwargs):

self.active = in_sigma_range(transformer_options, self.sigma_range)

if self.active:

with comfy.model_prefetch.pause_malloc_graph():

self.prepare(*x.shape[-2:])

else:

# outside the range the block patches are dropped so the model can use its prefix cache

dit = transformer_options.get("patches_replace", {}).get("dit", {})

dit = {k: p.previous if isinstance(p, QwenImage21FunControlBlockPatch) and p.control_patch is self else p for k, p in dit.items()}

dit = {k: p for k, p in dit.items() if p is not None}

transformer_options = {**transformer_options, "patches_replace": {**transformer_options["patches_replace"], "dit": dit}}

try:

return executor(x, timestep, context, ref_latents, image_slots, transformer_options, **kwargs)

finally:

self.stream = None

self.pristine = None

def before_block(self, block_index, args):

if self.active and block_index == self.injection_layers[0]:

# the base block updates its input in place

self.pristine = args["img"].clone()

def after_block(self, block_index, args, out):

if not self.active:

return out

model = self.model_patch.model

index = self.injection_layers.index(block_index)

if index == 0:

self.control = self.control.to(out["img"].device, out["img"].dtype)

self.stream = model.init_stream(self.pristine, self.control.flatten(2).transpose(1, 2), args["prefix_len"])

self.pristine = None

self.stream, skip = model.step(index, self.stream, args["mod"], args["pe"], args["attn_fn"], args["prefix_len"], args["transformer_options"])

out["img"].add_(skip, alpha=self.strength)

return out

def to(self, device_or_dtype):

if isinstance(device_or_dtype, torch.device):

if self.control is not None:

self.control = self.control.to(device_or_dtype)

self.stream = None

return self

def cleanup(self):

self.control = None

self.stream = None

self.pristine = None

self.active = False

def models(self):

return [self.model_patch]

def register(self, model):

# the control blocks pair with every (num_base / num_control)-th base block

num_base = len(model.get_model_object("diffusion_model").transformer_blocks)

self.injection_layers = list(range(0, num_base, num_base // len(self.model_patch.model.control_blocks)))

model.add_wrapper(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, self.diffusion_model_wrapper)

blocks_replace = model.model_options.get("transformer_options", {}).get("patches_replace", {}).get("dit", {})

for block_index in self.injection_layers:

previous = blocks_replace.get(("single_block", block_index))

def __init__(self, control_patch, block_index, previous):

self.control_patch = control_patch

self.block_index = block_index

self.previous = previous

def __call__(self, args, extra_args):

# control state stays outside the base block's allocation scope

with comfy.model_prefetch.pause_malloc_graph():

self.control_patch.before_block(self.block_index, args)

out = extra_args["original_block"](args) if self.previous is None else self.previous(args, extra_args)

with comfy.model_prefetch.pause_malloc_graph():

return self.control_patch.after_block(self.block_index, args, out)

def to(self, device_or_dtype):

self.control_patch.to(device_or_dtype)

if hasattr(self.previous, "to"):

self.previous = self.previous.to(device_or_dtype)

return self

def cleanup(self):

self.control_patch.cleanup()

if hasattr(self.previous, "cleanup"):

self.previous.cleanup()

def models(self):

models = self.control_patch.models()

if hasattr(self.previous, "models"):

models += self.previous.models()

@classmethod

def INPUT_TYPES(s):

return {"required": { "model": ("MODEL",),

"model_patch": ("MODEL_PATCH",),

"vae": ("VAE",),

"image": ("IMAGE",),

"strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),

},

"optional": {"mask": ("MASK",),

"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001, "advanced": True}),

"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001, "advanced": True})}}

RETURN_TYPES = ("MODEL",)

FUNCTION = "diffsynth_controlnet"

EXPERIMENTAL = True

CATEGORY = "model/patch/qwen"

def diffsynth_controlnet(self, model, model_patch, vae, image=None, strength=1.0, inpaint_image=None, mask=None, start_percent=0.0, end_percent=1.0):

if strength == 0 or (image is None and inpaint_image is None and mask is None):

return (model,)

model_patched = model.clone()

model_sampling = model.get_model_object("model_sampling")

sigma_range = (float(model_sampling.percent_to_sigma(start_percent)), float(model_sampling.percent_to_sigma(end_percent)))

if image is not None:

image = image[:, :, :, :3]

if inpaint_image is not None:

inpaint_image = inpaint_image[:, :, :, :3]

if isinstance(model_patch.model, comfy.ldm.qwen_image21.model.QwenImage21FunControl):

QwenImage21FunControlPatch(model_patch, vae, image, strength, inpaint_image=inpaint_image, mask=mask, sigma_range=sigma_range).register(model_patched)

return (model_patched,)

if mask is not None:

if mask.ndim == 3:

mask = mask.unsqueeze(1)

if mask.ndim == 4:

mask = mask.unsqueeze(2)

mask = 1.0 - mask

if isinstance(model_patch.model, comfy.ldm.lumina.controlnet.ZImage_Control):

patch = ZImageControlPatch(model_patch, vae, image, strength, inpaint_image=inpaint_image, mask=mask, sigma_range=sigma_range)

model_patched.set_model_noise_refiner_patch(patch)

model_patched.set_model_double_block_patch(patch)

else:

model_patched.set_model_double_block_patch(DiffSynthCnetPatch(model_patch, vae, image, strength, mask, sigma_range=sigma_range))

@classmethod

def INPUT_TYPES(s):

return {"required": { "model": ("MODEL",),

"model_patch": ("MODEL_PATCH",),

"vae": ("VAE",),

"strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),

},

"optional": {"image": ("IMAGE",), "inpaint_image": ("IMAGE",), "mask": ("MASK",),

"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001, "advanced": True}),

"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001, "advanced": True})}}

CATEGORY = "model/patch"

SEARCH_ALIASES = ["z-image controlnet", "qwen image 2.1 controlnet", "controlnet union", "fun controlnet"]

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