Update app.py

app.py

@@ -1,37 +1,43 @@
-from flask import Flask, send_from_directory
-from ctransformers import AutoModelForCausalLM
-import gradio as gr, io, base64
+from flask import Flask
+import gradio as gr
+from llama_cpp import Llama
+import base64, io
 from PIL import Image
 
-MODEL_REPO = "Nasa1423/RolmOCR-Q4_K_M-GGUF"
-print("Loading RolmOCR GGUF on CPU…")
+print("Loading RolmOCR (Q4 GGUF) via llama‑cpp...")
+model_path = "/home/user/.cache/huggingface/hub/models--Nasa1423--RolmOCR-Q4_K_M-GGUF/blobs/3831d4a42b3a054bbe0a6634e8ee5bb24d275d6104a39250cffeecbb1bbc3d19"
 
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_REPO,
-    model_type="llama",
-    context_length=2048,
-    gpu_layers=0,
-    lib="basic",
-)
+# if llama‑cpp cannot find the model automatically, use the HF download path like above,
+# or rely on repo_id below:
+try:
+    llm = Llama.from_pretrained(
+        repo_id="Nasa1423/RolmOCR-Q4_K_M-GGUF",
+        filename="rolmocr-q4_k_m.gguf",
+        n_ctx=2048,
+        n_threads=4,
+    )
+except Exception as e:
+    print("Auto load failed", e)
+    # fallback: load local path explicitly
+    llm = Llama(model_path=model_path, n_ctx=2048, n_threads=4)
 
 def ocr_infer(prompt, image):
-    if image is None: return "(No image given)"
-    _ = Image.fromarray(image)
-    return model(f"{prompt}\nConvert image to markdown text.", max_new_tokens=512)
+    if image is None:
+        return "(No image provided)"
+    # image validation only — the model processes text prompt
+    Image.fromarray(image)
+    q = f"You are an OCR system.\n{prompt}\nRead the uploaded image and output markdown text."
+    output = llm(q, max_tokens=512, temperature=0.1)
+    # llama‑cpp returns a dict with 'choices'
+    return output["choices"][0]["text"].strip()
 
 demo = gr.Interface(
     fn=ocr_infer,
     inputs=[gr.Textbox(label="Prompt"), gr.Image(label="Upload Image")],
     outputs="text",
-    title="RolmOCR‑GGUF CPU Space",
-    description="Runs quantized RolmOCR (Q4_K_M) locally on HF CPU infrastructure.",
+    title="RolmOCR GGUF (LLAMA‑CPP)",
+    description="Runs RolmOCR Q4_K_M on CPU via llama‑cpp‑python; fully Space compatible.",
 )
 
-app = Flask(__name__)
-
-@app.route("/")
-def index():
-    return send_from_directory(".", "index.html")
-
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)