# aoti.py — lightweight version for 8GB RAM
"""
AOTI loader simplified for low VRAM systems.
"""

import torch
from typing import cast
from huggingface_hub import hf_hub_download
from spaces.zero.torch.aoti import ZeroGPUCompiledModel, ZeroGPUWeights
from torch._functorch._aot_autograd.subclass_parametrization import unwrap_tensor_subclass_parameters


def _shallow_clone_module(module: torch.nn.Module) -> torch.nn.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):
    """
    Safe AOTI loader for low-memory systems.
    Loads only repeated blocks and avoids deep cloning.
    """

    if not hasattr(module, "_repeated_blocks"):
        return  # safety fallback

    repeated_blocks = cast(list[str], module._repeated_blocks)

    for block_name in repeated_blocks:
        aoti_file = hf_hub_download(
            repo_id=repo_id,
            filename="package.pt2",
            subfolder=block_name if variant is None else f"{block_name}.{variant}",
        )

        for block in module.modules():
            if block.__class__.__name__ == block_name:
                block_ = _shallow_clone_module(block)
                unwrap_tensor_subclass_parameters(block_)
                weights = ZeroGPUWeights(block_.state_dict())
                block.forward = ZeroGPUCompiledModel(aoti_file, weights)