Major QoL + Fixes Update 1.08

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "WanVideoLooper"
 description = "This is a set of custom nodes for ComfyUI designed for generating sequential long video clips based on the Wan 2.2 model architecture, handling continuity between segments and offering advanced control over the sampling process"
-version = "1.0.7"
+version = "1.0.8"
 license = {file = "LICENSE"}
 dependencies = [
     "color-match",

wan_looper_node.py

@@ -23,10 +23,7 @@ except ImportError:
 # ====================================================================================================
 def _log(msg: str):
     """Simple print logger for the console."""
-    try:
-        print(f"[WanVideoLooper] {msg}", flush=True)
-    except Exception:
-        pass
+    print(f"[WanVideoLooper] {msg}", flush=True)
 
 if not COLOR_MATCHER_AVAILABLE:
     _log("#####################################################################")
@@ -43,9 +40,11 @@ if not COLOR_MATCHER_AVAILABLE:
 def _align_to_multiple(x, multiple=16): return max(multiple, (int(x) // multiple) * multiple)
 def _get_model_device_info(model):
     try:
-        model_device = next(model.model.parameters()).device; model_dtype = next(model.model.parameters()).dtype
+        model_device = next(model.model.parameters()).device
+        model_dtype = next(model.model.parameters()).dtype
     except Exception:
-        model_device = comfy.model_management.intermediate_device(); model_dtype = torch.float32
+        model_device = comfy.model_management.intermediate_device()
+        model_dtype = torch.float32
     return model_device, model_dtype
 def _move_tensor(tensor: torch.Tensor, device, dtype):
     if tensor.device != device or tensor.dtype != dtype: return tensor.to(device=device, dtype=dtype, non_blocking=True)
@@ -86,14 +85,22 @@ def _encode_seed_image_vision(clip_vision, seed_image, cropping="center"):
     if image_tensor.dim() != 4 or image_tensor.shape[-1] not in (1, 3, 4): _log(f"Warning: Invalid image shape for CLIP Vision: {tuple(image_tensor.shape)}"); return None
     vision_output = clip_vision.encode_image(image_tensor, crop=(cropping == "center")); return vision_output
 def _tokenize_and_encode(clip, positive_prompt, negative_prompt):
-    pos_tokens = clip.tokenize(positive_prompt or ""); pos_data = clip.encode_from_tokens(pos_tokens, return_pooled=True, return_dict=True)
+    pos_tokens = clip.tokenize(positive_prompt or "")
+    pos_data = clip.encode_from_tokens(pos_tokens, return_pooled=True, return_dict=True)
     positive_cond = [(pos_data["cond"], {"pooled_output": pos_data.get("pooled_output")})]
-    neg_tokens = clip.tokenize(negative_prompt or ""); neg_data = clip.encode_from_tokens(neg_tokens, return_pooled=True, return_dict=True)
-    negative_cond = [(neg_data["cond"], {"pooled_output": neg_data.get("pooled_output")})]; return positive_cond, negative_cond
-def _prepare_latent_window(positive_cond, negative_cond, vae, width, height, frame_count, batch_size, seed_image=None, vision_output=None):
+    neg_tokens = clip.tokenize(negative_prompt or "")
+    neg_data = clip.encode_from_tokens(neg_tokens, return_pooled=True, return_dict=True)
+    negative_cond = [(neg_data["cond"], {"pooled_output": neg_data.get("pooled_output")})]
+    return positive_cond, negative_cond
+def _prepare_latent_window(positive_cond, negative_cond, vae, width, height, frame_count, batch_size, seed_image=None, vision_output=None, pre_allocated_latent=None):
     device = comfy.model_management.intermediate_device()
-    latent_height = height // 8; latent_width = width // 8; time_dim = ((frame_count - 1) // 4) + 1
-    empty_latent = torch.zeros([batch_size, 16, time_dim, latent_height, latent_width], device=device)
+    latent_height = height // 8; latent_width = width // 8; time_dim = (frame_count + 3) // 4
+    target_shape = [batch_size, 16, time_dim, latent_height, latent_width]
+    
+    if pre_allocated_latent is not None and pre_allocated_latent.shape == torch.Size(target_shape) and pre_allocated_latent.device == device:
+        empty_latent = pre_allocated_latent.zero_()
+    else:
+        empty_latent = torch.zeros(target_shape, device=device)
     if seed_image is not None:
         seed_image_upscaled = comfy.utils.common_upscale(seed_image[:frame_count].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
         image_padded = torch.ones((frame_count, height, width, seed_image_upscaled.shape[-1]), device=seed_image_upscaled.device, dtype=seed_image_upscaled.dtype) * 0.5
@@ -110,11 +117,164 @@ def _prepare_latent_window(positive_cond, negative_cond, vae, width, height, fra
         negative_cond = node_helpers.conditioning_set_values(negative_cond, {"clip_vision_output": vision_output})
     return positive_cond, negative_cond, {"samples": empty_latent}
 
+# ====================================================================================================
+# Color Match Helper (GPU Accelerated)
+# ====================================================================================================
+def _calculate_covariance(x):
+    """Calculate the covariance matrix of a tensor."""
+    x_mean = x.mean(dim=0)
+    x_centered = x - x_mean
+    return torch.matmul(x_centered.T, x_centered) / (x.shape[0] - 1)
+
+def _rgb_to_xyz(rgb):
+    """Converts a PyTorch tensor from RGB to XYZ color space."""
+    M = torch.tensor([
+        [0.412453, 0.357580, 0.180423],
+        [0.212671, 0.715160, 0.072169],
+        [0.019334, 0.119193, 0.950227]
+    ], device=rgb.device, dtype=rgb.dtype)
+    
+    if rgb.dim() == 3:
+        h, w = rgb.shape[1:]
+        xyz = torch.matmul(M, rgb.view(3, -1)).view(3, h, w)
+    elif rgb.dim() == 4:
+        b, h, w = rgb.shape[0], rgb.shape[2], rgb.shape[3]
+        xyz = torch.matmul(M, rgb.view(b, 3, -1)).view(b, 3, h, w)
+    else:
+        raise ValueError("Input RGB tensor must be 3D (C, H, W) or 4D (B, C, H, W)")
+    return xyz
+
+def _xyz_to_lab(xyz):
+    """Converts a PyTorch tensor from XYZ to Lab color space."""
+    Xn, Yn, Zn = 0.95047, 1.0, 1.08883
+    xyz_norm = torch.empty_like(xyz)
+    xyz_norm[:, 0] = xyz[:, 0] / Xn
+    xyz_norm[:, 1] = xyz[:, 1] / Yn
+    xyz_norm[:, 2] = xyz[:, 2] / Zn
+
+    delta = 6/29
+    f_t = torch.where(xyz_norm > delta**3, xyz_norm**(1/3), (xyz_norm / (3 * delta**2)) + (4/29))
+
+    L = (116 * f_t[:, 1]) - 16
+    a = 500 * (f_t[:, 0] - f_t[:, 1])
+    b = 200 * (f_t[:, 1] - f_t[:, 2])
+    return torch.stack([L, a, b], dim=1)
+
+def _lab_to_xyz(lab):
+    """Converts a PyTorch tensor from Lab to XYZ color space."""
+    Xn, Yn, Zn = 0.95047, 1.0, 1.08883
+    f_y = (lab[:, 0] + 16) / 116
+    f_x = lab[:, 1] / 500 + f_y
+    f_z = f_y - lab[:, 2] / 200
+
+    delta = 6/29
+    x_norm = torch.where(f_x > delta, f_x**3, (f_x - 16/116) / 7.787)
+    y_norm = torch.where(f_y > delta, f_y**3, (f_y - 16/116) / 7.787)
+    z_norm = torch.where(f_z > delta, f_z**3, (f_z - 16/116) / 7.787)
+
+    X = x_norm * Xn
+    Y = y_norm * Yn
+    Z = z_norm * Zn
+    return torch.stack([X, Y, Z], dim=1)
+
+def _xyz_to_rgb(xyz):
+    """Converts a PyTorch tensor from XYZ to RGB color space."""
+    M_inv = torch.tensor([
+        [ 3.240479, -1.537150, -0.498535],
+        [-0.969256,  1.875992,  0.041556],
+        [ 0.055648, -0.204043,  1.057311]
+    ], device=xyz.device, dtype=xyz.dtype)
+
+    if xyz.dim() == 3:
+        h, w = xyz.shape[1:]
+        rgb = torch.matmul(M_inv, xyz.view(3, -1)).view(3, h, w)
+    elif xyz.dim() == 4:
+        b, h, w = xyz.shape[0], xyz.shape[2], xyz.shape[3]
+        rgb = torch.matmul(M_inv, xyz.view(b, 3, -1)).view(b, 3, h, w)
+    else:
+        raise ValueError("Input XYZ tensor must be 3D (C, H, W) or 4D (B, C, H, W)")
+    return rgb
+
+def _rgb_to_lab_torch(rgb_image):
+    """Converts an RGB image tensor to Lab color space."""
+    xyz_image = _rgb_to_xyz(rgb_image)
+    lab_image = _xyz_to_lab(xyz_image)
+    return lab_image
+
+def _lab_to_rgb_torch(lab_image):
+    """Converts a Lab image tensor to RGB color space."""
+    xyz_image = _lab_to_xyz(lab_image)
+    rgb_image = _xyz_to_rgb(xyz_image)
+    return rgb_image
+
+def _reinhard_color_transfer(source_image_rgb, target_image_rgb):
+    """Performs Reinhard color transfer from target_image_rgb to source_image_rgb."""
+    if source_image_rgb.dim() == 3:
+        source_image_rgb = source_image_rgb.unsqueeze(0)
+    if target_image_rgb.dim() == 3:
+        target_image_rgb = target_image_rgb.unsqueeze(0)
+
+    source_lab = _rgb_to_lab_torch(source_image_rgb)
+    target_lab = _rgb_to_lab_torch(target_image_rgb)
+
+    mean_source = torch.mean(source_lab, dim=[2, 3], keepdim=True)
+    std_source = torch.std(source_lab, dim=[2, 3], keepdim=True)
+    mean_target = torch.mean(target_lab, dim=[2, 3], keepdim=True)
+    std_target = torch.std(target_lab, dim=[2, 3], keepdim=True)
+
+    transferred_lab = (source_lab - mean_source) * (std_target / (std_source + 1e-8)) + mean_target
+    transferred_rgb = _lab_to_rgb_torch(transferred_lab)
+    transferred_rgb = torch.clamp(transferred_rgb, 0, 1)
+
+    if source_image_rgb.shape[0] == 1:
+        return transferred_rgb.squeeze(0)
+    return transferred_rgb
+
+def _gpu_mkl(cov_source, cov_target):
+    """PyTorch implementation of the MKL algorithm."""
+    EPS = 1e-8
+    eigvals_source, eigvecs_source = torch.linalg.eigh(cov_source)
+    eigvals_source = eigvals_source.real
+    eigvecs_source = eigvecs_source.real
+
+    sqrt_eigvals_source = torch.sqrt(torch.clamp(eigvals_source, min=0)) + EPS
+
+    C = torch.matmul(torch.matmul(torch.diag(sqrt_eigvals_source), eigvecs_source.T), torch.matmul(cov_target, torch.matmul(eigvecs_source, torch.diag(sqrt_eigvals_source))))
+
+    eigvals_C, eigvecs_C = torch.linalg.eigh(C)
+    eigvals_C = eigvals_C.real
+    eigvecs_C = eigvecs_C.real
+
+    sqrt_eigvals_C = torch.sqrt(torch.clamp(eigvals_C, min=0)) + EPS
+    inv_sqrt_eigvals_source = torch.diag(1.0 / sqrt_eigvals_source)
+
+    transform_matrix = torch.matmul(torch.matmul(torch.matmul(torch.matmul(eigvecs_source, inv_sqrt_eigvals_source), eigvecs_C), torch.diag(sqrt_eigvals_C)), torch.matmul(eigvecs_C.T, torch.matmul(inv_sqrt_eigvals_source, eigvecs_source.T)))
+    return transform_matrix
+
+def _gpu_mkl_color_transfer(source, target):
+    """PyTorch implementation of the MKL color transfer."""
+    b, c, h, w = source.shape
+    source_reshaped = source.permute(0, 2, 3, 1).reshape(-1, 3)
+    target_reshaped = target.permute(0, 2, 3, 1).reshape(-1, 3)
+
+    cov_source = _calculate_covariance(source_reshaped)
+    cov_target = _calculate_covariance(target_reshaped)
+
+    transform_matrix = _gpu_mkl(cov_source, cov_target)
+
+    mean_source = source_reshaped.mean(dim=0)
+    mean_target = target_reshaped.mean(dim=0)
+
+    result = torch.matmul(source_reshaped - mean_source, transform_matrix) + mean_target
+    result = result.reshape(b, h, w, c).permute(0, 3, 1, 2)
+
+    return torch.clamp(result, 0, 1)
+
 # ====================================================================================================
 # Color Match Helper
 # ====================================================================================================
-def _apply_color_match(target_batch_tensor, reference_tensor, strength=1.0):
-    """Applies MKL color matching from reference to target batch."""
+def _apply_color_match(target_batch_tensor, reference_tensor, strength=1.0, method="Reinhard"):
+    """Applies color matching from reference to target batch using GPU."""
     if not COLOR_MATCHER_AVAILABLE:
         _log("Skipping color match: color-matcher library not available.")
         return target_batch_tensor
@@ -122,29 +282,32 @@ def _apply_color_match(target_batch_tensor, reference_tensor, strength=1.0):
     if reference_tensor is None or target_batch_tensor is None:
         return target_batch_tensor
 
-    ref_np = reference_tensor.squeeze(0).cpu().numpy().astype(np.float64)
-    target_np = target_batch_tensor.cpu().numpy().astype(np.float64)
-    batch_size = target_np.shape[0]
-    matched_results = []
+    if strength <= 0:
+        _log("Skipping color match: strength is 0 or less.")
+        return target_batch_tensor
 
-    _log(f"Applying internal color match (strength: {strength:.2f}) using MKL method.")
-    cm = ColorMatcher()
+    _log(f"Applying internal color match (strength: {strength:.2f}) using {method} method on GPU.")
+    
+    # Ensure tensors are in (B, C, H, W) format
+    if target_batch_tensor.dim() == 4 and target_batch_tensor.shape[-1] == 3:
+        target_batch_tensor = target_batch_tensor.permute(0, 3, 1, 2)
+    if reference_tensor.dim() == 4 and reference_tensor.shape[-1] == 3:
+        reference_tensor = reference_tensor.permute(0, 3, 1, 2)
 
-    for i in range(batch_size):
-        if (i + 1) % 25 == 0 or i == 0: # Log every 25 frames
-         _log(f"Color matching frame {i+1}/{batch_size}...")
-        try:
-            target_img = target_np[i]
-            transfer_result = cm.transfer(src=target_img, ref=ref_np, method='mkl')
-            final_result = target_img + strength * (transfer_result - target_img)
-            final_result = np.clip(final_result, 0.0, 1.0)
-            matched_results.append(torch.from_numpy(final_result))
-        except Exception as e:
-            _log(f"Error during color matching frame {i}: {e}. Using original frame.")
-            matched_results.append(torch.from_numpy(target_np[i]))
+    if method == "MKL":
+        matched_batch = _gpu_mkl_color_transfer(target_batch_tensor, reference_tensor)
+    else: # Default to Reinhard
+        matched_batch = _reinhard_color_transfer(target_batch_tensor, reference_tensor)
+    
+    # Blend the original and color-matched images
+    final_result = target_batch_tensor + strength * (matched_batch - target_batch_tensor)
+    final_result = torch.clamp(final_result, 0.0, 1.0)
 
-    matched_batch = torch.stack(matched_results).to(device=target_batch_tensor.device, dtype=torch.float32)
-    return matched_batch
+    # Permute back to (B, H, W, C)
+    if final_result.dim() == 4 and final_result.shape[1] == 3:
+        final_result = final_result.permute(0, 2, 3, 1)
+
+    return final_result
 
 
 # ====================================================================================================
@@ -165,10 +328,8 @@ class WanVideoLooper:
                 "negative_prompt": ("STRING", {"tooltip": "A single, global negative prompt applied to all loops."}),
                 "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "The random seed. This will be used for all loops."}),
                 "steps": ("INT", {"default": 6, "min": 1, "max": 10000, "tooltip": "The total number of sampling steps for each loop."}),
-                "enable_motion_cfg": ("BOOLEAN", {"default": False, "tooltip": "Enable a separate CFG for the very first step to control motion."}),
-                "cfg_motion_step": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01, "tooltip": "CFG for the first step ONLY. A value > 1.0 will engage the negative prompt for initial motion guidance."}),
-                "cfg_high_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01, "tooltip": "CFG for the high-noise model (structure/motion). 1.0 disables guidance."}),
-                "cfg_low_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01, "tooltip": "CFG for the low-noise model (details). 1.0 disables guidance."}),
+                "cfg_motion_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01, "tooltip": "CFG for the first step ONLY. A value > 1.0 will engage the negative prompt for initial motion guidance. Set to 1.0 to disable motion guidance."}),
+                "cfg_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01, "tooltip": "Classifier-Free Guidance scale for both high and low noise models. 1.0 disables guidance."}),
                 "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The sampling algorithm (e.g., dpmpp_2m_sde_gpu)."}),
                 "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "The noise scheduler (e.g., karras)."}),
                 "model_switch_point": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.001, "round": 0.001, "tooltip": "Timestep to switch from high-noise to low-noise model. 0.9 (I2V) or 0.875 (T2V) recommended."}),
@@ -179,10 +340,13 @@ class WanVideoLooper:
                 "height": ("INT", {"default": 480, "min": 64, "max": 8192, "step": 16, "tooltip": "The height of the output video."}),
                 "frame_merge": ("INT", {"default": 1, "min": 0, "max": 8, "step": 1, "tooltip": "Number of frames to overlap/merge between loops for smooth transitions."}),
                 "duration_sec": ("INT", {"default": 5, "min": 1, "max": 10, "step": 1, "tooltip": "Duration of each segment in seconds. 5s is recommended, max 10."}),
-                "color_match": ("BOOLEAN", {"default": False, "tooltip": "Apply color matching between segments internally using MKL method."}),
-                "color_match_lastframe": ("BOOLEAN", {"default": False, "tooltip": "If True, overrides segment-by-segment matching. Instead, matches all frames to the *very last* frame of the entire sequence."}),              
+                "color_match_method": (["Disabled", "Reinhard (GPU)", "MKL (GPU)"], {"default": "Disabled", "tooltip": "'Disabled': No color matching is applied. 'Reinhard (GPU)': Applies Reinhard color transfer. 'MKL (GPU)': Applies MKL color transfer."}),
+                "color_match_reference_frame": (["First Frame", "Last Frame", "Sequential", "Reference"], {"default": "First Frame", "tooltip": "'First Frame': Matches all frames to the very first frame of the entire sequence. 'Last Frame': Matches all frames to the very last frame of the entire sequence. 'Sequential': Matches each segment to the last frame of the previous segment. 'Reference': Matches all frames to the optional 'color_match_ref' image."}),
+
                 "color_match_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Strength of the internal color matching (1.0 = full match)."}),
                 "enable_dry_run": ("BOOLEAN", {"default": False, "tooltip": "If enabled, skips sampling and decoding to quickly check setup and logs."}),
+
+
             },
             "optional": {
                 "clip_vision": ("CLIP_VISION", {"tooltip": "(Optional) A CLIP Vision model for guiding the start image."}),
@@ -201,10 +365,10 @@ class WanVideoLooper:
         self,
         model_high, model_low, clip, vae,
         start_image, positive_prompts, negative_prompt,
-        seed, steps, enable_motion_cfg, cfg_motion_step, cfg_high_noise, cfg_low_noise,
+        seed, steps, cfg_motion_noise, cfg_noise,
         sampler_name, scheduler, model_switch_point, sigma_shift, denoise,
         width, height, frame_merge, duration_sec, enable_dry_run,
-        color_match, color_match_strength, color_match_lastframe,
+        color_match_strength, color_match_method, color_match_reference_frame,
         clip_vision=None,
         model_clip_sequence=None,
         color_match_ref=None 
@@ -225,18 +389,28 @@ class WanVideoLooper:
         duration = max(1, int(duration_sec))
         window_frame_count = (duration * 16) + 1
 
-        _log("Cloning base models and clip for looping.")
-        base_model_high = model_high.clone()
-        base_model_low = model_low.clone()
-        base_clip = clip.clone()
+        valid_sequence = isinstance(model_clip_sequence, list)
+
+        if valid_sequence:
+            _log("Using input models directly as base for sequence processing (no initial cloning).")
+            base_model_high = model_high
+            base_model_low = model_low
+            base_clip = clip
+        else:
+            _log("Cloning base models and clip for looping (no sequence provided).")
+            base_model_high = model_high.clone()
+            base_model_low = model_low.clone()
+            base_clip = clip.clone()
 
         _log(f"Starting loop with {len(positive_prompts)} prompts...")
         if enable_dry_run: _log("!!! DRY RUN MODE ENABLED - Sampling and decoding will be skipped !!!")
-        if color_match and not COLOR_MATCHER_AVAILABLE:
+        if color_match_method != "Disabled" and not COLOR_MATCHER_AVAILABLE:
              _log("Warning: Internal color match enabled, but 'color-matcher' library is missing. Feature disabled.")
-             color_match = False
-        if color_match and color_match_lastframe:
-             _log("'Last Frame' color match is enabled. Segment-by-segment matching will be skipped.")
+             color_match_method = "Disabled"
+        if color_match_method != "Disabled" and color_match_reference_frame in ["Last Frame", "First Frame", "Reference"]:
+             _log(f"'{color_match_reference_frame}' color match is enabled. Segment-by-segment matching will be skipped.")
+        elif color_match_method != "Disabled" and color_match_reference_frame == "Sequential":
+             _log(f"'Sequential' color match is enabled. Each segment will be matched to the previous segment's last frame.")
 
 
         model_device, model_dtype = _get_model_device_info(base_model_high)
@@ -245,10 +419,17 @@ class WanVideoLooper:
         render_height = _align_to_multiple(height, 16)
         current_seed = int(seed) & 0xFFFFFFFFFFFFFFFF
         total_steps = int(steps)
+        _log("Pre-allocating latent tensor for reuse.")
+        # Calculate target shape for pre-allocated latent
+        latent_height_for_prealloc = render_height // 8
+        latent_width_for_prealloc = render_width // 8
+        time_dim_for_prealloc = (window_frame_count + 3) // 4
+        pre_allocated_latent_tensor = torch.zeros([1, 16, time_dim_for_prealloc, latent_height_for_prealloc, latent_width_for_prealloc], device=comfy.model_management.intermediate_device())
 
         # --- Color Match Setup ---
         static_color_reference = color_match_ref[0:1] if color_match_ref is not None else None
         previous_loop_last_frame = None
+        first_frame_for_color_match = None
 
         # --- 2. Pre-Loop Preparations (Applied Once) ---
         _log(f"Applying Sigma Shift: {sigma_shift}")
@@ -266,19 +447,27 @@ class WanVideoLooper:
             _log(f"Prompt: {prompt_text[:80]}...")
 
             # --- 3a. Select/Clone Models for this Loop ---
-            _log(f"Segment {i+1}: Cloning fresh base models and clip.")
-            loop_model_high = base_model_high.clone()
-            loop_model_low = base_model_low.clone()
-            loop_clip = base_clip.clone()
-            segment_models_provided = False
-            using_provided_high, using_provided_low, using_provided_clip = False, False, False
-            if valid_sequence and i < len(model_clip_sequence):
-                segment_data = model_clip_sequence[i]
-                if segment_data and isinstance(segment_data, tuple) and len(segment_data) == 3:
-                    mh, ml, c = segment_data
-                    if mh is not None: loop_model_high = mh; segment_models_provided = True; using_provided_high = True
-                    if ml is not None: loop_model_low = ml; segment_models_provided = True; using_provided_low = True
-                    if c is not None: loop_clip = c; segment_models_provided = True; using_provided_clip = True
+            if not valid_sequence:
+                _log(f"Segment {i+1}: Using BASE models and clip (no sequence provided).")
+                loop_model_high = base_model_high
+                loop_model_low = base_model_low
+                loop_clip = base_clip
+                segment_models_provided = False
+                using_provided_high, using_provided_low, using_provided_clip = False, False, False
+            else:
+                _log(f"Segment {i+1}: Processing models and clip from sequence.")
+                loop_model_high = base_model_high # Start with base, will be overwritten if sequence provides
+                loop_model_low = base_model_low   # Start with base, will be overwritten if sequence provides
+                loop_clip = base_clip
+                segment_models_provided = False
+                using_provided_high, using_provided_low, using_provided_clip = False, False, False
+                if i < len(model_clip_sequence):
+                    segment_data = model_clip_sequence[i]
+                    if segment_data and isinstance(segment_data, tuple) and len(segment_data) == 3:
+                        model_high_seq, model_low_seq, clip_seq = segment_data
+                        if model_high_seq is not None: loop_model_high = model_high_seq; segment_models_provided = True; using_provided_high = True
+                        if model_low_seq is not None: loop_model_low = model_low_seq; segment_models_provided = True; using_provided_low = True
+                        if clip_seq is not None: loop_clip = clip_seq; segment_models_provided = True; using_provided_clip = True
 
             if segment_models_provided:
                 high_source = "PROVIDED (LoRA?)" if using_provided_high else "CLONED BASE"
@@ -294,7 +483,7 @@ class WanVideoLooper:
             loop_positive_cond, loop_negative_cond, latent_on_cpu = _prepare_latent_window(
                 global_positive_cond, global_negative_cond, vae,
                 render_width, render_height, window_frame_count, 1,
-                current_seed_image, vision_output
+                current_seed_image, vision_output, pre_allocated_latent=pre_allocated_latent_tensor
             )
             latent_on_gpu = {"samples": _move_tensor(latent_on_cpu["samples"], model_device, model_dtype)}
             final_latent = None
@@ -302,15 +491,15 @@ class WanVideoLooper:
             # --- 3c/3d. Run Samplers OR Skip for Dry Run ---
             if not enable_dry_run:
                 _log(f"Segment {i+1}: Running HIGH noise sampler (Steps 0 to {switch_step_index}).")
-                if enable_motion_cfg and switch_step_index > 1:
-                    _log(f"Segment {i+1}: Using special Motion CFG ({cfg_motion_step}) for step 0.")
-                    initial_latent = _run_sampler_step(loop_model_high, current_seed, total_steps, cfg_motion_step, sampler_name, scheduler, loop_positive_cond, loop_negative_cond, latent_on_gpu, denoise, disable_noise=False, start_step=0, last_step=1, force_full_denoise=False)
-                    high_noise_latent = _run_sampler_step(loop_model_high, current_seed, total_steps, cfg_high_noise, sampler_name, scheduler, loop_positive_cond, loop_negative_cond, initial_latent, denoise, disable_noise=True, start_step=1, last_step=switch_step_index, force_full_denoise=False)
+                if cfg_motion_noise != 1.0 and switch_step_index > 1:
+                    _log(f"Segment {i+1}: Using special Motion CFG ({cfg_motion_noise}) for step 0.")
+                    initial_latent = _run_sampler_step(loop_model_high, current_seed, total_steps, cfg_motion_noise, sampler_name, scheduler, loop_positive_cond, loop_negative_cond, latent_on_gpu, denoise, disable_noise=False, start_step=0, last_step=1, force_full_denoise=False)
+                    high_noise_latent = _run_sampler_step(loop_model_high, current_seed, total_steps, cfg_noise, sampler_name, scheduler, loop_positive_cond, loop_negative_cond, initial_latent, denoise, disable_noise=True, start_step=1, last_step=switch_step_index, force_full_denoise=False)
                 else:
-                    high_noise_latent = _run_sampler_step(loop_model_high, current_seed, total_steps, cfg_high_noise, sampler_name, scheduler, loop_positive_cond, loop_negative_cond, latent_on_gpu, denoise, disable_noise=False, start_step=0, last_step=switch_step_index, force_full_denoise=False)
+                    high_noise_latent = _run_sampler_step(loop_model_high, current_seed, total_steps, cfg_noise, sampler_name, scheduler, loop_positive_cond, loop_negative_cond, latent_on_gpu, denoise, disable_noise=False, start_step=0, last_step=switch_step_index, force_full_denoise=False)
 
                 _log(f"Segment {i+1}: Running LOW noise sampler (Steps {switch_step_index} to {total_steps}).")
-                final_latent = _run_sampler_step(loop_model_low, current_seed, total_steps, cfg_low_noise, sampler_name, scheduler, loop_positive_cond, loop_negative_cond, high_noise_latent, denoise, disable_noise=True, start_step=switch_step_index, last_step=total_steps, force_full_denoise=True)
+                final_latent = _run_sampler_step(loop_model_low, current_seed, total_steps, cfg_noise, sampler_name, scheduler, loop_positive_cond, loop_negative_cond, high_noise_latent, denoise, disable_noise=True, start_step=switch_step_index, last_step=total_steps, force_full_denoise=True)
                 _log(f"Segment {i+1}: Sampling finished.")
             else:
                 _log(f"Segment {i+1}: DRY RUN - Skipping HIGH noise sampler."); _log(f"Segment {i+1}: DRY RUN - Skipping LOW noise sampler.")
@@ -346,22 +535,27 @@ class WanVideoLooper:
                     reference_for_this_loop = previous_loop_last_frame
                 
                 # Only run segment matching if color_match is on AND color_match_lastframe is OFF
-                if color_match and not color_match_lastframe and reference_for_this_loop is not None:
+                if color_match_method != "Disabled" and color_match_reference_frame == "Sequential" and reference_for_this_loop is not None:
                     _log(f"Segment {i+1}: Applying segment-to-segment color match.")
                     decoded_image_batch = _apply_color_match(
                         target_batch_tensor=decoded_image_batch,
                         reference_tensor=reference_for_this_loop,
-                        strength=color_match_strength
+                        strength=color_match_strength,
+                        method=color_match_method
                     )
 
                 previous_loop_last_frame = next_reference_frame
             else:
-                 previous_loop_last_frame = None
+                 _log("Warning: Skipping frame storage due to missing decoded batch.")
+                 current_seed_image = None
 
             # --- 3g. Store Frames & Update Seed ---
             if decoded_image_batch is not None and decoded_image_batch.shape[0] > 0:
+                if i == 0 and color_match_reference_frame == "First Frame":
+                    first_frame_for_color_match = decoded_image_batch[0:1]
+
                 if decoded_image_batch.shape[0] >= merge_frame_count and merge_frame_count > 0:
-                    current_seed_image = decoded_image_batch[-merge_frame_count:-merge_frame_count+1 if merge_frame_count > 1 else None]
+                    current_seed_image = decoded_image_batch[-merge_frame_count:]
                 elif decoded_image_batch.shape[0] > 0: current_seed_image = decoded_image_batch[-1:]
                 else: _log("Warning: No frames decoded/generated..."); current_seed_image = None
                 frames_to_save = decoded_image_batch if i == 0 else decoded_image_batch[merge_frame_count:]
@@ -370,22 +564,8 @@ class WanVideoLooper:
                  _log("Warning: Skipping frame storage due to missing decoded batch.")
                  current_seed_image = None
 
-            # --- 3h. Per-Iteration Cleanup ---
             _log(f"Segment {i+1}: Cleaning up loop resources.")
-            del loop_model_high, loop_model_low, loop_clip
-            _log(f"Segment {i+1}: Deleted loop model/clip variables.")
-            if 'high_noise_latent' in locals(): del high_noise_latent
-            if final_latent is not None: del final_latent
-            if 'latent_on_gpu' in locals(): del latent_on_gpu
-            if 'final_latent_on_cpu' in locals(): del final_latent_on_cpu
-            if 'decoded_data' in locals(): del decoded_data
-            if decoded_image_batch is not None: del decoded_image_batch
-            if 'frames_to_save' in locals(): del frames_to_save
-            _log(f"Segment {i+1}: Forcing garbage collection.")
-            gc.collect()
-            if torch.cuda.is_available():
-                _log(f"Segment {i+1}: Emptying CUDA cache.")
-                torch.cuda.memory.empty_cache()
+
             _log(f"--- Finished Loop {i+1} ---")
 
         # --- 4. Finalize & Return ---
@@ -403,22 +583,35 @@ class WanVideoLooper:
             return (torch.zeros((1, 64, 64, 3)),)*2
 
         # --- 4b. Final Frame Color Match (New Logic) ---
-        if color_match and color_match_lastframe and final_batch.shape[0] > 1:
-            _log(f"Applying 'Last Frame' color match using frame {final_batch.shape[0]-1} as reference.")
-            last_frame_reference = final_batch[-1:] 
-            
-            frames_to_match = final_batch[:-1]
-            
-            matched_frames = _apply_color_match(
-                target_batch_tensor=frames_to_match,
-                reference_tensor=last_frame_reference,
-                strength=color_match_strength
-            )
-            
-            final_batch = torch.cat([matched_frames, last_frame_reference], dim=0)
-            _log("'Last Frame' color match complete.")
-        elif color_match and color_match_lastframe:
-             _log("Warning: 'Last Frame' color match enabled, but not enough frames to perform match (<= 1).")
+        if color_match_method != "Disabled" and final_batch.shape[0] > 1 and color_match_reference_frame in ["Last Frame", "First Frame", "Reference"]:
+            reference_frame_for_final_match = None
+            if color_match_reference_frame == "Last Frame":
+                _log(f"Applying 'Last Frame' color match using frame {final_batch.shape[0]-1} as reference.")
+                reference_frame_for_final_match = final_batch[-1:]
+            elif color_match_reference_frame == "First Frame":
+                _log(f"Applying 'First Frame' color match using frame 0 as reference.")
+                reference_frame_for_final_match = first_frame_for_color_match
+            elif color_match_reference_frame == "Reference":
+                if static_color_reference is not None:
+                    _log(f"Applying 'Reference' color match using provided reference image.")
+                    reference_frame_for_final_match = static_color_reference
+                else:
+                    _log(f"Warning: 'Reference' color match selected, but no reference image provided. Skipping final color match.")
+
+            if reference_frame_for_final_match is not None:
+                frames_to_match = final_batch[:-1]
+
+                matched_frames = _apply_color_match(
+                    target_batch_tensor=frames_to_match,
+                    reference_tensor=reference_frame_for_final_match,
+                    strength=color_match_strength,
+                    method=color_match_method
+                )
+
+                final_batch = torch.cat([matched_frames, reference_frame_for_final_match], dim=0)
+                _log(f"'{color_match_reference_frame}' color match complete.")
+        elif color_match_method != "Disabled" and final_batch.shape[0] <= 1 and color_match_reference_frame in ["Last Frame", "First Frame", "Reference"]:
+             _log(f"Warning: '{color_match_reference_frame}' color match enabled, but not enough frames to perform match (<= 1).")
 
 
         last_frame = final_batch[-1:] if final_batch.shape[0] > 0 else torch.zeros((1, render_height, render_width, 3))

wan_lora_sequencer.py

@@ -1,4 +1,4 @@
-import folder_paths
+
 
 # ====================================================================================================
 # Logging Utility

wan_prompts_node.py

@@ -1,4 +1,4 @@
-import json
+
 
 # ====================================================================================================
 # Logging Utility
@@ -21,7 +21,7 @@ class WanVideoLooperPrompts:
             "required": {
                 "prompts": ("STRING", {"multiline": True, "default": "", "tooltip": "Enter one positive prompt per line. Empty lines are ignored."}),
                 "negative_prompt": ("STRING", {"multiline": True, "default": "", "tooltip": "Enter the global negative prompt."}),
-                "enable_prefix_suffix": (["true", "false"], {"default": "true", "tooltip": "Enable the prefix and suffix string."}),
+                "enable_prefix_suffix": ("BOOLEAN", {"default": True, "tooltip": "Enable the prefix and suffix string."}),
                 "prefix": ("STRING", {"default": "", "tooltip": "Text to add BEFORE each prompt line (positive and negative)."}),
                 "suffix": ("STRING", {"default": "", "tooltip": "Text to add AFTER each prompt line (positive and negative)."}),
             }
@@ -40,7 +40,7 @@ class WanVideoLooperPrompts:
         if not prompt_text:
             return ""
 
-        is_prefix_suffix_enabled = (enable_prefix_suffix == "true")
+        is_prefix_suffix_enabled = enable_prefix_suffix
 
         prefix_text = prefix.strip() if is_prefix_suffix_enabled else ""
         suffix_text = suffix.strip() if is_prefix_suffix_enabled else ""