Update app.py

app.py

@@ -36,15 +36,21 @@ st.set_page_config(
 # ---------------------------
 # Global UI / Render constants (NOT args to set_page_config)
 # ---------------------------
-IMAGE_PREVIEW_WIDTH = 1000         # 5x larger preview
-PDF_RENDER_SCALE = 3.0             # higher-res PDF rasterization
+IMAGE_PREVIEW_WIDTH = 1000
+PDF_RENDER_SCALE = 3.0
 
 # ---------------------------
 # Secrets / Tokens
 # ---------------------------
+# OpenRouter + HF API
 OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")  # For OpenRouter models
 HF_TOKEN = os.getenv("HF_TOKEN")                      # For HF Inference API (LLaVA)
 
+# RunPod (secured, OpenAI-compatible)
+RUNPOD_SECURE_BASE_URL   = os.getenv("RUNPOD_SECURE_BASE_URL", "").rstrip("/")  # e.g. http://194.68.245.201:22156/v1
+RUNPOD_SECURE_API_KEY    = os.getenv("RUNPOD_SECURE_API_KEY")                   # optional
+RUNPOD_SECURE_MODEL      = os.getenv("RUNPOD_SECURE_MODEL", "qwen2.5:32b-instruct")  # set to your model id
+
 # ---------------------------
 # Helpers
 # ---------------------------
@@ -83,21 +89,21 @@ def extract_structured_data(content, fields):
         pass
     return structured_data
 
+def is_vision_model_name(name: str) -> bool:
+    """Heuristic: treat models containing 'vl', 'vision', 'mm', or 'multimodal' as vision-capable."""
+    n = (name or "").lower()
+    return any(k in n for k in ["vl", "vision", "mm", "multimodal"])
+
 # ---------------------------
 # OpenRouter client (multimodal chat)
 # ---------------------------
 def query_openrouter(prompt: str, image_base64: str, model_id: str) -> str:
-    """
-    Calls OpenRouter's /api/v1/chat/completions with a text prompt + one image.
-    Requires OPENROUTER_API_KEY.
-    """
     if not OPENROUTER_API_KEY:
         raise RuntimeError("Missing OPENROUTER_API_KEY. Add it in your Space → Settings → Variables & secrets.")
 
     data_url = f"data:image/jpeg;base64,{image_base64}"
-
     payload = {
-        "model": model_id,  # e.g., "google/gemma-3-4b-it", "openai/gpt-4.1"
+        "model": model_id,
         "messages": [
             {
                 "role": "user",
@@ -109,20 +115,14 @@ def query_openrouter(prompt: str, image_base64: str, model_id: str) -> str:
         ],
         "max_tokens": 800
     }
-
     headers = {
         "Authorization": f"Bearer {OPENROUTER_API_KEY}",
         "Content-Type": "application/json",
         "HTTP-Referer": st.secrets.get("SPACE_URL", "https://hf.space"),
         "X-Title": "EZOFIS AI OCR"
     }
-
-    r = requests.post(
-        "https://openrouter.ai/api/v1/chat/completions",
-        headers=headers,
-        json=payload,
-        timeout=120
-    )
+    r = requests.post("https://openrouter.ai/api/v1/chat/completions",
+                      headers=headers, json=payload, timeout=120)
     r.raise_for_status()
     data = r.json()
     return data["choices"][0]["message"]["content"]
@@ -139,53 +139,92 @@ def _hf_client(model_id: str):
     return InferenceClient(model=model_id, token=HF_TOKEN)
 
 def query_hf_llava_vqa(prompt: str, image_base64: str, model_id: str) -> str:
-    """
-    Calls Hugging Face Hosted Inference API for VQA without extra kwargs that
-    some client versions don’t support. Includes robust fallbacks for return types.
-    """
     client = _hf_client(model_id)
     image_bytes = base64.b64decode(image_base64)
-
-    # Primary: simple VQA call (most deployments support this signature)
     try:
-        result = client.visual_question_answering(
-            image=image_bytes,
-            question=prompt
-        )
+        result = client.visual_question_answering(image=image_bytes, question=prompt)
     except TypeError:
-        # Fallback for client variants that don’t expose the helper
         result = client.request(
             task="visual_question_answering",
             data={"inputs": {"question": prompt}},
             files={"image": image_bytes}
         )
 
-    # Normalize result into a string
     if isinstance(result, str):
         return result
     if isinstance(result, dict):
-        if "answer" in result:
-            return result["answer"]
-        if "generated_text" in result:
-            return result["generated_text"]
+        return result.get("answer") or result.get("generated_text") or json.dumps(result, ensure_ascii=False)
     if isinstance(result, list) and result:
         first = result[0]
         if isinstance(first, dict):
-            if "answer" in first:
-                return first["answer"]
-            if "generated_text" in first:
-                return first["generated_text"]
+            return first.get("answer") or first.get("generated_text") or json.dumps(first, ensure_ascii=False)
+        return str(first)
     return str(result)
 
+# ---------------------------
+# RunPod (secured, OpenAI-compatible)
+# ---------------------------
+def _secured_openai_compatible(prompt: str, image_base64: str) -> str:
+    """
+    Call your OpenAI-compatible server on RunPod/OpenWebUI/Ollama.
+    Works with base URLs that already include /v1 or not.
+    API key header is added only if provided.
+    """
+    if not RUNPOD_SECURE_BASE_URL:
+        raise RuntimeError("RUNPOD_SECURE_BASE_URL is missing.")
+
+    base = RUNPOD_SECURE_BASE_URL.rstrip("/")
+    if base.endswith("/v1"):
+        url = f"{base}/chat/completions"
+    else:
+        url = f"{base}/v1/chat/completions"
+
+    headers = {"Content-Type": "application/json"}
+    if RUNPOD_SECURE_API_KEY:
+        headers["Authorization"] = f"Bearer {RUNPOD_SECURE_API_KEY}"
+
+    # If the configured model isn't vision-capable, send text-only content.
+    model_name = RUNPOD_SECURE_MODEL
+    vision_ok = is_vision_model_name(model_name)
+
+    if vision_ok:
+        data_url = f"data:image/jpeg;base64,{image_base64}"
+        content = [
+            {"type": "text", "text": prompt},
+            {"type": "image_url", "image_url": {"url": data_url}}
+        ]
+    else:
+        # Text-only fallback: no image is sent.
+        content = [
+            {"type": "text", "text": f"{prompt}\n\n(Note: model configured as text-only; image not sent.)"}
+        ]
+
+    payload = {
+        "model": model_name,
+        "messages": [{"role": "user", "content": content}],
+        "max_tokens": 800
+    }
+
+    r = requests.post(url, headers=headers, json=payload, timeout=600)
+    r.raise_for_status()
+    js = r.json()
+    return js["choices"][0]["message"]["content"]
+
+def query_runpod_secured(prompt: str, image_base64: str) -> str:
+    return _secured_openai_compatible(prompt, image_base64)
+
 # ---------------------------
 # Router to pick the right backend by model selection
 # ---------------------------
 HF_LLaVA_LABEL = "llava-hf/llava-v1.6-mistral-7b-hf (HF API)"
 HF_LLaVA_ID = "llava-hf/llava-v1.6-mistral-7b-hf"
+RUNPOD_SECURE_LABEL = "RunPod (secured)"
 
 def run_vision_inference(prompt: str, img_b64: str, model_id: str) -> str:
     if model_id == HF_LLaVA_LABEL:
         return query_hf_llava_vqa(prompt, img_b64, HF_LLaVA_ID)
+    if model_id == RUNPOD_SECURE_LABEL:
+        return query_runpod_secured(prompt, img_b64)
     # All others go via OpenRouter
     return query_openrouter(prompt, img_b64, model_id)
 
@@ -306,12 +345,24 @@ with st.sidebar:
             "google/gemma-3-12b-it",
             "openai/gpt-4.1",
             "openai/gpt-4.1-mini",
-            "qwen/qwen2.5-vl-32b-instruct",
-            HF_LLaVA_LABEL  # LLaVA via HF API
+            "qwen/qwen2.5-vl-32b-instruct",  # OpenRouter vision option
+            HF_LLaVA_LABEL,                   # LLaVA via HF API
+            RUNPOD_SECURE_LABEL               # Your RunPod OpenAI-compatible server
         ],
-        help="OpenRouter models use OPENROUTER_API_KEY. LLaVA uses HF_TOKEN via HF Inference API."
+        help=("OpenRouter uses OPENROUTER_API_KEY. "
+              "LLaVA (HF API) uses HF_TOKEN. "
+              "RunPod (secured) uses RUNPOD_SECURE_* env vars. "
+              f"Current RunPod model: {RUNPOD_SECURE_MODEL}")
     )
 
+    # If RunPod model looks text-only, warn user
+    if selected_model == RUNPOD_SECURE_LABEL and not is_vision_model_name(RUNPOD_SECURE_MODEL):
+        st.warning(
+            f"RunPod model '{RUNPOD_SECURE_MODEL}' appears text-only. "
+            "Requests to this endpoint will NOT include images. "
+            "Use a VL model (e.g. 'qwen2.5-vl:32b-instruct') for vision."
+        )
+
     extraction_mode = "General description"
     pdf_process_mode = "Process each page separately"
     fields = None
@@ -345,22 +396,27 @@ with st.sidebar:
 
 # Processing loop
 if uploaded_files and process_button:
-    # Check tokens depending on model choice
+    # Token checks by route
+    can_run = False
     if selected_model == HF_LLaVA_LABEL:
         if not HF_CLIENT_AVAILABLE:
             st.error("huggingface_hub not installed. Add 'huggingface_hub' to requirements.txt.")
         elif not HF_TOKEN:
-            st.error("HF_TOKEN is not set. Add it in your Space → Settings → Variables & secrets.")
+            st.error("HF_TOKEN is not set.")
+        else:
+            can_run = True
+    elif selected_model == RUNPOD_SECURE_LABEL:
+        if not RUNPOD_SECURE_BASE_URL:
+            st.error("RUNPOD_SECURE_BASE_URL is not set.")
         else:
             can_run = True
     else:
         if not OPENROUTER_API_KEY:
-            st.error("OPENROUTER_API_KEY is not set. Add it in your Space → Settings → Variables & secrets.")
-            can_run = False
+            st.error("OPENROUTER_API_KEY is not set.")
         else:
             can_run = True
 
-    if 'can_run' in locals() and can_run:
+    if can_run:
         st.header("Processing Results")
         progress_bar = st.progress(0)
         status_text = st.empty()
@@ -472,7 +528,10 @@ if not uploaded_files:
     st.write("""
     How to use:
     1) Upload one or more images or PDFs
-    2) Choose a model
+    2) Choose a model:
+       - OpenRouter: Gemma-3 4B/12B, GPT-4.1/4.1-mini, Qwen2.5-VL-32B
+       - HF API: LLaVA v1.6 Mistral-7B
+       - RunPod (secured): OpenAI-compatible base URL (supports images only if the model is VL)
     3) Pick description or custom field extraction
     4) For PDFs, choose page-by-page or first page
     5) Click Process Files
@@ -487,4 +546,4 @@ st.markdown(
     </div>
     """,
     unsafe_allow_html=True
-)
+)