Update app.py

app.py

@@ -3,6 +3,7 @@ import io
 import time
 import sys
 import requests
+import subprocess
 from PIL import Image, ImageSequence
 import gradio as gr
 
@@ -12,9 +13,55 @@ try:
 except Exception as e:
     raise RuntimeError("llama-cpp-python import failed: " + str(e))
 
-MODEL_PATH = os.path.join("model", "llama-joycaption-q4_k_m.gguf")
+# Ensure model is present: try Q4_K_M, fall back to Q4_K_S
+def download_with_curl(url: str, path: str) -> bool:
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    try:
+        subprocess.check_call(["curl", "-L", "-C", "-", "-o", path, url])
+        return True
+    except Exception:
+        if os.path.exists(path):
+            try:
+                os.remove(path)
+            except Exception:
+                pass
+        return False
+
+def ensure_model() -> str:
+    target = os.path.join("model", "llama-joycaption-q4_k_m.gguf")
+    alt = os.path.join("model", "llama-joycaption-q4_k_s.gguf")
+    if os.path.exists(target):
+        return target
+    # URLs (update if upstream changes)
+    url_m = "https://huggingface.co/mradermacher/llama-joycaption-beta-one-hf-llava-GGUF/resolve/main/llama-joycaption-beta-one-hf-llava-q4_k_m.gguf"
+    url_s = "https://huggingface.co/mradermacher/llama-joycaption-beta-one-hf-llava-GGUF/resolve/main/llama-joycaption-beta-one-hf-llava-q4_k_s.gguf"
+    sys.stderr.write("Model not found locally, attempting download (this can be several GB)...\n")
+    # Try M first
+    if download_with_curl(url_m, target):
+        sys.stderr.write("Downloaded Q4_K_M model.\n")
+        return target
+    # Try S
+    if download_with_curl(url_s, alt):
+        # create symlink so code expecting the M filename still works
+        try:
+            if os.path.exists(target):
+                os.remove(target)
+            os.symlink(os.path.basename(alt), target)
+        except Exception:
+            # fallback: copy file
+            try:
+                import shutil
+                shutil.copyfile(alt, target)
+            except Exception:
+                pass
+        sys.stderr.write("Downloaded Q4_K_S model and linked as Q4_K_M filename.\n")
+        return target
+    raise FileNotFoundError("Failed to download both Q4_K_M and Q4_K_S GGUF models. Check logs and URLs.")
+
+# Ensure model is available before loading llama.cpp
+MODEL_PATH = ensure_model()
 if not os.path.exists(MODEL_PATH):
-    raise FileNotFoundError(f"Model not found at {MODEL_PATH}. Ensure start.sh downloaded the GGUF.")
+    raise FileNotFoundError(f"Model not found at {MODEL_PATH} after download attempt.")
 
 def download_bytes(url: str, timeout: int = 30) -> bytes:
     with requests.get(url, stream=True, timeout=timeout) as r:
@@ -35,6 +82,7 @@ def make_prompt_for_image(image_path: str, user_prompt: str = "Describe the imag
 
 # Initialize model (low-resource options)
 print("Loading GGUF model (this can take 30–120s)...", file=sys.stderr)
+# Adjust n_threads as appropriate for the Space CPU available
 llm = Llama(model_path=MODEL_PATH, n_ctx=2048, n_threads=2)
 
 def generate_caption_from_url(url: str, prompt: str = "Describe the image."):
@@ -84,7 +132,7 @@ iface = gr.Interface(
         gr.Textbox(label="Prompt (optional)", value="Describe the image."),
     ],
     outputs=gr.Textbox(label="Generated caption"),
-    title="JoyCaption GGUF (Q4_K_M)",
+    title="JoyCaption GGUF (Q4_K)",
     description="Runs a quantized JoyCaption GGUF locally via llama.cpp (no external API).",
 )