import os
import torch
from flask import Flask, render_template, request
from transformers import AutoProcessor, AutoModelForVision2Seq
from PIL import Image
from dotenv import load_dotenv

# โหลด token
load_dotenv()
token = os.getenv("HUGGINGFACE_TOKEN")
if not token:
    raise ValueError("HUGGINGFACE_TOKEN is not set or .env file not loaded")

# Flask setup
app = Flask(__name__)
UPLOAD_FOLDER = "static/uploads"
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# โหลดโมเดลและ processor
processor = AutoProcessor.from_pretrained("scb10x/typhoon-ocr-7b", token=token)
model = AutoModelForVision2Seq.from_pretrained(
    "scb10x/typhoon-ocr-7b",
    torch_dtype=torch.float16,
    device_map="auto",
    token=token
)

@app.route('/', methods=['GET', 'POST'])
def index():
    result_text = ""
    image_path = ""

    if request.method == 'POST':
        file = request.files['image']
        if file:
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
            file.save(filepath)
            image_path = filepath

            image = Image.open(filepath).convert("RGB")
            pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(model.device)
            generated_ids = model.generate(pixel_values, max_new_tokens=256)
            result_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

    return render_template("index.html", result=result_text, image_path=image_path)

if __name__ == '__main__':
    app.run(debug=True)