Update simple_test.py

simple_test.py

@@ -1,42 +1,33 @@
+import gradio as gr
 from transformers import AutoProcessor, AutoModelForImageTextToText
 from PIL import Image
-import torch
 
+# Load model & processor once at startup
+processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
+model = AutoModelForImageTextToText.from_pretrained("ds4sd/SmolDocling-256M-preview")
 
-
-def smoldocling_readimage(image, prompt_text="Convert this page to docling."):
-    # Load model and processor
-    processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
-    model = AutoModelForImageTextToText.from_pretrained("ds4sd/SmolDocling-256M-preview")
-
-    # Create input messages
+def smoldocling_readimage(image, prompt_text):
     messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image"},
-                {"type": "text", "text": prompt_text}
-            ]
-        },
+        {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt_text}]}
     ]
-
-
-    # Prepare inputs
     prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
     inputs = processor(text=prompt, images=[image], return_tensors="pt")
-    # inputs = inputs.to(device)
-
-
-    # Generate outputs
-    generated_ids = model.generate(**inputs, max_new_tokens=1024)  # Reduced for testing
+    outputs = model.generate(**inputs, max_new_tokens=1024)
     prompt_length = inputs.input_ids.shape[1]
-    trimmed_generated_ids = generated_ids[:, prompt_length:]
-    doctags = processor.batch_decode(
-        trimmed_generated_ids,
-        skip_special_tokens=False,
-    )[0].lstrip()
-
-    # Clean the output
-    doctags = doctags.replace("<end_of_utterance>", "").strip()
-        
-    return doctags
+    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=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="text",
+    title="SmolDocling Web App",
+    description="Upload a document image and convert it to structured docling format."
+)
+
+demo.launch()