fix gradio demo

app.py

@@ -3,22 +3,19 @@ import time
 import gradio as gr
 import torch
 import torch.nn.functional as F
+from torchvision.utils import save_image
 from huggingface_hub import hf_hub_download
 
 from digitnet import Model
 
 torch.set_grad_enabled(False)
 
-hf_hub_download("karanravindra/digitnet", filename="mnist", local_dir=".")
+hf_hub_download("karanravindra/digitnet", filename="model.ckpt", local_dir=".")
 
-model = Model(1, 16, num_classes=36)
-model.load_state_dict(
-    torch.load("models/emnist-model.pth", weights_only=True, map_location="cpu")
-)
+model = Model.load_from_checkpoint("model.ckpt", map_location="cpu")
 model.eval()
 
-all_classes = list("0123456789abcdefghijklmnopqrstuvwxyz")
-all_classes = map(str.upper, all_classes)
+all_classes = list(map(str.upper, "0123456789"))
 
 
 def predict(inputs: dict) -> dict[str, float]:
@@ -32,7 +29,6 @@ def predict(inputs: dict) -> dict[str, float]:
     logits = model(img)
 
     probs = torch.softmax(logits, dim=1).squeeze()
-
     return {c: p for c, p in zip(all_classes, probs.tolist())}
 
 
@@ -50,8 +46,10 @@ demo = gr.Interface(
     live=True,
     inputs=input,
     outputs="label",
+    submit_btn=gr.Button("Predict"),
     title="DigitNet",
     description="A simple handwritten number and letter classifier.\n\nDraw a digit or letter in the box below and see the model's predictions.",
+    flagging_mode="never"
 )
 
 if __name__ == "__main__":

requirements.txt

@@ -1,7 +1,6 @@
 digitnet @ git+https://github.com/karanravindra/digitnet@main
 
 # This file was autogenerated via `uv export`.
--e .
 aiofiles==23.2.1
 aiohappyeyeballs==2.4.3
 aiohttp==3.10.10