Update app.py

app.py

@@ -1,12 +1,10 @@
 # =========================
-# CAMEL-DOC-OCR (FAST)
-# Single-file version
+# CAMEL-DOC-OCR (HF Spaces SAFE)
+# Single-file – NO CUDA init at global scope
 # =========================
 
 import os
 import gc
-import json
-import re
 import torch
 import fitz
 import gradio as gr
@@ -23,22 +21,18 @@ from transformers.models.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
 MODEL_ID = "prithivMLmods/Camel-Doc-OCR-062825"
 DPI = 150
 MAX_IMAGE_SIZE = 2048
-GPU_MEMORY_FRACTION = 0.8
 
 
 # =========================
-# TORCH OPTIMIZATION
+# TORCH FLAGS (SAFE FOR SPACES)
 # =========================
 torch.set_grad_enabled(False)
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
 
-if torch.cuda.is_available():
-    torch.cuda.set_per_process_memory_fraction(GPU_MEMORY_FRACTION, device=0)
-
 
 # =========================
-# LOAD MODEL (ONCE)
+# LOAD MODEL (NO CUDA INIT HERE)
 # =========================
 bnb = BitsAndBytesConfig(
     load_in_4bit=True,
@@ -55,7 +49,7 @@ processor = AutoProcessor.from_pretrained(
 model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     MODEL_ID,
     quantization_config=bnb,
-    device_map="auto",
+    device_map="auto",            # HF Spaces will inject GPU here
     torch_dtype=torch.float16,
     trust_remote_code=True
 ).eval()
@@ -64,7 +58,7 @@ processor.tokenizer.pad_token_id = processor.tokenizer.eos_token_id
 
 
 # =========================
-# PDF → IMAGE
+# PDF → IMAGE (FAST & SAFE)
 # =========================
 def pdf_to_images(pdf_bytes):
     doc = fitz.open(stream=pdf_bytes, filetype="pdf")
@@ -86,7 +80,7 @@ def pdf_to_images(pdf_bytes):
 
 
 # =========================
-# OCR INFERENCE (FAST)
+# OCR INFERENCE (CUDA ONLY HERE)
 # =========================
 @spaces.GPU
 def run_inference(image, prompt, max_new_tokens):
@@ -181,7 +175,7 @@ OUTPUT FORMAT:
 # GRADIO UI
 # =========================
 with gr.Blocks(title="Camel-Doc-OCR") as demo:
-    gr.Markdown("## 🧾 Camel-Doc-OCR (Qwen2.5-VL – 4bit, Fast)")
+    gr.Markdown("## 🧾 Camel-Doc-OCR (Qwen2.5-VL – 4bit, HF Spaces Safe)")
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -218,13 +212,16 @@ with gr.Blocks(title="Camel-Doc-OCR") as demo:
 
 
 # =========================
-# CLEANUP & LAUNCH
+# CLEANUP
 # =========================
 def cleanup():
     torch.cuda.empty_cache()
     gc.collect()
 
 
+# =========================
+# LAUNCH
+# =========================
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",