Spaces:
Running on Zero
Running on Zero
File size: 5,153 Bytes
5570f76 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | """Qwen2.5-VL Smart Document, Blueprint & Invoice OCR Studio on ZeroGPU."""
import os
os.environ.setdefault("PYTORCH_ALLOC_CONF", "backend:cudaMallocAsync")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")
import spaces # Must be imported before torch
import gradio as gr
import PIL.Image
import torch
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
MODEL_ID = "Qwen/Qwen2.5-VL-3B-Instruct"
EXTRACTION_PRESETS = {
"Full Clean Markdown + Tables (Default)": (
"Extract all text, headings, and tables from this image into clean, well-formatted GitHub Flavored Markdown. "
"Preserve exact numbers, part numbers, prices, and table columns."
),
"Structured Invoice / Receipt JSON": (
"Extract all fields from this invoice or receipt as valid JSON with keys: vendor, date, invoice_number, "
"line_items (description, quantity, unit_price, total), subtotal, tax, and total_amount."
),
"Electrical Panel / Blueprint / Spec Analysis": (
"Analyze this electrical diagram, panel schedule, or technical blueprint. List all visible circuits, "
"breaker ratings, wire gauges, load calculations, labels, and technical notes in structured Markdown."
),
"Custom Question / Prompt": "",
}
print(f"[vision-ocr-studio] Preloading {MODEL_ID} at module scope for ZeroGPU...", flush=True)
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
MODEL_ID,
torch_dtype=torch.bfloat16,
)
model.to("cuda")
model.eval()
processor = AutoProcessor.from_pretrained(MODEL_ID)
print("[vision-ocr-studio] Qwen2.5-VL OCR engine ready on ZeroGPU!", flush=True)
@spaces.GPU(duration=60)
def analyze_document(
image_path: str | None,
preset: str = "Full Clean Markdown + Tables (Default)",
custom_instruction: str = "",
max_tokens: int = 1024,
) -> str:
"""Run vision-language OCR or structured extraction on an uploaded document/image."""
if not image_path:
raise gr.Error("Please upload a document, blueprint, invoice, or screenshot image first.")
instruction = (custom_instruction or "").strip() or EXTRACTION_PRESETS.get(
preset, EXTRACTION_PRESETS["Full Clean Markdown + Tables (Default)"]
)
img = PIL.Image.open(image_path).convert("RGB")
if max(img.size) > 1600:
scale = 1600.0 / max(img.size)
img = img.resize((int(img.width * scale), int(img.height * scale)), PIL.Image.Resampling.LANCZOS)
messages = [
{
"role": "user",
"content": [
{"type": "image", "image": img},
{"type": "text", "text": instruction},
],
}
]
text_prompt = processor.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
inputs = processor(
text=[text_prompt],
images=[img],
padding=True,
return_tensors="pt",
).to("cuda")
with torch.no_grad():
generated_ids = model.generate(**inputs, max_new_tokens=int(max_tokens))
trimmed_ids = [
out_ids[len(in_ids):]
for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
]
output_text = processor.batch_decode(
trimmed_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)[0]
return output_text.strip()
with gr.Blocks(title="Smart Document, Blueprint & Invoice OCR Studio", fill_width=True) as demo:
gr.Markdown(
"# 📄 Smart Document, Blueprint & Invoice OCR Studio\n"
"Powered by **[Qwen/Qwen2.5-VL-3B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct)** on **ZeroGPU (`zero-a10g`)** at **$0/hour**—extract clean **Markdown, JSON, tables, invoices, or electrical specs** from any image."
)
with gr.Row():
with gr.Column(scale=1):
img_in = gr.Image(
label="Upload Document, Invoice, Electrical Schedule, Blueprint, or Screenshot",
type="filepath",
)
preset_dd = gr.Dropdown(
choices=list(EXTRACTION_PRESETS),
value="Full Clean Markdown + Tables (Default)",
label="Extraction Mode",
)
custom_box = gr.Textbox(
label="Custom Instruction / Question (Optional — overrides preset when filled)",
lines=3,
placeholder="e.g., Extract all circuit breaker amps and wire sizes into a JSON array...",
)
max_tok = gr.Slider(256, 2048, value=1024, step=128, label="Max Output Tokens")
extract_btn = gr.Button("📄 Extract & Analyze Document (ZeroGPU)", variant="primary")
with gr.Column(scale=1):
md_out = gr.Textbox(
label="Extracted Output (Markdown / JSON)",
lines=18,
buttons=["copy"],
)
extract_btn.click(
fn=analyze_document,
inputs=[img_in, preset_dd, custom_box, max_tok],
outputs=md_out,
api_name="extract_document",
)
if __name__ == "__main__":
demo.launch(theme=gr.themes.Ocean(), mcp_server=True)
|