Update app.py

app.py

@@ -1,10 +1,13 @@
-# app.py — DeepSeek-OCR + BioMedLM-7B (GGUF llama.cpp local, ZeroGPU-safe) — Gradio 5
-# - OCR con DeepSeek-OCR (GPU solo en @spaces.GPU)
-# - Chat con BioMedLM-7B GGUF via llama.cpp (GPU solo en @spaces.GPU)
-# - Prompt reforzado (few-shot) y decodificación determinista
-# - Configurable por variables de entorno: GGUF_REPO, GGUF_FILE, N_CTX, N_BATCH, N_GPU_LAYERS
-
-import os, re, json, tempfile, traceback
+# app.py — OpScanIA: DeepSeek-OCR + BioMedLM-7B (GGUF local con llama.cpp, ZeroGPU-safe) — Gradio 5
+# -----------------------------------------------------------------------------------------------
+# • OCR: DeepSeek-OCR (GPU SOLO dentro del worker @spaces.GPU; sin inicializar CUDA en el main).
+# • Chat: BioMedLM-7B en formato GGUF con llama.cpp (también SOLO en worker GPU).
+# • Sin llamadas GPU anidadas; todo atrapado con try/except para evitar RuntimeError genéricos.
+# • Prompt reforzado en español y generación determinista (sensible a OCR y sin alucinaciones).
+# • Configurable por variables de entorno: GGUF_REPO, GGUF_FILE, N_CTX, N_BATCH, N_GPU_LAYERS, etc.
+# -----------------------------------------------------------------------------------------------
+
+import os, re, tempfile, traceback
 import gradio as gr
 import torch
 from PIL import Image
@@ -14,12 +17,12 @@ from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
 
 # =========================
-# CONFIG (env)
+# CONFIG (entorno)
 # =========================
-# --- Llama.cpp (BioMedLM-7B GGUF) ---
-GGUF_REPO = os.getenv("GGUF_REPO", "").strip()   # ej: "theuser/biomedlm-7b-gguf" (pon el tuyo)
-GGUF_FILE = os.getenv("GGUF_FILE", "").strip()   # ej: "BioMedLM-7B.Q4_K_M.gguf"
-# candidatos por defecto si no das GGUF_FILE
+# --- BioMedLM-7B (GGUF / llama.cpp) ---
+GGUF_REPO = os.getenv("GGUF_REPO", "").strip()   # p.ej.: "tu_usuario/biomedlm-7b-gguf" (si lo tienes en HF)
+GGUF_FILE = os.getenv("GGUF_FILE", "").strip()   # p.ej.: "BioMedLM-7B.Q4_K_M.gguf" (nombre exacto del archivo)
+# Candidatos comunes si no estableces GGUF_FILE:
 _GGUF_CANDIDATES = [
     "BioMedLM-7B.Q4_K_M.gguf",
     "BioMedLM-7B.Q5_K_M.gguf",
@@ -32,36 +35,37 @@ _GGUF_CANDIDATES = [
 ]
 GGUF_CANDIDATES = [GGUF_FILE] if GGUF_FILE else _GGUF_CANDIDATES
 
-# rendimiento / memoria
+# Rendimiento / memoria (ajusta según GPU del Space: T4 / A10G)
 N_CTX = int(os.getenv("N_CTX", "4096"))
 N_THREADS = int(os.getenv("N_THREADS", str(os.cpu_count() or 4)))
-N_GPU_LAYERS = int(os.getenv("N_GPU_LAYERS", "35"))  # 7B ~32 capas; 35 = "todas"
+N_GPU_LAYERS = int(os.getenv("N_GPU_LAYERS", "35"))  # 7B ~32 capas; 35 ≈ todas
 N_BATCH = int(os.getenv("N_BATCH", "512"))           # sube a 1024 si tu GPU lo permite
 
-# generación determinista para obediencia
+# Decodificación (determinista por defecto)
 GEN_TEMPERATURE = float(os.getenv("TEMPERATURE", "0.0"))
 GEN_TOP_P = float(os.getenv("TOP_P", "1.0"))
 GEN_MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", "384"))
 STOP_SEQS = ["\n###", "\nUser:", "\nAssistant:", "\nUsuario:", "\nAsistente:"]
 
-# DeepSeek-OCR revision opcional para evitar cambios inesperados
-DS_OCR_REV = os.getenv("DS_OCR_REV", None)  # e.g. hash de commit
+# Token para repos privados en HF (opcional)
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+# DeepSeek-OCR: fija una revisión/commit opcional para evitar cambios inesperados
+DS_OCR_REV = os.getenv("DS_OCR_REV", None)  # e.g., "2b6f6c2..."
 
 # =========================
-# Estado global (solo dentro de workers GPU)
+# Utilidades de texto / prompt
 # =========================
-_llm = None
-_llm_name = None
-
 def _truncate(s: str, n=3000):
     s = (s or "")
     return s if len(s) <= n else s[:n]
 
 def _clean_ocr(s: str) -> str:
-    if not s: return ""
-    s = re.sub(r'[^\S\r\n]+', ' ', s)           # colapsa espacios
-    s = re.sub(r'(\{#Sec\d+\}|#+\w*)', ' ', s)  # anchors/headers raros
-    s = re.sub(r'\s{2,}', ' ', s)
+    if not s:
+        return ""
+    s = re.sub(r"[^\S\r\n]+", " ", s)                # colapsa espacios
+    s = re.sub(r"(\{#Sec\d+\}|#+\w*)", " ", s)       # anchors/headers raros
+    s = re.sub(r"\s{2,}", " ", s)
     lines = []
     for par in s.splitlines():
         par = par.strip()
@@ -110,75 +114,92 @@ def build_user_prompt(ocr_md, ocr_txt, user_msg):
     )
     return prompt
 
+def _to_chatml(system_prompt, user_prompt):
+    # Formato minimalista tipo ChatML/llama.cpp
+    return [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_prompt},
+    ]
+
 # =========================
 # BioMedLM-7B GGUF — llama.cpp (GPU solo en worker)
 # =========================
+_llm = None
+_llm_name = None
+
 def _download_gguf_path():
+    """Busca el .gguf en HF (con token si hace falta) o en local (Files del Space)."""
     last_err = None
     if GGUF_REPO:
         for fname in GGUF_CANDIDATES:
             try:
-                path = hf_hub_download(repo_id=GGUF_REPO, filename=fname)
+                path = hf_hub_download(repo_id=GGUF_REPO, filename=fname, token=HF_TOKEN)
                 return path, f"{GGUF_REPO}:{fname}"
             except Exception as e:
                 last_err = e
-    # fallback: si subiste el gguf al Space (en la carpeta del repo)
+    # Fallback: archivo subido al Space (pestaña Files)
     for fname in GGUF_CANDIDATES:
         local_path = os.path.join(os.getcwd(), fname)
         if os.path.exists(local_path):
             return local_path, f"./{fname}"
-    raise RuntimeError(f"No se pudo localizar el GGUF. Configura GGUF_REPO/GGUF_FILE o sube el .gguf. Último error: {last_err}")
+    raise RuntimeError(f"No se encontró el GGUF. Configura GGUF_REPO/GGUF_FILE o sube el .gguf. Último error: {last_err}")
 
-@spaces.GPU
-def biomedlm_warmup():
-    """Inicializa llama.cpp dentro del worker GPU (evita CUDA en main)."""
+def _ensure_llm():
+    """Inicializa llama.cpp en el MISMO worker; nunca lanza excepción hacia arriba."""
     global _llm, _llm_name
     if _llm is not None:
-        return f"OK::warm (reusing {_llm_name})"
-    gguf_path, used = _download_gguf_path()
-    _llm = Llama(
-        model_path=gguf_path,
-        n_ctx=N_CTX,
-        n_threads=N_THREADS,
-        n_gpu_layers=N_GPU_LAYERS,
-        n_batch=N_BATCH,
-        # decodificación por defecto: greedy (sin sampling)
-        verbose=False,
-    )
-    _llm_name = used
-    return f"OK::loaded {used}"
+        return True, f"warm (reusing {_llm_name})"
+    try:
+        gguf_path, used = _download_gguf_path()
+        _llm = Llama(
+            model_path=gguf_path,
+            n_ctx=N_CTX,
+            n_threads=N_THREADS,
+            n_gpu_layers=N_GPU_LAYERS,
+            n_batch=N_BATCH,
+            verbose=False,
+        )
+        _llm_name = used
+        return True, f"loaded {used}"
+    except Exception as e:
+        return False, f"[{e.__class__.__name__}] {str(e) or repr(e)}"
 
-def _to_chatml(system_prompt, user_prompt):
-    # formato simple ChatML-compatible para llama.cpp
-    return [
-        {"role": "system", "content": system_prompt},
-        {"role": "user", "content": user_prompt},
-    ]
+@spaces.GPU
+def biomedlm_warmup():
+    """Warmup opcional (manual) — NO se llama desde otra función GPU."""
+    ok, msg = _ensure_llm()
+    return ("OK::" if ok else "ERR::") + msg
 
 @spaces.GPU
-def biomedlm_chat(ocr_md, ocr_txt, user_msg, temperature=GEN_TEMPERATURE, top_p=GEN_TOP_P, max_tokens=GEN_MAX_NEW_TOKENS):
-    """Generación dentro del worker GPU con el LLM ya inicializado."""
-    global _llm
-    if _llm is None:
-        status = biomedlm_warmup()
-        if not str(status).startswith("OK::"):
-            return "ERR::No se pudo inicializar el modelo GGUF"
-    prompt_user = build_user_prompt(ocr_md, ocr_txt, user_msg)
-    messages = _to_chatml(SYSTEM_INSTR, prompt_user)
+def biomedlm_chat(ocr_md, ocr_txt, user_msg,
+                  temperature=GEN_TEMPERATURE, top_p=GEN_TOP_P, max_tokens=GEN_MAX_NEW_TOKENS):
+    """Chat en GPU; TODO envuelto en try/except para evitar RuntimeError del worker."""
     try:
-        out = _llm.create_chat_completion(
-            messages=messages,
-            temperature=temperature,
-            top_p=top_p,
-            max_tokens=max_tokens,
-        )
-        ans = out["choices"][0]["message"]["content"]
-        return "OK::" + (ans or "").strip()
+        ok, msg = _ensure_llm()
+        if not ok:
+            return "ERR::No se pudo inicializar el modelo GGUF -> " + msg
+
+        prompt_user = build_user_prompt(ocr_md, ocr_txt, user_msg)
+        messages = _to_chatml(SYSTEM_INSTR, prompt_user)
+
+        try:
+            out = _llm.create_chat_completion(
+                messages=messages,
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                stop=STOP_SEQS,
+            )
+            ans = (out["choices"][0]["message"]["content"] or "").strip()
+            return "OK::" + ans
+        except Exception as e:
+            return f"ERR::[Inferencia] {e.__class__.__name__}: {str(e) or repr(e)}"
+
     except Exception as e:
-        return f"ERR::[{e.__class__.__name__}] {str(e) or repr(e)}"
+        return f"ERR::[Worker] {e.__class__.__name__}: {str(e) or repr(e)}"
 
 # =========================
-# DeepSeek-OCR (GPU solo dentro del worker)
+# DeepSeek-OCR (GPU solo en worker)
 # =========================
 def _load_ocr_model():
     model_name = "deepseek-ai/DeepSeek-OCR"
@@ -186,7 +207,7 @@ def _load_ocr_model():
     kwargs = dict(
         _attn_implementation=os.getenv("OCR_ATTN_IMPL", "flash_attention_2"),
         trust_remote_code=True,
-        use_safetensors=True
+        use_safetensors=True,
     )
     if DS_OCR_REV:
         kwargs["revision"] = DS_OCR_REV
@@ -194,6 +215,7 @@ def _load_ocr_model():
         mdl = AutoModel.from_pretrained(model_name, **kwargs).eval()
         return tok, mdl
     except Exception as e:
+        # Fallback si FA2 no está disponible
         if any(k in str(e).lower() for k in ["flash_attn", "flashattention2", "flash_attention_2"]):
             kwargs["_attn_implementation"] = "eager"
             mdl = AutoModel.from_pretrained(model_name, **kwargs).eval()
@@ -206,7 +228,8 @@ tokenizer, model = _load_ocr_model()
 def process_image(image, model_size, task_type, is_eval_mode):
     if image is None:
         return None, "Please upload an image first.", "Please upload an image first."
-    # mover a GPU SOLO dentro del worker
+
+    # Mover a GPU SOLO dentro del worker
     if torch.cuda.is_available():
         dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
         model_device = model.to(dtype).to("cuda")
@@ -251,36 +274,46 @@ def process_image(image, model_size, task_type, is_eval_mode):
 
         result_image = None
         if os.path.exists(image_result_path):
-            result_image = Image.open(image_result_path); result_image.load()
+            result_image = Image.open(image_result_path)
+            result_image.load()
 
         text_result = plain_text if plain_text else markdown_content
         return result_image, markdown_content, text_result
 
 # =========================
-# Chat wrapper para la UI
+# Orquestador de chat (NO GPU)
 # =========================
 def biomedlm_reply(user_msg, chat_msgs, ocr_md, ocr_txt):
     try:
-        res = biomedlm_chat(ocr_md, ocr_txt, user_msg, temperature=GEN_TEMPERATURE, top_p=GEN_TOP_P, max_tokens=GEN_MAX_NEW_TOKENS)
-        if str(res).startswith("OK::"):
-            answer = res[4:]
+        res = biomedlm_chat(
+            ocr_md,
+            ocr_txt,
+            user_msg,
+            temperature=GEN_TEMPERATURE,
+            top_p=GEN_TOP_P,
+            max_tokens=GEN_MAX_NEW_TOKENS,
+        )
+        s = str(res)
+        if s.startswith("OK::"):
+            answer = s[4:]
             updated = (chat_msgs or []) + [
                 {"role": "user", "content": user_msg or "(analizar solo OCR)"},
-                {"role": "assistant", "content": answer}
+                {"role": "assistant", "content": answer},
             ]
             return updated, "", gr.update(value="")
         else:
-            err_msg = res[5:] if str(res).startswith("ERR::") else str(res)
+            # Mostramos TODO el mensaje de error del worker en el panel Debug
+            err_msg = s[5:] if s.startswith("ERR::") else s
             updated = (chat_msgs or []) + [
                 {"role": "user", "content": user_msg or ""},
-                {"role": "assistant", "content": "⚠️ Error LLM (local). Revisa el panel de debug."}
+                {"role": "assistant", "content": "⚠️ Error LLM (local). Revisa el panel de debug."},
             ]
             return updated, "", gr.update(value=err_msg)
     except Exception as e:
         tb = traceback.format_exc(limit=2)
         updated = (chat_msgs or []) + [
             {"role": "user", "content": user_msg or ""},
-            {"role": "assistant", "content": f"⚠️ Error LLM: {e}"}
+            {"role": "assistant", "content": f"⚠️ Error LLM: {e}"},
         ]
         return updated, "", gr.update(value=f"{e}\n{tb}")
 
@@ -293,9 +326,9 @@ def clear_chat():
 with gr.Blocks(title="OpScanIA — DeepSeek-OCR + BioMedLM-7B (GGUF)", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
-        # DeepSeek-OCR → Chat Clínico con **BioMedLM-7B (GGUF local)**
+        # DeepSeek-OCR → Chat Clínico con **BioMedLM-7B (GGUF local, llama.cpp)**
         1) **Sube una imagen** y corre **OCR** (imagen anotada, Markdown y texto).  
-        2) **Chatea** con **BioMedLM-7B GGUF (llama.cpp)** usando automáticamente el **OCR** como contexto.  
+        2) **Chatea** con **BioMedLM-7B GGUF** usando automáticamente el **OCR** como contexto.  
         *Uso educativo; no reemplaza consejo médico.*
         """
     )
@@ -305,17 +338,26 @@ with gr.Blocks(title="OpScanIA — DeepSeek-OCR + BioMedLM-7B (GGUF)", theme=gr.
 
     with gr.Row():
         with gr.Column(scale=1):
-            image_input = gr.Image(type="pil", label="Upload Image", sources=["upload", "clipboard", "webcam"])
+            image_input = gr.Image(
+                type="pil",
+                label="Upload Image",
+                sources=["upload", "clipboard", "webcam"]
+            )
             model_size = gr.Dropdown(
                 choices=["Tiny", "Small", "Base", "Large", "Gundam (Recommended)"],
-                value="Gundam (Recommended)", label="Model Size"
+                value="Gundam (Recommended)",
+                label="Model Size"
             )
             task_type = gr.Dropdown(
                 choices=["Free OCR", "Convert to Markdown"],
-                value="Convert to Markdown", label="Task Type"
+                value="Convert to Markdown",
+                label="Task Type"
+            )
+            eval_mode_checkbox = gr.Checkbox(
+                value=False,
+                label="Enable Evaluation Mode",
+                info="Solo texto (más rápido). Desmárcalo para ver imagen anotada y markdown."
             )
-            eval_mode_checkbox = gr.Checkbox(value=False, label="Enable Evaluation Mode",
-                                             info="Solo texto (más rápido). Desmárcalo para ver imagen anotada y markdown.")
             submit_btn = gr.Button("Process Image", variant="primary")
             warm_btn = gr.Button("Warmup BioMedLM-7B (GGUF)")
 
@@ -335,7 +377,11 @@ with gr.Blocks(title="OpScanIA — DeepSeek-OCR + BioMedLM-7B (GGUF)", theme=gr.
     with gr.Row():
         with gr.Column(scale=2):
             chatbot = gr.Chatbot(label="Asistente OCR (BioMedLM-7B GGUF)", type="messages", height=420)
-            user_in = gr.Textbox(label="Mensaje", placeholder="Escribe tu consulta… (vacío = analiza solo el OCR)", lines=2)
+            user_in = gr.Textbox(
+                label="Mensaje",
+                placeholder="Escribe tu consulta… (vacío = analiza solo el OCR)",
+                lines=2
+            )
             with gr.Row():
                 send_btn = gr.Button("Enviar", variant="primary")
                 clear_btn = gr.Button("Limpiar")
@@ -353,7 +399,7 @@ with gr.Blocks(title="OpScanIA — DeepSeek-OCR + BioMedLM-7B (GGUF)", theme=gr.
         outputs=[ocr_md_state, ocr_txt_state, md_preview, txt_preview],
     )
 
-    # Warmup LLM (descarga/crea el objeto Llama en GPU)
+    # Warmup LLM (descarga/carga el GGUF y crea el objeto Llama en GPU)
     warm_btn.click(fn=biomedlm_warmup, outputs=[debug_box])
 
     # Chat