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

# Map tên model thành HuggingFace ID
model_map = {
    "Realistic Vision v5.1": "SG161222/Realistic_Vision_V5.1_noVAE",
    "DreamShaper v8": "Lykon/DreamShaper_8",
    "Anything v5": "stablediffusionapi/anything-v5",
    "Juggernaut XL v8": "RunDiffusion/juggernaut-xl-v8"
}

# Style presets
style_presets = {
    "Realistic": {
        "add": "ultra realistic, 4k photo, detailed skin texture, cinematic lighting",
        "neg": "blurry, cartoon, lowres, painting"
    },
    "Fantasy": {
        "add": "fantasy art, glowing light, ethereal atmosphere, magical",
        "neg": "photo, realistic skin, dirty"
    },
    "Cyberpunk": {
        "add": "cyberpunk, neon lights, rain, urban night, blade runner style",
        "neg": "natural light, low contrast"
    },
    "Default": {
        "add": "",
        "neg": ""
    }
}

# Preset image sizes
image_sizes = {
    "768x768": (768, 768),
    "1024x1024": (1024, 1024),
    "768x1024": (768, 1024),
    "1024x768": (1024, 768)
}

loaded_models = {}

def load_model(model_name):
    if model_name in loaded_models:
        return loaded_models[model_name]
    model_id = model_map[model_name]
    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
    pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
    loaded_models[model_name] = pipe
    return pipe

def generate(prompt, negative_prompt, style, model_name, size_label, batch_size):
    pipe = load_model(model_name)

    # Gộp prompt và style
    style_info = style_presets.get(style, style_presets["Default"])
    full_prompt = f"{prompt}, {style_info['add']}".strip(", ")
    negative_full = f"{negative_prompt}, {style_info['neg']}".strip(", ")

    width, height = image_sizes[size_label]

    images = pipe(
        prompt=[full_prompt] * batch_size,
        negative_prompt=[negative_full] * batch_size if negative_full else None,
        width=width,
        height=height,
        num_inference_steps=25,
        guidance_scale=7.5
    ).images

    # Lưu file để download
    output_paths = []
    for img in images:
        filename = f"{uuid.uuid4().hex}.png"
        path = f"/tmp/{filename}"
        img.save(path)
        output_paths.append((img, path))
    
    image_outputs = [img for img, _ in output_paths]
    file_outputs = [path for _, path in output_paths]
    return image_outputs, file_outputs

with gr.Blocks() as demo:
    gr.Markdown("🎨 **Prompt-to-Image Generator with Multiple Models**")

    with gr.Row():
        prompt = gr.Textbox(label="Prompt", placeholder="e.g. Vietnamese girl in ao dai, studio light")
        negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Không muốn có gì (tuỳ chọn)")

    with gr.Row():
        style = gr.Dropdown(choices=list(style_presets.keys()), label="Style", value="Default")
        model = gr.Dropdown(choices=list(model_map.keys()), label="Model", value="Realistic Vision v5.1")

    with gr.Row():
        size = gr.Dropdown(choices=list(image_sizes.keys()), label="Output Size", value="768x768")
        batch = gr.Slider(1, 2, step=1, label="Batch Size", value=1)

        generate_btn = gr.Button("Generate")

    gallery = gr.Gallery(label="Generated Images", columns=2, allow_preview=True)
    download_links = gr.File(label="Download Files", file_types=[".png"], visible=False, interactive=False)

    generate_btn.click(
        fn=generate,
        inputs=[prompt, negative_prompt, style, model, size, batch],
        outputs=[gallery, download_links]
    )

demo.launch()