Merge pull request #390 from AInVFX/main

v2.5.19: new logo, remove dead flash-attn wrapper, graceful DLL fallback, improved VRAM tracking

CONTRIBUTING.md

@@ -16,15 +16,9 @@ We welcome code contributions from the community. Here's how to propose changes:
 1. Fork this repository to your own GitHub account.
 2. Create a new branch on your fork for your changes.
 3. Make your changes in this branch.
-4. When you are ready, submit a pull request to either:
-   - **`main` branch** - For stable features and bug fixes
-   - **`nightly` branch** - For experimental features or if you want faster integration (note: this branch changes more frequently and may be less stable)
+4. When you are ready, submit a pull request to the **`main`** branch.
 
-### Branch Guidelines
-- **Main Branch**: Use this for production-ready changes, critical bug fixes, and well-tested features
-- **Nightly Branch**: Use this for experimental features, cutting-edge improvements, or if you need your changes integrated quickly. Be aware that this branch is actively under development and may contain unstable features.
-
-We use the GitHub Flow workflow. Choose the appropriate target branch based on the nature of your contribution.
+We use the GitHub Flow workflow.
 
 Before submitting a pull request, please make sure your code adheres to the project's coding conventions and it has passed all tests. If you are adding features, please also add appropriate tests.
 

README.md

@@ -36,6 +36,80 @@ We're actively working on improvements and new features. To stay informed:
 
 ## 🚀 Updates
 
+**2025.12.10 - Version 2.5.19**
+
+- **🎨 New header logo design** - Refreshed ASCII art banner *(thanks [@naxci1](https://github.com/naxci1))*
+- **🧹 Remove dead flash attention wrapper** - Removed legacy code from FP8CompatibleDiT; FlashAttentionVarlen already handles backend switching via its `attention_mode` attribute
+- **🛡️ Fix graceful fallback from flash-attn** - Add compatibility shims for corrupted flash_attn/xformers DLLs, preventing startup crashes when CUDA extensions are broken
+- **📊 Improved VRAM tracking** - Separate allocated vs reserved memory tracking, Windows-only overflow detection (WDDM paging behavior)
+- **♻️ Centralize backend detection** - Unified `is_mps_available()`, `is_cuda_available()`, `get_gpu_backend()` helpers across codebase
+- **🔄 Revert 2.5.14 VRAM limit enforcement** - Removed `set_per_process_memory_fraction` call; Overflow detection and warnings remain.
+
+**2025.12.09 - Version 2.5.18**
+
+- **🚀 CLI: Streaming mode for long videos** - New `--chunk_size` flag processes videos in memory-bounded chunks, enabling arbitrarily long videos without RAM limits. Works with model caching (`--cache_dit`/`--cache_vae`) for chunk-to-chunk reuse *(inspired by [disk02](https://github.com/disk02) PR contribution)*
+- **⚡ CLI: Multi-GPU streaming** - Each GPU now streams its segment internally with independent model caching, improving memory efficiency and enabling `--temporal_overlap` blending at GPU boundaries
+- **🔧 CLI: Fix large video MemoryError** - Shared memory transfer replaces numpy pickling, preventing crashes on high-resolution/long video outputs *(inspired by  [FurkanGozukara](https://github.com/FurkanGozukara) PR contribution)*
+
+**2025.12.05 - Version 2.5.17**
+
+- **🔧 Fix: Older GPU compatibility (GTX 970, etc.)** - Runtime bf16 CUBLAS probe replaces compute capability heuristics, correctly detecting unsupported GPUs without affecting RTX 20XX
+
+**2025.12.05 - Version 2.5.16**
+
+- **🔧 Fix: Older GPU compatibility (GTX 970, etc.)** - Automatic fallback for GPUs without bfloat16 support
+- **🐛 Fix: Quality regression** - Reverted bfloat16 detection that was causing artifact issues
+- **📋 Debug: Environment info display** - Shows system info in debug mode to help with issue reporting
+- **📚 Docs: Simplified contribution workflow** - Streamlined to main branch only
+
+**2025.12.03 - Version 2.5.15**
+
+- **🍎 Fix: MPS compatibility** - Disable antialias for MPS tensors and fix bfloat16 arange issues
+- **⚡ Fix: Autocast device type** - Use proper device type attribute to prevent autocast errors
+- **📊 Memory: Accurate VRAM tracking** - Use max_memory_reserved for more precise peak reporting
+- **🔧 Fix: Triton compatibility** - Add shim for bitsandbytes 0.45+ / triton 3.0+ (fixes PyTorch 2.7 installation errors)
+
+**2025.12.01 - Version 2.5.14**
+
+- **🍎 Fix: MPS device comparison** - Normalize device strings to prevent unnecessary tensor movements
+- **📊 Memory: VRAM swap detection** - Peak stats now show GPU+swap breakdown when overflow occurs, with warning when swap detected
+- **🛡️ Memory: Enforce physical VRAM limit** - PyTorch now OOMs instead of silently swapping to shared memory (prevents extreme slowdowns on Windows)
+
+**2025.11.30 - Version 2.5.13**
+
+- **🔧 Fix: PyTorch 2.7+ triton import error** - Resolved installation crash caused by triton.ops import chain on newer triton versions
+- **💾 Fix: OOM on float32 conversion for long videos** - Graceful fallback to native dtype when insufficient memory for float32 conversion
+- **🍎 Fix: CLI watermark error on macOS** - Resolved MPS-related watermark processing crash on Apple Silicon
+
+**2025.11.28 - Version 2.5.12**
+
+- **🐛 Fix: Color artifacts regression** - Reverted in-place tensor operations in video transform pipeline that caused color artifacts on some images
+
+**2025.11.28 - Version 2.5.11**
+
+- **⚡ Feature: CUDNN attention backend** - Added support for PyTorch 2.3+ CUDNN_ATTENTION backend with automatic fallback for older versions (thanks @eadwu)
+- **💾 Fix: Memory spike for long videos** - VAE decode now streams directly to pre-allocated tensor, eliminating OOM errors during long video processing
+- **🎨 Fix: LAB color correction artifacts** - Resolved tile boundary artifacts using wavelet reconstruction preprocessing
+- **🎨 Fix: Color reference misalignment** - Fixed color correction frame alignment with temporal overlap
+- **🍎 Fix: MPS detection reliability** - Switched to canonical `torch.backends.mps.is_available()` API for consistent Apple Silicon detection
+- **🖥️ Fix: Mac subprocess error** - CLI now uses direct processing on Mac to avoid MPS allocator failures in child processes
+- **🖥️ Fix: Multi-GPU device assignment** - CUDA_VISIBLE_DEVICES now set before spawn for proper worker inheritance
+- **📊 Fix: BlockSwap logging** - Now shows effective/total blocks (e.g., 32/32) instead of raw requested value
+- **🔧 Feature: Auto bfloat16 detection** - Automatically detects bfloat16 support to prevent CUBLAS errors on older GPUs
+- **📊 Feature: Peak RAM tracking** - Added RAM usage alongside VRAM in debug summary
+- **⚡ Performance: In-place tensor ops** - Reduced memory allocation overhead with in-place operations throughout pipeline
+- **📖 Docs: Multi-GPU clarification** - Clarified frame-level parallelism behavior expectations for multi-GPU setups
+
+**2025.11.13 - Version 2.5.10**
+
+- **🎯 Fix: Deterministic generation** - Identical images with the same seed now produce identical results across different sessions and batch positions
+- **🔧 Fix: Model caching with BlockSwap** - Resolved issue where cached DiT models wouldn't properly reload when VAE caching state changed
+- **💾 Fix: Runner caching optimization** - Runner templates now correctly cache whenever both DiT and VAE are cached, regardless of caching order
+- **📁 Fix: Case-insensitive model paths** - Extra model paths in YAML config now work regardless of case (seedvr2, SEEDVR2, SeedVR2, etc.)
+- **🐛 Fix: High resolution tile debug crash** - Fixed "NoneType has no attribute log" error when using maximum resolution with VAE tiling
+- **📊 Fix: Temporal overlap logging** - Corrected frame count reporting when temporal overlap is automatically adjusted
+- **🔍 Feature: Enhanced model path debugging** - Added detailed logging to help troubleshoot model loading issues (visible in debug mode)
+
 **2025.11.12 - Version 2.5.9**
 
 - **🐛 Fix: Tile debug visualization crash** - Fixed OpenCV error when using VAE tile debug mode on certain systems.
@@ -712,9 +786,18 @@ The CLI provides comprehensive options for single-GPU, multi-GPU, and batch proc
 # Basic image upscaling
 python inference_cli.py image.jpg
 
-# Basic video video upscaling with temporal consistency
+# Basic video upscaling with temporal consistency
 python inference_cli.py video.mp4 --resolution 720 --batch_size 33
 
+# Streaming mode for long videos (memory-efficient)
+# Processes video in chunks of 330 frames to avoid loading entire video into RAM
+# Use --temporal_overlap to ensure smooth transitions between chunks
+python inference_cli.py long_video.mp4 \
+    --resolution 1080 \
+    --batch_size 33 \
+    --chunk_size 330 \
+    --temporal_overlap 3
+
 # Multi-GPU processing with temporal overlap
 python inference_cli.py video.mp4 \
     --cuda_device 0,1 \
@@ -771,7 +854,8 @@ python inference_cli.py media_folder/ \
 - `--batch_size`: Frames per batch (must follow 4n+1: 1, 5, 9, 13, 17, 21...). Ideally matches shot length for best temporal consistency (default: 5)
 - `--seed`: Random seed for reproducibility (default: 42)
 - `--skip_first_frames`: Skip N initial frames (default: 0)
-- `--load_cap`: Load maximum N frames from video. 0 = load all (default: 0)
+- `--load_cap`: Maximum total frames to load from video. 0 = load all (default: 0)
+- `--chunk_size`: Frames per chunk for streaming mode. When > 0, processes video in memory-bounded chunks of N frames, writing each chunk before loading the next. Essential for long videos that would otherwise exceed RAM. Use with `--temporal_overlap` for seamless chunk transitions. 0 = load all frames at once (default: 0)
 - `--prepend_frames`: Prepend N reversed frames to reduce start artifacts (auto-removed) (default: 0)
 - `--temporal_overlap`: Frames to overlap between batches/GPUs for smooth blending (default: 0)
 
@@ -797,6 +881,7 @@ python inference_cli.py media_folder/ \
 - `--tile_debug`: Visualize tiles: 'false' (default), 'encode', or 'decode'
 
 **Performance Optimization:**
+- `--allow_vram_overflow`: Allow VRAM overflow to system RAM. Prevents OOM but may cause severe slowdown
 - `--attention_mode`: Attention backend: 'sdpa' (default, stable) or 'flash_attn' (faster, requires package)
 - `--compile_dit`: Enable torch.compile for DiT model (20-40% speedup, requires PyTorch 2.0+ and Triton)
 - `--compile_vae`: Enable torch.compile for VAE model (15-25% speedup, requires PyTorch 2.0+ and Triton)
@@ -819,26 +904,40 @@ python inference_cli.py media_folder/ \
 
 ### Multi-GPU Processing Explained
 
-The CLI's multi-GPU mode automatically distributes the workload across multiple GPUs with intelligent temporal overlap handling:
+The CLI's multi-GPU mode uses **frame-level parallelism**: the video is split into chunks and each GPU processes its chunk independently through all 4 phases (encode → upscale → decode → postprocess). This is ideal for long videos where you want to reduce total processing time by dividing the workload.
 
 **How it works:**
-1. Video is split into chunks, one per GPU
-2. Each GPU processes its chunk independently
-3. Chunks overlap by `--temporal_overlap` frames
-4. Results are blended together seamlessly using the overlap region
+1. Video frames are split evenly across GPUs (e.g., 100 frames on 2 GPUs → 50 frames each)
+2. Each GPU loads its own copy of the models and processes its chunk independently
+3. When `--temporal_overlap` is set, chunks include overlapping frames for seamless blending
+4. Results are concatenated (and blended at overlap regions) into the final video
 
-**Example for 2 GPUs with temporal_overlap=4:**
+**Example for 100 frames on 2 GPUs with temporal_overlap=4:**
 ```
-GPU 0: Frames 0-50 (includes 4 overlap frames at end)
-GPU 1: Frames 46-100 (includes 4 overlap frames at beginning)
-Result: Frames 0-100 with smooth transition at frame 48
+GPU 0: Frames 0-53 (50 base + 4 overlap at end, processed as independent video)
+GPU 1: Frames 50-99 (50 frames, 4 overlap at start, processed as independent video)
+Result: Frames 0-99 with smooth blending at the transition point
 ```
 
+**Important considerations:**
+- Each GPU processes its chunk as a separate video with its own batch splitting
+- `batch_size` controls batching *within* each GPU's chunk, not across GPUs
+- For short videos (< 100 frames), single GPU is often more efficient due to model loading overhead
+- Multi-GPU doubles VRAM usage (each GPU loads full models) but roughly halves processing time
+
+**When to use multi-GPU:**
+- Long videos (100+ frames) where splitting provides significant time savings
+- When you have multiple GPUs with sufficient VRAM each
+
+**When to use single GPU:**
+- Short videos where model loading overhead outweighs parallel gains
+- When you want all frames processed together for maximum temporal coherence
+
 **Best practices:**
-- Set `--temporal_overlap` to 2-8 frames for smooth blending
+- Set `--temporal_overlap` to 2-4 frames for smooth blending between GPU chunks
 - Higher overlap = smoother transitions but more redundant processing
 - Use `--prepend_frames` to reduce artifacts at video start
-- batch_size should divide evenly into chunk sizes for best results
+- For optimal quality on short videos, use single GPU with `batch_size` matching your shot length
 
 ## ⚠️ Limitations
 
@@ -886,7 +985,7 @@ For detailed contribution guidelines, see [CONTRIBUTING.md](CONTRIBUTING.md).
 2. Create your feature branch (`git checkout -b feature/AmazingFeature`)
 3. Commit your changes (`git commit -m 'Add some AmazingFeature'`)
 4. Push to the branch (`git push origin feature/AmazingFeature`)
-5. Open a Pull Request to **main** branch for stable features or **nightly** branch for experimental features
+5. Open a Pull Request to the **main** branch
 
 **Get Help:**
 - YouTube: [AInVFX Channel](https://www.youtube.com/@AInVFX)

__init__.py

@@ -3,6 +3,7 @@ ComfyUI-SeedVR2_VideoUpscaler
 Official SeedVR2 integration for ComfyUI
 """
 
+from .src.optimization.compatibility import ensure_triton_compat  # noqa: F401
 from .src.interfaces import comfy_entrypoint, SeedVR2Extension
 
 __all__ = ["comfy_entrypoint", "SeedVR2Extension"]

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "seedvr2_videoupscaler"
 description = "SeedVR2 official ComfyUI integration: ByteDance-Seed's one-step diffusion-based video/image upscaling with memory-efficient inference"
-version = "2.5.9"
+version = "2.5.19"
 authors = [
     {name = "numz"},
     {name = "adrientoupet"}

src/common/diffusion/timesteps/sampling/trailing.py

@@ -36,7 +36,7 @@ class UniformTrailingSamplingTimesteps(SamplingTimesteps):
         dtype: torch.dtype = torch.float32,
     ):
         # Create trailing timesteps with specified dtype
-        timesteps = torch.arange(1.0, 0.0, -1.0 / steps, device=device, dtype=dtype)
+        timesteps = torch.arange(1.0, 0.0, -1.0 / steps, device='cpu').to(device=device, dtype=dtype)
 
         # Shift timesteps.
         timesteps = shift * timesteps / (1 + (shift - 1) * timesteps)

src/common/distributed/basic.py

@@ -21,6 +21,7 @@ from datetime import timedelta
 import torch
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel
+from ...optimization.memory_manager import is_mps_available
 
 def get_global_rank() -> int:
     """
@@ -47,7 +48,7 @@ def get_device() -> torch.device:
     """
     Get current rank device.
     """
-    if hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
+    if is_mps_available():
         return torch.device("mps")
     return torch.device("cuda", get_local_rank())
 

src/core/generation_phases.py

@@ -268,6 +268,10 @@ def encode_all_batches(
     if step <= 0:
         step = batch_size
         temporal_overlap = 0
+        debug.log(f"temporal_overlap >= batch_size, resetting to 0", level="WARNING", category="setup", force=True)
+    
+    # Store actual temporal overlap used (may differ from parameter if reset)
+    ctx['actual_temporal_overlap'] = temporal_overlap
     
     # Calculate number of batches
     num_encode_batches = 0
@@ -306,6 +310,14 @@ def encode_all_batches(
                 runner.vae, ctx['cache_context']['vae_model'], debug
             )
             ctx['cache_context']['vae_newly_cached'] = True
+            
+            # If both models now cached, cache runner template
+            dit_is_cached = ctx['cache_context']['cached_dit'] or ctx['cache_context']['dit_newly_cached']
+            if dit_is_cached:
+                ctx['cache_context']['global_cache'].set_runner(
+                    ctx['cache_context']['dit_id'], ctx['cache_context']['vae_id'], 
+                    runner, debug
+                )
         
         # Set deterministic seed for VAE encoding (separate from diffusion noise)
         # Uses seed + 1,000,000 to avoid collision with upscaling batch seeds
@@ -620,6 +632,7 @@ def upscale_all_batches(
                 runner.dit, ctx['cache_context']['dit_model'], debug
             )
             ctx['cache_context']['dit_newly_cached'] = True
+            
             # If both models now cached, cache runner template
             vae_is_cached = ctx['cache_context']['cached_vae'] or ctx['cache_context']['vae_newly_cached']
             if vae_is_cached:
@@ -628,11 +641,6 @@ def upscale_all_batches(
                     runner, debug
                 )
         
-        # Set base seed for DiT noise generation
-        # Ensures deterministic noise across all batches in this upscaling phase
-        set_seed(seed)
-        debug.log(f"Using seed: {seed}", category="dit")
-        
         # Move DiT to GPU for upscaling (no-op if already there)
         manage_model_device(model=runner.dit, target_device=ctx['dit_device'], 
                             model_name="DiT", debug=debug, runner=runner)
@@ -646,6 +654,11 @@ def upscale_all_batches(
             check_interrupt(ctx)
             
             debug.log(f"Upscaling batch {upscale_idx+1}/{num_valid_latents}", category="generation", force=True)
+            # Reset seed for each batch to ensure identical RNG state
+            # This ensures identical inputs produce identical outputs regardless of batch position
+            set_seed(seed)
+            debug.log(f"Using seed: {seed} for deterministic generation", category="dit")
+
             debug.start_timer(f"upscale_batch_{upscale_idx+1}")
             
             # Move to DiT device with correct dtype for upscaling (no-op if already there)
@@ -698,7 +711,7 @@ def upscale_all_batches(
             debug.start_timer(f"dit_inference_{upscale_idx+1}")
             with torch.no_grad():
                 if dit_dtype != ctx['compute_dtype']:
-                    with torch.autocast(str(ctx['dit_device']), ctx['compute_dtype'], enabled=True):
+                    with torch.autocast(ctx['dit_device'].type, ctx['compute_dtype'], enabled=True):
                         upscaled_latents = runner.inference(
                             noises=noises,
                             conditions=conditions,
@@ -796,7 +809,10 @@ def decode_all_batches(
     """
     Phase 3: VAE Decoding.
     
-    Decodes all upscaled latents back to pixel space.
+    Decodes all upscaled latents back to pixel space and writes directly to
+    pre-allocated final_video tensor. This avoids memory duplication by not
+    storing intermediate batch_samples.
+    
     Requires context from upscale_all_batches with upscaled latents.
     
     Args:
@@ -808,7 +824,8 @@ def decode_all_batches(
         
     Returns:
         dict: Updated context containing:
-            - batch_samples: List of decoded samples ready for post-processing
+            - final_video: Pre-allocated tensor with decoded samples (unnormalized, in [-1,1])
+            - decode_batch_info: List of (start_idx, end_idx, ori_length) for Phase 4 processing
             - VAE cleanup completed
             
     Raises:
@@ -831,11 +848,36 @@ def decode_all_batches(
 
     # Count valid latents
     num_valid_latents = len([l for l in ctx['all_upscaled_latents'] if l is not None])
-    
-    # Pre-allocate to match original batches (use ori_lengths which is always available)
     num_batches = len([l for l in ctx['all_ori_lengths'] if l is not None])
-    ctx['batch_samples'] = [None] * num_batches
     
+    # Get output dimensions from context (set during Phase 1)
+    if 'true_target_dims' not in ctx:
+        raise ValueError("true_target_dims not found in context. Run encode_all_batches first.")
+    true_h, true_w = ctx['true_target_dims']
+    total_frames = ctx.get('total_frames', 0)
+    C = 4 if ctx.get('is_rgba', False) else 3
+    
+    # Pre-allocate final_video at the START of decode phase (before any batch processing)
+    # This ensures we only need memory for final_video + 1 batch, not final_video + all batch_samples
+    target_device = ctx['tensor_offload_device'] if ctx['tensor_offload_device'] is not None else 'cpu'
+    channels_str = "RGBA" if C == 4 else "RGB"
+    required_gb = (total_frames * true_h * true_w * C * 2) / (1024**3)
+    debug.log(f"Pre-allocating output tensor: {total_frames} frames, {true_w}x{true_h}px, {channels_str} ({required_gb:.2f}GB)", 
+              category="setup", force=True)
+    
+    ctx['final_video'] = torch.empty((total_frames, true_h, true_w, C), dtype=ctx['compute_dtype'], device=target_device)
+    
+    # Track batch write positions for Phase 4 processing
+    # Each entry: (write_start, write_end, batch_idx, ori_length)
+    ctx['decode_batch_info'] = []
+    
+    # Get temporal overlap from context (set during Phase 1)
+    temporal_overlap = ctx.get('actual_temporal_overlap', 0)
+    
+    # Track padding removed for final summary
+    total_padding_removed = 0
+    
+    current_write_idx = 0
     decode_idx = 0
     
     try:
@@ -884,41 +926,96 @@ def decode_all_batches(
             samples = runner.vae_decode([upscaled_latent])
             debug.end_timer("vae_decode", "VAE decode")
             
-            # Process samples
+            # Process samples - get the single decoded sample
             debug.start_timer("optimized_video_rearrange")
             samples = optimized_video_rearrange(samples)
             debug.end_timer("optimized_video_rearrange", "Video rearrange")
             
-            # Convert from VAE dtype to compute dtype and offload to avoid VRAM accumulation
-            if ctx['tensor_offload_device'] is not None:
-                # samples is always a single-element list from vae_decode([upscaled_latent])]
-                if samples[0].is_cuda or samples[0].is_mps:
-                    samples[0] = manage_tensor(
-                        tensor=samples[0],
-                        target_device=ctx['tensor_offload_device'],
-                        tensor_name=f"sample_{decode_idx+1}",
-                        dtype=ctx['compute_dtype'],
-                        debug=debug,
-                        reason="storing decoded samples for post-processing",
-                        indent_level=1
-                    )
-            else:
-                # No offload device, but still convert to compute dtype
-                samples[0] = manage_tensor(
-                    tensor=samples[0],
-                    target_device=samples[0].device,
-                    tensor_name=f"sample_{decode_idx+1}",
-                    dtype=ctx['compute_dtype'],
-                    debug=debug,
-                    reason="VAE dtype → compute dtype",
-                    indent_level=1
-                )
-            ctx['batch_samples'][decode_idx] = samples
+            # Get the decoded sample (always single-element list)
+            sample = samples[0]
+            del samples
             
-            # Free the upscaled latent and GPU samples
+            # Get original length for this batch (before any padding was added)
+            ori_length = ctx['all_ori_lengths'][decode_idx] if decode_idx < len(ctx['all_ori_lengths']) else sample.shape[0]
+            
+            # Trim temporal padding: sample is in [T, C, H, W] format after rearrange
+            if ori_length < sample.shape[0]:
+                padding_removed = sample.shape[0] - ori_length
+                debug.log(f"Trimming temporal padding: {padding_removed} frames removed ({sample.shape[0]} → {ori_length})", 
+                         category="video", indent_level=1)
+                sample = sample[:ori_length]
+                total_padding_removed += padding_removed
+            
+            # Trim spatial padding to true target dimensions
+            current_h, current_w = sample.shape[-2:]
+            if current_h != true_h or current_w != true_w:
+                debug.log(f"Trimming spatial padding: {current_w}x{current_h} → {true_w}x{true_h}", 
+                         category="video", indent_level=1)
+                sample = sample[:, :, :true_h, :true_w]
+            
+            # Convert to output format: [T, C, H, W] → [T, H, W, C]
+            # Note: We keep values in [-1, 1] range - normalization happens in Phase 4
+            sample = optimized_sample_to_image_format(sample)  # T, C, H, W → T, H, W, C
+            
+            # Calculate write position with temporal overlap handling
+            batch_frames = sample.shape[0]
+            if decode_idx == 0 or temporal_overlap == 0:
+                # First batch or no overlap: write all frames
+                write_start = current_write_idx
+                write_end = current_write_idx + batch_frames
+            else:
+                # Subsequent batches with overlap: blend overlapping region
+                if temporal_overlap < batch_frames and current_write_idx >= temporal_overlap:
+                    # Blend overlapping region in-place on final_video
+                    prev_tail = ctx['final_video'][current_write_idx - temporal_overlap:current_write_idx]
+                    cur_head = sample[:temporal_overlap]
+                    
+                    # Move to same device for blending if needed
+                    if prev_tail.device != cur_head.device:
+                        cur_head = cur_head.to(prev_tail.device)
+                    
+                    blended = blend_overlapping_frames(prev_tail, cur_head, temporal_overlap)
+                    ctx['final_video'][current_write_idx - temporal_overlap:current_write_idx] = blended
+                    
+                    debug.log(f"Blended {temporal_overlap} overlapping frames at positions {current_write_idx - temporal_overlap}-{current_write_idx}", 
+                             category="video", indent_level=1)
+                    
+                    # Write only non-overlapping part
+                    sample = sample[temporal_overlap:]
+                    batch_frames = sample.shape[0]
+                    del prev_tail, cur_head, blended
+                
+                write_start = current_write_idx
+                write_end = current_write_idx + batch_frames
+            
+            # Move sample to target device and write directly to final_video
+            sample = manage_tensor(
+                tensor=sample,
+                target_device=target_device,
+                tensor_name=f"sample_{decode_idx+1}",
+                dtype=ctx['compute_dtype'],
+                debug=debug,
+                reason="writing to final_video",
+                indent_level=1
+            )
+            
+            # Write to final_video - for RGBA, write only RGB channels (VAE outputs 3 channels)
+            if ctx.get('is_rgba', False):
+                ctx['final_video'][write_start:write_end, :, :, :3] = sample
+            else:
+                ctx['final_video'][write_start:write_end] = sample
+            
+            # Store batch info for Phase 4 processing
+            ctx['decode_batch_info'].append((write_start, write_end, decode_idx, ori_length))
+            current_write_idx = write_end
+            
+            debug.log(f"Wrote {batch_frames} frames to positions {write_start}-{write_end}", 
+                     category="video", indent_level=1)
+            
+            # Free memory immediately - no batch_samples storage
             release_tensor_memory(ctx['all_upscaled_latents'][batch_idx])
             ctx['all_upscaled_latents'][batch_idx] = None
-            del upscaled_latent, samples
+            del upscaled_latent, sample
             
             debug.end_timer(f"decode_batch_{decode_idx+1}", f"Decoded batch {decode_idx+1}")
             
@@ -927,6 +1024,9 @@ def decode_all_batches(
                                 1, "Phase 3: Decoding")
             
             decode_idx += 1
+        
+        # Store padding stats for Phase 4 final summary
+        ctx['total_padding_removed'] = total_padding_removed
             
     except Exception as e:
         debug.log(f"Error in Phase 3 (Decoding): {e}", level="ERROR", category="error", force=True)
@@ -958,14 +1058,12 @@ def postprocess_all_batches(
     """
     Phase 4: Post-processing and Final Assembly.
     
-    Applies color correction and assembles the final video from decoded batches.
-    Uses stream-based direct writing to minimize memory usage - processes each
-    batch and writes directly to pre-allocated output tensor without accumulating
-    intermediate results. Automatically applies temporal overlap blending and 
-    removes prepended frames if specified.
+    Processes final_video slices in-place: applies alpha upscaling, color correction,
+    and normalization. Reads from and writes back to the same final_video tensor
+    to avoid memory duplication.
     
     Args:
-        ctx: Context from decode_all_batches containing batch_samples (required)
+        ctx: Context from decode_all_batches containing final_video (required)
         debug: Debug instance for logging (required)
         progress_callback: Optional callback(current, total, frames, phase_name)
         color_correction: Color correction method - "wavelet", "adain", or "none" (default: "wavelet")
@@ -979,7 +1077,7 @@ def postprocess_all_batches(
             - All intermediate storage cleared for memory efficiency
             
     Raises:
-        ValueError: If context is missing or has no batch samples
+        ValueError: If context is missing or has no final_video
     """
     if debug is None:
         raise ValueError("Debug instance must be provided to postprocess_all_batches")
@@ -987,9 +1085,13 @@ def postprocess_all_batches(
     if ctx is None:
         raise ValueError("Context is required for postprocess_all_batches. Run decode_all_batches first.")
     
-    # Validate we have batch samples
-    if 'batch_samples' not in ctx or not ctx['batch_samples']:
-        raise ValueError("No batch samples found. Run decode_all_batches first.")
+    # Validate we have final_video (pre-allocated in decode_all_batches)
+    if 'final_video' not in ctx or ctx['final_video'] is None:
+        raise ValueError("final_video not found. Run decode_all_batches first.")
+    
+    # Validate we have batch info for processing
+    if 'decode_batch_info' not in ctx or not ctx['decode_batch_info']:
+        raise ValueError("decode_batch_info not found. Run decode_all_batches first.")
     
     debug.log("", category="none", force=True)
     debug.log("━━━━━━━━ Phase 4: Post-processing ━━━━━━━━", category="none", force=True)
@@ -1004,8 +1106,9 @@ def postprocess_all_batches(
         debug.log("No frames to process", level="WARNING", category="generation", force=True)
         return ctx
     
-    # Count valid samples for progress reporting
-    num_valid_samples = len([s for s in ctx['batch_samples'] if s is not None])
+    # Get batch info from decode phase
+    batch_info_list = ctx['decode_batch_info']
+    num_valid_samples = len(batch_info_list)
     
     # Calculate total post-processing work units
     # For RGBA: each batch needs 2 steps (alpha processing + color correction/assembly)
@@ -1022,15 +1125,12 @@ def postprocess_all_batches(
     
     current_postprocessing_step = 0
     
-    # Pre-allocation will happen after processing first sample to get exact dimensions
-    ctx['final_video'] = None
-    current_frame_idx = 0
-    
-    # Initialize padding tracking for final log message
-    total_padding_removed = 0
+    # Get padding stats from Phase 3
+    total_padding_removed = ctx.get('total_padding_removed', 0)
     
     # Alpha processing - handle RGBA inputs with edge-guided upscaling
-    if ctx.get('is_rgba', False) and 'all_alpha_channels' in ctx and 'all_input_rgb' in ctx:
+    # Process alpha on final_video slices in-place
+    if has_alpha_processing:
         debug.log("Processing Alpha channel with edge-guided upscaling...", category="alpha")
         
         # Validate alpha channel data exists
@@ -1038,10 +1138,7 @@ def postprocess_all_batches(
             debug.log("WARNING: Alpha channel data malformed, skipping alpha processing", 
                      level="WARNING", category="alpha", force=True)
         else:
-            for batch_idx in range(len(ctx['batch_samples'])):
-                if ctx['batch_samples'][batch_idx] is None:
-                    continue
-                
+            for write_start, write_end, batch_idx, ori_length in batch_info_list:
                 # Bounds checking for alpha channel lists
                 if batch_idx >= len(ctx['all_alpha_channels']) or ctx['all_alpha_channels'][batch_idx] is None:
                     continue
@@ -1055,9 +1152,14 @@ def postprocess_all_batches(
                 debug.log(f"Processing Alpha batch {batch_idx+1}/{num_valid_samples}", category="alpha", force=True)
                 debug.start_timer(f"alpha_batch_{batch_idx+1}")
 
+                # Get RGB slice from final_video for alpha processing
+                # final_video is [T, H, W, C], process_alpha_for_batch expects list of [T, C, H, W]
+                rgb_slice = ctx['final_video'][write_start:write_end, :, :, :3]  # Only RGB
+                rgb_tchw = rgb_slice.permute(0, 3, 1, 2)  # [T, H, W, 3] → [T, 3, H, W]
+                
                 # Process Alpha and merge with RGB
-                ctx['batch_samples'][batch_idx] = process_alpha_for_batch(
-                    rgb_samples=ctx['batch_samples'][batch_idx],
+                processed_samples = process_alpha_for_batch(
+                    rgb_samples=[rgb_tchw],
                     alpha_original=ctx['all_alpha_channels'][batch_idx],
                     rgb_original=ctx['all_input_rgb'][batch_idx],
                     device=ctx['vae_device'],
@@ -1065,6 +1167,27 @@ def postprocess_all_batches(
                     debug=debug
                 )
                 
+                # processed_samples[0] is [T, 4, H, W] (RGBA)
+                # Extract only the alpha channel and write to final_video's alpha slot
+                processed_rgba = processed_samples[0]  # [T, 4, H, W]
+                alpha_channel = processed_rgba[:, 3:4, :, :]  # [T, 1, H, W]
+                alpha_thwc = alpha_channel.permute(0, 2, 3, 1)  # [T, 1, H, W] → [T, H, W, 1]
+                
+                alpha_thwc = manage_tensor(
+                    tensor=alpha_thwc,
+                    target_device=ctx['final_video'].device,
+                    tensor_name=f"alpha_channel_{batch_idx+1}",
+                    dtype=ctx['compute_dtype'],
+                    debug=debug,
+                    reason="writing alpha channel to final_video",
+                    indent_level=1
+                )
+                
+                # Write only the alpha channel to the 4th channel slot
+                ctx['final_video'][write_start:write_end, :, :, 3:4] = alpha_thwc
+                
+                del rgb_slice, rgb_tchw, processed_samples, processed_rgba, alpha_channel, alpha_thwc
+                
                 # Free memory immediately
                 release_tensor_memory(ctx['all_alpha_channels'][batch_idx])
                 ctx['all_alpha_channels'][batch_idx] = None
@@ -1083,221 +1206,165 @@ def postprocess_all_batches(
         debug.log("Alpha processing complete for all batches", category="alpha")
     
     try:
-        # Stream-based processing: write directly to final_video without accumulation
-        for batch_idx, samples in enumerate(ctx['batch_samples']):
-            if samples is None:
-                continue
-                
+        # Process each batch slice in final_video in-place
+        for info_idx, (write_start, write_end, batch_idx, ori_length) in enumerate(batch_info_list):
             check_interrupt(ctx)
             
-            debug.log(f"Post-processing batch {batch_idx+1}/{num_valid_samples}", category="video", force=True)
-            debug.start_timer(f"postprocess_batch_{batch_idx+1}")
+            debug.log(f"Post-processing batch {info_idx+1}/{num_valid_samples}", category="video", force=True)
+            debug.start_timer(f"postprocess_batch_{info_idx+1}")
             
-            # Post-process each sample in the batch
-            for i, sample in enumerate(samples):
-                # Move to VAE device with correct dtype for processing (no-op if already there)
-                sample = manage_tensor(
-                    tensor=sample,
-                    target_device=ctx['vae_device'],
-                    tensor_name=f"sample_{batch_idx+1}_{i}",
-                    dtype=ctx['compute_dtype'],
-                    debug=debug,
-                    reason="post-processing",
-                    indent_level=1
-                )
-                
-                # Get original length for trimming (always available)
-                video_idx = min(batch_idx, len(ctx['all_ori_lengths']) - 1)
-                ori_length = ctx['all_ori_lengths'][video_idx] if 'all_ori_lengths' in ctx else sample.shape[0]
-                
-                # Reconstruct transformed video on-demand for color correction
-                input_video = None
-                if color_correction != "none" and ctx.get('batch_metadata') is not None:
-                    if video_idx < len(ctx['batch_metadata']) and ctx['batch_metadata'][video_idx] is not None:
-                        # Reconstruct transformation
-                        transformed_video = _reconstruct_and_transform_batch(ctx, video_idx, debug)
-                        input_video = optimized_single_video_rearrange(transformed_video)
-                        del transformed_video
-                
-                # Trim both sample and input_video to original length if necessary (handles temporal padding)
-                if ori_length < sample.shape[0]:
-                    padding_removed = sample.shape[0] - ori_length
-                    debug.log(f"Removing padding: {padding_removed} frame{'s' if padding_removed != 1 else ''} trimmed ({sample.shape[0]} -> {ori_length})", 
-                            category="video", force=True, indent_level=1)
-                    sample = sample[:ori_length]
-                    if input_video is not None:
-                        input_video = input_video[:ori_length]
+            # Get slice from final_video - currently in [T, H, W, C] format, values in [-1, 1]
+            sample_thwc = ctx['final_video'][write_start:write_end]
+            
+            # For RGBA, we only process RGB channels for color correction
+            # Alpha was already written during alpha processing above
+            if ctx.get('is_rgba', False) and sample_thwc.shape[-1] == 4:
+                sample_thwc_rgb = sample_thwc[..., :3]  # [T, H, W, 3]
+                sample = sample_thwc_rgb.permute(0, 3, 1, 2)  # [T, H, W, 3] → [T, 3, H, W]
+            else:
+                sample = sample_thwc.permute(0, 3, 1, 2)  # [T, H, W, C] → [T, C, H, W]
+            
+            # Move to VAE device for processing
+            sample = manage_tensor(
+                tensor=sample,
+                target_device=ctx['vae_device'],
+                tensor_name=f"sample_{info_idx+1}",
+                dtype=ctx['compute_dtype'],
+                debug=debug,
+                reason="post-processing",
+                indent_level=1
+            )
+            
+            # Reconstruct transformed video on-demand for color correction
+            input_video = None
+            if color_correction != "none" and ctx.get('batch_metadata') is not None:
+                if batch_idx < len(ctx['batch_metadata']) and ctx['batch_metadata'][batch_idx] is not None:
+                    # Reconstruct transformation
+                    transformed_video = _reconstruct_and_transform_batch(ctx, batch_idx, debug)
+                    input_video = optimized_single_video_rearrange(transformed_video)
+                    del transformed_video
                     
-                    # Accumulate padding removed across all batches
-                    total_padding_removed += padding_removed
-                
-                # Trim spatial dimensions to true target size (removes divisible padding)
-                if 'true_target_dims' in ctx:
-                    true_h, true_w = ctx['true_target_dims']
-                    current_h, current_w = sample.shape[-2:]
+                    # For batches after the first with temporal overlap, the overlap frames
+                    # were blended in Phase 3 and are not part of this slice. Skip them.
+                    actual_overlap = ctx.get('actual_temporal_overlap', 0)
+                    if info_idx > 0 and actual_overlap > 0:
+                        input_video = input_video[actual_overlap:]
                     
-                    if current_h != true_h or current_w != true_w:
-                        debug.log(f"Trimming spatial padding: {current_w}x{current_h} → {true_w}x{true_h}", 
-                                category="video", indent_level=1)
-                        
-                        # Trim from bottom-right to match padding location
-                        sample = sample[:, :, :true_h, :true_w]
-                        if input_video is not None:
+                    # Trim input_video to match sample length (handles padding differences)
+                    if input_video.shape[0] > sample.shape[0]:
+                        input_video = input_video[:sample.shape[0]]
+                    
+                    # Trim spatial dimensions to true target size
+                    if 'true_target_dims' in ctx:
+                        true_h, true_w = ctx['true_target_dims']
+                        if input_video.shape[-2] != true_h or input_video.shape[-1] != true_w:
                             input_video = input_video[:, :, :true_h, :true_w]
+            
+            # Apply color correction if enabled (RGB only)
+            if color_correction != "none" and input_video is not None:
+                # Check if RGBA (samples are in T, C, H, W format at this point)
+                has_alpha = ctx.get('is_rgba', False)
+                alpha_channel = None
                 
-                # Apply color correction if enabled (RGB only)
-                if color_correction != "none" and input_video is not None:
-                    # Check if RGBA (samples are in T, C, H, W format at this point)
-                    has_alpha = ctx.get('is_rgba', False)
-                    alpha_channel = None
-                    
-                    if has_alpha:
-                        # Check actual channel count
-                        if sample.shape[1] == 4:
-                            # Extract and temporarily store alpha for reattachment after color correction
-                            alpha_channel = sample[:, 3:4, :, :]  # (T, 1, H, W)
-                            sample = sample[:, :3, :, :]  # Keep only RGB (T, 3, H, W)
-                    
-                    # Ensure both tensors are on same device (GPU) for color correction
-                    if input_video.device != sample.device:
-                        input_video = manage_tensor(
-                            tensor=input_video,
-                            target_device=sample.device,
-                            tensor_name=f"input_video_{batch_idx+1}",
-                            debug=debug,
-                            reason="color correction",
-                            indent_level=1
-                        )
-                        
-                    # Apply selected color correction method
-                    debug.start_timer(f"color_correction_{color_correction}")
-                    
-                    if color_correction == "lab":
-                        debug.log("Applying LAB perceptual color transfer", category="video", force=True, indent_level=1)
-                        sample = lab_color_transfer(sample, input_video, debug, luminance_weight=0.8)
-                    elif color_correction == "wavelet_adaptive":
-                        debug.log("Applying wavelet with adaptive saturation correction", category="video", force=True, indent_level=1)
-                        sample = wavelet_adaptive_color_correction(sample, input_video, debug)
-                    elif color_correction == "wavelet":
-                        debug.log("Applying wavelet color reconstruction", category="video", force=True, indent_level=1)
-                        sample = wavelet_reconstruction(sample, input_video, debug)
-                    elif color_correction == "hsv":
-                        debug.log("Applying HSV hue-conditional saturation matching", category="video", force=True, indent_level=1)
-                        sample = hsv_saturation_histogram_match(sample, input_video, debug)
-                    elif color_correction == "adain":
-                        debug.log("Applying AdaIN color correction", category="video", force=True, indent_level=1)
-                        sample = adaptive_instance_normalization(sample, input_video)
-                    else:
-                        debug.log(f"Unknown color correction method: {color_correction}", level="WARNING", category="video", force=True, indent_level=1)
-                    
-                    debug.end_timer(f"color_correction_{color_correction}", f"Color correction ({color_correction})")
-                    
-                    # Free the reconstructed transformed video
-                    del input_video
-
-                    # Recombine with Alpha if it was present in input
-                    if has_alpha and alpha_channel is not None:
-                        # Concatenate in channels-first: (T, 3, H, W) + (T, 1, H, W) -> (T, 4, H, W)
-                        sample = torch.cat([sample, alpha_channel], dim=1)
+                if has_alpha:
+                    # Check actual channel count
+                    if sample.shape[1] == 4:
+                        # Extract and temporarily store alpha for reattachment after color correction
+                        alpha_channel = sample[:, 3:4, :, :]  # (T, 1, H, W)
+                        sample = sample[:, :3, :, :]  # Keep only RGB (T, 3, H, W)
                 
-                else:
-                    debug.log("Color correction disabled (set to none)", category="video", indent_level=1)
-                
-                # Free the original length entry
-                if 'all_ori_lengths' in ctx and video_idx < len(ctx['all_ori_lengths']):
-                    ctx['all_ori_lengths'][video_idx] = None
-                
-                # Convert to final format (still on GPU at this point)
-                sample = optimized_sample_to_image_format(sample)
-                
-                # Apply normalization only to RGB channels, preserve Alpha as-is
-                if ctx.get('is_rgba', False) and sample.shape[-1] == 4:
-                    # Split RGBA: sample is (T, H, W, C) format after optimized_sample_to_image_format
-                    rgb_channels = sample[..., :3]  # (T, H, W, 3)
-                    alpha_channel = sample[..., 3:4]  # (T, H, W, 1)
-                    
-                    # Normalize only RGB from [-1, 1] to [0, 1]
-                    rgb_channels = rgb_channels.clip(-1, 1).mul_(0.5).add_(0.5)
-                    
-                    # Merge back with unchanged Alpha
-                    sample = torch.cat([rgb_channels, alpha_channel], dim=-1)
-                else:
-                    # RGB only: apply normalization as usual
-                    sample = sample.clip(-1, 1).mul_(0.5).add_(0.5)
-                
-                # Draw tile boundaries for debugging (if tile info available)
-                for phase, attr in [('encode', 'encode_tile_boundaries'), ('decode', 'decode_tile_boundaries')]:
-                    tiles = getattr(debug, attr, None)
-                    if tiles:
-                        sample = _draw_tile_boundaries(sample, debug, tiles, phase)
-                        break
-                
-                # Move to tensor_offload_device if specified
-                if ctx['tensor_offload_device'] is not None:
-                    sample = manage_tensor(
-                        tensor=sample,
-                        target_device=ctx['tensor_offload_device'],
-                        tensor_name=f"sample_{batch_idx+1}_{i}_final",
+                # Ensure both tensors are on same device (GPU) for color correction
+                if input_video.device != sample.device:
+                    input_video = manage_tensor(
+                        tensor=input_video,
+                        target_device=sample.device,
+                        tensor_name=f"input_video_{info_idx+1}",
                         debug=debug,
-                        reason="storing final processed samples",
+                        reason="color correction",
                         indent_level=1
                     )
-                
-                # Get batch dimensions
-                batch_frames = sample.shape[0]
-                
-                # Pre-allocate output tensor on first write
-                if ctx['final_video'] is None:
-                    H, W = sample.shape[1], sample.shape[2]
-                    # Use input format, not first sample's shape, for consistent allocation
-                    C = 4 if ctx.get('is_rgba', False) else 3
-                    channels_str = "RGBA" if C == 4 else "RGB"
                     
-                    debug.log(f"Pre-allocating output tensor: {total_frames} frames, {W}x{H}px, {channels_str}", category="setup")
-                    
-                    target_device = ctx['tensor_offload_device'] if ctx['tensor_offload_device'] is not None else ctx['vae_device']
-                    ctx['final_video'] = torch.empty((total_frames, H, W, C), dtype=ctx['compute_dtype'], device=target_device)
-
-                # Write with temporal overlap blending if enabled
-                if batch_idx == 0 or temporal_overlap == 0:
-                    # First batch or no overlap: write all frames normally
-                    ctx['final_video'][current_frame_idx:current_frame_idx + batch_frames] = sample
-                    current_frame_idx += batch_frames
+                # Apply selected color correction method
+                debug.start_timer(f"color_correction_{color_correction}")
+                
+                if color_correction == "lab":
+                    debug.log("Applying LAB perceptual color transfer", category="video", force=True, indent_level=1)
+                    sample = lab_color_transfer(sample, input_video, debug, luminance_weight=0.8)
+                elif color_correction == "wavelet_adaptive":
+                    debug.log("Applying wavelet with adaptive saturation correction", category="video", force=True, indent_level=1)
+                    sample = wavelet_adaptive_color_correction(sample, input_video, debug)
+                elif color_correction == "wavelet":
+                    debug.log("Applying wavelet color reconstruction", category="video", force=True, indent_level=1)
+                    sample = wavelet_reconstruction(sample, input_video, debug)
+                elif color_correction == "hsv":
+                    debug.log("Applying HSV hue-conditional saturation matching", category="video", force=True, indent_level=1)
+                    sample = hsv_saturation_histogram_match(sample, input_video, debug)
+                elif color_correction == "adain":
+                    debug.log("Applying AdaIN color correction", category="video", force=True, indent_level=1)
+                    sample = adaptive_instance_normalization(sample, input_video)
                 else:
-                    # Subsequent batches with overlap: blend overlapping region in-place
-                    if temporal_overlap < batch_frames and current_frame_idx >= temporal_overlap:
-                        # Get overlapping regions
-                        prev_tail = ctx['final_video'][current_frame_idx - temporal_overlap:current_frame_idx]
-                        cur_head = sample[:temporal_overlap]
-                        
-                        # Log blending operation
-                        blend_range_start = current_frame_idx - temporal_overlap
-                        blend_range_end = current_frame_idx
-                        debug.log(f"Blending {temporal_overlap} overlapping frames (positions {blend_range_start}-{blend_range_end})", 
-                                 category="video", indent_level=1)
-                        
-                        # Blend and replace in-place
-                        blended = blend_overlapping_frames(prev_tail, cur_head, temporal_overlap)
-                        ctx['final_video'][current_frame_idx - temporal_overlap:current_frame_idx] = blended
-                        
-                        # Write non-overlapping part
-                        non_overlapping = sample[temporal_overlap:]
-                        non_overlapping_count = non_overlapping.shape[0]
-                        ctx['final_video'][current_frame_idx:current_frame_idx + non_overlapping_count] = non_overlapping
-                        current_frame_idx += non_overlapping_count
-                    else:
-                        # Edge case: write normally if overlap >= batch_frames
-                        ctx['final_video'][current_frame_idx:current_frame_idx + batch_frames] = sample
-                        current_frame_idx += batch_frames
+                    debug.log(f"Unknown color correction method: {color_correction}", level="WARNING", category="video", force=True, indent_level=1)
                 
-                # Immediately release sample memory
-                del sample
+                debug.end_timer(f"color_correction_{color_correction}", f"Color correction ({color_correction})")
                 
-            # Clear batch samples as we go to free memory progressively
-            release_tensor_memory(ctx['batch_samples'][batch_idx])
-            ctx['batch_samples'][batch_idx] = None
+                # Free the reconstructed transformed video
+                del input_video
+
+                # Recombine with Alpha if it was present in input
+                if has_alpha and alpha_channel is not None:
+                    # Concatenate in channels-first: (T, 3, H, W) + (T, 1, H, W) -> (T, 4, H, W)
+                    sample = torch.cat([sample, alpha_channel], dim=1)
             
-            debug.end_timer(f"postprocess_batch_{batch_idx+1}", f"Post-processed batch {batch_idx+1}")
+            else:
+                debug.log("Color correction disabled (set to none)", category="video", indent_level=1)
+            
+            # Convert to final format: [T, C, H, W] → [T, H, W, C]
+            sample = optimized_sample_to_image_format(sample)
+            
+            # Apply normalization only to RGB channels, preserve Alpha as-is
+            if ctx.get('is_rgba', False) and sample.shape[-1] == 4:
+                # Split RGBA: sample is (T, H, W, C) format after optimized_sample_to_image_format
+                rgb_channels = sample[..., :3]  # (T, H, W, 3)
+                alpha_channel = sample[..., 3:4]  # (T, H, W, 1)
+                
+                # Normalize only RGB from [-1, 1] to [0, 1]
+                rgb_channels.clamp_(-1, 1).mul_(0.5).add_(0.5)
+                
+                # Merge back with unchanged Alpha
+                sample = torch.cat([rgb_channels, alpha_channel], dim=-1)
+            else:
+                # RGB only: apply normalization as usual
+                sample.clamp_(-1, 1).mul_(0.5).add_(0.5)
+            
+            # Draw tile boundaries for debugging (if tile info available)
+            for phase, attr in [('encode', 'encode_tile_boundaries'), ('decode', 'decode_tile_boundaries')]:
+                tiles = getattr(debug, attr, None)
+                if tiles:
+                    sample = _draw_tile_boundaries(sample, debug, tiles, phase)
+                    break
+            
+            # Move to final_video device and write back in-place
+            sample = manage_tensor(
+                tensor=sample,
+                target_device=ctx['final_video'].device,
+                tensor_name=f"sample_{info_idx+1}_final",
+                dtype=ctx['compute_dtype'],
+                debug=debug,
+                reason="writing processed result to final_video",
+                indent_level=1
+            )
+            
+            # Write back to final_video in-place
+            # For RGBA, write only RGB channels (alpha already written during alpha processing)
+            if ctx.get('is_rgba', False) and ctx['final_video'].shape[-1] == 4:
+                ctx['final_video'][write_start:write_end, :, :, :3] = sample
+            else:
+                ctx['final_video'][write_start:write_end] = sample
+            
+            # Free sample memory
+            del sample, sample_thwc
+            
+            debug.end_timer(f"postprocess_batch_{info_idx+1}", f"Post-processed batch {info_idx+1}")
             
             # Update progress for main processing step
             current_postprocessing_step += 1
@@ -1332,24 +1399,23 @@ def postprocess_all_batches(
             if total_padding_removed > 0:
                 adjustments.append(f"{total_padding_removed} padding")
             
+            # Use actual temporal overlap from encoding (may have been reset)
+            actual_overlap = ctx.get('actual_temporal_overlap', temporal_overlap)
+            
             # Calculate and include temporal overlap blending info
-            if temporal_overlap > 0:
-                frames_blended = (num_valid_samples - 1) * temporal_overlap
+            if actual_overlap > 0:
+                frames_blended = (num_valid_samples - 1) * actual_overlap
                 adjustments.append(f"{frames_blended} overlap")
 
             if adjustments:
                 # Add back all removed/blended frames to get true computed count
                 total_computed = frames_before_removal + total_padding_removed
-                if temporal_overlap > 0:
-                    total_computed += (num_valid_samples - 1) * temporal_overlap
+                if actual_overlap > 0:
+                    total_computed += (num_valid_samples - 1) * actual_overlap
                 frame_info += f" ({total_computed} computed with {' + '.join(adjustments)} removed)"
             
-            debug.log(f"Final output assembled: {frame_info}, Resolution: {Wf}x{Hf}px, Channels: {channels_str}", 
+            debug.log(f"Output assembled: {frame_info}, Resolution: {Wf}x{Hf}px, Channels: {channels_str}", 
                     category="generation", force=True)
-            
-            if current_frame_idx != total_frames:
-                debug.log(f"WARNING: Frame count mismatch - expected {total_frames}, wrote {current_frame_idx}", 
-                        level="WARNING", category="generation", force=True)
         else:
             ctx['final_video'] = torch.empty((0, 0, 0, 0), dtype=ctx['compute_dtype'])
             debug.log("No frames were processed", level="WARNING", category="generation", force=True)
@@ -1358,11 +1424,11 @@ def postprocess_all_batches(
         debug.log(f"Error in Phase 4 (Post-processing): {e}", level="ERROR", category="generation", force=True)
         raise
     finally:
-        # 1. Clean up batch_samples from context (already mostly freed during processing)
-        if 'batch_samples' in ctx and ctx['batch_samples']:
-            release_tensor_collection(ctx['batch_samples'])
-            ctx['batch_samples'].clear()
-            del ctx['batch_samples']
+        # 1. Clean up decode_batch_info and padding stats
+        if 'decode_batch_info' in ctx:
+            del ctx['decode_batch_info']
+        if 'total_padding_removed' in ctx:
+            del ctx['total_padding_removed']
         
         # 2. Clean up video transform caches
         if 'video_transform' in ctx and ctx['video_transform'] is not None:

src/core/generation_utils.py

@@ -36,6 +36,7 @@ from .model_configuration import configure_runner
 from .infer import VideoDiffusionInfer
 from ..data.image.transforms.divisible_crop import DivisiblePad
 from ..data.image.transforms.na_resize import NaResize
+from ..optimization.compatibility import COMPUTE_DTYPE, BFLOAT16_SUPPORTED
 from ..optimization.memory_manager import manage_tensor
 from ..utils.constants import get_script_directory
 
@@ -371,7 +372,7 @@ def setup_generation_context(
         'dit_offload_device': dit_offload_device,
         'vae_offload_device': vae_offload_device,
         'tensor_offload_device': tensor_offload_device,
-        'compute_dtype': torch.bfloat16, # Hardcoded - gives the best compromise between memory & quality without artifacts
+        'compute_dtype': COMPUTE_DTYPE,
         'interrupt_fn': interrupt_fn,
         'video_transform': None,
         'text_embeds': None,
@@ -401,7 +402,12 @@ def setup_generation_context(
             f"LOCAL_RANK={os.environ['LOCAL_RANK']}",
             category="setup"
         )
-        reason = "quality" if ctx['compute_dtype'] == torch.float32 else "compatibility"
+        if ctx['compute_dtype'] == torch.float32:
+            reason = "quality"
+        elif not BFLOAT16_SUPPORTED:
+            reason = "compatibility (GPU lacks bfloat16 CUBLAS - 7B models unsupported, 3B may have artifacts)"
+        else:
+            reason = "performance"
         debug.log(f"Unified compute dtype: {ctx['compute_dtype']} across entire pipeline for maximum {reason}", category="precision")
     
     return ctx

src/core/infer.py

@@ -26,41 +26,13 @@ from ..common.diffusion import (
 from ..common.distributed import (
     get_device,
 )
+from ..optimization.performance import (
+    optimized_channels_to_last,
+    optimized_channels_to_second
+)
 from ..models.dit_3b import na
 
 
-def optimized_channels_to_last(tensor):
-    """🚀 Optimized replacement for rearrange(tensor, 'b c ... -> b ... c')
-    Moves channels from position 1 to last position using PyTorch native operations.
-    """
-    if tensor.ndim == 3:  # [batch, channels, spatial]
-        return tensor.permute(0, 2, 1)
-    elif tensor.ndim == 4:  # [batch, channels, height, width]
-        return tensor.permute(0, 2, 3, 1)
-    elif tensor.ndim == 5:  # [batch, channels, depth, height, width]
-        return tensor.permute(0, 2, 3, 4, 1)
-    else:
-        # Fallback for other dimensions - move channel (dim=1) to last
-        dims = list(range(tensor.ndim))
-        dims = [dims[0]] + dims[2:] + [dims[1]]  # [0, 2, 3, ..., 1]
-        return tensor.permute(*dims)
-
-def optimized_channels_to_second(tensor):
-    """🚀 Optimized replacement for rearrange(tensor, 'b ... c -> b c ...')
-    Moves channels from last position to position 1 using PyTorch native operations.
-    """
-    if tensor.ndim == 3:  # [batch, spatial, channels]
-        return tensor.permute(0, 2, 1)
-    elif tensor.ndim == 4:  # [batch, height, width, channels]
-        return tensor.permute(0, 3, 1, 2)
-    elif tensor.ndim == 5:  # [batch, depth, height, width, channels]
-        return tensor.permute(0, 4, 1, 2, 3)
-    else:
-        # Fallback for other dimensions - move last dim to position 1
-        dims = list(range(tensor.ndim))
-        dims = [dims[0], dims[-1]] + dims[1:-1]  # [0, -1, 1, 2, ..., -2]
-        return tensor.permute(*dims)
-
 class VideoDiffusionInfer():
     def __init__(self, config: DictConfig, debug: 'Debug',
                  encode_tiled: bool = False, encode_tile_size: Tuple[int, int] = (512, 512), 
@@ -183,7 +155,7 @@ class VideoDiffusionInfer():
 
                 # Use autocast if VAE dtype differs from input dtype
                 if vae_dtype != sample.dtype:
-                    with torch.autocast(str(device), sample.dtype, enabled=True):
+                    with torch.autocast(device.type, sample.dtype, enabled=True):
                         if use_sample:
                             latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size, 
                                                     tile_overlap=self.encode_tile_overlap).latent
@@ -201,7 +173,7 @@ class VideoDiffusionInfer():
                                             tile_overlap=self.encode_tile_overlap).posterior.mode().squeeze(2)
 
                 latent = latent.unsqueeze(2) if latent.ndim == 4 else latent
-                latent = rearrange(latent, "b c ... -> b ... c")
+                latent = optimized_channels_to_last(latent)
                 latent = (latent - shift) * scale
                 latents.append(latent)
 
@@ -248,7 +220,7 @@ class VideoDiffusionInfer():
 
             for i, latent in enumerate(latents):
                 latent = latent / scale + shift
-                latent = rearrange(latent, "b ... c -> b c ...")
+                latent = optimized_channels_to_second(latent)
                 latent = latent.squeeze(2)
 
                 # Detect VAE model dtype
@@ -259,7 +231,7 @@ class VideoDiffusionInfer():
 
                 # Use autocast if VAE dtype differs from latent dtype
                 if vae_dtype != latent.dtype:
-                    with torch.autocast(str(device), latent.dtype, enabled=True):
+                    with torch.autocast(device.type, latent.dtype, enabled=True):
                         sample = self.vae.decode(
                             latent,
                             tiled=self.decode_tiled, tile_size=self.decode_tile_size,

src/data/image/transforms/area_resize.py

@@ -19,6 +19,7 @@ import torch
 from PIL import Image
 from torchvision.transforms import functional as TVF
 from torchvision.transforms.functional import InterpolationMode
+from ....optimization.memory_manager import is_mps_available
 
 
 class AreaResize:
@@ -31,7 +32,7 @@ class AreaResize:
         self.max_area = max_area
         self.downsample_only = downsample_only
         self.interpolation = interpolation
-        if hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
+        if is_mps_available():
             self.interpolation = InterpolationMode.BILINEAR
 
     def __call__(self, image: Union[torch.Tensor, Image.Image]):
@@ -50,10 +51,12 @@ class AreaResize:
 
         resized_height, resized_width = round(height * scale), round(width * scale)
 
+        antialias = not (isinstance(image, torch.Tensor) and image.device.type == 'mps')
         return TVF.resize(
             image,
             size=(resized_height, resized_width),
             interpolation=self.interpolation,
+            antialias=antialias,
         )
 
 

src/data/image/transforms/na_resize.py

@@ -18,6 +18,7 @@ from torchvision.transforms import CenterCrop, Compose, InterpolationMode, Resiz
 
 from .area_resize import AreaResize
 from .side_resize import SideResize
+from ....optimization.memory_manager import is_mps_available
 
 def NaResize(
     resolution: int,
@@ -26,7 +27,7 @@ def NaResize(
     max_resolution: int = 0,
     interpolation: InterpolationMode = InterpolationMode.BICUBIC,
 ):
-    Interpolation = InterpolationMode.BILINEAR if (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()) else interpolation
+    Interpolation = InterpolationMode.BILINEAR if is_mps_available() else interpolation
     if mode == "area":
         return AreaResize(
             max_area=resolution**2,

src/data/image/transforms/side_resize.py

@@ -17,6 +17,7 @@ import torch
 from PIL import Image
 from torchvision.transforms import InterpolationMode
 from torchvision.transforms import functional as TVF
+from ....optimization.memory_manager import is_mps_available
 
 class SideResize:
     def __init__(
@@ -30,7 +31,7 @@ class SideResize:
         self.max_size = max_size
         self.downsample_only = downsample_only
         self.interpolation = interpolation
-        if hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
+        if is_mps_available():
             self.interpolation = InterpolationMode.BILINEAR
 
     def __call__(self, image: Union[torch.Tensor, Image.Image]):
@@ -56,8 +57,9 @@ class SideResize:
         else:
             size = self.size
 
-        # Resize to shortest edge
-        resized = TVF.resize(image, size, self.interpolation)
+        # Resize to shortest edge (disable antialias only for MPS tensors - not supported)
+        antialias = not (isinstance(image, torch.Tensor) and image.device.type == 'mps')
+        resized = TVF.resize(image, size, self.interpolation, antialias=antialias)
         
         # Apply max_size constraint if specified
         if self.max_size > 0:
@@ -69,6 +71,6 @@ class SideResize:
             if max(h, w) > self.max_size:
                 scale = self.max_size / max(h, w)
                 new_h, new_w = round(h * scale), round(w * scale)
-                resized = TVF.resize(resized, (new_h, new_w), self.interpolation)
+                resized = TVF.resize(resized, (new_h, new_w), self.interpolation, antialias=antialias)
         
         return resized

src/interfaces/video_upscaler.py

@@ -509,15 +509,21 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
             )
 
             sample = ctx['final_video']
-            
+            debug.log("", category="none", force=True)
+
             # Ensure CPU tensor in float32 for maximum ComfyUI compatibility
             if torch.is_tensor(sample):
                 if sample.is_cuda or sample.is_mps:
                     sample = sample.cpu()
                 if sample.dtype != torch.float32:
-                    sample = sample.to(torch.float32)
+                    src_dtype = sample.dtype
+                    try:
+                        sample = sample.to(torch.float32)
+                        debug.log(f"Converted output from {src_dtype} to float32", category="precision")
+                    except Exception as e:
+                        debug.log(f"Could not convert to float32: {e}. Output is {src_dtype}, compatibility with other nodes not guaranteed", 
+                                  level="WARNING", category="precision", force=True)
 
-            debug.log("", category="none", force=True)
             debug.log("Upscaling completed successfully!", category="success", force=True)
             debug.end_timer("generation", "Video generation")
 
@@ -529,7 +535,7 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
             debug.log_memory_state("After all phases complete", show_tensors=False, detailed_tensors=False)
             
             # Final peak vram summary
-            debug.log_peak_vram_summary()
+            debug.log_peak_memory_summary()
 
             # Final timing summary
             debug.log("", category="none")

src/models/video_vae_v3/modules/attn_video_vae.py

@@ -1308,7 +1308,8 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
         if H <= tile_h and W <= tile_w:
             return self.slicing_encode(x)
         else:
-            self.debug.log(f"Using VAE tiled encoding (Tile: {tile_size}, Overlap: {tile_overlap})", category="vae", force=True, indent_level=1)
+            if self.debug:
+                self.debug.log(f"Using VAE tiled encoding (Tile: {tile_size}, Overlap: {tile_overlap})", category="vae", force=True, indent_level=1)
 
         # Spatial scale factor (output/latent)
         scale_factor = self.spatial_downsample_factor
@@ -1481,7 +1482,8 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
         if H <= latent_tile_h and W <= latent_tile_w:
             return self.slicing_decode(z)
         else:
-            self.debug.log(f"Using VAE tiled decoding (Tile: {tile_size}, Overlap: {tile_overlap})", category="vae", force=True, indent_level=1)
+            if self.debug:
+                self.debug.log(f"Using VAE tiled decoding (Tile: {tile_size}, Overlap: {tile_overlap})", category="vae", force=True, indent_level=1)
         
         latent_overlap_h = max(0, min((overlap_h // scale_factor), latent_tile_h - 1))
         latent_overlap_w = max(0, min((overlap_w // scale_factor), latent_tile_w - 1))

src/optimization/blockswap.py

@@ -144,22 +144,21 @@ def apply_block_swap_to_dit(
 
     total_blocks = len(model.blocks)
     
-    # Log configuration clearly based on what's enabled
-    block_text = "block" if blocks_to_swap <= 1 else "blocks"
-    if blocks_to_swap > 0 and swap_io_components:
-        debug.log(f"BlockSwap: {blocks_to_swap} transformer {block_text} + I/O components offloaded to {str(offload_device).upper()}", category="blockswap", force=True)
-    elif blocks_to_swap > 0:
-        debug.log(f"BlockSwap: {blocks_to_swap} transformer {block_text} offloaded to {str(offload_device).upper()}", category="blockswap", force=True)
-    elif swap_io_components:
-        debug.log(f"BlockSwap: I/O components offloaded to {str(offload_device).upper()} (blocks remain on GPU)", category="blockswap", force=True)
+    # Clamp blocks_to_swap to available blocks BEFORE logging
+    effective_blocks = min(blocks_to_swap, total_blocks) if blocks_to_swap > 0 else 0
     
-    debug.log(f"Model has {total_blocks} transformer blocks", category="blockswap")
+    # Log configuration clearly based on what's enabled
+    block_text = "block" if effective_blocks <= 1 else "blocks"
+    if effective_blocks > 0 and swap_io_components:
+        debug.log(f"BlockSwap: {effective_blocks}/{total_blocks} transformer {block_text} + I/O components offloaded to {str(offload_device).upper()}", category="blockswap", force=True)
+    elif effective_blocks > 0:
+        debug.log(f"BlockSwap: {effective_blocks}/{total_blocks} transformer {block_text} offloaded to {str(offload_device).upper()}", category="blockswap", force=True)
+    elif swap_io_components:
+        debug.log(f"BlockSwap: I/O components offloaded to {str(offload_device).upper()} (0/{total_blocks} blocks swapped)", category="blockswap", force=True)
     
     # Configure model with blockswap attributes
     if blocks_to_swap > 0:
-        blocks_to_swap = min(blocks_to_swap, total_blocks)
-        model.blocks_to_swap = blocks_to_swap - 1  # Convert to 0-indexed
-        debug.log(f"Transformer blocks to swap: {blocks_to_swap}/{total_blocks}", category="blockswap")
+        model.blocks_to_swap = effective_blocks - 1  # Convert to 0-indexed
     else:
         # No block swapping, set to -1 so no blocks match the swap condition
         model.blocks_to_swap = -1
@@ -193,7 +192,7 @@ def apply_block_swap_to_dit(
     runner._blockswap_active = True
 
     # Store configuration for debugging and cleanup
-    runner._block_swap_config = {
+    model._block_swap_config = {
         "blocks_swapped": blocks_to_swap,
         "swap_io_components": swap_io_components,
         "total_blocks": total_blocks,
@@ -650,11 +649,11 @@ def _protect_model_from_move(
     
     Wraps model.to() method to prevent other code from accidentally moving
     the entire model to GPU, which would defeat BlockSwap's memory savings.
-    Allows movement only when explicitly bypassed via runner flag.
+    Allows movement only when explicitly bypassed via model flag.
     
     Args:
         model: DiT model to protect
-        runner: VideoDiffusionInfer instance (stored as weak reference)
+        runner: VideoDiffusionInfer instance (for active status check)
         debug: Debug instance for logging (required)
     """
     if not hasattr(model, '_original_to'):
@@ -665,34 +664,46 @@ def _protect_model_from_move(
         # Define the protected method without closures
         def protected_model_to(self, device, *args, **kwargs):
             # Check if protection is temporarily bypassed for offloading
+            # Flag is stored on model itself (not runner) to survive runner recreation
+            if getattr(self, "_blockswap_bypass_protection", False):
+                # Protection bypassed, allow movement
+                if hasattr(self, '_original_to'):
+                    return self._original_to(device, *args, **kwargs)
+            
+            # Get configured offload device directly from model
+            blockswap_offload_device = "cpu"  # default
+            if hasattr(self, "_block_swap_config"):
+                blockswap_offload_device = self._block_swap_config.get("offload_device", "cpu")
+            
+            # Check if BlockSwap is currently active via runner weak reference
             runner_ref = getattr(self, '_blockswap_runner_ref', None)
+            blockswap_is_active = False
             if runner_ref:
                 runner_obj = runner_ref()
-                if runner_obj and getattr(runner_obj, "_blockswap_bypass_protection", False):
-                    # Protection bypassed, allow movement
-                    if hasattr(self, '_original_to'):
-                        return self._original_to(device, *args, **kwargs)
+                if runner_obj and hasattr(runner_obj, "_blockswap_active"):
+                    blockswap_is_active = runner_obj._blockswap_active
             
-            # Check blockswap status using weak reference
-            # Get configured offload device from runner
-            blockswap_offload_device = "cpu" # default
-            if runner_ref:
-                runner_obj = runner_ref()
-                if runner_obj and hasattr(runner_obj, "_block_swap_config"):
-                    blockswap_offload_device = runner_obj._block_swap_config.get("offload_device", "cpu")
+            # Block attempts to move model away from configured offload device when active
+            if blockswap_is_active and str(device) != str(blockswap_offload_device):
+                # Get debug instance from runner if available
+                debug_instance = None
+                if runner_ref:
+                    runner_obj = runner_ref()
+                    if runner_obj and hasattr(runner_obj, 'debug'):
+                        debug_instance = runner_obj.debug
                 
-                # Block attempts to move model away from configured offload device
-                if str(device) != str(blockswap_offload_device):
-                    if runner_obj and hasattr(runner_obj, "_blockswap_active") and runner_obj._blockswap_active:
-                        debug.log(f"Blocked attempt to move blockswapped model from {blockswap_offload_device} to {device}", 
-                                level="WARNING", category="blockswap", force=True)
-                        return self
+                if debug_instance:
+                    debug_instance.log(
+                        f"Blocked attempt to move BlockSwap model from {blockswap_offload_device} to {device}",
+                        level="WARNING", category="blockswap", force=True
+                    )
+                return self
             
-            # Use original method stored as attribute
+            # Allow movement (either bypass is enabled or target is offload device)
             if hasattr(self, '_original_to'):
                 return self._original_to(device, *args, **kwargs)
             else:
-                # This shouldn't happen, but fallback to super().to()
+                # Fallback - shouldn't happen
                 return super(type(self), self).to(device, *args, **kwargs)
         
         # Bind as a method to the model instance
@@ -712,7 +723,13 @@ def set_blockswap_bypass(runner, bypass: bool, debug):
     if not hasattr(runner, "_blockswap_active") or not runner._blockswap_active:
         return
     
-    runner._blockswap_bypass_protection = bypass
+    # Get the actual model (handle FP8CompatibleDiT wrapper)
+    model = runner.dit
+    if hasattr(model, "dit_model"):
+        model = model.dit_model
+    
+    # Store on model so it survives runner recreation during caching
+    model._blockswap_bypass_protection = bypass
     
     if bypass:
         debug.log("BlockSwap protection disabled to allow model DiT offloading", category="success")
@@ -741,11 +758,16 @@ def cleanup_blockswap(runner, keep_state_for_cache=False):
     
     debug = runner.debug
     
-    # Check if there's any BlockSwap state to clean up
+    # Get the actual model (handle FP8CompatibleDiT wrapper)
+    model = runner.dit
+    if hasattr(model, "dit_model"):
+        model = model.dit_model
+    
+    # Check if there's any BlockSwap state to clean up (check both runner and model)
     has_blockswap_state = (
         hasattr(runner, "_blockswap_active") or 
-        hasattr(runner, "_block_swap_config") or
-        hasattr(runner, "_blockswap_bypass_protection")
+        hasattr(model, "_block_swap_config") or
+        hasattr(model, "_blockswap_bypass_protection")
     )
     
     if not has_blockswap_state:
@@ -757,7 +779,7 @@ def cleanup_blockswap(runner, keep_state_for_cache=False):
         # Minimal cleanup for caching - just mark as inactive and allow offloading
         # Everything else stays intact for fast reactivation
         if hasattr(runner, "_blockswap_active") and runner._blockswap_active:
-            if not getattr(runner, "_blockswap_bypass_protection", False):
+            if not getattr(model, "_blockswap_bypass_protection", False):
                 set_blockswap_bypass(runner=runner, bypass=True, debug=debug)
             runner._blockswap_active = False
         debug.log("BlockSwap deactivated for caching (configuration preserved)", category="success")
@@ -829,7 +851,7 @@ def cleanup_blockswap(runner, keep_state_for_cache=False):
 
     # 5. Clean up BlockSwap-specific attributes
     for attr in ['_blockswap_runner_ref', 'blocks_to_swap', 'main_device', 
-                 'offload_device', '_blockswap_configured']:
+                 'offload_device']:
         if hasattr(model, attr):
             delattr(model, attr)
 

src/optimization/compatibility.py

@@ -5,16 +5,99 @@ Contains FP8/FP16 compatibility layers and wrappers for different model architec
 Extracted from: seedvr2.py (lines 1045-1630)
 """
 
-import torch
+# Compatibility shims - Must run before any torch/diffusers import
+import sys
 import types
+
+
+def ensure_triton_compat():
+    """Create minimal triton.ops stubs only if missing, to allow bitsandbytes import."""
+    if 'triton.ops.matmul_perf_model' in sys.modules:
+        return
+    
+    try:
+        from triton.ops.matmul_perf_model import early_config_prune  # noqa: F401
+        return
+    except (ImportError, ModuleNotFoundError, AttributeError):
+        pass
+    
+    if 'triton.ops' not in sys.modules:
+        sys.modules['triton.ops'] = types.ModuleType('triton.ops')
+    
+    matmul_perf = types.ModuleType('triton.ops.matmul_perf_model')
+    matmul_perf.early_config_prune = lambda configs, *a, **kw: configs
+    matmul_perf.estimate_matmul_time = lambda *a, **kw: 0.0
+    
+    sys.modules['triton.ops'].matmul_perf_model = matmul_perf
+    sys.modules['triton.ops.matmul_perf_model'] = matmul_perf
+
+
+def ensure_flash_attn_safe():
+    """
+    Pre-test flash_attn package; stub if DLL is broken.
+    Prevents diffusers from crashing when flash_attn has broken DLLs.
+    """
+    if 'flash_attn' in sys.modules:
+        return  # Already loaded
+    
+    try:
+        import flash_attn
+    except (ImportError, OSError):
+        # DLL broken or not installed - create stub with proper __spec__
+        import importlib.machinery
+        
+        stub = types.ModuleType('flash_attn')
+        stub.__spec__ = importlib.machinery.ModuleSpec('flash_attn', None)
+        stub.__file__ = None
+        stub.__path__ = []
+        stub.__loader__ = None
+        # Provide attributes that diffusers/transformers import
+        stub.flash_attn_func = None
+        stub.flash_attn_varlen_func = None
+        sys.modules['flash_attn'] = stub
+
+
+def ensure_xformers_flash_compat():
+    """
+    Pre-test xformers._C_flashattention; stub if DLL is broken.
+    Prevents xformers.ops.fmha.flash from crashing on import.
+    """
+    if 'xformers._C_flashattention' in sys.modules:
+        return  # Already loaded
+    
+    try:
+        from xformers import _C_flashattention  # noqa: F401
+    except (ImportError, OSError):
+        # DLL broken or not installed - create stub that fails gracefully
+        class _FailingStub(types.ModuleType):
+            """Stub that lets xformers gracefully disable its flash backend."""
+            def __getattr__(self, name):
+                # Dunder attributes: raise AttributeError (normal Python behavior)
+                if name.startswith('__') and name.endswith('__'):
+                    raise AttributeError(name)
+                # xformers functional attributes: raise ImportError so xformers catches it
+                raise ImportError("_C_flashattention unavailable")
+        sys.modules['xformers._C_flashattention'] = _FailingStub('xformers._C_flashattention')
+
+
+# Run all shims immediately on import, before torch/diffusers
+ensure_triton_compat()
+ensure_flash_attn_safe()
+ensure_xformers_flash_compat()
+
+
+import torch
 import os
 
+
 # Flash Attention & Triton Compatibility Layer
 # 1. Flash Attention - speedup for attention operations
 try:
     from flash_attn import flash_attn_varlen_func
+    # Force load the CUDA extension to verify it's not corrupted
+    import flash_attn_2_cuda  # noqa: F401
     FLASH_ATTN_AVAILABLE = True
-except ImportError:
+except (ImportError, AttributeError, OSError):
     flash_attn_varlen_func = None
     FLASH_ATTN_AVAILABLE = False
 
@@ -171,6 +254,24 @@ if not os.environ.get("SEEDVR2_OPTIMIZATIONS_LOGGED"):
         print(f"🔧 Conv3d workaround active: PyTorch {torch_ver}, cuDNN {cudnn_ver} (fixing VAE 3x memory bug)")
 
 
+# Bfloat16 CUBLAS support
+def _probe_bfloat16_support() -> bool:
+    if not torch.cuda.is_available():
+        return True
+    try:
+        a = torch.randn(8, 8, dtype=torch.bfloat16, device='cuda:0')
+        _ = torch.matmul(a, a)
+        del a
+        return True
+    except RuntimeError as e:
+        if "CUBLAS_STATUS_NOT_SUPPORTED" in str(e):
+            return False
+        raise
+
+BFLOAT16_SUPPORTED = _probe_bfloat16_support()
+COMPUTE_DTYPE = torch.bfloat16 if BFLOAT16_SUPPORTED else torch.float16
+
+
 def call_rope_with_stability(method, *args, **kwargs):
     """
     Call RoPE method with stability fixes:
@@ -221,7 +322,7 @@ class FP8CompatibleDiT(torch.nn.Module):
             self._convert_rope_freqs(target_dtype=self.compute_dtype)
             self.debug.end_timer("_convert_rope_freqs", "RoPE freqs conversion")
             
-            if hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
+            if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                 self.debug.log(f"Also converting NaDiT parameters/buffers for MPS backend", category="setup", force=True)
                 self.debug.start_timer("_force_nadit_precision")
                 self._force_nadit_precision(target_dtype=self.compute_dtype)
@@ -232,11 +333,6 @@ class FP8CompatibleDiT(torch.nn.Module):
         self.debug.start_timer("_stabilize_rope_computations")
         self._stabilize_rope_computations()
         self.debug.end_timer("_stabilize_rope_computations", "RoPE stabilization")
-
-        # 🚀 FLASH ATTENTION OPTIMIZATION (Phase 2)
-        self.debug.start_timer("_apply_flash_attention_optimization")
-        self._apply_flash_attention_optimization()
-        self.debug.end_timer("_apply_flash_attention_optimization", "Flash Attention application")
     
     def _detect_model_dtype(self) -> torch.dtype:
         """Detect main model dtype"""
@@ -359,205 +455,6 @@ class FP8CompatibleDiT(torch.nn.Module):
         
         if rope_count > 0:
             self.debug.log(f"Stabilized {rope_count} RoPE modules", category="success")
-
-    def _apply_flash_attention_optimization(self) -> None:
-        """🚀 FLASH ATTENTION OPTIMIZATION - 30-50% speedup of attention layers"""
-        attention_layers_optimized = 0
-        flash_attention_available = self._check_flash_attention_support()
-        
-        for name, module in self.dit_model.named_modules():
-            # Identify all attention layers
-            if self._is_attention_layer(name, module):
-                # Apply optimization based on availability
-                if self._optimize_attention_layer(name, module, flash_attention_available):
-                    attention_layers_optimized += 1
-        
-        if not flash_attention_available:
-            self.debug.log("Flash Attention not available, using PyTorch SDPA as fallback", category="info", force=True)
-    
-    def _check_flash_attention_support(self) -> bool:
-        """Check if Flash Attention is available"""
-        # Check PyTorch SDPA (includes Flash Attention on H100/A100)
-        if hasattr(torch.nn.functional, 'scaled_dot_product_attention'):
-            return True
-        
-        # Check flash-attn package (uses module-level check from top of file)
-        return FLASH_ATTN_AVAILABLE
-    
-    def _is_attention_layer(self, name: str, module: torch.nn.Module) -> bool:
-        """Identify if a module is an attention layer"""
-        attention_keywords = [
-            'attention', 'attn', 'self_attn', 'cross_attn', 'mhattn', 'multihead',
-            'transformer_block', 'dit_block'
-        ]
-        
-        # Check by name
-        if any(keyword in name.lower() for keyword in attention_keywords):
-            return True
-        
-        # Check by module type
-        module_type = type(module).__name__.lower()
-        if any(keyword in module_type for keyword in attention_keywords):
-            return True
-        
-        # Check by attributes (modules with q, k, v projections)
-        if hasattr(module, 'q_proj') or hasattr(module, 'qkv') or hasattr(module, 'to_q'):
-            return True
-        
-        return False
-    
-    def _optimize_attention_layer(self, name: str, module: torch.nn.Module, flash_attention_available: bool) -> bool:
-        """Optimize a specific attention layer"""
-        try:
-            # Save original forward method
-            if not hasattr(module, '_original_forward'):
-                module._original_forward = module.forward
-            
-            # Create new optimized forward method
-            if flash_attention_available:
-                optimized_forward = self._create_flash_attention_forward(module, name)
-            else:
-                optimized_forward = self._create_sdpa_forward(module, name)
-            
-            # Replace forward method
-            module.forward = optimized_forward
-            return True
-            
-        except Exception as e:
-            self.debug.log(f"Failed to optimize attention layer '{name}': {e}", level="WARNING", category="dit", force=True)
-            return False
-    
-    def _create_flash_attention_forward(self, module: torch.nn.Module, layer_name: str):
-        """Create optimized forward with Flash Attention"""
-        original_forward = module._original_forward
-        
-        def flash_attention_forward(*args, **kwargs):
-            try:
-                # Try to use Flash Attention via SDPA
-                return self._sdpa_attention_forward(original_forward, module, *args, **kwargs)
-            except Exception as e:
-                # Fallback to original implementation
-                self.debug.log(f"Flash Attention failed for {layer_name}, using original: {e}", level="WARNING", category="dit", force=True)
-                return original_forward(*args, **kwargs)
-        
-        return flash_attention_forward
-    
-    def _create_sdpa_forward(self, module: torch.nn.Module, layer_name: str):
-        """Create optimized forward with PyTorch SDPA"""
-        original_forward = module._original_forward
-        
-        def sdpa_forward(*args, **kwargs):
-            try:
-                return self._sdpa_attention_forward(original_forward, module, *args, **kwargs)
-            except Exception as e:
-                # Fallback to original implementation
-                return original_forward(*args, **kwargs)
-        
-        return sdpa_forward
-    
-    def _sdpa_attention_forward(self, original_forward, module: torch.nn.Module, *args, **kwargs):
-        """Optimized forward pass using SDPA (Scaled Dot Product Attention)"""
-        # Detect if we can intercept and optimize this layer
-        if len(args) >= 1 and isinstance(args[0], torch.Tensor):
-            input_tensor = args[0]
-            
-            # Check dimensions to ensure it's standard attention
-            if len(input_tensor.shape) >= 3:  # [batch, seq_len, hidden_dim] or similar
-                try:
-                    return self._optimized_attention_computation(module, input_tensor, *args[1:], **kwargs)
-                except:
-                    pass
-        
-        # Fallback to original implementation
-        return original_forward(*args, **kwargs)
-    
-    def _optimized_attention_computation(self, module: torch.nn.Module, input_tensor: torch.Tensor, *args, **kwargs):
-        """Optimized attention computation with SDPA"""
-        # Try to detect standard attention format
-        batch_size, seq_len = input_tensor.shape[:2]
-        
-        # Check if module has standard Q, K, V projections
-        if hasattr(module, 'qkv') or (hasattr(module, 'q_proj') and hasattr(module, 'k_proj') and hasattr(module, 'v_proj')):
-            return self._compute_sdpa_attention(module, input_tensor, *args, **kwargs)
-        
-        # If no standard format detected, use original
-        return module._original_forward(input_tensor, *args, **kwargs)
-    
-    def _compute_sdpa_attention(self, module: torch.nn.Module, x: torch.Tensor, *args, **kwargs):
-        """Optimized SDPA computation for standard attention modules"""
-        try:
-            # Case 1: Module with combined QKV projection
-            if hasattr(module, 'qkv'):
-                qkv = module.qkv(x)
-                # Reshape to separate Q, K, V
-                batch_size, seq_len, _ = qkv.shape
-                qkv = qkv.reshape(batch_size, seq_len, 3, -1)
-                q, k, v = qkv.unbind(dim=2)
-                
-            # Case 2: Separate Q, K, V projections
-            elif hasattr(module, 'q_proj') and hasattr(module, 'k_proj') and hasattr(module, 'v_proj'):
-                q = module.q_proj(x)
-                k = module.k_proj(x)
-                v = module.v_proj(x)
-            else:
-                # Unsupported format, use original
-                return module._original_forward(x, *args, **kwargs)
-            
-            # Detect number of heads
-            head_dim = getattr(module, 'head_dim', None)
-            num_heads = getattr(module, 'num_heads', None)
-            
-            if head_dim is None or num_heads is None:
-                # Try to guess from dimensions
-                hidden_dim = q.shape[-1]
-                if hasattr(module, 'num_heads'):
-                    num_heads = module.num_heads
-                    head_dim = hidden_dim // num_heads
-                else:
-                    # Reasonable defaults
-                    head_dim = 64
-                    num_heads = hidden_dim // head_dim
-            
-            # Reshape for multi-head attention
-            batch_size, seq_len = q.shape[:2]
-            q = q.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
-            k = k.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
-            v = v.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
-            
-            if hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
-                attn_output = torch.nn.functional.scaled_dot_product_attention(
-                    q, k, v,
-                    dropout_p=0.0,
-                    is_causal=False
-                )
-            else:
-                # Use optimized SDPA
-                with torch.backends.cuda.sdp_kernel(
-                    enable_flash=True,
-                    enable_math=True,
-                    enable_mem_efficient=True
-                ):
-                    attn_output = torch.nn.functional.scaled_dot_product_attention(
-                        q, k, v,
-                        dropout_p=0.0,
-                        is_causal=False
-                    )
-            
-            # Reshape back
-            attn_output = attn_output.transpose(1, 2).contiguous().view(
-                batch_size, seq_len, num_heads * head_dim
-            )
-            
-            # Output projection if it exists
-            if hasattr(module, 'out_proj') or hasattr(module, 'o_proj'):
-                proj = getattr(module, 'out_proj', None) or getattr(module, 'o_proj', None)
-                attn_output = proj(attn_output)
-            
-            return attn_output
-            
-        except Exception as e:
-            # In case of error, use original implementation
-            return module._original_forward(x, *args, **kwargs)
     
     def forward(self, *args, **kwargs):
         """

src/optimization/memory_manager.py

@@ -10,8 +10,40 @@ import gc
 import sys
 import time
 import psutil
+import platform
 from typing import Tuple, Dict, Any, Optional, List, Union
+
+
+def _device_str(device: Union[torch.device, str]) -> str:
+    """Normalized uppercase device string for comparison and logging. MPS variants → 'MPS'."""
+    s = str(device).upper()
+    return 'MPS' if s.startswith('MPS') else s
+
+
+def is_mps_available() -> bool:
+    """Check if MPS (Apple Metal) backend is available."""
+    return hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()
+
+
+def is_cuda_available() -> bool:
+    """Check if CUDA backend is available."""
+    return torch.cuda.is_available()
+
+
+def get_gpu_backend() -> str:
+    """Get the active GPU backend type.
     
+    Returns:
+        'cuda': NVIDIA CUDA
+        'mps': Apple Metal Performance Shaders
+        'cpu': No GPU backend available
+    """
+    if is_cuda_available():
+        return 'cuda'
+    if is_mps_available():
+        return 'mps'
+    return 'cpu'
+
 
 def get_device_list(include_none: bool = False, include_cpu: bool = False) -> List[str]:
     """
@@ -31,14 +63,14 @@ def get_device_list(include_none: bool = False, include_cpu: bool = False) -> Li
     has_mps = False
     
     try:
-        if hasattr(torch, "cuda") and hasattr(torch.cuda, "is_available") and torch.cuda.is_available():
+        if is_cuda_available():
             devs += [f"cuda:{i}" for i in range(torch.cuda.device_count())]
             has_cuda = True
     except Exception:
         pass
     
     try:
-        if hasattr(torch, "mps") and hasattr(torch.mps, "is_available") and torch.mps.is_available():
+        if is_mps_available():
             devs.append("mps")  # MPS doesn't use device indices
             has_mps = True
     except Exception:
@@ -60,7 +92,7 @@ def get_device_list(include_none: bool = False, include_cpu: bool = False) -> Li
     result.extend(devs)
     
     return result if result else []
-    
+
 
 def get_basic_vram_info(device: Optional[torch.device] = None) -> Dict[str, Any]:
     """
@@ -74,13 +106,13 @@ def get_basic_vram_info(device: Optional[torch.device] = None) -> Dict[str, Any]
         dict: {"free_gb": float, "total_gb": float} or {"error": str}
     """
     try:
-        if torch.cuda.is_available():
+        if is_cuda_available():
             if device is None:
                 device = torch.device("cuda:0")
             elif not isinstance(device, torch.device):
                 device = torch.device(device)
             free_memory, total_memory = torch.cuda.mem_get_info(device)
-        elif hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
+        elif is_mps_available():
             # MPS doesn't support per-device queries or mem_get_info
             # Use system memory as proxy
             mem = psutil.virtual_memory()
@@ -100,13 +132,13 @@ def get_basic_vram_info(device: Optional[torch.device] = None) -> Dict[str, Any]
 # Initial VRAM check at module load
 vram_info = get_basic_vram_info(device=None)
 if "error" not in vram_info:
-    backend = "MPS" if (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()) else "CUDA"
+    backend = "MPS" if is_mps_available() else "CUDA"
     print(f"📊 Initial {backend} memory: {vram_info['free_gb']:.2f}GB free / {vram_info['total_gb']:.2f}GB total")
 else:
     print(f"⚠️ Memory check failed: {vram_info['error']} - No available backend!")
 
 
-def get_vram_usage(device: Optional[torch.device] = None, debug: Optional['Debug'] = None) -> Tuple[float, float, float]:
+def get_vram_usage(device: Optional[torch.device] = None, debug: Optional['Debug'] = None) -> Tuple[float, float, float, float]:
     """
     Get current VRAM usage metrics for monitoring.
     Used for tracking memory consumption during processing.
@@ -116,29 +148,30 @@ def get_vram_usage(device: Optional[torch.device] = None, debug: Optional['Debug
         debug: Optional debug instance for logging
     
     Returns:
-        tuple: (allocated_gb, reserved_gb, max_allocated_gb)
-               Returns (0, 0, 0) if no GPU available
+        tuple: (allocated_gb, reserved_gb, peak_allocated_gb, peak_reserved_gb)
+               Returns (0, 0, 0, 0) if no GPU available
     """
     try:
-        if torch.cuda.is_available():
+        if is_cuda_available():
             if device is None:
                 device = torch.device("cuda:0")
             elif not isinstance(device, torch.device):
                 device = torch.device(device)
             allocated = torch.cuda.memory_allocated(device) / (1024**3)
             reserved = torch.cuda.memory_reserved(device) / (1024**3)
-            max_allocated = torch.cuda.max_memory_allocated(device) / (1024**3)
-            return allocated, reserved, max_allocated
-        elif hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
+            peak_allocated = torch.cuda.max_memory_allocated(device) / (1024**3)
+            peak_reserved = torch.cuda.max_memory_reserved(device) / (1024**3)
+            return allocated, reserved, peak_allocated, peak_reserved
+        elif is_mps_available():
             # MPS doesn't support per-device queries - uses global memory tracking
             allocated = torch.mps.current_allocated_memory() / (1024**3)
             reserved = torch.mps.driver_allocated_memory() / (1024**3)
-            max_allocated = allocated  # MPS doesn't track peak separately
-            return allocated, reserved, max_allocated
+            # MPS doesn't track peak separately
+            return allocated, reserved, allocated, reserved
     except Exception as e:
         if debug:
             debug.log(f"Failed to get VRAM usage: {e}", level="WARNING", category="memory", force=True)
-    return 0.0, 0.0, 0.0
+    return 0.0, 0.0, 0.0, 0.0
 
 
 def get_ram_usage(debug: Optional['Debug'] = None) -> Tuple[float, float, float, float]:
@@ -229,17 +262,17 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo
         # Use existing function for memory info
         mem_info = get_basic_vram_info(device=None)
         
-        if "error" not in mem_info:
+        if "error" not in mem_info and mem_info["total_gb"] > 0:
             # Check VRAM/MPS memory pressure (5% free threshold)
             free_ratio = mem_info["free_gb"] / mem_info["total_gb"]
             if free_ratio < 0.05:
                 should_clear = True
                 if debug:
-                    backend = "MPS" if (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()) else "VRAM"
+                    backend = "Unified Memory" if is_mps_available() else "VRAM"
                     debug.log(f"{backend} pressure: {mem_info['free_gb']:.2f}GB free of {mem_info['total_gb']:.2f}GB", category="memory")
         
         # For non-MPS systems, also check system RAM separately
-        if not should_clear and not (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()):
+        if not should_clear and not is_mps_available():
             mem = psutil.virtual_memory()
             if mem.available < mem.total * 0.05:
                 should_clear = True
@@ -262,10 +295,10 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo
     if debug:
         debug.start_timer(gpu_timer)
     
-    if torch.cuda.is_available():
+    if is_cuda_available():
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
-    elif hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
+    elif is_mps_available():
         torch.mps.empty_cache()
     
     if debug:
@@ -302,7 +335,7 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo
                 handle = _os_memory_lib.GetCurrentProcess()
                 _os_memory_lib.SetProcessWorkingSetSize(handle, -1, -1)
                 
-            elif hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available():
+            elif is_mps_available():
                 # macOS with MPS
                 import ctypes  # Import only when needed
                 import ctypes.util
@@ -379,7 +412,7 @@ def reset_vram_peak(device: Optional[torch.device] = None, debug: Optional['Debu
     if debug and debug.enabled:
         debug.log("Resetting VRAM peak memory statistics", category="memory")
     try:
-        if torch.cuda.is_available():
+        if is_cuda_available():
             if device is None:
                 device = torch.device("cuda:0")
             elif not isinstance(device, torch.device):
@@ -591,7 +624,7 @@ def manage_tensor(
     target_dtype = dtype if dtype is not None else current_dtype
     
     # Check if movement is actually needed
-    needs_device_move = current_device != target_device
+    needs_device_move = _device_str(current_device) != _device_str(target_device)
     needs_dtype_change = dtype is not None and current_dtype != target_dtype
     
     if not needs_device_move and not needs_dtype_change:
@@ -609,8 +642,8 @@ def manage_tensor(
     
     # Log the movement
     if debug:
-        current_device_str = str(current_device).upper()
-        target_device_str = str(target_device).upper()
+        current_device_str = _device_str(current_device)
+        target_device_str = _device_str(target_device)
         
         dtype_info = ""
         if needs_dtype_change:
@@ -681,8 +714,8 @@ def manage_model_device(model: torch.nn.Module, target_device: torch.device, mod
     
     # Extract device type for comparison (both are torch.device objects)
     target_type = target_device.type
-    current_device_upper = str(current_device).upper()
-    target_device_upper = str(target_device).upper()
+    current_device_upper = _device_str(current_device)
+    target_device_upper = _device_str(target_device)
 
     # Compare normalized device types
     if current_device_upper == target_device_upper and not is_blockswap_model:
@@ -737,10 +770,10 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
                 actual_source_device = param.device
                 break
         
-        source_device_desc = str(actual_source_device).upper() if actual_source_device else str(target_device).upper()
+        source_device_desc = _device_str(actual_source_device) if actual_source_device else _device_str(target_device)
         
         if debug:
-            debug.log(f"Moving {model_name} from {source_device_desc} to {str(target_device).upper()} ({reason or 'model caching'})", category="general")
+            debug.log(f"Moving {model_name} from {source_device_desc} to {_device_str(target_device)} ({reason or 'model caching'})", category="general")
         
         # Enable bypass to allow movement
         set_blockswap_bypass(runner=runner, bypass=True, debug=debug)
@@ -755,14 +788,14 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
         model.zero_grad(set_to_none=True)
         
         if debug:
-            debug.end_timer(timer_name, f"BlockSwap model offloaded to {str(target_device).upper()}")
+            debug.end_timer(timer_name, f"BlockSwap model offloaded to {_device_str(target_device)}")
         
         return True
         
     else:
         # Moving to GPU (reload)
         # Check if we're in bypass mode (coming from offload)
-        if not getattr(runner, "_blockswap_bypass_protection", False):
+        if not getattr(model, "_blockswap_bypass_protection", False):
             # Not in bypass mode, blocks are already configured
             if debug:
                 debug.log(f"{model_name} with BlockSwap active - blocks already distributed across devices, skipping movement", category="general")
@@ -775,10 +808,10 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
                 actual_current_device = param.device
                 break
         
-        current_device_desc = str(actual_current_device).upper() if actual_current_device else "OFFLOAD"
+        current_device_desc = _device_str(actual_current_device) if actual_current_device else "OFFLOAD"
         
         if debug:
-            debug.log(f"Moving {model_name} from {current_device_desc} to {str(target_device).upper()} ({reason or 'inference requirement'})", category="general")
+            debug.log(f"Moving {model_name} from {current_device_desc} to {_device_str(target_device)} ({reason or 'inference requirement'})", category="general")
         
         timer_name = f"{model_name.lower()}_to_gpu"
         if debug:
@@ -787,7 +820,7 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
         # Restore blocks to their configured devices
         if hasattr(model, "blocks") and hasattr(model, "blocks_to_swap"):
             # Use configured offload_device from BlockSwap config
-            offload_device = runner._block_swap_config.get("offload_device")
+            offload_device = model._block_swap_config.get("offload_device")
             if not offload_device:
                 raise ValueError("BlockSwap config missing offload_device")
             
@@ -801,7 +834,7 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
                     block.to(offload_device)
             
             # Handle I/O components
-            if not runner._block_swap_config.get("swap_io_components", False):
+            if not model._block_swap_config.get("swap_io_components", False):
                 # I/O components should be on GPU if not offloaded
                 for name, module in model.named_children():
                     if name != "blocks":
@@ -814,11 +847,11 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
             
             if debug:
                 # Get actual configuration from runner
-                if hasattr(runner, '_block_swap_config'):
-                    blocks_on_gpu = runner._block_swap_config.get('total_blocks', 32) - runner._block_swap_config.get('blocks_swapped', 16)
-                    total_blocks = runner._block_swap_config.get('total_blocks', 32)
-                    main_device = runner._block_swap_config.get('main_device', 'GPU')
-                    debug.log(f"BlockSwap blocks restored to configured devices ({blocks_on_gpu}/{total_blocks} blocks on {str(main_device).upper()})", category="success")
+                if hasattr(model, '_block_swap_config'):
+                    blocks_on_gpu = model._block_swap_config.get('total_blocks', 32) - model._block_swap_config.get('blocks_swapped', 16)
+                    total_blocks = model._block_swap_config.get('total_blocks', 32)
+                    main_device = model._block_swap_config.get('main_device', 'GPU')
+                    debug.log(f"BlockSwap blocks restored to configured devices ({blocks_on_gpu}/{total_blocks} blocks on {_device_str(main_device)})", category="success")
                 else:
                     debug.log("BlockSwap blocks restored to configured devices", category="success")
 
@@ -865,8 +898,8 @@ def _standard_model_movement(model: torch.nn.Module, current_device: torch.devic
     
     # Log the movement with full device strings
     if debug:
-        current_device_str = str(current_device).upper()
-        target_device_str = str(target_device).upper()
+        current_device_str = _device_str(current_device)
+        target_device_str = _device_str(target_device)
         debug.log(f"Moving {model_name} from {current_device_str} to {target_device_str} ({reason})", category="general")
 
     # Start timer based on direction
@@ -891,7 +924,7 @@ def _standard_model_movement(model: torch.nn.Module, current_device: torch.devic
     
     # End timer
     if debug:
-        debug.end_timer(timer_name, f"{model_name} moved to {str(target_device).upper()}")
+        debug.end_timer(timer_name, f"{model_name} moved to {_device_str(target_device)}")
     
     return True
 

src/optimization/performance.py

@@ -9,6 +9,40 @@ import torch
 from typing import List
 
 
+def optimized_channels_to_last(tensor):
+    """🚀 Optimized replacement for rearrange(tensor, 'b c ... -> b ... c')
+    Moves channels from position 1 to last position using PyTorch native operations.
+    """
+    if tensor.ndim == 3:  # [batch, channels, spatial]
+        return tensor.permute(0, 2, 1)
+    elif tensor.ndim == 4:  # [batch, channels, height, width]
+        return tensor.permute(0, 2, 3, 1)
+    elif tensor.ndim == 5:  # [batch, channels, depth, height, width]
+        return tensor.permute(0, 2, 3, 4, 1)
+    else:
+        # Fallback for other dimensions - move channel (dim=1) to last
+        dims = list(range(tensor.ndim))
+        dims = [dims[0]] + dims[2:] + [dims[1]]  # [0, 2, 3, ..., 1]
+        return tensor.permute(*dims)
+
+
+def optimized_channels_to_second(tensor):
+    """🚀 Optimized replacement for rearrange(tensor, 'b ... c -> b c ...')
+    Moves channels from last position to position 1 using PyTorch native operations.
+    """
+    if tensor.ndim == 3:  # [batch, spatial, channels]
+        return tensor.permute(0, 2, 1)
+    elif tensor.ndim == 4:  # [batch, height, width, channels]
+        return tensor.permute(0, 3, 1, 2)
+    elif tensor.ndim == 5:  # [batch, depth, height, width, channels]
+        return tensor.permute(0, 4, 1, 2, 3)
+    else:
+        # Fallback for other dimensions - move last dim to position 1
+        dims = list(range(tensor.ndim))
+        dims = [dims[0], dims[-1]] + dims[1:-1]  # [0, -1, 1, 2, ..., -2]
+        return tensor.permute(*dims)
+
+
 def optimized_video_rearrange(video_tensors: List[torch.Tensor]) -> List[torch.Tensor]:
     """
     🚀 OPTIMIZED version of video rearrangement
@@ -132,34 +166,4 @@ def optimized_sample_to_image_format(sample: torch.Tensor) -> torch.Tensor:
         return sample.permute(0, 2, 3, 1)
 
 
-def temporal_latent_blending(latents1: torch.Tensor, latents2: torch.Tensor, blend_frames: int) -> torch.Tensor:
-    """
-    🎨 Temporal blending in latent space to avoid discontinuities
-    
-    Args:
-        latents1: Latents from previous batch (end frames)
-        latents2: Latents from current batch (start frames)
-        blend_frames: Number of frames to blend
-    
-    Returns:
-        Blended latents for smooth transition
-    """
-    if latents1.shape[0] != latents2.shape[0]:
-        # Adjust dimensions if necessary
-        min_frames = min(latents1.shape[0], latents2.shape[0])
-        latents1 = latents1[:min_frames]
-        latents2 = latents2[:min_frames]
-    
-    # Create linear blending weights
-    # Frame 0: 100% latents1, 0% latents2
-    # Frame n: 0% latents1, 100% latents2
-    weights1 = torch.linspace(1.0, 0.0, blend_frames).view(-1, 1, 1, 1).to(latents1.device)
-    weights2 = torch.linspace(0.0, 1.0, blend_frames).view(-1, 1, 1, 1).to(latents2.device)
-    
-    # Apply blending
-    blended_latents = weights1 * latents1 + weights2 * latents2
-    
-    return blended_latents
-
-
 

src/utils/color_fix.py

@@ -178,7 +178,7 @@ def wavelet_decomposition(image: Tensor, levels: int = 5) -> tuple[Tensor, Tenso
     for i in range(levels):
         radius = 2 ** i
         low_freq = wavelet_blur(image, radius)
-        high_freq += (image - low_freq)
+        high_freq.add_(image).sub_(low_freq)
         image = low_freq
     
     return high_freq, low_freq
@@ -237,13 +237,13 @@ def wavelet_reconstruction(content_feat: Tensor, style_feat: Tensor, debug: Opti
             align_corners=False
         )
     
-    # Reconstruct: content details + style color
-    result = content_high_freq + style_low_freq
+    # Reconstruct: content details + style color (in-place on content_high_freq)
+    content_high_freq.add_(style_low_freq)
     
-    # Safety clamp for normalized SDR range
+    # Safety clamp for normalized SDR range (in-place)
     # This prevents numerical errors from propagating
     # Note: For HDR support, this would need to be removed
-    return torch.clamp(result, -1.0, 1.0)
+    return content_high_freq.clamp_(-1.0, 1.0)
 
 
 def lab_color_transfer(
@@ -255,17 +255,15 @@ def lab_color_transfer(
     """
     Perceptually-accurate color transfer using CIELAB color space.
     
-    LAB provides superior perceptual uniformity compared to RGB/HSV, enabling
-    highly accurate color matching that preserves the original image's appearance.
-    This is the RECOMMENDED method for color correction.
+    Combines wavelet reconstruction (for spatial continuity) with LAB histogram
+    matching (for precise color matching). This eliminates tile artifacts while
+    providing superior color accuracy.
     
     Algorithm:
-    1. Convert both images to LAB color space (D65 illuminant)
-    2. Apply histogram matching to all LAB channels:
-       - L* (luminance): Weighted blend to preserve detail
-       - a* (green-red): Full histogram matching
-       - b* (blue-yellow): Full histogram matching
-    3. Convert back to RGB
+    1. Apply wavelet reconstruction to get artifact-free base
+    2. Convert both wavelet result and style to LAB color space
+    3. Apply histogram matching to LAB channels
+    4. Convert back to RGB
     
     Args:
         content_feat: Target tensor [B, C, H, W] in [-1, 1] with upscaled details
@@ -278,7 +276,10 @@ def lab_color_transfer(
     Returns:
         Color-corrected tensor [B, C, H, W] in [-1, 1]
     """
-    # Handle spatial dimension mismatch
+    # Step 1: Apply wavelet to get artifact-free base with correct spatial structure
+    content_feat = wavelet_reconstruction(content_feat, style_feat, debug=None)
+    
+    # Handle spatial dimension mismatch (should already match after wavelet)
     if content_feat.shape != style_feat.shape:
         debug.log(
             f"LAB: Resizing style {style_feat.shape} to match content {content_feat.shape}",
@@ -291,14 +292,14 @@ def lab_color_transfer(
             align_corners=False
         )
     
-    # Store device
+    # Store device and convert to float32
     device = content_feat.device
     
     # Convert to float32 for accurate color space conversion
     content_feat, original_dtype = ensure_float32_precision(content_feat)
     style_feat, _ = ensure_float32_precision(style_feat)
     
-    # Precompute color space conversion matrices (once per batch, not per pixel)
+    # Precompute color space conversion matrices
     rgb_to_xyz_matrix = torch.tensor([
         [0.4124564, 0.3575761, 0.1804375],
         [0.2126729, 0.7151522, 0.0721750],
@@ -315,17 +316,16 @@ def lab_color_transfer(
     epsilon = 6.0 / 29.0
     kappa = (29.0 / 3.0) ** 3
     
-    # Convert from [-1, 1] to [0, 1] range (in-place where possible)
-    content_rgb = content_feat.add(1.0).mul_(0.5).clamp_(0.0, 1.0)
-    style_rgb = style_feat.add(1.0).mul_(0.5).clamp_(0.0, 1.0)
-    del content_feat, style_feat
+    # Convert from [-1, 1] to [0, 1] range (in-place)
+    content_feat.add_(1.0).mul_(0.5).clamp_(0.0, 1.0)
+    style_feat.add_(1.0).mul_(0.5).clamp_(0.0, 1.0)
     
     # Convert to LAB color space
-    content_lab = _rgb_to_lab_batch(content_rgb, device, rgb_to_xyz_matrix, epsilon, kappa)
-    del content_rgb
+    content_lab = _rgb_to_lab_batch(content_feat, device, rgb_to_xyz_matrix, epsilon, kappa)
+    del content_feat
     
-    style_lab = _rgb_to_lab_batch(style_rgb, device, rgb_to_xyz_matrix, epsilon, kappa)
-    del style_rgb, rgb_to_xyz_matrix
+    style_lab = _rgb_to_lab_batch(style_feat, device, rgb_to_xyz_matrix, epsilon, kappa)
+    del style_feat, rgb_to_xyz_matrix
     
     # Match chrominance channels (a*, b*) for accurate color transfer
     matched_a = _histogram_matching_channel(content_lab[:, 1], style_lab[:, 1], device)

src/utils/constants.py

@@ -4,7 +4,7 @@ Only includes constants actually used in the codebase
 """
 
 # Version information
-__version__ = "2.5.9"
+__version__ = "2.5.19"
 
 import os
 import warnings
@@ -55,14 +55,33 @@ def get_base_cache_dir() -> str:
 
 
 def get_all_model_paths() -> list:
-    """Get all registered model paths including those from extra_model_paths.yaml"""
+    """Get all registered model paths including those from extra_model_paths.yaml (case-insensitive)"""
     try:
         import folder_paths
         # Ensure default path is registered first
         get_base_cache_dir()
-        # Get all paths registered for seedvr2 model type
-        paths = folder_paths.get_folder_paths(SEEDVR2_MODEL_TYPE)
-        return paths if paths else [get_base_cache_dir()]
+        
+        # Case-insensitive lookup: search through all registered folder types
+        # This handles any case variation users might use in extra_model_paths.yaml
+        all_paths = []
+        target_lower = SEEDVR2_MODEL_TYPE.lower()
+        
+        # folder_paths.folder_names_and_paths is the underlying dict: {type: ([paths], extensions)}
+        if hasattr(folder_paths, 'folder_names_and_paths'):
+            for folder_type, (paths, _) in folder_paths.folder_names_and_paths.items():
+                if folder_type.lower() == target_lower:
+                    all_paths.extend(paths)
+        
+        # Remove duplicates while preserving order (os.path.normpath handles Windows/Linux path differences)
+        seen = set()
+        unique_paths = []
+        for path in all_paths:
+            normalized = os.path.normpath(path.lower())
+            if normalized not in seen:
+                seen.add(normalized)
+                unique_paths.append(path)
+        
+        return unique_paths if unique_paths else [get_base_cache_dir()]
     except:
         return [get_base_cache_dir()]
 

src/utils/debug.py

@@ -10,10 +10,35 @@ import torch
 import gc
 from typing import Optional, List, Dict, Any, Union
 from datetime import datetime
-from ..optimization.memory_manager import get_vram_usage, get_basic_vram_info, get_ram_usage, reset_vram_peak
+import platform
+from ..optimization.memory_manager import (
+    get_vram_usage, 
+    get_basic_vram_info, 
+    get_ram_usage, 
+    reset_vram_peak,
+    is_mps_available,
+    is_cuda_available
+)
 from ..utils.constants import __version__
 
 
+def _format_peak_with_overflow(peak_gb: float, total_vram_gb: float) -> str:
+    """Format peak reserved memory, showing overflow breakdown on Windows.
+    
+    Args:
+        peak_gb: Peak reserved memory from PyTorch
+        total_vram_gb: Physical GPU VRAM capacity
+    """
+    if total_vram_gb <= 0:
+        return f"{peak_gb:.2f}GB reserved"
+    
+    overflow_gb = peak_gb - total_vram_gb
+    if overflow_gb <= 0 or platform.system() != 'Windows':
+        return f"{peak_gb:.2f}GB reserved"
+    
+    return f"{peak_gb:.2f}GB reserved ({total_vram_gb:.0f}GB GPU + {overflow_gb:.2f}GB overflow)"
+
+
 class Debug:
     """
     Unified debug logging for generation pipeline and BlockSwap monitoring
@@ -72,7 +97,9 @@ class Debug:
         self.vram_history: List[float] = []
         self.active_timer_stack: List[str] = [] 
         self.timer_namespace: str = ""
-        self.phase_peaks: Dict[str, float] = {}
+        self.phase_vram_peaks_alloc: Dict[str, float] = {}
+        self.phase_vram_peaks_rsv: Dict[str, float] = {}
+        self.phase_ram_peaks: Dict[str, float] = {}
         
     @torch._dynamo.disable  # Skip tracing to avoid datetime.now() warnings
     def log(self, message: str, level: str = "INFO", category: str = "general", force: bool = False, indent_level: int = 0) -> None:
@@ -116,26 +143,99 @@ class Debug:
 
     def print_header(self, cli: bool = False) -> None:
         """Print the header with banner - always displayed"""
-        # Intro logo
-        self.log("", category="none", force=True)
-        self.log(" ╔══════════════════════════════════════════════════════════╗", category="none", force=True)
-        self.log(" ║ ███████ ███████ ███████ ██████  ██    ██ ██████  ███████ ║", category="none", force=True)
-        self.log(" ║ ██      ██      ██      ██   ██ ██    ██ ██   ██      ██ ║", category="none", force=True)
-        self.log(" ║ ███████ █████   █████   ██   ██ ██    ██ ██████  █████   ║", category="none", force=True)
-        self.log(" ║      ██ ██      ██      ██   ██  ██  ██  ██   ██ ██      ║", category="none", force=True)
-        self.log(" ║ ███████ ███████ ███████ ██████    ████   ██   ██ ███████ ║", category="none", force=True)
+        # Temporarily disable timestamps for clean header display
+        original_timestamps = self.show_timestamps
+        self.show_timestamps = False
         
-        # Version number with dynamic padding to maintain visual alignment with any version length
-        version_text = f"v{__version__}"
-        prefix = " 💻 CLI mode · " if cli else " "
-        suffix = "© ByteDance Seed · NumZ · AInVFX "
-        emoji_compensation = 1 if cli else 0
-        padding_width = 59 - len(prefix) - len(version_text) - len(suffix) - 2 - emoji_compensation
-        padding = " " * max(1, padding_width)
-        self.log(f" ║{prefix}{version_text}{padding} {suffix}║", category="none", force=True)
-
-        self.log(" ╚══════════════════════════════════════════════════════════╝", category="none", force=True)
+        # ASCII art logo
         self.log("", category="none", force=True)
+        self.log("", category="none", force=True)
+        self.log("███████╗███████╗███████╗██████╗ ██╗   ██╗██████╗     ██████╗       ███████╗", category="none", force=True, indent_level=1)
+        self.log("██╔════╝██╔════╝██╔════╝██╔══██╗██║   ██║██╔══██╗    ╚════██╗      ██╔════╝", category="none", force=True, indent_level=1)
+        self.log("███████╗█████╗  █████╗  ██║  ██║██║   ██║██████╔╝     █████╔╝      ███████╗", category="none", force=True, indent_level=1)
+        self.log("╚════██║██╔══╝  ██╔══╝  ██║  ██║╚██╗ ██╔╝██╔══██╗    ██╔═══╝       ╚════██║", category="none", force=True, indent_level=1)
+        self.log("███████║███████╗███████╗██████╔╝ ╚████╔╝ ██║  ██║    ███████╗  ██╗ ███████║", category="none", force=True, indent_level=1)
+        self.log("╚══════╝╚══════╝╚══════╝╚═════╝   ╚═══╝  ╚═╝  ╚═╝    ╚══════╝  ╚═╝ ╚══════╝", category="none", force=True, indent_level=1)
+        # Version and credits - left/right aligned to logo width
+        version_text = f"v{__version__}"
+        cli_indicator = "💻 CLI · " if cli else ""
+        left_part = f"{cli_indicator}{version_text}"
+        right_part = "© ByteDance Seed · NumZ · AInVFX"
+        logo_width = 75
+        emoji_compensation = 1 if cli else 0
+        padding = logo_width - len(left_part) - len(right_part) - emoji_compensation
+        self.log(f"{left_part}{' ' * max(1, padding)}{right_part}", category="none", force=True, indent_level=1)
+        self.log("━" * logo_width, category="none", force=True, indent_level=1)
+        self.log("", category="none", force=True)
+        
+        # Restore timestamps setting
+        self.show_timestamps = original_timestamps
+        
+        # Environment info - only in debug mode
+        if self.enabled:
+            self._print_environment_info(cli)
+
+    def _print_environment_info(self, cli: bool = False) -> None:
+        """Print concise environment info for bug reports - zero cost when debug disabled"""
+        import platform
+        import sys
+        
+        # OS
+        os_name = platform.system()
+        if os_name == "Windows":
+            os_str = f"Windows ({platform.version()})"
+        elif os_name == "Darwin":
+            os_str = f"macOS {platform.mac_ver()[0]}"
+        else:
+            try:
+                distro = platform.freedesktop_os_release()
+                os_str = f"{distro.get('NAME', 'Linux')} {distro.get('VERSION_ID', '')}"
+            except (OSError, AttributeError):
+                os_str = f"Linux {platform.release()}"
+        
+        # Python & PyTorch & CUDA
+        py_ver = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
+        torch_ver = torch.__version__
+        cuda_ver = getattr(torch.version, 'cuda', None) or "N/A"
+        
+        # GPU
+        if is_cuda_available():
+            try:
+                props = torch.cuda.get_device_properties(0)
+                gpu_str = f"{props.name} ({round(props.total_memory / (1024**3))}GB)"
+                cudnn_ver = str(torch.backends.cudnn.version()) if torch.backends.cudnn.is_available() else "N/A"
+            except Exception:
+                gpu_str = "CUDA"
+                cudnn_ver = "N/A"
+        elif is_mps_available():
+            gpu_str = "Apple Silicon (MPS)"
+            cudnn_ver = "N/A"
+        else:
+            gpu_str = "CPU"
+            cudnn_ver = "N/A"
+        
+        # Flash Attn & Triton - reuse existing module constants
+        try:
+            from ..optimization.compatibility import FLASH_ATTN_AVAILABLE, TRITON_AVAILABLE
+            flash_str, triton_str = ("✓" if FLASH_ATTN_AVAILABLE else "✗"), ("✓" if TRITON_AVAILABLE else "✗")
+        except ImportError:
+            flash_str = triton_str = "?"
+        
+        # ComfyUI version
+        comfy_str = None
+        if not cli:
+            try:
+                from comfyui_version import __version__ as comfy_ver
+                comfy_str = comfy_ver
+            except ImportError:
+                pass
+        
+        # Print
+        self.log(f"OS: {os_str} | GPU: {gpu_str}", category="info")
+        self.log(f"Python: {py_ver} | PyTorch: {torch_ver} | Flash Attn: {flash_str} | Triton: {triton_str}", category="info")
+        cuda_line = f"CUDA: {cuda_ver} | cuDNN: {cudnn_ver}"
+        self.log(f"{cuda_line} | ComfyUI: {comfy_str}" if comfy_str else cuda_line, category="info")
+        self.log("", category="none")
 
     def print_footer(self) -> None:
         """Print the footer with links - always displayed"""
@@ -306,31 +406,55 @@ class Debug:
         if show_diff and self.memory_checkpoints:
             self._log_memory_diff(current_metrics=memory_info, force=force)
 
-       # Log detailed analysis if requested
+        # Overflow warning (Windows only - WDDM can page to system RAM)
+        overflow = memory_info.get('vram_overflow', 0.0)
+        
+        if overflow > 0 and platform.system() == 'Windows':
+            self.log(f"VRAM overflow: {overflow:.2f}GB paged to system RAM - severe slowdown expected. "
+                     "Consider optimizing (e.g., reduce resolution, batch size, enable BlockSwap, VAE tiling...).",
+                     level="WARNING", category="memory", force=True)
+
+        # Log detailed analysis if requested
         if detailed_tensors and tensor_stats.get('details'):
             self._log_detailed_tensor_analysis(details=tensor_stats['details'], force=force)
                 
         # Store checkpoint with memory limit
         self._store_checkpoint(label, memory_info)
 
-        # Update phase peak if we're in an active phase
-        if self.current_phase and memory_info['vram_peak_since_last'] > 0:
-            self.phase_peaks[self.current_phase] = max(
-                self.phase_peaks.get(self.current_phase, 0),
-                memory_info['vram_peak_since_last']
-            )
+        # Update phase peaks if we're in an active phase
+        if self.current_phase:
+            if memory_info['vram_peak_alloc'] > 0:
+                self.phase_vram_peaks_alloc[self.current_phase] = max(
+                    self.phase_vram_peaks_alloc.get(self.current_phase, 0),
+                    memory_info['vram_peak_alloc']
+                )
+            if memory_info['vram_peak_rsv'] > 0:
+                self.phase_vram_peaks_rsv[self.current_phase] = max(
+                    self.phase_vram_peaks_rsv.get(self.current_phase, 0),
+                    memory_info['vram_peak_rsv']
+                )
+            if memory_info['ram_process'] > 0:
+                self.phase_ram_peaks[self.current_phase] = max(
+                    self.phase_ram_peaks.get(self.current_phase, 0),
+                    memory_info['ram_process']
+                )
 
         # Reset PyTorch's peak memory stats for next interval
         reset_vram_peak(device=None, debug=self)
     
     def _collect_memory_metrics(self) -> Dict[str, Any]:
-        """Collect current memory metrics efficiently."""
+        """Collect current memory metrics."""
+        is_mps = is_mps_available()
+        has_gpu = is_mps or is_cuda_available()
+        
         metrics = {
             'vram_allocated': 0.0,
             'vram_reserved': 0.0,
             'vram_free': 0.0,
             'vram_total': 0.0,
-            'vram_peak_since_last': 0.0,
+            'vram_peak_alloc': 0.0,
+            'vram_peak_rsv': 0.0,
+            'vram_overflow': 0.0,
             'ram_process': 0.0,
             'ram_available': 0.0,
             'ram_total': 0.0,
@@ -339,45 +463,36 @@ class Debug:
             'summary_ram': ""
         }
         
-        # VRAM metrics
-        if torch.cuda.is_available() or (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()):
-            metrics['vram_allocated'], metrics['vram_reserved'], current_global_peak = get_vram_usage(device=None, debug=self)
-
-            # Calculate peak since last log_memory_state
-            # This captures the actual peak that occurred between calls
-            metrics['vram_peak_since_last'] = current_global_peak
-            
+        if has_gpu:
+            metrics['vram_allocated'], metrics['vram_reserved'], metrics['vram_peak_alloc'], metrics['vram_peak_rsv'] = get_vram_usage(device=None, debug=self)
             vram_info = get_basic_vram_info(device=None)
             
-            if "error" not in vram_info:
+            if "error" not in vram_info and vram_info["total_gb"] > 0:
                 metrics['vram_free'] = vram_info["free_gb"]
                 metrics['vram_total'] = vram_info["total_gb"]
+                metrics['vram_overflow'] = max(0.0, metrics['vram_peak_rsv'] - metrics['vram_total'])
                 
-                backend = "MPS" if (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()) else "VRAM"
-                metrics['summary_vram'] = (f"  [{backend}] {metrics['vram_allocated']:.2f}GB allocated / "
-                        f"{metrics['vram_reserved']:.2f}GB reserved / "
-                        f"Peak: {metrics['vram_peak_since_last']:.2f}GB / "
-                        f"{metrics['vram_free']:.2f}GB free / "
-                        f"{metrics['vram_total']:.2f}GB total")
-            else:
-                metrics['summary_vram'] = ""
-        else:
-            metrics['summary_vram'] = ""
+                backend = "Unified Memory" if is_mps else "VRAM"
+                metrics['summary_vram'] = (
+                    f"  [{backend}] {metrics['vram_allocated']:.2f}GB allocated / "
+                    f"{metrics['vram_reserved']:.2f}GB reserved / "
+                    f"Peak: {metrics['vram_peak_alloc']:.2f}GB / "
+                    f"{metrics['vram_free']:.2f}GB free / "
+                    f"{metrics['vram_total']:.2f}GB total"
+                )
+            
+            self.vram_history.append(metrics['vram_reserved'])
         
-        # RAM metrics using new function
+        # RAM metrics
         metrics['ram_process'], metrics['ram_available'], metrics['ram_total'], metrics['ram_others'] = get_ram_usage(debug=self)
         
         if metrics['ram_total'] > 0:
-            metrics['summary_ram'] = (f"  [RAM] {metrics['ram_process']:.2f}GB process / "
-                      f"{metrics['ram_others']:.2f}GB others / "
-                      f"{metrics['ram_available']:.2f}GB free / "
-                      f"{metrics['ram_total']:.2f}GB total")
-        else:
-            metrics['summary_ram'] = ""
-        
-        # Update VRAM history for tracking
-        if torch.cuda.is_available() or (hasattr(torch, 'mps') and callable(getattr(torch.mps, 'is_available', None)) and torch.mps.is_available()):
-            self.vram_history.append(metrics['vram_allocated'])
+            metrics['summary_ram'] = (
+                f"  [RAM] {metrics['ram_process']:.2f}GB process / "
+                f"{metrics['ram_others']:.2f}GB others / "
+                f"{metrics['ram_available']:.2f}GB free / "
+                f"{metrics['ram_total']:.2f}GB total"
+            )
         
         return metrics
     
@@ -504,9 +619,9 @@ class Debug:
         if diffs:
             self.log(f"Memory changes: {', '.join(diffs)}", category="memory", force=force, indent_level=1)
     
-    def log_peak_vram_summary(self, force: bool = True) -> None:
-        """Display peak VRAM usage across all phases"""
-        if not self.phase_peaks:
+    def log_peak_memory_summary(self, force: bool = True) -> None:
+        """Display peak memory usage across all phases."""
+        if not self.phase_vram_peaks_alloc and not self.phase_ram_peaks:
             return
         
         phase_names = {
@@ -516,17 +631,42 @@ class Debug:
             'phase4': 'Post-processing'
         }
         
+        is_mps = is_mps_available()
+        
+        # Get total VRAM for overflow formatting (Windows only)
+        total_vram_gb = 0.0
+        if not is_mps:
+            vram_info = get_basic_vram_info(device=None)
+            if "error" not in vram_info:
+                total_vram_gb = vram_info["total_gb"]
+        
         self.log("", category="none", force=force)
         self.log("────────────────────────", category="none", force=force)
-        self.log("Peak VRAM by Phase:", category="memory", force=force)
+        self.log("Peak memory by phase:", category="memory", force=force)
         
-        for phase_key in sorted(self.phase_peaks.keys()):
+        all_phases = sorted(set(self.phase_vram_peaks_alloc.keys()) | set(self.phase_ram_peaks.keys()))
+        for phase_key in all_phases:
             phase_num = phase_key[-1]
-            self.log(f"  Phase {phase_num}: {phase_names[phase_key]}: {self.phase_peaks[phase_key]:.2f}GB", 
-                    category="memory", force=force)
+            phase_name = phase_names.get(phase_key, phase_key)
+            alloc = self.phase_vram_peaks_alloc.get(phase_key, 0)
+            rsv = self.phase_vram_peaks_rsv.get(phase_key, 0)
+            ram = self.phase_ram_peaks.get(phase_key, 0)
+            
+            if is_mps:
+                self.log(f"{phase_num}. {phase_name}: {alloc:.2f}GB", category="memory", indent_level=1, force=force)
+            else:
+                rsv_str = _format_peak_with_overflow(rsv, total_vram_gb)
+                self.log(f"{phase_num}. {phase_name}: VRAM {alloc:.2f}GB allocated, {rsv_str} | RAM {ram:.2f}GB", category="memory", indent_level=1, force=force)
         
-        overall_max = max(self.phase_peaks.values())
-        self.log(f"Overall Peak VRAM: {overall_max:.2f}GB", category="memory", force=force)
+        overall_alloc = max(self.phase_vram_peaks_alloc.values()) if self.phase_vram_peaks_alloc else 0
+        overall_rsv = max(self.phase_vram_peaks_rsv.values()) if self.phase_vram_peaks_rsv else 0
+        overall_ram = max(self.phase_ram_peaks.values()) if self.phase_ram_peaks else 0
+        
+        if is_mps:
+            self.log(f"Overall peak: {overall_alloc:.2f}GB", category="memory", force=force)
+        else:
+            overall_rsv_str = _format_peak_with_overflow(overall_rsv, total_vram_gb)
+            self.log(f"Overall peak: VRAM {overall_alloc:.2f}GB allocated, {overall_rsv_str} | RAM {overall_ram:.2f}GB", category="memory", force=force)
     
     @torch._dynamo.disable  # Skip tracing to avoid time.time() warnings
     def _store_checkpoint(self, label: str, metrics: Dict[str, Any]) -> None:
@@ -636,5 +776,7 @@ class Debug:
         self.timer_durations.clear()
         self.timer_messages.clear()
         self.active_timer_stack.clear()
-        self.phase_peaks.clear()
+        self.phase_vram_peaks_alloc.clear()
+        self.phase_vram_peaks_rsv.clear()
+        self.phase_ram_peaks.clear()
         self.current_phase = None

src/utils/downloads.py

@@ -14,6 +14,7 @@ import time
 from .model_registry import MODEL_REGISTRY, DEFAULT_VAE
 from .constants import (
     get_base_cache_dir,
+    get_all_model_paths,
     find_model_file,
     get_validation_cache_path,
     HUGGINGFACE_BASE_URL,
@@ -179,20 +180,36 @@ def download_weight(dit_model: str, vae_model: str, model_dir: Optional[str] = N
                          level="WARNING", category="setup")
             continue
         
+        # Get model type for logging
+        model_type = "VAE" if model_info.category == "vae" else "DiT"
+        
         # Check if file exists in any registered path first
         existing_filepath = find_model_file(filename, fallback_dir=cache_dir)
         
         # Use existing file path if it exists, otherwise download to cache_dir
         if os.path.exists(existing_filepath):
             filepath = existing_filepath
+            # Debug log: Model found
+            if debug:
+                debug.log(f"{model_type} model found: {filepath}", category="setup")
         else:
             filepath = os.path.join(cache_dir, filename)
+            # Debug log: Model not found, will need to download
+            if debug:
+                searched_paths = get_all_model_paths()
+                debug.log(f"{model_type} model not found: {filename}", category="setup")
+                debug.log(f"Searched in {len(searched_paths)} location(s):", category="setup")
+                for i, path in enumerate(searched_paths, 1):
+                    debug.log(f"[{i}] {path}", category="setup", indent_level=1)
             
         expected_hash = model_info.sha256
         repo = model_info.repo
         
         ## Quick cache check first
         if is_file_validated_cached(filepath, cache_dir):
+            # Debug log: Model already validated (using cache)
+            if debug:
+                debug.log(f"{model_type} model already validated (cache): {filepath}", category="setup")
             continue
         
         # File exists - validate it
@@ -201,6 +218,9 @@ def download_weight(dit_model: str, vae_model: str, model_dir: Optional[str] = N
                 debug.log(f"Validating {filename}...", category="setup", force=True)
             
             if validate_file(filepath, expected_hash, cache_dir):
+                # Debug log: Model validated successfully
+                if debug:
+                    debug.log(f"{model_type} model validated successfully: {filepath}", category="setup")
                 continue
             else:
                 # File is corrupted