"""Minimal inference with the EditPPT vision-validator LoRA on Qwen3.5-9B.

Input order matters: [image BEFORE the edit, image AFTER the edit, prompt]. Render slides at 960x540
(1 slide pt = 1 px for a standard 16:9 deck). Fill prompt_template.txt with str.replace, not str.format
(the template contains literal JSON braces).
"""
import json

import torch
from peft import PeftModel
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor

BASE, ADAPTER = "Qwen/Qwen3.5-9B", "EditPPT/editppt-vision-validator-lora"
PIXELS = 960 * 540

processor = AutoProcessor.from_pretrained(BASE, min_pixels=PIXELS, max_pixels=PIXELS)
model = AutoModelForImageTextToText.from_pretrained(BASE, dtype=torch.bfloat16, device_map="cuda:0")
model = PeftModel.from_pretrained(model, ADAPTER).eval()


def validate(before_png, after_png, instruction, tools_used, slide_json):
    prompt = (open("prompt_template.txt").read()
              .replace("{instruction}", instruction)
              .replace("{tools_used}", str(tools_used))   # python-repr list of {'name', 'arguments'}
              .replace("{slide_json}", str(slide_json)))  # python-repr dict: slide_w, slide_h, current_slide, objects
    msgs = [{"role": "user", "content": [{"type": "image"}, {"type": "image"}, {"type": "text", "text": prompt}]}]
    text = processor.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
    imgs = [Image.open(p).convert("RGB") for p in (before_png, after_png)]
    enc = processor(text=[text], images=imgs, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(**enc, max_new_tokens=1024, do_sample=False)
    raw = processor.tokenizer.decode(out[0, enc["input_ids"].shape[1]:], skip_special_tokens=True)
    return json.loads(raw[raw.index("{"): raw.rindex("}") + 1])