import gradio as gr
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from peft import PeftModel
import torch
from PIL import Image

# 1. Aapka apna model (Jisme naya knowledge hai)
adapter_id = "ocilab/flowdex-sketch-model"
# 2. Original Base Model (Jiske upar knowledge attach hoga)
base_model_id = "unsloth/Qwen2-VL-2B-Instruct"

print("AI Pipeline Setup ho rahi hai... (Takes a moment on CPU)")

# Processor aapke adapter se hi load hoga
processor = AutoProcessor.from_pretrained(adapter_id)

# Pehle Base Model load karein
base_model = Qwen2VLForConditionalGeneration.from_pretrained(
    base_model_id, 
    torch_dtype=torch.bfloat16, 
    device_map="cpu"
)

# Ab base model ke upar aapki training attach karein
model = PeftModel.from_pretrained(base_model, adapter_id)

def generate_ui(image):
    if image is None:
        return "Please upload a sketch image first."
    
    # Memory bachane ke liye resize
    image.thumbnail((448, 448))
    
    prompt = "You are an expert frontend developer. Convert this hand-drawn sketch into clean and functional components for an editable UI canvas."
    messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt}]}]
    
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt")
    
    outputs = model.generate(
        **inputs, 
        max_new_tokens=1500,
        do_sample=True,
        temperature=0.3,
        top_p=0.9
    )
    
    input_length = inputs["input_ids"].shape[1]
    generated_tokens = outputs[0][input_length:]
    generated_code = processor.decode(generated_tokens, skip_special_tokens=True)
    
    return generated_code

iface = gr.Interface(
    fn=generate_ui,
    inputs=gr.Image(type="pil", label="Upload Sketch"),
    outputs=gr.Code(language="html", label="Generated Editable UI Code"),
    title="Flowdex AI - Sketch to Design",
    description="Upload a hand-drawn wireframe to generate Tailwind/HTML components."
)

iface.launch()