import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
from PIL import Image

# Model nhỏ, chạy nhanh trên Hugging Face CPU
model_id = "segmind/tiny-sd"

pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
pipe = pipe.to("cpu")

# Hàm sinh ảnh
def generate_image(person_img, cloth_img, prompt):
    # Nếu không có ảnh thì dùng prompt đơn thuần
    full_prompt = "A person wearing fashionable clothes, high quality, realistic"

    if prompt:
        full_prompt += f", {prompt}"

    # Nếu có ảnh nhân vật
    if person_img:
        full_prompt += ", base character reference provided"

    # Nếu có ảnh quần áo
    if cloth_img:
        full_prompt += ", outfit reference provided"

    # Sinh ảnh
    image = pipe(full_prompt).images[0]
    return image

# Giao diện
with gr.Blocks() as demo:
    gr.Markdown("## 👗 AI Fashion Try-On (Tiny Model)")
    gr.Markdown("Upload ảnh nhân vật + quần áo + nhập mô tả để thử phối đồ AI")

    with gr.Row():
        person = gr.Image(label="Ảnh nhân vật", type="pil")
        cloth = gr.Image(label="Ảnh quần áo", type="pil")

    prompt = gr.Textbox(label="Mô tả thêm (prompt)", placeholder="Ví dụ: phong cách streetwear, ảnh ngoài trời...")

    output = gr.Image(label="Kết quả AI")

    btn = gr.Button("✨ Tạo ảnh ngay")
    btn.click(fn=generate_image, inputs=[person, cloth, prompt], outputs=output)

demo.launch()