Update main.py

main.py

@@ -25,26 +25,27 @@ def read_root():
     return {"message": "API is live. Use the /predict endpoint."}
 
 @app.get("/predict")
-def predict():
-    
-    return {"got query"}
-    # messages = [
-    #     {"role": "system", "content": "You are a helpful assistant with vision abilities."},
-    #     {"role": "user", "content": [{"type": "image", "image": image_url}, {"type": "text", "text": prompt}]},
-    # ]
-    # text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    # image_inputs, video_inputs = process_vision_info(messages)
-    # inputs = processor(
-    #     text=[text],
-    #     images=image_inputs,
-    #     videos=video_inputs,
-    #     padding=True,
-    #     return_tensors="pt",
-    # ).to(model.device)
-    # with torch.no_grad():
-    #     generated_ids = model.generate(**inputs, max_new_tokens=128)
-    # generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
-    # output_texts = processor.batch_decode(
-    #     generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-    # )
-    # return {"response": output_texts[0]}
+def predict(
+    image_url: str = Query(..., description="URL of the image"),
+    prompt: str = Query(..., description="Prompt for the image")
+):
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant with vision abilities."},
+        {"role": "user", "content": [{"type": "image", "image": image_url}, {"type": "text", "text": prompt}]},
+    ]
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
+        return_tensors="pt",
+    ).to(model.device)
+    with torch.no_grad():
+        generated_ids = model.generate(**inputs, max_new_tokens=128)
+    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+    output_texts = processor.batch_decode(
+        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )
+    return {"response": output_texts[0]}