Fix aoti

aoti.py

@@ -1,6 +1,8 @@
 """
 """
 
+from typing import cast
+
 import torch
 from huggingface_hub import hf_hub_download
 from spaces.zero.torch.aoti import ZeroGPUCompiledModel
@@ -13,14 +15,17 @@ def _shallow_clone_module(module: torch.nn.Module) -> torch.nn.Module:
     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()}
+    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 = module._repeated_blocks
-    subfolder = name if variant is None else f'{name}.{variant}'
-    aoti_files = {name: hf_hub_download(repo_id, 'package.pt2', subfolder=subfolder) for name in repeated_blocks}
+    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: