Update app.py

app.py

@@ -64,85 +64,35 @@
 # demo.launch()
 
 import gradio as gr
-from transformers import AutoProcessor, AutoModelForVision2Seq, TextIteratorStreamer
-from transformers.image_utils import load_image
-from threading import Thread
-import torch
-import html
-import re
-from PIL import Image, ImageOps
+from transformers import AutoProcessor, AutoModelForImageTextToText
+from PIL import Image
 
 # Load model & processor once at startup
 processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
-model = AutoModelForVision2Seq.from_pretrained("ds4sd/SmolDocling-256M-preview")
-
-def add_random_padding(image, min_percent=0.1, max_percent=0.10):
-    image = image.convert("RGB")
-    width, height = image.size
-    pad_w_percent = random.uniform(min_percent, max_percent)
-    pad_h_percent = random.uniform(min_percent, max_percent)
-    pad_w = int(width * pad_w_percent)
-    pad_h = int(height * pad_h_percent)
-    corner_pixel = image.getpixel((0, 0))  # Top-left corner
-    padded_image = ImageOps.expand(image, border=(pad_w, pad_h, pad_w, pad_h), fill=corner_pixel)
-    return padded_image
-
-def extract_table(image_file):
-    # Load image
-    image = load_image(image_file)
-    
-    # Optionally add padding if needed for model robustness (optional)
-    image = add_random_padding(image)
-
-    # Fixed prompt to extract table only (modify if needed)
-    text = "Convert this table to OTSL."
-
-    # Build the message structure for processor
-    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=[image], return_tensors="pt").to('cuda')
-
-    generation_args = {
-        "input_ids": inputs.input_ids,
-        "pixel_values": inputs.pixel_values,
-        "attention_mask": inputs.attention_mask,
-        "max_new_tokens": 8192,
-        "num_return_sequences": 1,
-    }
-
-    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-    generation_args["streamer"] = streamer
-
-    thread = Thread(target=model.generate, kwargs=generation_args)
-    thread.start()
-
-    output_text = ""
-    for new_text in streamer:
-        output_text += new_text
-
-    # Clean and return output
-    cleaned_output = output_text.replace("<end_of_utterance>", "").strip()
-    
-    # Optionally convert <chart> tags to <otsl> if present
-    if "<chart>" in cleaned_output:
-        cleaned_output = cleaned_output.replace("<chart>", "<otsl>").replace("</chart>", "</otsl>")
-        cleaned_output = re.sub(r'(<loc_500>)(?!.*<loc_500>)<[^>]+>', r'\1', cleaned_output)
-
-    return cleaned_output or "No table found or unable to extract."
+model = AutoModelForImageTextToText.from_pretrained("ds4sd/SmolDocling-256M-preview")
+
+def smoldocling_readimage(image, prompt_text):
+    messages = [
+        {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt_text}]}
+    ]
+    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
+    inputs = processor(text=prompt, images=[image], return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=1024)
+    prompt_length = inputs.input_ids.shape[1]
+    generated = outputs[:, prompt_length:]
+    result = processor.batch_decode(generated, skip_special_tokens=False)[0]
+    return result.replace("<end_of_utterance>", "").strip()
 
 # Gradio UI
 demo = gr.Interface(
-    fn=extract_table,
-    inputs=gr.Image(type="filepath", label="Upload Table Image"),
-    outputs=gr.Textbox(label="Extracted Table (OTSL Format)"),
-    title="Table Extraction from Image using SmolDocling-256M",
-    description="Upload an image containing a table. The model will extract the table and output it in OTSL format."
+    fn=smoldocling_readimage,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Textbox(lines=1, placeholder="Enter prompt (e.g. Convert to docling)", label="Prompt"),
+    ],
+    outputs="ostl",
+    title="SmolDocling Web App",
+    description="Upload a document image and convert it to structured docling format."
 )
 
-demo.launch(debug=True)
-
-
+demo.launch()