Phase 1

aoti_base_example.py

@@ -1,37 +0,0 @@
-"""
-Modified from https://docs.pytorch.org/tutorials/recipes/torch_export_aoti_python.html
-"""
-
-import os
-import torch
-import torch._inductor
-from torchvision.models import ResNet18_Weights, resnet18
-
-model = resnet18(weights=ResNet18_Weights.DEFAULT)
-model.eval()
-
-package_path = os.path.join(os.getcwd(), "resnet18.pt2")
-inductor_configs = {'max_autotune': True}
-device = "cuda"
-
-# Compile
-with torch.inference_mode():
-    model = model.to(device=device)
-    example_inputs = (torch.randn(2, 3, 224, 224, device=device),)
-    exported_program = torch.export.export(
-        model,
-        example_inputs,
-    )
-    torch._inductor.aoti_compile_and_package(
-        exported_program,
-        package_path=package_path,
-        inductor_configs=inductor_configs
-    )
-
-# Load
-compiled_model = torch._inductor.aoti_load_package(package_path)
-example_inputs = (torch.randn(2, 3, 224, 224, device=device),)
-
-# Run
-with torch.inference_mode():
-    output = compiled_model(example_inputs)

app.py

@@ -3,23 +3,54 @@
 
 # Upgrade PyTorch
 import os
-os.system('pip install --upgrade torch torchvision')
+os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 torch torchvision spaces')
 
 # CUDA toolkit install
 from utils.cuda_toolkit import install_cuda_toolkit; install_cuda_toolkit()
 
-# Try AOTI compile + load
+# Actual app.py
+import os
+
+import gradio as gr
 import spaces
+import torch
+import torch._inductor
+from torchvision.models import ResNet18_Weights, resnet18
+
+
+model = resnet18(weights=ResNet18_Weights.DEFAULT)
+model.eval()
+model.to('cuda')
+
+package_path = os.path.join(os.getcwd(), 'resnet18.pt2')
+inductor_configs = {'max_autotune': True}
+example_inputs = (torch.randn(2, 3, 224, 224, device='cuda'),)
+
 @spaces.GPU
-def run():
-    import aoti_base_example
-run()
+def compile_model():
+    with torch.inference_mode():
+        exported_program = torch.export.export(
+            model,
+            example_inputs,
+        )
+        torch._inductor.aoti_compile_and_package(
+            exported_program,
+            package_path=package_path,
+            inductor_configs=inductor_configs
+        )
+    return "compiled"
 
-# Base demo
-import gradio as gr
+@spaces.GPU
+def run_model():
+    compiled_model = torch._inductor.aoti_load_package(package_path)
+    with torch.inference_mode():
+        return str(compiled_model(example_inputs))
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+gr.TabbedInterface([
+    gr.Interface(compile_model, [], "text", clear_btn=None, flagging_mode='never'),
+    gr.Interface(run_model, [], "text", clear_btn=None, flagging_mode='never'),
+], [
+    "Compile",
+    "Run"
+]).launch(show_error=True)