1252 lines
54 KiB
Python
1252 lines
54 KiB
Python
import torch
|
||
import math
|
||
import copy
|
||
import logging
|
||
from typing import Tuple, Dict, Any, List, Optional
|
||
|
||
import comfy.samplers
|
||
import comfy.sample
|
||
import comfy.model_management
|
||
import comfy.latent_formats
|
||
import latent_preview
|
||
import node_helpers
|
||
|
||
# Set up logging for FreeLong
|
||
logger = logging.getLogger("comfyui.longlook")
|
||
logger.setLevel(logging.INFO)
|
||
|
||
|
||
class WanContinuationConditioning:
|
||
"""
|
||
Simple video continuation conditioning for Wan 2.2.
|
||
|
||
Takes the last frame from a previous video chunk and creates i2v conditioning
|
||
for generating the next chunk. Equivalent to WanImageToVideo but designed
|
||
for easy chaining in continuation workflows.
|
||
|
||
Usage: Connect decoded images from chunk 1 → this node → KSampler for chunk 2
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
||
return {
|
||
"required": {
|
||
"positive": ("CONDITIONING",),
|
||
"negative": ("CONDITIONING",),
|
||
"anchor_images": ("IMAGE",),
|
||
"vae": ("VAE",),
|
||
"width": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 16}),
|
||
"height": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 16}),
|
||
"video_length": ("INT", {
|
||
"default": 81,
|
||
"min": 1,
|
||
"max": 1024,
|
||
"step": 4,
|
||
"tooltip": "Output video length in frames. Must match your sampler settings."
|
||
}),
|
||
},
|
||
"optional": {
|
||
"end_images": ("IMAGE",),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
|
||
RETURN_NAMES = ("positive", "negative", "latent")
|
||
FUNCTION = "modify_conditioning"
|
||
CATEGORY = "video/wan"
|
||
DESCRIPTION = "Creates i2v conditioning from last frame with optional end-frame guidance"
|
||
|
||
def modify_conditioning(
|
||
self,
|
||
positive,
|
||
negative,
|
||
anchor_images: torch.Tensor,
|
||
vae,
|
||
width: int,
|
||
height: int,
|
||
video_length: int,
|
||
end_images: Optional[torch.Tensor] = None
|
||
) -> Tuple[list, list, Dict[str, torch.Tensor]]:
|
||
"""
|
||
Modify conditioning to use anchor frames for continuation.
|
||
"""
|
||
import comfy.utils
|
||
|
||
# Calculate latent dimensions using Wan's formula
|
||
latent_frames = ((video_length - 1) // 4) + 1
|
||
latent_h = height // 8
|
||
latent_w = width // 8
|
||
|
||
# Get the LAST frame from anchor images (single frame continuation like standard i2v)
|
||
start_frame = anchor_images[-1:] # Always use just the last frame
|
||
|
||
end_frame = None
|
||
if end_images is not None:
|
||
end_frame = end_images[-1:] # Use last frame as end anchor
|
||
|
||
# Resize to target dimensions if needed
|
||
if start_frame.shape[1] != height or start_frame.shape[2] != width:
|
||
start_frame = comfy.utils.common_upscale(
|
||
start_frame.movedim(-1, 1),
|
||
width, height, "bilinear", "center"
|
||
).movedim(1, -1)
|
||
|
||
if end_frame is not None and (end_frame.shape[1] != height or end_frame.shape[2] != width):
|
||
end_frame = comfy.utils.common_upscale(
|
||
end_frame.movedim(-1, 1),
|
||
width, height, "bilinear", "center"
|
||
).movedim(1, -1)
|
||
if end_frame is not None and (end_frame.device != start_frame.device or end_frame.dtype != start_frame.dtype):
|
||
end_frame = end_frame.to(device=start_frame.device, dtype=start_frame.dtype)
|
||
|
||
# Create full video tensor for VAE encoding (matches WanImageToVideo exactly)
|
||
# Gray (0.5) for frames to generate, actual image for first frame
|
||
full_video = torch.ones(
|
||
video_length, height, width, 3,
|
||
device=start_frame.device, dtype=start_frame.dtype
|
||
) * 0.5
|
||
|
||
# Place the start anchor at the beginning
|
||
full_video[:1] = start_frame[:, :, :, :3]
|
||
|
||
# Optionally place end anchor at the end
|
||
if end_frame is not None:
|
||
full_video[-1:] = end_frame[:, :, :, :3]
|
||
|
||
# VAE encode the entire video (anchor frames + gray neutral frames)
|
||
concat_latent_image = vae.encode(full_video)
|
||
if concat_latent_image.ndim == 4:
|
||
concat_latent_image = concat_latent_image.unsqueeze(0)
|
||
|
||
encoded_latent_frames = concat_latent_image.shape[2]
|
||
latent_channels = concat_latent_image.shape[1]
|
||
|
||
# Calculate how many latent frames to preserve (1 video frame = 1 latent frame)
|
||
# Formula: ((1 - 1) // 4) + 1 = 1
|
||
anchor_latent_frames = 1
|
||
|
||
if end_frame is None:
|
||
# Create concat_mask using LATENT frames (like WanImageToVideo does - simpler, no reshape)
|
||
# Shape: [1, 1, latent_frames, h, w]
|
||
# Wan convention: 0.0 = preserve (don't denoise), 1.0 = generate (full denoise)
|
||
concat_mask = torch.ones(
|
||
1, 1, encoded_latent_frames,
|
||
concat_latent_image.shape[-2],
|
||
concat_latent_image.shape[-1],
|
||
device=start_frame.device, dtype=start_frame.dtype
|
||
)
|
||
# Preserve anchor latent frames
|
||
concat_mask[:, :, :anchor_latent_frames] = 0.0
|
||
else:
|
||
# ComfyCore-style frame mask with end-frame preservation
|
||
# Uses expanded [1, 4, latent_frames, h, w] format for per-subframe control
|
||
# (vs simple [1, 1, latent_frames, h, w] when no end_frame)
|
||
frame_mask = torch.ones(
|
||
1, 1, encoded_latent_frames * 4,
|
||
concat_latent_image.shape[-2],
|
||
concat_latent_image.shape[-1],
|
||
device=start_frame.device, dtype=start_frame.dtype
|
||
)
|
||
frame_mask[:, :, :start_frame.shape[0] + 3] = 0.0
|
||
frame_mask[:, :, -end_frame.shape[0]:] = 0.0
|
||
concat_mask = frame_mask.view(
|
||
1, frame_mask.shape[2] // 4, 4,
|
||
frame_mask.shape[-2], frame_mask.shape[-1]
|
||
).movedim(1, 2)
|
||
|
||
# Move to CPU for conditioning storage
|
||
concat_latent_image = concat_latent_image.cpu()
|
||
concat_mask = concat_mask.cpu()
|
||
|
||
# Modify conditioning with our anchor data
|
||
# This overwrites any existing concat_latent_image/concat_mask from WanImageToVideo
|
||
positive_out = node_helpers.conditioning_set_values(positive, {
|
||
"concat_latent_image": concat_latent_image,
|
||
"concat_mask": concat_mask
|
||
})
|
||
negative_out = node_helpers.conditioning_set_values(negative, {
|
||
"concat_latent_image": concat_latent_image,
|
||
"concat_mask": concat_mask
|
||
})
|
||
|
||
# Create empty latent for sampling
|
||
latent = torch.zeros(
|
||
1, latent_channels, latent_frames, latent_h, latent_w,
|
||
device=comfy.model_management.intermediate_device()
|
||
)
|
||
|
||
return (positive_out, negative_out, {"samples": latent})
|
||
|
||
|
||
# ============================================================================
|
||
# FreeLong Helper Functions
|
||
# ============================================================================
|
||
|
||
def freelong_enforced_blend(
|
||
global_features: torch.Tensor,
|
||
local_features: torch.Tensor,
|
||
motion_lock_ratio: float = 0.15,
|
||
low_freq_ratio: float = 0.5,
|
||
blend_strength: float = 1.0
|
||
) -> torch.Tensor:
|
||
"""
|
||
Enforced spectral blend with 3-tier frequency handling:
|
||
- Ultra-low (0 to motion_lock_ratio): 100% global - LOCKED motion skeleton
|
||
- Mid (motion_lock_ratio to low_freq_ratio): Blended global/local
|
||
- High (low_freq_ratio to 1.0): 100% local - fine details
|
||
|
||
This ensures motion trajectory cannot be overridden by local attention.
|
||
"""
|
||
if blend_strength == 0.0:
|
||
return local_features
|
||
|
||
orig_dtype = global_features.dtype
|
||
global_float = global_features.float()
|
||
local_float = local_features.float()
|
||
|
||
# Single FFT pass
|
||
global_freq = torch.fft.rfft(global_float, dim=1)
|
||
local_freq = torch.fft.rfft(local_float, dim=1)
|
||
del global_float, local_float
|
||
|
||
num_freq = global_freq.shape[1]
|
||
lock_cutoff = int(num_freq * motion_lock_ratio)
|
||
blend_cutoff = int(num_freq * low_freq_ratio)
|
||
|
||
# Build 3-tier mask: 1.0 = global, 0.0 = local
|
||
global_mask = torch.zeros(num_freq, device=global_features.device, dtype=torch.float32)
|
||
|
||
# Tier 1: Ultra-low frequencies - 100% global (locked)
|
||
global_mask[:lock_cutoff] = 1.0
|
||
|
||
# Tier 2: Mid frequencies - smooth transition from global to local
|
||
if blend_cutoff > lock_cutoff:
|
||
mid_len = blend_cutoff - lock_cutoff
|
||
transition = torch.linspace(1, 0, mid_len + 2, device=global_features.device, dtype=torch.float32)[1:-1]
|
||
global_mask[lock_cutoff:blend_cutoff] = transition
|
||
del transition
|
||
|
||
# Tier 3: High frequencies - 100% local (already 0.0)
|
||
|
||
global_mask = global_mask.view(1, -1, *([1] * (len(global_freq.shape) - 2)))
|
||
|
||
# Blend: global * mask + local * (1 - mask)
|
||
blended_freq = global_freq * global_mask + local_freq * (1.0 - global_mask)
|
||
del global_freq, local_freq, global_mask
|
||
|
||
result = torch.fft.irfft(blended_freq, n=global_features.shape[1], dim=1)
|
||
del blended_freq
|
||
|
||
if blend_strength < 1.0:
|
||
result = local_features.float() * (1.0 - blend_strength) + result * blend_strength
|
||
|
||
return result.to(orig_dtype)
|
||
|
||
|
||
def freelong_spectral_blend(
|
||
global_features: torch.Tensor,
|
||
local_features: torch.Tensor,
|
||
low_freq_ratio: float = 0.25,
|
||
blend_strength: float = 1.0
|
||
) -> torch.Tensor:
|
||
"""
|
||
Blend global and local features using spectral decomposition.
|
||
Low frequencies from global (structure/motion direction),
|
||
High frequencies from local (details/sharpness).
|
||
|
||
Memory-optimized: Uses single FFT pass instead of two separate filter calls.
|
||
|
||
Args:
|
||
global_features: Features from full-sequence attention
|
||
local_features: Features from windowed attention
|
||
low_freq_ratio: Ratio of frequencies to take from global (default 0.25)
|
||
blend_strength: How much to apply the blend (0=passthrough, 1=full blend)
|
||
|
||
Returns:
|
||
Blended features
|
||
"""
|
||
if blend_strength == 0.0:
|
||
return local_features
|
||
|
||
# Convert to float32 for FFT (cuFFT doesn't support half precision for non-power-of-two)
|
||
orig_dtype = global_features.dtype
|
||
global_float = global_features.float()
|
||
local_float = local_features.float()
|
||
|
||
# Single FFT pass for both streams
|
||
global_freq = torch.fft.rfft(global_float, dim=1)
|
||
local_freq = torch.fft.rfft(local_float, dim=1)
|
||
del global_float, local_float # Free immediately
|
||
|
||
# Build low-pass mask
|
||
num_freq = global_freq.shape[1]
|
||
cutoff = int(num_freq * low_freq_ratio)
|
||
transition_width = max(1, min(num_freq // 4, cutoff // 2))
|
||
|
||
low_mask = torch.ones(num_freq, device=global_features.device, dtype=torch.float32)
|
||
transition_start = max(0, cutoff - transition_width // 2)
|
||
transition_end = min(num_freq, cutoff + transition_width // 2)
|
||
|
||
if transition_end > transition_start:
|
||
transition_idx = torch.arange(transition_start, transition_end, device=global_features.device, dtype=torch.float32)
|
||
transition_values = 0.5 * (1 + torch.cos(torch.pi * (transition_idx - transition_start) / (transition_end - transition_start)))
|
||
low_mask[transition_start:transition_end] = transition_values
|
||
del transition_idx, transition_values
|
||
|
||
low_mask[transition_end:] = 0.0
|
||
low_mask = low_mask.view(1, -1, *([1] * (len(global_freq.shape) - 2)))
|
||
|
||
# Blend in frequency domain: low from global, high from local
|
||
# high_mask = 1 - low_mask, so: global*low + local*high = global*low + local*(1-low)
|
||
blended_freq = global_freq * low_mask + local_freq * (1.0 - low_mask)
|
||
del global_freq, local_freq, low_mask # Free immediately
|
||
|
||
# Inverse FFT
|
||
result = torch.fft.irfft(blended_freq, n=global_features.shape[1], dim=1)
|
||
del blended_freq
|
||
|
||
# Mix with original based on blend_strength
|
||
if blend_strength < 1.0:
|
||
result = local_features.float() * (1.0 - blend_strength) + result * blend_strength
|
||
|
||
return result.to(orig_dtype)
|
||
|
||
|
||
# ============================================================================
|
||
# FreeLong Model Patcher
|
||
# ============================================================================
|
||
|
||
class WanFreeLong:
|
||
"""
|
||
FreeLong model patcher for Wan 2.2 video generation.
|
||
|
||
Implements SpectralBlend temporal attention to extend video generation
|
||
beyond the training context window while maintaining motion consistency.
|
||
|
||
Based on "FreeLong: Training-Free Long Video Generation with
|
||
SpectralBlend Temporal Attention" (NeurIPS 2024)
|
||
|
||
How it works:
|
||
1. Hooks into Wan's transformer blocks
|
||
2. Runs attention twice: global (full sequence) and local (windowed)
|
||
3. Blends results using FFT: low-freq from global, high-freq from local
|
||
4. This preserves motion direction (low-freq) while keeping details (high-freq)
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
||
return {
|
||
"required": {
|
||
"model": ("MODEL",),
|
||
"enabled": ("BOOLEAN", {
|
||
"default": True,
|
||
"tooltip": "Turn FreeLong on or off. Useful for comparing results with and without."
|
||
}),
|
||
"blend_strength": ("FLOAT", {
|
||
"default": 0.8,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.1,
|
||
"tooltip": "How strongly FreeLong affects the video. Higher = more consistent motion but slightly softer. Lower = sharper but may drift. 0.8 is optimal for most cases."
|
||
}),
|
||
"low_freq_ratio": ("FLOAT", {
|
||
"default": 0.8,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.05,
|
||
"tooltip": "Balance between motion consistency and fine detail. Higher = smoother, more consistent motion. Lower = more dynamic but may drift. 0.8 is a good balance."
|
||
}),
|
||
"local_window_frames": ("INT", {
|
||
"default": 33,
|
||
"min": 9,
|
||
"max": 241,
|
||
"step": 4,
|
||
"tooltip": "How many video frames the detail-preservation looks at. Smaller = sharper details but may cause morphing. Larger = smoother transitions. 33 works well for 81-frame videos. For longer videos, try ~40% of total frames."
|
||
}),
|
||
"blend_start_block": ("INT", {
|
||
"default": 0,
|
||
"min": -1,
|
||
"max": 40,
|
||
"tooltip": "Advanced: Which model layer to start applying FreeLong. 0 = from the beginning. Most users should leave this at 0."
|
||
}),
|
||
"blend_end_block": ("INT", {
|
||
"default": -1,
|
||
"min": -1,
|
||
"max": 40,
|
||
"tooltip": "Advanced: Which model layer to stop applying FreeLong. -1 = apply to all layers. Most users should leave this at -1."
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("MODEL",)
|
||
RETURN_NAMES = ("model",)
|
||
FUNCTION = "patch_model"
|
||
CATEGORY = "video/wan"
|
||
DESCRIPTION = "Apply FreeLong spectral blending for better motion consistency in long videos"
|
||
|
||
def patch_model(
|
||
self,
|
||
model,
|
||
enabled: bool,
|
||
blend_strength: float,
|
||
low_freq_ratio: float,
|
||
local_window_frames: int,
|
||
blend_start_block: int,
|
||
blend_end_block: int
|
||
):
|
||
"""
|
||
Patch the model to use FreeLong spectral blending in temporal attention.
|
||
"""
|
||
# Convert video frames to latent frames
|
||
local_window_size = ((local_window_frames - 1) // 4) + 1
|
||
# If disabled, return model unchanged (no patches)
|
||
if not enabled:
|
||
logger.info("[FreeLong] DISABLED - Model returned without patching")
|
||
return (model,)
|
||
|
||
logger.info("=" * 60)
|
||
logger.info("[FreeLong] APPLYING SPECTRAL BLENDING PATCHES")
|
||
logger.info("=" * 60)
|
||
logger.info(f" Blend Strength: {blend_strength:.2f}")
|
||
logger.info(f" Low Freq Ratio: {low_freq_ratio:.2f}")
|
||
logger.info(f" Local Window: {local_window_frames} video frames → {local_window_size} latent frames")
|
||
logger.info(f" Block Range: {blend_start_block} to {blend_end_block if blend_end_block != -1 else 'ALL'}")
|
||
|
||
model = model.clone()
|
||
|
||
# Store settings for the patch function
|
||
freelong_settings = {
|
||
"blend_strength": blend_strength,
|
||
"low_freq_ratio": low_freq_ratio,
|
||
"local_window_size": local_window_size,
|
||
"blend_start_block": blend_start_block,
|
||
"blend_end_block": blend_end_block,
|
||
"first_run": True, # Track first execution to log details once
|
||
"blocks_patched": set(), # Track which blocks executed
|
||
}
|
||
|
||
def freelong_block_patch(block_args, extra_args):
|
||
"""
|
||
Patch function that wraps transformer blocks with FreeLong blending.
|
||
|
||
Implements dual-stream processing:
|
||
1. Global stream: Full sequence attention (original behavior)
|
||
2. Local stream: Windowed attention (process temporal chunks)
|
||
3. Spectral blend: Low-freq from global, high-freq from local
|
||
"""
|
||
original_block = extra_args["original_block"]
|
||
transformer_options = block_args.get("transformer_options", {})
|
||
block_index = transformer_options.get("block_index", 0)
|
||
total_blocks = transformer_options.get("total_blocks", 40)
|
||
|
||
# Determine if this block should use FreeLong
|
||
start = freelong_settings["blend_start_block"]
|
||
end = freelong_settings["blend_end_block"]
|
||
if end == -1:
|
||
end = total_blocks
|
||
|
||
# For blocks outside our range, just run original
|
||
if block_index < start or block_index >= end:
|
||
return original_block(block_args)
|
||
|
||
# Skip if blend_strength is 0
|
||
strength = freelong_settings["blend_strength"]
|
||
if strength == 0.0:
|
||
return original_block(block_args)
|
||
|
||
# Track that this block executed (for logging summary)
|
||
freelong_settings["blocks_patched"].add(block_index)
|
||
|
||
try:
|
||
x = block_args["img"] # [batch, seq_len, dim]
|
||
batch_size, seq_len, hidden_dim = x.shape
|
||
|
||
window_size = freelong_settings["local_window_size"]
|
||
low_ratio = freelong_settings["low_freq_ratio"]
|
||
|
||
# Log details on first execution
|
||
if freelong_settings["first_run"]:
|
||
freelong_settings["first_run"] = False
|
||
logger.info("")
|
||
logger.info("[FreeLong] DUAL-STREAM PROCESSING STARTING")
|
||
logger.info(f" Input shape: batch={batch_size}, seq_len={seq_len}, hidden_dim={hidden_dim}")
|
||
logger.info(f" Window size: {window_size} latent frames")
|
||
|
||
# ============ GLOBAL STREAM ============
|
||
# Full sequence attention - preserves global motion patterns
|
||
global_out = original_block(block_args)
|
||
global_x = global_out["img"]
|
||
|
||
# ============ LOCAL STREAM ============
|
||
# Process sequence in temporal windows for local detail preservation
|
||
#
|
||
# Wan's sequence is flattened [frames * spatial, dim]
|
||
# We need to estimate the temporal structure
|
||
# For 81 frame video: 21 latent frames, so spatial = seq_len / 21
|
||
|
||
# Use cached values if available (avoids repeated detection across 40 blocks)
|
||
cached_seq_len = freelong_settings.get("cached_seq_len")
|
||
if cached_seq_len == seq_len:
|
||
# Use cached values
|
||
num_frames = freelong_settings["cached_num_frames"]
|
||
spatial_size = freelong_settings["cached_spatial_size"]
|
||
freqs_seq_dim = freelong_settings.get("cached_freqs_seq_dim")
|
||
overlap = freelong_settings["cached_overlap"]
|
||
stride = freelong_settings["cached_stride"]
|
||
else:
|
||
# First time or seq_len changed - detect and cache
|
||
# Try to infer frame count from sequence length
|
||
# Common Wan resolutions and their spatial token counts:
|
||
# 1280x720 -> 160x90/4 = 40x22.5 -> ~900 tokens per frame (after 2x2 patch)
|
||
# 848x480 -> 106x60/4 = 26.5x15 -> ~400 tokens per frame
|
||
# We'll try common latent frame counts: 21, 17, 13, 9, 5
|
||
|
||
num_frames = None
|
||
for candidate_frames in [21, 17, 13, 9, 5]:
|
||
if seq_len % candidate_frames == 0:
|
||
num_frames = candidate_frames
|
||
break
|
||
|
||
if num_frames is None or num_frames <= window_size:
|
||
# Can't determine temporal structure or window covers all frames
|
||
# Fall back to just using global output with spectral smoothing
|
||
if freelong_settings.get("log_fallback", True):
|
||
freelong_settings["log_fallback"] = False
|
||
if num_frames is None:
|
||
logger.warning(f"[FreeLong] Could not determine temporal structure (seq_len={seq_len})")
|
||
else:
|
||
logger.warning(f"[FreeLong] Window size >= frames ({window_size} >= {num_frames})")
|
||
logger.warning("[FreeLong] Using fallback: spectral smoothing without windowing")
|
||
blended_x = freelong_spectral_blend(global_x, global_x, low_ratio, strength)
|
||
return {"img": blended_x}
|
||
|
||
spatial_size = seq_len // num_frames
|
||
overlap = window_size // 2
|
||
stride = window_size - overlap
|
||
|
||
# Detect RoPE sequence dimension
|
||
freqs = block_args.get("pe")
|
||
freqs_seq_dim = None
|
||
if freqs is not None:
|
||
for dim_idx, dim_size in enumerate(freqs.shape):
|
||
if dim_size == seq_len:
|
||
freqs_seq_dim = dim_idx
|
||
break
|
||
|
||
# Cache all computed values
|
||
freelong_settings["cached_seq_len"] = seq_len
|
||
freelong_settings["cached_num_frames"] = num_frames
|
||
freelong_settings["cached_spatial_size"] = spatial_size
|
||
freelong_settings["cached_freqs_seq_dim"] = freqs_seq_dim
|
||
freelong_settings["cached_overlap"] = overlap
|
||
freelong_settings["cached_stride"] = stride
|
||
|
||
# Precompute blend ramps (avoid torch.linspace per window)
|
||
if overlap > 0:
|
||
ramp_up = torch.linspace(0, 1, overlap + 1, device=x.device, dtype=x.dtype)[1:]
|
||
ramp_down = torch.linspace(1, 0, overlap + 1, device=x.device, dtype=x.dtype)[:-1]
|
||
freelong_settings["cached_ramp_up"] = ramp_up
|
||
freelong_settings["cached_ramp_down"] = ramp_down
|
||
|
||
# Log temporal structure detection on first run
|
||
if freelong_settings.get("log_structure", True):
|
||
freelong_settings["log_structure"] = False
|
||
logger.info(f" Detected temporal structure: {num_frames} frames × {spatial_size} spatial tokens")
|
||
|
||
# Calculate and log number of windows
|
||
num_windows = 0
|
||
temp_start = 0
|
||
while temp_start < num_frames:
|
||
num_windows += 1
|
||
temp_start += stride
|
||
if temp_start >= num_frames:
|
||
break
|
||
|
||
if freelong_settings.get("log_windows", True):
|
||
freelong_settings["log_windows"] = False
|
||
logger.info(f" Window processing: {num_windows} windows with {overlap} frame overlap")
|
||
logger.info(f" Overlap strategy: 50% crossfade blending between windows")
|
||
|
||
# Get RoPE embeddings
|
||
freqs = block_args.get("pe")
|
||
|
||
# Reshape to [batch, frames, spatial, dim]
|
||
x_temporal = x.view(batch_size, num_frames, spatial_size, hidden_dim)
|
||
|
||
# Accumulator for blended output and weights
|
||
local_x_temporal = torch.zeros(batch_size, num_frames, spatial_size, hidden_dim,
|
||
device=x.device, dtype=x.dtype)
|
||
weight_accumulator = torch.zeros(batch_size, num_frames, 1, 1,
|
||
device=x.device, dtype=x.dtype)
|
||
|
||
# Get cached ramps if available
|
||
cached_ramp_up = freelong_settings.get("cached_ramp_up")
|
||
cached_ramp_down = freelong_settings.get("cached_ramp_down")
|
||
|
||
# Save original values for in-place modification (avoid dict copy per window)
|
||
original_img = block_args.get("img")
|
||
original_pe = block_args.get("pe")
|
||
|
||
# Process overlapping windows
|
||
start_frame = 0
|
||
while start_frame < num_frames:
|
||
end_frame = min(start_frame + window_size, num_frames)
|
||
window_frames = end_frame - start_frame
|
||
|
||
# Extract window and flatten for block processing
|
||
window_x = x_temporal[:, start_frame:end_frame, :, :]
|
||
window_flat = window_x.reshape(batch_size, window_frames * spatial_size, hidden_dim)
|
||
|
||
# Modify block_args in-place (restore after)
|
||
block_args["img"] = window_flat
|
||
|
||
# Slice RoPE embeddings to match window
|
||
if freqs is not None and freqs_seq_dim is not None:
|
||
start_token = start_frame * spatial_size
|
||
end_token = end_frame * spatial_size
|
||
slices = [slice(None)] * len(freqs.shape)
|
||
slices[freqs_seq_dim] = slice(start_token, end_token)
|
||
block_args["pe"] = freqs[tuple(slices)]
|
||
|
||
# Run block on window
|
||
window_out = original_block(block_args)
|
||
window_result = window_out["img"]
|
||
|
||
# Reshape back to temporal
|
||
window_result = window_result.view(batch_size, window_frames, spatial_size, hidden_dim)
|
||
|
||
# Create blending weights using cached ramps
|
||
# This creates smooth crossfades in overlapping regions
|
||
weights = torch.ones(window_frames, device=x.device, dtype=x.dtype)
|
||
|
||
if start_frame > 0 and overlap > 0 and cached_ramp_up is not None:
|
||
# Ramp up at the start (we're overlapping with previous window)
|
||
ramp_len = min(overlap, window_frames)
|
||
weights[:ramp_len] = cached_ramp_up[:ramp_len]
|
||
|
||
if end_frame < num_frames and overlap > 0 and cached_ramp_down is not None:
|
||
# Ramp down at the end (next window will overlap with us)
|
||
ramp_len = min(overlap, window_frames)
|
||
weights[-ramp_len:] = weights[-ramp_len:] * cached_ramp_down[:ramp_len]
|
||
|
||
weights = weights.view(1, window_frames, 1, 1)
|
||
|
||
# Accumulate weighted results (in-place for memory efficiency)
|
||
weighted_result = window_result * weights
|
||
local_x_temporal[:, start_frame:end_frame].add_(weighted_result)
|
||
weight_accumulator[:, start_frame:end_frame].add_(weights)
|
||
|
||
# Free window tensors immediately
|
||
del window_x, window_flat, window_out
|
||
del window_result, weights, weighted_result
|
||
|
||
# Move to next window position
|
||
start_frame += stride
|
||
if start_frame >= num_frames:
|
||
break
|
||
|
||
# Restore original block_args
|
||
block_args["img"] = original_img
|
||
if original_pe is not None:
|
||
block_args["pe"] = original_pe
|
||
|
||
# Normalize by accumulated weights
|
||
weight_accumulator = weight_accumulator.clamp(min=1e-8)
|
||
local_x_temporal.div_(weight_accumulator) # In-place division
|
||
del weight_accumulator
|
||
local_x = local_x_temporal.reshape(batch_size, seq_len, hidden_dim)
|
||
del local_x_temporal, x_temporal
|
||
|
||
# ============ SPECTRAL BLEND ============
|
||
# Low frequencies from global (motion direction/consistency)
|
||
# High frequencies from local (sharp details)
|
||
if freelong_settings.get("log_blend", True):
|
||
freelong_settings["log_blend"] = False
|
||
logger.info("")
|
||
logger.info("[FreeLong] SPECTRAL BLENDING")
|
||
logger.info(f" Low-freq (global): {low_ratio:.0%} of spectrum (motion consistency)")
|
||
logger.info(f" High-freq (local): {1-low_ratio:.0%} of spectrum (sharp details)")
|
||
logger.info(f" Blend strength: {strength:.0%}")
|
||
|
||
blended_x = freelong_spectral_blend(global_x, local_x, low_ratio, strength)
|
||
|
||
# Free large intermediate tensors
|
||
del global_x, local_x
|
||
|
||
return {"img": blended_x}
|
||
|
||
except Exception as e:
|
||
# If anything fails, fall back to original block
|
||
logger.error(f"[FreeLong] Error in block {block_index}: {e}", exc_info=True)
|
||
return original_block(block_args)
|
||
|
||
# Wan 2.2 typically has 40 blocks
|
||
num_blocks = 40
|
||
|
||
# Register the patch for transformer blocks
|
||
for i in range(num_blocks):
|
||
model.set_model_patch_replace(
|
||
freelong_block_patch,
|
||
"dit",
|
||
"double_block",
|
||
i
|
||
)
|
||
|
||
actual_end = blend_end_block if blend_end_block != -1 else num_blocks
|
||
blocks_to_patch = actual_end - blend_start_block
|
||
logger.info(f" Registered patches for {blocks_to_patch} transformer blocks ({blend_start_block} to {actual_end-1})")
|
||
logger.info("=" * 60)
|
||
logger.info("[FreeLong] Patching complete. FreeLong will activate during sampling.")
|
||
logger.info("=" * 60)
|
||
|
||
return (model,)
|
||
|
||
|
||
# ============================================================================
|
||
# FreeLong Enforcer - Stricter Motion Locking
|
||
# ============================================================================
|
||
|
||
class WanFreeLongEnforcer:
|
||
"""
|
||
FreeLong Enforcer: Stricter motion consistency for Wan 2.2 video generation.
|
||
|
||
An enhanced version of FreeLong with stronger motion locking:
|
||
1. Early blocks use global-only attention (faster, locks motion early)
|
||
2. 3-tier frequency blending: locked ultra-low, blended mid, local high
|
||
3. Motion skeleton is protected from local attention override
|
||
|
||
Use this when standard FreeLong still shows motion drift or direction reversals.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
||
return {
|
||
"required": {
|
||
"model": ("MODEL",),
|
||
"enabled": ("BOOLEAN", {
|
||
"default": True,
|
||
"tooltip": "Turn FreeLong Enforcer on or off."
|
||
}),
|
||
"motion_lock_ratio": ("FLOAT", {
|
||
"default": 0.15,
|
||
"min": 0.0,
|
||
"max": 0.5,
|
||
"step": 0.05,
|
||
"tooltip": "Ultra-low frequencies that are 100% locked to global. Higher = stronger motion lock but may reduce dynamism. 0.15 is a good starting point."
|
||
}),
|
||
"blend_strength": ("FLOAT", {
|
||
"default": 0.8,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.1,
|
||
"tooltip": "How strongly the enforcer affects the video. Higher = more consistent but slightly softer."
|
||
}),
|
||
"low_freq_ratio": ("FLOAT", {
|
||
"default": 0.5,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.05,
|
||
"tooltip": "Upper bound for global influence (above motion_lock_ratio, frequencies are blended)."
|
||
}),
|
||
"local_window_frames": ("INT", {
|
||
"default": 33,
|
||
"min": 9,
|
||
"max": 241,
|
||
"step": 4,
|
||
"tooltip": "Window size for local attention. Smaller = sharper details, larger = smoother."
|
||
}),
|
||
"motion_lock_blocks": ("INT", {
|
||
"default": 5,
|
||
"min": 0,
|
||
"max": 20,
|
||
"step": 1,
|
||
"tooltip": "First N blocks use global-only (no local processing). Establishes motion early. 5 is a good default. 0 disables this feature."
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("MODEL",)
|
||
RETURN_NAMES = ("model",)
|
||
FUNCTION = "patch_model"
|
||
CATEGORY = "video/wan"
|
||
DESCRIPTION = "Stricter FreeLong with locked motion skeleton - use when standard FreeLong shows drift"
|
||
|
||
def patch_model(
|
||
self,
|
||
model,
|
||
enabled: bool,
|
||
motion_lock_ratio: float,
|
||
blend_strength: float,
|
||
low_freq_ratio: float,
|
||
local_window_frames: int,
|
||
motion_lock_blocks: int
|
||
):
|
||
"""
|
||
Patch the model with enforced motion locking.
|
||
"""
|
||
local_window_size = ((local_window_frames - 1) // 4) + 1
|
||
|
||
if not enabled:
|
||
logger.info("[FreeLong Enforcer] DISABLED - Model returned without patching")
|
||
return (model,)
|
||
|
||
logger.info("=" * 60)
|
||
logger.info("[FreeLong Enforcer] APPLYING ENFORCED MOTION LOCKING")
|
||
logger.info("=" * 60)
|
||
logger.info(f" Motion Lock Ratio: {motion_lock_ratio:.0%} (ultra-low frequencies locked)")
|
||
logger.info(f" Blend Ratio: {motion_lock_ratio:.0%} - {low_freq_ratio:.0%} (mid frequencies blended)")
|
||
logger.info(f" Local Ratio: {low_freq_ratio:.0%} - 100% (high frequencies from local)")
|
||
logger.info(f" Motion Lock Blocks: first {motion_lock_blocks} blocks (global-only)")
|
||
logger.info(f" Blend Strength: {blend_strength:.2f}")
|
||
logger.info(f" Local Window: {local_window_frames} video frames → {local_window_size} latent frames")
|
||
|
||
model = model.clone()
|
||
|
||
enforcer_settings = {
|
||
"motion_lock_ratio": motion_lock_ratio,
|
||
"low_freq_ratio": low_freq_ratio,
|
||
"blend_strength": blend_strength,
|
||
"local_window_size": local_window_size,
|
||
"motion_lock_blocks": motion_lock_blocks,
|
||
"first_run": True,
|
||
}
|
||
|
||
def enforcer_block_patch(block_args, extra_args):
|
||
"""
|
||
Enforcer patch with early-block global-only and 3-tier spectral blend.
|
||
"""
|
||
original_block = extra_args["original_block"]
|
||
transformer_options = block_args.get("transformer_options", {})
|
||
block_index = transformer_options.get("block_index", 0)
|
||
|
||
strength = enforcer_settings["blend_strength"]
|
||
if strength == 0.0:
|
||
return original_block(block_args)
|
||
|
||
try:
|
||
x = block_args["img"]
|
||
batch_size, seq_len, hidden_dim = x.shape
|
||
|
||
window_size = enforcer_settings["local_window_size"]
|
||
lock_ratio = enforcer_settings["motion_lock_ratio"]
|
||
low_ratio = enforcer_settings["low_freq_ratio"]
|
||
lock_blocks = enforcer_settings["motion_lock_blocks"]
|
||
|
||
if enforcer_settings["first_run"]:
|
||
enforcer_settings["first_run"] = False
|
||
logger.info("")
|
||
logger.info("[FreeLong Enforcer] PROCESSING STARTING")
|
||
logger.info(f" Input shape: batch={batch_size}, seq_len={seq_len}, hidden_dim={hidden_dim}")
|
||
|
||
# ============ GLOBAL STREAM ============
|
||
global_out = original_block(block_args)
|
||
global_x = global_out["img"]
|
||
|
||
# Early blocks: global-only (faster, establishes motion)
|
||
if block_index < lock_blocks:
|
||
if enforcer_settings.get("log_lock_blocks", True) and block_index == 0:
|
||
enforcer_settings["log_lock_blocks"] = False
|
||
logger.info(f" Blocks 0-{lock_blocks-1}: Global-only (motion locking phase)")
|
||
return {"img": global_x}
|
||
|
||
# ============ LOCAL STREAM (same as FreeLong) ============
|
||
cached_seq_len = enforcer_settings.get("cached_seq_len")
|
||
if cached_seq_len == seq_len:
|
||
num_frames = enforcer_settings["cached_num_frames"]
|
||
spatial_size = enforcer_settings["cached_spatial_size"]
|
||
freqs_seq_dim = enforcer_settings.get("cached_freqs_seq_dim")
|
||
overlap = enforcer_settings["cached_overlap"]
|
||
stride = enforcer_settings["cached_stride"]
|
||
else:
|
||
num_frames = None
|
||
for candidate_frames in [21, 17, 13, 9, 5]:
|
||
if seq_len % candidate_frames == 0:
|
||
num_frames = candidate_frames
|
||
break
|
||
|
||
if num_frames is None or num_frames <= window_size:
|
||
# Fallback: apply enforced blend to global only
|
||
blended_x = freelong_enforced_blend(global_x, global_x, lock_ratio, low_ratio, strength)
|
||
return {"img": blended_x}
|
||
|
||
spatial_size = seq_len // num_frames
|
||
overlap = window_size // 2
|
||
stride = window_size - overlap
|
||
|
||
freqs = block_args.get("pe")
|
||
freqs_seq_dim = None
|
||
if freqs is not None:
|
||
for dim_idx, dim_size in enumerate(freqs.shape):
|
||
if dim_size == seq_len:
|
||
freqs_seq_dim = dim_idx
|
||
break
|
||
|
||
enforcer_settings["cached_seq_len"] = seq_len
|
||
enforcer_settings["cached_num_frames"] = num_frames
|
||
enforcer_settings["cached_spatial_size"] = spatial_size
|
||
enforcer_settings["cached_freqs_seq_dim"] = freqs_seq_dim
|
||
enforcer_settings["cached_overlap"] = overlap
|
||
enforcer_settings["cached_stride"] = stride
|
||
|
||
if overlap > 0:
|
||
ramp_up = torch.linspace(0, 1, overlap + 1, device=x.device, dtype=x.dtype)[1:]
|
||
ramp_down = torch.linspace(1, 0, overlap + 1, device=x.device, dtype=x.dtype)[:-1]
|
||
enforcer_settings["cached_ramp_up"] = ramp_up
|
||
enforcer_settings["cached_ramp_down"] = ramp_down
|
||
|
||
if enforcer_settings.get("log_structure", True):
|
||
enforcer_settings["log_structure"] = False
|
||
logger.info(f" Blocks {lock_blocks}+: Dual-stream with enforced 3-tier blend")
|
||
logger.info(f" Detected: {num_frames} frames × {spatial_size} spatial tokens")
|
||
|
||
freqs = block_args.get("pe")
|
||
x_temporal = x.view(batch_size, num_frames, spatial_size, hidden_dim)
|
||
|
||
local_x_temporal = torch.zeros(batch_size, num_frames, spatial_size, hidden_dim,
|
||
device=x.device, dtype=x.dtype)
|
||
weight_accumulator = torch.zeros(batch_size, num_frames, 1, 1,
|
||
device=x.device, dtype=x.dtype)
|
||
|
||
cached_ramp_up = enforcer_settings.get("cached_ramp_up")
|
||
cached_ramp_down = enforcer_settings.get("cached_ramp_down")
|
||
|
||
original_img = block_args.get("img")
|
||
original_pe = block_args.get("pe")
|
||
|
||
start_frame = 0
|
||
while start_frame < num_frames:
|
||
end_frame = min(start_frame + window_size, num_frames)
|
||
window_frames = end_frame - start_frame
|
||
|
||
window_x = x_temporal[:, start_frame:end_frame, :, :]
|
||
window_flat = window_x.reshape(batch_size, window_frames * spatial_size, hidden_dim)
|
||
|
||
block_args["img"] = window_flat
|
||
|
||
if freqs is not None and freqs_seq_dim is not None:
|
||
start_token = start_frame * spatial_size
|
||
end_token = end_frame * spatial_size
|
||
slices = [slice(None)] * len(freqs.shape)
|
||
slices[freqs_seq_dim] = slice(start_token, end_token)
|
||
block_args["pe"] = freqs[tuple(slices)]
|
||
|
||
window_out = original_block(block_args)
|
||
window_result = window_out["img"]
|
||
window_result = window_result.view(batch_size, window_frames, spatial_size, hidden_dim)
|
||
|
||
weights = torch.ones(window_frames, device=x.device, dtype=x.dtype)
|
||
|
||
if start_frame > 0 and overlap > 0 and cached_ramp_up is not None:
|
||
ramp_len = min(overlap, window_frames)
|
||
weights[:ramp_len] = cached_ramp_up[:ramp_len]
|
||
|
||
if end_frame < num_frames and overlap > 0 and cached_ramp_down is not None:
|
||
ramp_len = min(overlap, window_frames)
|
||
weights[-ramp_len:] = weights[-ramp_len:] * cached_ramp_down[:ramp_len]
|
||
|
||
weights = weights.view(1, window_frames, 1, 1)
|
||
|
||
weighted_result = window_result * weights
|
||
local_x_temporal[:, start_frame:end_frame].add_(weighted_result)
|
||
weight_accumulator[:, start_frame:end_frame].add_(weights)
|
||
|
||
del window_x, window_flat, window_out
|
||
del window_result, weights, weighted_result
|
||
|
||
start_frame += stride
|
||
if start_frame >= num_frames:
|
||
break
|
||
|
||
block_args["img"] = original_img
|
||
if original_pe is not None:
|
||
block_args["pe"] = original_pe
|
||
|
||
weight_accumulator = weight_accumulator.clamp(min=1e-8)
|
||
local_x_temporal.div_(weight_accumulator)
|
||
del weight_accumulator
|
||
local_x = local_x_temporal.reshape(batch_size, seq_len, hidden_dim)
|
||
del local_x_temporal, x_temporal
|
||
|
||
# ============ ENFORCED 3-TIER BLEND ============
|
||
blended_x = freelong_enforced_blend(global_x, local_x, lock_ratio, low_ratio, strength)
|
||
|
||
del global_x, local_x
|
||
|
||
return {"img": blended_x}
|
||
|
||
except Exception as e:
|
||
logger.error(f"[FreeLong Enforcer] Error in block {block_index}: {e}", exc_info=True)
|
||
return original_block(block_args)
|
||
|
||
num_blocks = 40
|
||
for i in range(num_blocks):
|
||
model.set_model_patch_replace(
|
||
enforcer_block_patch,
|
||
"dit",
|
||
"double_block",
|
||
i
|
||
)
|
||
|
||
logger.info(f" Registered enforcer patches for {num_blocks} transformer blocks")
|
||
logger.info("=" * 60)
|
||
logger.info("[FreeLong Enforcer] Patching complete.")
|
||
logger.info("=" * 60)
|
||
|
||
return (model,)
|
||
|
||
|
||
# ============================================================================
|
||
# WanMotionScale - Temporal RoPE Scaling
|
||
# ============================================================================
|
||
|
||
class WanMotionScale:
|
||
"""
|
||
Scale temporal position embeddings to control motion speed in Wan video generation.
|
||
|
||
Uses ComfyUI's built-in rope_options to scale temporal positions at the source,
|
||
making the model "think" frames are further apart in time.
|
||
|
||
Use case: Generate with scale_t > 1, then RIFE interpolate to fill gaps,
|
||
effectively getting longer videos with more motion coverage.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
||
return {
|
||
"required": {
|
||
"model": ("MODEL",),
|
||
"enabled": ("BOOLEAN", {
|
||
"default": True,
|
||
"tooltip": "Enable/disable motion scaling."
|
||
}),
|
||
"scale_t": ("FLOAT", {
|
||
"default": 1.5,
|
||
"min": -10.0,
|
||
"max": 10.0,
|
||
"step": 0.05,
|
||
"tooltip": "Temporal scale. >1 = faster motion, <1 = slower, <0 = reversed time positions. Experiment freely!"
|
||
}),
|
||
},
|
||
"optional": {
|
||
"scale_y": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": -10.0,
|
||
"max": 10.0,
|
||
"step": 0.05,
|
||
"tooltip": "Height/Y scale. Affects vertical spatial encoding."
|
||
}),
|
||
"scale_x": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": -10.0,
|
||
"max": 10.0,
|
||
"step": 0.05,
|
||
"tooltip": "Width/X scale. Affects horizontal spatial encoding."
|
||
}),
|
||
# shift_t, shift_y, shift_x support exists in backend but hidden from UI
|
||
# (RoPE uses relative positions, so shifts have minimal effect)
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("MODEL",)
|
||
RETURN_NAMES = ("model",)
|
||
FUNCTION = "patch_model"
|
||
CATEGORY = "video/wan"
|
||
DESCRIPTION = "Scale temporal positions to control motion speed - use with RIFE for longer videos"
|
||
|
||
def patch_model(
|
||
self,
|
||
model,
|
||
enabled: bool,
|
||
scale_t: float,
|
||
scale_y: float = 1.0,
|
||
scale_x: float = 1.0,
|
||
shift_t: float = 0.0,
|
||
shift_y: float = 0.0,
|
||
shift_x: float = 0.0,
|
||
):
|
||
"""
|
||
Patch the model to scale/shift RoPE positions via transformer_options.
|
||
"""
|
||
if not enabled:
|
||
logger.info("[WanMotionScale] DISABLED - Model returned without patching")
|
||
return (model,)
|
||
|
||
# Check if anything is actually being changed
|
||
no_scale = (scale_t == 1.0 and scale_y == 1.0 and scale_x == 1.0)
|
||
no_shift = (shift_t == 0.0 and shift_y == 0.0 and shift_x == 0.0)
|
||
if no_scale and no_shift:
|
||
logger.info("[WanMotionScale] All values at default - Model returned without patching")
|
||
return (model,)
|
||
|
||
logger.info("=" * 60)
|
||
logger.info("[WanMotionScale] TEMPORAL ROPE SCALING/SHIFTING")
|
||
logger.info("=" * 60)
|
||
if scale_t != 1.0:
|
||
logger.info(f" Temporal Scale (scale_t): {scale_t:.2f}x")
|
||
if scale_y != 1.0:
|
||
logger.info(f" Height Scale (scale_y): {scale_y:.2f}x")
|
||
if scale_x != 1.0:
|
||
logger.info(f" Width Scale (scale_x): {scale_x:.2f}x")
|
||
if shift_t != 0.0:
|
||
logger.info(f" Temporal Shift (shift_t): {shift_t:.1f}")
|
||
if shift_y != 0.0:
|
||
logger.info(f" Height Shift (shift_y): {shift_y:.1f}")
|
||
if shift_x != 0.0:
|
||
logger.info(f" Width Shift (shift_x): {shift_x:.1f}")
|
||
logger.info("=" * 60)
|
||
|
||
model = model.clone()
|
||
|
||
# Set rope_options in model_options which gets passed to transformer_options
|
||
rope_options = {
|
||
"scale_t": scale_t,
|
||
"scale_y": scale_y,
|
||
"scale_x": scale_x,
|
||
"shift_t": shift_t,
|
||
"shift_y": shift_y,
|
||
"shift_x": shift_x,
|
||
}
|
||
|
||
# Add rope_options to the model's transformer options
|
||
model.model_options.setdefault("transformer_options", {})["rope_options"] = rope_options
|
||
|
||
logger.info("[WanMotionScale] rope_options added to transformer_options")
|
||
logger.info("=" * 60)
|
||
|
||
return (model,)
|
||
|
||
|
||
# ============================================================================
|
||
# WanMotionScaleAdvanced - Experimental RoPE controls including theta
|
||
# ============================================================================
|
||
|
||
class WanMotionScaleAdvanced:
|
||
"""
|
||
Advanced motion scaling with theta control for Wan video generation.
|
||
|
||
Extends the basic Motion Scale with access to RoPE theta parameter,
|
||
which controls the base frequency of position embeddings.
|
||
|
||
WARNING: This is experimental. Theta modification changes fundamental
|
||
model behavior and may produce unexpected results.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
||
return {
|
||
"required": {
|
||
"model": ("MODEL",),
|
||
"enabled": ("BOOLEAN", {
|
||
"default": True,
|
||
"tooltip": "Enable/disable all modifications."
|
||
}),
|
||
},
|
||
"optional": {
|
||
"scale_t": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": -10.0,
|
||
"max": 10.0,
|
||
"step": 0.05,
|
||
"tooltip": "Temporal scale. >1 = faster motion, <1 = slower. Same as basic Motion Scale."
|
||
}),
|
||
"scale_y": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": -10.0,
|
||
"max": 10.0,
|
||
"step": 0.05,
|
||
"tooltip": "Height/Y position scale."
|
||
}),
|
||
"scale_x": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": -10.0,
|
||
"max": 10.0,
|
||
"step": 0.05,
|
||
"tooltip": "Width/X position scale."
|
||
}),
|
||
"theta": ("FLOAT", {
|
||
"default": 10000.0,
|
||
"min": 100.0,
|
||
"max": 1000000.0,
|
||
"step": 100.0,
|
||
"tooltip": "RoPE base frequency (default 10000 = Wan's training). Higher theta (20000-50000) = longer effective context, frames appear 'closer together' - may help >81 frame generations stay coherent. Lower theta (2000-5000) = frames seem 'further apart' temporally, may increase motion intensity or cause instability. EXPERIMENTAL - start with small changes."
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("MODEL",)
|
||
RETURN_NAMES = ("model",)
|
||
FUNCTION = "patch_model"
|
||
CATEGORY = "video/wan"
|
||
DESCRIPTION = "Advanced motion control with theta - EXPERIMENTAL"
|
||
|
||
def patch_model(
|
||
self,
|
||
model,
|
||
enabled: bool,
|
||
scale_t: float = 1.0,
|
||
scale_y: float = 1.0,
|
||
scale_x: float = 1.0,
|
||
theta: float = 10000.0,
|
||
):
|
||
"""
|
||
Patch the model with advanced RoPE modifications including theta.
|
||
"""
|
||
if not enabled:
|
||
logger.info("[WanMotionScaleAdvanced] DISABLED - Model returned without patching")
|
||
return (model,)
|
||
|
||
# Check if anything is actually being changed
|
||
no_scale = (scale_t == 1.0 and scale_y == 1.0 and scale_x == 1.0)
|
||
default_theta = (theta == 10000.0)
|
||
|
||
if no_scale and default_theta:
|
||
logger.info("[WanMotionScaleAdvanced] All values at default - Model returned without patching")
|
||
return (model,)
|
||
|
||
logger.info("=" * 60)
|
||
logger.info("[WanMotionScaleAdvanced] ADVANCED ROPE MODIFICATIONS")
|
||
logger.info("=" * 60)
|
||
if scale_t != 1.0:
|
||
logger.info(f" Temporal Scale (scale_t): {scale_t:.2f}x")
|
||
if scale_y != 1.0:
|
||
logger.info(f" Height Scale (scale_y): {scale_y:.2f}x")
|
||
if scale_x != 1.0:
|
||
logger.info(f" Width Scale (scale_x): {scale_x:.2f}x")
|
||
if not default_theta:
|
||
logger.info(f" Theta: {theta:.0f} (default: 10000)")
|
||
if theta > 10000:
|
||
logger.info(f" → Higher theta = longer effective context window")
|
||
else:
|
||
logger.info(f" → Lower theta = positions cycle faster")
|
||
logger.info("=" * 60)
|
||
|
||
model = model.clone()
|
||
|
||
# Apply scale options via rope_options (same mechanism as basic Motion Scale)
|
||
if not no_scale:
|
||
rope_options = {
|
||
"scale_t": scale_t,
|
||
"scale_y": scale_y,
|
||
"scale_x": scale_x,
|
||
}
|
||
model.model_options.setdefault("transformer_options", {})["rope_options"] = rope_options
|
||
logger.info("[WanMotionScaleAdvanced] rope_options (scales) added to transformer_options")
|
||
|
||
# Apply theta modification directly to the rope_embedder
|
||
if not default_theta:
|
||
try:
|
||
# Access the diffusion model's rope_embedder
|
||
diffusion_model = model.model.diffusion_model
|
||
if hasattr(diffusion_model, 'rope_embedder'):
|
||
original_theta = diffusion_model.rope_embedder.theta
|
||
diffusion_model.rope_embedder.theta = theta
|
||
logger.info(f"[WanMotionScaleAdvanced] rope_embedder.theta: {original_theta} → {theta}")
|
||
else:
|
||
logger.warning("[WanMotionScaleAdvanced] Model does not have rope_embedder - theta not applied")
|
||
except Exception as e:
|
||
logger.error(f"[WanMotionScaleAdvanced] Failed to set theta: {e}")
|
||
|
||
logger.info("=" * 60)
|
||
|
||
return (model,)
|