Update app.py

app.py

@@ -33,29 +33,45 @@ def model_inference(
     if isinstance(images, Image.Image):
         images = [images]
 
-    if isinstance(text, str):
-        text = "<image>" + text
-        text = [text]
 
-    inputs = processor(text=text, images=images, padding=True, return_tensors="pt").to("cuda")
+    resulting_messages = [
+                {
+                    "role": "user",
+                    "content": [{"type": "image"}] + [
+                        {"type": "text", "text": text}
+                    ]
+                }
+            ]
+
+
+    prompt = processor.apply_chat_template(resulting_messages, add_generation_prompt=True)
+    inputs = processor(text=prompt, images=[images], return_tensors="pt")
+    inputs = {k: v.to("cuda") for k, v in inputs.items()}
+
+    generation_args = {
+        "max_new_tokens": max_new_tokens,
+        "repetition_penalty": repetition_penalty,
+
+    }
 
     assert decoding_strategy in [
         "Greedy",
         "Top P Sampling",
     ]
     if decoding_strategy == "Greedy":
-        do_sample = False
+        generation_args["do_sample"] = False
     elif decoding_strategy == "Top P Sampling":
-        do_sample = True
+        generation_args["temperature"] = temperature
+        generation_args["do_sample"] = True
+        generation_args["top_p"] = top_p
+
+
+    generation_args.update(inputs)
 
     # Generate
-    
-    generated_ids = model.generate(**inputs, bad_words_ids=BAD_WORDS_IDS, max_new_tokens=max_new_tokens,
-                                temperature=temperature, do_sample=do_sample, repetition_penalty=repetition_penalty,
-                                top_p=top_p),
-    #generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)
-    generated_texts = processor.batch_decode(generated_ids[:, inputs["input_ids"].size(1):], skip_special_tokens=True)
-    print("INPUT:", text, "|OUTPUT:", generated_texts)
+    generated_ids = model.generate(**generation_args)
+
+    generated_texts = processor.batch_decode(generated_ids[:, generation_args["input_ids"].size(1):], skip_special_tokens=True)
     return generated_texts[0]