Update app.py

app.py

@@ -1,26 +1,25 @@
 import os
 import io
-import sys
 import time
+import sys
 import requests
 from PIL import Image, ImageSequence
 import gradio as gr
 
-# Try to import llama-cpp-python
+# llama-cpp-python import
 try:
     from llama_cpp import Llama
 except Exception as e:
-    raise RuntimeError("llama-cpp-python import failed; ensure requirements installed and wheel built: " + str(e))
+    raise RuntimeError("llama-cpp-python import failed: " + str(e))
 
-MODEL_PATH = os.path.join("model", "model.gguf")  # start.sh places GGUF here
+MODEL_PATH = os.path.join("model", "llama-joycaption-q4_k_m.gguf")
 if not os.path.exists(MODEL_PATH):
-    raise FileNotFoundError(f"Model not found at {MODEL_PATH}. Set correct GGUF in start.sh and redeploy.")
+    raise FileNotFoundError(f"Model not found at {MODEL_PATH}. Ensure start.sh downloaded the GGUF.")
 
-# Helper: load first frame and convert to JPEG bytes
 def download_bytes(url: str, timeout: int = 30) -> bytes:
-    with requests.get(url, stream=True, timeout=timeout) as resp:
-        resp.raise_for_status()
-        return resp.content
+    with requests.get(url, stream=True, timeout=timeout) as r:
+        r.raise_for_status()
+        return r.content
 
 def load_first_frame_from_bytes(raw: bytes):
     img = Image.open(io.BytesIO(raw))
@@ -30,15 +29,12 @@ def load_first_frame_from_bytes(raw: bytes):
         img = img.convert("RGB")
     return img
 
-# Minimal image caption prompt template — adjust for your model's expected prompt
 def make_prompt_for_image(image_path: str, user_prompt: str = "Describe the image."):
-    # Many llama.cpp-based multimodal ggufs accept: "<img>{path}</img>\nUser: {prompt}\nAssistant:"
-    # We'll use that pattern.
+    # JoyCaption-style multimodal GGUFs accept <img>{path}</img>
     return f"<img>{image_path}</img>\nUser: {user_prompt}\nAssistant:"
 
-# Start model (llama-cpp-python will mmap model and run inference)
-# Use low-memory opts: n_ctx small, use_mlock=0, n_gpu_layers=0
-print("Loading model (this may take a minute)...", file=sys.stderr)
+# Initialize model (low-resource options)
+print("Loading GGUF model (this can take 30–120s)...", file=sys.stderr)
 llm = Llama(model_path=MODEL_PATH, n_ctx=2048, n_threads=2)
 
 def generate_caption_from_url(url: str, prompt: str = "Describe the image."):
@@ -53,7 +49,6 @@ def generate_caption_from_url(url: str, prompt: str = "Describe the image."):
     except Exception as e:
         return f"Image processing error: {e}"
 
-    # Save a temporary JPEG locally so the gguf image token loader can access it
     tmp_dir = "/tmp/joycap"
     os.makedirs(tmp_dir, exist_ok=True)
     ts = int(time.time() * 1000)
@@ -65,7 +60,6 @@ def generate_caption_from_url(url: str, prompt: str = "Describe the image."):
 
     prompt_full = make_prompt_for_image(tmp_path, prompt)
     try:
-        # llama-cpp-python generate call
         resp = llm.create(
             prompt=prompt_full,
             max_tokens=256,
@@ -86,12 +80,12 @@ def generate_caption_from_url(url: str, prompt: str = "Describe the image."):
 iface = gr.Interface(
     fn=generate_caption_from_url,
     inputs=[
-        gr.Textbox(label="Image / GIF / MP4 URL", placeholder="https://example.com/photo.jpg"),
+        gr.Textbox(label="Image URL", placeholder="https://example.com/photo.jpg"),
         gr.Textbox(label="Prompt (optional)", value="Describe the image."),
     ],
     outputs=gr.Textbox(label="Generated caption"),
-    title="JoyCaption - local GGUF (Q4)",
-    description="Runs a quantized GGUF model locally via llama.cpp (no external APIs). Ensure the GGUF in start.sh is a multimodal model that supports <img> tags.",
+    title="JoyCaption GGUF (Q4_K_M)",
+    description="Runs a quantized JoyCaption GGUF locally via llama.cpp (no external API).",
 )
 
 if __name__ == "__main__":