custom_nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/compatibility.py

"""
Compatibility module for SeedVR2
Contains FP8/FP16 compatibility layers and wrappers for different model architectures

Extracted from: seedvr2.py (lines 1045-1630)
"""

# Compatibility shims - Must run before any torch/diffusers import
import sys
import types
import importlib.machinery


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__
        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 with proper __spec__
        class _FailingStub(types.ModuleType):
            """Stub that lets xformers gracefully disable its flash backend."""
            def __getattr__(self, name):
                raise ImportError("_C_flashattention unavailable")
        
        stub = _FailingStub('xformers._C_flashattention')
        stub.__spec__ = importlib.machinery.ModuleSpec('xformers._C_flashattention', None)
        stub.__file__ = None
        stub.__path__ = []
        stub.__loader__ = None
        sys.modules['xformers._C_flashattention'] = stub


def ensure_bitsandbytes_safe():
    """
    Pre-test bitsandbytes; stub if broken to prevent import conflicts.
    
    On some systems (e.g., ROCm without proper binaries), bitsandbytes registers
    PyTorch kernels during import then fails. If another node already triggered 
    this partial load, re-importing causes kernel registration conflicts.
    
    This shim catches such failures and stubs the module so diffusers can load
    gracefully without bitsandbytes quantization support.
    """
    if 'bitsandbytes' in sys.modules:
        return  # Already loaded or stubbed
    
    try:
        import bitsandbytes
        # Success - bitsandbytes works, other nodes can use it
    except (ImportError, OSError, RuntimeError, ValueError):
        # Installation broken, not present, or version detection failed - create stub
        stub = types.ModuleType('bitsandbytes')
        stub.__spec__ = importlib.machinery.ModuleSpec('bitsandbytes', None)
        stub.__file__ = None
        stub.__path__ = []
        stub.__version__ = "0.0.0"
        sys.modules['bitsandbytes'] = stub


# Run all shims immediately on import, before torch/diffusers
ensure_triton_compat()
ensure_flash_attn_safe()
ensure_xformers_flash_compat()
ensure_bitsandbytes_safe()


import torch
import os


# Flash/Sage Attention & Triton Compatibility Layer

# 1. Flash Attention 3 (Hopper+, faster, no dropout/window support)
flash_attn_3_varlen_func = None
FLASH_ATTN_3_AVAILABLE = False
try:
    import flash_attn_interface
    flash_attn_3_varlen_func = flash_attn_interface.flash_attn_varlen_func
    FLASH_ATTN_3_AVAILABLE = True
except (ImportError, AttributeError, OSError):
    pass

# 2. Flash Attention 2 (wider compatibility, supports dropout/window)
flash_attn_2_varlen_func = None
FLASH_ATTN_2_AVAILABLE = False
try:
    from flash_attn import flash_attn_varlen_func as _fa2_varlen
    import flash_attn_2_cuda  # noqa: F401
    flash_attn_2_varlen_func = _fa2_varlen
    FLASH_ATTN_2_AVAILABLE = True
except (ImportError, AttributeError, OSError):
    pass

FLASH_ATTN_AVAILABLE = FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE

# 3. SageAttention 2 (varlen support)
sageattn_varlen = None
SAGE_ATTN_2_AVAILABLE = False
try:
    from sageattention import sageattn_varlen as _sa2_varlen
    sageattn_varlen = _sa2_varlen
    SAGE_ATTN_2_AVAILABLE = True
except (ImportError, AttributeError, OSError):
    pass

# 4. SageAttention 3 / Blackwell (RTX 50xx only, batched attention)
sageattn_blackwell = None
SAGE_ATTN_3_AVAILABLE = False
try:
    from sageattn3 import sageattn3_blackwell as _sa3_blackwell
    sageattn_blackwell = _sa3_blackwell
    SAGE_ATTN_3_AVAILABLE = True
except (ImportError, AttributeError, OSError):
    try:
        from sageattention import sageattn_blackwell as _sa3_blackwell
        sageattn_blackwell = _sa3_blackwell
        SAGE_ATTN_3_AVAILABLE = True
    except (ImportError, AttributeError, OSError):
        pass

SAGE_ATTN_AVAILABLE = SAGE_ATTN_2_AVAILABLE or SAGE_ATTN_3_AVAILABLE


def validate_attention_mode(requested_mode: str, debug=None) -> str:
    """
    Validate attention mode availability with automatic fallback.
    
    Args:
        requested_mode: 'sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3'
        debug: Optional debug instance for logging
        
    Returns:
        Validated mode that is available
    """
    # Flash Attention 3
    if requested_mode == 'flash_attn_3':
        if FLASH_ATTN_3_AVAILABLE:
            return requested_mode
        if FLASH_ATTN_2_AVAILABLE:
            if debug:
                debug.log(
                    "Flash Attention 3 not available (requires Hopper+ GPU and flash-attn with FA3 support).\n"
                    "Falling back to Flash Attention 2.",
                    level="WARNING", category="setup", force=True
                )
            return 'flash_attn_2'
        error_msg = (
            "Cannot use 'flash_attn_3' attention mode: Flash Attention is not installed.\n"
            "\n"
            "Flash Attention 3 provides maximum speedup on Hopper+ GPUs through optimized CUDA kernels.\n"
            "Falling back to PyTorch SDPA (scaled dot-product attention).\n"
            "\n"
            "To fix this issue:\n"
            "  1. Install Flash Attention: pip install flash-attn\n"
            "  2. OR change attention_mode to 'sdpa' (default, always available)\n"
            "\n"
            "For more info: https://github.com/Dao-AILab/flash-attention"
        )
        if debug:
            debug.log(error_msg, level="WARNING", category="setup", force=True)
        return 'sdpa'
    
    # Flash Attention 2
    if requested_mode == 'flash_attn_2':
        if FLASH_ATTN_2_AVAILABLE:
            return requested_mode
        error_msg = (
            "Cannot use 'flash_attn_2' attention mode: Flash Attention 2 is not installed.\n"
            "\n"
            "Flash Attention 2 provides speedup on Ampere+ GPUs through optimized CUDA kernels.\n"
            "Falling back to PyTorch SDPA (scaled dot-product attention).\n"
            "\n"
            "To fix this issue:\n"
            "  1. Install Flash Attention: pip install flash-attn\n"
            "  2. OR change attention_mode to 'sdpa' (default, always available)\n"
            "\n"
            "For more info: https://github.com/Dao-AILab/flash-attention"
        )
        if debug:
            debug.log(error_msg, level="WARNING", category="setup", force=True)
        return 'sdpa'
    
    # SageAttention 3 (Blackwell)
    if requested_mode == 'sageattn_3':
        if SAGE_ATTN_3_AVAILABLE:
            return requested_mode
        if SAGE_ATTN_2_AVAILABLE:
            if debug:
                debug.log(
                    "SageAttention 3 (Blackwell) not available (requires RTX 50xx GPU and sageattn3 package).\n"
                    "Falling back to SageAttention 2.",
                    level="WARNING", category="setup", force=True
                )
            return 'sageattn_2'
        error_msg = (
            "Cannot use 'sageattn_3' attention mode: SageAttention is not installed.\n"
            "\n"
            "SageAttention 3 provides maximum speedup on Blackwell (RTX 50xx) GPUs.\n"
            "Falling back to PyTorch SDPA (scaled dot-product attention).\n"
            "\n"
            "To fix this issue:\n"
            "  1. Install SageAttention: pip install sageattention\n"
            "  2. For SA3 Blackwell support: pip install sageattn3\n"
            "  3. OR change attention_mode to 'flash_attn_2' or 'sdpa'\n"
            "\n"
            "For more info: https://github.com/thu-ml/SageAttention"
        )
        if debug:
            debug.log(error_msg, level="WARNING", category="setup", force=True)
        return 'sdpa'
    
    # SageAttention 2
    if requested_mode == 'sageattn_2':
        if SAGE_ATTN_2_AVAILABLE:
            return requested_mode
        error_msg = (
            "Cannot use 'sageattn_2' attention mode: SageAttention is not installed.\n"
            "\n"
            "SageAttention provides speedup on NVIDIA GPUs through optimized CUDA kernels.\n"
            "Falling back to PyTorch SDPA (scaled dot-product attention).\n"
            "\n"
            "To fix this issue:\n"
            "  1. Install SageAttention: pip install sageattention\n"
            "  2. OR change attention_mode to 'flash_attn_2' or 'sdpa'\n"
            "\n"
            "For more info: https://github.com/thu-ml/SageAttention"
        )
        if debug:
            debug.log(error_msg, level="WARNING", category="setup", force=True)
        return 'sdpa'
    
    return requested_mode


@torch._dynamo.disable
def call_flash_attn_2_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
    """
    Wrapper for Flash Attention 2 flash_attn_varlen_func that handles tensor-to-scalar conversion.
    
    Flash Attention 2 supports dropout_p and window_size parameters.
    Works on Ampere+ GPUs (RTX 30xx, 40xx, A100, etc.).
    
    This function is excluded from torch.compile because:
    1. flash_attn is a C++ extension that can't be compiled anyway
    2. It requires Python int scalars for max_seqlen parameters
    3. Disabling compilation here keeps the rest of the model compilable
    
    Args:
        q: Query tensor (total_seq, heads, head_dim)
        k: Key tensor (total_seq, heads, head_dim)
        v: Value tensor (total_seq, heads, head_dim)
        cu_seqlens_q: Cumulative sequence lengths for queries
        cu_seqlens_k: Cumulative sequence lengths for keys
        max_seqlen_q: Maximum query sequence length (can be tensor or int)
        max_seqlen_k: Maximum key sequence length (can be tensor or int)
        **kwargs: Additional arguments (dropout_p, softmax_scale, causal, window_size, deterministic)
        
    Returns:
        Attention output tensor (total_seq, heads, head_dim)
    """
    if not FLASH_ATTN_2_AVAILABLE:
        raise ImportError("Flash Attention 2 is not available")
    
    # Convert tensor max_seqlen to Python int if needed
    if torch.is_tensor(max_seqlen_q):
        max_seqlen_q = int(max_seqlen_q.item())
    if torch.is_tensor(max_seqlen_k):
        max_seqlen_k = int(max_seqlen_k.item())
    
    return flash_attn_2_varlen_func(
        q=q,
        k=k,
        v=v,
        cu_seqlens_q=cu_seqlens_q,
        cu_seqlens_k=cu_seqlens_k,
        max_seqlen_q=max_seqlen_q,
        max_seqlen_k=max_seqlen_k,
        **kwargs
    )


@torch._dynamo.disable
def call_flash_attn_3_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
    """
    Wrapper for Flash Attention 3 flash_attn_varlen_func that handles tensor-to-scalar conversion.
    
    Flash Attention 3 is faster than FA2 but does NOT support dropout_p and window_size.
    Works on Hopper+ GPUs (H100, etc.) - requires flash_attn_interface package.
    
    This function is excluded from torch.compile because:
    1. flash_attn is a C++ extension that can't be compiled anyway
    2. It requires Python int scalars for max_seqlen parameters
    3. Disabling compilation here keeps the rest of the model compilable
    
    Args:
        q: Query tensor (total_seq, heads, head_dim)
        k: Key tensor (total_seq, heads, head_dim)
        v: Value tensor (total_seq, heads, head_dim)
        cu_seqlens_q: Cumulative sequence lengths for queries
        cu_seqlens_k: Cumulative sequence lengths for keys
        max_seqlen_q: Maximum query sequence length (can be tensor or int)
        max_seqlen_k: Maximum key sequence length (can be tensor or int)
        **kwargs: Additional arguments (softmax_scale, causal, deterministic)
                  Note: dropout_p and window_size are ignored (not supported by FA3)
        
    Returns:
        Attention output tensor (total_seq, heads, head_dim)
    """
    if not FLASH_ATTN_3_AVAILABLE:
        raise ImportError("Flash Attention 3 is not available")
    
    # Convert tensor max_seqlen to Python int if needed
    if torch.is_tensor(max_seqlen_q):
        max_seqlen_q = int(max_seqlen_q.item())
    if torch.is_tensor(max_seqlen_k):
        max_seqlen_k = int(max_seqlen_k.item())
    
    # FA3 doesn't support dropout_p and window_size - filter them out
    fa3_kwargs = {key: val for key, val in kwargs.items() if key not in ('dropout_p', 'window_size')}
    
    # FA3 returns a tuple (output, softmax_lse), we only need output
    return flash_attn_3_varlen_func(
        q=q,
        k=k,
        v=v,
        cu_seqlens_q=cu_seqlens_q,
        cu_seqlens_k=cu_seqlens_k,
        max_seqlen_q=max_seqlen_q,
        max_seqlen_k=max_seqlen_k,
        seqused_q=None,
        seqused_k=None,
        **fa3_kwargs
    )[0]


@torch._dynamo.disable
def call_sage_attn_2_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
    """
    Wrapper for SageAttention 2 sageattn_varlen that handles tensor-to-scalar conversion.
    
    SageAttention 2 provides optimized attention for NVIDIA GPUs with native varlen support.
    Works on most modern NVIDIA GPUs.
    
    This function is excluded from torch.compile because:
    1. SageAttention is a C++ extension that can't be compiled anyway
    2. It requires Python int scalars for max_seqlen parameters
    3. Disabling compilation here keeps the rest of the model compilable
    
    Args:
        q: Query tensor (total_seq, heads, head_dim)
        k: Key tensor (total_seq, heads, head_dim)
        v: Value tensor (total_seq, heads, head_dim)
        cu_seqlens_q: Cumulative sequence lengths for queries
        cu_seqlens_k: Cumulative sequence lengths for keys
        max_seqlen_q: Maximum query sequence length (can be tensor or int)
        max_seqlen_k: Maximum key sequence length (can be tensor or int)
        **kwargs: Additional arguments (causal supported, others ignored)
        
    Returns:
        Attention output tensor (total_seq, heads, head_dim)
    """
    if not SAGE_ATTN_2_AVAILABLE:
        raise ImportError("SageAttention 2 is not available")
    
    # Convert tensor max_seqlen to Python int if needed
    if torch.is_tensor(max_seqlen_q):
        max_seqlen_q = int(max_seqlen_q.item())
    if torch.is_tensor(max_seqlen_k):
        max_seqlen_k = int(max_seqlen_k.item())
    
    # SageAttention requires half precision (fp16/bf16)
    out_dtype = q.dtype
    half_dtypes = (torch.float16, torch.bfloat16)
    
    if not (q.dtype == k.dtype == v.dtype):
        k = k.to(q.dtype)
        v = v.to(q.dtype)
    
    if q.dtype not in half_dtypes:
        q = q.to(torch.bfloat16)
        k = k.to(torch.bfloat16)
        v = v.to(torch.bfloat16)
    
    is_causal = kwargs.get('causal', False)
    sm_scale = 1.0 / (q.shape[-1] ** 0.5)
    
    out = sageattn_varlen(
        q, k, v,
        cu_seqlens_q, cu_seqlens_k,
        max_seqlen_q, max_seqlen_k,
        is_causal, sm_scale
    )
    
    return out.to(out_dtype) if out.dtype != out_dtype else out


@torch._dynamo.disable
def call_sage_attn_3_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
    """
    Wrapper for SageAttention 3 (Blackwell) that converts varlen format to batched format.
    
    SageAttention 3 / Blackwell provides maximum performance on RTX 50xx series GPUs.
    However, it only supports batched attention (uniform sequence lengths), not varlen.
    
    This wrapper detects uniform-length batches and reshapes accordingly.
    For variable-length sequences, it automatically falls back to SageAttention 2.
    
    This function is excluded from torch.compile because:
    1. SageAttention is a C++ extension that can't be compiled anyway
    2. It requires Python int scalars for max_seqlen parameters
    3. The varlen-to-batched conversion involves dynamic shapes
    4. Disabling compilation here keeps the rest of the model compilable
    
    Args:
        q: Query tensor (total_seq, heads, head_dim)
        k: Key tensor (total_seq, heads, head_dim)
        v: Value tensor (total_seq, heads, head_dim)
        cu_seqlens_q: Cumulative sequence lengths for queries
        cu_seqlens_k: Cumulative sequence lengths for keys
        max_seqlen_q: Maximum query sequence length (can be tensor or int)
        max_seqlen_k: Maximum key sequence length (can be tensor or int)
        **kwargs: Additional arguments (passed to SA2 fallback if needed)
        
    Returns:
        Attention output tensor (total_seq, heads, head_dim)
    """
    if not SAGE_ATTN_3_AVAILABLE:
        raise ImportError("SageAttention 3 (Blackwell) is not available")
    
    # Convert tensor max_seqlen to Python int if needed
    if torch.is_tensor(max_seqlen_q):
        max_seqlen_q = int(max_seqlen_q.item())
    if torch.is_tensor(max_seqlen_k):
        max_seqlen_k = int(max_seqlen_k.item())
    
    # Check if all sequences have uniform length (required for SA3 batched API)
    # SA3/Blackwell uses batched attention, not varlen, so we need uniform lengths
    seq_lens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1]
    seq_lens_k = cu_seqlens_k[1:] - cu_seqlens_k[:-1]
    
    uniform_q = (seq_lens_q == seq_lens_q[0]).all()
    uniform_k = (seq_lens_k == seq_lens_k[0]).all()
    
    if not (uniform_q and uniform_k):
        # Fall back to SA2 for variable-length sequences
        # This is expected behavior - SA3 Blackwell doesn't support varlen natively
        if SAGE_ATTN_2_AVAILABLE:
            return call_sage_attn_2_varlen(
                q, k, v, cu_seqlens_q, cu_seqlens_k,
                max_seqlen_q, max_seqlen_k, **kwargs
            )
        raise RuntimeError(
            "SageAttention 3 (Blackwell) requires uniform sequence lengths, "
            "and SageAttention 2 is not available as fallback. "
            "Please install sageattention package or use flash_attn/sdpa instead."
        )
    
    # Extract batch dimensions
    batch_size = len(cu_seqlens_q) - 1
    seq_len_q = int(seq_lens_q[0].item())
    seq_len_k = int(seq_lens_k[0].item())
    heads = q.shape[1]
    dim = q.shape[2]
    
    # SageAttention requires half precision (fp16/bf16)
    out_dtype = q.dtype
    half_dtypes = (torch.float16, torch.bfloat16)
    
    if not (q.dtype == k.dtype == v.dtype):
        k = k.to(q.dtype)
        v = v.to(q.dtype)
    
    if q.dtype not in half_dtypes:
        q = q.to(torch.bfloat16)
        k = k.to(torch.bfloat16)
        v = v.to(torch.bfloat16)
    
    # Reshape varlen (total_seq, heads, dim) -> batched (batch, seq, heads, dim)
    q_batched = q.view(batch_size, seq_len_q, heads, dim)
    k_batched = k.view(batch_size, seq_len_k, heads, dim)
    v_batched = v.view(batch_size, seq_len_k, heads, dim)
    
    # SA3/Blackwell expects (batch, heads, seq, dim) layout
    q_batched = q_batched.transpose(1, 2)  # (batch, heads, seq, dim)
    k_batched = k_batched.transpose(1, 2)
    v_batched = v_batched.transpose(1, 2)
    
    # Call SA3 Blackwell
    out = sageattn_blackwell(q_batched, k_batched, v_batched, per_block_mean=False)
    
    # Reshape back to varlen format (total_seq, heads, dim)
    out = out.transpose(1, 2).reshape(-1, heads, dim).contiguous()
    
    return out.to(out_dtype) if out.dtype != out_dtype else out


# 2. Triton - Required for torch.compile with inductor backend
try:
    import triton
    TRITON_AVAILABLE = True
except ImportError:
    TRITON_AVAILABLE = False


# 3. GGUF - Required for quantized model loading
try:
    import gguf
    from gguf import GGMLQuantizationType
    GGUF_AVAILABLE = True
except ImportError:
    GGUF_AVAILABLE = False
    gguf = None
    GGMLQuantizationType = None


def validate_gguf_availability(operation: str = "load GGUF model", debug=None) -> None:
    """
    Validate GGUF availability and raise error if not installed.
    
    Args:
        operation: Description of the operation requiring GGUF
        debug: Optional debug instance for logging
        
    Raises:
        RuntimeError: If GGUF is not available
    """
    if not GGUF_AVAILABLE:
        error_msg = (
            f"Cannot {operation}: GGUF library is not installed.\n"
            f"\n"
            f"GGUF provides quantized model support for memory-efficient loading.\n"
            f"\n"
            f"To fix this issue:\n"
            f"  1. Install GGUF: pip install gguf\n"
            f"  2. OR use a non-quantized model format (.safetensors)\n"
            f"\n"
            f"For more info: https://github.com/ggerganov/ggml"
        )
        if debug:
            debug.log(error_msg, level="ERROR", category="setup", force=True)
        raise RuntimeError(f"GGUF library required to {operation}")


# 4. NVIDIA Conv3d Memory Bug - Workaround for PyTorch >= 2.9 + cuDNN >= 91002
def _check_conv3d_memory_bug():
    """
    Check if Conv3d memory bug workaround needed.
    Bug: PyTorch 2.9+ with cuDNN >= 91002 uses 3x memory for Conv3d 
    with fp16/bfloat16 due to buggy dispatch layer.
    """
    try:
        # Exclude AMD ROCm/HIP builds (they use MIOpen, not cuDNN)
        if hasattr(torch.version, 'hip') and torch.version.hip is not None:
            return False
        
        # Must have CUDA available
        if not (hasattr(torch, 'cuda') and torch.cuda.is_available()):
            return False
        
        # Must have cuDNN actually available (not just the attribute)
        if not (hasattr(torch.backends.cudnn, 'is_available') and 
                torch.backends.cudnn.is_available()):
            return False
        
        # Check device capability (NVIDIA GPUs)
        if torch.cuda.get_device_capability()[0] < 3:
            return False
        
        # Parse torch version
        version_str = torch.__version__.split('+')[0]
        parts = version_str.split('.')
        torch_version = tuple(int(p) for p in parts[:2])
        
        # Bug affects PyTorch 2.9 and later versions
        if torch_version < (2, 9):
            return False
        
        if not hasattr(torch.backends.cudnn, 'version'):
            return False
            
        cudnn_version = torch.backends.cudnn.version()
        if cudnn_version is None or cudnn_version < 91002:
            return False
        
        return True
    except:
        return False

NVIDIA_CONV3D_MEMORY_BUG_WORKAROUND = _check_conv3d_memory_bug()


# Log all optimization status once globally (cross-process) using environment variable
if not os.environ.get("SEEDVR2_OPTIMIZATIONS_LOGGED"):
    os.environ["SEEDVR2_OPTIMIZATIONS_LOGGED"] = "1"
    
    # Build status strings
    sage_status = "✅" if SAGE_ATTN_AVAILABLE else "❌"
    flash_status = "✅" if FLASH_ATTN_AVAILABLE else "❌"
    triton_status = "✅" if TRITON_AVAILABLE else "❌"
    
    # Count available optimizations
    available = [SAGE_ATTN_AVAILABLE, FLASH_ATTN_AVAILABLE, TRITON_AVAILABLE]
    num_available = sum(available)
    
    if num_available == 3:
        print(f"⚡ SeedVR2 optimizations check: SageAttention {sage_status} | Flash Attention {flash_status} | Triton {triton_status}")
    elif num_available == 0:
        print(f"⚠️  SeedVR2 optimizations check: SageAttention {sage_status} | Flash Attention {flash_status} | Triton {triton_status}")
        print("💡 For best performance: pip install sageattention flash-attn triton")
    else:
        icon = "⚡" if num_available >= 2 else "⚠️ "
        print(f"{icon} SeedVR2 optimizations check: SageAttention {sage_status} | Flash Attention {flash_status} | Triton {triton_status}")
        
        # Build install suggestions for missing packages
        missing = []
        if not SAGE_ATTN_AVAILABLE:
            missing.append("sageattention")
        if not FLASH_ATTN_AVAILABLE:
            missing.append("flash-attn")
        if not TRITON_AVAILABLE:
            missing.append("triton")
        if missing:
            print(f"💡 Optional: pip install {' '.join(missing)}")
    
    # Conv3d workaround status (if applicable)
    if NVIDIA_CONV3D_MEMORY_BUG_WORKAROUND:
        torch_ver = torch.__version__.split('+')[0]
        cudnn_ver = torch.backends.cudnn.version()
        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:
    1. Clear cache if available
    2. Disable autocast to prevent numerical issues (CUDA only)
    This prevents artifacts in FP8/mixed precision models.
    """
    if hasattr(method, 'cache_clear'):
        method.cache_clear()
    
    # Only use CUDA autocast context on CUDA devices
    # MPS has no CUDA autocast to disable
    if torch.cuda.is_available():
        with torch.cuda.amp.autocast(enabled=False):
            return method(*args, **kwargs)
    else:
        return method(*args, **kwargs)
    
    
class CompatibleDiT(torch.nn.Module):
    """
    Wrapper for DiT models with automatic compatibility management + advanced optimizations
    
    Precision Handling:
    - FP8: Keeps native FP8 parameters (memory efficient), converts inputs/outputs to compute_dtype for arithmetic
    - FP16/BFloat16/Float32: Uses native precision throughout
    - GGUF: On-the-fly dequantization to compute_dtype
    - MPS: Forces all parameters to compute_dtype (unified memory requires dtype consistency)
    - RoPE: Converted from FP8 to compute_dtype for numerical consistency
    
    Optimizations:
    - RoPE Stabilization: Error handling for numerical stability in mixed precision
    - MPS Compatibility: Unified dtype conversion for Apple Silicon backends
    """
    
    def __init__(self, dit_model, debug: 'Debug', compute_dtype: torch.dtype = torch.bfloat16, skip_conversion: bool = False):
        super().__init__()
        self.dit_model = dit_model
        self.debug = debug
        self.compute_dtype = compute_dtype
        self.model_dtype = self._detect_model_dtype()
        self.is_fp8_model = self.model_dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
        self.is_fp16_model = self.model_dtype == torch.float16

        # Only convert if not already done (e.g., when reusing cached weights)
        if not skip_conversion and self.is_fp8_model:
            # FP8 models need RoPE frequency conversion to compute dtype
            model_variant = self._get_model_variant()
            self.debug.log(f"Detected NaDiT {model_variant} FP8 - Converting RoPE freqs for FP8 compatibility", 
                        category="precision")
            self.debug.start_timer("_convert_rope_freqs")
            self._convert_rope_freqs(target_dtype=self.compute_dtype)
            self.debug.end_timer("_convert_rope_freqs", "RoPE freqs conversion")
        
        # MPS requires unified dtype for all parameters/buffers (no autocast fallback)
        # Apply to ALL model types (FP8, FP16, GGUF) when dtype differs from compute_dtype
        if not skip_conversion and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
            if self.model_dtype != self.compute_dtype:
                self.debug.log(f"Converting NaDiT parameters/buffers to {self.compute_dtype} for MPS backend", category="setup", force=True)
                self.debug.start_timer("_force_nadit_precision")
                self._force_nadit_precision(target_dtype=self.compute_dtype)
                self.debug.end_timer("_force_nadit_precision", "NaDiT parameters/buffers conversion")
            
        # Apply RoPE stabilization for numerical stability
        self.debug.log(f"Stabilizing RoPE computations for numerical stability", category="setup")
        self.debug.start_timer("_stabilize_rope_computations")
        self._stabilize_rope_computations()
        self.debug.end_timer("_stabilize_rope_computations", "RoPE stabilization")
    
    def _detect_model_dtype(self) -> torch.dtype:
        """Detect main model dtype"""
        try:
            return next(self.dit_model.parameters()).dtype
        except:
            return torch.bfloat16
    
    def _get_model_variant(self) -> str:
        """Detect model variant from module path"""
        model_module = str(self.dit_model.__class__.__module__).lower()
        if 'dit_7b' in model_module:
            return "7B"
        elif 'dit_3b' in model_module:
            return "3B"
        else:
            return "Unknown"
        
    def _convert_rope_freqs(self, target_dtype: torch.dtype = torch.bfloat16) -> None:
        """
        Convert RoPE frequency buffers from FP8 to target dtype for compatibility.
        
        Args:
            target_dtype: Target dtype for RoPE freqs (default: bfloat16 for stability)
        """
        converted = 0
        for module in self.dit_model.modules():
            if 'RotaryEmbedding' in type(module).__name__:
                if hasattr(module, 'rope') and hasattr(module.rope, 'freqs'):
                    if module.rope.freqs.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
                        if module.rope.freqs.device.type == "mps":
                            module.rope.freqs.data = module.rope.freqs.to("cpu").to(target_dtype).to("mps")
                        else:
                            module.rope.freqs.data = module.rope.freqs.to(target_dtype)
                        converted += 1
        self.debug.log(f"Converted {converted} RoPE frequency buffers from FP8 to {target_dtype} for compatibility", category="success")
                        
    def _force_nadit_precision(self, target_dtype: torch.dtype = torch.bfloat16) -> None:
        """
        Force ALL NaDiT parameters to target dtype to avoid promotion errors (MPS requirement).
        
        Args:
            target_dtype: Target dtype for all parameters (default: bfloat16 for MPS compatibility)
        """
        converted_count = 0
        original_dtype = None
        
        # Convert ALL parameters to target dtype
        for name, param in self.dit_model.named_parameters():
            if original_dtype is None:
                original_dtype = param.dtype
            if param.dtype != target_dtype:
                if param.device.type == "mps":
                    temp_cpu = param.data.to("cpu")
                    temp_converted = temp_cpu.to(target_dtype)
                    param.data = temp_converted.to("mps")
                    del temp_cpu, temp_converted
                else:
                    param.data = param.data.to(target_dtype)
                converted_count += 1
                
        # Also convert buffers (skip GGUF quantized buffers - they have tensor_type attribute)
        for name, buffer in self.dit_model.named_buffers():
            # Skip GGUF quantized buffers - these must stay in packed format for on-the-fly dequantization
            if hasattr(buffer, 'tensor_type'):
                continue
            if buffer.dtype != target_dtype:
                if buffer.device.type == "mps":
                    temp_cpu = buffer.data.to("cpu")
                    temp_converted = temp_cpu.to(target_dtype)
                    buffer.data = temp_converted.to("mps")
                    del temp_cpu, temp_converted
                else:
                    buffer.data = buffer.data.to(target_dtype)
                converted_count += 1
        
        self.debug.log(f"Converted {converted_count} NaDiT parameters/buffers to {target_dtype} for MPS", category="success")
        
        # Update detected dtype
        self.model_dtype = target_dtype
        self.is_fp8_model = (target_dtype in (torch.float8_e4m3fn, torch.float8_e5m2))

    def _stabilize_rope_computations(self):
        """
        Add error handling to RoPE computations to prevent artifacts.
        
        Wraps the get_axial_freqs method of RoPE modules with a try-except handler.
        During normal operation, uses the original cached method for performance.
        Only on exceptions (e.g., numerical instability, NaN propagation) does it
        intervene by clearing the cache and retrying the computation through
        call_rope_with_stability.
        
        This prevents artifacts in FP8, mixed precision, and edge cases while
        maintaining optimal performance for normal operations.
        """
        if not hasattr(self.dit_model, 'blocks'):
            return
        
        rope_count = 0
        
        # Wrap RoPE modules to handle numerical instability
        for name, module in self.dit_model.named_modules():
            if "rope" in name.lower() and hasattr(module, "get_axial_freqs"):
                # Check if already wrapped
                if hasattr(module, '_rope_wrapped'):
                    continue
                    
                original_method = module.get_axial_freqs
                
                # Mark as wrapped and store original
                module._rope_wrapped = 'stability'
                module._original_get_axial_freqs = original_method
                
                # Error handler that prevents NaN propagation
                def stable_rope_computation(self, *args, **kwargs):
                    try:
                        return original_method(*args, **kwargs)
                    except Exception:
                        return call_rope_with_stability(original_method, *args, **kwargs)
                
                module.get_axial_freqs = types.MethodType(stable_rope_computation, module)
                rope_count += 1
        
        if rope_count > 0:
            self.debug.log(f"Stabilized {rope_count} RoPE modules", category="success")
    
    def forward(self, *args, **kwargs):
        """
        Forward pass with minimal dtype conversion overhead
        
        Conversion strategy:
            - FP16/BFloat16/Float32 models: Use native precision (no conversion needed)
            - FP8 models: Convert FP8 tensors to compute_dtype for arithmetic operations
            (FP8 parameters stay in FP8 for memory efficiency, only converted for computation)
        """
        
        # Only convert if we have an FP8 model for arithmetic operations 
        if self.is_fp8_model:
            fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2)
            target_dtype = self.compute_dtype
            
            # Convert args
            converted_args = []
            for arg in args:
                if isinstance(arg, torch.Tensor) and arg.dtype in fp8_dtypes:
                    converted_args.append(arg.to(target_dtype))
                else:
                    converted_args.append(arg)
            
            # Convert kwargs
            converted_kwargs = {}
            for key, value in kwargs.items():
                if isinstance(value, torch.Tensor) and value.dtype in fp8_dtypes:
                    converted_kwargs[key] = value.to(target_dtype)
                else:
                    converted_kwargs[key] = value
            
            args = tuple(converted_args)
            kwargs = converted_kwargs
        
        # Execute forward pass
        try:
            return self.dit_model(*args, **kwargs)
        except Exception as e:
            self.debug.log(f"Forward pass error: {e}", level="ERROR", category="generation", force=True)
            if self.is_fp8_model:
                self.debug.log(f"FP8 model - converted FP8 tensors to {self.compute_dtype}", category="info", force=True)
            else:
                self.debug.log(f"{self.model_dtype} model - no conversion applied", category="info", force=True)
            raise
    
    def __getattr__(self, name):
        """Redirect all other attributes to original model"""
        if name in ['dit_model', 'model_dtype', 'is_fp8_model', 'is_fp16_model']:
            return super().__getattr__(name)
        return getattr(self.dit_model, name)
    
    def __setattr__(self, name, value):
        """Redirect assignments to original model except for our attributes"""
        if name in ['dit_model', 'model_dtype', 'is_fp8_model', 'is_fp16_model']:
            super().__setattr__(name, value)
        else:
            if hasattr(self, 'dit_model'):
                setattr(self.dit_model, name, value)
            else:
                super().__setattr__(name, value)