""" """ from typing import cast import torch from huggingface_hub import hf_hub_download from spaces.zero.torch.aoti import ZeroGPUCompiledModel from spaces.zero.torch.aoti import 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): repeated_blocks = cast(list[str], module._repeated_blocks) aoti_files = {name: hf_hub_download( repo_id=repo_id, filename='package.pt2', subfolder=name if variant is None else f'{name}.{variant}', ) for name in repeated_blocks} for block_name, aoti_file in aoti_files.items(): 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)