app.py

"""
"""

# Upgrade PyTorch
import os
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()

# 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 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"

@spaces.GPU
def run_model():
    compiled_model = torch._inductor.aoti_load_package(package_path)
    with torch.inference_mode():
        return str(compiled_model(example_inputs))


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)