fix jit 2

app.py

@@ -8,15 +8,16 @@ import faiss
 
 
 # Init similarity search AI model and processor
-torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-dino_v2_model = AutoModel.from_pretrained("./dinov2-large").to(torch_device)
-dino_v2_image_processor = AutoImageProcessor.from_pretrained("./dinov2-large")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+processor = AutoImageProcessor.from_pretrained("./dinov2-large")
+model = AutoModel.from_pretrained("./dinov2-large")
+model.config.return_dict = False  # Set return_dict to False for JIT tracing
+model.to(device)
 
-# Provide a sample input for tracing
-sample_input = dino_v2_image_processor(
-    images=Image.new("RGB", (224, 224)), return_tensors="pt"
-).to(torch_device)
-traced_dino_v2_model = torch.jit.trace(dino_v2_model, sample_input["pixel_values"])
+# Prepare an example input for tracing
+example_input = torch.rand(1, 3, 224, 224).to(device)  # Adjust size if needed
+traced_model = torch.jit.trace(model, example_input)
+traced_model = traced_model.to(device)
 
 
 def process_image(image):
@@ -49,12 +50,10 @@ def process_image(image):
 
     # Extract the features from the uploaded image
     with torch.no_grad():
-        inputs = dino_v2_image_processor(images=image, return_tensors="pt").to(
-            torch_device
-        )
+        inputs = processor(images=image, return_tensors="pt").to(device)
 
         # Use the traced model for inference
-        outputs = traced_dino_v2_model(**inputs)
+        outputs = traced_model(**inputs)
 
     # Normalize the features before search, whatever that means
     embeddings = outputs.last_hidden_state