Update app.py

app.py

@@ -39,29 +39,27 @@ def fig_to_pil(fig):
     return Image.open(buf)
 
 @spaces.GPU
-def run_example(task_prompt, image, text_input=None, model_id='J-LAB/Florence-Idesire'):
-    model = models[model_id]
-    processor = processors[model_id]
-    if text_input is None:
-        prompt = task_prompt
+def process_image(image, task_prompt, text_input=None, model_id='J-LAB/Florence_2_B_FluxiAI_Product_Caption'):
+    image = Image.fromarray(image)  # Convert NumPy array to PIL Image
+    if task_prompt == 'Product Caption':
+        task_prompt = '<PC>'
+        results = run_example(task_prompt, image, model_id=model_id)
+    elif task_prompt == 'More Detailed Caption':
+        task_prompt = '<MORE_DETAILED_CAPTION>'
+        results = run_example(task_prompt, image, model_id=model_id)
+    else:
+        return "", None  # Return empty string and None for unknown task prompts
+
+    # Remove the key and get the text value
+    if results and task_prompt in results:
+        output_text = results[task_prompt]
     else:
-        prompt = task_prompt + text_input
-    inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda")
-    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]
-    parsed_answer = processor.post_process_generation(
-        generated_text,
-        task=task_prompt,
-        image_size=(image.width, image.height)
-    )
-    return parsed_answer
+        output_text = ""
+
+    # Convert newline characters to HTML line breaks
+    output_text = output_text.replace("\n\n", "<br><br>").replace("\n", "<br>")
+
+    return output_text, None
 
 def plot_bbox(image, data):
     fig, ax = plt.subplots()
@@ -117,27 +115,7 @@ def draw_ocr_bboxes(image, prediction):
                   fill=color)
     return image
 
-def process_image(image, task_prompt, text_input=None, model_id='J-LAB/Florence_2_B_FluxiAI_Product_Caption'):
-    image = Image.fromarray(image)  # Convert NumPy array to PIL Image
-    if task_prompt == 'Product Caption':
-        task_prompt = '<PC>'
-        results = run_example(task_prompt, image, model_id=model_id)
-    elif task_prompt == 'More Detailed Caption':
-        task_prompt = '<MORE_DETAILED_CAPTION>'
-        results = run_example(task_prompt, image, model_id=model_id)
-    else:
-        return "", None  # Return empty string and None for unknown task prompts
 
-    # Remove the key and get the text value
-    if results and task_prompt in results:
-        output_text = results[task_prompt]
-    else:
-        output_text = ""
-
-    # Convert newline characters to HTML line breaks
-    output_text = output_text.replace("\n\n", "<br><br>").replace("\n", "<br>")
-
-    return output_text, None
 
 
 css = """