Update app.py

app.py

@@ -10,6 +10,7 @@ import torch
 from transformers import Qwen2VLForConditionalGeneration, GenerationConfig, AutoProcessor
 import spaces
 
+from vllm import LLM, SamplingParams
 
 def extract_answer_content(text: str) -> str:
     """
@@ -62,6 +63,10 @@ SYSTEM_PROMPT = (
 processor = AutoProcessor.from_pretrained("JosephZ/qwen2vl-7b-sft-grpo-close-sgg", max_pixels=1024*28*28)
 
 device='cuda' if torch.cuda.is_available() else "cpu"
+model_name = "JosephZ/qwen2vl-7b-sft-grpo-close-sgg"
+
+
+"""
 model = Qwen2VLForConditionalGeneration.from_pretrained("JosephZ/qwen2vl-7b-sft-grpo-close-sgg",
              torch_dtype=torch.bfloat16,
              device_map=device)
@@ -75,9 +80,25 @@ generation_config=GenerationConfig(
         max_new_tokens=2048,
         use_cache=True 
 )
+"""
+model = LLM(
+            model=model_name, 
+            limit_mm_per_prompt={"image": 1},
+            dtype='bfloat16',
+            device=device,
+            max_model_len=4096,
+            mm_processor_kwargs= { "max_pixels": 1024*28*28, "min_pixels": 4*28*28},
+)
+sampling_params = SamplingParams(
+        temperature=0.01,
+        top_k=1,
+        top_p=0.001,
+        repetition_penalty=1.0,
+        max_tokens=2048,
+)
 
 def build_prompt(image, user_text):
-    #base64_image = encode_image_to_base64(image)
+    base64_image = encode_image_to_base64(image)
     messages = [
         {
             "role": "system",
@@ -86,8 +107,8 @@ def build_prompt(image, user_text):
         {
             "role": "user",
             "content": [
-                #{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}},
-                 {"type": "image"},
+                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}},
+                # {"type": "image"},
                 {"type": "text", "text": user_text},
             ],
         },
@@ -161,24 +182,14 @@ def generate_sgg(image):
     
     iw, ih = image.size
     scale_factors = (iw / 1000.0, ih / 1000.0)
-
+    
     conversation = build_prompt(image, PROMPT_CLOSE)
-    text_prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
 
-    inputs = processor(
-        text=[text_prompt], images=[image], padding=True, return_tensors="pt"
-    )
-    inputs = inputs.to(model.device)
-    
-    output_ids = model.generate(**inputs, generation_config=generation_config)
-    generated_ids = [
-        output_ids[len(input_ids) :]
-        for input_ids, output_ids in zip(inputs.input_ids, output_ids)
-    ]
-    output_text = processor.batch_decode(
-        generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
-    )[0]    
+    with torch.no_grad():
+        outputs = model.chat([conversation], sampling_params=sampling_params)
+        output_texts = [output.outputs[0].text for output in outputs] 
 
+    output_text = output_texts[0]
     resp = extract_answer_content(output_text)
     
     try: