Update app.py

app.py

@@ -1,45 +1,63 @@
-import gradio as gr
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+from transformers import MllamaForConditionalGeneration, AutoProcessor
 from PIL import Image
 import torch
+import gradio as gr
 import spaces
 
-# Load TrOCR model
-processor = TrOCRProcessor.from_pretrained("microsoft/trocr-large-handwritten")
-model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-large-handwritten")
+# Initialize model and processor
+ocr = "unsloth/Llama-3.2-11B-Vision-Instruct"
+model = MllamaForConditionalGeneration.from_pretrained(
+    ocr,
+    torch_dtype=torch.bfloat16
+).to("cuda")
+processor = AutoProcessor.from_pretrained(ocr)
 
 @spaces.GPU
-def recognize_text(image):
-    try:
-        # Convert image to RGB
-        image = image.convert("RGB")
-        print("Image converted to RGB.")
-
-        # Preprocess the image using the processor
-        pixel_values = processor(image, return_tensors="pt").pixel_values
-        print("Image preprocessed. Pixel values shape:", pixel_values.shape)
-
-        # Generate text from the image
-        with torch.no_grad():
-            generated_ids = model.generate(pixel_values)
-            print("Generated IDs:", generated_ids)
-
-        # Decode the generated IDs to text
-        text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        print("Decoded text:", text)
-
-        return text
-    except Exception as e:
-        print(f"Error: {str(e)}")
-        return f"Error: {str(e)}"
-
-# Gradio UI
+def extract_text(image):
+    # Convert image to RGB
+    image = Image.open(image).convert("RGB")
+    
+    # Create message structure
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Extract handwritten text from the image and output only the extracted text without any additional description or commentary in output"},
+                {"type": "image"}
+            ]
+        }
+    ]
+    
+    # Process input
+    texts = processor.apply_chat_template(messages, add_generation_prompt=True)
+    inputs = processor(text=texts, images=[image], return_tensors="pt").to("cuda")
+
+    
+    # Generate output
+    outputs = model.generate(**inputs, max_new_tokens=250)
+    result = processor.decode(outputs[0], skip_special_tokens=True)
+
+    print(result)
+    
+    # Clean up the output to remove the prompt and assistant text
+    if "assistant" in result.lower():
+        result = result[result.lower().find("assistant") + len("assistant"):].strip()
+    
+    # Remove any remaining conversation markers
+    result = result.replace("user", "").replace("Extract handwritten text from the image and output only the extracted text without any additional description or commentary in output", "").strip()
+
+    print(result)
+    
+    return result
+
+# Create Gradio interface
 note = gr.Interface(
-    fn=recognize_text,
-    inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="Handwritten Note to Digital Text",
-    description="Upload an image of handwritten text, and the AI will convert it to digital text."
+    fn=extract_text,
+    inputs=gr.Image(type="filepath", label="Upload Image"),
+    outputs=gr.Textbox(label="Extracted Text"),
+    title="Handwritten Text Extractor",
+    description="Upload an image containing handwritten text to extract its content.",
 )
 
-note.launch()
+# Launch the app
+note.launch(debug=True)