app.py

import gradio as gr
import torch
from transformers import AutoModel, AutoTokenizer
from PIL import Image
import os
import spaces
import tempfile

# Set CUDA device
os.environ["CUDA_VISIBLE_DEVICES"] = '0'

# Load model and tokenizer
model_name = "deepseek-ai/DeepSeek-OCR"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(
    model_name,
    _attn_implementation="flash_attention_2",
    trust_remote_code=True,
    use_safetensors=True,
)
model = model.eval()


@spaces.GPU(duration=120)
def ocr_process(
    image_input: Image.Image,
    task_type: str = "ocr",
    preset: str = "gundam",
) -> str:
    """
    Process image and extract text using DeepSeek-OCR model.

    Args:
        image_input: Input image
        task_type: Type of task - "ocr" for text extraction or "markdown" for document conversion
        preset: Preset configuration for model parameters

    Returns:
        Extracted text or markdown content
    """
    if image_input is None:
        return "Please upload an image first."

    # Move model to GPU and set dtype
    model.cuda().to(torch.bfloat16)
    
    # Create temp directory for this session
    with tempfile.TemporaryDirectory() as temp_dir:
        # Save image with proper format
        temp_image_path = os.path.join(temp_dir, "input_image.jpg")
        # Convert RGBA to RGB if necessary
        if image_input.mode in ('RGBA', 'LA', 'P'):
            rgb_image = Image.new('RGB', image_input.size, (255, 255, 255))
            # Handle different image modes
            if image_input.mode == 'RGBA':
                rgb_image.paste(image_input, mask=image_input.split()[3])
            else:
                rgb_image.paste(image_input)
            rgb_image.save(temp_image_path, 'JPEG', quality=95)
        else:
            image_input.save(temp_image_path, 'JPEG', quality=95)
        
        # Set parameters based on preset
        presets = {
            "tiny": {"base_size": 512, "image_size": 512, "crop_mode": False},
            "small": {"base_size": 640, "image_size": 640, "crop_mode": False},
            "base": {"base_size": 1024, "image_size": 1024, "crop_mode": False},
            "large": {"base_size": 1280, "image_size": 1280, "crop_mode": False},
            "gundam": {"base_size": 1024, "image_size": 640, "crop_mode": True},
        }
        
        config = presets[preset]

        # Set prompt based on task type
        if task_type == "markdown":
            prompt = "<image>\n<|grounding|>Convert the document to markdown. "
        else:
            prompt = "<image>\nFree OCR. "

        # Run inference - return the result directly
        result = model.infer(
            tokenizer,
            prompt=prompt,
            image_file=temp_image_path,
            output_path=temp_dir,
            base_size=config["base_size"],
            image_size=config["image_size"],
            crop_mode=config["crop_mode"],
            save_results=False,
            test_compress=False,
        )

    # Move model back to CPU to free GPU memory
    model.to("cpu")
    torch.cuda.empty_cache()

    # Return the result directly - the model returns the extracted text
    return result


# Create Gradio interface
with gr.Blocks(title="DeepSeek OCR", theme=gr.themes.Soft()) as demo:
    gr.HTML(
        """
        <div style="text-align: center; margin-bottom: 20px;">
            <h1>🔍 DeepSeek OCR</h1>
            <p>Extract text and convert documents to markdown using DeepSeek-OCR</p>
            <p>Built with <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: #0066cc; text-decoration: none;">anycoder</a></p>
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Upload Image")
            image_input = gr.Image(
                label="Input Image",
                type="pil",
                sources=["upload", "webcam", "clipboard"],
                height=300,
            )

            gr.Markdown("### ⚙️ Settings")
            task_type = gr.Radio(
                choices=["ocr", "markdown"],
                value="ocr",
                label="Task Type",
                info="OCR: Extract plain text | Markdown: Convert to formatted markdown",
            )

            preset = gr.Radio(
                choices=["gundam", "base", "large", "small", "tiny"],
                value="gundam",
                label="Model Preset",
                info="Start with 'gundam' - it's optimized for most documents",
            )

            with gr.Accordion("ℹ️ Preset Details", open=False):
                gr.Markdown("""
                - **Gundam** (Recommended): Balanced performance with crop mode
                - **Base**: Standard quality without cropping
                - **Large**: Highest quality for complex documents
                - **Small**: Faster processing, good for simple text
                - **Tiny**: Fastest, suitable for clear printed text
                """)

            submit_btn = gr.Button("🚀 Extract Text", variant="primary", size="lg")
            clear_btn = gr.ClearButton([image_input], value="🗑️ Clear")

        with gr.Column(scale=1):
            gr.Markdown("### 📝 Extracted Text")
            output_text = gr.Textbox(
                label="Output",
                lines=15,
                max_lines=30,
                interactive=False,
                placeholder="Extracted text will appear here...",
                show_copy_button=True,
            )

    # Event handlers
    submit_btn.click(
        fn=ocr_process,
        inputs=[image_input, task_type, preset],
        outputs=output_text,
    )

    # Example section with receipt image
    gr.Markdown("### 📚 Example")
    gr.Examples(
        examples=[
            ["https://upload.wikimedia.org/wikipedia/commons/thumb/0/0b/ReceiptSwiss.jpg/800px-ReceiptSwiss.jpg", "ocr", "gundam"],
        ],
        inputs=[image_input, task_type, preset],
        label="Try this receipt example",
    )

    gr.Markdown("""
    ### 💡 Tips for Best Results
    - **For receipts**: Use "ocr" mode with "gundam" or "base" preset
    - **For documents with tables**: Use "markdown" mode with "large" preset
    - **If text is not detected**: Try different presets in this order: gundam → base → large
    - **For handwritten text**: Use "large" preset for better accuracy
    - Ensure images are clear and well-lit for optimal results
    """)


if __name__ == "__main__":
    demo.launch(share=False)