Update app.py

app.py

@@ -5,6 +5,7 @@ import random
 from datetime import datetime
 from typing import List, Tuple
 
+import spaces
 import gradio as gr
 import torch
 import torch.nn as nn
@@ -15,7 +16,7 @@ from PIL import Image
 
 
 # ============================================================
-# Paths / Device
+# Paths / basic config
 # ============================================================
 BASE_DIR = os.path.dirname(os.path.abspath(__file__)) if "__file__" in globals() else os.getcwd()
 DATA_DIR = os.path.join(BASE_DIR, "data")
@@ -26,8 +27,6 @@ os.makedirs(DATA_DIR, exist_ok=True)
 os.makedirs(MODEL_DIR, exist_ok=True)
 os.makedirs(META_DIR, exist_ok=True)
 
-# Force CPU on Hugging Face Spaces for this lightweight demo
-DEVICE = torch.device("cpu")
 CLASS_NAMES = [str(i) for i in range(10)]
 
 
@@ -56,8 +55,7 @@ class SimpleCNN(nn.Module):
             nn.MaxPool2d(2),
         )
 
-        # 28x28 -> 14x14 -> 7x7
-        flattened_dim = conv2_channels * 7 * 7
+        flattened_dim = conv2_channels * 7 * 7  # 28x28 -> 14x14 -> 7x7
 
         self.classifier = nn.Sequential(
             nn.Flatten(),
@@ -132,7 +130,9 @@ def list_saved_models() -> List[str]:
 
 
 def save_model(model: nn.Module, model_name: str, config: dict, training_summary: dict):
-    torch.save(model.state_dict(), model_weight_path(model_name))
+    cpu_state_dict = {k: v.detach().cpu() for k, v in model.state_dict().items()}
+    torch.save(cpu_state_dict, model_weight_path(model_name))
+
     payload = {
         "model_name": model_name,
         "config": config,
@@ -143,7 +143,7 @@ def save_model(model: nn.Module, model_name: str, config: dict, training_summary
         json.dump(payload, f, indent=2, ensure_ascii=False)
 
 
-def load_model(model_name: str) -> Tuple[nn.Module, dict]:
+def load_model(model_name: str, device: torch.device) -> Tuple[nn.Module, dict]:
     meta_file = model_meta_path(model_name)
     weight_file = model_weight_path(model_name)
 
@@ -164,40 +164,23 @@ def load_model(model_name: str) -> Tuple[nn.Module, dict]:
         dropout=cfg["dropout"],
         fc_dim=cfg["fc_dim"],
     )
-    state_dict = torch.load(weight_file, map_location=DEVICE)
+
+    state_dict = torch.load(weight_file, map_location="cpu")
     model.load_state_dict(state_dict)
-    model.to(DEVICE)
+    model.to(device)
     model.eval()
     return model, meta
 
 
 # ============================================================
-# Train / Eval
+# ZeroGPU helpers
 # ============================================================
-def evaluate(model: nn.Module, loader: DataLoader, criterion: nn.Module):
-    model.eval()
-    total_loss = 0.0
-    total = 0
-    correct = 0
-
-    with torch.no_grad():
-        for images, labels in loader:
-            images, labels = images.to(DEVICE), labels.to(DEVICE)
-
-            outputs = model(images)
-            loss = criterion(outputs, labels)
-
-            total_loss += loss.item() * images.size(0)
-            preds = outputs.argmax(dim=1)
-            correct += (preds == labels).sum().item()
-            total += labels.size(0)
+def get_runtime_device() -> torch.device:
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    avg_loss = total_loss / total if total else 0.0
-    acc = correct / total if total else 0.0
-    return avg_loss, acc
 
-
-def train_model(
+@spaces.GPU(duration=120)
+def _train_on_gpu(
     dataset_name: str,
     conv1_channels: int,
     conv2_channels: int,
@@ -209,6 +192,8 @@ def train_model(
     epochs: int,
     model_tag: str,
 ):
+    device = get_runtime_device()
+
     train_loader, val_loader, test_loader = make_loaders(dataset_name, batch_size)
 
     model = SimpleCNN(
@@ -217,7 +202,7 @@ def train_model(
         kernel_size=kernel_size,
         dropout=dropout,
         fc_dim=fc_dim,
-    ).to(DEVICE)
+    ).to(device)
 
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.Adam(model.parameters(), lr=learning_rate)
@@ -226,6 +211,27 @@ def train_model(
     logs = []
     start_time = time.time()
 
+    def evaluate(loader):
+        model.eval()
+        total_loss = 0.0
+        total = 0
+        correct = 0
+
+        with torch.no_grad():
+            for images, labels in loader:
+                images, labels = images.to(device), labels.to(device)
+                outputs = model(images)
+                loss = criterion(outputs, labels)
+
+                total_loss += loss.item() * images.size(0)
+                preds = outputs.argmax(dim=1)
+                correct += (preds == labels).sum().item()
+                total += labels.size(0)
+
+        avg_loss = total_loss / total if total else 0.0
+        acc = correct / total if total else 0.0
+        return avg_loss, acc
+
     for epoch in range(1, epochs + 1):
         model.train()
         running_loss = 0.0
@@ -233,7 +239,7 @@ def train_model(
         correct = 0
 
         for images, labels in train_loader:
-            images, labels = images.to(DEVICE), labels.to(DEVICE)
+            images, labels = images.to(device), labels.to(device)
 
             optimizer.zero_grad()
             outputs = model(images)
@@ -248,7 +254,7 @@ def train_model(
 
         train_loss = running_loss / total if total else 0.0
         train_acc = correct / total if total else 0.0
-        val_loss, val_acc = evaluate(model, val_loader, criterion)
+        val_loss, val_acc = evaluate(val_loader)
 
         row = {
             "epoch": epoch,
@@ -265,13 +271,7 @@ def train_model(
             f"val_loss={val_loss:.4f}, val_acc={val_acc:.4f}"
         )
 
-        yield (
-            "\n".join(logs),
-            history,
-            gr.update(),
-        )
-
-    test_loss, test_acc = evaluate(model, test_loader, criterion)
+    test_loss, test_acc = evaluate(test_loader)
     elapsed = time.time() - start_time
 
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
@@ -298,7 +298,7 @@ def train_model(
         "test_loss": round(test_loss, 4),
         "test_acc": round(test_acc, 4),
         "elapsed_seconds": round(elapsed, 2),
-        "device": str(DEVICE),
+        "device": str(device),
     }
 
     save_model(model, model_name, config, training_summary)
@@ -306,27 +306,24 @@ def train_model(
     logs.append("")
     logs.append("Training finished.")
     logs.append(f"Saved model: {model_name}")
-    logs.append(f"Device: {DEVICE}")
+    logs.append(f"Device: {device}")
     logs.append(f"Test loss: {test_loss:.4f}")
     logs.append(f"Test accuracy: {test_acc:.4f}")
     logs.append(f"Elapsed time: {elapsed:.1f}s")
 
-    models = list_saved_models()
-    selected = model_name if model_name in models else (models[0] if models else None)
+    return "\n".join(logs), history, training_summary, model_name
 
-    yield (
-        "\n".join(logs),
-        history,
-        gr.update(choices=models, value=selected),
-    )
 
+@spaces.GPU(duration=60)
+def _predict_uploaded_image_gpu(model_name: str, image: Image.Image):
+    if not model_name:
+        return "Please select a model.", None
 
-# ============================================================
-# Inference
-# ============================================================
-def preprocess_uploaded_image(image: Image.Image):
     if image is None:
-        raise ValueError("Please upload an image.")
+        return "Please upload an image.", None
+
+    device = get_runtime_device()
+    model, meta = load_model(model_name, device)
 
     transform = transforms.Compose(
         [
@@ -336,37 +333,32 @@ def preprocess_uploaded_image(image: Image.Image):
             transforms.Normalize((0.5,), (0.5,))
         ]
     )
-    tensor = transform(image).unsqueeze(0)
-    return tensor
 
-
-def predict_uploaded_image(model_name: str, image: Image.Image):
-    if not model_name:
-        return "Please select a model.", None
-
-    model, meta = load_model(model_name)
-    tensor = preprocess_uploaded_image(image).to(DEVICE)
+    tensor = transform(image).unsqueeze(0).to(device)
 
     with torch.no_grad():
         logits = model(tensor)
-        probs = torch.softmax(logits, dim=1).squeeze(0).cpu().tolist()
+        probs = torch.softmax(logits, dim=1).squeeze(0).detach().cpu().tolist()
         pred_idx = int(torch.argmax(logits, dim=1).item())
 
     result_text = (
         f"Prediction: {CLASS_NAMES[pred_idx]}\n"
         f"Confidence: {max(probs):.4f}\n\n"
         f"Model: {model_name}\n"
-        f"Dataset: {meta['config']['dataset_name']}"
+        f"Dataset: {meta['config']['dataset_name']}\n"
+        f"Runtime device: {device}"
     )
     prob_dict = {CLASS_NAMES[i]: float(probs[i]) for i in range(10)}
     return result_text, prob_dict
 
 
-def test_random_sample(model_name: str):
+@spaces.GPU(duration=60)
+def _test_random_sample_gpu(model_name: str):
     if not model_name:
         return None, "Please select a model.", None
 
-    model, meta = load_model(model_name)
+    device = get_runtime_device()
+    model, meta = load_model(model_name, device)
     dataset_name = meta["config"]["dataset_name"]
 
     _, test_dataset = get_datasets(dataset_name)
@@ -374,8 +366,8 @@ def test_random_sample(model_name: str):
     image_tensor, label = test_dataset[idx]
 
     with torch.no_grad():
-        logits = model(image_tensor.unsqueeze(0).to(DEVICE))
-        probs = torch.softmax(logits, dim=1).squeeze(0).cpu().tolist()
+        logits = model(image_tensor.unsqueeze(0).to(device))
+        probs = torch.softmax(logits, dim=1).squeeze(0).detach().cpu().tolist()
         pred_idx = int(torch.argmax(logits, dim=1).item())
 
     display_img = image_tensor.squeeze(0).cpu().numpy()
@@ -385,12 +377,62 @@ def test_random_sample(model_name: str):
         f"Ground truth: {label}\n"
         f"Prediction: {pred_idx}\n"
         f"Confidence: {max(probs):.4f}\n"
-        f"Model dataset: {dataset_name}"
+        f"Model dataset: {dataset_name}\n"
+        f"Runtime device: {device}"
     )
     prob_dict = {CLASS_NAMES[i]: float(probs[i]) for i in range(10)}
     return display_img, result_text, prob_dict
 
 
+# ============================================================
+# UI callbacks
+# ============================================================
+def train_callback(
+    dataset_name,
+    conv1_channels,
+    conv2_channels,
+    kernel_size,
+    dropout,
+    fc_dim,
+    learning_rate,
+    batch_size,
+    epochs,
+    model_tag,
+):
+    try:
+        logs, history, summary, model_name = _train_on_gpu(
+            dataset_name,
+            int(conv1_channels),
+            int(conv2_channels),
+            int(kernel_size),
+            float(dropout),
+            int(fc_dim),
+            float(learning_rate),
+            int(batch_size),
+            int(epochs),
+            model_tag,
+        )
+        models = list_saved_models()
+        selected = model_name if model_name in models else (models[0] if models else None)
+        return logs, history, summary, gr.update(choices=models, value=selected)
+    except Exception as e:
+        return f"Training failed:\n{str(e)}", None, None, gr.update()
+
+
+def predict_uploaded_image_callback(model_name, image):
+    try:
+        return _predict_uploaded_image_gpu(model_name, image)
+    except Exception as e:
+        return f"Prediction failed:\n{str(e)}", None
+
+
+def test_random_sample_callback(model_name):
+    try:
+        return _test_random_sample_gpu(model_name)
+    except Exception as e:
+        return None, f"Random test failed:\n{str(e)}", None
+
+
 def get_model_info(model_name: str):
     if not model_name:
         return {"message": "No model selected."}
@@ -428,7 +470,7 @@ with gr.Blocks(title="Image Classification") as demo:
                     dataset_name = gr.Dropdown(
                         choices=["MNIST", "FashionMNIST"],
                         value="MNIST",
-                        label="Dataset"
+                        label="Dataset",
                     )
                     conv1_channels = gr.Slider(8, 64, value=16, step=8, label="Conv1 Channels")
                     conv2_channels = gr.Slider(16, 128, value=32, step=16, label="Conv2 Channels")
@@ -444,6 +486,7 @@ with gr.Blocks(title="Image Classification") as demo:
                 with gr.Column():
                     train_status = gr.Textbox(label="Training Log", lines=18)
                     train_history = gr.JSON(label="Training History")
+                    train_summary = gr.JSON(label="Training Summary")
 
         with gr.Tab("Test"):
             with gr.Row():
@@ -451,7 +494,7 @@ with gr.Blocks(title="Image Classification") as demo:
                     model_selector = gr.Dropdown(
                         choices=initial_models,
                         value=initial_models[0] if initial_models else None,
-                        label="Select Saved Model"
+                        label="Select Saved Model",
                     )
                     refresh_btn = gr.Button("Refresh Model List")
                     load_info_btn = gr.Button("Show Model Info")
@@ -460,7 +503,7 @@ with gr.Blocks(title="Image Classification") as demo:
                 with gr.Column():
                     upload_image = gr.Image(type="pil", label="Upload Image")
                     predict_btn = gr.Button("Predict Uploaded Image", variant="primary")
-                    predict_text = gr.Textbox(label="Prediction Result", lines=6)
+                    predict_text = gr.Textbox(label="Prediction Result", lines=7)
                     predict_probs = gr.Label(label="Class Probabilities")
 
             with gr.Row():
@@ -468,11 +511,11 @@ with gr.Blocks(title="Image Classification") as demo:
 
             with gr.Row():
                 random_sample_image = gr.Image(type="numpy", label="Random Test Image")
-                random_sample_text = gr.Textbox(label="Random Sample Result", lines=6)
+                random_sample_text = gr.Textbox(label="Random Sample Result", lines=7)
                 random_sample_probs = gr.Label(label="Random Sample Probabilities")
 
     train_btn.click(
-        fn=train_model,
+        fn=train_callback,
         inputs=[
             dataset_name,
             conv1_channels,
@@ -485,7 +528,7 @@ with gr.Blocks(title="Image Classification") as demo:
             epochs,
             model_tag,
         ],
-        outputs=[train_status, train_history, model_selector],
+        outputs=[train_status, train_history, train_summary, model_selector],
     )
 
     refresh_btn.click(
@@ -501,17 +544,17 @@ with gr.Blocks(title="Image Classification") as demo:
     )
 
     predict_btn.click(
-        fn=predict_uploaded_image,
+        fn=predict_uploaded_image_callback,
         inputs=[model_selector, upload_image],
         outputs=[predict_text, predict_probs],
     )
 
     random_test_btn.click(
-        fn=test_random_sample,
+        fn=test_random_sample_callback,
         inputs=[model_selector],
         outputs=[random_sample_image, random_sample_text, random_sample_probs],
     )
 
 
 if __name__ == "__main__":
-    demo.launch(ssr_mode=False)
+    demo.launch()