debug

app.py

@@ -2,6 +2,7 @@ from typing import List
 
 import gradio as gr
 import numpy as np
+import torch
 from transformers import CLIPProcessor, CLIPModel
 
 IMAGENET_CLASSES_FILE = "imagenet-classes.txt"
@@ -23,14 +24,15 @@ def load_text_lines(file_path: str) -> List[str]:
         return [line.rstrip() for line in lines]
 
 
-model = CLIPModel.from_pretrained("facebook/metaclip-b32-400m")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = CLIPModel.from_pretrained("facebook/metaclip-b32-400m").to(device)
 processor = CLIPProcessor.from_pretrained("facebook/metaclip-b32-400m")
 imagenet_classes = load_text_lines(IMAGENET_CLASSES_FILE)
 
 
 def classify_image(input_image) -> str:
     inputs = processor(
-        text=['dog', 'person'],
+        text=imagenet_classes,
         images=input_image,
         return_tensors="pt",
         padding=True)