src/optimization/blockswap.py

"""
BlockSwap Module for SeedVR2

This module implements dynamic block swapping between GPU and CPU memory
to enable running large models on limited VRAM systems.

Key Features:
- Dynamic transformer block offloading during inference
- Non-blocking GPU transfers for optimal performance
- RoPE computation fallback to CPU on OOM
- Minimal performance overhead with intelligent caching
- I/O component offloading for maximum memory savings
"""

import time
import types
import torch
import weakref

from typing import Dict, Any, List, Optional
from .memory_manager import clear_memory
from .compatibility import call_rope_with_stability
from ..common.distributed import get_device


def is_blockswap_enabled(config: Optional[Dict[str, Any]]) -> bool:
    """
    Check if BlockSwap configuration indicates BlockSwap should be enabled.
    
    BlockSwap is enabled if either blocks_to_swap > 0 OR swap_io_components is True.
    This is the authoritative function for determining BlockSwap status from configuration.
    
    Args:
        config: BlockSwap configuration dictionary with optional keys:
            - blocks_to_swap: Number of blocks to offload (0 = disabled)
            - swap_io_components: Whether to offload I/O components
        
    Returns:
        True if BlockSwap should be active, False otherwise
    """
    if not config:
        return False
    
    blocks_to_swap = config.get("blocks_to_swap", 0)
    swap_io_components = config.get("swap_io_components", False)
    
    return blocks_to_swap > 0 or swap_io_components


def validate_blockswap_config(
    block_swap_config: Optional[Dict[str, Any]],
    dit_device: 'torch.device',
    dit_offload_device: Optional['torch.device'],
    debug: 'Debug'
) -> Optional[Dict[str, Any]]:
    """
    Validate and potentially modify BlockSwap configuration.
    
    Performs platform-specific validation and configuration adjustment:
    - On macOS (MPS): Auto-disables BlockSwap since unified memory makes it meaningless
    - On other platforms: Validates that offload_device is properly configured
    
    This is the single authoritative validation point for BlockSwap configuration,
    called early in configure_runner() before any model loading.
    
    Args:
        block_swap_config: BlockSwap configuration dictionary (may be None)
        dit_device: Target device for DiT model inference
        dit_offload_device: Device for offloading DiT blocks (may be None)
        debug: Debug instance for logging warnings/errors
        
    Returns:
        Validated/modified block_swap_config (may be None or modified copy)
        
    Raises:
        ValueError: If BlockSwap is enabled but offload_device is invalid (non-MPS only)
    """
    if not is_blockswap_enabled(block_swap_config):
        return block_swap_config
    
    blocks_to_swap = block_swap_config.get("blocks_to_swap", 0)
    swap_io_components = block_swap_config.get("swap_io_components", False)
    
    # Check for macOS unified memory - BlockSwap is meaningless there
    if dit_device.type == "mps":
        debug.log(
            f"BlockSwap disabled: macOS uses unified memory (no separate VRAM/RAM). "
            f"Ignoring blocks_to_swap={blocks_to_swap}, swap_io_components={swap_io_components}",
            level="WARNING", category="blockswap", force=True
        )
        # Return disabled config
        return {
            **block_swap_config,
            "blocks_to_swap": 0,
            "swap_io_components": False
        }
    
    # Validate offload_device is set and different from dit_device
    offload_device_valid = (
        dit_offload_device is not None and 
        str(dit_offload_device) != str(dit_device)
    )
    
    if not offload_device_valid:
        config_details = []
        if blocks_to_swap > 0:
            config_details.append(f"blocks_to_swap={blocks_to_swap}")
        if swap_io_components:
            config_details.append("swap_io_components=True")
        
        offload_str = str(dit_offload_device) if dit_offload_device else "none"
        raise ValueError(
            f"BlockSwap enabled ({', '.join(config_details)}) but dit_offload_device is invalid. "
            f"Current: device='{dit_device}', dit_offload_device='{offload_str}'. "
            f"BlockSwap requires offload_device on the DiT Model to be set and different from device. "
            f"Set --dit_offload_device cpu or disable BlockSwap."
        )
    
    return block_swap_config


# Timing helpers marked to skip torch.compile tracing
# These functions are excluded from Dynamo's graph tracing to avoid warnings
# about non-traceable builtins like time.time(), but they still execute normally
@torch._dynamo.disable
def _get_swap_start_time(debug, enabled: bool) -> Optional[float]:
    """Get start time for swap operation if debug is enabled."""
    return time.time() if debug and enabled else None


@torch._dynamo.disable  
def _log_swap_timing(debug, t_start: Optional[float], component_id, component_type: str) -> None:
    """Log swap timing if start time was captured."""
    if debug and t_start is not None:
        debug.log_swap_time(
            component_id=component_id,
            duration=time.time() - t_start,
            component_type=component_type
        )


def get_module_memory_mb(module: torch.nn.Module) -> float:
    """
    Calculate memory usage of a module in MB.
    
    Args:
        module: PyTorch module to measure
        
    Returns:
        Memory usage in megabytes
    """
    total_bytes = sum(
        param.nelement() * param.element_size() 
        for param in module.parameters() 
        if param.data is not None
    )
    return total_bytes / (1024 * 1024)


def apply_block_swap_to_dit(
    runner: 'VideoDiffusionInfer',
    block_swap_config: Dict[str, Any],
    debug: 'Debug'
) -> None:
    """
    Apply block swapping configuration to a DiT model with OOM protection.
    
    This is the main entry point for configuring block swapping on a model.
    Handles block selection, I/O component offloading, device placement, and
    forward method wrapping for dynamic memory management.
    
    Args:
        runner: VideoDiffusionInfer instance containing the model
        block_swap_config: Configuration dictionary with keys:
            - blocks_to_swap: Number of blocks to swap (from the start)
            - swap_io_components: Whether to offload I/O components  
            - enable_debug: Whether to enable debug logging
            - offload_device: Device to offload to (default: 'cpu')
        debug: Debug instance for logging (required)
    """
    # Early return if BlockSwap not enabled
    if not is_blockswap_enabled(block_swap_config):
        return

    blocks_to_swap = block_swap_config.get("blocks_to_swap", 0)
    swap_io_components = block_swap_config.get("swap_io_components", False)
    
    # Early return only if both block swap and I/O swap are disabled
    if blocks_to_swap <= 0 and not swap_io_components:
        return
    
    if debug is None:
        if hasattr(runner, 'debug') and runner.debug is not None:
            debug = runner.debug
        else:
            raise ValueError("Debug instance must be provided to apply_block_swap_to_dit")
    
    debug.start_timer("apply_blockswap")

    # Get the actual model (handle CompatibleDiT wrapper)
    model = runner.dit
    if hasattr(model, "dit_model"):
        model = model.dit_model
    
    # Determine devices
    if hasattr(runner, '_dit_device'):
        device = runner._dit_device
    else:
        device = get_device()
    offload_device = block_swap_config.get("offload_device", torch.device('cpu'))

    # Validate model structure
    if not hasattr(model, "blocks"):
        debug.log("Model doesn't have 'blocks' attribute for BlockSwap", level="ERROR", category="blockswap", force=True)
        return

    total_blocks = len(model.blocks)
    
    # Clamp blocks_to_swap to available blocks BEFORE logging
    effective_blocks = min(blocks_to_swap, total_blocks) if blocks_to_swap > 0 else 0
    
    # 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:
        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
    
    model.main_device = device
    model.offload_device = offload_device

    # Configure I/O components
    io_config = _configure_io_components(model, device, offload_device, 
                                        swap_io_components, debug)
    memory_stats = _configure_blocks(model, device, offload_device, debug)
    memory_stats['io_components'] = io_config['components']
    memory_stats['io_memory_mb'] = io_config['memory_mb']
    memory_stats['gpu_components'] = io_config['gpu_components']
    memory_stats['io_gpu_memory_mb'] = io_config['gpu_memory_mb']

    # Log memory summary
    _log_memory_summary(memory_stats, offload_device, device, swap_io_components, 
                       debug)

    # Initialize Nunchaku-style async management object
    if blocks_to_swap > 0:
        # normalize device objects
        if isinstance(device, str):
            device = torch.device(device)
        model._swap_stream = torch.cuda.Stream(device=device)
        model._block_ready_events = {}

        # Preload first swapped block to seed pipeline (non-blocking on swap_stream)
        try:
            first_idx = 0
            if first_idx <= model.blocks_to_swap:
                with torch.cuda.stream(model._swap_stream):
                    model.blocks[first_idx].to(device, non_blocking=True)
                    ev = torch.cuda.Event(blocking=False)
                    ev.record(model._swap_stream) # record on swap_stream -> event gets device-bound here
                    model._block_ready_events[first_idx] = ev
        except Exception as e:
            debug.log(f"Failed to initialize swap-stream prefetch: {e}", level="WARNING", category="blockswap", force=True)


    # Wrap block forward methods for dynamic swapping (only if blocks_to_swap > 0)
    if blocks_to_swap > 0:
        for b, block in enumerate(model.blocks):
            if b <= model.blocks_to_swap:
                _wrap_block_forward(block, b, model, debug)

    # Patch RoPE modules for robust error handling
    _patch_rope_for_blockswap(model, debug)

    # Mark BlockSwap as active
    runner._blockswap_active = True

    # Store configuration for debugging and cleanup
    model._block_swap_config = {
        "blocks_swapped": blocks_to_swap,
        "swap_io_components": swap_io_components,
        "total_blocks": total_blocks,
        "offload_device": offload_device,
        "main_device": device,
        "offload_memory": memory_stats['offload_memory'],
        "main_memory": memory_stats['main_memory']
    }

    # Protect model from being moved entirely
    _protect_model_from_move(model, runner, debug)

    debug.log("BlockSwap configuration complete", category="success")
    debug.end_timer("apply_blockswap", "BlockSwap configuration application")
    

def _configure_io_components(
    model: torch.nn.Module,
    device: torch.device,
    offload_device: torch.device,
    swap_io_components: bool,
    debug: 'Debug'
) -> Dict[str, Any]:
    """
    Configure I/O component placement and wrapping with memory tracking.
    
    Handles all non-block modules (embeddings, normalization layers, etc.) by
    either keeping them on GPU or offloading them with dynamic swapping wrappers.
    
    Args:
        model: DiT model containing named children to configure
        device: Main computation device (typically GPU)
        offload_device: Device for offloaded components (typically CPU)
        swap_io_components: If True, offload I/O components with dynamic swapping
        debug: Debug instance for logging (required)
        
    Returns:
        Dictionary containing:
            - components: List of offloaded component names
            - memory_mb: Total memory of offloaded components in MB
            - gpu_components: List of components remaining on GPU
            - gpu_memory_mb: Total memory of GPU components in MB
    """
    io_components_offloaded = []
    io_components_on_gpu = []
    io_memory_mb = 0.0
    io_gpu_memory_mb = 0.0
    
    # Check for pin memory condition
    use_pin_memory = (offload_device == "cpu") if isinstance(offload_device, str) else (offload_device.type == "cpu")

    # Handle I/O modules with dynamic swapping
    for name, module in model.named_children():
        if name != "blocks":
            module_memory = get_module_memory_mb(module)
            
            if swap_io_components:
                module.to(offload_device)
                
                # Enable Pin Memory for I/O components
                if use_pin_memory:
                    for p in module.parameters():
                        if not p.is_pinned():
                            p.data = p.data.pin_memory()
                    for buf in module.buffers():
                        if not buf.is_pinned():
                            buf.data = buf.data.pin_memory()

                _wrap_io_forward(module, name, model, debug)
                io_components_offloaded.append(name)
                io_memory_mb += module_memory
                debug.log(f"{name} → {str(offload_device).upper()} ({module_memory:.2f}MB, dynamic swapping)", category="blockswap", indent_level=1)
            else:
                module.to(device)
                io_components_on_gpu.append(name)
                io_gpu_memory_mb += module_memory
                debug.log(f"{name} → {str(device).upper()} ({module_memory:.2f}MB)", category="blockswap", indent_level=1)

    return {
        'components': io_components_offloaded,
        'memory_mb': io_memory_mb,
        'gpu_components': io_components_on_gpu,
        'gpu_memory_mb': io_gpu_memory_mb
    }


def _configure_blocks(
    model: torch.nn.Module,
    device: torch.device,
    offload_device: torch.device,
    debug: 'Debug'
) -> Dict[str, float]:
    """
    Configure transformer block placement and calculate memory statistics.
    
    Moves blocks to their designated devices based on model.blocks_to_swap
    attribute. Blocks with index <= blocks_to_swap go to offload device,
    others stay on main device.
    
    Args:
        model: DiT model with blocks attribute and blocks_to_swap configured
        device: Main computation device for non-swapped blocks
        offload_device: Device for swapped blocks  
        debug: Debug instance for logging (required)
        
    Returns:
        Dictionary containing:
            - offload_memory: Total memory of offloaded blocks in MB
            - main_memory: Total memory of blocks on main device in MB
            - io_components: Empty list (populated by caller)
    """
    total_offload_memory = 0.0
    total_main_memory = 0.0
    
    # Check if we should pin memory (if offloading to CPU)
    # Nunchaku uses pinned memory for faster async transfers
    use_pin_memory = (offload_device == "cpu") if isinstance(offload_device, str) else (offload_device.type == "cpu")

    # Move blocks based on swap configuration
    for b, block in enumerate(model.blocks):
        block_memory = get_module_memory_mb(block)

        if b > model.blocks_to_swap:
            block.to(device)
            total_main_memory += block_memory
        else:
            block.to(offload_device, non_blocking=False)
            total_offload_memory += block_memory
            
            # Enable Pin Memory optimization for CPU Offload transfer speed
            if use_pin_memory:
                for p in block.parameters():
                    if not p.is_pinned():
                        p.data = p.data.pin_memory()
                for buf in block.buffers():
                    if not buf.is_pinned():
                        buf.data = buf.data.pin_memory()
                        
    # Ensure all buffers match their containing module's device
    for b, block in enumerate(model.blocks):
        target_device = device if b > model.blocks_to_swap else offload_device
        for name, buffer in block.named_buffers():
            if buffer.device != torch.device(target_device):
                # Apply pinning if needed
                if use_pin_memory and target_device.type == "cpu" and not buffer.is_pinned():
                    buffer.data = buffer.data.pin_memory()
                buffer.data = buffer.data.to(target_device, non_blocking=False)

    return {
        "offload_memory": total_offload_memory,
        "main_memory": total_main_memory,
        "io_components": []  # Will be populated by caller
    }


def _log_memory_summary(
    memory_stats: Dict[str, float],
    offload_device: torch.device,
    device: torch.device,
    swap_io_components: bool,
    debug: 'Debug'
) -> None:
    """
    Log comprehensive memory usage summary for BlockSwap configuration.
    
    Displays detailed breakdown of memory distribution across devices,
    including transformer blocks and I/O components.
    
    Args:
        memory_stats: Dictionary containing:
            - offload_memory: Memory offloaded from blocks (MB)
            - main_memory: Memory remaining on main device (MB)
            - io_memory_mb: Memory from offloaded I/O components (MB)
            - io_gpu_memory_mb: Memory from I/O components on GPU (MB)
        offload_device: Device used for offloading
        device: Main computation device
        swap_io_components: Whether I/O components are being swapped
        debug: Debug instance for logging (required)
    """
    debug.log("BlockSwap memory configuration:", category="blockswap")
    
    # Log transformer blocks memory
    blocks_offloaded = memory_stats['offload_memory']
    blocks_on_gpu = memory_stats['main_memory']
    
    offload_str = str(offload_device)
    device_str = str(device)
    
    if blocks_on_gpu == 0:
        debug.log(f"Transformer blocks: {blocks_offloaded:.2f}MB on {offload_str} (dynamic swapping)", category="blockswap", indent_level=1)
    else:
        debug.log(f"Transformer blocks: {blocks_on_gpu:.2f}MB on {device_str}, {blocks_offloaded:.2f}MB on {offload_str}", category="blockswap", indent_level=1)
    
    # Always log I/O components (whether swapping or not)
    io_memory = memory_stats.get('io_memory_mb', 0.0)
    io_gpu_memory = memory_stats.get('io_gpu_memory_mb', 0.0)
    
    if swap_io_components and io_memory > 0:
        io_components = memory_stats.get('io_components', [])
        debug.log(f"I/O components: {io_memory:.2f}MB on {offload_str} (dynamic swapping)", category="blockswap", indent_level=1)
        debug.log(f"{', '.join(io_components)}", category="blockswap", indent_level=2)
    elif io_gpu_memory > 0:
        io_gpu_components = memory_stats.get('gpu_components', [])
        debug.log(f"I/O components: {io_gpu_memory:.2f}MB on {device_str}", category="blockswap", indent_level=1)
        debug.log(f"{', '.join(io_gpu_components)}", category="blockswap", indent_level=2)
    
    # Log total VRAM savings
    total_offloaded = blocks_offloaded + (io_memory if swap_io_components else 0)
    if total_offloaded > 0:
        debug.log(f"Total VRAM saved: {total_offloaded:.2f}MB (~{total_offloaded/1024:.2f}GB)", category="blockswap", indent_level=1)


def _wrap_block_forward(
    block: torch.nn.Module,
    block_idx: int,
    model: torch.nn.Module,
    debug: 'Debug'
) -> None:
    """
    Wrap individual transformer block forward for dynamic device swapping.

    Implements Nunchaku-style pipelining: Prefetch Next -> Compute Current -> Offload Current.
    https://github.com/nunchaku-tech/nunchaku/blob/main/nunchaku/models/utils.py
    
    Creates a wrapped forward method that automatically:
    1. Moves block to GPU before computation
    2. Executes original forward pass
    3. Moves block back to offload device after computation
    4. Logs timing and manages memory pressure
    
    Uses weak references to prevent memory leaks from closure retention.
    
    Args:
        block: Individual transformer block to wrap
        block_idx: Index of this block in model.blocks
        model: Parent DiT model (used for device references)
        debug: Debug instance for logging (required)
    """
    if hasattr(block, '_original_forward'):
        return  # Already wrapped

    # Store original forward method
    original_forward = block.forward
    
    # Create weak references
    model_ref = weakref.ref(model)
    debug_ref = weakref.ref(debug) if debug is not None else (lambda: None)
    
    # Store block_idx on the block itself to avoid closure issues
    block._block_idx = block_idx
    
    def wrapped_forward(self, *args, **kwargs):
        # Retrieve weak references
        model = model_ref()
        debug = debug_ref()
        
        if not model:
            # Model has been garbage collected, fall back to original
            return original_forward(*args, **kwargs)

        # Check if block swap is active for this block
        if hasattr(model, 'blocks_to_swap') and self._block_idx <= model.blocks_to_swap:
            # Use dynamo-disabled helper to get start time (avoids compilation warnings)
            t_start = _get_swap_start_time(debug, debug.enabled if debug else False)
            
            # Only move to GPU if necessary
            current_device = next(self.parameters()).device
            target_device = torch.device(model.main_device)
            
            # 1. Ensure CURRENT block is ready on GPU
            # Check if we have a prefetch event waiting
            if hasattr(model, '_block_ready_events') and self._block_idx in model._block_ready_events:
                # Wait for the swap stream to finish moving this block
                torch.cuda.current_stream().wait_event(model._block_ready_events[self._block_idx])
                # Cleanup event
                del model._block_ready_events[self._block_idx]
            elif current_device != target_device:
                # Fallback: First block or missed prefetch, move synchronously (but non-blocking)
                debug.log(f"[blockswap] Block {self._block_idx} missing prefetch event, moving synchronously", level="WARNING", category="blockswap", force=True)
                self.to(model.main_device, non_blocking=True)
            
            # 2. Trigger Prefetch for NEXT block (Pipelining)
            # Nunchaku logic: Start moving i+1 while i is computing
            next_idx = self._block_idx + 1
            if next_idx <= model.blocks_to_swap:
                next_block = model.blocks[next_idx]
                # Use the dedicated swap stream
                with torch.cuda.stream(model._swap_stream):
                    next_block.to(model.main_device, non_blocking=True)
                    # Record event so next iteration knows when to wait
                    event = torch.cuda.Event(blocking=False)
                    event.record(model._swap_stream)
                    model._block_ready_events[next_idx] = event

            # 3. Execute forward pass (Compute)
            # This runs on the default stream, overlapping with the prefetch above
            output = original_forward(*args, **kwargs)

            # 4. Offload CURRENT block (Async)
            # We record an event on compute stream to ensure we don't move data while it's being used
            compute_done_event = torch.cuda.Event(blocking=False)
            compute_done_event.record(torch.cuda.current_stream())
            
            with torch.cuda.stream(model._swap_stream):
                # Wait for compute to finish before moving memory out
                model._swap_stream.wait_event(compute_done_event)
                # Move back to offload device
                self.to(model.offload_device, non_blocking=True)

            # Use dynamo-disabled helper to log timing (avoids compilation warnings)
            _log_swap_timing(debug, t_start, self._block_idx, "block (pipelined)")

            # Only clear cache under memory pressure
            clear_memory(debug=debug, deep=False, force=False, timer_name="wrap_block_forward")
        else:
            output = original_forward(*args, **kwargs)

        return output
    
    # Bind the wrapped function as a method to the block
    block.forward = types.MethodType(wrapped_forward, block)
    
    # Store reference to original forward for cleanup
    block._original_forward = original_forward


def _wrap_io_forward(
    module: torch.nn.Module,
    module_name: str,
    model: torch.nn.Module,
    debug: 'Debug'
) -> None:
    """
    Wrap I/O component forward for dynamic device swapping.
    
    Similar to _wrap_block_forward but for I/O components (embeddings,
    normalization layers, etc.). Handles swapping between GPU and CPU
    during forward passes.
    
    Uses weak references to prevent circular dependencies and memory leaks.
    
    Args:
        module: I/O component module to wrap
        module_name: Name identifier for logging (e.g., 'x_embedder')
        model: Parent DiT model (used for device references)
        debug: Debug instance for logging (required)
    """
    if hasattr(module, '_is_io_wrapped') and module._is_io_wrapped:
        debug.log(f"Reusing existing I/O wrapper for {module_name}", category="reuse")
        return  # Already wrapped

    # Store original forward method
    original_forward = module.forward
    
    # Create weak references
    model_ref = weakref.ref(model)
    debug_ref = weakref.ref(debug) if debug else lambda: None
    
    # Store module name on the module itself
    module._module_name = module_name
    module._original_forward = original_forward
    
    def wrapped_io_forward(self, *args, **kwargs):
        # Retrieve weak references
        model = model_ref()
        debug = debug_ref()
        
        if not model:
            # Model has been garbage collected, fall back to original
            return self._original_forward(*args, **kwargs)

        # Use dynamo-disabled helper to get start time (avoids compilation warnings)
        t_start = _get_swap_start_time(debug, debug.enabled if debug else False)
        
        # Check current device to avoid unnecessary moves
        current_device = next(self.parameters()).device
        target_device = torch.device(model.main_device)
        
        # Move to GPU for computation if needed
        if current_device != target_device:
            self.to(model.main_device, non_blocking=False)

        # Execute forward pass
        output = self._original_forward(*args, **kwargs)

        # Move back to offload device
        self.to(model.offload_device, non_blocking=False)
        
        # Use dynamo-disabled helper to log timing (avoids compilation warnings)
        _log_swap_timing(debug, t_start, self._module_name, "I/O")

        # Only clear cache under memory pressure
        clear_memory(debug=debug, deep=False, force=False, timer_name="wrap_block_forward")

        return output
    
    # Bind as a method
    module.forward = types.MethodType(wrapped_io_forward, module)
    module._is_io_wrapped = True
    
    # Store module reference for restoration
    if not hasattr(model, '_io_swappers'):
        model._io_swappers = []
    model._io_swappers.append((module, module_name))


def _patch_rope_for_blockswap(
    model: torch.nn.Module,
    debug: 'Debug'
) -> None:
    """
    Patch RoPE (Rotary Position Embedding) modules for device-aware fallback.
    
    Adds CPU fallback logic to RoPE modules to handle device mismatch errors
    that can occur during BlockSwap operations. Complements the stability
    wrapper from compatibility.py with device-specific error handling.
    
    Args:
        model: DiT model containing RoPE modules to patch
        debug: Debug instance for logging (required)
    """
    rope_patches = []
    
    for name, module in model.named_modules():
        if "rope" in name.lower() and hasattr(module, "get_axial_freqs"):
            # Skip if already wrapped by blockswap
            if hasattr(module, '_blockswap_wrapped') and module._blockswap_wrapped:
                continue
            
            # Get current method (might be stability-wrapped)
            current_method = module.get_axial_freqs
            
            # Create device-aware wrapper with proper closure handling
            def make_device_aware_wrapper(module_name, current_fn):
                def device_aware_rope_wrapper(self, *args, **kwargs):
                    try:
                        # Try current method (original or stability-wrapped)
                        return current_fn(*args, **kwargs)
                    except (RuntimeError, torch.cuda.OutOfMemoryError) as e:
                        error_msg = str(e).lower()
                        # Only handle device/memory specific errors
                        if any(x in error_msg for x in ["device", "memory", "allocation"]):
                            debug.log(f"RoPE OOM for {module_name}", level="WARNING", category="rope", force=True)
                            debug.log(f"Clearing RoPE cache and retrying", category="info", force=True)
                            
                            # Get current device from parameters
                            try:
                                current_device = next(self.parameters()).device
                            except StopIteration:
                                # Fallback: use model's main_device if BlockSwap has set it, else use offload_device
                                if hasattr(model, 'main_device'):
                                    current_device = torch.device(model.main_device)
                                elif hasattr(model, 'offload_device'):
                                    current_device = torch.device(model.offload_device)
                            
                            # Try clearing cache first (non-invasive fix)
                            if hasattr(current_fn, 'cache_clear'):
                                current_fn.cache_clear()
                                try:
                                    # Retry on same device after clearing cache
                                    return current_fn(*args, **kwargs)
                                except Exception as retry_error:
                                    # Cache clear wasn't enough, need more drastic measures
                                    debug.log(f"Cache clear insufficient for {module_name}, falling back to CPU", level="WARNING", category="rope", force=True)
                            
                            # Fallback to CPU computation with stability
                            self.cpu()
                            
                            try:
                                # Use call_rope_with_stability for CPU computation
                                # This ensures cache is cleared and autocast disabled
                                original_fn = getattr(self, '_original_get_axial_freqs', current_fn)
                                result = call_rope_with_stability(original_fn, *args, **kwargs)
                                
                                # Move module back to original device
                                self.to(current_device)
                                
                                # Move result to appropriate device if it's a tensor
                                if hasattr(result, 'to'):
                                    target_device = args[0].device if len(args) > 0 and hasattr(args[0], 'device') else current_device
                                    return result.to(target_device)
                                return result
                                
                            except Exception as cpu_error:
                                # Always restore device even on error
                                self.to(current_device)
                                raise cpu_error
                        else:
                            # Not a device error, let it bubble up
                            raise
                
                return device_aware_rope_wrapper
            
            # Apply wrapper
            module.get_axial_freqs = types.MethodType(
                make_device_aware_wrapper(name, current_method), 
                module
            )
            module._blockswap_wrapped = True
            
            # Store for cleanup (use original or previously stored)
            original_method = getattr(module, '_original_get_axial_freqs', current_method)
            rope_patches.append((module, original_method))
    
    if rope_patches:
        model._rope_patches = rope_patches
        debug.log(f"Patched {len(rope_patches)} RoPE modules with device handling", category="success")


def _protect_model_from_move(
    model: torch.nn.Module,
    runner: 'VideoDiffusionInfer',
    debug: 'Debug'
) -> None:
    """
    Protect model from unintended full device movement during BlockSwap.
    
    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 model flag.
    
    Args:
        model: DiT model to protect
        runner: VideoDiffusionInfer instance (for active status check)
        debug: Debug instance for logging (required)
    """
    if not hasattr(model, '_original_to'):
        # Store runner reference as weak reference to avoid circular refs
        model._blockswap_runner_ref = weakref.ref(runner)
        model._original_to = model.to
        
        # 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 hasattr(runner_obj, "_blockswap_active"):
                    blockswap_is_active = runner_obj._blockswap_active
            
            # 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
                
                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
            
            # Allow movement (either bypass is enabled or target is offload device)
            if hasattr(self, '_original_to'):
                return self._original_to(device, *args, **kwargs)
            else:
                # Fallback - shouldn't happen
                return super(type(self), self).to(device, *args, **kwargs)
        
        # Bind as a method to the model instance
        model.to = types.MethodType(protected_model_to, model)


def set_blockswap_bypass(runner, bypass: bool, debug):
    """
    Set or unset bypass flag for BlockSwap protection.
    Used for offloading to temporarily allow model movement.
    
    Args:
        runner: Runner instance with BlockSwap
        bypass: True to bypass protection, False to enforce it
        debug: Debug instance for logging
    """
    if not hasattr(runner, "_blockswap_active") or not runner._blockswap_active:
        return
    
    # Get the actual model (handle CompatibleDiT 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")
    else:
        debug.log("BlockSwap protection renabled to avoid accidentally offloading the entire DiT model", category="success")


def cleanup_blockswap(runner, keep_state_for_cache=False):
    """
    Clean up BlockSwap configuration based on caching mode.
    
    When caching (keep_state_for_cache=True):
    - Keep all BlockSwap configuration intact
    - Only mark as inactive for safety during non-inference operations
    
    When not caching (keep_state_for_cache=False):
    - Full cleanup of all BlockSwap state
    
    Args:
        runner: VideoDiffusionInfer instance to clean up
        keep_state_for_cache: If True, preserve BlockSwap state for reuse
    """
    # Get debug instance from runner
    if not hasattr(runner, 'debug') or runner.debug is None:
        raise ValueError("Debug instance must be available on runner for cleanup_blockswap")
    
    debug = runner.debug
    
    # Get the actual model (handle CompatibleDiT 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(model, "_block_swap_config") or
        hasattr(model, "_blockswap_bypass_protection")
    )
    
    if not has_blockswap_state:
        return

    debug.log("Starting BlockSwap cleanup", category="cleanup")

    if keep_state_for_cache:
        # 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(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")
        return

    # Full cleanup when not caching
    # Get the actual model (handle CompatibleDiT wrapper)
    model = runner.dit
    if hasattr(model, "dit_model"):
        model = model.dit_model

    # 1. Restore block forward methods
    if hasattr(model, 'blocks'):
        restored_count = 0
        for block in model.blocks:
            if hasattr(block, '_original_forward'):
                block.forward = block._original_forward
                delattr(block, '_original_forward')
                restored_count += 1
                
                # Clean up wrapper attributes
                for attr in ['_block_idx', '_model_ref', '_debug_ref', '_blockswap_wrapped']:
                    if hasattr(block, attr):
                        delattr(block, attr)
        
        if restored_count > 0:
            debug.log(f"Restored {restored_count} block forward methods", category="success")

    # 2. Restore RoPE patches
    if hasattr(model, '_rope_patches'):
        for module, original_method in model._rope_patches:
            module.get_axial_freqs = original_method
            # Clean up wrapper attributes
            for attr in ['_rope_wrapped', '_original_get_axial_freqs']:
                if hasattr(module, attr):
                    delattr(module, attr)
        debug.log(f"Restored {len(model._rope_patches)} RoPE methods", category="success")
        delattr(model, '_rope_patches')

    # 3. Restore I/O component forward methods and move to offload device
    if hasattr(model, '_io_swappers'):
        for module, module_name in model._io_swappers:
            if hasattr(module, '_original_forward'):
                module.forward = module._original_forward
                # Clean up wrapper attributes
                for attr in ['_original_forward', '_model_ref', '_debug_ref', 
                           '_module_name', '_is_io_wrapped']:
                    if hasattr(module, attr):
                        delattr(module, attr)
        debug.log(f"Restored {len(model._io_swappers)} I/O components", category="success")
        delattr(model, '_io_swappers')
    
    # Move all IO components to offload device during full cleanup
    if hasattr(model, 'offload_device'):
        offload_device = model.offload_device
        moved_count = 0
        for name, module in model.named_children():
            if name != "blocks":
                module.to(offload_device)
                moved_count += 1
        if moved_count > 0:
            debug.log(f"Moved {moved_count} IO components to offload device", category="success")

    # 4. Restore original .to() method
    if hasattr(model, '_original_to'):
        model.to = model._original_to
        delattr(model, '_original_to')
        debug.log("Restored original .to() method", category="success")

    # 5. Clean up BlockSwap-specific attributes
    for attr in ['_blockswap_runner_ref', 'blocks_to_swap', 'main_device', 
                 'offload_device']:
        if hasattr(model, attr):
            delattr(model, attr)

    # 6. Clean up runner attributes
    runner._blockswap_active = False
    
    # Clean up pipelining resources on model (synchronize first)
    if hasattr(model, '_swap_stream'):
        try:
            model._swap_stream.synchronize()
        except Exception:
            pass
    for attr in ['_swap_stream', '_block_ready_events']:
        if hasattr(model, attr):
            delattr(model, attr)

    # Remove all config attributes
    for attr in ['_cached_blockswap_config', '_block_swap_config', '_blockswap_debug']:
        if hasattr(runner, attr):
            delattr(runner, attr)
    
    debug.log("BlockSwap cleanup complete", category="success")