diff --git a/.gitignore b/.gitignore index 1584c81..da9642b 100644 --- a/.gitignore +++ b/.gitignore @@ -7,9 +7,7 @@ nodes/__pycache__/ .env2 opencode.json wanvideo_cached_samples.pt -wanvideo_cached_videos.pt ref_image.jpg QWEN.md -ComfyUI-WanVideoWrapper/ rgthree-comfy-main/ -utility/userdata/info/ \ No newline at end of file +utility/userdata/info/ diff --git a/__init__.py b/__init__.py index d0de23b..43b887e 100644 --- a/__init__.py +++ b/__init__.py @@ -22,6 +22,7 @@ from .nodes.image_nodes import * from .nodes.image_nodes_extra import * from .nodes.image_to_video import * from .nodes.video_nodes_extra import * +from .nodes.wan_first_middle_last import * from .nodes.vace_utils import * from .nodes.cache_samples import * diff --git a/nodes/cache_samples.py b/nodes/cache_samples.py index e3d6efa..1e47880 100644 --- a/nodes/cache_samples.py +++ b/nodes/cache_samples.py @@ -1,5 +1,80 @@ import os +import sys import torch +import numpy as np +import json +import subprocess +import tempfile +from PIL import Image +import datetime +import re +import folder_paths + +def get_ffmpeg_path(): + """Minimal ffmpeg detection based on VideoHelperSuite approach""" + # Check environment variable first + ffmpeg_path = os.environ.get("VHS_FORCE_FFMPEG_PATH") + if ffmpeg_path and os.path.exists(ffmpeg_path): + return ffmpeg_path + + ffmpeg_paths = [] + + # Try to use imageio-ffmpeg (same as VideoHelperSuite) + try: + from imageio_ffmpeg import get_ffmpeg_exe + imageio_ffmpeg_path = get_ffmpeg_exe() + if imageio_ffmpeg_path: + ffmpeg_paths.append(imageio_ffmpeg_path) + except ImportError: + pass + + # Try system PATH + if os.name == 'nt': # Windows + ffmpeg_names = ['ffmpeg.exe'] + else: # Linux/Mac + ffmpeg_names = ['ffmpeg'] + + for name in ffmpeg_names: + try: + result = subprocess.run([name, '-version'], capture_output=True, check=True, timeout=5) + ffmpeg_paths.append(name) + break + except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired): + continue + + # Try common locations (same as VideoHelperSuite) + if os.name == 'nt': + common_paths = [ + os.path.abspath("ffmpeg.exe"), + os.path.abspath(os.path.join(os.path.dirname(__file__), "ffmpeg.exe")), + ] + else: + common_paths = [ + os.path.abspath("ffmpeg"), + os.path.abspath(os.path.join(os.path.dirname(__file__), "ffmpeg")), + ] + + for path in common_paths: + if os.path.exists(path): + ffmpeg_paths.append(path) + + # Return the best path (prioritize system PATH over bundled) + if len(ffmpeg_paths) == 0: + return None + elif len(ffmpeg_paths) == 1: + return ffmpeg_paths[0] + else: + # Prefer system ffmpeg over bundled ones + return ffmpeg_paths[-1] + +def tensor_to_bytes(tensor): + """Convert tensor to bytes""" + tensor = tensor.cpu().numpy() + tensor = np.clip(tensor * 255, 0, 255).astype(np.uint8) + return tensor + +# Get ffmpeg path using minimal detection +FFMPEG_PATH = get_ffmpeg_path() # Use a simple string for imageOrLatent to avoid JSON serialization issues imageOrLatent = "IMAGE" @@ -14,6 +89,7 @@ class WanVideoCacheSamples: "required": { "cache_path": ("STRING", {"default": "cache_data", "multiline": False}), "cache_name": ("STRING", {"default": "wanvideo_cached_samples.pt", "multiline": False}), + "save_over": ("BOOLEAN", {"default": True}), }, "optional": { "samples": ("LATENT", {"forceInput": True}), @@ -25,33 +101,42 @@ class WanVideoCacheSamples: FUNCTION = "cache_or_load_samples" CATEGORY = "WanVideoWrapper_QQ/utils" - def cache_or_load_samples(self, cache_path, cache_name, samples=None): - # Get the path to the ComfyUI-WanVideoWrapper_QQ directory - qq_wrapper_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + def cache_or_load_samples(self, cache_path, cache_name, save_over, samples=None): + # Get ComfyUI's output directory + output_dir = folder_paths.get_output_directory() # Sanitize cache_path to prevent it from being an absolute path # and ensure it's just a directory name. sanitized_cache_path = os.path.basename(cache_path) - # The cache directory will be inside the node's directory. - full_cache_dir = os.path.join(qq_wrapper_path, sanitized_cache_path) + # The cache directory will be inside ComfyUI's output directory. + full_cache_dir = os.path.join(output_dir, sanitized_cache_path) os.makedirs(full_cache_dir, exist_ok=True) - + full_cache_file_path = os.path.join(full_cache_dir, cache_name) - if samples is not None: - # Samples are provided, cache them and pass them through - print(f"WanVideoCacheSamples: Caching samples to {full_cache_file_path}") - torch.save(samples, full_cache_file_path) - return (samples,) + if save_over: + # Save over mode: use input samples and cache them + if samples is not None: + print(f"WanVideoCacheSamples: Caching samples to {full_cache_file_path}") + torch.save(samples, full_cache_file_path) + return (samples,) + else: + # No samples provided, try to load from cache as fallback + if os.path.exists(full_cache_file_path): + print(f"WanVideoCacheSamples: No input provided, loading samples from cache {full_cache_file_path}") + loaded_samples = torch.load(full_cache_file_path) + return (loaded_samples,) + else: + raise FileNotFoundError(f"WanVideoCacheSamples: No samples provided and no cached samples found at {full_cache_file_path}") else: - # Samples are not provided, try to load from cache + # Load-only mode: ignore input, only load from cache if os.path.exists(full_cache_file_path): - print(f"WanVideoCacheSamples: Loading samples from {full_cache_file_path}") + print(f"WanVideoCacheSamples: Load-only mode - loading samples from cache {full_cache_file_path}") loaded_samples = torch.load(full_cache_file_path) return (loaded_samples,) else: - raise FileNotFoundError(f"WanVideoCacheSamples: No samples provided and no cached samples found at {full_cache_file_path}") + raise FileNotFoundError(f"WanVideoCacheSamples: Load-only mode - no cached samples found at {full_cache_file_path}") class WanVideoCacheVideos: @@ -63,59 +148,315 @@ class WanVideoCacheVideos: return { "required": { "cache_path": ("STRING", {"default": "cache_data", "multiline": False}), - "cache_name": ("STRING", {"default": "wanvideo_cached_videos.pt", "multiline": False}), + "cache_name": ("STRING", {"default": "wanvideo_cached_video.mp4", "multiline": False}), + "save_over": ("BOOLEAN", {"default": True}), }, "optional": { "images": (imageOrLatent, {"forceInput": True}), - "frame_count": ("INT", {"forceInput": True}), - "audio": ("AUDIO", {"forceInput": True}), - "video_info": ("VHS_VIDEOINFO", {"forceInput": True}), + "save_multiple": ("BOOLEAN", {"default": False}), + "use_multi_cached": ("BOOLEAN", {"default": False}), + "cache_num": ("INT", {"default": 1, "min": 1, "max": 999, "step": 1}), } } - RETURN_TYPES = (imageOrLatent, "INT", "AUDIO", "VHS_VIDEOINFO") - RETURN_NAMES = ("images", "frame_count", "audio", "video_info") + RETURN_TYPES = (imageOrLatent,) + RETURN_NAMES = ("images",) FUNCTION = "cache_or_load_videos" CATEGORY = "WanVideoWrapper_QQ/utils" - def cache_or_load_videos(self, cache_path, cache_name, images=None, frame_count=None, audio=None, video_info=None): - # Get the path to the ComfyUI-WanVideoWrapper_QQ directory - qq_wrapper_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + def cache_or_load_videos(self, cache_path, cache_name, save_over, images=None, save_multiple=False, use_multi_cached=False, cache_num=1): + # Get ComfyUI's output directory + output_dir = folder_paths.get_output_directory() # Sanitize cache_path to prevent it from being an absolute path # and ensure it's just a directory name. sanitized_cache_path = os.path.basename(cache_path) - # The cache directory will be inside the node's directory. - full_cache_dir = os.path.join(qq_wrapper_path, sanitized_cache_path) + # The cache directory will be inside ComfyUI's output directory. + full_cache_dir = os.path.join(output_dir, sanitized_cache_path) os.makedirs(full_cache_dir, exist_ok=True) - + full_cache_file_path = os.path.join(full_cache_dir, cache_name) - if images is not None: - # Video data is provided, cache them and pass them through - print(f"WanVideoCacheVideos: Caching video data to {full_cache_file_path}") - video_data = { - "images": images, - "frame_count": frame_count, - "audio": audio, - "video_info": video_info - } - torch.save(video_data, full_cache_file_path) - return (images, frame_count, audio, video_info) - else: - # Video data is not provided, try to load from cache - if os.path.exists(full_cache_file_path): - print(f"WanVideoCacheVideos: Loading video data from {full_cache_file_path}") - loaded_video_data = torch.load(full_cache_file_path) - return ( - loaded_video_data["images"], - loaded_video_data["frame_count"], - loaded_video_data["audio"], - loaded_video_data["video_info"] - ) + if FFMPEG_PATH is None: + raise Exception("ffmpeg is required for MP4 video caching but VideoHelperSuite could not find it.\n\nTo fix this:\n1. Make sure VideoHelperSuite is properly installed (Video Combine 🎥🅥🅗🅢 node should work)\n2. If Video Combine works but this doesn't, there may be a path issue\n3. Try installing imageio-ffmpeg: pip install imageio-ffmpeg\n4. Or install ffmpeg system-wide from https://ffmpeg.org/download.html") + + if save_over: + # Save over mode: use input images and cache them as MP4 + if images is not None: + # Determine the actual filename to save + if save_multiple: + actual_save_path = self._get_cache_filename(full_cache_file_path, save_multiple=True) + print(f"WanVideoCacheVideos: Saving images as MP4 to {actual_save_path} (multiple cache mode)") + else: + actual_save_path = full_cache_file_path + print(f"WanVideoCacheVideos: Saving images as MP4 to {actual_save_path}") + + self._save_images_as_mp4(images, actual_save_path) + return (images,) else: - raise FileNotFoundError(f"WanVideoCacheVideos: No video data provided and no cached video found at {full_cache_file_path}") + # No images provided, try to load from cache as fallback + cache_file = self._find_cached_file(full_cache_file_path, use_multi_cached, cache_num) + if cache_file: + print(f"WanVideoCacheVideos: No input provided, loading images from MP4 {cache_file}") + loaded_images = self._load_images_from_mp4(cache_file) + return (loaded_images,) + else: + raise FileNotFoundError(f"WanVideoCacheVideos: No images provided and no cached MP4 file found (searched for numbered versions up to {cache_num})") + else: + # Load-only mode: ignore input, only load from MP4 cache + cache_file = self._find_cached_file(full_cache_file_path, use_multi_cached, cache_num) + if cache_file: + print(f"WanVideoCacheVideos: Load-only mode - loading images from MP4 {cache_file}") + loaded_images = self._load_images_from_mp4(cache_file) + return (loaded_images,) + else: + if use_multi_cached: + raise FileNotFoundError(f"WanVideoCacheVideos: Load-only mode - no cached MP4 file found for cache_num {cache_num} (tried base filename too)") + else: + raise FileNotFoundError(f"WanVideoCacheVideos: Load-only mode - no cached MP4 file found at {full_cache_file_path}") + + def _save_images_as_mp4(self, images, output_path): + """Save images as MP4 video with 16 fps, H.264, CRF 18""" + if FFMPEG_PATH is None: + raise Exception("ffmpeg is required for MP4 caching but could not be found") + + try: + # Handle both IMAGE and LATENT formats + if isinstance(images, dict) and "samples" in images: + # LATENT format - convert to images (basic conversion) + latent_samples = images["samples"] + if latent_samples.dim() == 4: + # Convert from latent space to pixel space + images_np = ((latent_samples * 0.5 + 0.5) * 255).clamp(0, 255).cpu().numpy().astype(np.uint8) + images_np = np.transpose(images_np, (0, 2, 3, 1)) # NCHW -> NHWC + else: + raise ValueError(f"Unexpected latent tensor shape: {latent_samples.shape}") + else: + # IMAGE format - ensure it's in the right format + if torch.is_tensor(images): + images_np = images.cpu().numpy() + else: + images_np = images + + if images_np.ndim == 4: + # Handle different tensor formats + if images_np.shape[1] in [3, 4]: # NCHW format + images_np = np.transpose(images_np, (0, 2, 3, 1)) # NCHW -> NHWC + images_np = (images_np * 255).astype(np.uint8) + + # Get dimensions from first frame + if len(images_np) == 0: + raise ValueError("No images to save") + + first_frame = images_np[0] + height, width = first_frame.shape[:2] + has_alpha = first_frame.shape[-1] == 4 + + # Choose pixel format + if has_alpha: + i_pix_fmt = 'rgba' + o_pix_fmt = 'yuva420p' + else: + i_pix_fmt = 'rgb24' + o_pix_fmt = 'yuv420p' + + # Prepare ffmpeg arguments (16 fps, H.264, CRF 18) + frame_rate = 16 + args = [ + FFMPEG_PATH, "-y", # Overwrite output file + "-f", "rawvideo", + "-pix_fmt", i_pix_fmt, + "-s", f"{width}x{height}", + "-r", str(frame_rate), + "-i", "-", + "-c:v", "libx264", + "-crf", "18", + "-pix_fmt", o_pix_fmt, + "-r", str(frame_rate), + output_path + ] + + # Process frames and send to ffmpeg + process = subprocess.Popen(args, + stdin=subprocess.PIPE, + stderr=subprocess.PIPE, + env=os.environ.copy()) + + try: + for img_np in images_np: + # Convert frame to bytes and send to ffmpeg + frame_bytes = img_np.astype(np.uint8).tobytes() + process.stdin.write(frame_bytes) + + process.stdin.close() + stderr_output = process.stderr.read() + + return_code = process.wait() + if return_code != 0: + error_msg = stderr_output.decode('utf-8', errors='ignore') if stderr_output else "Unknown ffmpeg error" + raise Exception(f"FFmpeg error (code {return_code}): {error_msg}") + elif stderr_output: + print(f"FFmpeg warning: {stderr_output.decode('utf-8', errors='ignore')}") + + print(f"WanVideoCacheVideos: Successfully saved MP4 to {output_path}") + + except BrokenPipeError as e: + process.stdin.close() + stderr_output = process.stderr.read() + error_msg = stderr_output.decode('utf-8', errors='ignore') if stderr_output else "Broken pipe error" + raise Exception(f"FFmpeg broken pipe: {error_msg}") + + except Exception as e: + print(f"WanVideoCacheVideos: Error saving MP4: {e}") + raise e + + def _get_cache_filename(self, base_path, save_multiple=False): + """Generate cache filename with number suffix if save_multiple is True""" + if not save_multiple: + return base_path + + # Split filename and extension + name, ext = os.path.splitext(base_path) + + # Find the next available number + num = 1 + while True: + if num < 10: + suffix = f"_0{num}" + else: + suffix = f"_{num}" + + new_filename = f"{name}{suffix}{ext}" + if not os.path.exists(new_filename): + return new_filename + num += 1 + + def _find_cached_file(self, base_path, use_multi_cached=False, cache_num=1): + """Find cached file with optional number suffix""" + if not use_multi_cached: + # Try base filename first + if os.path.exists(base_path): + return base_path + else: + return None + + # Try numbered version + name, ext = os.path.splitext(base_path) + if cache_num < 10: + suffix = f"_0{cache_num}" + else: + suffix = f"_{cache_num}" + + numbered_filename = f"{name}{suffix}{ext}" + if os.path.exists(numbered_filename): + return numbered_filename + + # If numbered version not found, fall back to base filename + if os.path.exists(base_path): + return base_path + + return None + + def _load_images_from_mp4(self, video_path): + """Load frames from MP4 video""" + if FFMPEG_PATH is None: + raise Exception("ffmpeg is required for MP4 loading but could not be found") + + try: + # First get video info to get dimensions + info_cmd = [FFMPEG_PATH, '-i', video_path] + result = subprocess.run(info_cmd, capture_output=True, text=True) + + width, height = None, None + + # Parse ffmpeg output for dimensions + import re + dimension_match = re.search(r'(\d{3,4})x(\d{3,4})', result.stderr) + if dimension_match: + width, height = int(dimension_match.group(1)), int(dimension_match.group(2)) + print(f"Detected video dimensions: {width}x{height}") + + if width is None or height is None: + # Try ffprobe as backup + ffprobe_path = FFMPEG_PATH.replace('ffmpeg', 'ffprobe') if 'ffmpeg' in FFMPEG_PATH else None + if ffprobe_path and os.path.exists(ffprobe_path): + try: + probe_cmd = [ffprobe_path, '-v', 'error', '-select_streams', 'v:0', + '-show_entries', 'stream=width,height', '-of', 'csv=s=x:p=0', video_path] + probe_result = subprocess.run(probe_cmd, capture_output=True, text=True, timeout=10) + if probe_result.returncode == 0 and probe_result.stdout.strip(): + dims = probe_result.stdout.strip().split(',') + if len(dims) >= 2: + width, height = int(dims[0]), int(dims[1]) + print(f"Probed video dimensions: {width}x{height}") + except Exception as probe_error: + print(f"Warning: ffprobe failed: {probe_error}") + + # If still no dimensions, try inference from common sizes + if width is None or height is None: + print("Could not detect dimensions, trying inference...") + # First, get a small sample of frame data + sample_cmd = [FFMPEG_PATH, '-i', video_path, '-frames:v', '1', + '-f', 'image2pipe', '-pix_fmt', 'rgb24', '-vcodec', 'rawvideo', '-'] + sample_result = subprocess.run(sample_cmd, capture_output=True) + sample_data = sample_result.stdout + + # Try common video dimensions + common_sizes = [(320, 240), (480, 270), (512, 512), (640, 360), (640, 480), + (720, 405), (720, 480), (854, 480), (1280, 720), (1920, 1080)] + + for w, h in common_sizes: + frame_size = w * h * 3 # RGB24 + if len(sample_data) == frame_size: + width, height = w, h + print(f"Inferred video dimensions from sample: {width}x{height}") + break + + if width is None or height is None: + raise Exception(f"Could not determine video dimensions. Sample data size: {len(sample_data)} bytes") + + # Now extract all frames at 16 fps + cmd = [ + FFMPEG_PATH, '-i', video_path, + '-vf', 'fps=16', # Ensure 16 fps loading + '-f', 'image2pipe', + '-pix_fmt', 'rgb24', + '-vcodec', 'rawvideo', + '-' + ] + + result = subprocess.run(cmd, check=True, capture_output=True) + frame_data = result.stdout + + if not frame_data: + raise Exception("No frame data received from ffmpeg") + + # Calculate frame count + frame_size = width * height * 3 # RGB24 + frame_count = len(frame_data) // frame_size + + if frame_count == 0: + raise Exception(f"No complete frames found in video data. Data size: {len(frame_data)}, Frame size: {frame_size}") + + # Verify frame data integrity + if len(frame_data) % frame_size != 0: + print(f"Warning: Frame data size ({len(frame_data)}) is not perfectly divisible by frame size ({frame_size})") + frame_count = len(frame_data) // frame_size + print(f"Using {frame_count} complete frames, discarding {len(frame_data) % frame_size} bytes") + + # Reshape into numpy array + frames = np.frombuffer(frame_data, dtype=np.uint8).reshape((frame_count, height, width, 3)) + + # Convert to torch tensor (NHWC format as expected by ComfyUI) + images_tensor = torch.from_numpy(frames).float() / 255.0 + + print(f"WanVideoCacheVideos: Successfully loaded {frame_count} frames ({width}x{height}) from {video_path}") + return images_tensor + + except Exception as e: + print(f"WanVideoCacheVideos: Error loading MP4: {e}") + raise e # Update the node mappings to include the new class NODE_CLASS_MAPPINGS = { diff --git a/nodes/wan_first_middle_last.py b/nodes/wan_first_middle_last.py new file mode 100644 index 0000000..4bbcba2 --- /dev/null +++ b/nodes/wan_first_middle_last.py @@ -0,0 +1,294 @@ +# -*- coding: utf-8 -*- + +import torch +import node_helpers +import comfy +import comfy.utils +import comfy.clip_vision +from nodes import MAX_RESOLUTION +from typing import Optional, Tuple, Any + + +class WanFMLF: + """ + 3-frame reference node for Wan2.2 A14B I2V with dual MoE conditioning. + + Features: + - First, middle, and last frame reference + - Dual conditioning outputs for high-noise and low-noise stages + - Adjustable constraint strengths for MoE dual-phase sampling + - Designed for LightX2V distilled model (8 steps: 4 high-noise + 4 low-noise) + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + "vae": ("VAE",), + "width": ("INT", {"default": 832, "min": 16, "max": MAX_RESOLUTION, "step": 16}), + "height": ("INT", {"default": 480, "min": 16, "max": MAX_RESOLUTION, "step": 16}), + "length": ("INT", {"default": 81, "min": 1, "max": MAX_RESOLUTION, "step": 4}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), + }, + "optional": { + "start_image": ("IMAGE",), + "middle_image": ("IMAGE",), + "end_image": ("IMAGE",), + "middle_frame_ratio": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "display": "slider", + }), + # Dual-stage strength control + "high_noise_strength": ("FLOAT", { + "default": 0.8, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "display": "slider", + "tooltip": "Middle frame constraint strength for high-noise stage (determines motion trajectory)" + }), + "low_noise_strength": ("FLOAT", { + "default": 0.2, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "display": "slider", + "tooltip": "Middle frame constraint strength for low-noise stage (prevents detail flickering)" + }), + "motion_amplitude": ("FLOAT", { + "default": 1.15, + "min": 1.0, + "max": 2.0, + "step": 0.05, + "display": "slider", + "tooltip": "Motion amplitude enhancement (fixes 4-step LoRA slow-motion issues)" + }), + "clip_vision_start_image": ("CLIP_VISION_OUTPUT",), + "clip_vision_middle_image": ("CLIP_VISION_OUTPUT",), + "clip_vision_end_image": ("CLIP_VISION_OUTPUT",), + }, + } + + # 🎯 Three outputs: positive high-noise, positive low-noise, latent + # Negative conditioning uses original input (following ComfyUI conventions) + RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "CONDITIONING", "LATENT") + RETURN_NAMES = ("positive_high_noise", "positive_low_noise", "negative", "latent") + FUNCTION = "generate" + CATEGORY = "ComfyUI-Wan22FMLF/video" + + def generate(self, positive: Tuple[Any, ...], + negative: Tuple[Any, ...], + vae: Any, + width: int, + height: int, + length: int, + batch_size: int, + start_image: Optional[torch.Tensor] = None, + middle_image: Optional[torch.Tensor] = None, + end_image: Optional[torch.Tensor] = None, + middle_frame_ratio: float = 0.5, + high_noise_strength: float = 0.8, + low_noise_strength: float = 0.2, + motion_amplitude: float = 1.15, + clip_vision_start_image: Optional[Any] = None, + clip_vision_middle_image: Optional[Any] = None, + clip_vision_end_image: Optional[Any] = None) -> Tuple[Tuple[Any, ...], Tuple[Any, ...], Tuple[Any, ...], dict]: + + spacial_scale = vae.spacial_compression_encode() + latent_channels = vae.latent_channels + latent_t = ((length - 1) // 4) + 1 + + device = comfy.model_management.intermediate_device() + + latent = torch.zeros([batch_size, latent_channels, latent_t, + height // spacial_scale, width // spacial_scale], + device=device) + + # Image preprocessing + if start_image is not None: + start_image = comfy.utils.common_upscale( + start_image[:length].movedim(-1, 1), width, height, + "bilinear", "center").movedim(1, -1) + + if middle_image is not None: + middle_image = comfy.utils.common_upscale( + middle_image[:1].movedim(-1, 1), width, height, + "bilinear", "center").movedim(1, -1) + + if end_image is not None: + end_image = comfy.utils.common_upscale( + end_image[-length:].movedim(-1, 1), width, height, + "bilinear", "center").movedim(1, -1) + + # Create timeline and base mask + image = torch.ones((length, height, width, 3), device=device) * 0.5 + mask_base = torch.ones((1, 1, latent_t * 4, latent.shape[-2], latent.shape[-1]), + device=device) + + def calculate_aligned_position(ratio: float, total_frames: int) -> Tuple[int, int]: + desired_pixel_idx = int(total_frames * ratio) + latent_idx = desired_pixel_idx // 4 + aligned_pixel_idx = latent_idx * 4 + aligned_pixel_idx = max(0, min(aligned_pixel_idx, total_frames - 1)) + return aligned_pixel_idx, latent_idx + + middle_idx, middle_latent_idx = calculate_aligned_position(middle_frame_ratio, length) + middle_idx = max(4, min(middle_idx, length - 5)) + + # 🎯 Key fix: create masks in advance to avoid only creating when middle_image exists + mask_high_noise = mask_base.clone() + mask_low_noise = mask_base.clone() + + # Place reference frames + if start_image is not None: + image[:start_image.shape[0]] = start_image + mask_base[:, :, :start_image.shape[0] + 3] = 0.0 + mask_high_noise[:, :, :start_image.shape[0] + 3] = 0.0 + mask_low_noise[:, :, :start_image.shape[0] + 3] = 0.0 + + if middle_image is not None: + image[middle_idx:middle_idx + 1] = middle_image + + # 🎯 Fix: masks are now created, directly set middle frame masks + start_range = max(0, middle_idx) + end_range = min(length, middle_idx + 4) + high_noise_mask_value = 1.0 - high_noise_strength + mask_high_noise[:, :, start_range:end_range] = high_noise_mask_value + + # Low noise mask (weak constraint) + low_noise_mask_value = 1.0 - low_noise_strength + mask_low_noise[:, :, start_range:end_range] = low_noise_mask_value + + if end_image is not None: + image[-end_image.shape[0]:] = end_image + if middle_image is not None: + mask_high_noise[:, :, -end_image.shape[0]:] = 0.0 + mask_low_noise[:, :, -end_image.shape[0]:] = 0.0 + + # 🎯 Separate high-noise and low-noise latent images + # High noise stage: includes middle frame + concat_latent_image_high = vae.encode(image[:, :, :, :3]) + + # Low noise stage: skip middle frame if strength is 0 + if low_noise_strength == 0.0: + # 🎯 Low noise strength is 0: create latent without middle frame + image_low_only = torch.ones((length, height, width, 3), device=device) * 0.5 + + # Only place start and end frames + if start_image is not None: + image_low_only[:start_image.shape[0]] = start_image + if end_image is not None: + image_low_only[-end_image.shape[0]:] = end_image + + concat_latent_image_low = vae.encode(image_low_only[:, :, :, :3]) + else: + # Low noise strength > 0: use complete image + concat_latent_image_low = vae.encode(image[:, :, :, :3]) + + # Motion amplitude enhancement processing (fixes 4-step LoRA slow-motion issues) + if motion_amplitude > 1.0: + concat_latent_image_high = self._apply_motion_amplitude(concat_latent_image_high, motion_amplitude) + concat_latent_image_low = self._apply_motion_amplitude(concat_latent_image_low, motion_amplitude) + + # Mask reshaping + mask_high_reshaped = mask_high_noise.view(1, mask_high_noise.shape[2] // 4, 4, mask_high_noise.shape[3], mask_high_noise.shape[4]).transpose(1, 2) + mask_low_reshaped = mask_low_noise.view(1, mask_low_noise.shape[2] // 4, 4, mask_low_noise.shape[3], mask_low_noise.shape[4]).transpose(1, 2) + + # 🎯 Create three conditioning settings + # High noise stage: strong constraint, determines motion trajectory + positive_high_noise = node_helpers.conditioning_set_values(positive, { + "concat_latent_image": concat_latent_image_high, + "concat_mask": mask_high_reshaped + }) + + # Low noise stage: decide whether to use middle frame based on strength + positive_low_noise = node_helpers.conditioning_set_values(positive, { + "concat_latent_image": concat_latent_image_low, # 🎯 Separated latent image + "concat_mask": mask_low_reshaped + }) + + # Negative conditioning uses original input + negative_out = negative + + # CLIP Vision processing (mainly for detail optimization in low-noise stage) + clip_vision_output = self._merge_clip_vision_outputs( + clip_vision_start_image, + clip_vision_middle_image, + clip_vision_end_image + ) + + if clip_vision_output is not None: + # Only add CLIP Vision in low-noise stage (better detail understanding) + positive_low_noise = node_helpers.conditioning_set_values(positive_low_noise, + {"clip_vision_output": clip_vision_output}) + + out_latent = {"samples": latent} + + return (positive_high_noise, positive_low_noise, negative_out, out_latent) + + def _apply_motion_amplitude(self, concat_latent_image: torch.Tensor, motion_amplitude: float) -> torch.Tensor: + """ + Apply motion amplitude enhancement (brightness preservation core algorithm) + + Args: + concat_latent_image: VAE-encoded latent image + motion_amplitude: Motion amplitude multiplier (>1.0 enhances motion) + + Returns: + Enhanced latent image + """ + if motion_amplitude <= 1.0: + return concat_latent_image + + # Separate first frame and subsequent frames + base_latent = concat_latent_image[:, :, 0:1] # First frame + gray_latent = concat_latent_image[:, :, 1:] # Subsequent gray frames + + if gray_latent.shape[2] == 0: + return concat_latent_image + + # Calculate motion vectors and separate brightness + diff = gray_latent - base_latent + diff_mean = diff.mean(dim=(1, 3, 4), keepdim=True) # Preserve brightness mean + diff_centered = diff - diff_mean # Pure motion vector with brightness removed + + # Apply motion amplitude enhancement + scaled_latent = base_latent + diff_centered * motion_amplitude + diff_mean + + # Clamp range to prevent anomalies + scaled_latent = torch.clamp(scaled_latent, -6, 6) + + # Reassemble: first frame + enhanced subsequent frames + result = torch.cat([base_latent, scaled_latent], dim=2) + + return result + + def _merge_clip_vision_outputs(self, *outputs: Any) -> Optional[Any]: + valid_outputs = [o for o in outputs if o is not None] + + if not valid_outputs: + return None + + if len(valid_outputs) == 1: + return valid_outputs[0] + + all_states = [o.penultimate_hidden_states for o in valid_outputs] + combined_states = torch.cat(all_states, dim=-2) + + result = comfy.clip_vision.Output() + result.penultimate_hidden_states = combined_states + return result + + +NODE_CLASS_MAPPINGS = { + "WanFMLF": WanFMLF +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "WanFMLF": "WanFMLF (Dual MoE) 🎬" +}