[OCR API] Added initial API code.

Dockerfile

@@ -1,16 +1,36 @@
-# Read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
-# you will also find guides on how best to write your Dockerfile
+# Use Python 3.12 as the base image
+FROM python:3.12
 
-FROM python:3.9
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    ffmpeg \
+    git \
+    && rm -rf /var/lib/apt/lists/*
 
+# Create a non-root user
 RUN useradd -m -u 1000 user
-USER user
-ENV PATH="/home/user/.local/bin:$PATH"
-
 WORKDIR /app
 
-COPY --chown=user ./requirements.txt requirements.txt
-RUN pip install --no-cache-dir --upgrade -r requirements.txt
+# Install Python dependencies directly
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir \
+        torch \
+        torchvision \
+        git+https://github.com/huggingface/transformers \
+        accelerate \
+        qwen-vl-utils[decord]==0.0.8 \
+        fastapi \
+        uvicorn[standard]
 
+# Copy application files
 COPY --chown=user . /app
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+
+# Switch to the non-root user
+USER user
+
+# Set environment variables
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+# Command to run the application
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,7 +1,48 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Query
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+from qwen_vl_utils import process_vision_info
+import torch
 
 app = FastAPI()
 
+checkpoint = "Qwen/Qwen2.5-VL-3B-Instruct"
+min_pixels = 256*28*28
+max_pixels = 1280*28*28
+processor = AutoProcessor.from_pretrained(
+    checkpoint,
+    min_pixels=min_pixels,
+    max_pixels=max_pixels
+)
+model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    checkpoint,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    # attn_implementation="flash_attention_2",
+)
+
 @app.get("/")
-def greet_json():
-    return {"Hello": "World!"}
+def read_root():
+    return {"message": "API is live. Use the /predict endpoint."}
+
+@app.get("/predict")
+def predict(image_url: str = Query(...), prompt: str = Query(...)):
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant with vision abilities."},
+        {"role": "user", "content": [{"type": "image", "image": image_url}, {"type": "text", "text": prompt}]},
+    ]
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
+        return_tensors="pt",
+    ).to(model.device)
+    with torch.no_grad():
+        generated_ids = model.generate(**inputs, max_new_tokens=128)
+    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+    output_texts = processor.batch_decode(
+        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )
+    return {"response": output_texts[0]}