File size: 1,139 Bytes
9ceea25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | 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} |