Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
 import sys
 import time
+import json
 import base64
 import mimetypes
 import subprocess
@@ -12,63 +13,328 @@ import gradio as gr
 
 
 # ============================================================
-# CONFIG
+# CONFIGURATION
 # ============================================================
 
-MODEL_PATH = "/app/models/Qwen2.5-VL-3B-Instruct.Q4_K_M.gguf"
-MMPROJ_PATH = "/app/models/mmproj-Qwen2.5-VL-3B-Instruct-Q8_0.gguf"
+MODEL_DIR = "/app/models"
+
+HF_REPO = "ggml-org/Qwen2.5-VL-3B-Instruct-GGUF"
+
+# Main language model
+MODEL_FILE = "Qwen2.5-VL-3B-Instruct-Q4_K_M.gguf"
+
+# Multimodal projector
+MMPROJ_FILE = "mmproj-Qwen2.5-VL-3B-Instruct-Q8_0.gguf"
+
+MODEL_PATH = os.path.join(
+    MODEL_DIR,
+    MODEL_FILE
+)
+
+MMPROJ_PATH = os.path.join(
+    MODEL_DIR,
+    MMPROJ_FILE
+)
 
 SERVER_HOST = "127.0.0.1"
 SERVER_PORT = 8080
-SERVER_URL = f"http://{SERVER_HOST}:{SERVER_PORT}"
 
-MODEL_NAME = "Qwen2.5-VL-3B-Instruct"
+SERVER_URL = (
+    f"http://{SERVER_HOST}:{SERVER_PORT}"
+)
 
-# Number of concurrent llama.cpp slots
-N_SLOTS = 4
-
-# Context per slot
 N_CTX = 4096
+N_SLOTS = 2
 
-# CPU threads
-CPU_THREADS = max(1, os.cpu_count() or 4)
+CPU_THREADS = max(
+    1,
+    os.cpu_count() or 4
+)
 
-# Change to a smaller number if GPU memory is limited.
+# GPU offload
 N_GPU_LAYERS = 99
 
 
+server_process = None
+
+
 # ============================================================
-# SERVER PROCESS
+# PRINT CONFIG
 # ============================================================
 
-server_process = None
+def print_config():
 
+    print()
+    print("=" * 70)
+    print("QWEN2.5-VL-3B IMAGE READER")
+    print("=" * 70)
 
-def start_llama_server():
-    global server_process
+    print("HF repository:")
+    print(HF_REPO)
 
-    if not os.path.exists(MODEL_PATH):
-        raise FileNotFoundError(
-            f"Model not found:\n{MODEL_PATH}"
+    print()
+    print("Model:")
+    print(MODEL_PATH)
+
+    print()
+    print("MMProj:")
+    print(MMPROJ_PATH)
+
+    print()
+    print("Context:", N_CTX)
+    print("Slots:", N_SLOTS)
+    print("CPU threads:", CPU_THREADS)
+    print("GPU layers:", N_GPU_LAYERS)
+
+    print("=" * 70)
+    print()
+
+
+# ============================================================
+# CHECK LLAMA SERVER
+# ============================================================
+
+def find_llama_server():
+
+    possible = [
+        "llama-server",
+        "/usr/local/bin/llama-server",
+        "/usr/bin/llama-server",
+        "/app/llama-server",
+    ]
+
+    for path in possible:
+
+        if path == "llama-server":
+            import shutil
+
+            found = shutil.which(
+                "llama-server"
+            )
+
+            if found:
+                return found
+
+        elif os.path.isfile(path):
+            return path
+
+    raise FileNotFoundError(
+        "llama-server executable was not found."
+    )
+
+
+# ============================================================
+# DOWNLOAD MODEL
+# ============================================================
+
+def download_file_from_hf(filename):
+
+    os.makedirs(
+        MODEL_DIR,
+        exist_ok=True
+    )
+
+    destination = os.path.join(
+        MODEL_DIR,
+        filename
+    )
+
+    if os.path.exists(destination):
+
+        size = os.path.getsize(
+            destination
         )
 
-    if not os.path.exists(MMPROJ_PATH):
-        raise FileNotFoundError(
-            f"Multimodal projector not found:\n{MMPROJ_PATH}"
+        print(
+            f"[MODEL] Already exists: "
+            f"{destination} "
+            f"({size / 1024 / 1024:.1f} MB)"
         )
 
-    print("=" * 60)
-    print("STARTING LLAMA.CPP MULTIMODAL SERVER")
-    print("=" * 60)
+        return destination
+
+    url = (
+        f"https://huggingface.co/"
+        f"{HF_REPO}/resolve/main/"
+        f"{filename}?download=true"
+    )
+
+    print()
+    print("=" * 70)
+    print("DOWNLOADING MODEL")
+    print("=" * 70)
+
+    print("URL:")
+    print(url)
+
+    print()
+    print("Destination:")
+    print(destination)
+
+    print()
+    print("This may take several minutes.")
+    print("=" * 70)
+
+    temporary = destination + ".part"
+
+    try:
+
+        with requests.get(
+            url,
+            stream=True,
+            timeout=60,
+        ) as response:
+
+            response.raise_for_status()
+
+            total = int(
+                response.headers.get(
+                    "content-length",
+                    0
+                )
+            )
+
+            downloaded = 0
+            last_print = time.time()
+
+            with open(
+                temporary,
+                "wb"
+            ) as f:
+
+                for chunk in response.iter_content(
+                    chunk_size=1024 * 1024
+                ):
+
+                    if not chunk:
+                        continue
+
+                    f.write(chunk)
+
+                    downloaded += len(chunk)
+
+                    now = time.time()
+
+                    if (
+                        now - last_print >= 2
+                        or (
+                            total
+                            and downloaded >= total
+                        )
+                    ):
+
+                        if total:
+
+                            percent = (
+                                downloaded
+                                / total
+                                * 100
+                            )
+
+                            print(
+                                f"[DOWNLOAD] "
+                                f"{percent:.1f}% "
+                                f""
+                                f"{downloaded / 1024 / 1024:.0f} MB"
+                                f" / "
+                                f"{total / 1024 / 1024:.0f} MB"
+                            )
+
+                        else:
+
+                            print(
+                                f"[DOWNLOAD] "
+                                f"{downloaded / 1024 / 1024:.0f} MB"
+                            )
+
+                        last_print = now
+
+        os.replace(
+            temporary,
+            destination
+        )
+
+        print()
+        print(
+            "[MODEL] Download complete:"
+        )
+        print(destination)
+
+        return destination
+
+    except Exception:
+
+        if os.path.exists(
+            temporary
+        ):
+
+            try:
+                os.remove(
+                    temporary
+                )
+            except Exception:
+                pass
+
+        raise
+
+
+# ============================================================
+# PREPARE MODELS
+# ============================================================
+
+def prepare_models():
+
+    print()
+    print("=" * 70)
+    print("PREPARING QWEN2.5-VL MODELS")
+    print("=" * 70)
+
+    model = download_file_from_hf(
+        MODEL_FILE
+    )
+
+    mmproj = download_file_from_hf(
+        MMPROJ_FILE
+    )
+
+    print()
+    print("MODEL READY")
+    print()
+    print("LLM:")
+    print(model)
+
+    print()
+    print("MMPROJ:")
+    print(mmproj)
+
+    print("=" * 70)
+
+    return model, mmproj
+
+
+# ============================================================
+# START LLAMA.CPP
+# ============================================================
+
+def start_llama_server():
+
+    global server_process
+
+    model_path, mmproj_path = (
+        prepare_models()
+    )
+
+    llama_server = find_llama_server()
 
     command = [
-        "llama-server",
+
+        llama_server,
 
         "-m",
-        MODEL_PATH,
+        model_path,
 
         "--mmproj",
-        MMPROJ_PATH,
+        mmproj_path,
 
         "--host",
         SERVER_HOST,
@@ -88,63 +354,108 @@ def start_llama_server():
         "-ngl",
         str(N_GPU_LAYERS),
 
-        # Enable prompt caching
+        # KV cache quantization
         "--cache-type-k",
         "q8_0",
 
         "--cache-type-v",
         "q8_0",
 
-        # Faster startup / processing
+        # Do not start llama.cpp's own web UI
         "--no-webui",
-
-        # Disable thinking when supported by template.
-        "--reasoning-format",
-        "none",
     ]
 
-    print("Command:")
-    print(" ".join(command))
     print()
+    print("=" * 70)
+    print("STARTING LLAMA.CPP")
+    print("=" * 70)
+
+    print()
+
+    print(
+        " ".join(command)
+    )
+
+    print()
+    print("=" * 70)
 
     server_process = subprocess.Popen(
+
         command,
+
         stdout=subprocess.PIPE,
+
         stderr=subprocess.STDOUT,
+
         text=True,
+
         bufsize=1,
     )
 
-    def read_output():
-        for line in iter(server_process.stdout.readline, ""):
-            print("[llama.cpp]", line.rstrip())
+    def read_server_output():
+
+        try:
+
+            for line in iter(
+                server_process.stdout.readline,
+                ""
+            ):
+
+                print(
+                    "[llama.cpp]",
+                    line.rstrip()
+                )
+
+        except Exception as e:
+
+            print(
+                "[llama.cpp output error]",
+                e
+            )
 
     threading.Thread(
-        target=read_output,
-        daemon=True
+        target=read_server_output,
+        daemon=True,
     ).start()
 
     wait_for_server()
 
 
-def wait_for_server(timeout=180):
-    print("Waiting for llama.cpp server...")
+# ============================================================
+# WAIT FOR SERVER
+# ============================================================
+
+def wait_for_server(
+    timeout=180
+):
+
+    print()
+    print(
+        "[SYSTEM] Waiting for llama.cpp..."
+    )
 
     start = time.time()
 
-    while time.time() - start < timeout:
+    while (
+        time.time() - start
+        < timeout
+    ):
+
         try:
-            r = requests.get(
+
+            response = requests.get(
                 f"{SERVER_URL}/health",
                 timeout=2,
             )
 
-            if r.status_code == 200:
+            if response.status_code == 200:
+
                 print()
-                print("=" * 60)
+                print("=" * 70)
                 print("LLAMA.CPP SERVER READY")
-                print("=" * 60)
+                print("=" * 70)
                 print()
+
                 return
 
         except Exception:
@@ -153,77 +464,139 @@ def wait_for_server(timeout=180):
         time.sleep(0.5)
 
     raise RuntimeError(
-        "llama.cpp server did not become ready."
+        "llama.cpp did not become ready "
+        f"within {timeout} seconds."
     )
 
 
 # ============================================================
-# IMAGE → BASE64
+# IMAGE → DATA URL
 # ============================================================
 
-def image_to_data_url(image_path):
+def image_to_data_url(
+    image_path
+):
+
     if not image_path:
-        return None
+        raise ValueError(
+            "No image was provided."
+        )
 
-    path = Path(image_path)
+    path = Path(
+        image_path
+    )
 
     if not path.exists():
+
         raise FileNotFoundError(
             f"Image not found: {image_path}"
         )
 
-    mime_type, _ = mimetypes.guess_type(str(path))
+    mime_type, _ = (
+        mimetypes.guess_type(
+            str(path)
+        )
+    )
+
+    if mime_type is None:
 
-    if not mime_type:
         mime_type = "image/jpeg"
 
-    with open(path, "rb") as f:
+    with open(
+        path,
+        "rb"
+    ) as f:
+
         encoded = base64.b64encode(
             f.read()
-        ).decode("utf-8")
+        ).decode(
+            "utf-8"
+        )
 
-    return f"data:{mime_type};base64,{encoded}"
+    return (
+        f"data:{mime_type};base64,{encoded}"
+    )
 
 
 # ============================================================
-# QWEN VISION REQUEST
+# SEND IMAGE + TEXT TO LLAMA.CPP
 # ============================================================
 
-def ask_qwen(image_path, instruction, max_tokens):
+def ask_qwen(
+    image_path,
+    instruction,
+    max_tokens,
+):
+
     if not image_path:
-        raise gr.Error("Please upload an image.")
 
-    if not instruction or not instruction.strip():
+        raise gr.Error(
+            "Please upload an image."
+        )
+
+    if (
+        instruction is None
+        or not instruction.strip()
+    ):
+
         instruction = (
             "Read all visible text in this image. "
-            "Read it in natural reading order, "
-            "from top to bottom and left to right. "
+            "Read it from top to bottom and left to right. "
+            "Preserve the natural reading order. "
             "Return only the text."
         )
 
-    image_data_url = image_to_data_url(image_path)
+    try:
+
+        image_url = (
+            image_to_data_url(
+                image_path
+            )
+        )
+
+    except Exception as e:
+
+        raise gr.Error(
+            str(e)
+        )
+
+    # ========================================================
+    # MULTIMODAL MESSAGE
+    # ========================================================
 
     messages = [
+
         {
             "role": "system",
-            "content": (
-                "You are an image reading and OCR assistant. "
-                "Follow the user's instruction exactly. "
-                "When extracting text, preserve the natural "
-                "reading order from top to bottom and left to right."
-            ),
+
+            "content":
+                (
+                    "You are a fast image reading "
+                    "and OCR assistant. "
+                    "Follow the user's instruction exactly. "
+                    "When reading text, use natural "
+                    "reading order from top to bottom "
+                    "and left to right."
+                ),
         },
+
         {
             "role": "user",
+
             "content": [
+
                 {
                     "type": "image_url",
+
                     "image_url": {
-                        "url": image_data_url
+
+                        "url": image_url
                     },
                 },
+
                 {
                     "type": "text",
+
                     "text": instruction,
                 },
             ],
@@ -231,115 +604,181 @@ def ask_qwen(image_path, instruction, max_tokens):
     ]
 
     payload = {
-        "model": MODEL_NAME,
-        "messages": messages,
 
-        "temperature": 0.0,
+        "model":
+            "Qwen2.5-VL-3B-Instruct",
+
+        "messages":
+            messages,
 
-        "max_tokens": int(max_tokens),
+        "temperature":
+            0.0,
 
-        "stream": True,
+        "top_p":
+            1.0,
 
-        # llama.cpp options
-        "cache_prompt": True,
+        "max_tokens":
+            int(max_tokens),
 
-        "reasoning_format": "none",
+        "stream":
+            True,
+
+        "cache_prompt":
+            True,
     }
 
     print()
-    print("=" * 60)
-    print("VISION REQUEST")
-    print("=" * 60)
-    print("Instruction:", instruction)
-    print()
+    print("=" * 70)
+    print("NEW VISION REQUEST")
+    print("=" * 70)
+
+    print(
+        "Instruction:",
+        instruction
+    )
+
+    print("=" * 70)
 
     try:
+
         response = requests.post(
+
             f"{SERVER_URL}/v1/chat/completions",
+
             json=payload,
+
             stream=True,
+
             timeout=600,
         )
 
-        if response.status_code != 200:
-            error_text = response.text
+    except Exception as e:
+
+        raise gr.Error(
+            f"Could not connect to llama.cpp:\n{e}"
+        )
+
+    if response.status_code != 200:
 
-            raise RuntimeError(
-                f"llama.cpp returned HTTP "
-                f"{response.status_code}:\n{error_text}"
+        try:
+            error_body = (
+                response.text
             )
+        except Exception:
+            error_body = "Unknown error"
 
-        output = ""
+        raise gr.Error(
+            "llama.cpp error "
+            f"HTTP {response.status_code}:\n"
+            f"{error_body}"
+        )
+
+    output = ""
+
+    try:
 
-        for line in response.iter_lines(
+        for raw_line in response.iter_lines(
             decode_unicode=True
         ):
+
+            if not raw_line:
+                continue
+
+            line = raw_line.strip()
+
             if not line:
                 continue
 
-            if line.startswith("data:"):
-                data = line[5:].strip()
+            if line.startswith(
+                "data:"
+            ):
 
-                if data == "[DONE]":
-                    break
+                data = line[
+                    len("data:"):
+                ].strip()
 
-                try:
-                    import json
+            else:
 
-                    obj = json.loads(data)
+                data = line
 
-                    choices = obj.get(
-                        "choices",
-                        []
-                    )
+            if data == "[DONE]":
+                break
 
-                    if not choices:
-                        continue
+            try:
 
-                    delta = choices[0].get(
-                        "delta",
-                        {}
-                    )
+                obj = json.loads(
+                    data
+                )
 
-                    token = delta.get(
-                        "content",
-                        ""
-                    )
+            except json.JSONDecodeError:
+
+                continue
+
+            choices = obj.get(
+                "choices",
+                []
+            )
 
-                    if token:
-                        output += token
-                        yield output
+            if not choices:
+                continue
+
+            choice = choices[0]
 
-                except Exception:
-                    continue
+            delta = choice.get(
+                "delta",
+                {}
+            )
+
+            token = delta.get(
+                "content",
+                ""
+            )
 
-        if not output:
-            yield "No response was generated."
+            if token:
+
+                output += token
+
+                yield output
+
+    except Exception as e:
 
-    except requests.exceptions.RequestException as e:
         raise gr.Error(
-            f"Connection to llama.cpp failed:\n{e}"
+            f"Streaming error:\n{e}"
         )
 
-    except Exception as e:
-        raise gr.Error(str(e))
+    if not output:
+
+        yield (
+            "No text was returned by "
+            "the vision model."
+        )
 
 
 # ============================================================
-# SIMPLE IMAGE OCR
+# DEFAULT OCR PROMPT
 # ============================================================
 
 DEFAULT_INSTRUCTION = """Read all text visible in this image.
 
-Reading order:
-1. Top to bottom.
-2. Within the same line, left to right.
-3. Preserve paragraphs and line breaks where possible.
-4. Do not describe the image.
-5. Return only the text you can read."""
+Read the text in natural reading sequence:
+1. From top to bottom.
+2. Within each line, from left to right.
+3. Follow columns in their natural order.
+4. Preserve paragraphs and line breaks where possible.
+5. Do not describe the image.
+6. Return only the text you can read."""
+
+
+# ============================================================
+# RUN
+# ============================================================
 
+def run_vision(
+    image,
+    instruction,
+    max_tokens,
+):
 
-def run_vision(image, instruction, max_tokens):
     yield from ask_qwen(
         image,
         instruction,
@@ -352,11 +791,16 @@ def run_vision(image, instruction, max_tokens):
 # ============================================================
 
 def clear_all():
-    return None, DEFAULT_INSTRUCTION, ""
+
+    return (
+        None,
+        DEFAULT_INSTRUCTION,
+        "",
+    )
 
 
 # ============================================================
-# UI
+# GRADIO UI
 # ============================================================
 
 def create_ui():
@@ -369,22 +813,29 @@ def create_ui():
             """
 # Qwen2.5-VL 3B Image Reader
 
-Upload an image and enter an instruction.
+Upload an image and give the model an instruction.
 
-The image and text instruction are sent directly to the
-local **llama.cpp Qwen2.5-VL** server.
+The image and instruction are sent to the local
+**llama.cpp + Qwen2.5-VL-3B** server.
 """
         )
 
         with gr.Row():
 
+            # =================================================
+            # LEFT
+            # =================================================
+
             with gr.Column(
                 scale=1
             ):
 
                 image = gr.Image(
-                    label="Upload Image",
+
+                    label="Image",
+
                     type="filepath",
+
                     sources=[
                         "upload",
                         "clipboard",
@@ -392,27 +843,43 @@ local **llama.cpp Qwen2.5-VL** server.
                 )
 
                 instruction = gr.Textbox(
+
                     label="Instruction",
-                    value=DEFAULT_INSTRUCTION,
-                    lines=8,
+
+                    value=(
+                        DEFAULT_INSTRUCTION
+                    ),
+
+                    lines=9,
+
                     placeholder=(
-                        "Example: Read all text from "
+                        "Example: "
+                        "Read all text from "
                         "top to bottom."
                     ),
                 )
 
                 max_tokens = gr.Slider(
+
                     minimum=64,
+
                     maximum=4096,
+
                     value=1024,
+
                     step=64,
-                    label="Maximum Output Tokens",
+
+                    label=(
+                        "Maximum output tokens"
+                    ),
                 )
 
                 with gr.Row():
 
                     run_button = gr.Button(
+
                         "Read Image",
+
                         variant="primary",
                     )
 
@@ -420,30 +887,48 @@ local **llama.cpp Qwen2.5-VL** server.
                         "Clear"
                     )
 
+            # =================================================
+            # RIGHT
+            # =================================================
+
             with gr.Column(
                 scale=1
             ):
 
                 output = gr.Textbox(
+
                     label="Result",
+
                     lines=25,
+
                     show_copy_button=True,
+
                     autoscroll=True,
                 )
 
+        # =====================================================
+        # EVENTS
+        # =====================================================
+
         run_button.click(
+
             fn=run_vision,
+
             inputs=[
                 image,
                 instruction,
                 max_tokens,
             ],
+
             outputs=output,
         )
 
         clear_button.click(
+
             fn=clear_all,
+
             inputs=[],
+
             outputs=[
                 image,
                 instruction,
@@ -451,76 +936,4 @@ local **llama.cpp Qwen2.5-VL** server.
             ],
         )
 
-        gr.Markdown(
-            """
-### Example instructions
-
-**OCR**
-> Read all text in this image from top to bottom and left to right. Return only the text.
-
-**Specific region**
-> Read the text in the upper half of the image.
-
-**Form**
-> Extract the form fields and their values in reading order.
-
-**Question**
-> What is the name shown in the image?
-"""
-        )
-
-    return demo
-
-
-# ============================================================
-# MAIN
-# ============================================================
-
-if __name__ == "__main__":
-
-    print()
-    print("=" * 60)
-    print("QWEN2.5-VL 3B + LLAMA.CPP")
-    print("=" * 60)
-    print("Model:", MODEL_PATH)
-    print("MMProj:", MMPROJ_PATH)
-    print("Context:", N_CTX)
-    print("Slots:", N_SLOTS)
-    print("CPU threads:", CPU_THREADS)
-    print("=" * 60)
-    print()
-
-    try:
-        start_llama_server()
-
-        demo = create_ui()
-
-        demo.queue(
-            max_size=20,
-            default_concurrency_limit=N_SLOTS,
-        )
-
-        demo.launch(
-            server_name="0.0.0.0",
-            server_port=7860,
-            show_error=True,
-        )
-
-    except KeyboardInterrupt:
-        print("Stopping...")
-
-    except Exception as e:
-        print()
-        print("=" * 60)
-        print("STARTUP FAILED")
-        print("=" * 60)
-        print(repr(e))
-        print()
-
-        if server_process is not None:
-            try:
-                server_process.terminate()
-            except Exception:
-                pass
-
-        raise
+