Create aoti.py

aoti.py

@@ -0,0 +1,15 @@
+import torch
+from huggingface_hub import hf_hub_download
+from spaces.zero.torch.aoti import ZeroGPUCompiledModel
+from spaces.zero.torch.aoti import ZeroGPUWeights
+from spaces.zero.torch.aoti import drain_module_parameters
+
+
+def aoti_load_(module: torch.nn.Module, repo_id: str, filename: str):
+    compiled_graph_file = hf_hub_download(repo_id, filename)
+    state_dict = module.state_dict()
+    zerogpu_weights = ZeroGPUWeights({name: weight for name, weight in state_dict.items()})
+    compiled = ZeroGPUCompiledModel(compiled_graph_file, zerogpu_weights)
+
+    setattr(module, "forward", compiled)
+    drain_module_parameters(module)