Update app.py

app.py

@@ -1,70 +1,29 @@
-import os
-import shutil
-import tempfile
-import json
-from fastapi import FastAPI, UploadFile, File, HTTPException
-
-# --- CPU OPTIMIZATION FLAGS ---
-# Limit threads for HF Free Tier (2 vCPUs)
-os.environ["OMP_NUM_THREADS"] = "2" 
-
-# THE MAGIC FIX: Disable the buggy PIR compiler but KEEP MKLDNN!
-os.environ["FLAGS_enable_pir_api"] = "0"
-os.environ["FLAGS_use_mkldnn"] = "1" 
-
-# IMPORT THE NEW V3 PIPELINE (Must be imported AFTER setting the flags)
-from paddleocr import PPStructureV3
-
-app = FastAPI(
-    title="Document Ingestion API", 
-    description="Lightweight PP-StructureV3 extraction"
-)
-
-print("Initializing PP-StructureV3 (MKLDNN Enabled via AST Executor)...")
-pipeline = PPStructureV3()
-print("Pipeline ready!")
-
-@app.post("/ingest")
-async def ingest_document(file: UploadFile = File(...)):
-    if not file.filename:
-        raise HTTPException(status_code=400, detail="No file provided")
-
-    with tempfile.TemporaryDirectory() as temp_dir:
-        file_path = os.path.join(temp_dir, file.filename)
-        
-        try:
-            # 1. Save file
-            with open(file_path, "wb") as buffer:
-                shutil.copyfileobj(file.file, buffer)
-            
-            # 2. Predict 
-            output = pipeline.predict(file_path)
-            
-            parsed_pages = []
-            for page_num, res in enumerate(output):
-                md_path = os.path.join(temp_dir, f"page_{page_num + 1}.md")
-                json_path = os.path.join(temp_dir, f"page_{page_num + 1}.json")
-                
-                res.save_to_markdown(save_path=md_path)
-                res.save_to_json(save_path=json_path)
-                
-                with open(md_path, "r", encoding="utf-8") as f:
-                    md_content = f.read()
-                    
-                with open(json_path, "r", encoding="utf-8") as f:
-                    json_content = json.load(f)
-                    
-                parsed_pages.append({
-                    "page": page_num + 1,
-                    "markdown": md_content,
-                    "json_data": json_content
-                })
-                
-            return {
-                "status": "success",
-                "filename": file.filename,
-                "data": parsed_pages
-            }
-
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
+from fastapi import FastAPI, UploadFile, File
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from PIL import Image
+import torch
+import io
+
+app = FastAPI()
+
+# Load model (optimized for CPU/Low RAM)
+model_id = "THUDM/glm-4v-9b" # Or the specific 0.9B GLM-OCR variant
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = model_id.from_pretrained(model_id, trust_remote_code=True).cpu().eval()
+
+@app.post("/convert")
+async def convert_image(file: UploadFile = File(...)):
+    image_data = await file.read()
+    image = Image.open(io.BytesIO(image_data)).convert("RGB")
+    
+    # Specific prompt to trigger Chart-to-HTML
+    prompt = "Read this chart and output the data as a clean HTML table with headers."
+    
+    inputs = tokenizer.apply_chat_template([{"role": "user", "image": image, "content": prompt}], 
+                                          add_generation_prompt=True, tokenize=True, return_tensors="pt")
+    
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=1000)
+    
+    response = tokenizer.decode(outputs[0])
+    return {"html_result": response}