first init

app.py

@@ -18,10 +18,9 @@ import spaces
 # ==================== DEEPSEEK OCR SETUP ====================
 OCR_MODEL_NAME = 'deepseek-ai/DeepSeek-OCR'
 
-print("🔄 Loading OCR tokenizer...")
+print("🔄 Loading OCR model...")
 ocr_tokenizer = AutoTokenizer.from_pretrained(OCR_MODEL_NAME, trust_remote_code=True)
 
-print("🔄 Loading OCR model...")
 try:
     ocr_model = AutoModel.from_pretrained(
         OCR_MODEL_NAME, 
@@ -41,6 +40,7 @@ except (ImportError, ValueError):
         use_safetensors=True
     )
 
+# Don't move model to GPU here - let @spaces.GPU handle it
 ocr_model = ocr_model.eval()
 
 MODEL_CONFIGS = {
@@ -124,12 +124,12 @@ def clean_output(text, include_images=False, remove_labels=False):
     return text.strip()
 
 # ==================== OCR FUNCTIONS ====================
-@spaces.GPU(duration=60)
+@spaces.GPU
 def ocr_process_image(image, mode="Crab"):
     if image is None:
         return "Error: Upload image"
     
-    # Move model to GPU inside @spaces.GPU decorated function
+    # Move model to GPU inside the @spaces.GPU decorated function
     device = "cuda" if torch.cuda.is_available() else "cpu"
     ocr_model.to(device)
     
@@ -281,9 +281,10 @@ def split_by_sentences(text: str, max_words: int = 100):
     
     return chunks
 
-@spaces.GPU(duration=60)
+@spaces.GPU
 def translate_chunk(chunk_text):
     device = "cuda" if torch.cuda.is_available() else "cpu"
+    # Ensure translator is on correct device
     if hasattr(translator, 'model') and hasattr(translator.model, 'to'):
         translator.model.to(device)
     return translator.translate(chunk_text, max_new_tokens=2048).strip()
@@ -417,12 +418,12 @@ footer { visibility: hidden }
 }
 """
 
-with gr.Blocks(theme=gr.themes.Soft(), css=css, title="🦀 MedCrab Translation") as demo:
+with gr.Blocks(theme=gr.themes.Soft(), css=css, title="OCR + Translation") as demo:
     
     gr.Markdown("""
     <div class="main-title">
     <h1>🦀 MedCrab Translation</h1>
-    <p><b>Quét PDF Y khoa → Dịch trực tiếp sang tiếng Việt</b></p>
+    <p><b>Quét PDF Y khoa → Dịch trực tiếp sang tiếng Việt (Streaming)</b></p>
     </div>
     """)
     
@@ -456,7 +457,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, title="🦀 MedCrab Translation"
         - **Crab**: 1024 base + 640 tiles (Tốt nhất, cân bằng)
         - **Base**: 1024×1024 (Nhanh hơn)
         
-        **Lưu ý:** Space này sử dụng GPU miễn phí của Hugging Face. Có thể mất vài giây để khởi động.
+        **Lưu ý:** Space này sử dụng GPU miễn phí của Hugging Face, có thể mất vài giây để khởi động.
         """)
     
     file_in.change(load_image, [file_in, page_input], [input_img])