commit de736c5b712ed2ddfe7e4d93ec307a4e26767e55 Author: ozbayb <17261091+ozbayb@users.noreply.github.com> Date: Fri Aug 8 10:07:15 2025 -0600 Initial commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..bd13e80 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +.DS_Store +*pyc +.vscode +__pycache__ +*.egg-info +*.bak +checkpoints +results +backup \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..2a1e8ba --- /dev/null +++ b/__init__.py @@ -0,0 +1,55 @@ +import importlib.util +import sys +from pathlib import Path +import traceback + +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + + +def load_module(module_name, file_path, package_name=None): + """Loads a module from a .py file with proper package context.""" + spec = importlib.util.spec_from_file_location(module_name, file_path, submodule_search_locations=[]) + if package_name: + spec.submodule_search_locations = [str(Path(file_path).parent.parent)] + module = importlib.util.module_from_spec(spec) + + # Set up the module in sys.modules before execution to handle relative imports + sys.modules[module_name] = module + if package_name: + module.__package__ = package_name + + try: + spec.loader.exec_module(module) + return module + except Exception: + # Clean up on failure + if module_name in sys.modules: + del sys.modules[module_name] + raise + +def load_nodes(): + """Loads nodes from the nodes folder.""" + # Get the package name for this custom node + package_name = __name__ + + # Ensure the package structure is in sys.modules + if package_name not in sys.modules: + sys.modules[package_name] = sys.modules[__name__] + + nodes_path = Path(__file__).parent / 'nodes' + for file in nodes_path.glob("*.py"): + if file.name != "__init__.py": + try: + module_name = f"{package_name}.nodes.{file.stem}" + module = load_module(module_name, str(file), f"{package_name}.nodes") + NODE_CLASS_MAPPINGS.update(getattr(module, 'NODE_CLASS_MAPPINGS', {})) + NODE_DISPLAY_NAME_MAPPINGS.update(getattr(module, 'NODE_DISPLAY_NAME_MAPPINGS', {})) + except Exception as e: + print(f"Error loading {file.name}: {e}\n{traceback.format_exc()}") + +# Execute functions +load_nodes() + +WEB_DIRECTORY = "./js" +__all__ = ["NODE_CLASS_MAPPINGS", 'NODE_DISPLAY_NAME_MAPPINGS', "WEB_DIRECTORY"] \ No newline at end of file diff --git a/nodes/vace_custom_nodes.py b/nodes/vace_custom_nodes.py new file mode 100644 index 0000000..1ec632d --- /dev/null +++ b/nodes/vace_custom_nodes.py @@ -0,0 +1,587 @@ +import nodes +import node_helpers +import torch +import comfy.model_management +import comfy.utils +import comfy.latent_formats + + +class WanVacePhantomExperimental: + def __init__(self) -> None: + pass + + @classmethod + def INPUT_TYPES(s): + return {"required": { + "positive": ("CONDITIONING", ), + "negative": ("CONDITIONING", ), + "vae": ("VAE", ), + "width": ("INT", {"default": 832, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}), + "height": ("INT", {"default": 480, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}), + "length": ("INT", {"default": 81, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 4}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), + }, + "optional": { + "control_video": ("IMAGE", ), + "control_video2": ("IMAGE", {"tooltip": "Second control video input"}), + "control_masks": ("MASK", ), + "control_masks2": ("MASK", {"tooltip": "Second control masks input"}), + "vace_reference": ("IMAGE", ), + "vace_reference_2": ("IMAGE", ), + "vace_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1000.0, "step": 0.01}), + "vace_strength2": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1000.0, "step": 0.01}), + "vace_ref_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1000.0, "step": 0.01}), + "vace_ref2_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1000.0, "step": 0.01}), + # Phantom inputs + "phantom_images": ("IMAGE", ), + "phantom_mask_value": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Vace mask value for the Phantom embed region."}), + "phantom_control_value": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Padded vace embedded latents value for the Phantom embed region."}), + }} + + RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "CONDITIONING", "LATENT", "INT") + RETURN_NAMES = ("positive", "negative", "neg_phant_img", "latent", "trim_latent") + FUNCTION = "encode" + + CATEGORY = "WanVaceAdvanced" + + EXPERIMENTAL = True + + def encode(self, positive, negative, vae, width, height, length, batch_size, + vace_strength=1.0, vace_strength2=1.0, vace_ref_strength=None, vace_ref2_strength=None, + control_video=None, control_masks=None, vace_reference=None, + control_video2=None, control_masks2=None, vace_reference_2=None, + phantom_images=None, phantom_mask_value=0.0, phantom_control_value=0.0, + ): + + if control_video is not None: + control_video = control_video.clone() + if control_masks is not None: + control_masks = control_masks.clone() + if vace_reference is not None: + vace_reference = vace_reference.clone() + if phantom_images is not None: + phantom_images = phantom_images.clone() + if control_video2 is not None: + control_video2 = control_video2.clone() + if control_masks2 is not None: + control_masks2 = control_masks2.clone() + if vace_reference_2 is not None: + vace_reference_2 = vace_reference_2.clone() + + if vace_ref_strength is None: + if isinstance(vace_strength, list): + vace_ref_strength = vace_strength[0] + else: + vace_ref_strength = vace_strength + if vace_ref2_strength is None: + if isinstance(vace_strength2, list): + vace_ref2_strength = vace_strength2[0] + else: + vace_ref2_strength = vace_strength2 + + num_phantom_images = 0 if phantom_images is None else len(phantom_images) + # Calculate the additional length needed for Phantom images in Vace embeds + vace_length = length + num_phantom_images * 4 + + ######################## + # execute WanVaceToVideo logic + + vace_latent_length = ((vace_length - 1) // 4) + 1 + original_vace_latent_length = vace_latent_length # Save original before it gets modified + control_video_original_length = 0 if control_video is None else control_video.shape[0] + if control_video is not None: + control_video = comfy.utils.common_upscale(control_video[:control_video_original_length].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1) + if control_video.shape[0] < vace_length: + control_video = torch.nn.functional.pad(control_video, (0, 0, 0, 0, 0, 0, 0, vace_length - control_video.shape[0]), value=0.5) + else: + control_video = torch.ones((vace_length, height, width, 3)) * 0.5 + + if isinstance(vace_strength, list): + if len(vace_strength) != vace_latent_length: + raise ValueError(f"If vace_strength is a list, it must have length {vace_latent_length}, got {len(vace_strength)}") + + vace_references_encoded = None + if vace_reference is not None: + vace_references_scaled = comfy.utils.common_upscale(vace_reference[:].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1) + vace_references_encoded = vae.encode(vace_references_scaled[:, :, :, :3]) + vace_references_encoded = torch.cat([vace_references_encoded, comfy.latent_formats.Wan21().process_out(torch.zeros_like(vace_references_encoded))], dim=1) + if not isinstance(vace_strength, list): + if vace_ref_strength != vace_strength: + vace_strength_list = [[vace_ref_strength, vace_strength]] * batch_size + else: + # If vace_strength is a single value and equal to vace_ref_strength, we need to expand it to a list + inner_strength_list = [vace_strength] * (vace_latent_length + vace_references_encoded.shape[2]) + vace_strength_list = [inner_strength_list] * batch_size + else: + # If vace_strength is already a list, we assume it has the correct length + vace_strength_list = [[vace_ref_strength] + vace_strength] * batch_size + else: + # If no references are provided, expand the vace_strength to match the vace_latent_length + if isinstance(vace_strength, list): + # If vace_strength is already a list, use it directly (should already be correct length) + vace_strength_list = [vace_strength] * batch_size + else: + # If vace_strength is a single value, expand it + vace_strength_list = [[vace_strength] * vace_latent_length] * batch_size + + def format_strength_list(s_list): + if not s_list: + return "[]" + if len(s_list) == 0: + return "[]" + + # Format all numbers to two decimal places + formatted_list = [f"{x:.2f}" for x in s_list] + + # If all elements are the same and there are more than 2, summarize + if len(formatted_list) > 2 and len(set(formatted_list)) == 1: + return f"[{formatted_list[0]} ... {formatted_list[-1]}]" + else: + return f"[{', '.join(formatted_list)}]" + + # Format the inner list for printing + formatted_inner_list = format_strength_list(vace_strength_list[0]) + print(f"Using vace_strength as list: {formatted_inner_list} (length {len(vace_strength_list[0])})") + + # Print the reference, control, and phantom strength parts + vace_reference_length = vace_references_encoded.shape[2] if vace_references_encoded is not None else 0 + reference_strength_part = vace_strength_list[0][:vace_reference_length] if vace_references_encoded is not None else [] + control_strength_part = vace_strength_list[0][vace_reference_length:-(num_phantom_images)] if num_phantom_images > 0 else vace_strength_list[0][vace_reference_length:] + phantom_strength_part = vace_strength_list[0][-(num_phantom_images):] if num_phantom_images > 0 else [] + print(f"Reference strength: {format_strength_list(reference_strength_part)} (length {len(reference_strength_part)})") + print(f"Control strength: {format_strength_list(control_strength_part)} (length {len(control_strength_part)})") + print(f"Phantom images strength: {format_strength_list(phantom_strength_part)} (length {len(phantom_strength_part)})") + + control_masks_original_length = 0 if control_masks is None else control_masks.shape[0] + if control_masks is not None: + mask = control_masks + if mask.ndim == 3: + mask = mask.unsqueeze(1) + mask = comfy.utils.common_upscale(mask[:control_masks_original_length], width, height, "bilinear", "center").movedim(1, -1) + if mask.shape[0] < vace_length: + mask = torch.nn.functional.pad(mask, (0, 0, 0, 0, 0, 0, 0, vace_length - mask.shape[0]), value=1.0) + else: + mask = torch.ones((vace_length, height, width, 1)) + + # Modify the phantom-overlapping portions if phantom images exist + if num_phantom_images > 0: + # modify control video values for phantom padding region + control_video[control_video_original_length:, :, :, :] = phantom_control_value + # modify mask values for phantom padding region + mask[control_masks_original_length:, :, :, :] = phantom_mask_value + + control_video = control_video - 0.5 + inactive = (control_video * (1 - mask)) + 0.5 + reactive = (control_video * mask) + 0.5 + + inactive = vae.encode(inactive[:, :, :, :3]) + reactive = vae.encode(reactive[:, :, :, :3]) + control_video_latent = torch.cat((inactive, reactive), dim=1) + if vace_references_encoded is not None: + control_video_latent = torch.cat((vace_references_encoded, control_video_latent), dim=2) + + vae_stride = 8 + height_mask = height // vae_stride + width_mask = width // vae_stride + mask = mask.view(vace_length, height_mask, vae_stride, width_mask, vae_stride) + mask = mask.permute(2, 4, 0, 1, 3) + mask = mask.reshape(vae_stride * vae_stride, vace_length, height_mask, width_mask) + mask = torch.nn.functional.interpolate(mask.unsqueeze(0), size=(vace_latent_length, height_mask, width_mask), mode='nearest-exact').squeeze(0) + + trim_latent = 0 + if vace_references_encoded is not None: + mask_pad = torch.zeros_like(mask[:, :vace_references_encoded.shape[2], :, :]) + mask = torch.cat((mask_pad, mask), dim=1) + vace_latent_length += vace_references_encoded.shape[2] + trim_latent = vace_references_encoded.shape[2] + + mask = mask.unsqueeze(0) + + # Prepare lists for conditioning (first VACE is already processed) + vace_frames_list = [control_video_latent] + vace_mask_list = [mask] + vace_strength_list_all = [vace_strength_list] + + # Process second VACE operation if provided + if control_video2 is not None or control_masks2 is not None: + print(f"Processing second VACE with strength {vace_strength2}") + + # Process second control video + control_video2_original_length = 0 if control_video2 is None else control_video2.shape[0] + if control_video2 is not None: + control_video2 = comfy.utils.common_upscale(control_video2[:control_video2_original_length].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1) + if control_video2.shape[0] < vace_length: + control_video2 = torch.nn.functional.pad(control_video2, (0, 0, 0, 0, 0, 0, 0, vace_length - control_video2.shape[0]), value=0.5) + else: + control_video2 = torch.ones((vace_length, height, width, 3)) * 0.5 + + # Process second control masks + control_masks2_original_length = 0 if control_masks2 is None else control_masks2.shape[0] + if control_masks2 is not None: + mask2 = control_masks2 + if mask2.ndim == 3: + mask2 = mask2.unsqueeze(1) + mask2 = comfy.utils.common_upscale(mask2[:control_masks2_original_length], width, height, "bilinear", "center").movedim(1, -1) + if mask2.shape[0] < vace_length: + mask2 = torch.nn.functional.pad(mask2, (0, 0, 0, 0, 0, 0, 0, vace_length - mask2.shape[0]), value=1.0) + else: + mask2 = torch.ones((vace_length, height, width, 1)) + + # Modify phantom regions for second VACE + if num_phantom_images > 0: + control_video2[control_video2_original_length:, :, :, :] = phantom_control_value + mask2[control_masks2_original_length:, :, :, :] = phantom_mask_value + + # Process second control video same as first + control_video2 = control_video2 - 0.5 + inactive2 = (control_video2 * (1 - mask2)) + 0.5 + reactive2 = (control_video2 * mask2) + 0.5 + + inactive2 = vae.encode(inactive2[:, :, :, :3]) + reactive2 = vae.encode(reactive2[:, :, :, :3]) + control_video_latent2 = torch.cat((inactive2, reactive2), dim=1) + + # Process references for second VACE (use vace_reference_2 if provided, otherwise reuse first references) + vace_references_encoded2 = None + if vace_reference_2 is not None: + vace_reference_2_scaled = comfy.utils.common_upscale(vace_reference_2[:].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1) + vace_references_encoded2 = vae.encode(vace_reference_2_scaled[:, :, :, :3]) + vace_references_encoded2 = torch.cat([vace_references_encoded2, comfy.latent_formats.Wan21().process_out(torch.zeros_like(vace_references_encoded2))], dim=1) + + # Validate that both VACE operations have the same number of reference frames for proper stacking + if vace_references_encoded is not None and vace_references_encoded2.shape[2] != vace_references_encoded.shape[2]: + raise ValueError(f"First VACE has {vace_references_encoded.shape[2]} reference frames, but second VACE has {vace_references_encoded2.shape[2]} reference frames. They must match for proper tensor stacking.") + + control_video_latent2 = torch.cat((vace_references_encoded2, control_video_latent2), dim=2) + elif vace_references_encoded is not None: + # Reuse first references if no second references provided + vace_references_encoded2 = vace_references_encoded + control_video_latent2 = torch.cat((vace_references_encoded, control_video_latent2), dim=2) + + # Create strength list for second VACE using its specific references + if isinstance(vace_strength2, list): + if len(vace_strength2) != original_vace_latent_length: + raise ValueError(f"If vace_strength2 is a list, it must have length {original_vace_latent_length}, got {len(vace_strength2)}") + + if vace_references_encoded2 is not None: + if not isinstance(vace_strength2, list): + if vace_ref2_strength != vace_strength2: + vace_strength_list2 = [[vace_ref2_strength, vace_strength2]] * batch_size + else: + inner_strength_list = [vace_strength2] * (original_vace_latent_length + vace_references_encoded2.shape[2]) + vace_strength_list2 = [inner_strength_list] * batch_size + else: + vace_strength_list2 = [[vace_ref2_strength] + vace_strength2] * batch_size + else: + if isinstance(vace_strength2, list): + vace_strength_list2 = [vace_strength2] * batch_size + else: + vace_strength_list2 = [[vace_strength2] * original_vace_latent_length] * batch_size + + # Process mask same as first + mask2 = mask2.view(vace_length, height_mask, vae_stride, width_mask, vae_stride) + mask2 = mask2.permute(2, 4, 0, 1, 3) + mask2 = mask2.reshape(vae_stride * vae_stride, vace_length, height_mask, width_mask) + mask2 = torch.nn.functional.interpolate(mask2.unsqueeze(0), size=(original_vace_latent_length, height_mask, width_mask), mode='nearest-exact').squeeze(0) + + if vace_references_encoded2 is not None: + mask2_pad = torch.zeros_like(mask2[:, :vace_references_encoded2.shape[2], :, :]) + mask2 = torch.cat((mask2_pad, mask2), dim=1) + + mask2 = mask2.unsqueeze(0) + + # Add to lists + vace_frames_list.append(control_video_latent2) + vace_mask_list.append(mask2) + vace_strength_list_all.append(vace_strength_list2) + + positive = node_helpers.conditioning_set_values(positive, { + "vace_frames": vace_frames_list, + "vace_mask": vace_mask_list, + "vace_strength": vace_strength_list_all}, + append=True) + negative = node_helpers.conditioning_set_values(negative, { + "vace_frames": vace_frames_list, + "vace_mask": vace_mask_list, + "vace_strength": vace_strength_list_all}, + append=True) + + ###################### + # execute WanPhantomSubjectToVideo logic + phantom_length = length + # Create the latent from WanPhantomSubjectToVideo (this will be our output latent) + latent = torch.zeros([batch_size, 16, ((phantom_length - 1) // 4) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device()) + + # WanPhantomSubjectToVideo uses the negative from WanVace as input and creates two outputs + neg_phant_img = negative # This becomes negative_img_text (with zeros) + + if phantom_images is not None: + phantom_images_scaled = comfy.utils.common_upscale(phantom_images[:num_phantom_images].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1) + latent_images = [] + for i in phantom_images_scaled: + latent_images += [vae.encode(i.unsqueeze(0)[:, :, :, :3])] + concat_latent_image = torch.cat(latent_images, dim=2) + + # Apply phantom logic: positive gets the phantom images + positive = node_helpers.conditioning_set_values(positive, {"time_dim_concat": concat_latent_image}) + # negative_text (cond2 in original) gets the phantom images + negative = node_helpers.conditioning_set_values(negative, {"time_dim_concat": concat_latent_image}) + # neg_phant_img (negative in original) gets zeros instead of phantom images + neg_phant_img = node_helpers.conditioning_set_values(negative, {"time_dim_concat": comfy.latent_formats.Wan21().process_out(torch.zeros_like(concat_latent_image))}) + + # Adjust latent size if reference image was provided (WanVaceToVideo logic) + if vace_references_encoded is not None: + # Prepend zeros to match the reference image dimensions + latent_pad = torch.zeros([batch_size, 16, vace_references_encoded.shape[2], height // 8, width // 8], device=latent.device, dtype=latent.dtype) + latent = torch.cat([latent_pad, latent], dim=2) + + out_latent = {} + out_latent["samples"] = latent + + return (positive, negative, neg_phant_img, out_latent, trim_latent) + + +from ..patches import wrap_vace_phantom_wan_model, unwrap_vace_phantom_wan_model + +class WanVaceToVideoLatent: + @classmethod + def INPUT_TYPES(s): + return {"required": {"positive": ("CONDITIONING", ), + "negative": ("CONDITIONING", ), + "vae": ("VAE", ), + "width": ("INT", {"default": 832, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}), + "height": ("INT", {"default": 480, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}), + "length": ("INT", {"default": 81, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 4}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), + "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1000.0, "step": 0.01}), + }, + "optional": {"control_latent_input": ("LATENT", ), + "reference_latent": ("LATENT", ), + }} + + RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT", "INT") + RETURN_NAMES = ("positive", "negative", "latent", "trim_latent") + FUNCTION = "encode" + + CATEGORY = "WanVaceAdvanced" + + EXPERIMENTAL = True + + def encode(self, positive, negative, vae, width, height, length, batch_size, strength, control_latent_input=None, reference_latent=None): + # Calculate dimensions for latent space + vae_stride = 8 + height_latent = height // vae_stride + width_latent = width // vae_stride + latent_length = ((length - 1) // 4) + 1 + + # Process control_latent_input if provided + if control_latent_input is not None: + # Extract the control frames from the latent input + control_frames_latent = control_latent_input["samples"] + + # Get the number of latent frames in the input + input_latent_length = control_frames_latent.shape[2] + + # TARGETED FIX: Create a full tensor of 0.5s in pixel space for the entire length + # This will be encoded to create the proper inactive tensor + full_pixel_tensor = torch.ones((length, height, width, 3), + device=comfy.model_management.intermediate_device()) * 0.5 + + # Encode this full tensor to get the proper inactive latent + inactive = vae.encode(full_pixel_tensor) + + # For reactive, we need to start with the same encoded 0.5s + reactive = inactive.clone() + + # Calculate the corresponding latent frames in the reactive tensor + latent_frames_to_replace = min(input_latent_length, latent_length) + + # Replace the relevant portion of the reactive tensor with our input latent + reactive[:, :, :latent_frames_to_replace] = control_frames_latent[:, :, :latent_frames_to_replace] + + else: + # If no control latent provided, create blank latents + # Create pixel-space frames with value 0.5 + blank_frames = torch.ones((length, height, width, 3), + device=comfy.model_management.intermediate_device()) * 0.5 + + # Encode through VAE + inactive = vae.encode(blank_frames) + reactive = inactive.clone() + + # Combine inactive and reactive + control_video_latent = torch.cat((inactive, reactive), dim=1) + + # Create a mask of ones with the same temporal dimension as control_video_latent + mask_temporal_length = control_video_latent.shape[2] + mask = torch.ones([1, vae_stride * vae_stride, mask_temporal_length, height_latent, width_latent], + device=control_video_latent.device) + + # Process reference latent if provided + trim_latent = 0 + if reference_latent is not None: + ref_s = reference_latent["samples"] + + # Process the reference latent to have 32 channels like in the original + # This matches how the original node processes reference images + ref_s = torch.cat([ref_s, comfy.latent_formats.Wan21().process_out(torch.zeros_like(ref_s))], dim=1) + + # Concatenate with control video latent along temporal dimension + control_video_latent = torch.cat((ref_s, control_video_latent), dim=2) + trim_latent = ref_s.shape[2] # Set trim amount to reference frame count + + # If we have a reference latent, we need to add zeros to mask for those frames + if reference_latent is not None: + mask_pad = torch.zeros([1, vae_stride * vae_stride, trim_latent, height_latent, width_latent], + device=control_video_latent.device) + mask = torch.cat((mask_pad, mask), dim=2) + + # Set conditioning values + positive = node_helpers.conditioning_set_values(positive, {"vace_frames": [control_video_latent], "vace_mask": [mask], "vace_strength": [strength]}, append=True) + negative = node_helpers.conditioning_set_values(negative, {"vace_frames": [control_video_latent], "vace_mask": [mask], "vace_strength": [strength]}, append=True) + + # Create output latent + output_latent_length = mask_temporal_length + if reference_latent is not None: + output_latent_length += trim_latent + + latent = torch.zeros([batch_size, 16, output_latent_length, height_latent, width_latent], + device=comfy.model_management.intermediate_device()) + + out_latent = {} + out_latent["samples"] = latent + + return (positive, negative, out_latent, trim_latent) + + +class VaceAdvancedModelPatch: + """ + Node to enable per-frame strength support for VaceWanModel. + This patches the model to support separate strength values for reference and control frames. + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "enable": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch_model" + CATEGORY = "WanVaceAdvanced" + + def patch_model(self, model, enable): + """ + Wrap or unwrap the VaceWanModel to enable/disable per-frame strength support. + + Args: + model: The model to patch + enable: Whether to enable or disable per-frame strength support + + Returns: + The patched model + """ + # Clone the model to avoid modifying the original + m = model.clone() + + if enable: + m = wrap_vace_phantom_wan_model(m) + else: + m = unwrap_vace_phantom_wan_model(m) + + return (m,) + + +class VaceStrengthTester: + """ + A utility node to test different strength combinations for VACE models. + This node helps visualize how different reference and control strengths affect the output. + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "positive": ("CONDITIONING",), + "reference_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}), + "control_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}), + "frame_count": ("INT", {"default": 20, "min": 1, "max": 1000}), + }, + "optional": { + "custom_strength_list": ("STRING", {"default": "", "multiline": True}), + } + } + + RETURN_TYPES = ("CONDITIONING", "STRING") + RETURN_NAMES = ("conditioning", "strength_info") + FUNCTION = "create_test_conditioning" + CATEGORY = "WanVaceAdvanced" + + def create_test_conditioning(self, positive, reference_strength, control_strength, frame_count, custom_strength_list=""): + """ + Create conditioning with test strength values. + + Args: + positive: Input conditioning + reference_strength: Strength for reference frame + control_strength: Strength for control frames + frame_count: Number of frames + custom_strength_list: Optional custom list of strength values (comma-separated) + + Returns: + Modified conditioning and info string + """ + + # Parse custom strength list if provided + if custom_strength_list.strip(): + try: + strength_list = [float(x.strip()) for x in custom_strength_list.split(",")] + # Pad or truncate to frame_count + if len(strength_list) < frame_count: + strength_list.extend([strength_list[-1]] * (frame_count - len(strength_list))) + else: + strength_list = strength_list[:frame_count] + + info = f"Using custom strength list: {strength_list[:10]}{'...' if len(strength_list) > 10 else ''}" + + except ValueError: + # Fall back to reference + control pattern + strength_list = [reference_strength] + [control_strength] * (frame_count - 1) + info = f"Invalid custom list, using ref:{reference_strength}, ctrl:{control_strength} for {frame_count} frames" + else: + # Standard reference + control pattern + strength_list = [reference_strength] + [control_strength] * (frame_count - 1) + info = f"Using ref:{reference_strength}, ctrl:{control_strength} for {frame_count} frames" + + # Clone conditioning to avoid modifying original + import copy + new_positive = copy.deepcopy(positive) + + # Find and update vace_strength in conditioning + for i, cond_item in enumerate(new_positive): + if len(cond_item) > 1 and isinstance(cond_item[1], dict): + if 'vace_strength' in cond_item[1]: + cond_item[1]['vace_strength'] = strength_list + cond_item[1]['vace_has_reference'] = True + info += f" | Applied to conditioning item {i}" + + return (new_positive, info) + + +NODE_CLASS_MAPPINGS = { + "WanVacePhantomExperimental": WanVacePhantomExperimental, + "WanVaceToVideoLatent": WanVaceToVideoLatent, + "VaceAdvancedModelPatch": VaceAdvancedModelPatch, + "VaceStrengthTester": VaceStrengthTester, + +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "WanVacePhantomAdvanced": "WanVacePhantomAdvanced", + "WanVaceToVideoLatent": "WanVaceToVideoLatent", + "VaceAdvancedModelPatch": "VaceAdvancedModelPatch", + "VaceStrengthTester": "VaceStrengthTester", +} \ No newline at end of file diff --git a/patches.py b/patches.py new file mode 100644 index 0000000..04ad992 --- /dev/null +++ b/patches.py @@ -0,0 +1,381 @@ +import comfy.model_management +import comfy.model_patcher +import logging + +""" +VaceWanModel patching system for handling per-frame strength values. +This section provides a wrapper for the VaceWanModel that supports separate +strength values for reference frames and control frames. +""" + +# Test mode flags +TEST_MODE_1_ZEROS_INSTEAD_OF_PATCH_EMBEDDING = False +TEST_MODE_2_USE_ONES_VALUES_INSTEAD_OF_PATCH_EMBEDDING = False +TEST_MODE_3_USE_SCALED_VALUES_INSTEAD_OF_PATCH_EMBEDDING = False +TEST_SCALE = 1.0 + +import torch +import comfy.ldm.wan.model +from comfy.ldm.wan.model import sinusoidal_embedding_1d + +def wrap_vace_phantom_wan_model(model): + """ + Patch a VaceWanModel to support per-frame strength and Phantom embeds. + Based on the WanVideoTeaCacheKJ implementation in comfyui-kjnodes. + """ + model_clone = model.clone() + + # Check if this is a VaceWanModel + diffusion_model = model_clone.get_model_object("diffusion_model") + if not isinstance(diffusion_model, comfy.ldm.wan.model.VaceWanModel): + logging.info("Not a VaceWanModel, skipping per-frame strength patch") + return model_clone + + def outer_wrapper(): + def unet_wrapper_function(model_function, kwargs): + # Extract parameters from kwargs + input_data = kwargs["input"] + timestep = kwargs["timestep"] + c = kwargs["c"] + + vace_strength = c.get("vace_strength", [1.0]) + + # Check if we have nested lists (our separate reference strength format) + use_patched_forward_orig = ( + isinstance(vace_strength, list) and + len(vace_strength) > 0 and + isinstance(vace_strength[0], list) + ) + + # Also check if we have Phantom concats (c.get("time_dim_concat") exists and is a tensor) + use_patched_forward_orig = use_patched_forward_orig or ( + c.get("time_dim_concat") is not None and + isinstance(c.get("time_dim_concat"), torch.Tensor) + ) + + if use_patched_forward_orig: + # Extract frame count information from conditioning + phantom_frames = 0 + reference_frames = 0 + + # Check for phantom frames from time_dim_concat + time_dim_concat = c.get("time_dim_concat") + if time_dim_concat is not None and isinstance(time_dim_concat, torch.Tensor): + phantom_frames = time_dim_concat.shape[2] # temporal dimension + + # Check for reference frames from c['vace_context'] + vace_context = c.get("vace_context") + if vace_context is None or not isinstance(vace_context, torch.Tensor): + raise ValueError("vace_context is missing or not a tensor") + else: + total_frames = vace_context.shape[-3] + control_frames = input_data.shape[-3] + reference_frames = total_frames - control_frames - phantom_frames + if reference_frames < 0: + raise ValueError("Reference frames count cannot be negative. Check vace_context and input data shapes.") + + # Store frame counts in conditioning for the patched function + c["_vace_reference_frames"] = reference_frames + c["_vace_phantom_frames"] = phantom_frames + + # Use the patched forward_orig method + from unittest.mock import patch + + forward_function = _vaceph_forward_orig + context = patch.multiple( + diffusion_model, + forward_orig=forward_function.__get__(diffusion_model, diffusion_model.__class__) + ) + + with context: + out = model_function(input_data, timestep, **c) + else: + # Use original behavior + out = model_function(input_data, timestep, **c) + + return out + return unet_wrapper_function + + model_clone.set_model_unet_function_wrapper(outer_wrapper()) + + logging.info("VaceWanModel patched for separate reference strength support using unet wrapper") + return model_clone + + +def unwrap_vace_phantom_wan_model(model): + """ + Remove the patch from a VaceWanModel. + """ + model_clone = model.clone() + logging.info("VaceWanModel patch removed") + return model_clone + +def _vaceph_forward_orig( + self, + x, + t, + context, + vace_context, + vace_strength, + clip_fea=None, + freqs=None, + transformer_options={}, + **kwargs, +): + """ + Custom forward implementation that applies separate strength values for all Vace reference and control frames. + Also allows for separate Phantom frames compatibility. + The Vace reference frames are at the beginning of the temporal dimension (dim 2). + The Phantom frames are concatenated at the end of the temporal dimension (dim 2). + This is adapted from the original VaceWanModel.forward_orig method. + """ + + # Extract frame count information for strength application + phantom_frames = 0 + time_dim_concat = kwargs.get("time_dim_concat") + if time_dim_concat is not None and isinstance(time_dim_concat, torch.Tensor): + phantom_frames = time_dim_concat.shape[2] + + # embeddings + x = self.patch_embedding(x.float()).to(x.dtype) + grid_sizes = x.shape[2:] + x = x.flatten(2).transpose(1, 2) + + # time embeddings + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, t).to(dtype=x.dtype)) + e0 = self.time_projection(e).unflatten(1, (6, self.dim)) + + # context + context = self.text_embedding(context) + + context_img_len = None + if clip_fea is not None: + if self.img_emb is not None: + context_clip = self.img_emb(clip_fea) # bs x 257 x dim + context = torch.concat([context_clip, context], dim=1) + context_img_len = clip_fea.shape[-2] + + # Process vace_context + orig_shape = list(vace_context.shape) + vace_context = vace_context.movedim(0, 1).reshape([-1] + orig_shape[2:]) + c = self.vace_patch_embedding(vace_context.float()).to(vace_context.dtype) + vace_grid_sizes = c.shape[2:] # Get Vace patch dimensions (t_patches, h_patches, w_patches) + c = c.flatten(2).transpose(1, 2) + + # Store expected sequence length from x for comparison + expected_seq_length = x.shape[1] + + # Extract reference frame count by analyzing the vace context structure + # Reference frames have 32 channels (16 + 16 zeros), control frames have 32 channels (16 inactive + 16 reactive) + # But we need to determine this from the strength list structure when we have it + reference_frames = 0 + + # Split into batch segments + c = list(c.split(orig_shape[0], dim=0)) + + # arguments + x_orig = x + + patches_replace = transformer_options.get("patches_replace", {}) + blocks_replace = patches_replace.get("dit", {}) + + for i, block in enumerate(self.blocks): + if ("double_block", i) in blocks_replace: + def block_wrap(args): + out = {} + out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], context_img_len=context_img_len) + return out + out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs}, {"original_block": block_wrap}) + x = out["img"] + else: + x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len) + + ii = self.vace_layers_mapping.get(i, None) + if ii is not None: + for iii in range(len(c)): + # Check and fix sequence length mismatch before passing to vace_blocks + if c[iii].shape[1] < expected_seq_length: + # Calculate how many frames to add + seq_length_per_frame = grid_sizes[1] * grid_sizes[2] + extra_seq_length = expected_seq_length - c[iii].shape[1] + extra_frames = extra_seq_length // seq_length_per_frame + # Create a tensor of shape [1, 16, extra_frames, orig_shape[-2], orig_shape[-1]] + # Keep new tensors on CPU until added to the sequence + SCALE_FACTOR = 0.5 # Scale factor (what should this be?!) + solid_frames = torch.ones(1, 16, extra_frames, orig_shape[-2], orig_shape[-1], device="cpu", dtype=c[iii].dtype) * SCALE_FACTOR + + from comfy.latent_formats import Wan21 + processed_frames = Wan21().process_out(solid_frames) + processed_frames = self.patch_embedding(processed_frames.float()).to(processed_frames.dtype) + c_padding = processed_frames.flatten(2).transpose(1, 2) + + # Concatenate along sequence dimension to extend c + c[iii] = torch.cat([c[iii], c_padding.to(c[iii].device)], dim=1) + + # if i == 0 and iii == 0: + # logging.info(f"\nVace-Phantom compatibility: Extended Vace frames by adding {extra_frames} frames to match sequence length {expected_seq_length}.\n") + + del processed_frames, c_padding, solid_frames + + elif c[iii].shape[1] > expected_seq_length: + # If c[iii] is longer than expected, truncate it + logging.warning(f"Truncating Vace frames for batch {iii} from {c[iii].shape[1]} to {expected_seq_length} frames.") + c[iii] = c[iii][:, :expected_seq_length, :] + + # Continue with original processing + c_skip, c[iii] = self.vace_blocks[ii]( + c[iii], x=x_orig, e=e0, freqs=freqs, + context=context, context_img_len=context_img_len + ) + + # CUSTOM LOGIC: Apply separate reference strength if provided + # iii is the batch index, vace_strength[iii] is the strength for this batch item + if iii < len(vace_strength): + batch_strength = vace_strength[iii] + else: + # Fallback to last strength if list is shorter + batch_strength = vace_strength[-1] + + # Extract frame counts from conditioning + reference_frames = kwargs.get("_vace_reference_frames", 0) + phantom_frames = kwargs.get("_vace_phantom_frames", 0) + + # Handle nested list (separate reference/control strengths) or single value + if isinstance(batch_strength, list): + # We have separate strengths as a list + x = _apply_separate_strengths(x, c_skip, batch_strength, reference_frames, phantom_frames, vace_grid_sizes) + else: + # Single strength value - original behavior + x += c_skip * batch_strength + + del c_skip + + # head + x = self.head(x, e) + + # unpatchify + x = self.unpatchify(x, grid_sizes) + return x + + +def _apply_separate_strengths(x, c_skip, strength_list, reference_frames=0, phantom_frames=0, vace_grid_sizes=None): + """ + Apply separate strength values for Vace frames with element-wise application. + + The sequence layout is: [reference frames] + [control frames] + [phantom frames] + Each portion gets element-wise strength application based on the strength_list. + + Args: + x: Main feature tensor [B, sequence_length, hidden_dim] + c_skip: Vace skip connection [B, sequence_length, hidden_dim] + strength_list: List of strengths matching the total frame count + reference_frames: Number of reference frames (0 if no reference) + phantom_frames: Number of phantom frames (0 if no phantom) + vace_grid_sizes: Actual patch grid sizes (t_patches, h_patches, w_patches) from Vace embedding + """ + # Get tensor dimensions + batch_size, sequence_length, hidden_dim = c_skip.shape + + # Simple case: single strength for everything + if len(strength_list) == 1 and not isinstance(strength_list[0], list): + x += c_skip * strength_list[0] + return x + + # Calculate control frames + control_frames = len(strength_list[0]) - reference_frames - phantom_frames + + # Calculate patches per frame if grid sizes available + if vace_grid_sizes is not None and len(vace_grid_sizes) >= 3: + t_patches, h_patches, w_patches = vace_grid_sizes + patches_per_frame = h_patches * w_patches + else: + # Fallback: assume equal distribution + total_frames = reference_frames + control_frames + phantom_frames + if total_frames > 0: + patches_per_frame = sequence_length // total_frames + else: + patches_per_frame = sequence_length + + try: + # Calculate sequence lengths for each portion + reference_seq_length = reference_frames * patches_per_frame + control_seq_length = control_frames * patches_per_frame + phantom_seq_length = phantom_frames * patches_per_frame + + # Ensure we don't exceed the actual sequence length + total_calculated = reference_seq_length + control_seq_length + phantom_seq_length + if total_calculated > sequence_length: + # Adjust phantom sequence length to fit + phantom_seq_length = sequence_length - reference_seq_length - control_seq_length + phantom_seq_length = max(0, phantom_seq_length) + + # Extract strength portions + flat_strength_list = strength_list[0] # Get the flat list for this batch + reference_strengths = flat_strength_list[:reference_frames] if reference_frames > 0 else [] + control_strengths = flat_strength_list[reference_frames:reference_frames + control_frames] + phantom_strengths = flat_strength_list[reference_frames + control_frames:] if phantom_frames > 0 else [] + + # Split c_skip into portions + current_pos = 0 + weighted_portions = [] + + # Process reference frames + if reference_seq_length > 0: + ref_portion = c_skip[:, current_pos:current_pos + reference_seq_length, :] + # Apply element-wise strengths to reference frames + weighted_ref = torch.zeros_like(ref_portion) + for i, strength in enumerate(reference_strengths): + start_idx = i * patches_per_frame + end_idx = min((i + 1) * patches_per_frame, reference_seq_length) + weighted_ref[:, start_idx:end_idx, :] = ref_portion[:, start_idx:end_idx, :] * strength + weighted_portions.append(weighted_ref) + current_pos += reference_seq_length + + # Process control frames + if control_seq_length > 0: + ctrl_portion = c_skip[:, current_pos:current_pos + control_seq_length, :] + # Apply element-wise strengths to control frames + weighted_ctrl = torch.zeros_like(ctrl_portion) + for i, strength in enumerate(control_strengths): + start_idx = i * patches_per_frame + end_idx = min((i + 1) * patches_per_frame, control_seq_length) + weighted_ctrl[:, start_idx:end_idx, :] = ctrl_portion[:, start_idx:end_idx, :] * strength + weighted_portions.append(weighted_ctrl) + current_pos += control_seq_length + + # Process phantom frames (if any remaining sequence) + if current_pos < sequence_length: + phantom_portion = c_skip[:, current_pos:, :] + if phantom_frames > 0 and phantom_strengths: + # Apply element-wise strengths to phantom frames + weighted_phantom = torch.zeros_like(phantom_portion) + for i, strength in enumerate(phantom_strengths): + start_idx = i * patches_per_frame + end_idx = min((i + 1) * patches_per_frame, phantom_portion.shape[1]) + if start_idx < phantom_portion.shape[1]: + weighted_phantom[:, start_idx:end_idx, :] = phantom_portion[:, start_idx:end_idx, :] * strength + weighted_portions.append(weighted_phantom) + else: + # No phantom frames expected, use last control strength + last_strength = control_strengths[-1] if control_strengths else 1.0 + weighted_portions.append(phantom_portion * last_strength) + + # Concatenate all weighted portions + if weighted_portions: + weighted_c_skip = torch.cat(weighted_portions, dim=1) + else: + # Fallback + weighted_c_skip = c_skip * strength_list[0] + + except Exception as e: + # Fallback to uniform application + logging.warning(f"Failed to apply element-wise strengths, using uniform: {e}") + flat_strength_list = strength_list[0] if isinstance(strength_list[0], list) else strength_list + avg_strength = sum(flat_strength_list) / len(flat_strength_list) + weighted_c_skip = c_skip * avg_strength + + # Add to main features + x += weighted_c_skip + + return x \ No newline at end of file