Delete aoti.py

aoti.py

@@ -1,99 +0,0 @@
-"""
-AOTI (Ahead-of-Time Compilation) loading utilities for ZeroGPU
-Based on official HuggingFace zerogpu-aoti implementation
-Author: cbensimon (HF Staff)
-"""
-
-from typing import cast
-import torch
-from huggingface_hub import hf_hub_download
-
-try:
-    from spaces.zero.torch.aoti import ZeroGPUCompiledModel, ZeroGPUWeights
-    from torch._functorch._aot_autograd.subclass_parametrization import unwrap_tensor_subclass_parameters
-    AOTI_AVAILABLE = True
-except ImportError:
-    AOTI_AVAILABLE = False
-    ZeroGPUCompiledModel = None
-    ZeroGPUWeights = None
-
-
-def _shallow_clone_module(module: torch.nn.Module) -> torch.nn.Module:
-    """Create a shallow clone of a PyTorch module."""
-    clone = object.__new__(module.__class__)
-    clone.__dict__ = module.__dict__.copy()
-    clone._parameters = module._parameters.copy()
-    clone._buffers = module._buffers.copy()
-    clone._modules = {
-        k: _shallow_clone_module(v)
-        for k, v in module._modules.items()
-        if v is not None
-    }
-    return clone
-
-
-def aoti_blocks_load(module: torch.nn.Module, repo_id: str, variant: str | None = None):
-    """
-    Load AOTI compiled blocks for a transformer module.
-
-    Args:
-        module: The transformer module to optimize
-        repo_id: HuggingFace repo containing AOTI files (e.g., 'zerogpu-aoti/Wan2')
-        variant: Variant name (e.g., 'fp8da' for FP8 dynamic activation)
-
-    Raises:
-        ImportError: If AOTI dependencies are not available
-        Exception: If AOTI files cannot be downloaded or loaded
-    """
-    if not AOTI_AVAILABLE:
-        raise ImportError(
-            "AOTI is not available. Please upgrade 'spaces' package:\n"
-            "pip install --upgrade spaces"
-        )
-
-    # Get repeated blocks from module
-    if hasattr(module, '_repeated_blocks'):
-        repeated_blocks = cast(list[str], module._repeated_blocks)
-    else:
-        # Fallback: assume WanTransformerBlock for Wan models
-        print("⚠ Module doesn't have _repeated_blocks, using default 'WanTransformerBlock'")
-        repeated_blocks = ['WanTransformerBlock']
-
-    # Download AOTI files for each block type
-    aoti_files = {}
-    for name in repeated_blocks:
-        try:
-            subfolder = name if variant is None else f'{name}.{variant}'
-            aoti_file = hf_hub_download(
-                repo_id=repo_id,
-                filename='package.pt2',
-                subfolder=subfolder,
-            )
-            aoti_files[name] = aoti_file
-            print(f"  ✓ Downloaded AOTI for {name}")
-        except Exception as e:
-            print(f"  ✗ Failed to download AOTI for {name}: {e}")
-            raise
-
-    # Replace forward methods with compiled versions
-    for block_name, aoti_file in aoti_files.items():
-        replaced_count = 0
-        for block in module.modules():
-            if block.__class__.__name__ == block_name:
-                # Clone and unwrap
-                block_ = _shallow_clone_module(block)
-                unwrap_tensor_subclass_parameters(block_)
-
-                # Extract weights
-                weights = ZeroGPUWeights(block_.state_dict())
-
-                # Replace forward with compiled version
-                block.forward = ZeroGPUCompiledModel(aoti_file, weights)
-                replaced_count += 1
-
-        print(f"  ✓ Replaced {replaced_count} blocks of type {block_name}")
-
-
-def check_aoti_available() -> bool:
-    """Check if AOTI is available in the current environment."""
-    return AOTI_AVAILABLE