Buckets:
| import json | |
| import torch | |
| from torch import nn | |
| import folder_paths | |
| import comfy.utils | |
| import comfy.ops | |
| import comfy.model_management | |
| import comfy.model_prefetch | |
| import comfy.patcher_extension | |
| import comfy.storage | |
| import comfy.ldm.common_dit | |
| import comfy.latent_formats | |
| import comfy.ldm.lumina.controlnet | |
| import comfy.ldm.supir.supir_modules | |
| import comfy.ldm.anima.lllite | |
| import comfy.ldm.minimax.controlnet | |
| import comfy.ldm.qwen_image21.model | |
| import comfy.ldm.wan.uni3c | |
| import comfy.ldm.lightricks.duration_head | |
| from comfy.ldm.wan.model_multitalk import WanMultiTalkAttentionBlock, MultiTalkAudioProjModel | |
| from comfy_api.latest import io | |
| from comfy.ldm.supir.supir_patch import SUPIRPatch | |
| class BlockWiseControlBlock(torch.nn.Module): | |
| # [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) | |
| return x | |
| class QwenImageBlockWiseControlNet(torch.nn.Module): | |
| 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): | |
| return self.controlnet_blocks[block_id](img, controlnet_conditioning) | |
| class SigLIPMultiFeatProjModel(torch.nn.Module): | |
| """ | |
| SigLIP Multi-Feature Projection Model for processing style features from different layers | |
| and projecting them into a unified hidden space. | |
| Args: | |
| siglip_token_nums (int): Number of SigLIP tokens, default 257 | |
| style_token_nums (int): Number of style tokens, default 256 | |
| siglip_token_dims (int): Dimension of SigLIP tokens, default 1536 | |
| hidden_size (int): Hidden layer size, default 3072 | |
| context_layer_norm (bool): Whether to use context layer normalization, default False | |
| """ | |
| 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): | |
| """ | |
| Forward pass function | |
| Args: | |
| siglip_outputs: Output from SigLIP model, containing hidden_states | |
| Returns: | |
| torch.Tensor: Concatenated multi-layer features with shape [bs, 3*style_token_nums, hidden_size] | |
| """ | |
| 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: | |
| """ | |
| Helper function to process features from a single layer | |
| Args: | |
| hidden_states: Input hidden states [bs, seq_len, dim] | |
| embedding_linear: Embedding linear layer | |
| layer_norm: Layer normalization | |
| projection: Projection layer | |
| dtype: Target data type | |
| Returns: | |
| torch.Tensor: Processed features [bs, style_token_nums, hidden_size] | |
| """ | |
| # 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) | |
| return embedding | |
| def z_image_convert(sd): | |
| 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 | |
| return out_sd | |
| def dit_patch_operations(sd): | |
| # 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]) | |
| return dtype, comfy.ops.pick_operations(dtype, manual_cast_dtype, load_device=load_device) | |
| class ModelPatchLoader: | |
| 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()) | |
| return (model_patcher,) | |
| class AnimaLLLiteApply: | |
| 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") | |
| return (model_patched,) | |
| def in_sigma_range(transformer_options, sigma_range): | |
| sigma = float(transformer_options["sigmas"].flatten()[0]) | |
| return sigma_range[1] <= sigma <= sigma_range[0] | |
| class DiffSynthCnetPatch: | |
| 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): | |
| return [self.model_patch] | |
| class ZImageControlPatch: | |
| 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): | |
| return [self.model_patch] | |
| class QwenImage21FunControlPatch: | |
| 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)) | |
| model.set_model_patch_replace(QwenImage21FunControlBlockPatch(self, block_index, previous), "dit", "single_block", block_index) | |
| class QwenImage21FunControlBlockPatch: | |
| 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() | |
| return models | |
| class QwenImageDiffsynthControlnet: | |
| 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)) | |
| return (model_patched,) | |
| class ZImageFunControlnet(QwenImageDiffsynthControlnet): | |
| 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"] | |
| DESCRIPTION = "Applies a Z-Image or Qwen Image 2.1 Fun ControlNet loaded with Load Model Patch." | |
| 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 | |
| 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: | |
| 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): | |
| if not isinstance(model_patch.model, comfy.ldm.wan.uni3c.WanUni3CControlnet): | |
| raise ValueError("The connected model patch is not a Uni3C ControlNet.") | |
| cnet_dim = model_patch.model.controlnet_blocks[0].norm1.linear.in_features | |
| model_dim = getattr(model.get_model_object("diffusion_model"), "dim", None) | |
| if model_dim is None: | |
| raise ValueError("The Uni3C ControlNet only works with Wan models.") | |
| if model_dim != cnet_dim: | |
| raise ValueError("This Uni3C ControlNet expects a Wan model with dim {}, the loaded model has dim {}.".format(cnet_dim, model_dim)) | |
| 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 | |
| self.encoded_image = encoded_image | |
| def __call__(self, kwargs): | |
| txt_ids = kwargs.get("txt_ids") | |
| txt = kwargs.get("txt") | |
| siglip_embedding = self.model_patch.model(self.encoded_image.to(txt.dtype)).to(txt.dtype) | |
| txt = torch.cat([siglip_embedding, txt], dim=1) | |
| kwargs['txt'] = txt | |
| kwargs['txt_ids'] = torch.cat([torch.zeros(siglip_embedding.shape[0], siglip_embedding.shape[1], 3, dtype=txt_ids.dtype, device=txt_ids.device), txt_ids], dim=1) | |
| 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): | |
| return [self.model_patch] | |
| class USOStyleReference: | |
| def INPUT_TYPES(s): | |
| return {"required": {"model": ("MODEL",), | |
| "model_patch": ("MODEL_PATCH",), | |
| "clip_vision_output": ("CLIP_VISION_OUTPUT", ), | |
| }} | |
| RETURN_TYPES = ("MODEL",) | |
| FUNCTION = "apply_patch" | |
| EXPERIMENTAL = True | |
| CATEGORY = "model/patch/flux" | |
| def apply_patch(self, model, model_patch, clip_vision_output): | |
| encoded_image = torch.stack((clip_vision_output.all_hidden_states[:, -20], clip_vision_output.all_hidden_states[:, -11], clip_vision_output.penultimate_hidden_states)) | |
| model_patched = model.clone() | |
| model_patched.set_model_post_input_patch(UsoStyleProjectorPatch(model_patch, encoded_image)) | |
| return (model_patched,) | |
| class MultiTalkModelPatch(torch.nn.Module): | |
| def __init__( | |
| self, | |
| audio_window: int = 5, | |
| intermediate_dim: int = 512, | |
| in_dim: int = 5120, | |
| out_dim: int = 768, | |
| context_tokens: int = 32, | |
| vae_scale: int = 4, | |
| num_layers: int = 40, | |
| device=None, dtype=None, operations=None | |
| ): | |
| super().__init__() | |
| self.audio_proj = MultiTalkAudioProjModel( | |
| seq_len=audio_window, | |
| seq_len_vf=audio_window+vae_scale-1, | |
| intermediate_dim=intermediate_dim, | |
| out_dim=out_dim, | |
| context_tokens=context_tokens, | |
| device=device, | |
| dtype=dtype, | |
| operations=operations | |
| ) | |
| self.blocks = torch.nn.ModuleList( | |
| [ | |
| WanMultiTalkAttentionBlock(in_dim, out_dim, device=device, dtype=dtype, operations=operations) | |
| for _ in range(num_layers) | |
| ] | |
| ) | |
| class SUPIRApply(io.ComfyNode): | |
| def define_schema(cls) -> io.Schema: | |
| return io.Schema( | |
| node_id="SUPIRApply", | |
| category="model/patch/supir", | |
| is_experimental=True, | |
| inputs=[ | |
| io.Model.Input("model"), | |
| io.ModelPatch.Input("model_patch"), | |
| io.Vae.Input("vae"), | |
| io.Image.Input("image"), | |
| io.Float.Input("strength_start", default=1.0, min=0.0, max=10.0, step=0.01, | |
| tooltip="Control strength at the start of sampling (high sigma)."), | |
| io.Float.Input("strength_end", default=1.0, min=0.0, max=10.0, step=0.01, | |
| tooltip="Control strength at the end of sampling (low sigma). Linearly interpolated from start."), | |
| io.Float.Input("restore_cfg", default=4.0, min=0.0, max=20.0, step=0.1, advanced=True, | |
| tooltip="Pulls denoised output toward the input latent. Higher = stronger fidelity to input. 0 to disable."), | |
| io.Float.Input("restore_cfg_s_tmin", default=0.05, min=0.0, max=1.0, step=0.01, advanced=True, | |
| tooltip="Sigma threshold below which restore_cfg is disabled."), | |
| ], | |
| outputs=[io.Model.Output()], | |
| ) | |
| def _encode_with_denoise_encoder(cls, vae, model_patch, image): | |
| """Encode using denoise_encoder weights from SUPIR checkpoint if available.""" | |
| denoise_sd = getattr(model_patch.model, 'denoise_encoder_sd', None) | |
| if not denoise_sd: | |
| return vae.encode(image) | |
| # Clone VAE patcher, apply denoise_encoder weights to clone, encode | |
| orig_patcher = vae.patcher | |
| vae.patcher = orig_patcher.clone() | |
| patches = {f"encoder.{k}": (v,) for k, v in denoise_sd.items()} | |
| vae.patcher.add_patches(patches, strength_patch=1.0, strength_model=0.0) | |
| try: | |
| return vae.encode(image) | |
| finally: | |
| vae.patcher = orig_patcher | |
| def execute(cls, *, model: io.Model.Type, model_patch: io.ModelPatch.Type, vae: io.Vae.Type, image: io.Image.Type, | |
| strength_start: float, strength_end: float, restore_cfg: float, restore_cfg_s_tmin: float) -> io.NodeOutput: | |
| model_patched = model.clone() | |
| hint_latent = model.get_model_object("latent_format").process_in( | |
| cls._encode_with_denoise_encoder(vae, model_patch, image[:, :, :, :3])) | |
| patch = SUPIRPatch(model_patch, model_patch.model.project_modules, hint_latent, strength_start, strength_end) | |
| patch.register(model_patched) | |
| if restore_cfg > 0.0: | |
| # Round-trip to match original pipeline: decode hint, re-encode with regular VAE | |
| latent_format = model.get_model_object("latent_format") | |
| decoded = vae.decode(latent_format.process_out(hint_latent)) | |
| x_center = latent_format.process_in(vae.encode(decoded[:, :, :, :3])) | |
| sigma_max = 14.6146 | |
| def restore_cfg_function(args): | |
| denoised = args["denoised"] | |
| sigma = args["sigma"] | |
| if sigma.dim() > 0: | |
| s = sigma[0].item() | |
| else: | |
| s = sigma.item() | |
| if s > restore_cfg_s_tmin: | |
| ref = x_center.to(device=denoised.device, dtype=denoised.dtype) | |
| b = denoised.shape[0] | |
| if ref.shape[0] != b: | |
| ref = ref.expand(b, -1, -1, -1) if ref.shape[0] == 1 else ref.repeat((b + ref.shape[0] - 1) // ref.shape[0], 1, 1, 1)[:b] | |
| sigma_val = sigma.view(-1, 1, 1, 1) if sigma.dim() > 0 else sigma | |
| d_center = denoised - ref | |
| denoised = denoised - d_center * ((sigma_val / sigma_max) ** restore_cfg) | |
| return denoised | |
| model_patched.set_model_sampler_post_cfg_function(restore_cfg_function) | |
| return io.NodeOutput(model_patched) | |
| NODE_CLASS_MAPPINGS = { | |
| "ModelPatchLoader": ModelPatchLoader, | |
| "QwenImageDiffsynthControlnet": QwenImageDiffsynthControlnet, | |
| "ZImageFunControlnet": ZImageFunControlnet, | |
| "WanUni3CControlnetApply": WanUni3CControlnetApply, | |
| "USOStyleReference": USOStyleReference, | |
| "SUPIRApply": SUPIRApply, | |
| "AnimaLLLiteApply": AnimaLLLiteApply, | |
| } | |
| NODE_DISPLAY_NAME_MAPPINGS = { | |
| "ModelPatchLoader": "Load Model Patch", | |
| "QwenImageDiffsynthControlnet": "Apply Qwen Image DiffSynth ControlNet", | |
| "ZImageFunControlnet": "Apply Fun ControlNet", | |
| "WanUni3CControlnetApply": "Apply Wan Uni3C ControlNet", | |
| "USOStyleReference": "Apply USO Style Reference", | |
| "SUPIRApply": "Apply SUPIR Patch", | |
| "AnimaLLLiteApply": "Apply Anima LLLite", | |
| } | |
Xet Storage Details
- Size:
- 55 kB
- Xet hash:
- 43f9e5cf6dbb260c3bcc19bd19419d7674def8fbb802e16300f98937885efaf8
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.