import gradio as gr
import numpy as np
import random
import torch
from diffusers import DiffusionPipeline

# -----------------------------
# CPU MODE ONLY
# -----------------------------
device = "cpu"
torch_dtype = torch.float32

MODEL_CONFIGS = {
    "FLUX.1-dev (CPU mode)": {
        "repo_id": "black-forest-labs/FLUX.1-dev",
        "width": 512,
        "height": 512,
        "guidance": 3.0,
        "steps": 15,
    },
    "SDXL 1.0 (CPU mode)": {
        "repo_id": "stabilityai/stable-diffusion-xl-base-1.0",
        "width": 768,
        "height": 768,
        "guidance": 5.0,
        "steps": 20,
    },
}

PIPELINES = {}
MAX_SEED = np.iinfo(np.int32).max


def get_pipeline(model_label):
    if model_label in PIPELINES:
        return PIPELINES[model_label]

    cfg = MODEL_CONFIGS[model_label]

    pipe = DiffusionPipeline.from_pretrained(
        cfg["repo_id"],
        torch_dtype=torch_dtype,
        low_cpu_mem_usage=True,
    )

    pipe.to(device)
    pipe.enable_model_cpu_offload()

    PIPELINES[model_label] = pipe
    return pipe


def build_prompt(prompt, style):
    styles = {
        "Tanpa gaya": "",
        "Studio": "product photography, clean studio background, soft lighting",
        "E-commerce": "white background, catalog photo, sharp, high quality",
        "Pastel": "pastel colors, soft light, aesthetic instagram style",
        "Lifestyle": "realistic lifestyle photography, natural light",
    }
    suffix = styles.get(style, "")
    return f"{prompt}, {suffix}" if suffix else prompt


def infer(prompt, negative_prompt, seed, randomize_seed,
          width, height, guidance_scale, steps,
          model_label, style, num_images):

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)
    pipe = get_pipeline(model_label)

    full_prompt = build_prompt(prompt, style)

    images = []
    for _ in range(num_images):
        out = pipe(
            prompt=full_prompt,
            negative_prompt=negative_prompt or None,
            width=width,
            height=height,
            guidance_scale=guidance_scale,
            num_inference_steps=steps,
            generator=generator,
        )
        images.append(out.images[0])

    return images, seed


with gr.Blocks(title="RuangAI CPU Mode") as demo:
    gr.Markdown("# 🧴 RuangAI – CPU Mode Product Visualizer")

    with gr.Row():
        prompt = gr.Textbox(label="Prompt", placeholder="Deskripsi produk...")
        run_btn = gr.Button("Generate")

    with gr.Row():
        model_label = gr.Dropdown(
            list(MODEL_CONFIGS.keys()),
            value="SDXL 1.0 (CPU mode)",
            label="Model"
        )
        style = gr.Dropdown(
            ["Tanpa gaya", "Studio", "E-commerce", "Pastel", "Lifestyle"],
            value="Studio",
            label="Gaya visual"
        )
        num_images = gr.Slider(1, 3, value=1, step=1, label="Jumlah gambar")

    gallery = gr.Gallery(label="Hasil", columns=2, height=512)

    with gr.Accordion("Advanced", open=False):
        negative_prompt = gr.Textbox(label="Negative prompt")
        seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
        randomize_seed = gr.Checkbox(True, label="Randomize seed")
        width = gr.Slider(256, 768, value=512, step=32, label="Width")
        height = gr.Slider(256, 768, value=512, step=32, label="Height")
        guidance_scale = gr.Slider(0, 10, value=5, step=0.5, label="Guidance")
        steps = gr.Slider(5, 40, value=20, step=1, label="Steps")

    run_btn.click(
        infer,
        inputs=[
            prompt, negative_prompt, seed, randomize_seed,
            width, height, guidance_scale, steps,
            model_label, style, num_images
        ],
        outputs=[gallery, seed]
    )

demo.launch()