Baverne-comfyUI-TiledWan/nodes.py

import torch
import comfy.utils
import comfy.model_management
import node_helpers
import random
import os
import sys
import importlib
import math
import nodes
import numpy as np
import torchvision.transforms.functional as F
from PIL import Image
from scipy.ndimage import gaussian_filter, grey_dilation, binary_closing, binary_fill_holes


class ImageToMask:
    """
    Convert an image to a mask by extracting a specific channel.
    Supports red, green, blue, and alpha channels with optional clamping and normalization.
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE", {"tooltip": "Input image to convert to mask. Can be any standard image format."}),
                "channel": (["red", "green", "blue", "alpha"], {"tooltip": "Which color channel to extract as the mask. Alpha channel only available if image has transparency."}),
                "clamp_output": ("BOOLEAN", {
                    "default": True,
                    "label_on": "enabled",
                    "label_off": "disabled",
                    "tooltip": "Clamp output values to 0-1 range. Recommended to keep enabled for proper mask values."
                }),
                "normalize_output": ("BOOLEAN", {
                    "default": True,
                    "label_on": "enabled", 
                    "label_off": "disabled",
                    "tooltip": "Normalize the output to use full 0-1 range."
                }),
            },
        }

    CATEGORY = "TiledWan"
    DESCRIPTION = """
    Extracts a specific color channel from an image to create a mask.
    
    Useful for converting colored regions, alpha channels, or channel-specific data into masks
    for use with inpainting, compositing, or other masking operations. The output mask values
    represent the intensity of the selected channel, with proper normalization and clamping.
    """
    
    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("mask",)
    OUTPUT_TOOLTIPS = ("Grayscale mask derived from the selected image channel. White areas indicate high channel values, black areas indicate low values.",)
    FUNCTION = "image_to_mask"

    def image_to_mask(self, image, channel, clamp_output, normalize_output):
        """
        Convert image to mask by extracting specified channel.
        
        Args:
            image: Input image tensor [batch, height, width, channels]
            channel: Channel to extract ("red", "green", "blue", "alpha")
            clamp_output: Whether to clamp values between 0 and 1
            normalize_output: Whether to normalize values to 0-1 range
            
        Returns:
            tuple: Mask tensor [batch, height, width]
        """
        channels = ["red", "green", "blue", "alpha"]
        channel_index = channels.index(channel)
        
        # Handle case where alpha channel is requested but doesn't exist
        if channel_index == 3 and image.shape[-1] < 4:
            # Create an alpha channel filled with ones (fully opaque)
            mask = torch.ones(image.shape[0], image.shape[1], image.shape[2], dtype=image.dtype, device=image.device)
        else:
            # Extract the specified channel
            mask = image[:, :, :, channel_index]
        
        # Apply clamping if requested
        if clamp_output:
            mask = torch.clamp(mask, 0.0, 1.0)
        
        # Apply normalization if requested
        if normalize_output:
            mask_min = mask.min()
            mask_max = mask.max()
            if mask_max > mask_min:  # Avoid division by zero
                mask = (mask - mask_min) / (mask_max - mask_min)
        
        
        
        return (mask,)


class ImageStatistics:
    """
    Calculate and display comprehensive statistics for an image.
    Shows min, max, mean, variance, median, and other useful statistics.
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE", {"tooltip": "Input image to analyze. Statistics will be calculated across all pixels and channels."}),
                "show_per_channel": ("BOOLEAN", {
                    "default": True,
                    "label_on": "enabled",
                    "label_off": "disabled",
                    "tooltip": "Display separate statistics for each color channel (R, G, B) in addition to overall statistics."
                }),
            },
        }

    RETURN_TYPES = ("IMAGE", "FLOAT", "FLOAT", "FLOAT", "FLOAT", "FLOAT")
    RETURN_NAMES = ("image", "min_value", "max_value", "mean_value", "variance", "median_value")
    OUTPUT_TOOLTIPS = ("Pass-through of the input image", 
                      "Minimum pixel value found in the image", 
                      "Maximum pixel value found in the image", 
                      "Average pixel value across the entire image", 
                      "Variance of pixel values (measure of spread/contrast)", 
                      "Median pixel value (middle value when sorted)")
    FUNCTION = "calculate_statistics"
    OUTPUT_NODE = True  # This allows the node to display output in the console
    CATEGORY = "TiledWan"
    DESCRIPTION = """
    Analyzes image pixel values and provides comprehensive statistical information.
    
    Calculates key metrics including minimum, maximum, mean, variance, and median values.
    Useful for understanding image characteristics, checking normalization, detecting
    issues with image processing, and debugging workflows. Can show per-channel
    statistics for color analysis.
    """

    def calculate_statistics(self, image: torch.Tensor, show_per_channel: bool):
        """
        Calculate comprehensive statistics for the input image.
        
        Args:
            image: Input image tensor [batch, height, width, channels]
            show_per_channel: Whether to show per-channel statistics
            
        Returns:
            tuple: (original_image, min, max, mean, variance, median)
        """
        print("\n=== TiledWan Image Statistics ===")
        
        # Basic image info
        batch_size, height, width, channels = image.shape
        total_pixels = height * width * channels * batch_size
        
        print(f"Image shape: {image.shape} (B×H×W×C)")
        print(f"Total pixels: {total_pixels:,}")
        print(f"Memory usage: {image.numel() * image.element_size() / (1024*1024):.2f} MB")
        print(f"Data type: {image.dtype}")
        print(f"Device: {image.device}")
        
        # Calculate global statistics
        min_val = image.min().item()
        max_val = image.max().item()
        mean_val = image.mean().item()
        variance_val = image.var().item()
        std_val = image.std().item()
        
        # Calculate median (flatten the tensor first)
        flattened = image.flatten()
        # Sort and find median manually to avoid compatibility issues
        sorted_values = torch.sort(flattened).values
        n = len(sorted_values)
        if n % 2 == 0:
            median_val = (sorted_values[n//2 - 1] + sorted_values[n//2]).item() / 2.0
        else:
            median_val = sorted_values[n//2].item()
        
        # Additional useful statistics
        non_zero_count = torch.count_nonzero(image).item()
        zero_count = total_pixels - non_zero_count
        
        print("\n--- Global Statistics ---")
        print(f"Min value: {min_val:.6f}")
        print(f"Max value: {max_val:.6f}")
        print(f"Mean value: {mean_val:.6f}")
        print(f"Median value: {median_val:.6f}")
        print(f"Variance: {variance_val:.6f}")
        print(f"Standard deviation: {std_val:.6f}")
        print(f"Value range: {max_val - min_val:.6f}")
        print(f"Non-zero pixels: {non_zero_count:,} ({non_zero_count/total_pixels:.2%})")
        print(f"Zero pixels: {zero_count:,} ({zero_count/total_pixels:.2%})")
        
        # Per-channel statistics if requested
        if show_per_channel and channels > 1:
            print("\n--- Per-Channel Statistics ---")
            channel_names = ["Red", "Green", "Blue", "Alpha"][:channels]
            
            for c in range(channels):
                channel_data = image[:, :, :, c]
                ch_min = channel_data.min().item()
                ch_max = channel_data.max().item()
                ch_mean = channel_data.mean().item()
                # Calculate median manually for compatibility
                ch_flattened = channel_data.flatten()
                ch_sorted = torch.sort(ch_flattened).values
                ch_n = len(ch_sorted)
                if ch_n % 2 == 0:
                    ch_median = (ch_sorted[ch_n//2 - 1] + ch_sorted[ch_n//2]).item() / 2.0
                else:
                    ch_median = ch_sorted[ch_n//2].item()
                ch_std = channel_data.std().item()
                
                print(f"{channel_names[c]} Channel:")
                print(f"  Min: {ch_min:.6f}, Max: {ch_max:.6f}")
                print(f"  Mean: {ch_mean:.6f}, Median: {ch_median:.6f}")
                print(f"  Std: {ch_std:.6f}, Range: {ch_max - ch_min:.6f}")
        
        # Value distribution analysis
        print("\n--- Value Distribution ---")
        # Create bins for histogram analysis
        bins = torch.linspace(min_val, max_val, 11)  # 10 bins
        hist = torch.histc(flattened, bins=10, min=min_val, max=max_val)
        
        print("Value distribution (10 bins):")
        for i in range(len(hist)):
            bin_start = bins[i].item()
            bin_end = bins[i+1].item() if i < len(bins)-1 else max_val
            count = hist[i].item()
            percentage = count / total_pixels * 100
            print(f"  [{bin_start:.3f} - {bin_end:.3f}]: {count:,.0f} pixels ({percentage:.1f}%)")
        
        # Percentiles
        print("\n--- Percentiles ---")
        percentiles = [1, 5, 10, 25, 50, 75, 90, 95, 99]
        # Use the already sorted values from median calculation
        for p in percentiles:
            idx = int((p / 100.0) * (len(sorted_values) - 1))
            value = sorted_values[idx].item()
            print(f"  {p:2d}th percentile: {value:.6f}")
        
        # Potential issues detection
        print("\n--- Analysis ---")
        if min_val < 0:
            print("⚠️  Warning: Image contains negative values!")
        if max_val > 1:
            print("⚠️  Warning: Image contains values > 1.0!")
        if min_val == max_val:
            print("⚠️  Warning: Image has constant values (no variation)!")
        if std_val < 0.01:
            print("⚠️  Notice: Very low standard deviation - image might be nearly uniform")
        if non_zero_count / total_pixels < 0.1:
            print("⚠️  Notice: Most pixels are zero - sparse image")
        
        print("================================\n")
        
        # Return the image and key statistics
        return (image, min_val, max_val, mean_val, variance_val, median_val)


class MaskStatistics:
    """
    Calculate and display comprehensive statistics for a mask.
    Shows min, max, mean, variance, median, and mask-specific statistics.
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "mask": ("MASK", {"tooltip": "Input mask to analyze. Should be a grayscale mask with values typically between 0-1."}),
                "analyze_connected_components": ("BOOLEAN", {
                    "default": True,
                    "label_on": "enabled",
                    "label_off": "disabled",
                    "tooltip": "Analyze connected regions in the mask to count separate masked areas and find the largest connected component."
                }),
            },
        }

    RETURN_TYPES = ("MASK", "FLOAT", "FLOAT", "FLOAT", "FLOAT", "FLOAT", "INT")
    RETURN_NAMES = ("mask", "min_value", "max_value", "mean_value", "variance", "median_value", "white_pixel_count")
    OUTPUT_TOOLTIPS = ("Pass-through of the input mask", 
                      "Minimum mask value found", 
                      "Maximum mask value found", 
                      "Average mask value (coverage ratio)", 
                      "Variance of mask values", 
                      "Median mask value", 
                      "Number of white/masked pixels (where value > 0.5)")
    FUNCTION = "calculate_mask_statistics"
    OUTPUT_NODE = True  # This allows the node to display output in the console
    CATEGORY = "TiledWan"
    DESCRIPTION = """
    Analyzes mask properties and provides detailed statistical information.
    
    Calculates standard statistics plus mask-specific metrics like coverage ratio,
    connected component analysis, and white pixel count. Essential for understanding
    mask quality, coverage area, and structure. Helps debug masking operations and
    validate mask inputs for inpainting or compositing workflows.
    """

    def calculate_mask_statistics(self, mask: torch.Tensor, analyze_connected_components: bool):
        """
        Calculate comprehensive statistics for the input mask.
        
        Args:
            mask: Input mask tensor [batch, height, width]
            analyze_connected_components: Whether to analyze connected components
            
        Returns:
            tuple: (original_mask, min, max, mean, variance, median, white_pixel_count)
        """
        print("\n=== TiledWan Mask Statistics ===")
        
        # Basic mask info
        if len(mask.shape) == 3:
            batch_size, height, width = mask.shape
            channels = 1
        else:
            batch_size, height, width, channels = mask.shape
            
        total_pixels = height * width * batch_size
        
        print(f"Mask shape: {mask.shape} (B×H×W)")
        print(f"Total pixels: {total_pixels:,}")
        print(f"Memory usage: {mask.numel() * mask.element_size() / (1024*1024):.2f} MB")
        print(f"Data type: {mask.dtype}")
        print(f"Device: {mask.device}")
        
        # Calculate global statistics
        min_val = mask.min().item()
        max_val = mask.max().item()
        mean_val = mask.mean().item()
        variance_val = mask.var().item()
        std_val = mask.std().item()
        
        # Calculate median (flatten the tensor first)
        flattened = mask.flatten()
        # Sort and find median manually to avoid compatibility issues
        sorted_values = torch.sort(flattened).values
        n = len(sorted_values)
        if n % 2 == 0:
            median_val = (sorted_values[n//2 - 1] + sorted_values[n//2]).item() / 2.0
        else:
            median_val = sorted_values[n//2].item()
        
        # Mask-specific statistics
        white_pixels = torch.sum(mask > 0.5).item()  # Pixels considered "white" (active)
        black_pixels = total_pixels - white_pixels
        coverage = white_pixels / total_pixels
        
        # Count unique values (useful for masks)
        unique_values = torch.unique(mask)
        unique_count = len(unique_values)
        
        print("\n--- Global Statistics ---")
        print(f"Min value: {min_val:.6f}")
        print(f"Max value: {max_val:.6f}")
        print(f"Mean value: {mean_val:.6f}")
        print(f"Median value: {median_val:.6f}")
        print(f"Variance: {variance_val:.6f}")
        print(f"Standard deviation: {std_val:.6f}")
        print(f"Value range: {max_val - min_val:.6f}")
        
        print("\n--- Mask-Specific Statistics ---")
        print(f"White pixels (>0.5): {white_pixels:,} ({coverage:.2%})")
        print(f"Black pixels (≤0.5): {black_pixels:,} ({(1-coverage):.2%})")
        print(f"Unique values: {unique_count}")
        if unique_count <= 10:
            print(f"Unique values list: {[f'{val:.3f}' for val in unique_values.tolist()]}")
        
        # Mask density analysis
        print(f"Mask density: {coverage:.4f}")
        if coverage > 0.8:
            print("  → High density mask (mostly white)")
        elif coverage > 0.5:
            print("  → Medium-high density mask")
        elif coverage > 0.2:
            print("  → Medium-low density mask")
        elif coverage > 0.05:
            print("  → Low density mask")
        else:
            print("  → Very sparse mask")
        
        # Value distribution analysis
        print("\n--- Value Distribution ---")
        # For masks, show distribution more relevant to binary/grayscale nature
        if unique_count <= 20:  # If few unique values, show exact counts
            print("Exact value counts:")
            for val in unique_values:
                count = torch.sum(mask == val).item()
                percentage = count / total_pixels * 100
                print(f"  Value {val:.3f}: {count:,} pixels ({percentage:.1f}%)")
        else:
            # Create bins for histogram analysis
            bins = torch.linspace(min_val, max_val, 11)  # 10 bins
            hist = torch.histc(flattened, bins=10, min=min_val, max=max_val)
            
            print("Value distribution (10 bins):")
            for i in range(len(hist)):
                bin_start = bins[i].item()
                bin_end = bins[i+1].item() if i < len(bins)-1 else max_val
                count = hist[i].item()
                percentage = count / total_pixels * 100
                print(f"  [{bin_start:.3f} - {bin_end:.3f}]: {count:,.0f} pixels ({percentage:.1f}%)")
        
        # Percentiles
        print("\n--- Percentiles ---")
        percentiles = [1, 5, 10, 25, 50, 75, 90, 95, 99]
        # Use the already sorted values from median calculation
        for p in percentiles:
            idx = int((p / 100.0) * (len(sorted_values) - 1))
            value = sorted_values[idx].item()
            print(f"  {p:2d}th percentile: {value:.6f}")
        
        # Connected components analysis (if requested)
        if analyze_connected_components and coverage > 0 and coverage < 1:
            print("\n--- Connected Components Analysis ---")
            try:
                # Simple connected components analysis
                # For each batch item
                total_components = 0
                largest_component = 0
                smallest_component = float('inf')
                
                for b in range(batch_size):
                    mask_slice = mask[b] if len(mask.shape) == 3 else mask[b, :, :, 0]
                    binary_mask = (mask_slice > 0.5).float()
                    
                    if binary_mask.sum() > 0:  # Only analyze if there are white pixels
                        # Simple component counting (not perfect but gives an idea)
                        # Count clusters by looking at transitions
                        transitions = 0
                        prev_row_white = torch.zeros(width, dtype=torch.bool, device=mask.device)
                        
                        for row in range(height):
                            current_row = binary_mask[row] > 0.5
                            # Count horizontal transitions in this row
                            if row == 0:
                                transitions += torch.sum(current_row[1:] != current_row[:-1]).item()
                            
                            # Count vertical transitions
                            transitions += torch.sum(current_row != prev_row_white).item()
                            prev_row_white = current_row
                        
                        # Rough estimate of components (very approximate)
                        estimated_components = max(1, transitions // 4)
                        total_components += estimated_components
                        
                        component_size = binary_mask.sum().item()
                        largest_component = max(largest_component, component_size)
                        smallest_component = min(smallest_component, component_size)
                
                if total_components > 0:
                    avg_components = total_components / batch_size
                    print(f"Estimated components per image: {avg_components:.1f}")
                    print(f"Largest component: {largest_component:,} pixels")
                    if smallest_component != float('inf'):
                        print(f"Smallest component: {smallest_component:,} pixels")
                else:
                    print("No connected components found")
                    
            except Exception as e:
                print(f"Connected components analysis failed: {str(e)}")
        
        # Potential issues detection
        print("\n--- Analysis ---")
        if min_val < 0:
            print("⚠️  Warning: Mask contains negative values!")
        if max_val > 1:
            print("⚠️  Warning: Mask contains values > 1.0!")
        if min_val == max_val:
            print("⚠️  Warning: Mask has constant values (uniform mask)!")
        if unique_count == 2 and min_val == 0 and max_val == 1:
            print("✓  Perfect binary mask (0 and 1 only)")
        elif unique_count <= 5:
            print(f"✓  Nearly binary mask ({unique_count} unique values)")
        if coverage < 0.01:
            print("⚠️  Notice: Very sparse mask - almost empty")
        if coverage > 0.99:
            print("⚠️  Notice: Very dense mask - almost full")
        if std_val < 0.01:
            print("⚠️  Notice: Very low standard deviation - mask might be nearly uniform")
        
        print("================================\n")
        
        # Return the mask and key statistics
        return (mask, min_val, max_val, mean_val, variance_val, median_val, white_pixels)


class Tile:
    """
    Represents a single tile with its content and position in the tiling grid
    """
    def __init__(self, content, temporal_index, line_index, column_index, 
                 temporal_range, spatial_range_h, spatial_range_w):
        self.content = content  # The actual tile tensor
        self.temporal_index = temporal_index  # Index in temporal dimension
        self.line_index = line_index  # Index in height/line dimension
        self.column_index = column_index  # Index in width/column dimension
        
        # Store the ranges for overlap calculations
        self.temporal_range = temporal_range  # (start_frame, end_frame)
        self.spatial_range_h = spatial_range_h  # (start_h, end_h)
        self.spatial_range_w = spatial_range_w  # (start_w, end_w)
    
    def __str__(self):
        return f"Tile[T{self.temporal_index}:L{self.line_index}:C{self.column_index}] " \
               f"Shape: {self.content.shape} " \
               f"Temporal: {self.temporal_range} " \
               f"Spatial: H{self.spatial_range_h} W{self.spatial_range_w}"


class TileAndStitchBack:
    """
    Tile and Stitch Back node with proper dimension-wise stitching.
    
    New approach:
    1. Temporal tiling: Split video into chunks with overlap
    2. Spatial tiling: Split each chunk into tiles with overlap  
    3. Processing: Apply transformations to each tile
    4. Dimension-wise stitching:
       a) Column-wise: Stitch tiles in same column (same temporal chunk)
       b) Line-wise: Stitch columns together (complete temporal chunks)
       c) Time-wise: Stitch temporal chunks together
    
    This approach eliminates the black-to-color ramp issue by handling
    overlaps systematically in each dimension.
    """
    
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        """
        Tile and Stitch Back inputs
        """
        return {
            "required": {
                "video": ("IMAGE",),
                "target_frames": ("INT", {"default": 81, "min": 16, "max": 200, "step": 1}),
                "target_width": ("INT", {"default": 832, "min": 64, "max": 2048, "step": 8}),
                "target_height": ("INT", {"default": 480, "min": 64, "max": 2048, "step": 8}),
                "frame_overlap": ("INT", {"default": 10, "min": 0, "max": 40, "step": 1}),
                "spatial_overlap": ("INT", {"default": 20, "min": 0, "max": 100, "step": 4}),
                "color_shift_strength": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01}),
                "debug_mode": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("stitched_video", "tile_info")
    FUNCTION = "tile_and_stitch"
    OUTPUT_NODE = True
    CATEGORY = "TiledWan"

    def tile_and_stitch(self, video, target_frames, target_width, target_height, 
                       frame_overlap, spatial_overlap, color_shift_strength, debug_mode):
        """
        Tile video temporally and spatially, apply transformations, then stitch back using dimension-wise approach
        """
        
        print("\n" + "="*80)
        print("                  TILE AND STITCH BACK (NEW ALGORITHM)")
        print("="*80)
        print("🧩 Starting dimension-wise video tiling and stitching...")
        
        try:
            import random
            import traceback
            
            # Input video info
            batch_size, height, width, channels = video.shape
            print(f"📹 Input video shape: {video.shape} (B×H×W×C)")
            print(f"🎯 Target tile size: {target_frames} frames × {target_width}×{target_height}")
            print(f"🔗 Overlaps: {frame_overlap} frames, {spatial_overlap} pixels")
            
            # Calculate tile ranges
            temporal_tiles = self._calculate_temporal_tiles(batch_size, target_frames, frame_overlap)
            spatial_tiles_h = self._calculate_spatial_tiles(height, target_height, spatial_overlap)
            spatial_tiles_w = self._calculate_spatial_tiles(width, target_width, spatial_overlap)
            
            print(f"⏱️  Temporal chunks: {len(temporal_tiles)}")
            print(f"🗺️  Spatial grid: {len(spatial_tiles_h)}×{len(spatial_tiles_w)} tiles per frame")
            print(f"📦 Total tiles: {len(temporal_tiles) * len(spatial_tiles_h) * len(spatial_tiles_w)}")
            
            # STEP 1: Extract and process all tiles
            print(f"\n🔸 STEP 1: Extracting and processing tiles...")
            all_tiles = self._extract_and_process_tiles(
                video, temporal_tiles, spatial_tiles_h, spatial_tiles_w,
                color_shift_strength, debug_mode
            )
            
            # STEP 2: Column-wise stitching (stitch tiles in same column, same temporal chunk)
            print(f"\n🔸 STEP 2: Column-wise stitching...")
            column_strips = self._stitch_columns(all_tiles, spatial_overlap, debug_mode)
            
            # STEP 3: Line-wise stitching (stitch column strips to complete temporal chunks)
            print(f"\n🔸 STEP 3: Line-wise stitching...")
            temporal_chunks = self._stitch_lines(column_strips, spatial_overlap, debug_mode)
            
            # STEP 4: Time-wise stitching (stitch temporal chunks together)
            print(f"\n🔸 STEP 4: Time-wise stitching...")
            stitched_video = self._stitch_temporal_chunks_new(temporal_chunks, temporal_tiles, frame_overlap)
            
            # STEP 5: Crop back to original dimensions
            print(f"\n🔸 STEP 5: Cropping to original dimensions...")
            print(f"   📐 Stitched video shape: {stitched_video.shape}")
            print(f"   🎯 Target shape: [{batch_size}, {height}, {width}, {channels}]")
            
            # Crop spatial dimensions back to original size
            final_video = stitched_video[:batch_size, :height, :width, :channels]
            
            print(f"   ✂️  Cropped to: {final_video.shape}")
            
            # Generate summary
            tile_info_summary = self._generate_tile_info_summary_new(
                all_tiles, temporal_tiles, spatial_tiles_h, spatial_tiles_w
            )
            
            print(f"✅ Dimension-wise tiling and stitching completed!")
            print(f"📤 Final video shape: {final_video.shape} (cropped to match input)")
            print(f"🧩 Total tiles processed: {len(all_tiles)}")
            print(f"📏 Original input: {batch_size}×{height}×{width}×{channels}")
            print(f"📏 After stitching: {stitched_video.shape}")
            print(f"📏 After cropping: {final_video.shape}")
            print("="*80 + "\n")
            
            return (final_video, tile_info_summary)
            
        except Exception as e:
            print(f"❌ Error in tile and stitch: {str(e)}")
            print(f"📋 Full traceback:")
            print(traceback.format_exc())
            print("="*80 + "\n")
            
            # Return original video in case of error
            dummy_info = f"Error during tiling: {str(e)}"
            return (video, dummy_info)
    
    def _extract_and_process_tiles(self, video, temporal_tiles, spatial_tiles_h, spatial_tiles_w,
                                  color_shift_strength, debug_mode):
        """
        Extract all tiles and apply transformations, storing them in Tile objects
        """
        all_tiles = []
        
        for temporal_idx, (t_start, t_end) in enumerate(temporal_tiles):
            if debug_mode:
                print(f"   🎬 Temporal chunk {temporal_idx + 1}/{len(temporal_tiles)} (frames {t_start}-{t_end-1})")
            
            # Extract temporal chunk
            chunk = video[t_start:t_end]
            
            # Process each spatial tile within this temporal chunk
            for h_idx, (h_start, h_end) in enumerate(spatial_tiles_h):
                for w_idx, (w_start, w_end) in enumerate(spatial_tiles_w):
                    if debug_mode and temporal_idx == 0:  # Only show spatial info for first temporal chunk
                        print(f"      🧩 Spatial tile L{h_idx}:C{w_idx} H[{h_start}:{h_end}] × W[{w_start}:{w_end}]")
                    
                    # Extract spatial tile
                    tile_content = chunk[:, h_start:h_end, w_start:w_end, :]
                    
                    # Apply color transformation (for debugging)
                    tile_transformed = self._apply_color_shift(tile_content, color_shift_strength, 
                                                             temporal_idx, h_idx, w_idx)
                    
                    # Create Tile object
                    tile = Tile(
                        content=tile_transformed,
                        temporal_index=temporal_idx,
                        line_index=h_idx,
                        column_index=w_idx,
                        temporal_range=(t_start, t_end),
                        spatial_range_h=(h_start, h_end),
                        spatial_range_w=(w_start, w_end)
                    )
                    
                    all_tiles.append(tile)
        
        print(f"   ✅ Extracted {len(all_tiles)} tiles")
        return all_tiles
    
    def _stitch_columns(self, all_tiles, spatial_overlap, debug_mode):
        """
        Stitch tiles vertically (same column, same temporal chunk) to create column strips
        """
        # Group tiles by temporal_index and column_index
        column_groups = {}
        for tile in all_tiles:
            key = (tile.temporal_index, tile.column_index)
            if key not in column_groups:
                column_groups[key] = []
            column_groups[key].append(tile)
        
        # Sort each group by line_index
        for key in column_groups:
            column_groups[key].sort(key=lambda t: t.line_index)
        
        column_strips = []
        
        for (temporal_idx, col_idx), column_tiles in column_groups.items():
            if debug_mode and temporal_idx == 0:  # Only show info for first temporal chunk
                print(f"   📏 Column {col_idx}: {len(column_tiles)} tiles")
            
            # Stitch tiles in this column vertically
            if len(column_tiles) == 1:
                # Single tile - no stitching needed
                column_strip = column_tiles[0].content
            else:
                # Multiple tiles - stitch vertically with overlap
                column_strip = self._stitch_tiles_vertically(column_tiles, spatial_overlap)
            
            # Store column strip with its position info
            strip_info = {
                'content': column_strip,
                'temporal_index': temporal_idx,
                'column_index': col_idx,
                'spatial_range_w': column_tiles[0].spatial_range_w,  # All tiles in column have same W range
                'spatial_range_h': (column_tiles[0].spatial_range_h[0], column_tiles[-1].spatial_range_h[1])
            }
            column_strips.append(strip_info)
        
        print(f"   ✅ Created {len(column_strips)} column strips")
        return column_strips
    
    def _stitch_lines(self, column_strips, spatial_overlap, debug_mode):
        """
        Stitch column strips horizontally (same temporal chunk) to create complete temporal chunks
        """
        # Group column strips by temporal_index
        temporal_groups = {}
        for strip in column_strips:
            temporal_idx = strip['temporal_index']
            if temporal_idx not in temporal_groups:
                temporal_groups[temporal_idx] = []
            temporal_groups[temporal_idx].append(strip)
        
        # Sort each group by column_index
        for temporal_idx in temporal_groups:
            temporal_groups[temporal_idx].sort(key=lambda s: s['column_index'])
        
        temporal_chunks = []
        
        for temporal_idx, strips in temporal_groups.items():
            if debug_mode and temporal_idx == 0:  # Only show info for first temporal chunk
                print(f"   📐 Temporal chunk {temporal_idx}: {len(strips)} column strips")
            
            # Stitch column strips horizontally
            if len(strips) == 1:
                # Single strip - no stitching needed
                chunk_content = strips[0]['content']
            else:
                # Multiple strips - stitch horizontally with overlap
                chunk_content = self._stitch_strips_horizontally(strips, spatial_overlap)
            
            # Store temporal chunk with its info
            chunk_info = {
                'content': chunk_content,
                'temporal_index': temporal_idx,
                'spatial_range_h': strips[0]['spatial_range_h'],  # All strips have same H range
                'spatial_range_w': (strips[0]['spatial_range_w'][0], strips[-1]['spatial_range_w'][1])
            }
            temporal_chunks.append(chunk_info)
        
        # Sort by temporal index
        temporal_chunks.sort(key=lambda c: c['temporal_index'])
        
        print(f"   ✅ Created {len(temporal_chunks)} temporal chunks")
        return temporal_chunks
    
    def _stitch_tiles_vertically(self, column_tiles, spatial_overlap):
        """
        Stitch tiles vertically (in height dimension) with proper fade blending
        """
        if len(column_tiles) == 1:
            return column_tiles[0].content
        
        # Get dimensions
        first_tile = column_tiles[0].content
        total_height = sum(tile.spatial_range_h[1] - tile.spatial_range_h[0] for tile in column_tiles)
        # Subtract overlaps
        total_height -= spatial_overlap * (len(column_tiles) - 1)
        
        width = first_tile.shape[2]
        channels = first_tile.shape[3]
        frames = first_tile.shape[0]
        
        # Create result tensor
        result = torch.zeros((frames, total_height, width, channels), 
                           dtype=first_tile.dtype, device=first_tile.device)
        
        current_h = 0
        for i, tile in enumerate(column_tiles):
            tile_h = tile.content.shape[1]
            
            if i == 0:
                # First tile - place directly
                result[:, current_h:current_h+tile_h, :, :] = tile.content
                current_h += tile_h
            else:
                # Subsequent tiles - check for last tile with large overlap
                is_last_tile = (i == len(column_tiles) - 1)
                
                # Calculate actual overlap by comparing tile positions
                expected_start = tile.spatial_range_h[0]
                actual_overlap = current_h - expected_start
                actual_overlap = max(0, actual_overlap)  # Ensure non-negative
                
                if is_last_tile and actual_overlap > spatial_overlap:
                    # Last tile with large overlap - fade across entire overlapping region
                    print(f"      🔥 Last tile in column: Large overlap detected ({actual_overlap} > {spatial_overlap})")
                    overlap_start = expected_start
                    overlap_size = actual_overlap
                    non_overlap_start = current_h
                    non_overlap_size = tile_h - actual_overlap
                    
                    # Place non-overlapping part
                    if non_overlap_size > 0:
                        result[:, non_overlap_start:non_overlap_start+non_overlap_size, :, :] = \
                            tile.content[:, actual_overlap:, :, :]
                    
                    # Handle large overlapping region with fade blending
                    if overlap_size > 0:
                        existing_region = result[:, overlap_start:overlap_start+overlap_size, :, :]
                        new_region = tile.content[:, :overlap_size, :, :]
                        
                        # Create fade mask for entire overlapping region
                        fade_mask = self._create_vertical_fade_mask(overlap_size, existing_region.dtype, existing_region.device)
                        
                        # Apply fade blending
                        blended = existing_region * (1.0 - fade_mask) + new_region * fade_mask
                        result[:, overlap_start:overlap_start+overlap_size, :, :] = blended
                    
                    current_h += non_overlap_size
                else:
                    # Normal tile with standard overlap handling
                    overlap_start = current_h - spatial_overlap
                    
                    # Place non-overlapping part first
                    non_overlap_h = tile_h - spatial_overlap
                    result[:, current_h:current_h+non_overlap_h, :, :] = tile.content[:, spatial_overlap:, :, :]
                    
                    # Handle overlapping region with fade blending
                    if spatial_overlap > 0:
                        existing_region = result[:, overlap_start:current_h, :, :]
                        new_region = tile.content[:, :spatial_overlap, :, :]
                        
                        # Create vertical fade mask
                        fade_mask = self._create_vertical_fade_mask(spatial_overlap, existing_region.dtype, existing_region.device)
                        
                        # Apply fade blending
                        blended = existing_region * (1.0 - fade_mask) + new_region * fade_mask
                        result[:, overlap_start:current_h, :, :] = blended
                    
                    current_h += non_overlap_h
        
        return result
    
    def _stitch_strips_horizontally(self, strips, spatial_overlap):
        """
        Stitch column strips horizontally (in width dimension) with proper fade blending
        """
        if len(strips) == 1:
            return strips[0]['content']
        
        # Get dimensions
        first_strip = strips[0]['content']
        total_width = sum(s['spatial_range_w'][1] - s['spatial_range_w'][0] for s in strips)
        # Subtract overlaps
        total_width -= spatial_overlap * (len(strips) - 1)
        
        height = first_strip.shape[1]
        channels = first_strip.shape[3]
        frames = first_strip.shape[0]
        
        # Create result tensor
        result = torch.zeros((frames, height, total_width, channels), 
                           dtype=first_strip.dtype, device=first_strip.device)
        
        current_w = 0
        for i, strip in enumerate(strips):
            strip_content = strip['content']
            strip_w = strip_content.shape[2]
            
            if i == 0:
                # First strip - place directly
                result[:, :, current_w:current_w+strip_w, :] = strip_content
                current_w += strip_w
            else:
                # Subsequent strips - check for last strip with large overlap
                is_last_strip = (i == len(strips) - 1)
                
                # Calculate actual overlap by comparing strip positions
                expected_start = strip['spatial_range_w'][0]
                actual_overlap = current_w - expected_start
                actual_overlap = max(0, actual_overlap)  # Ensure non-negative
                
                if is_last_strip and actual_overlap > spatial_overlap:
                    # Last strip with large overlap - fade across entire overlapping region
                    print(f"      🔥 Last strip in line: Large overlap detected ({actual_overlap} > {spatial_overlap})")
                    overlap_start = expected_start
                    overlap_size = actual_overlap
                    non_overlap_start = current_w
                    non_overlap_size = strip_w - actual_overlap
                    
                    # Place non-overlapping part
                    if non_overlap_size > 0:
                        result[:, :, non_overlap_start:non_overlap_start+non_overlap_size, :] = \
                            strip_content[:, :, actual_overlap:, :]
                    
                    # Handle large overlapping region with fade blending
                    if overlap_size > 0:
                        existing_region = result[:, :, overlap_start:overlap_start+overlap_size, :]
                        new_region = strip_content[:, :, :overlap_size, :]
                        
                        # Create fade mask for entire overlapping region
                        fade_mask = self._create_horizontal_fade_mask(overlap_size, existing_region.dtype, existing_region.device)
                        
                        # Apply fade blending
                        blended_region = existing_region * (1 - fade_mask) + new_region * fade_mask
                        result[:, :, overlap_start:overlap_start+overlap_size, :] = blended_region
                    
                    current_w += non_overlap_size
                else:
                    # Normal strip with standard overlap handling
                    overlap_start = current_w - spatial_overlap
                    
                    # Place non-overlapping part first
                    non_overlap_w = strip_w - spatial_overlap
                    result[:, :, current_w:current_w+non_overlap_w, :] = strip_content[:, :, spatial_overlap:, :]
                    
                    # Handle overlapping region with fade blending
                    if spatial_overlap > 0:
                        existing_region = result[:, :, overlap_start:current_w, :]
                        new_region = strip_content[:, :, :spatial_overlap, :]
                        
                        # Create horizontal fade mask
                        fade_mask = self._create_horizontal_fade_mask(spatial_overlap, existing_region.dtype, existing_region.device)
                        
                        # Apply fade blending
                        blended_region = existing_region * (1 - fade_mask) + new_region * fade_mask
                        result[:, :, overlap_start:current_w, :] = blended_region
                    
                    current_w += non_overlap_w
            
        return result
            
        return result
    
    def _create_vertical_fade_mask(self, fade_size, dtype, device):
        """Create a vertical fade mask for blending tiles vertically."""
        fade_mask = torch.linspace(0, 1, fade_size, dtype=dtype, device=device)
        # Shape: (1, fade_size, 1, 1) for broadcasting
        return fade_mask.view(1, fade_size, 1, 1)
    
    def _create_horizontal_fade_mask(self, fade_size, dtype, device):
        """Create a horizontal fade mask for blending tiles horizontally."""
        fade_mask = torch.linspace(0, 1, fade_size, dtype=dtype, device=device)
        # Shape: (1, 1, fade_size, 1) for broadcasting
        return fade_mask.view(1, 1, fade_size, 1)
    
    def _stitch_temporal_chunks_new(self, temporal_chunks, temporal_tiles, frame_overlap):
        """Stitch temporal chunks to create the final output with temporal blending."""
        if not temporal_chunks:
            return None
            
        if len(temporal_chunks) == 1:
            return temporal_chunks[0]['content']
        
        # Get dimensions from first chunk
        first_chunk_content = temporal_chunks[0]['content']
        
        # Calculate total frames accounting for overlaps
        total_frames = sum(chunk['content'].shape[0] for chunk in temporal_chunks)
        if frame_overlap > 0:
            total_frames -= (len(temporal_chunks) - 1) * frame_overlap
        
        # Get spatial dimensions from the stitched chunks (not original input)
        height = first_chunk_content.shape[1]
        width = first_chunk_content.shape[2]
        channels = first_chunk_content.shape[3]
        
        print(f"   📏 Temporal stitching dimensions:")
        print(f"      • Total frames: {total_frames}")
        print(f"      • Spatial size: {height}×{width}")
        print(f"      • Channels: {channels}")
        print(f"      • Chunk sizes: {[chunk['content'].shape for chunk in temporal_chunks]}")
        
        # Initialize result tensor with correct dimensions
        result = torch.zeros(
            (total_frames, height, width, channels),
            dtype=first_chunk_content.dtype,
            device=first_chunk_content.device
        )
        
        current_t = 0
        for i, chunk_info in enumerate(temporal_chunks):
            chunk_content = chunk_info['content']
            chunk_frames = chunk_content.shape[0]
            
            if i == 0:
                # First chunk - place entirely
                result[current_t:current_t+chunk_frames] = chunk_content
                current_t += chunk_frames
            else:
                # Subsequent chunks - check for last chunk with large overlap
                is_last_chunk = (i == len(temporal_chunks) - 1)
                
                # Get the temporal range for this chunk from temporal_tiles
                temporal_tile_info = temporal_tiles[i]
                expected_start = temporal_tile_info[0]
                actual_overlap = current_t - expected_start
                actual_overlap = max(0, actual_overlap)  # Ensure non-negative
                
                if is_last_chunk and actual_overlap > frame_overlap:
                    # Last chunk with large overlap - fade across entire overlapping region
                    print(f"      🔥 Last temporal chunk: Large overlap detected ({actual_overlap} > {frame_overlap})")
                    overlap_start = expected_start
                    overlap_size = actual_overlap
                    non_overlap_start = current_t
                    non_overlap_size = chunk_frames - actual_overlap
                    
                    # Place non-overlapping part
                    if non_overlap_size > 0:
                        result[non_overlap_start:non_overlap_start+non_overlap_size] = \
                            chunk_content[actual_overlap:]
                    
                    # Handle large overlapping region with fade blending
                    if overlap_size > 0:
                        existing_region = result[overlap_start:overlap_start+overlap_size]
                        new_region = chunk_content[:overlap_size]
                        
                        # Create temporal fade mask for entire overlapping region
                        fade_mask = self._create_temporal_fade_mask(overlap_size, existing_region.dtype, existing_region.device)
                        
                        # Apply fade blending
                        blended_region = existing_region * (1 - fade_mask) + new_region * fade_mask
                        result[overlap_start:overlap_start+overlap_size] = blended_region
                    
                    current_t += non_overlap_size
                else:
                    # Normal chunk with standard overlap handling
                    non_overlap_frames = chunk_frames - frame_overlap
                    overlap_start = current_t - frame_overlap
                    
                    # Place non-overlapping part
                    result[current_t:current_t+non_overlap_frames] = chunk_content[frame_overlap:]
                    
                    # Handle overlapping region with temporal fade blending
                    if frame_overlap > 0:
                        existing_region = result[overlap_start:current_t]
                        new_region = chunk_content[:frame_overlap]
                        
                        # Create temporal fade mask
                        fade_mask = self._create_temporal_fade_mask(frame_overlap, existing_region.dtype, existing_region.device)
                        
                        # Apply fade blending
                        blended_region = existing_region * (1 - fade_mask) + new_region * fade_mask
                        result[overlap_start:current_t] = blended_region
                    
                    current_t += non_overlap_frames
        
        return result
    
    def _create_temporal_fade_mask(self, fade_size, dtype, device):
        """Create a temporal fade mask for blending across time dimension."""
        fade_mask = torch.linspace(0, 1, fade_size, dtype=dtype, device=device)
        # Shape: (fade_size, 1, 1, 1) for broadcasting
        return fade_mask.view(fade_size, 1, 1, 1)
    
    def _generate_tile_info_summary_new(self, tiles, temporal_tiles, spatial_tiles_h, spatial_tiles_w):
        """Generate a summary of tile processing for the new dimension-wise algorithm."""
        summary_lines = []
        summary_lines.append("=== Tile Processing Summary (Dimension-wise Algorithm) ===")
        summary_lines.append(f"Input video - Temporal chunks: {len(temporal_tiles)}")
        summary_lines.append(f"Input video - Spatial grid: {len(spatial_tiles_h)}×{len(spatial_tiles_w)}")
        summary_lines.append(f"Total tiles processed: {len(tiles)}")
        
        # Group tiles by temporal chunk
        temporal_chunks = {}
        for tile in tiles:
            t_idx = tile.temporal_index
            if t_idx not in temporal_chunks:
                temporal_chunks[t_idx] = []
            temporal_chunks[t_idx].append(tile)
        
        summary_lines.append(f"Temporal chunks: {len(temporal_chunks)}")
        
        # Analyze grid structure
        for t_idx, chunk_tiles in temporal_chunks.items():
            line_groups = {}
            for tile in chunk_tiles:
                l_idx = tile.line_index
                if l_idx not in line_groups:
                    line_groups[l_idx] = []
                line_groups[l_idx].append(tile)
            
            max_cols = max(len(line_tiles) for line_tiles in line_groups.values())
            summary_lines.append(f"  Chunk {t_idx}: {len(line_groups)} lines × {max_cols} columns")
        
        summary_lines.append("Processing order: Column-wise → Line-wise → Temporal → Crop")
        summary_lines.append("Last tiles in each dimension use full-region fade blending")
        summary_lines.append("Final video cropped back to original input dimensions")
        summary_lines.append("="*60)
        
        return "\n".join(summary_lines)
    
    def _calculate_temporal_tiles(self, total_frames, target_frames, overlap):
        """Calculate temporal tile ranges with overlap handling"""
        tiles = []
        
        if total_frames <= target_frames:
            # If video is shorter than target, use the whole video
            tiles.append((0, total_frames))
        else:
            stride = target_frames - overlap
            current = 0
            
            while current < total_frames:
                end = min(current + target_frames, total_frames)
                
                # If this would be the last tile and it doesn't cover remaining frames
                remaining = total_frames - end
                if remaining > 0:
                    # Add this tile normally
                    tiles.append((current, end))
                    # Check if we need another tile to cover the remaining frames
                    if remaining < stride:
                        # Add a final tile that is exactly target_frames but starts from the end
                        final_start = total_frames - target_frames
                        tiles.append((final_start, total_frames))
                        break
                else:
                    # This tile reaches exactly to the end
                    tiles.append((current, end))
                    break
                    
                current += stride
        
        return tiles
    
    def _calculate_spatial_tiles(self, total_size, target_size, overlap):
        """Calculate spatial tile ranges with overlap handling"""
        tiles = []
        
        if total_size <= target_size:
            # If dimension is smaller than target, use the whole dimension
            tiles.append((0, total_size))
        else:
            stride = target_size - overlap
            current = 0
            
            while current < total_size:
                end = min(current + target_size, total_size)
                
                # If this would be the last tile and it doesn't cover remaining pixels
                remaining = total_size - end
                if remaining > 0:
                    # Add this tile normally
                    tiles.append((current, end))
                    # Check if we need another tile to cover the remaining pixels
                    if remaining < stride:
                        # Add a final tile that is exactly target_size but starts from the end
                        final_start = total_size - target_size
                        tiles.append((final_start, total_size))
                        break
                else:
                    # This tile reaches exactly to the end
                    tiles.append((current, end))
                    break
                    
                current += stride
        
        return tiles
    
    def _apply_color_shift(self, tile, strength, temporal_idx, h_idx, w_idx):
        """Apply color transformation to a tile for debugging purposes"""
        if strength <= 0:
            return tile.clone()
        
        # Generate deterministic but varied color shifts based on tile indices
        random.seed(temporal_idx * 1000 + h_idx * 100 + w_idx)
        
        # Random RGB shifts
        r_shift = (random.random() - 0.5) * 2 * strength
        g_shift = (random.random() - 0.5) * 2 * strength  
        b_shift = (random.random() - 0.5) * 2 * strength
        
        transformed = tile.clone()
        
        if tile.shape[-1] >= 3:  # RGB channels
            transformed[:, :, :, 0] = torch.clamp(transformed[:, :, :, 0] + r_shift, 0, 1)
            transformed[:, :, :, 1] = torch.clamp(transformed[:, :, :, 1] + g_shift, 0, 1)
            transformed[:, :, :, 2] = torch.clamp(transformed[:, :, :, 2] + b_shift, 0, 1)
            
            # Debug output to confirm color shift is being applied
            if strength > 0:
                print(f"      🎨 Color shift applied: R{r_shift:+.3f}, G{g_shift:+.3f}, B{b_shift:+.3f}")
        
        return transformed
    
    def _place_tile_with_overlap(self, target, tile, h_start, h_end, w_start, w_end, spatial_overlap):
        """Place tile into target tensor, handling overlaps with spatial fade blending"""
        tile_h, tile_w = tile.shape[1:3]
        
        # Get the target region
        target_region = target[:, h_start:h_start+tile_h, w_start:w_start+tile_w, :]
        
        # Check if there's existing content (non-zero) in the target area
        if target_region.sum() > 0:
            # Always use production-quality fade blending regardless of color shift strength
            # Color shift is purely for visual debugging and shouldn't affect blending quality
            fade_mask = self._create_spatial_fade_mask(tile_h, tile_w, target_region, tile, spatial_overlap)
            
            # Apply fade blending: existing * (1 - fade_mask) + tile * fade_mask
            blended = target_region * (1.0 - fade_mask) + tile * fade_mask
            target[:, h_start:h_start+tile_h, w_start:w_start+tile_w, :] = blended
        else:
            # First tile in this area or non-overlapping area - always place the tile
            # (this includes color-shifted tiles when debugging is enabled)
            target[:, h_start:h_start+tile_h, w_start:w_start+tile_w, :] = tile
    
    def _create_spatial_fade_mask(self, tile_h, tile_w, existing, new_tile, spatial_overlap):
        """Create a spatial fade mask for smooth blending between tiles"""
        # Create base mask (1.0 = use new tile, 0.0 = use existing)
        mask = torch.ones(1, tile_h, tile_w, 1, dtype=existing.dtype, device=existing.device)
        
        # Define fade distance using spatial_overlap parameter
        fade_h = min(spatial_overlap, tile_h // 2)  # Fade over spatial_overlap pixels or half tile height
        fade_w = min(spatial_overlap, tile_w // 2)  # Fade over spatial_overlap pixels or half tile width
        
        # Check which edges have existing content to determine fade direction
        has_top = existing[:, :fade_h, :, :].sum() > 0
        has_bottom = existing[:, -fade_h:, :, :].sum() > 0
        has_left = existing[:, :, :fade_w, :].sum() > 0
        has_right = existing[:, :, -fade_w:, :].sum() > 0
        
        # Create fade gradients for each edge that has existing content
        if has_top:
            # Fade from 0 at top to 1 after fade_h pixels
            for i in range(fade_h):
                alpha = i / fade_h
                mask[:, i, :, :] = alpha
        
        if has_bottom:
            # Fade from 1 before last fade_h pixels to 0 at bottom
            for i in range(fade_h):
                alpha = 1.0 - (i / fade_h)
                mask[:, tile_h - 1 - i, :, :] = alpha
        
        if has_left:
            # Fade from 0 at left to 1 after fade_w pixels
            for i in range(fade_w):
                alpha = i / fade_w
                mask[:, :, i, :] = torch.minimum(mask[:, :, i, :], torch.tensor(alpha, dtype=mask.dtype, device=mask.device))
        
        if has_right:
            # Fade from 1 before last fade_w pixels to 0 at right
            for i in range(fade_w):
                alpha = 1.0 - (i / fade_w)
                mask[:, :, tile_w - 1 - i, :] = torch.minimum(mask[:, :, tile_w - 1 - i, :], torch.tensor(alpha, dtype=mask.dtype, device=mask.device))
        
        # Broadcast mask to match tile dimensions
        return mask.expand_as(new_tile)
    
    def _stitch_temporal_chunks(self, chunks, temporal_tiles, total_frames, height, width, channels, overlap):
        """Stitch temporal chunks back together with temporal fade blending"""
        result = torch.zeros((total_frames, height, width, channels), dtype=chunks[0].dtype, device=chunks[0].device)
        
        for i, ((t_start, t_end), chunk) in enumerate(zip(temporal_tiles, chunks)):
            chunk_frames = chunk.shape[0]
            
            if i == 0:
                # First chunk, place directly
                result[t_start:t_start+chunk_frames] = chunk
            else:
                # Handle overlap with previous chunk using temporal fade
                prev_end = temporal_tiles[i-1][1]
                overlap_start = max(t_start, prev_end - overlap)
                overlap_end = min(t_end, prev_end)
                
                if overlap_start < overlap_end:
                    # There's an overlap, apply temporal fade blending
                    overlap_frames = overlap_end - overlap_start
                    chunk_overlap_start = overlap_start - t_start
                    
                    # Get overlapping regions
                    existing_frames = result[overlap_start:overlap_end]
                    new_frames = chunk[chunk_overlap_start:chunk_overlap_start+overlap_frames]
                    
                    # Create temporal fade mask
                    fade_mask = self._create_temporal_fade_mask(overlap_frames, existing_frames.dtype, existing_frames.device)
                    
                    # Apply temporal fade: existing * (1 - fade) + new * fade
                    blended_frames = existing_frames * (1.0 - fade_mask) + new_frames * fade_mask
                    result[overlap_start:overlap_end] = blended_frames
                    
                    # Place non-overlapping frames
                    if overlap_end < t_start + chunk_frames:
                        non_overlap_start = overlap_end
                        chunk_offset = non_overlap_start - t_start
                        result[non_overlap_start:t_start+chunk_frames] = chunk[chunk_offset:]
                else:
                    # No overlap, place directly
                    result[t_start:t_start+chunk_frames] = chunk
        
        return result
    
    def _create_temporal_fade_mask(self, overlap_frames, dtype, device):
        """Create a temporal fade mask for smooth frame transitions"""
        # Create fade mask: 0.0 at start (keep existing) to 1.0 at end (use new)
        fade_values = torch.linspace(0.0, 1.0, overlap_frames, dtype=dtype, device=device)
        
        # Reshape to broadcast properly: [frames, 1, 1, 1]
        fade_mask = fade_values.view(overlap_frames, 1, 1, 1)
        
        return fade_mask
    
    def _generate_tile_info_summary(self, tile_info_list, temporal_tiles, spatial_tiles_h, spatial_tiles_w):
        """Generate a summary of the tiling process"""
        summary = f"=== TILE AND STITCH SUMMARY ===\n"
        summary += f"Total tiles processed: {len(tile_info_list)}\n"
        summary += f"Temporal chunks: {len(temporal_tiles)}\n"
        summary += f"Spatial tiles per frame: {len(spatial_tiles_h)}×{len(spatial_tiles_w)}\n\n"
        
        summary += "Temporal chunks:\n"
        for i, (start, end) in enumerate(temporal_tiles):
            summary += f"  Chunk {i+1}: frames {start}-{end-1} ({end-start} frames)\n"
        
        summary += f"\nSpatial tiles (Height):\n"
        for i, (start, end) in enumerate(spatial_tiles_h):
            summary += f"  Row {i+1}: pixels {start}-{end-1} ({end-start} pixels)\n"
        
        summary += f"\nSpatial tiles (Width):\n"
        for i, (start, end) in enumerate(spatial_tiles_w):
            summary += f"  Col {i+1}: pixels {start}-{end-1} ({end-start} pixels)\n"
        
        summary += f"\nProcessing completed successfully!"
        
        return summary


class TiledWanVideoVACEpipe:
    """
    Dimension-wise Tiled WanVideo VACE Pipeline - The ultimate node for processing large videos.
    
    This node combines the sophisticated dimension-wise tiling system from TileAndStitchBack with 
    the complete WanVideo VACE pipeline. It handles large videos using the advanced algorithm:
    
    1. Temporal tiling: Split video into chunks with overlap (default: 81 frames, 10-frame overlap)
    2. Spatial tiling: Split each chunk into tiles with overlap (default: 832×480, 20-pixel overlap)
    3. WanVideo VACE processing: Process each tile through the complete pipeline
    4. Dimension-wise stitching:
       a) Column-wise: Stitch tiles vertically within each temporal chunk
       b) Line-wise: Stitch columns horizontally to complete temporal chunks  
       c) Temporal: Stitch temporal chunks together across time
    5. STRICT OVERLAP CONTROL: Overlap blending is limited to specified parameters
       - Spatial overlap never exceeds spatial_overlap pixels
       - Temporal overlap never exceeds frame_overlap frames
       - Extra pixels/frames are ignored (they help generation but don't affect output)
    6. Cropping: Final video cropped to exact input dimensions
    7. Memory management: Proper model offloading between tiles to prevent leaks
    
    This enables processing of arbitrarily large videos that would otherwise exceed memory limits,
    while maintaining WanVideo's optimal performance and eliminating overlap artifacts through
    sophisticated dimension-wise fade blending with strict overlap control.
    """
    
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        """
        Complete inputs combining tiling parameters with WanVideo VACE pipeline
        """
        return {
            "required": {
                # Input video and mask
                "video": ("IMAGE", {"tooltip": "Input video frames to process. Can be very large videos that will be automatically tiled for memory efficiency."}),
                "mask": ("MASK", {"tooltip": "Mask defining areas to be processed. Should match the video dimensions and frame count."}),
                
                # Core WanVideo pipeline inputs
                "model": ("WANVIDEOMODEL", {"tooltip": "WanVideo model for video generation/processing."}),
                "vae": ("WANVAE", {"tooltip": "Video VAE for encoding/decoding frames."}),
                
                # Tiling parameters
                "target_frames": ("INT", {"default": 81, "min": 16, "max": 200, "step": 1, "tooltip": "Target number of frames per temporal chunk. The model works best with 81 frames. If different, might generate artefact or frame drops"}),
                "target_width": ("INT", {"default": 832, "min": 64, "max": 2048, "step": 8, "tooltip": "Target width for spatial tiles. Works best around 832."}),
                "target_height": ("INT", {"default": 480, "min": 64, "max": 2048, "step": 8, "tooltip": "Target height for spatial tiles. Works best around 480."}),
                "frame_overlap": ("INT", {"default": 10, "min": 0, "max": 40, "step": 1, "tooltip": "Number of overlapping frames between temporal chunks."}),
                "spatial_overlap": ("INT", {"default": 20, "min": 0, "max": 100, "step": 4, "tooltip": "Pixel overlap between spatial tiles."}),
                
                # WanVideoSampler core parameters
                "steps": ("INT", {"default": 30, "min": 1, "tooltip": "Number of diffusion sampling steps."}),
                "cfg": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 30.0, "step": 0.01, "tooltip": "Classifier-free guidance strength. Higher values follow the conditioning more closely."}),
                "shift": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 1000.0, "step": 0.01, "tooltip": "Shift parameter for the sampling schedule."}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Random seed for reproducible results. Use the same seed to get identical outputs."}),
                "scheduler": (["unipc", "unipc/beta", "dpm++", "dpm++/beta","dpm++_sde", "dpm++_sde/beta", "euler", "euler/beta", "euler/accvideo", "deis", "lcm", "lcm/beta", "flowmatch_causvid", "flowmatch_distill", "multitalk"],
                    {"default": 'unipc', "tooltip": "Sampling scheduler algorithm."}),
                
                # WanVideoVACEEncode parameters  
                "vace_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "VACE encoding strength. Controls how much the video is modified during processing. 1.0 mean total reconstruction, 0.0 means no change."}),
                "vace_start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "vace_end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                
                # WanVideoDecode parameters
                "decode_enable_vae_tiling": ("BOOLEAN", {"default": False}),
                "decode_tile_x": ("INT", {"default": 272, "min": 40, "max": 2048, "step": 8}),
                "decode_tile_y": ("INT", {"default": 272, "min": 40, "max": 2048, "step": 8}),
                "decode_tile_stride_x": ("INT", {"default": 144, "min": 32, "max": 2040, "step": 8}),
                "decode_tile_stride_y": ("INT", {"default": 128, "min": 32, "max": 2040, "step": 8}),
                
                # Processing parameters
                "debug_mode": ("BOOLEAN", {"default": True}),
                "debug_color_shift": ("BOOLEAN", {"default": True}),
                "force_offload_between_tiles": ("BOOLEAN", {"default": True}),
            },
            "optional": {
                # WanVideoSampler optional inputs
                "text_embeds": ("WANVIDEOTEXTEMBEDS",),
                "samples": ("LATENT",),
                "riflex_freq_index": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
                "denoise_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                "force_offload": ("BOOLEAN", {"default": True}),
                "batched_cfg": ("BOOLEAN", {"default": False}),
                "rope_function": (["default", "comfy"], {"default": "comfy"}),
                
                # WanVideoSampler advanced optional inputs
                "feta_args": ("FETAARGS",),
                "context_options": ("CONTEXTOPTIONS",),
                "loop_args": ("LOOPARGS",),
                "sigmas": ("SIGMAS",),
                "unianimate_poses": ("UNIANIMATE_POSES",),
                "fantasytalking_embeds": ("FANTASYTALKING_EMBEDS",),
                "uni3c_embeds": ("UNI3C_EMBEDS",),
                "multitalk_embeds": ("MULTITALK_EMBEDS",),
                "freeinit_args": ("FREEINIT_ARGS",),
                
                # External pre-built arguments (for modularity)
                "cache_args": ("CACHEARGS",),
                "slg_args": ("SLGARGS",),
                "experimental_args": ("EXPERIMENTALARGS",),
                
                # WanVideoVACEEncode optional inputs  
                "vace_ref_images": ("IMAGE",),
                "vace_tiled_vae": ("BOOLEAN", {"default": False}),
                
                # WanVideoDecode optional inputs
                "decode_normalization": (["default", "minmax"], {"default": "default"}),
            },
        }

    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("processed_video", "processing_info")
    OUTPUT_TOOLTIPS = ("Processed video with WanVideo VACE pipeline applied through tiled processing. Maintains original dimensions and quality.", 
                      "Detailed information about the processing including tile counts, timing, and processing summary.")
    FUNCTION = "process_tiled_wanvideo"
    OUTPUT_NODE = True
    CATEGORY = "TiledWan"
    DESCRIPTION = """
    Video processing node that tiles WanVideo VACE pipeline.
    
    Enables processing of arbitrarily large videos that would otherwise exceed model limits.
    Uses dimension-wise tiling with temporal and spatial overlap to maintain
    quality and eliminate artifacts. The complete WanVideo VACE pipeline is applied to each
    tile with proper stitching and blending for seamless results.
    Give last chunk frame as reference for the upcoming generation in order to maintain consistency across temporal chunks

    Key Features:
    - Handles any video size through tiling
    - Temporal consistency across chunks thanks to cross-chunk reference
    - Spatial seamless stitching with fade blending
    - STRICT OVERLAP CONTROL: Overlap never exceeds specified parameters
    - Memory-efficient processing with model offloading
    - Complete WanVideo VACE pipeline integration
    - Debug modes for monitoring and troubleshooting
    
    Overlap Behavior:
    - Overlap blending is limited to frame_overlap and spatial_overlap parameters
    - When actual overlap exceeds specified limits, blending starts from the correct position
    - Extra pixels/frames are skipped (they help generation but don't affect output)
    - This ensures consistent results regardless of tile positioning
    """

    def process_tiled_wanvideo(self, video, mask, model, vae, target_frames, target_width, target_height, 
                              frame_overlap, spatial_overlap, steps, cfg, shift, seed, scheduler,
                              vace_strength, vace_start_percent, vace_end_percent, 
                              decode_enable_vae_tiling, decode_tile_x, decode_tile_y,
                              decode_tile_stride_x, decode_tile_stride_y, debug_mode, debug_color_shift,
                              force_offload_between_tiles, **kwargs):
        """
        Process large video through tiled WanVideo VACE pipeline with dimension-wise stitching
        """
        
        print("\n" + "="*80)
        print("              TILED WANVIDEO VACE PIPELINE (DIMENSION-WISE)")
        print("="*80)
        print("🚀 Starting dimension-wise tiled WanVideo VACE processing...")
        
        try:
            # Import WanVideo nodes
            custom_nodes_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "ComfyUI-WanVideoWrapper")
            if custom_nodes_path not in sys.path:
                sys.path.append(custom_nodes_path)
            
            parent_path = os.path.dirname(custom_nodes_path)
            if parent_path not in sys.path:
                sys.path.insert(0, parent_path)
            
            package_name = os.path.basename(custom_nodes_path)
            wanvideo_package = importlib.import_module(f"{package_name}.nodes")
            
            WanVideoSampler = wanvideo_package.WanVideoSampler
            WanVideoVACEEncode = wanvideo_package.WanVideoVACEEncode
            WanVideoDecode = wanvideo_package.WanVideoDecode
            print("✅ WanVideo nodes imported successfully!")
            
            # Input validation and preprocessing
            batch_size, height, width, channels = video.shape
            mask_batch, mask_height, mask_width = mask.shape
            
            print(f"📹 Input video shape: {video.shape} (B×H×W×C)")
            print(f"🎭 Input mask shape: {mask.shape} (B×H×W)")
            print(f"🎯 Target tile size: {target_frames} frames × {target_width}×{target_height}")
            print(f"🔗 Overlaps: {frame_overlap} frames, {spatial_overlap} pixels")
            
            # Validate mask dimensions
            if mask_batch != batch_size or mask_height != height or mask_width != width:
                print(f"⚠️  Warning: Mask dimensions {mask.shape} don't match video {video.shape}")
                print("   Resizing mask to match video...")
                mask = torch.nn.functional.interpolate(
                    mask.unsqueeze(1), size=(height, width), mode='nearest'
                ).squeeze(1)
                print(f"✅ Mask resized to: {mask.shape}")
            
            # STEP 1: Calculate tile dimensions using improved algorithm
            print(f"\n📏 STEP 1: Calculating optimal tile layout...")
            temporal_tiles = self._calculate_temporal_tiles(batch_size, target_frames, frame_overlap)
            spatial_tiles_h = self._calculate_spatial_tiles(height, target_height, spatial_overlap) 
            spatial_tiles_w = self._calculate_spatial_tiles(width, target_width, spatial_overlap)
            
            total_tiles = len(temporal_tiles) * len(spatial_tiles_h) * len(spatial_tiles_w)
            print(f"⏱️  Temporal chunks: {len(temporal_tiles)}")
            print(f"🗺️  Spatial tiles per frame: {len(spatial_tiles_h)}×{len(spatial_tiles_w)}")
            print(f"📦 Total tiles to process: {total_tiles}")
            
            # STEP 2: Extract all tiles and process through WanVideo VACE pipeline
            print(f"\n🧩 STEP 2: Extracting and processing {total_tiles} tiles through WanVideo VACE...")
            all_tiles = self._extract_and_process_wanvideo_tiles(
                video, mask, temporal_tiles, spatial_tiles_h, spatial_tiles_w,
                model, vae, WanVideoVACEEncode, WanVideoSampler, WanVideoDecode,
                steps, cfg, shift, seed, scheduler, vace_strength, vace_start_percent, vace_end_percent,
                decode_enable_vae_tiling, decode_tile_x, decode_tile_y, decode_tile_stride_x, decode_tile_stride_y,
                debug_color_shift, force_offload_between_tiles, debug_mode, frame_overlap, spatial_overlap, kwargs
            )
            
            # STEP 3: Dimension-wise stitching - Column-wise (vertical stitching)
            print(f"\n🔄 STEP 3: Column-wise stitching (vertical)...")
            column_strips = self._stitch_columns(all_tiles, spatial_overlap, debug_mode)
            
            # STEP 4: Dimension-wise stitching - Line-wise (horizontal stitching)  
            print(f"\n🔄 STEP 4: Line-wise stitching (horizontal)...")
            temporal_chunks = self._stitch_lines(column_strips, spatial_overlap, debug_mode)
            
            # STEP 5: Dimension-wise stitching - Temporal stitching
            print(f"\n🔄 STEP 5: Temporal stitching...")
            stitched_video = self._stitch_temporal_chunks_new(temporal_chunks, temporal_tiles, frame_overlap)
            
            # STEP 6: Crop to original dimensions (ensure exact input size)
            print(f"\n✂️ STEP 6: Cropping to original dimensions...")
            print(f"📐 Stitched video shape: {stitched_video.shape}")
            print(f"🎯 Target shape: {video.shape}")
            final_video = stitched_video[:batch_size, :height, :width, :channels]
            print(f"✂️ Cropped to: {final_video.shape}")
            
            # Generate comprehensive processing info
            processing_summary = self._generate_wanvideo_processing_summary(
                all_tiles, temporal_tiles, spatial_tiles_h, spatial_tiles_w, total_tiles
            )
            
            print(f"✅ Dimension-wise tiled WanVideo VACE processing completed!")
            print(f"📤 Final video shape: {final_video.shape}")
            print(f"🧩 Total tiles processed: {len(all_tiles)}")
            successful_tiles = sum(1 for tile in all_tiles if hasattr(tile, 'processing_status') and tile.processing_status == 'success')
            print(f"✅ Successful tiles: {successful_tiles}/{total_tiles}")
            print("="*80 + "\n")
            
            return (final_video, processing_summary)
            
        except Exception as e:
            print(f"❌ Error in tiled WanVideo VACE pipeline: {str(e)}")
            print(f"📋 Full traceback:")
            import traceback
            print(traceback.format_exc())
            print("="*80 + "\n")
            
            # Return original video in case of error
            error_info = f"Error during tiled WanVideo processing: {str(e)}"
            return (video, error_info)
    
    
    def _extract_and_process_wanvideo_tiles(self, video, mask, temporal_tiles, spatial_tiles_h, spatial_tiles_w,
                                        model, vae, WanVideoVACEEncode, WanVideoSampler, WanVideoDecode,
                                        steps, cfg, shift, seed, scheduler, vace_strength, vace_start_percent, vace_end_percent,
                                        decode_enable_vae_tiling, decode_tile_x, decode_tile_y, decode_tile_stride_x, decode_tile_stride_y,
                                        debug_color_shift, force_offload_between_tiles, debug_mode, frame_overlap, spatial_overlap, kwargs):
        """
        Extract all tiles and process through WanVideo VACE pipeline with frame-wise temporal consistency
        """
        all_tiles = []
        
        # Store previous temporal chunk's COMPLETE STITCHED FRAME for reference
        previous_chunk_stitched_frame = None
        
        for temporal_idx, (t_start, t_end) in enumerate(temporal_tiles):
            if debug_mode:
                print(f"   🎬 Processing temporal chunk {temporal_idx + 1}/{len(temporal_tiles)} (frames {t_start}-{t_end-1})")
            
            # Extract temporal chunk
            video_chunk = video[t_start:t_end]
            mask_chunk = mask[t_start:t_end]
            
            # Store current chunk tiles for stitching the reference frame
            current_chunk_tiles = []
            
            # Process each spatial tile within this temporal chunk
            for h_idx, (h_start, h_end) in enumerate(spatial_tiles_h):
                for w_idx, (w_start, w_end) in enumerate(spatial_tiles_w):
                    tile_idx = len(all_tiles)  # Global tile index
                    
                    if debug_mode:
                        print(f"      🧩 Tile {tile_idx + 1}: T{temporal_idx}:L{h_idx}:C{w_idx} "
                            f"H[{h_start}:{h_end}] × W[{w_start}:{w_end}]")
                    
                    # Extract spatial tiles from temporal chunk
                    video_tile = video_chunk[:, h_start:h_end, w_start:w_end, :]
                    mask_tile = mask_chunk[:, h_start:h_end, w_start:w_end]
                    
                    # ========== TEMPORAL OVERWRITING FOR CONSISTENCY ==========
                    if temporal_idx > 0 and previous_chunk_stitched_frame is not None:
                        # Determine how many frames to overwrite
                        frames_to_overwrite = frame_overlap
                        if temporal_idx == len(temporal_tiles) - 1:
                            # Last chunk: use more frames if chunk is large enough
                            available_frames = video_tile.shape[0]
                            frames_to_overwrite = min(available_frames - 1, max(frame_overlap, available_frames // 3))
                        
                        # Extract corresponding frames from previous chunk's stitched result
                        prev_chunk_frames = previous_chunk_stitched_frame.shape[0]
                        source_start_idx = max(0, prev_chunk_frames - frames_to_overwrite)
                        source_frames = previous_chunk_stitched_frame[source_start_idx:, h_start:h_end, w_start:w_end, :]
                        
                        # Overwrite first frames of current tile with last frames from previous chunk
                        actual_overwrite = min(frames_to_overwrite, source_frames.shape[0], video_tile.shape[0])
                        if actual_overwrite > 0:
                            video_tile[:actual_overwrite] = source_frames[:actual_overwrite]
                            # Zero out corresponding mask areas
                            mask_tile[:actual_overwrite] = 0.0
                            
                            if debug_mode:
                                print(f"         🔄 Temporal overwrite: {actual_overwrite} frames from prev chunk")
                    
                    # ========== SPATIAL OVERWRITING FOR CONSISTENCY ==========
                    # Look for already processed tiles to overwrite overlapping regions
                    
                    # Check for LEFT neighbor (same temporal and line, previous column)
                    if w_idx > 0:
                        left_neighbor = None
                        for tile in current_chunk_tiles:
                            if (tile.temporal_index == temporal_idx and 
                                tile.line_index == h_idx and 
                                tile.column_index == w_idx - 1):
                                left_neighbor = tile
                                break
                        
                        if left_neighbor is not None and hasattr(left_neighbor, 'content'):
                            # Calculate overlap region
                            left_tile_end = left_neighbor.spatial_range_w[1]
                            current_tile_start = w_start
                            overlap_width = left_tile_end - current_tile_start
                            
                            if overlap_width > 0:
                                # Extract overlap from left neighbor's right edge
                                left_content = left_neighbor.content
                                left_overlap = left_content[:, :, -overlap_width:, :]  # Last overlap_width columns
                                
                                # Overwrite left edge of current tile
                                video_tile[:, :, :overlap_width, :] = left_overlap
                                # Zero out corresponding mask areas
                                mask_tile[:, :, :overlap_width] = 0.0
                                
                                if debug_mode:
                                    print(f"         ↔️ Spatial overwrite LEFT: {overlap_width} pixels from neighbor")
                    
                    # Check for TOP neighbor (same temporal and column, previous line)
                    if h_idx > 0:
                        top_neighbor = None
                        for tile in current_chunk_tiles:
                            if (tile.temporal_index == temporal_idx and 
                                tile.line_index == h_idx - 1 and 
                                tile.column_index == w_idx):
                                top_neighbor = tile
                                break
                        
                        if top_neighbor is not None and hasattr(top_neighbor, 'content'):
                            # Calculate overlap region
                            top_tile_end = top_neighbor.spatial_range_h[1]
                            current_tile_start = h_start
                            overlap_height = top_tile_end - current_tile_start
                            
                            if overlap_height > 0:
                                # Extract overlap from top neighbor's bottom edge
                                top_content = top_neighbor.content
                                top_overlap = top_content[:, -overlap_height:, :, :]  # Last overlap_height rows
                                
                                # Overwrite top edge of current tile
                                video_tile[:, :overlap_height, :, :] = top_overlap
                                # Zero out corresponding mask areas
                                mask_tile[:, :overlap_height, :] = 0.0
                                
                                if debug_mode:
                                    print(f"         ↕️ Spatial overwrite TOP: {overlap_height} pixels from neighbor")
                    
                    # Handle TOP-LEFT corner conflict (if both top and left neighbors exist)
                    if h_idx > 0 and w_idx > 0:
                        # Find top-left diagonal neighbor
                        topleft_neighbor = None
                        for tile in current_chunk_tiles:
                            if (tile.temporal_index == temporal_idx and 
                                tile.line_index == h_idx - 1 and 
                                tile.column_index == w_idx - 1):
                                topleft_neighbor = tile
                                break
                        
                        if topleft_neighbor is not None and hasattr(topleft_neighbor, 'content'):
                            # Calculate corner overlap
                            top_tile_end_h = topleft_neighbor.spatial_range_h[1]
                            left_tile_end_w = topleft_neighbor.spatial_range_w[1]
                            current_start_h = h_start
                            current_start_w = w_start
                            
                            overlap_h = top_tile_end_h - current_start_h
                            overlap_w = left_tile_end_w - current_start_w
                            
                            if overlap_h > 0 and overlap_w > 0:
                                # Extract corner from top-left neighbor
                                topleft_content = topleft_neighbor.content
                                corner_overlap = topleft_content[:, -overlap_h:, -overlap_w:, :]
                                
                                # Overwrite corner of current tile
                                video_tile[:, :overlap_h, :overlap_w, :] = corner_overlap
                                # Zero out corresponding mask areas
                                mask_tile[:, :overlap_h, :overlap_w] = 0.0
                                
                                if debug_mode:
                                    print(f"         ↖️ Spatial overwrite CORNER: {overlap_h}x{overlap_w} from diagonal neighbor")
                    
                    # FRAME-WISE TEMPORAL CONSISTENCY: Determine reference image for this tile
                    tile_ref_images = None
                    
                    if temporal_idx == 0:
                        # First temporal chunk: Use user-provided ref_images (if any)
                        tile_ref_images = kwargs.get("vace_ref_images")
                        if debug_mode and h_idx == 0 and w_idx == 0:  # Only print once per chunk
                            print(f"      🎯 First chunk: Using user-provided ref_images")
                    else:
                        # Subsequent chunks: Use COMPLETE STITCHED FRAME from previous chunk as reference
                        if previous_chunk_stitched_frame is not None:
                            # Extract reference frame: frame_overlap frames before the last frame
                            ref_frame_idx = previous_chunk_stitched_frame.shape[0] - frame_overlap - 1
                            ref_frame_idx = max(0, ref_frame_idx)  # Ensure non-negative
                            
                            # Extract single COMPLETE FRAME as reference image
                            tile_ref_images = previous_chunk_stitched_frame[ref_frame_idx:ref_frame_idx+1]  # Shape: [1, H, W, C]
                            
                            if debug_mode and h_idx == 0 and w_idx == 0:  # Only print once per chunk
                                print(f"      🔗 Frame-wise temporal chain: Using complete stitched frame {ref_frame_idx} from previous chunk T{temporal_idx-1}")
                                print(f"         Reference frame shape: {tile_ref_images.shape}")
                        else:
                            # Fallback: Use user-provided ref_images
                            tile_ref_images = kwargs.get("vace_ref_images")
                            if debug_mode and h_idx == 0 and w_idx == 0:
                                print(f"      ⚠️  No previous stitched frame found, using user ref_images as fallback")
                    
                    # Process tile through WanVideo VACE pipeline with frame-wise temporal reference
                    try:
                        # Create kwargs with tile-specific reference image (COMPLETE FRAME)
                        tile_kwargs = kwargs.copy()
                        tile_kwargs["vace_ref_images"] = tile_ref_images
                        
                        processed_tile, tile_latents = self._process_tile_through_wanvideo(
                            video_tile, mask_tile, model, vae, 
                            WanVideoVACEEncode, WanVideoSampler, WanVideoDecode,
                            steps, cfg, shift, seed + tile_idx,  # Unique seed per tile
                            scheduler, vace_strength, vace_start_percent, vace_end_percent,
                            decode_enable_vae_tiling, decode_tile_x, decode_tile_y,
                            decode_tile_stride_x, decode_tile_stride_y, debug_color_shift, tile_kwargs
                        )
                        
                        # Create Tile object with processed content
                        tile = Tile(
                            content=processed_tile,
                            temporal_index=temporal_idx,
                            line_index=h_idx,
                            column_index=w_idx,
                            temporal_range=(t_start, t_end),
                            spatial_range_h=(h_start, h_end),
                            spatial_range_w=(w_start, w_end)
                        )
                        
                        # Store processing status and reference info for tracking
                        tile.processing_status = 'success'
                        tile.seed_used = seed + tile_idx
                        tile.latents = tile_latents
                        tile.ref_frame_source = f"complete_frame_T{temporal_idx-1}" if temporal_idx > 0 and previous_chunk_stitched_frame is not None else "user_provided"
                        
                        all_tiles.append(tile)
                        current_chunk_tiles.append(tile)
                        
                        # Force memory cleanup between tiles if requested
                        if force_offload_between_tiles:
                            self._force_memory_cleanup(model, vae)
                            
                    except Exception as tile_error:
                        print(f"         ❌ Error processing tile {tile_idx + 1}: {str(tile_error)}")
                        
                        # Create fallback Tile object with original content
                        tile = Tile(
                            content=video_tile,  # Use original video tile as fallback
                            temporal_index=temporal_idx,
                            line_index=h_idx,
                            column_index=w_idx,
                            temporal_range=(t_start, t_end),
                            spatial_range_h=(h_start, h_end),
                            spatial_range_w=(w_start, w_end)
                        )
                        
                        # Store error status for tracking
                        tile.processing_status = f'failed: {str(tile_error)}'
                        tile.seed_used = seed + tile_idx
                        tile.latents = None
                        tile.ref_frame_source = "error_fallback"
                        
                        all_tiles.append(tile)
                        current_chunk_tiles.append(tile)
            
            # STITCH THE CURRENT CHUNK TO CREATE REFERENCE FOR NEXT CHUNK
            if debug_mode:
                print(f"      🔧 Stitching current chunk T{temporal_idx} for next chunk's reference...")
            
            # Group current chunk tiles for stitching
            current_chunk_tiles_for_stitching = [tile for tile in current_chunk_tiles]
            
            # Perform dimension-wise stitching for this temporal chunk
            try:
                # Step 1: Column-wise stitching
                column_strips = self._stitch_columns(current_chunk_tiles_for_stitching, spatial_overlap, False)  # No debug for reference stitching
                
                # Step 2: Line-wise stitching  
                temporal_chunks = self._stitch_lines(column_strips, spatial_overlap, False)  # No debug for reference stitching
                
                # Get the stitched chunk content (should be only one chunk)
                if temporal_chunks and len(temporal_chunks) > 0:
                    current_chunk_stitched = temporal_chunks[0]['content']
                    
                    # Store this stitched chunk as reference for next temporal chunk
                    previous_chunk_stitched_frame = current_chunk_stitched
                    
                    if debug_mode:
                        print(f"      ✅ Chunk T{temporal_idx} stitched for reference: {current_chunk_stitched.shape}")
                else:
                    print(f"      ⚠️  Warning: Failed to stitch chunk T{temporal_idx} for reference")
                    previous_chunk_stitched_frame = None
                    
            except Exception as stitching_error:
                print(f"      ❌ Error stitching chunk T{temporal_idx} for reference: {str(stitching_error)}")
                previous_chunk_stitched_frame = None
            
            if debug_mode:
                chain_status = "user_ref" if temporal_idx == 0 else f"complete_frame_from_T{temporal_idx-1}"
                print(f"      ✅ Temporal chunk {temporal_idx} completed with {chain_status} references")
        
        print(f"   ✅ Extracted and processed {len(all_tiles)} tiles through WanVideo VACE pipeline")
        successful_tiles = sum(1 for tile in all_tiles if tile.processing_status == 'success')
        print(f"   🎯 Success rate: {successful_tiles}/{len(all_tiles)} ({(successful_tiles/len(all_tiles))*100:.1f}%)")
        
        # Debug: Print overwriting statistics
        if debug_mode:
            temporal_overwrites = sum(1 for t_idx in range(1, len(temporal_tiles)))
            spatial_overwrites = 0
            corner_overwrites = 0
            
            for temporal_idx in range(len(temporal_tiles)):
                for h_idx in range(len(spatial_tiles_h)):
                    for w_idx in range(len(spatial_tiles_w)):
                        if w_idx > 0:  # Has left neighbor
                            spatial_overwrites += 1
                        if h_idx > 0:  # Has top neighbor
                            spatial_overwrites += 1
                        if h_idx > 0 and w_idx > 0:  # Has corner conflict
                            corner_overwrites += 1
            
            print(f"   🔄 Overwriting consistency enhancements:")
            print(f"      • Temporal overwrites: {temporal_overwrites * len(spatial_tiles_h) * len(spatial_tiles_w)} tile operations")
            print(f"      • Spatial overwrites: {spatial_overwrites} edge operations") 
            print(f"      • Corner overwrites: {corner_overwrites} corner conflict resolutions")
        
        # Debug: Print frame-wise temporal consistency chain summary
        if debug_mode:
            print(f"   🖼️  Frame-wise temporal consistency chain summary:")
            print(f"      • Chunk 0: Used user-provided reference images")
            for t_idx in range(1, len(temporal_tiles)):
                print(f"      • Chunk {t_idx}: Used complete stitched frame from Chunk {t_idx-1}")
        
        return all_tiles
    
    def _process_tile_through_wanvideo(self, video_tile, mask_tile, model, vae,
                                    WanVideoVACEEncode, WanVideoSampler, WanVideoDecode,
                                    steps, cfg, shift, seed, scheduler, vace_strength,
                                    vace_start_percent, vace_end_percent, decode_enable_vae_tiling,
                                    decode_tile_x, decode_tile_y, decode_tile_stride_x, 
                                    decode_tile_stride_y, debug_color_shift, kwargs):
        """Process a single tile through the complete WanVideo VACE pipeline with tile-specific reference"""
        
        tile_frames, tile_height, tile_width = video_tile.shape[:3]
        
        # Get tile-specific reference images (from temporal consistency chain)
        tile_ref_images = kwargs.get("vace_ref_images")
        
        # Step 1: Create VACE embeds for this tile with its specific reference
        vace_node = WanVideoVACEEncode()
        vace_embeds = vace_node.process(
            vae=vae,
            width=tile_width,
            height=tile_height, 
            num_frames=tile_frames,
            strength=vace_strength,
            vace_start_percent=vace_start_percent,
            vace_end_percent=vace_end_percent,
            input_frames=video_tile,        # Use tile as input frames
            ref_images=tile_ref_images,     # Use tile-specific reference (from chain or user)
            input_masks=mask_tile,          # Use tile mask
            tiled_vae=kwargs.get("vace_tiled_vae", False)
        )[0]
        
        # Step 2: Run WanVideoSampler on this tile
        sampler_node = WanVideoSampler()
        latent_samples = sampler_node.process(
            model=model,
            image_embeds=vace_embeds,
            steps=steps,
            cfg=cfg,
            shift=shift,
            seed=seed,
            scheduler=scheduler,
            riflex_freq_index=kwargs.get("riflex_freq_index", 0),
            text_embeds=kwargs.get("text_embeds"),
            samples=kwargs.get("samples"),
            denoise_strength=kwargs.get("denoise_strength", 1.0),
            force_offload=kwargs.get("force_offload", True),
            
            # External arguments
            cache_args=kwargs.get("cache_args"),
            slg_args=kwargs.get("slg_args"),
            experimental_args=kwargs.get("experimental_args"),
            
            # Other arguments
            feta_args=kwargs.get("feta_args"),
            context_options=kwargs.get("context_options"),
            flowedit_args=None,
            batched_cfg=kwargs.get("batched_cfg", False),
            rope_function=kwargs.get("rope_function", "comfy"),
            loop_args=kwargs.get("loop_args"),
            sigmas=kwargs.get("sigmas"),
            unianimate_poses=kwargs.get("unianimate_poses"),
            fantasytalking_embeds=kwargs.get("fantasytalking_embeds"),
            uni3c_embeds=kwargs.get("uni3c_embeds"),
            multitalk_embeds=kwargs.get("multitalk_embeds"),
            freeinit_args=kwargs.get("freeinit_args")
        )[0]
        
        # Step 3: Decode latents back to video for this tile
        decode_node = WanVideoDecode()
        processed_tile = decode_node.decode(
            vae=vae,
            samples=latent_samples,
            enable_vae_tiling=decode_enable_vae_tiling,
            tile_x=decode_tile_x,
            tile_y=decode_tile_y,
            tile_stride_x=decode_tile_stride_x,
            tile_stride_y=decode_tile_stride_y,
            normalization=kwargs.get("decode_normalization", "default")
        )[0]
        
        # Step 4: Apply debug color shift to identify tiles visually (if enabled)
        if debug_color_shift:
            debug_color_shifted_tile = self._apply_debug_color_shift(processed_tile, seed)
            return debug_color_shifted_tile, latent_samples
        else:
            return processed_tile, latent_samples
    
    def _apply_debug_color_shift(self, tile, seed):
        """Apply a random color shift to each tile for debugging visualization"""
        import random
        
        # Generate truly random RGB shifts (different each time, not based on seed)
        r_shift = (random.random() - 0.5) * 0.3  # ±0.15 range
        g_shift = (random.random() - 0.5) * 0.3  # ±0.15 range  
        b_shift = (random.random() - 0.5) * 0.3  # ±0.15 range
        
        print(f"      🎨 Debug color shift - Tile: R{r_shift:+.3f}, G{g_shift:+.3f}, B{b_shift:+.3f}")
        
        # Apply color shift to the tile
        shifted_tile = tile.clone()
        
        if tile.shape[-1] >= 3:  # Ensure we have RGB channels
            shifted_tile[:, :, :, 0] = torch.clamp(shifted_tile[:, :, :, 0] + r_shift, 0, 1)
            shifted_tile[:, :, :, 1] = torch.clamp(shifted_tile[:, :, :, 1] + g_shift, 0, 1)
            shifted_tile[:, :, :, 2] = torch.clamp(shifted_tile[:, :, :, 2] + b_shift, 0, 1)
        
        return shifted_tile

    def _stitch_columns(self, all_tiles, spatial_overlap, debug_mode):
        """
        Stitch tiles vertically (same column, same temporal chunk) to create column strips
        """
        # Group tiles by temporal_index and column_index
        column_groups = {}
        for tile in all_tiles:
            key = (tile.temporal_index, tile.column_index)
            if key not in column_groups:
                column_groups[key] = []
            column_groups[key].append(tile)
        
        # Sort each group by line_index
        for key in column_groups:
            column_groups[key].sort(key=lambda t: t.line_index)
        
        column_strips = []
        
        for (temporal_idx, col_idx), column_tiles in column_groups.items():
            if debug_mode and temporal_idx == 0:
                print(f"      🔄 Column {col_idx}: {len(column_tiles)} tiles")
            
            # Stitch tiles in this column vertically
            if len(column_tiles) == 1:
                column_strip = column_tiles[0].content
            else:
                column_strip = self._stitch_tiles_vertically(column_tiles, spatial_overlap)
            
            # Store column strip with its position info
            strip_info = {
                'content': column_strip,
                'temporal_index': temporal_idx,
                'column_index': col_idx,
                'spatial_range_w': column_tiles[0].spatial_range_w,
                'spatial_range_h': (column_tiles[0].spatial_range_h[0], column_tiles[-1].spatial_range_h[1])
            }
            column_strips.append(strip_info)
        
        print(f"   ✅ Created {len(column_strips)} column strips")
        return column_strips
    

    def _stitch_lines(self, column_strips, spatial_overlap, debug_mode):
        """
        Stitch column strips horizontally (same temporal chunk) to create complete temporal chunks
        """
        # Group column strips by temporal_index
        temporal_groups = {}
        for strip in column_strips:
            temporal_idx = strip['temporal_index']
            if temporal_idx not in temporal_groups:
                temporal_groups[temporal_idx] = []
            temporal_groups[temporal_idx].append(strip)
        
        # Sort each group by column_index
        for temporal_idx in temporal_groups:
            temporal_groups[temporal_idx].sort(key=lambda s: s['column_index'])
        
        temporal_chunks = []
        
        for temporal_idx, strips in temporal_groups.items():
            if debug_mode and temporal_idx == 0:
                print(f"      🔄 Temporal chunk {temporal_idx}: {len(strips)} strips")
            
            # Stitch column strips horizontally
            if len(strips) == 1:
                chunk_content = strips[0]['content']
            else:
                chunk_content = self._stitch_strips_horizontally(strips, spatial_overlap)
            
            # Store temporal chunk with its info
            chunk_info = {
                'content': chunk_content,
                'temporal_index': temporal_idx,
                'spatial_range_h': strips[0]['spatial_range_h'],
                'spatial_range_w': (strips[0]['spatial_range_w'][0], strips[-1]['spatial_range_w'][1])
            }
            temporal_chunks.append(chunk_info)
        
        # Sort by temporal index
        temporal_chunks.sort(key=lambda c: c['temporal_index'])
        
        print(f"   ✅ Created {len(temporal_chunks)} temporal chunks")
        return temporal_chunks

    def _stitch_tiles_vertically(self, column_tiles, spatial_overlap):
        """
        Stitch tiles vertically (in height dimension) with proper fade blending.
        Always respects spatial_overlap limit - starts blending from the correct position
        within the overlapping region when actual overlap exceeds spatial_overlap.
        """
        if len(column_tiles) == 1:
            return column_tiles[0].content
        
        # Get dimensions
        first_tile = column_tiles[0].content
        total_height = sum(tile.spatial_range_h[1] - tile.spatial_range_h[0] for tile in column_tiles)
        # Subtract overlaps
        total_height -= spatial_overlap * (len(column_tiles) - 1)
        
        width = first_tile.shape[2]
        channels = first_tile.shape[3]
        frames = first_tile.shape[0]
        
        # Create result tensor
        result = torch.zeros((frames, total_height, width, channels), 
                           dtype=first_tile.dtype, device=first_tile.device)
        
        current_h = 0
        for i, tile in enumerate(column_tiles):
            tile_h = tile.content.shape[1]
            
            if i == 0:
                # First tile - place entirely
                result[:, current_h:current_h+tile_h, :, :] = tile.content
                current_h += tile_h
            else:
                # Subsequent tiles - ALWAYS limit overlap to spatial_overlap maximum
                # Any extra pixels are ignored (they were just for generation context)
                
                # Calculate the effective overlap, but limit it to spatial_overlap
                expected_start = tile.spatial_range_h[0]
                actual_overlap = current_h - expected_start
                actual_overlap = max(0, actual_overlap)  # Ensure non-negative
                
                # STRICT LIMIT: Never exceed spatial_overlap
                overlap_h = min(spatial_overlap, actual_overlap, tile_h // 2, current_h)
                
                if actual_overlap > spatial_overlap:
                    print(f"🎯 Tile {i}: Large overlap {actual_overlap} detected, limiting to {spatial_overlap} (extra pixels ignored)")
                
                if overlap_h > 0:
                    # Calculate where to start within the tile based on actual overlap
                    tile_start_offset = actual_overlap - overlap_h  # Skip extra pixels at the beginning
                    
                    # Get regions for blending (limited to spatial_overlap size)
                    result_overlap = result[:, current_h-overlap_h:current_h, :, :]
                    tile_overlap = tile.content[:, tile_start_offset:tile_start_offset+overlap_h, :, :]
                    
                    # Create vertical fade mask
                    fade_mask = self._create_vertical_fade_mask(overlap_h, result_overlap.dtype, result_overlap.device)
                    
                    # Apply fade blending only within the limited overlap region
                    blended_overlap = result_overlap * (1.0 - fade_mask) + tile_overlap * fade_mask
                    result[:, current_h-overlap_h:current_h, :, :] = blended_overlap
                    
                    # Place the non-overlapping part (skip the overlap region + offset)
                    remaining_start = tile_start_offset + overlap_h
                    if tile_h > remaining_start:
                        result[:, current_h:current_h+tile_h-remaining_start, :, :] = tile.content[:, remaining_start:, :, :]
                    current_h += tile_h - remaining_start
                else:
                    # No overlap - place directly
                    result[:, current_h:current_h+tile_h, :, :] = tile.content
                    current_h += tile_h
        
        return result
    
    def _stitch_strips_horizontally(self, strips, spatial_overlap):
        """
        Stitch column strips horizontally (in width dimension) with proper fade blending.
        Always respects spatial_overlap limit - starts blending from the correct position
        within the overlapping region when actual overlap exceeds spatial_overlap.
        """
        if len(strips) == 1:
            return strips[0]['content']
        
        # Get dimensions
        first_strip = strips[0]['content']
        total_width = sum(s['spatial_range_w'][1] - s['spatial_range_w'][0] for s in strips)
        # Subtract overlaps
        total_width -= spatial_overlap * (len(strips) - 1)
        
        height = first_strip.shape[1]
        channels = first_strip.shape[3]
        frames = first_strip.shape[0]
        
        # Create result tensor
        result = torch.zeros((frames, height, total_width, channels), 
                           dtype=first_strip.dtype, device=first_strip.device)
        
        current_w = 0
        for i, strip in enumerate(strips):
            strip_content = strip['content']
            strip_w = strip_content.shape[2]
            
            if i == 0:
                # First strip - place entirely
                result[:, :, current_w:current_w+strip_w, :] = strip_content
                current_w += strip_w
            else:
                # Subsequent strips - ALWAYS limit overlap to spatial_overlap maximum
                # Any extra pixels are ignored (they were just for generation context)
                
                # Calculate the effective overlap, but limit it to spatial_overlap
                expected_start = strip['spatial_range_w'][0]
                actual_overlap = current_w - expected_start
                actual_overlap = max(0, actual_overlap)  # Ensure non-negative
                
                # STRICT LIMIT: Never exceed spatial_overlap
                overlap_w = min(spatial_overlap, actual_overlap, strip_w // 2, current_w)
                
                if actual_overlap > spatial_overlap:
                    print(f"🎯 Strip {i}: Large overlap {actual_overlap} detected, limiting to {spatial_overlap} (extra pixels ignored)")
                
                if overlap_w > 0:
                    # Calculate where to start within the strip based on actual overlap
                    strip_start_offset = actual_overlap - overlap_w  # Skip extra pixels at the beginning
                    
                    # Get regions for blending (limited to spatial_overlap size)
                    result_overlap = result[:, :, current_w-overlap_w:current_w, :]
                    strip_overlap = strip_content[:, :, strip_start_offset:strip_start_offset+overlap_w, :]
                    
                    # Create horizontal fade mask
                    fade_mask = self._create_horizontal_fade_mask(overlap_w, result_overlap.dtype, result_overlap.device)
                    
                    # Apply fade blending only within the limited overlap region
                    blended_overlap = result_overlap * (1.0 - fade_mask) + strip_overlap * fade_mask
                    result[:, :, current_w-overlap_w:current_w, :] = blended_overlap
                    
                    # Place the non-overlapping part (skip the overlap region + offset)
                    remaining_start = strip_start_offset + overlap_w
                    if strip_w > remaining_start:
                        result[:, :, current_w:current_w+strip_w-remaining_start, :] = strip_content[:, :, remaining_start:, :]
                    current_w += strip_w - remaining_start
                else:
                    # No overlap - place directly
                    result[:, :, current_w:current_w+strip_w, :] = strip_content
                    current_w += strip_w
        
        return result
    
    def _create_vertical_fade_mask(self, fade_size, dtype, device):
        """Create a vertical fade mask for blending tiles vertically."""
        fade_mask = torch.linspace(0, 1, fade_size, dtype=dtype, device=device)
        # Shape: (1, fade_size, 1, 1) for broadcasting
        return fade_mask.view(1, fade_size, 1, 1)
    
    def _create_horizontal_fade_mask(self, fade_size, dtype, device):
        """Create a horizontal fade mask for blending tiles horizontally."""
        fade_mask = torch.linspace(0, 1, fade_size, dtype=dtype, device=device)
        # Shape: (1, 1, fade_size, 1) for broadcasting
        return fade_mask.view(1, 1, fade_size, 1)
    
    def _stitch_temporal_chunks_new(self, temporal_chunks, temporal_tiles, frame_overlap):
        """
        Stitch temporal chunks to create the final output with temporal blending.
        Always respects frame_overlap limit - starts blending from the correct position
        within the overlapping region when actual overlap exceeds frame_overlap.
        """
        if not temporal_chunks:
            return None
            
        if len(temporal_chunks) == 1:
            return temporal_chunks[0]['content']
        
        # Get dimensions from first chunk
        first_chunk_content = temporal_chunks[0]['content']
        
        # Calculate total frames accounting for overlaps
        total_frames = sum(chunk['content'].shape[0] for chunk in temporal_chunks)
        if frame_overlap > 0:
            total_frames -= (len(temporal_chunks) - 1) * frame_overlap
        
        # Get spatial dimensions from the stitched chunks (not original input)
        height = first_chunk_content.shape[1]
        width = first_chunk_content.shape[2]
        channels = first_chunk_content.shape[3]
        
        print(f"   📏 Temporal stitching dimensions:")
        print(f"      • Total frames: {total_frames}")
        print(f"      • Spatial size: {height}×{width}")
        print(f"      • Channels: {channels}")
        print(f"      • Chunk sizes: {[chunk['content'].shape for chunk in temporal_chunks]}")
        
        # Initialize result tensor with correct dimensions
        result = torch.zeros(
            (total_frames, height, width, channels),
            dtype=first_chunk_content.dtype,
            device=first_chunk_content.device
        )
        
        current_t = 0
        for i, chunk_info in enumerate(temporal_chunks):
            chunk_content = chunk_info['content']
            chunk_frames = chunk_content.shape[0]
            
            if i == 0:
                # First chunk - place entirely
                result[current_t:current_t+chunk_frames] = chunk_content
                current_t += chunk_frames
            else:
                # Subsequent chunks - ALWAYS limit overlap to frame_overlap maximum
                # Any extra frames are ignored (they were just for generation context)
                
                # Calculate the effective overlap, but limit it to frame_overlap
                temporal_tile_info = temporal_tiles[i]
                expected_start = temporal_tile_info[0]
                actual_overlap = current_t - expected_start
                actual_overlap = max(0, actual_overlap)  # Ensure non-negative
                
                # STRICT LIMIT: Never exceed frame_overlap
                overlap_frames = min(frame_overlap, actual_overlap, chunk_frames // 2, current_t)
                
                if actual_overlap > frame_overlap:
                    print(f"🎯 Chunk {i}: Large overlap {actual_overlap} detected, limiting to {frame_overlap} (extra frames ignored)")
                
                if overlap_frames > 0:
                    # Calculate where to start within the chunk based on actual overlap
                    chunk_start_offset = actual_overlap - overlap_frames  # Skip extra frames at the beginning
                    
                    # Get regions for blending (limited to frame_overlap size)
                    result_overlap = result[current_t-overlap_frames:current_t]
                    chunk_overlap = chunk_content[chunk_start_offset:chunk_start_offset+overlap_frames]
                    
                    # Create temporal fade mask
                    fade_mask = self._create_temporal_fade_mask(overlap_frames, result_overlap.dtype, result_overlap.device)
                    
                    # Apply fade blending only within the limited overlap region
                    blended_overlap = result_overlap * (1.0 - fade_mask) + chunk_overlap * fade_mask
                    result[current_t-overlap_frames:current_t] = blended_overlap
                    
                    # Place the non-overlapping part (skip the overlap region + offset)
                    remaining_start = chunk_start_offset + overlap_frames
                    if chunk_frames > remaining_start:
                        result[current_t:current_t+chunk_frames-remaining_start] = chunk_content[remaining_start:]
                    current_t += chunk_frames - remaining_start
                else:
                    # No overlap - place directly
                    result[current_t:current_t+chunk_frames] = chunk_content
                    current_t += chunk_frames
        
        return result
    
    def _create_temporal_fade_mask(self, fade_size, dtype, device):
        """Create a temporal fade mask for blending across time dimension."""
        fade_mask = torch.linspace(0, 1, fade_size, dtype=dtype, device=device)
        # Shape: (fade_size, 1, 1, 1) for broadcasting
        return fade_mask.view(fade_size, 1, 1, 1)
    
    def _force_memory_cleanup(self, model, vae):
        """Force memory cleanup between tile processing"""
        try:
            import gc
            
            # Force garbage collection
            gc.collect()
            
            # Clear CUDA cache if available
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                torch.cuda.synchronize()
            
            # Try to offload models if they have offload methods
            if hasattr(model, 'cpu_offload'):
                model.cpu_offload()
            elif hasattr(model, 'to'):
                model.to('cpu')
                
            if hasattr(vae, 'cpu_offload'):
                vae.cpu_offload()
            elif hasattr(vae, 'to'):
                vae.to('cpu')
                
        except Exception as cleanup_error:
            print(f"      ⚠️  Warning: Memory cleanup failed: {str(cleanup_error)}")
    
    def _calculate_temporal_tiles(self, total_frames, target_frames, overlap):
        """Calculate temporal tile ranges with overlap handling"""
        tiles = []
        
        if total_frames <= target_frames:
            tiles.append((0, total_frames))
        else:
            stride = target_frames - overlap
            current = 0
            
            while current < total_frames:
                end = min(current + target_frames, total_frames)
                remaining = total_frames - end
                
                if remaining > 0:
                    tiles.append((current, end))
                    if remaining < stride:
                        final_start = total_frames - target_frames
                        tiles.append((final_start, total_frames))
                        break
                else:
                    tiles.append((current, end))
                    break
                    
                current += stride
        
        return tiles
    
    def _calculate_spatial_tiles(self, total_size, target_size, overlap):
        """Calculate spatial tile ranges with overlap handling"""
        tiles = []
        
        if total_size <= target_size:
            tiles.append((0, total_size))
        else:
            stride = target_size - overlap
            current = 0
            
            while current < total_size:
                end = min(current + target_size, total_size)
                remaining = total_size - end
                
                if remaining > 0:
                    tiles.append((current, end))
                    if remaining < stride:
                        final_start = total_size - target_size
                        tiles.append((final_start, total_size))
                        break
                else:
                    tiles.append((current, end))
                    break
                    
                current += stride
        
        return tiles
    
    def _place_tile_with_overlap(self, target, tile, h_start, h_end, w_start, w_end, spatial_overlap):
        """Place tile into target tensor, handling overlaps with spatial fade blending"""
        tile_h, tile_w = tile.shape[1:3]
        
        # Get the target region
        target_region = target[:, h_start:h_start+tile_h, w_start:w_start+tile_w, :]
        
        # Check if there's existing content (non-zero) in the target area
        if target_region.sum() > 0:
            # Use production-quality fade blending
            fade_mask = self._create_spatial_fade_mask(tile_h, tile_w, target_region, tile, spatial_overlap)
            
            # Apply fade blending: existing * (1 - fade_mask) + tile * fade_mask
            blended = target_region * (1.0 - fade_mask) + tile * fade_mask
            target[:, h_start:h_start+tile_h, w_start:w_start+tile_w, :] = blended
        else:
            # First tile in this area - place directly
            target[:, h_start:h_start+tile_h, w_start:w_start+tile_w, :] = tile
    
    def _create_spatial_fade_mask(self, tile_h, tile_w, existing, new_tile, spatial_overlap):
        """Create a spatial fade mask for smooth blending between tiles"""
        # Create base mask (1.0 = use new tile, 0.0 = use existing)
        mask = torch.ones(1, tile_h, tile_w, 1, dtype=existing.dtype, device=existing.device)
        
        # Define fade distance using spatial_overlap parameter
        fade_h = min(spatial_overlap, tile_h // 2)  # Fade over spatial_overlap pixels or half tile height
        fade_w = min(spatial_overlap, tile_w // 2)  # Fade over spatial_overlap pixels or half tile width
        
        # Check which edges have existing content
        has_top = existing[:, :fade_h, :, :].sum() > 0
        has_bottom = existing[:, -fade_h:, :, :].sum() > 0
        has_left = existing[:, :, :fade_w, :].sum() > 0
        has_right = existing[:, :, -fade_w:, :].sum() > 0
        
        # Create fade gradients for each edge
        if has_top:
            for i in range(fade_h):
                alpha = i / fade_h
                mask[:, i, :, :] = alpha
        
        if has_bottom:
            for i in range(fade_h):
                alpha = 1.0 - (i / fade_h)
                mask[:, tile_h - 1 - i, :, :] = alpha
        
        if has_left:
            for i in range(fade_w):
                alpha = i / fade_w
                mask[:, :, i, :] = torch.minimum(mask[:, :, i, :], torch.tensor(alpha, dtype=mask.dtype, device=mask.device))
        
        if has_right:
            for i in range(fade_w):
                alpha = 1.0 - (i / fade_w)
                mask[:, :, tile_w - 1 - i, :] = torch.minimum(mask[:, :, tile_w - 1 - i, :], torch.tensor(alpha, dtype=mask.dtype, device=mask.device))
        
        return mask.expand_as(new_tile)
    
    def _stitch_temporal_chunks(self, chunks, temporal_tiles, total_frames, height, width, channels, overlap):
        """Stitch temporal chunks back together with temporal fade blending"""
        result = torch.zeros((total_frames, height, width, channels), dtype=chunks[0].dtype, device=chunks[0].device)
        
        for i, ((t_start, t_end), chunk) in enumerate(zip(temporal_tiles, chunks)):
            chunk_frames = chunk.shape[0]
            
            if i == 0:
                result[t_start:t_start+chunk_frames] = chunk
            else:
                prev_end = temporal_tiles[i-1][1]
                overlap_start = max(t_start, prev_end - overlap)
                overlap_end = min(t_end, prev_end)
                
                if overlap_start < overlap_end:
                    overlap_frames = overlap_end - overlap_start
                    chunk_overlap_start = overlap_start - t_start
                    
                    existing_frames = result[overlap_start:overlap_end]
                    new_frames = chunk[chunk_overlap_start:chunk_overlap_start+overlap_frames]
                    
                    fade_mask = self._create_temporal_fade_mask(overlap_frames, existing_frames.dtype, existing_frames.device)
                    
                    blended_frames = existing_frames * (1.0 - fade_mask) + new_frames * fade_mask
                    result[overlap_start:overlap_end] = blended_frames
                    
                    if overlap_end < t_start + chunk_frames:
                        non_overlap_start = overlap_end
                        chunk_offset = non_overlap_start - t_start
                        result[non_overlap_start:t_start+chunk_frames] = chunk[chunk_offset:]
                else:
                    result[t_start:t_start+chunk_frames] = chunk
        
        return result
    
    def _create_temporal_fade_mask(self, overlap_frames, dtype, device):
        """Create a temporal fade mask for smooth frame transitions"""
        fade_values = torch.linspace(0.0, 1.0, overlap_frames, dtype=dtype, device=device)
        fade_mask = fade_values.view(overlap_frames, 1, 1, 1)
        return fade_mask
    
    def _generate_wanvideo_processing_summary(self, all_tiles, temporal_tiles, spatial_tiles_h, spatial_tiles_w, total_tiles):
        """Generate comprehensive summary of the WanVideo tiled processing"""
        successful_tiles = sum(1 for tile in all_tiles if hasattr(tile, 'processing_status') and tile.processing_status == 'success')
        failed_tiles = total_tiles - successful_tiles
        
        summary = f"=== DIMENSION-WISE TILED WANVIDEO VACE PROCESSING SUMMARY ===\n"
        summary += f"Total tiles processed: {total_tiles}\n"
        summary += f"Successful WanVideo tiles: {successful_tiles}\n"
        summary += f"Failed tiles: {failed_tiles}\n"
        summary += f"WanVideo success rate: {(successful_tiles/total_tiles)*100:.1f}%\n\n"
        
        # Temporal consistency chain information
        summary += f"TEMPORAL CONSISTENCY CHAIN:\n"
        summary += f"Temporal chunks: {len(temporal_tiles)}\n"
        for i, (start, end) in enumerate(temporal_tiles):
            if i == 0:
                ref_source = "User-provided reference images"
            else:
                ref_source = f"Frame references from Chunk {i-1} (frame_overlap-based)"
            summary += f"  Chunk {i+1}: frames {start}-{end-1} ({end-start} frames) - {ref_source}\n"
        
        summary += f"\nSpatial tiles per frame: {len(spatial_tiles_h)}×{len(spatial_tiles_w)}\n"
        summary += f"Height tiles: {len(spatial_tiles_h)}\n"
        for i, (start, end) in enumerate(spatial_tiles_h):
            summary += f"  Row {i+1}: pixels {start}-{end-1} ({end-start} pixels)\n"
        
        summary += f"\nWidth tiles: {len(spatial_tiles_w)}\n"
        for i, (start, end) in enumerate(spatial_tiles_w):
            summary += f"  Col {i+1}: pixels {start}-{end-1} ({end-start} pixels)\n"
        
        # Analyze temporal consistency chain usage
        chain_usage = {}
        for tile in all_tiles:
            if hasattr(tile, 'ref_frame_source'):
                source = tile.ref_frame_source
                if source not in chain_usage:
                    chain_usage[source] = 0
                chain_usage[source] += 1
        
        summary += f"\nTemporal reference chain usage:\n"
        for source, count in chain_usage.items():
            summary += f"  {source}: {count} tiles\n"
        
        # Group tiles by processing status
        if failed_tiles > 0:
            summary += f"\nFailed tiles details:\n"
            for tile in all_tiles:
                if hasattr(tile, 'processing_status') and tile.processing_status != 'success':
                    summary += f"  T{tile.temporal_index}:L{tile.line_index}:C{tile.column_index}: {tile.processing_status}\n"
        
        # Add dimension-wise processing information
        summary += f"\nDimension-wise processing algorithm with temporal consistency:\n"
        summary += f"  1. Temporal consistency chain: Each chunk uses references from previous chunk\n"
        summary += f"  2. WanVideo VACE processing: Each tile processed with spatial-specific references\n"
        summary += f"  3. Column-wise stitching: Vertical fade blending within temporal chunks\n"
        summary += f"  4. Line-wise stitching: Horizontal fade blending across columns\n"
        summary += f"  5. Temporal stitching: Frame-based fade blending across time\n"
        summary += f"  6. Cropping: Final video cropped to original input dimensions\n"
        
        summary += f"\nTemporal consistency ensures each tile position maintains coherent evolution across chunks\n"
        summary += f"Reference frame extracted from frame_overlap position for smooth temporal transitions\n"
        summary += f"Large video processed through {total_tiles} individual WanVideo VACE operations with temporal continuity\n"
        
        return summary

def rescale_i(samples, width, height, algorithm: str):
    samples = samples.movedim(-1, 1)
    algorithm = getattr(Image, algorithm.upper())  # i.e. Image.BICUBIC
    samples_pil: Image.Image = F.to_pil_image(samples[0].cpu()).resize((width, height), algorithm)
    samples = F.to_tensor(samples_pil).unsqueeze(0)
    samples = samples.movedim(1, -1)
    return samples


def rescale_m(samples, width, height, algorithm: str):
    samples = samples.unsqueeze(1)
    algorithm = getattr(Image, algorithm.upper())  # i.e. Image.BICUBIC
    samples_pil: Image.Image = F.to_pil_image(samples[0].cpu()).resize((width, height), algorithm)
    samples = F.to_tensor(samples_pil).unsqueeze(0)
    samples = samples.squeeze(1)
    return samples


def preresize_imm(image, mask, optional_context_mask, downscale_algorithm, upscale_algorithm, preresize_mode, preresize_min_width, preresize_min_height, preresize_max_width, preresize_max_height):
    current_width, current_height = image.shape[2], image.shape[1]  # Image size [batch, height, width, channels]
    
    if preresize_mode == "ensure minimum resolution":
        if current_width >= preresize_min_width and current_height >= preresize_min_height:
            return image, mask, optional_context_mask

        scale_factor_min_width = preresize_min_width / current_width
        scale_factor_min_height = preresize_min_height / current_height

        scale_factor = max(scale_factor_min_width, scale_factor_min_height)

        target_width = int(current_width * scale_factor)
        target_height = int(current_height * scale_factor)

        image = rescale_i(image, target_width, target_height, upscale_algorithm)
        mask = rescale_m(mask, target_width, target_height, 'bilinear')
        optional_context_mask = rescale_m(optional_context_mask, target_width, target_height, 'bilinear')
        
        assert target_width >= preresize_min_width and target_height >= preresize_min_height, \
            f"Internal error: After resizing, target size {target_width}x{target_height} is smaller than min size {preresize_min_width}x{preresize_min_height}"

    elif preresize_mode == "ensure minimum and maximum resolution":
        if preresize_min_width <= current_width <= preresize_max_width and preresize_min_height <= current_height <= preresize_max_height:
            return image, mask, optional_context_mask

        scale_factor_min_width = preresize_min_width / current_width
        scale_factor_min_height = preresize_min_height / current_height
        scale_factor_min = max(scale_factor_min_width, scale_factor_min_height)

        scale_factor_max_width = preresize_max_width / current_width
        scale_factor_max_height = preresize_max_height / current_height
        scale_factor_max = min(scale_factor_max_width, scale_factor_max_height)

        if scale_factor_min > 1 and scale_factor_max < 1:
            assert False, "Cannot meet both minimum and maximum resolution requirements with aspect ratio preservation."
        
        if scale_factor_min > 1:  # We're upscaling to meet min resolution
            scale_factor = scale_factor_min
            rescale_algorithm = upscale_algorithm  # Use upscale algorithm for min resolution
        else:  # We're downscaling to meet max resolution
            scale_factor = scale_factor_max
            rescale_algorithm = downscale_algorithm  # Use downscale algorithm for max resolution

        target_width = int(current_width * scale_factor)
        target_height = int(current_height * scale_factor)

        image = rescale_i(image, target_width, target_height, rescale_algorithm)
        mask = rescale_m(mask, target_width, target_height, 'nearest') # Always nearest for efficiency
        optional_context_mask = rescale_m(optional_context_mask, target_width, target_height, 'nearest') # Always nearest for efficiency
        
        assert preresize_min_width <= target_width <= preresize_max_width, \
            f"Internal error: Target width {target_width} is outside the range {preresize_min_width} - {preresize_max_width}"
        assert preresize_min_height <= target_height <= preresize_max_height, \
            f"Internal error: Target height {target_height} is outside the range {preresize_min_height} - {preresize_max_height}"

    elif preresize_mode == "ensure maximum resolution":
        if current_width <= preresize_max_width and current_height <= preresize_max_height:
            return image, mask, optional_context_mask

        scale_factor_max_width = preresize_max_width / current_width
        scale_factor_max_height = preresize_max_height / current_height
        scale_factor_max = min(scale_factor_max_width, scale_factor_max_height)

        target_width = int(current_width * scale_factor_max)
        target_height = int(current_height * scale_factor_max)

        image = rescale_i(image, target_width, target_height, downscale_algorithm)
        mask = rescale_m(mask, target_width, target_height, 'nearest')  # Always nearest for efficiency
        optional_context_mask = rescale_m(optional_context_mask, target_width, target_height, 'nearest')  # Always nearest for efficiency

        assert target_width <= preresize_max_width and target_height <= preresize_max_height, \
            f"Internal error: Target size {target_width}x{target_height} is greater than max size {preresize_max_width}x{preresize_max_height}"

    return image, mask, optional_context_mask


def fillholes_iterative_hipass_fill_m(samples):
    thresholds = [1, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1]

    mask_np = samples.squeeze(0).cpu().numpy()

    for threshold in thresholds:
        thresholded_mask = mask_np >= threshold
        closed_mask = binary_closing(thresholded_mask, structure=np.ones((3, 3)), border_value=1)
        filled_mask = binary_fill_holes(closed_mask)
        mask_np = np.maximum(mask_np, np.where(filled_mask != 0, threshold, 0))

    final_mask = torch.from_numpy(mask_np.astype(np.float32)).unsqueeze(0)

    return final_mask


def hipassfilter_m(samples, threshold):
    filtered_mask = samples.clone()
    filtered_mask[filtered_mask < threshold] = 0
    return filtered_mask


def expand_m(mask, pixels):
    sigma = pixels / 4
    mask_np = mask.squeeze(0).cpu().numpy()
    kernel_size = math.ceil(sigma * 1.5 + 1)
    kernel = np.ones((kernel_size, kernel_size), dtype=np.uint8)
    dilated_mask = grey_dilation(mask_np, footprint=kernel)
    dilated_mask = dilated_mask.astype(np.float32)
    dilated_mask = torch.from_numpy(dilated_mask)
    dilated_mask = torch.clamp(dilated_mask, 0.0, 1.0)
    return dilated_mask.unsqueeze(0)


def invert_m(samples):
    inverted_mask = samples.clone()
    inverted_mask = 1.0 - inverted_mask
    return inverted_mask


def blur_m(samples, pixels):
    mask = samples.squeeze(0)
    sigma = pixels / 4 
    mask_np = mask.cpu().numpy()
    blurred_mask = gaussian_filter(mask_np, sigma=sigma)
    blurred_mask = torch.from_numpy(blurred_mask).float()
    blurred_mask = torch.clamp(blurred_mask, 0.0, 1.0)
    return blurred_mask.unsqueeze(0)


def extend_imm(image, mask, optional_context_mask, extend_up_factor, extend_down_factor, extend_left_factor, extend_right_factor):
    B, H, W, C = image.shape

    new_H = int(H * (1.0 + extend_up_factor - 1.0 + extend_down_factor - 1.0))
    new_W = int(W * (1.0 + extend_left_factor - 1.0 + extend_right_factor - 1.0))

    assert new_H >= 0, f"Error: Trying to crop too much, height ({new_H}) must be >= 0"
    assert new_W >= 0, f"Error: Trying to crop too much, width ({new_W}) must be >= 0"

    expanded_image = torch.zeros(1, new_H, new_W, C, device=image.device)
    expanded_mask = torch.ones(1, new_H, new_W, device=mask.device)
    expanded_optional_context_mask = torch.zeros(1, new_H, new_W, device=optional_context_mask.device)

    up_padding = int(H * (extend_up_factor - 1.0))
    down_padding = new_H - H - up_padding
    left_padding = int(W * (extend_left_factor - 1.0))
    right_padding = new_W - W - left_padding

    slice_target_up = max(0, up_padding)
    slice_target_down = min(new_H, up_padding + H)
    slice_target_left = max(0, left_padding)
    slice_target_right = min(new_W, left_padding + W)

    slice_source_up = max(0, -up_padding)
    slice_source_down = min(H, new_H - up_padding)
    slice_source_left = max(0, -left_padding)
    slice_source_right = min(W, new_W - left_padding)

    image = image.permute(0, 3, 1, 2)  # [B, H, W, C] -> [B, C, H, W]
    expanded_image = expanded_image.permute(0, 3, 1, 2)  # [B, H, W, C] -> [B, C, H, W]

    expanded_image[:, :, slice_target_up:slice_target_down, slice_target_left:slice_target_right] = image[:, :, slice_source_up:slice_source_down, slice_source_left:slice_source_right]
    if up_padding > 0:
        expanded_image[:, :, :up_padding, slice_target_left:slice_target_right] = image[:, :, 0:1, slice_source_left:slice_source_right].repeat(1, 1, up_padding, 1)
    if down_padding > 0:
        expanded_image[:, :, -down_padding:, slice_target_left:slice_target_right] = image[:, :, -1:, slice_source_left:slice_source_right].repeat(1, 1, down_padding, 1)
    if left_padding > 0:
        expanded_image[:, :, :, :left_padding] = expanded_image[:, :, :, left_padding:left_padding+1].repeat(1, 1, 1, left_padding)
    if right_padding > 0:
        expanded_image[:, :, :, -right_padding:] = expanded_image[:, :, :, -right_padding-1:-right_padding].repeat(1, 1, 1, right_padding)

    expanded_mask[:, slice_target_up:slice_target_down, slice_target_left:slice_target_right] = mask[:, slice_source_up:slice_source_down, slice_source_left:slice_source_right]
    expanded_optional_context_mask[:, slice_target_up:slice_target_down, slice_target_left:slice_target_right] = optional_context_mask[:, slice_source_up:slice_source_down, slice_source_left:slice_source_right]

    expanded_image = expanded_image.permute(0, 2, 3, 1)  # [B, C, H, W] -> [B, H, W, C]
    image = image.permute(0, 2, 3, 1)  # [B, C, H, W] -> [B, H, W, C]

    return expanded_image, expanded_mask, expanded_optional_context_mask


def debug_context_location_in_image(image, x, y, w, h):
    debug_image = image.clone()
    debug_image[:, y:y+h, x:x+w, :] = 1.0 - debug_image[:, y:y+h, x:x+w, :]
    return debug_image


def findcontextarea_m(mask):
    mask_squeezed = mask[0]  # Now shape is [H, W]
    non_zero_indices = torch.nonzero(mask_squeezed)

    H, W = mask_squeezed.shape

    if non_zero_indices.numel() == 0:
        x, y = -1, -1
        w, h = -1, -1
    else:
        y = torch.min(non_zero_indices[:, 0]).item()
        x = torch.min(non_zero_indices[:, 1]).item()
        y_max = torch.max(non_zero_indices[:, 0]).item()
        x_max = torch.max(non_zero_indices[:, 1]).item()
        w = x_max - x + 1  # +1 to include the max index
        h = y_max - y + 1  # +1 to include the max index

    context = mask[:, y:y+h, x:x+w]
    return context, x, y, w, h


def growcontextarea_m(context, mask, x, y, w, h, extend_factor):
    img_h, img_w = mask.shape[1], mask.shape[2]

    # Compute intended growth in each direction
    grow_left = int(round(w * (extend_factor-1.0) / 2.0))
    grow_right = int(round(w * (extend_factor-1.0) / 2.0))
    grow_up = int(round(h * (extend_factor-1.0) / 2.0))
    grow_down = int(round(h * (extend_factor-1.0) / 2.0))

    # Try to grow left, but clamp at 0
    new_x = x - grow_left
    if new_x < 0:
        new_x = 0

    # Try to grow up, but clamp at 0
    new_y = y - grow_up
    if new_y < 0:
        new_y = 0

    # Right edge
    new_x2 = x + w + grow_right
    if new_x2 > img_w:
        new_x2 = img_w

    # Bottom edge
    new_y2 = y + h + grow_down
    if new_y2 > img_h:
        new_y2 = img_h

    # New width and height
    new_w = new_x2 - new_x
    new_h = new_y2 - new_y

    # Extract the context
    new_context = mask[:, new_y:new_y+new_h, new_x:new_x+new_w]

    if new_h < 0 or new_w < 0:
        new_x = 0
        new_y = 0
        new_w = mask.shape[2]
        new_h = mask.shape[1]

    return new_context, new_x, new_y, new_w, new_h


def combinecontextmask_m(context, mask, x, y, w, h, optional_context_mask):
    _, x_opt, y_opt, w_opt, h_opt = findcontextarea_m(optional_context_mask)
    if x == -1:
        x, y, w, h = x_opt, y_opt, w_opt, h_opt
    if x_opt == -1:
        x_opt, y_opt, w_opt, h_opt = x, y, w, h
    if x == -1:
        return torch.zeros(1, 0, 0, device=mask.device), -1, -1, -1, -1
    new_x = min(x, x_opt)
    new_y = min(y, y_opt)
    new_x_max = max(x + w, x_opt + w_opt)
    new_y_max = max(y + h, y_opt + h_opt)
    new_w = new_x_max - new_x
    new_h = new_y_max - new_y
    combined_context = mask[:, new_y:new_y+new_h, new_x:new_x+new_w]
    return combined_context, new_x, new_y, new_w, new_h


def pad_to_multiple(value, multiple):
    return int(math.ceil(value / multiple) * multiple)


def crop_magic_im(image, mask, x, y, w, h, target_w, target_h, padding, downscale_algorithm, upscale_algorithm):
    image = image.clone()
    mask = mask.clone()
    
    # Ok this is the most complex function in this node. The one that does the magic after all the preparation done by the other nodes.
    # Basically this function determines the right context area that encompasses the whole context area (mask+optional_context_mask),
    # that is ideally within the bounds of the original image, and that has the right aspect ratio to match target width and height.
    # It may grow the image if the aspect ratio wouldn't fit in the original image.
    # It keeps track of that growing to then be able to crop the image in the stitch node.
    # Finally, it crops the context area and resizes it to be exactly target_w and target_h.
    # It keeps track of that resize to be able to revert it in the stitch node.

    # Check for invalid inputs
    if target_w <= 0 or target_h <= 0 or w == 0 or h == 0:
        return image, 0, 0, image.shape[2], image.shape[1], image, mask, 0, 0, image.shape[2], image.shape[1]

    # Step 1: Pad target dimensions to be multiples of padding
    if padding != 0:
        target_w = pad_to_multiple(target_w, padding)
        target_h = pad_to_multiple(target_h, padding)

    # Step 2: Calculate target aspect ratio
    target_aspect_ratio = target_w / target_h

    # Step 3: Grow current context area to meet the target aspect ratio
    B, image_h, image_w, C = image.shape
    context_aspect_ratio = w / h
    if context_aspect_ratio < target_aspect_ratio:
        # Grow width to meet aspect ratio
        new_w = int(h * target_aspect_ratio)
        new_h = h
        new_x = x - (new_w - w) // 2
        new_y = y

        # Adjust new_x to keep within bounds
        if new_x < 0:
            shift = -new_x
            if new_x + new_w + shift <= image_w:
                new_x += shift
            else:
                overflow = (new_w - image_w) // 2
                new_x = -overflow
        elif new_x + new_w > image_w:
            overflow = new_x + new_w - image_w
            if new_x - overflow >= 0:
                new_x -= overflow
            else:
                overflow = (new_w - image_w) // 2
                new_x = -overflow

    else:
        # Grow height to meet aspect ratio
        new_w = w
        new_h = int(w / target_aspect_ratio)
        new_x = x
        new_y = y - (new_h - h) // 2

        # Adjust new_y to keep within bounds
        if new_y < 0:
            shift = -new_y
            if new_y + new_h + shift <= image_h:
                new_y += shift
            else:
                overflow = (new_h - image_h) // 2
                new_y = -overflow
        elif new_y + new_h > image_h:
            overflow = new_y + new_h - image_h
            if new_y - overflow >= 0:
                new_y -= overflow
            else:
                overflow = (new_h - image_h) // 2
                new_y = -overflow

    # Step 4: Grow the image to accommodate the new context area
    up_padding, down_padding, left_padding, right_padding = 0, 0, 0, 0

    expanded_image_w = image_w
    expanded_image_h = image_h

    # Adjust width for left overflow (x < 0) and right overflow (x + w > image_w)
    if new_x < 0:
        left_padding = -new_x
        expanded_image_w += left_padding
    if new_x + new_w > image_w:
        right_padding = (new_x + new_w - image_w)
        expanded_image_w += right_padding
    # Adjust height for top overflow (y < 0) and bottom overflow (y + h > image_h)
    if new_y < 0:
        up_padding = -new_y
        expanded_image_h += up_padding 
    if new_y + new_h > image_h:
        down_padding = (new_y + new_h - image_h)
        expanded_image_h += down_padding

    # Step 5: Create the new image and mask
    expanded_image = torch.zeros((image.shape[0], expanded_image_h, expanded_image_w, image.shape[3]), device=image.device)
    expanded_mask = torch.ones((mask.shape[0], expanded_image_h, expanded_image_w), device=mask.device)

    # Reorder the tensors to match the required dimension format for padding
    image = image.permute(0, 3, 1, 2)  # [B, H, W, C] -> [B, C, H, W]
    expanded_image = expanded_image.permute(0, 3, 1, 2)  # [B, H, W, C] -> [B, C, H, W]

    # Ensure the expanded image has enough room to hold the padded version of the original image
    expanded_image[:, :, up_padding:up_padding + image_h, left_padding:left_padding + image_w] = image

    # Fill the new extended areas with the edge values of the image
    if up_padding > 0:
        expanded_image[:, :, :up_padding, left_padding:left_padding + image_w] = image[:, :, 0:1, left_padding:left_padding + image_w].repeat(1, 1, up_padding, 1)
    if down_padding > 0:
        expanded_image[:, :, -down_padding:, left_padding:left_padding + image_w] = image[:, :, -1:, left_padding:left_padding + image_w].repeat(1, 1, down_padding, 1)
    if left_padding > 0:
        expanded_image[:, :, up_padding:up_padding + image_h, :left_padding] = expanded_image[:, :, up_padding:up_padding + image_h, left_padding:left_padding+1].repeat(1, 1, 1, left_padding)
    if right_padding > 0:
        expanded_image[:, :, up_padding:up_padding + image_h, -right_padding:] = expanded_image[:, :, up_padding:up_padding + image_h, -right_padding-1:-right_padding].repeat(1, 1, 1, right_padding)

    # Reorder the tensors back to [B, H, W, C] format
    expanded_image = expanded_image.permute(0, 2, 3, 1)  # [B, C, H, W] -> [B, H, W, C]
    image = image.permute(0, 2, 3, 1)  # [B, C, H, W] -> [B, H, W, C]

    # Same for the mask
    expanded_mask[:, up_padding:up_padding + image_h, left_padding:left_padding + image_w] = mask

    # Record the cto values (canvas to original)
    cto_x = left_padding
    cto_y = up_padding
    cto_w = image_w
    cto_h = image_h

    # The final expanded image and mask
    canvas_image = expanded_image
    canvas_mask = expanded_mask

    # Step 6: Crop the image and mask around x, y, w, h
    ctc_x = new_x+left_padding
    ctc_y = new_y+up_padding
    ctc_w = new_w
    ctc_h = new_h

    # Crop the image and mask
    cropped_image = canvas_image[:, ctc_y:ctc_y + ctc_h, ctc_x:ctc_x + ctc_w]
    cropped_mask = canvas_mask[:, ctc_y:ctc_y + ctc_h, ctc_x:ctc_x + ctc_w]

    # Step 7: Resize image and mask to the target width and height
    # Decide which algorithm to use based on the scaling direction
    if target_w > ctc_w or target_h > ctc_h:  # Upscaling
        cropped_image = rescale_i(cropped_image, target_w, target_h, upscale_algorithm)
        cropped_mask = rescale_m(cropped_mask, target_w, target_h, upscale_algorithm)
    else:  # Downscaling
        cropped_image = rescale_i(cropped_image, target_w, target_h, downscale_algorithm)
        cropped_mask = rescale_m(cropped_mask, target_w, target_h, downscale_algorithm)

    return canvas_image, cto_x, cto_y, cto_w, cto_h, cropped_image, cropped_mask, ctc_x, ctc_y, ctc_w, ctc_h


def stitch_magic_im(canvas_image, inpainted_image, mask, ctc_x, ctc_y, ctc_w, ctc_h, cto_x, cto_y, cto_w, cto_h, downscale_algorithm, upscale_algorithm):
    canvas_image = canvas_image.clone()
    inpainted_image = inpainted_image.clone()
    mask = mask.clone()

    # Resize inpainted image and mask to match the context size
    _, h, w, _ = inpainted_image.shape
    if ctc_w > w or ctc_h > h:  # Upscaling
        resized_image = rescale_i(inpainted_image, ctc_w, ctc_h, upscale_algorithm)
        resized_mask = rescale_m(mask, ctc_w, ctc_h, upscale_algorithm)
    else:  # Downscaling
        resized_image = rescale_i(inpainted_image, ctc_w, ctc_h, downscale_algorithm)
        resized_mask = rescale_m(mask, ctc_w, ctc_h, downscale_algorithm)

    # Clamp mask to [0, 1] and expand to match image channels
    resized_mask = resized_mask.clamp(0, 1).unsqueeze(-1)  # shape: [1, H, W, 1]

    # Extract the canvas region we're about to overwrite
    canvas_crop = canvas_image[:, ctc_y:ctc_y + ctc_h, ctc_x:ctc_x + ctc_w]

    # Blend: new = mask * inpainted + (1 - mask) * canvas
    blended = resized_mask * resized_image + (1.0 - resized_mask) * canvas_crop

    # Paste the blended region back onto the canvas
    canvas_image[:, ctc_y:ctc_y + ctc_h, ctc_x:ctc_x + ctc_w] = blended

    # Final crop to get back the original image area
    output_image = canvas_image[:, cto_y:cto_y + cto_h, cto_x:cto_x + cto_w]

    return output_image


class InpaintCropImproved:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # Required inputs
                "image": ("IMAGE", {"tooltip": "Input image(s) to be processed. Can be a single image or batch of images for video processing."}),

                # Resize algorithms
                "downscale_algorithm": (["nearest", "bilinear", "bicubic", "lanczos", "box", "hamming"], {"default": "bilinear", "tooltip": "Algorithm used when downscaling images. Bilinear provides good balance of quality and speed."}),
                "upscale_algorithm": (["nearest", "bilinear", "bicubic", "lanczos", "box", "hamming"], {"default": "bicubic", "tooltip": "Algorithm used when upscaling images. Bicubic provides better quality for upscaling."}),

                # Pre-resize input image
                "preresize": ("BOOLEAN", {"default": False, "tooltip": "Enable to resize the original image before processing. Useful for normalizing input sizes."}),
                "preresize_mode": (["ensure minimum resolution", "ensure maximum resolution", "ensure minimum and maximum resolution"], {"default": "ensure minimum resolution", "tooltip": "Mode for pre-resizing: minimum ensures image is at least the specified size, maximum caps the size, both enforces a range."}),
                "preresize_min_width": ("INT", {"default": 1024, "min": 0, "max": nodes.MAX_RESOLUTION, "step": 1, "tooltip": "Minimum width for pre-resize operation. Image will be upscaled if smaller."}),
                "preresize_min_height": ("INT", {"default": 1024, "min": 0, "max": nodes.MAX_RESOLUTION, "step": 1, "tooltip": "Minimum height for pre-resize operation. Image will be upscaled if smaller."}),
                "preresize_max_width": ("INT", {"default": nodes.MAX_RESOLUTION, "min": 0, "max": nodes.MAX_RESOLUTION, "step": 1, "tooltip": "Maximum width for pre-resize operation. Image will be downscaled if larger."}),
                "preresize_max_height": ("INT", {"default": nodes.MAX_RESOLUTION, "min": 0, "max": nodes.MAX_RESOLUTION, "step": 1, "tooltip": "Maximum height for pre-resize operation. Image will be downscaled if larger."}),

                # Mask manipulation
                "mask_fill_holes": ("BOOLEAN", {"default": True, "tooltip": "Fill holes in the mask using iterative morphological operations. Helps create more complete masked regions."}),
                "mask_expand_pixels": ("INT", {"default": 0, "min": 0, "max": nodes.MAX_RESOLUTION, "step": 1, "tooltip": "Expand the mask by this many pixels before processing. Useful for including more context around masked areas."}),
                "mask_invert": ("BOOLEAN", {"default": False, "tooltip": "Invert the mask so that masked areas become unmasked and vice versa."}),
                "mask_blend_pixels": ("INT", {"default": 32, "min": 0, "max": 64, "step": 1, "tooltip": "Create a soft transition zone around mask edges. Higher values create smoother blending during stitching."}),
                "mask_hipass_filter": ("FLOAT", {"default": 0.1, "min": 0, "max": 1, "step": 0.01, "tooltip": "Remove mask values below this threshold. Helps eliminate weak mask areas and noise."}),

                # Extend image for outpainting
                "extend_for_outpainting": ("BOOLEAN", {"default": False, "tooltip": "Extend the image canvas for outpainting. Adds padding around the image for generating content beyond original boundaries."}),
                "extend_up_factor": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.01, "tooltip": "Factor to extend image upward. 1.0 = no extension, 1.5 = 50% extension upward."}),
                "extend_down_factor": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.01, "tooltip": "Factor to extend image downward. 1.0 = no extension, 1.5 = 50% extension downward."}),
                "extend_left_factor": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.01, "tooltip": "Factor to extend image leftward. 1.0 = no extension, 1.5 = 50% extension leftward."}),
                "extend_right_factor": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0, "step": 0.01, "tooltip": "Factor to extend image rightward. 1.0 = no extension, 1.5 = 50% extension rightward."}),

                # Context
                "context_from_mask_extend_factor": ("FLOAT", {"default": 1.2, "min": 1.0, "max": 100.0, "step": 0.01, "tooltip": "Expand the context area around the mask by this factor. 1.2 = 20% expansion in all directions. Larger values include more surrounding context for better inpainting."}),

                # Output
                "output_resize_to_target_size": ("BOOLEAN", {"default": True, "tooltip": "Resize the output to specific dimensions. When disabled, output size depends on mask area and extend factor."}),
                "output_target_width": ("INT", {"default": 512, "min": 64, "max": nodes.MAX_RESOLUTION, "step": 1, "tooltip": "Target width for output image when resize to target size is enabled. Should match your inpainting model's preferred resolution."}),
                "output_target_height": ("INT", {"default": 512, "min": 64, "max": nodes.MAX_RESOLUTION, "step": 1, "tooltip": "Target height for output image when resize to target size is enabled. Should match your inpainting model's preferred resolution."}),
                "output_padding": (["0", "8", "16", "32", "64", "128", "256", "512"], {"default": "32", "tooltip": "Padding to ensure output dimensions are multiples of this value. Important for models that require specific dimension alignment (e.g., 8 for VAE, 32 for some diffusion models)."}),
                
                # Batch consistency
                "keep_window_size": ("BOOLEAN", {"default": False, "tooltip": "Maintain consistent crop window size across all images in batch. Essential for video processing to avoid flickering. Uses maximum dimensions found and interpolates missing coordinates."}),
           },
           "optional": {
                # Optional inputs
                "mask": ("MASK", {"tooltip": "Mask defining areas to be inpainted. White areas will be inpainted, black areas will be preserved. If not provided, the entire image will be processed."}),
                "optional_context_mask": ("MASK", {"tooltip": "Additional mask defining extra context areas to include in the crop. Useful for ensuring important surrounding details are preserved during inpainting."}),
           }
        }

    FUNCTION = "inpaint_crop"
    CATEGORY = "inpaint"
    DESCRIPTION = """
    Advanced inpainting crop node that intelligently crops images around masked areas for optimal inpainting results.
    
    Features:
    - Batch processing with temporal consistency for video sequences
    - Automatic context area detection and expansion
    - Smart resizing with aspect ratio preservation
    - Optional pre-processing (resize, mask manipulation, outpainting extension)
    - Window size consistency across batches for video processing
    - Linear interpolation for smooth mask transitions in video sequences
    
    The node processes masks to find the optimal crop area, applies various preprocessing steps, 
    and outputs cropped images ready for inpainting along with stitcher data for reconstruction.
    """


    # Remove the following # to turn on debug mode (extra outputs, print statements)
    #'''
    DEBUG_MODE = False
    RETURN_TYPES = ("STITCHER", "IMAGE", "MASK")
    RETURN_NAMES = ("stitcher", "cropped_image", "cropped_mask")
    OUTPUT_TOOLTIPS = ("Stitcher data containing all information needed to reconstruct the original image after inpainting. Pass this to InpaintStitchImproved.", 
                      "Cropped and processed image ready for inpainting. Contains the masked area plus surrounding context.", 
                      "Processed mask corresponding to the cropped image. Shows which areas need to be inpainted.")

    '''
    
    DEBUG_MODE = True # TODO
    RETURN_TYPES = ("STITCHER", "IMAGE", "MASK",
        # DEBUG
        "IMAGE",
        "MASK",
        "MASK",
        "MASK",
        "MASK",
        "MASK",
        "MASK",
        "IMAGE",
        "MASK",
        "MASK",
        "IMAGE",
        "MASK",
        "IMAGE",
        "MASK",
        "IMAGE",
        "MASK",
        "IMAGE",
        "IMAGE",
        "MASK",
        "IMAGE",
        "IMAGE",
        "IMAGE",
        "MASK",
    )
    RETURN_NAMES = ("stitcher", "cropped_image", "cropped_mask",
        # DEBUG
        "DEBUG_preresize_image",
        "DEBUG_preresize_mask",
        "DEBUG_fillholes_mask",
        "DEBUG_expand_mask",
        "DEBUG_invert_mask",
        "DEBUG_blur_mask",
        "DEBUG_hipassfilter_mask",
        "DEBUG_extend_image",
        "DEBUG_extend_mask",
        "DEBUG_context_from_mask",
        "DEBUG_context_from_mask_location",
        "DEBUG_context_expand",
        "DEBUG_context_expand_location",
        "DEBUG_context_with_context_mask",
        "DEBUG_context_with_context_mask_location",
        "DEBUG_context_to_target",
        "DEBUG_context_to_target_location",
        "DEBUG_context_to_target_image",
        "DEBUG_context_to_target_mask",
        "DEBUG_canvas_image",
        "DEBUG_orig_in_canvas_location",
        "DEBUG_cropped_in_canvas_location",
        "DEBUG_cropped_mask_blend",
    )
    #'''

 
    def inpaint_crop(self, image, downscale_algorithm, upscale_algorithm, preresize, preresize_mode, preresize_min_width, preresize_min_height, preresize_max_width, preresize_max_height, extend_for_outpainting, extend_up_factor, extend_down_factor, extend_left_factor, extend_right_factor, mask_hipass_filter, mask_fill_holes, mask_expand_pixels, mask_invert, mask_blend_pixels, context_from_mask_extend_factor, output_resize_to_target_size, output_target_width, output_target_height, output_padding, keep_window_size, mask=None, optional_context_mask=None):
        image = image.clone()
        if mask is not None:
            mask = mask.clone()
        if optional_context_mask is not None:
            optional_context_mask = optional_context_mask.clone()

        output_padding = int(output_padding)
        
        # Check that some parameters make sense
        if preresize and preresize_mode == "ensure minimum and maximum resolution":
            assert preresize_max_width >= preresize_min_width, "Preresize maximum width must be greater than or equal to minimum width"
            assert preresize_max_height >= preresize_min_height, "Preresize maximum height must be greater than or equal to minimum height"

        if self.DEBUG_MODE:
            print('Inpaint Crop Batch input')
            print(image.shape, type(image), image.dtype)
            if mask is not None:
                print(mask.shape, type(mask), mask.dtype)
            if optional_context_mask is not None:
                print(optional_context_mask.shape, type(optional_context_mask), optional_context_mask.dtype)

        if image.shape[0] > 1 and not keep_window_size:
            assert output_resize_to_target_size, "output_resize_to_target_size must be enabled when input is a batch of images without keep_window_size enabled, given all images in the batch output have to be the same size (unless keep_window_size is enabled)"

        # When a LoadImage node passes a mask without user editing, it may be the wrong shape.
        # Detect and fix that to avoid shape mismatch errors.
        if mask is not None and (image.shape[0] == 1 or mask.shape[0] == 1 or mask.shape[0] == image.shape[0]):
            if mask.shape[1] != image.shape[1] or mask.shape[2] != image.shape[2]:
                if torch.count_nonzero(mask) == 0:
                    mask = torch.zeros((mask.shape[0], image.shape[1], image.shape[2]), device=image.device, dtype=image.dtype)

        if optional_context_mask is not None and (image.shape[0] == 1 or optional_context_mask.shape[0] == 1 or optional_context_mask.shape[0] == image.shape[0]):
            if optional_context_mask.shape[1] != image.shape[1] or optional_context_mask.shape[2] != image.shape[2]:
                if torch.count_nonzero(optional_context_mask) == 0:
                    optional_context_mask = torch.zeros((optional_context_mask.shape[0], image.shape[1], image.shape[2]), device=image.device, dtype=image.dtype)

        # If no mask is provided, create one with the shape of the image
        if mask is None:
            mask = torch.zeros_like(image[:, :, :, 0])
    
        # If there is only one image for many masks, replicate it for all masks
        if mask.shape[0] > 1 and image.shape[0] == 1:
            assert image.dim() == 4, f"Expected 4D BHWC image tensor, got {image.shape}"
            image = image.expand(mask.shape[0], -1, -1, -1).clone()

        # If there is only one mask for many images, replicate it for all images
        if image.shape[0] > 1 and mask.shape[0] == 1:
            assert mask.dim() == 3, f"Expected 3D BHW mask tensor, got {mask.shape}"
            mask = mask.expand(image.shape[0], -1, -1).clone()

        # If no optional_context_mask is provided, create one with the shape of the image
        if optional_context_mask is None:
            optional_context_mask = torch.zeros_like(image[:, :, :, 0])

        # If there is only one optional_context_mask for many images, replicate it for all images
        if image.shape[0] > 1 and optional_context_mask.shape[0] == 1:
            assert optional_context_mask.dim() == 3, f"Expected 3D BHW optional_context_mask tensor, got {optional_context_mask.shape}"
            optional_context_mask = optional_context_mask.expand(image.shape[0], -1, -1).clone()

        if self.DEBUG_MODE:
            print('Inpaint Crop Batch ready')
            print(image.shape, type(image), image.dtype)
            print(mask.shape, type(mask), mask.dtype)
            print(optional_context_mask.shape, type(optional_context_mask), optional_context_mask.dtype)

         # Validate data
        assert image.ndimension() == 4, f"Expected 4 dimensions for image, got {image.ndimension()}"
        assert mask.ndimension() == 3, f"Expected 3 dimensions for mask, got {mask.ndimension()}"
        assert optional_context_mask.ndimension() == 3, f"Expected 3 dimensions for optional_context_mask, got {optional_context_mask.ndimension()}"
        assert mask.shape[1:] == image.shape[1:3], f"Mask dimensions do not match image dimensions. Expected {image.shape[1:3]}, got {mask.shape[1:]}"
        assert optional_context_mask.shape[1:] == image.shape[1:3], f"optional_context_mask dimensions do not match image dimensions. Expected {image.shape[1:3]}, got {optional_context_mask.shape[1:]}"
        assert mask.shape[0] == image.shape[0], f"Mask batch does not match image batch. Expected {image.shape[0]}, got {mask.shape[0]}"
        assert optional_context_mask.shape[0] == image.shape[0], f"Optional context mask batch does not match image batch. Expected {image.shape[0]}, got {optional_context_mask.shape[0]}"

        # Run for each image separately
        result_stitcher = {
            'downscale_algorithm': downscale_algorithm,
            'upscale_algorithm': upscale_algorithm,
            'blend_pixels': mask_blend_pixels,
            'canvas_to_orig_x': [],
            'canvas_to_orig_y': [],
            'canvas_to_orig_w': [],
            'canvas_to_orig_h': [],
            'canvas_image': [],
            'cropped_to_canvas_x': [],
            'cropped_to_canvas_y': [],
            'cropped_to_canvas_w': [],
            'cropped_to_canvas_h': [],
            'cropped_mask_for_blend': [],
        }
        
        result_image = []
        result_mask = []

        debug_outputs = {name: [] for name in self.RETURN_NAMES if name.startswith("DEBUG_")}

        batch_size = image.shape[0]
        
        # Compute all contexts first
        batch_contexts, max_w, max_h = self.compute_batch_contexts(
            image, mask, optional_context_mask, downscale_algorithm, upscale_algorithm,
            preresize, preresize_mode, preresize_min_width, preresize_min_height, 
            preresize_max_width, preresize_max_height, extend_for_outpainting,
            extend_up_factor, extend_down_factor, extend_left_factor, extend_right_factor,
            mask_hipass_filter, mask_fill_holes, mask_expand_pixels, mask_invert,
            mask_blend_pixels, context_from_mask_extend_factor, keep_window_size
        )
        
        # Process each image with pre-computed contexts
        for b in range(batch_size):
            ctx = batch_contexts[b]
            
            outputs = self.inpaint_crop_single_image_with_context(
                ctx['processed_image'], ctx['processed_mask'], ctx['processed_optional_context_mask'],
                ctx['x'], ctx['y'], ctx['w'], ctx['h'], ctx['context'],
                downscale_algorithm, upscale_algorithm, mask_blend_pixels,
                output_resize_to_target_size, output_target_width, output_target_height, output_padding,
                keep_window_size, max_w, max_h, context_from_mask_extend_factor)

            stitcher, cropped_image, cropped_mask = outputs[:3]
            for key in ['canvas_to_orig_x', 'canvas_to_orig_y', 'canvas_to_orig_w', 'canvas_to_orig_h', 'canvas_image', 'cropped_to_canvas_x', 'cropped_to_canvas_y', 'cropped_to_canvas_w', 'cropped_to_canvas_h', 'cropped_mask_for_blend']:
                result_stitcher[key].append(stitcher[key])

            cropped_image = cropped_image.clone().squeeze(0)
            result_image.append(cropped_image)
            cropped_mask = cropped_mask.clone().squeeze(0)
            result_mask.append(cropped_mask)

            # Handle the DEBUG_ fields dynamically
            for name, output in zip(self.RETURN_NAMES[3:], outputs[3:]):  # Start from index 3 since first 3 are fixed
                if name.startswith("DEBUG_"):
                    output_array = output.squeeze(0)  # Assuming output needs to be squeezed similar to image/mask
                    debug_outputs[name].append(output_array)

        result_image = torch.stack(result_image, dim=0)
        result_mask = torch.stack(result_mask, dim=0)

        if self.DEBUG_MODE:
            print('Inpaint Crop Batch output')
            print(result_image.shape, type(result_image), result_image.dtype)
            print(result_mask.shape, type(result_mask), result_mask.dtype)

        debug_outputs = {name: torch.stack(values, dim=0) for name, values in debug_outputs.items()}

        return result_stitcher, result_image, result_mask, [debug_outputs[name] for name in self.RETURN_NAMES if name.startswith("DEBUG_")]

    def compute_batch_contexts(self, image, mask, optional_context_mask, downscale_algorithm, upscale_algorithm, 
                             preresize, preresize_mode, preresize_min_width, preresize_min_height, 
                             preresize_max_width, preresize_max_height, extend_for_outpainting, 
                             extend_up_factor, extend_down_factor, extend_left_factor, extend_right_factor,
                             mask_hipass_filter, mask_fill_holes, mask_expand_pixels, mask_invert, 
                             mask_blend_pixels, context_from_mask_extend_factor, keep_window_size):
        """
        Compute contexts for all images in the batch before processing.
        Returns processed masks, contexts, and coordinates for each image.
        """
        batch_size = image.shape[0]
        batch_contexts = []
        
        # Step 1: Apply mask preprocessing for all images
        for b in range(batch_size):
            one_image = image[b].unsqueeze(0)
            one_mask = mask[b].unsqueeze(0)
            one_optional_context_mask = optional_context_mask[b].unsqueeze(0)
            
            # Apply all mask preprocessing steps
            if preresize:
                one_image, one_mask, one_optional_context_mask = preresize_imm(
                    one_image, one_mask, one_optional_context_mask, downscale_algorithm, 
                    upscale_algorithm, preresize_mode, preresize_min_width, preresize_min_height, 
                    preresize_max_width, preresize_max_height)
            
            if mask_fill_holes:
                one_mask = fillholes_iterative_hipass_fill_m(one_mask)
            
            if mask_expand_pixels > 0:
                one_mask = expand_m(one_mask, mask_expand_pixels)
            
            if mask_invert:
                one_mask = invert_m(one_mask)
            
            if mask_blend_pixels > 0:
                one_mask = expand_m(one_mask, mask_blend_pixels)
                one_mask = blur_m(one_mask, mask_blend_pixels*0.5)
            
            if mask_hipass_filter >= 0.01:
                one_mask = hipassfilter_m(one_mask, mask_hipass_filter)
                one_optional_context_mask = hipassfilter_m(one_optional_context_mask, mask_hipass_filter)
            
            if extend_for_outpainting:
                one_image, one_mask, one_optional_context_mask = extend_imm(
                    one_image, one_mask, one_optional_context_mask, extend_up_factor, 
                    extend_down_factor, extend_left_factor, extend_right_factor)
            
            batch_contexts.append({
                'processed_image': one_image,
                'processed_mask': one_mask, 
                'processed_optional_context_mask': one_optional_context_mask
            })
        
        # Step 2: Find initial contexts without applying context operations yet
        initial_contexts = []
        for ctx in batch_contexts:
            context, x, y, w, h = findcontextarea_m(ctx['processed_mask'])
            initial_contexts.append({'x': x, 'y': y, 'w': w, 'h': h})
        
        # Step 3: Apply keep_window_size logic if enabled
        max_w = None
        max_h = None
        if keep_window_size and batch_size > 1:
            # Find maximum dimensions from initial contexts
            max_w = max(ctx['w'] for ctx in initial_contexts)
            max_h = max(ctx['h'] for ctx in initial_contexts)
            
  
            # Update all initial contexts to use consistent window size
            # First, interpolate x and y coordinates for frames with invalid values (-1)
            def interpolate_coordinates(coords):
                """Interpolate missing coordinates (-1) using linear interpolation between valid values"""
                coords = coords.copy()
                n = len(coords)
                
                # Find all valid indices (where coord != -1)
                valid_indices = [i for i, coord in enumerate(coords) if coord != -1]
                
                if not valid_indices:
                    # No valid coordinates, fill with zeros
                    return [0] * n
                
                # Handle leading -1s: use first valid value
                first_valid_idx = valid_indices[0]
                first_valid_value = coords[first_valid_idx]
                for i in range(first_valid_idx):
                    coords[i] = first_valid_value
                
                # Handle trailing -1s: use last valid value
                last_valid_idx = valid_indices[-1]
                last_valid_value = coords[last_valid_idx]
                for i in range(last_valid_idx + 1, n):
                    coords[i] = last_valid_value
                
                # Interpolate between valid values
                for i in range(len(valid_indices) - 1):
                    start_idx = valid_indices[i]
                    end_idx = valid_indices[i + 1]
                    start_val = coords[start_idx]
                    end_val = coords[end_idx]
                    
                    # Linear interpolation for indices between start_idx and end_idx
                    for j in range(start_idx + 1, end_idx):
                        alpha = (j - start_idx) / (end_idx - start_idx)
                        coords[j] = int(start_val + alpha * (end_val - start_val))
                
                return coords
            
            # Extract x and y coordinates
            x_coords = [ctx['x'] for ctx in initial_contexts]
            y_coords = [ctx['y'] for ctx in initial_contexts]
            
            # Interpolate invalid coordinates
            x_coords = interpolate_coordinates(x_coords)
            y_coords = interpolate_coordinates(y_coords)
            
            # Update contexts with interpolated coordinates
            for i, ctx in enumerate(initial_contexts):
                ctx['x'] = x_coords[i]
                ctx['y'] = y_coords[i]
                
                # Update dimensions to maximum
                ctx['w'] = max_w
                ctx['h'] = max_h
                
                # Ensure we don't go outside image boundaries
                img_h, img_w = batch_contexts[i]['processed_image'].shape[1], batch_contexts[i]['processed_image'].shape[2]
                if ctx['x'] + ctx['w'] > img_w:
                    ctx['x'] = max(0, img_w - ctx['w'])
                if ctx['y'] + ctx['h'] > img_h:
                    ctx['y'] = max(0, img_h - ctx['h'])
        
        # Step 4: Now apply context operations (grow context, combine with optional mask) after window size keeping
        for i, ctx in enumerate(batch_contexts):
            x, y, w, h = initial_contexts[i]['x'], initial_contexts[i]['y'], initial_contexts[i]['w'], initial_contexts[i]['h']
            
            # Find initial context
            context, x, y, w, h = findcontextarea_m(ctx['processed_mask'])
            if x == -1 or w == -1 or h == -1 or y == -1:
                x, y, w, h = initial_contexts[i]['x'], initial_contexts[i]['y'], initial_contexts[i]['w'], initial_contexts[i]['h']
                context = ctx['processed_mask'][:, y:y+h, x:x+w]
            else:
                # Use the coordinates from keep_window_size if enabled
                if keep_window_size and batch_size > 1:
                    x, y, w, h = initial_contexts[i]['x'], initial_contexts[i]['y'], initial_contexts[i]['w'], initial_contexts[i]['h']
                    context = ctx['processed_mask'][:, y:y+h, x:x+w]
            
            # Grow context if needed
            if context_from_mask_extend_factor >= 1.01:
                context, x, y, w, h = growcontextarea_m(context, ctx['processed_mask'], x, y, w, h, context_from_mask_extend_factor)
            if x == -1 or w == -1 or h == -1 or y == -1:
                x, y, w, h = 0, 0, ctx['processed_image'].shape[2], ctx['processed_image'].shape[1]
                context = ctx['processed_mask'][:, y:y+h, x:x+w]
            
            # Combine with optional context mask
            context, x, y, w, h = combinecontextmask_m(context, ctx['processed_mask'], x, y, w, h, ctx['processed_optional_context_mask'])
            if x == -1 or w == -1 or h == -1 or y == -1:
                x, y, w, h = 0, 0, ctx['processed_image'].shape[2], ctx['processed_image'].shape[1]
                context = ctx['processed_mask'][:, y:y+h, x:x+w]
            
            # Update the context with final values
            ctx['context'] = context
            ctx['x'] = x
            ctx['y'] = y
            ctx['w'] = w
            ctx['h'] = h
        
        return batch_contexts, max_w, max_h

    def inpaint_crop_single_image_with_context(self, image, mask, optional_context_mask, x, y, w, h, context,
                                             downscale_algorithm, upscale_algorithm, mask_blend_pixels,
                                             output_resize_to_target_size, output_target_width, output_target_height, output_padding, 
                                             keep_window_size=False, max_w=None, max_h=None, context_from_mask_extend_factor=1.0):
        """
        Process a single image with pre-computed context coordinates.
        This is the simplified version that skips all the mask preprocessing since it's already done.
        """
        if self.DEBUG_MODE:
            DEBUG_preresize_image = image.clone()
            DEBUG_preresize_mask = mask.clone()
            DEBUG_fillholes_mask = mask.clone()
            DEBUG_expand_mask = mask.clone()
            DEBUG_invert_mask = mask.clone()
            DEBUG_blur_mask = mask.clone()
            DEBUG_hipassfilter_mask = mask.clone()
            DEBUG_extend_image = image.clone()
            DEBUG_extend_mask = mask.clone()
            DEBUG_context_from_mask = context.clone()
            DEBUG_context_from_mask_location = debug_context_location_in_image(image, x, y, w, h)
            DEBUG_context_expand = context.clone()
            DEBUG_context_expand_location = debug_context_location_in_image(image, x, y, w, h)
            DEBUG_context_with_context_mask = context.clone()
            DEBUG_context_with_context_mask_location = debug_context_location_in_image(image, x, y, w, h)

        if not output_resize_to_target_size:
            # If keep_window_size is enabled and we have max dimensions, use them instead of w,h
            if keep_window_size and max_w is not None and max_h is not None:
                # Dont't forget to consider context_from_mask_extend_factor
                extended_w = int(max_w * context_from_mask_extend_factor)
                extended_h = int(max_h * context_from_mask_extend_factor)
                canvas_image, cto_x, cto_y, cto_w, cto_h, cropped_image, cropped_mask, ctc_x, ctc_y, ctc_w, ctc_h = crop_magic_im(image, mask, x, y, w, h, extended_w, extended_h, output_padding, downscale_algorithm, upscale_algorithm)
            else:
                # Use the actual context dimensions w,h which already include the extend factor
                canvas_image, cto_x, cto_y, cto_w, cto_h, cropped_image, cropped_mask, ctc_x, ctc_y, ctc_w, ctc_h = crop_magic_im(image, mask, x, y, w, h, w, h, output_padding, downscale_algorithm, upscale_algorithm)
        else: # if output_resize_to_target_size:
            canvas_image, cto_x, cto_y, cto_w, cto_h, cropped_image, cropped_mask, ctc_x, ctc_y, ctc_w, ctc_h = crop_magic_im(image, mask, x, y, w, h, output_target_width, output_target_height, output_padding, downscale_algorithm, upscale_algorithm)
        
        if self.DEBUG_MODE:
            DEBUG_context_to_target = context.clone()
            DEBUG_context_to_target_location = debug_context_location_in_image(image, x, y, w, h)
            DEBUG_context_to_target_image = image.clone()
            DEBUG_context_to_target_mask = mask.clone()
            DEBUG_canvas_image = canvas_image.clone()
            DEBUG_orig_in_canvas_location = debug_context_location_in_image(canvas_image, cto_x, cto_y, cto_w, cto_h)
            DEBUG_cropped_in_canvas_location = debug_context_location_in_image(canvas_image, ctc_x, ctc_y, ctc_w, ctc_h)

        # For blending, grow the mask even further and make it blurrier.
        cropped_mask_blend = cropped_mask.clone()
        if mask_blend_pixels > 0:
           cropped_mask_blend = blur_m(cropped_mask_blend, mask_blend_pixels*0.5)
        if self.DEBUG_MODE:
            DEBUG_cropped_mask_blend = cropped_mask_blend.clone()

        stitcher = {
            'canvas_to_orig_x': cto_x,
            'canvas_to_orig_y': cto_y,
            'canvas_to_orig_w': cto_w,
            'canvas_to_orig_h': cto_h,
            'canvas_image': canvas_image,
            'cropped_to_canvas_x': ctc_x,
            'cropped_to_canvas_y': ctc_y,
            'cropped_to_canvas_w': ctc_w,
            'cropped_to_canvas_h': ctc_h,
            'cropped_mask_for_blend': cropped_mask_blend,
        }

        if not self.DEBUG_MODE:
            return stitcher, cropped_image, cropped_mask
        else:
            return stitcher, cropped_image, cropped_mask, DEBUG_preresize_image, DEBUG_preresize_mask, DEBUG_fillholes_mask, DEBUG_expand_mask, DEBUG_invert_mask, DEBUG_blur_mask, DEBUG_hipassfilter_mask, DEBUG_extend_image, DEBUG_extend_mask, DEBUG_context_from_mask, DEBUG_context_from_mask_location, DEBUG_context_expand, DEBUG_context_expand_location, DEBUG_context_with_context_mask, DEBUG_context_with_context_mask_location, DEBUG_context_to_target, DEBUG_context_to_target_location, DEBUG_context_to_target_image, DEBUG_context_to_target_mask, DEBUG_canvas_image, DEBUG_orig_in_canvas_location, DEBUG_cropped_in_canvas_location, DEBUG_cropped_mask_blend


    def inpaint_crop_single_image(self, image, downscale_algorithm, upscale_algorithm, preresize, preresize_mode, preresize_min_width, preresize_min_height, preresize_max_width, preresize_max_height, extend_for_outpainting, extend_up_factor, extend_down_factor, extend_left_factor, extend_right_factor, mask_hipass_filter, mask_fill_holes, mask_expand_pixels, mask_invert, mask_blend_pixels, context_from_mask_extend_factor, output_resize_to_target_size, output_target_width, output_target_height, output_padding, mask, optional_context_mask):
        if preresize:
            image, mask, optional_context_mask = preresize_imm(image, mask, optional_context_mask, downscale_algorithm, upscale_algorithm, preresize_mode, preresize_min_width, preresize_min_height, preresize_max_width, preresize_max_height)
        if self.DEBUG_MODE:
            DEBUG_preresize_image = image.clone()
            DEBUG_preresize_mask = mask.clone()
       
        if mask_fill_holes:
           mask = fillholes_iterative_hipass_fill_m(mask)
        if self.DEBUG_MODE:
            DEBUG_fillholes_mask = mask.clone()

        if mask_expand_pixels > 0:
            mask = expand_m(mask, mask_expand_pixels)
        if self.DEBUG_MODE:
            DEBUG_expand_mask = mask.clone()

        if mask_invert:
            mask = invert_m(mask)
        if self.DEBUG_MODE:
            DEBUG_invert_mask = mask.clone()

        if mask_blend_pixels > 0:
            mask = expand_m(mask, mask_blend_pixels)
            mask = blur_m(mask, mask_blend_pixels*0.5)
        if self.DEBUG_MODE:
            DEBUG_blur_mask = mask.clone()

        if mask_hipass_filter >= 0.01:
            mask = hipassfilter_m(mask, mask_hipass_filter)
            optional_context_mask = hipassfilter_m(optional_context_mask, mask_hipass_filter)
        if self.DEBUG_MODE:
            DEBUG_hipassfilter_mask = mask.clone()

        if extend_for_outpainting:
            image, mask, optional_context_mask = extend_imm(image, mask, optional_context_mask, extend_up_factor, extend_down_factor, extend_left_factor, extend_right_factor)
        if self.DEBUG_MODE:
            DEBUG_extend_image = image.clone()
            DEBUG_extend_mask = mask.clone()

        context, x, y, w, h = findcontextarea_m(mask)
        # If no mask, mask everything for some inpainting.
        if x == -1 or w == -1 or h == -1 or y == -1:
            x, y, w, h = 0, 0, image.shape[2], image.shape[1]
            context = mask[:, y:y+h, x:x+w]
        if self.DEBUG_MODE:
            DEBUG_context_from_mask = context.clone()
            DEBUG_context_from_mask_location = debug_context_location_in_image(image, x, y, w, h)

        if context_from_mask_extend_factor >= 1.01:
            context, x, y, w, h = growcontextarea_m(context, mask, x, y, w, h, context_from_mask_extend_factor)
        # If no mask, mask everything for some inpainting.
        if x == -1 or w == -1 or h == -1 or y == -1:
            x, y, w, h = 0, 0, image.shape[2], image.shape[1]
            context = mask[:, y:y+h, x:x+w]
        if self.DEBUG_MODE:
            DEBUG_context_expand = context.clone()
            DEBUG_context_expand_location = debug_context_location_in_image(image, x, y, w, h)

        context, x, y, w, h = combinecontextmask_m(context, mask, x, y, w, h, optional_context_mask)
        # If no mask, mask everything for some inpainting.
        if x == -1 or w == -1 or h == -1 or y == -1:
            x, y, w, h = 0, 0, image.shape[2], image.shape[1]
            context = mask[:, y:y+h, x:x+w]
        if self.DEBUG_MODE:
            DEBUG_context_with_context_mask = context.clone()
            DEBUG_context_with_context_mask_location = debug_context_location_in_image(image, x, y, w, h)

        if not output_resize_to_target_size:
            canvas_image, cto_x, cto_y, cto_w, cto_h, cropped_image, cropped_mask, ctc_x, ctc_y, ctc_w, ctc_h = crop_magic_im(image, mask, x, y, w, h, w, h, output_padding, downscale_algorithm, upscale_algorithm)
        else: # if output_resize_to_target_size:
            canvas_image, cto_x, cto_y, cto_w, cto_h, cropped_image, cropped_mask, ctc_x, ctc_y, ctc_w, ctc_h = crop_magic_im(image, mask, x, y, w, h, output_target_width, output_target_height, output_padding, downscale_algorithm, upscale_algorithm)
        if self.DEBUG_MODE:
            DEBUG_context_to_target = context.clone()
            DEBUG_context_to_target_location = debug_context_location_in_image(image, x, y, w, h)
            DEBUG_context_to_target_image = image.clone()
            DEBUG_context_to_target_mask = mask.clone()
            DEBUG_canvas_image = canvas_image.clone()
            DEBUG_orig_in_canvas_location = debug_context_location_in_image(canvas_image, cto_x, cto_y, cto_w, cto_h)
            DEBUG_cropped_in_canvas_location = debug_context_location_in_image(canvas_image, ctc_x, ctc_y, ctc_w, ctc_h)

        # For blending, grow the mask even further and make it blurrier.
        cropped_mask_blend = cropped_mask.clone()
        if mask_blend_pixels > 0:
           cropped_mask_blend = blur_m(cropped_mask_blend, mask_blend_pixels*0.5)
        if self.DEBUG_MODE:
            DEBUG_cropped_mask_blend = cropped_mask_blend.clone()

        stitcher = {
            'canvas_to_orig_x': cto_x,
            'canvas_to_orig_y': cto_y,
            'canvas_to_orig_w': cto_w,
            'canvas_to_orig_h': cto_h,
            'canvas_image': canvas_image,
            'cropped_to_canvas_x': ctc_x,
            'cropped_to_canvas_y': ctc_y,
            'cropped_to_canvas_w': ctc_w,
            'cropped_to_canvas_h': ctc_h,
            'cropped_mask_for_blend': cropped_mask_blend,
        }

        if not self.DEBUG_MODE:
            return stitcher, cropped_image, cropped_mask
        else:
            return stitcher, cropped_image, cropped_mask, DEBUG_preresize_image, DEBUG_preresize_mask, DEBUG_fillholes_mask, DEBUG_expand_mask, DEBUG_invert_mask, DEBUG_blur_mask, DEBUG_hipassfilter_mask, DEBUG_extend_image, DEBUG_extend_mask, DEBUG_context_from_mask, DEBUG_context_from_mask_location, DEBUG_context_expand, DEBUG_context_expand_location, DEBUG_context_with_context_mask, DEBUG_context_with_context_mask_location, DEBUG_context_to_target, DEBUG_context_to_target_location, DEBUG_context_to_target_image, DEBUG_context_to_target_mask, DEBUG_canvas_image, DEBUG_orig_in_canvas_location, DEBUG_cropped_in_canvas_location, DEBUG_cropped_mask_blend


class InpaintStitchImproved:
    """
    Advanced inpainting stitch node that reconstructs full images from cropped inpainted regions.
    
    This node takes the output from InpaintCropImproved and stitches the inpainted results back 
    into the original image context. It supports multiple blending modes and handles batch processing
    for video sequences while maintaining temporal consistency.
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "stitcher": ("STITCHER", {"tooltip": "Stitcher data from InpaintCropImproved containing reconstruction information including canvas coordinates, dimensions, and blending masks."}),
                "inpainted_image": ("IMAGE", {"tooltip": "The inpainted image(s) to be stitched back. Should correspond to the cropped output from InpaintCropImproved after inpainting."}),
                "stitch_mode": (["mask_only", "entire_crop", "blend_entire"], {"default": "mask_only", "tooltip": "Stitching method: 'mask_only' blends only masked areas, 'entire_crop' replaces the entire cropped region, 'blend_entire' smoothly blends the entire crop with surrounding areas."}),
                "blend_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Strength of blending operation. 1.0 = full inpainted result, 0.0 = original image preserved. Only affects blend modes."}),
            }
        }

    CATEGORY = "inpaint"
    DESCRIPTION = """
    Reconstructs full-resolution images by stitching inpainted crops back into their original context.
    
    Features:
    - Multiple stitching modes for different use cases
    - Batch processing support for video sequences  
    - Smart blending to avoid seams and artifacts
    - Automatic coordinate transformation and scaling
    - Preserves original image quality in non-inpainted areas
    
    The node handles all the complex coordinate transformations, scaling, and blending
    needed to seamlessly integrate inpainted content back into the source images.
    """

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    OUTPUT_TOOLTIPS = ("Reconstructed full-resolution image(s) with inpainted content seamlessly integrated. Maintains original dimensions and quality.",)

    FUNCTION = "inpaint_stitch"

    def inpaint_stitch(self, stitcher, inpainted_image, stitch_mode, blend_strength):
        inpainted_image = inpainted_image.clone()
        results = []

        batch_size = inpainted_image.shape[0]
        assert len(stitcher['cropped_to_canvas_x']) == batch_size or len(stitcher['cropped_to_canvas_x']) == 1, "Stitch batch size doesn't match image batch size"
        override = False
        if len(stitcher['cropped_to_canvas_x']) != batch_size and len(stitcher['cropped_to_canvas_x']) == 1:
            override = True
        for b in range(batch_size):
            one_image = inpainted_image[b]
            one_stitcher = {}
            for key in ['downscale_algorithm', 'upscale_algorithm', 'blend_pixels']:
                one_stitcher[key] = stitcher[key]
            for key in ['canvas_to_orig_x', 'canvas_to_orig_y', 'canvas_to_orig_w', 'canvas_to_orig_h', 'canvas_image', 'cropped_to_canvas_x', 'cropped_to_canvas_y', 'cropped_to_canvas_w', 'cropped_to_canvas_h', 'cropped_mask_for_blend']:
                if override: # One stitcher for many images, always read 0.
                    one_stitcher[key] = stitcher[key][0]
                else:
                    one_stitcher[key] = stitcher[key][b]
            one_image = one_image.unsqueeze(0)
            one_image, = self.inpaint_stitch_single_image(one_stitcher, one_image, stitch_mode, blend_strength)
            one_image = one_image.squeeze(0)
            one_image = one_image.clone()
            results.append(one_image)

        result_batch = torch.stack(results, dim=0)

        return (result_batch,)

    def inpaint_stitch_single_image(self, stitcher, inpainted_image, stitch_mode, blend_strength):
        downscale_algorithm = stitcher['downscale_algorithm']
        upscale_algorithm = stitcher['upscale_algorithm']
        canvas_image = stitcher['canvas_image']

        ctc_x = stitcher['cropped_to_canvas_x']
        ctc_y = stitcher['cropped_to_canvas_y']
        ctc_w = stitcher['cropped_to_canvas_w']
        ctc_h = stitcher['cropped_to_canvas_h']

        cto_x = stitcher['canvas_to_orig_x']
        cto_y = stitcher['canvas_to_orig_y']
        cto_w = stitcher['canvas_to_orig_w']
        cto_h = stitcher['canvas_to_orig_h']

        mask = stitcher['cropped_mask_for_blend']  # shape: [1, H, W]

        if stitch_mode == "mask_only":
            # Original behavior - use mask for blending
            output_image = stitch_magic_im(canvas_image, inpainted_image, mask, ctc_x, ctc_y, ctc_w, ctc_h, cto_x, cto_y, cto_w, cto_h, downscale_algorithm, upscale_algorithm)
        
        elif stitch_mode == "entire_crop":
            # Replace entire cropped region - create a full mask
            full_mask = torch.ones_like(mask)
            output_image = stitch_magic_im(canvas_image, inpainted_image, full_mask, ctc_x, ctc_y, ctc_w, ctc_h, cto_x, cto_y, cto_w, cto_h, downscale_algorithm, upscale_algorithm)
        
        elif stitch_mode == "blend_entire":
            # Blend entire cropped region with adjustable strength
            full_mask = torch.ones_like(mask) * blend_strength
            output_image = stitch_magic_im(canvas_image, inpainted_image, full_mask, ctc_x, ctc_y, ctc_w, ctc_h, cto_x, cto_y, cto_w, cto_h, downscale_algorithm, upscale_algorithm)

        return (output_image,)


# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
NODE_CLASS_MAPPINGS = {
    "TiledWanImageToMask": ImageToMask,
    "TiledWanImageStatistics": ImageStatistics,
    "TiledWanMaskStatistics": MaskStatistics,
    "TileAndStitchBack": TileAndStitchBack,
    "TiledWanVideoVACEpipe": TiledWanVideoVACEpipe,
    "TiledWanInpaintCropImproved": InpaintCropImproved,
    "TiledWanInpaintStitchImproved": InpaintStitchImproved,
}

# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
    "TiledWanImageToMask": "TiledWan Image To Mask",
    "TiledWanImageStatistics": "TiledWan Image Statistics",
    "TiledWanMaskStatistics": "TiledWan Mask Statistics",
    "TileAndStitchBack": "Tile and Stitch Back",
    "TiledWanVideoVACEpipe": "TiledWan Video VACE Pipeline",
    "TiledWanInpaintCropImproved": "TiledWan Inpaint Crop Improved",
    "TiledWanInpaintStitchImproved": "TiledWan Inpaint Stitch Improved"
}