Update

app.py

@@ -1,5 +1,9 @@
 import gradio as gr
 from transformers import pipeline
+import torch
+
+# Check if CUDA is available
+device = 0 if torch.cuda.is_available() else -1
 
 model_names = [
     "apple/mobilevit-small",
@@ -15,11 +19,15 @@ model_names = [
     "shi-labs/dinat-base-in1k-224",
 ]
 
+# Cache for pipelines to avoid reloading models
+pipelines = {}
 
 def process(image_file, top_k):
     labels = []
     for m in model_names:
-        p = pipeline("image-classification", model=m)
+        if m not in pipelines:
+            pipelines[m] = pipeline("image-classification", model=m, device=device)
+        p = pipelines[m]
         pred = p(image_file)
         labels.append({x["label"]: x["score"] for x in pred[:top_k]})
     return labels
@@ -32,7 +40,11 @@ top_k = gr.Slider(minimum=1, maximum=5, step=1, value=5, label="Top k classes")
 # Output
 labels = [gr.Label(label=m) for m in model_names]
 
-description = "This Space lets you quickly compare the most popular image classifiers available on the hub, including the recent NAT and DINAT models. All of them have been fine-tuned on the ImageNet-1k dataset. Anecdotally, the three sample images have been generated with a Stable Diffusion model :)"
+description = (
+    "This Space compares popular image classifiers available on the Hugging Face hub, "
+    "including NAT and DINAT models. All models have been fine-tuned on ImageNet-1k. "
+    "The sample images were generated with Stable Diffusion."
+)
 
 iface = gr.Interface(
     theme="huggingface",

requirements.txt

@@ -1,3 +1,6 @@
-torch
-transformers
-natten
+torch==2.7.0
+transformers==4.51.3
+gradio>=3.50.0
+--no-build-isolation
+natten==0.17.5
+Pillow>=9.0.0