refactor

app.py

@@ -1,13 +1,16 @@
 import spaces
 import gradio as gr
 
-from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
+from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
 from qwen_vl_utils import process_vision_info
 import torch
 import os
 import json
+from pydantic import BaseModel
+from typing import Tuple
 
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+
 model = Qwen2VLForConditionalGeneration.from_pretrained(
     "Qwen/Qwen2-VL-7B-Instruct",
     torch_dtype=torch.bfloat16,
@@ -15,8 +18,7 @@ model = Qwen2VLForConditionalGeneration.from_pretrained(
     device_map="auto",
 )
 processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
-from pydantic import BaseModel
-from typing import Tuple
+
 
 class GeneralRetrievalQuery(BaseModel):
     broad_topical_query: str
@@ -26,6 +28,7 @@ class GeneralRetrievalQuery(BaseModel):
     visual_element_query: str
     visual_element_explanation: str
 
+
 def get_retrieval_prompt(prompt_name: str) -> Tuple[str, GeneralRetrievalQuery]:
     if prompt_name != "general":
         raise ValueError("Only 'general' prompt is available in this version")
@@ -66,11 +69,11 @@ Generate the queries based on this image and provide the response in the specifi
     return prompt, GeneralRetrievalQuery
 
 
-
+# defined like this so we can later add more prompting options
 prompt, pydantic_model = get_retrieval_prompt("general")
 
-@spaces.GPU
-def generate_response(image):
+
+def _prep_data_for_input(image):
     messages = [
         {
             "role": "user",
@@ -97,6 +100,12 @@ def generate_response(image):
         padding=True,
         return_tensors="pt",
     )
+    return inputs
+
+
+@spaces.GPU
+def generate_response(image):
+    inputs = _prep_data_for_input(image)
     inputs = inputs.to("cuda")
 
     generated_ids = model.generate(**inputs, max_new_tokens=200)
@@ -116,5 +125,6 @@ def generate_response(image):
     except Exception:
         return {}
 
-demo = gr.Interface(fn=generate_response, inputs=gr.Image(type='pil'), outputs="json")
+
+demo = gr.Interface(fn=generate_response, inputs=gr.Image(type="pil"), outputs="json")
 demo.launch()