Update app.py

app.py

@@ -42,13 +42,36 @@ def initialize_model():
 # Initialize the model and processor
 model, processor, device = initialize_model()
 
-def run_example(task_prompt, image, text_input=None):
-    prompt = task_prompt if text_input is None else task_prompt + text_input
-    inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
-    with torch.inference_mode():
-        generated_ids = model.generate(**inputs, max_new_tokens=1024, early_stopping=False, do_sample=False, num_beams=3)
+# def run_example(task_prompt, image, text_input=None):
+#     prompt = task_prompt if text_input is None else task_prompt + text_input
+#     inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
+#     with torch.inference_mode():
+#         generated_ids = model.generate(**inputs, max_new_tokens=1024, early_stopping=False, do_sample=False, num_beams=3)
+#     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+#     return processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.size[0], image.size[1]))
+
+def run_example(task_prompt, text_input=None):
+    if text_input is None:
+        prompt = task_prompt
+    else:
+        prompt = task_prompt + text_input
+    inputs = processor(text=prompt, images=image, return_tensors="pt")
+    generated_ids = model.generate(
+      input_ids=inputs["input_ids"],
+      pixel_values=inputs["pixel_values"],
+      max_new_tokens=1024,
+      early_stopping=False,
+      do_sample=False,
+      num_beams=3,
+    )
     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-    return processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.size[0], image.size[1]))
+    parsed_answer = processor.post_process_generation(
+        generated_text,
+        task=task_prompt,
+        image_size=(image.width, image.height)
+    )
+
+    return parsed_answer
 
 def fig_to_pil(fig):
     buf = io.BytesIO()