import torch
import gradio as gr
from diffusers import DiffusionPipeline

# ================== DEVICE ==================
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# ================== LOAD PIPELINE ==================
print("Loading Z-Image-Turbo pipeline...")
pipe = DiffusionPipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    low_cpu_mem_usage=False,
)
pipe = pipe.to(device)
print("Pipeline loaded!")

# ================== IMAGE GENERATION ==================
def generate_image(
    prompt,
    height,
    width,
    num_inference_steps,
    seed,
    randomize_seed,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = torch.randint(0, 2**32 - 1, (1,)).item()

    generator = torch.Generator(device).manual_seed(int(seed))

    image = pipe(
        prompt=prompt,
        height=int(height),
        width=int(width),
        num_inference_steps=int(num_inference_steps),
        guidance_scale=0.0,
        generator=generator,
    ).images[0]

    return image, seed


# ================== EXAMPLES ==================
examples = [
    ["Young Chinese woman in red Hanfu, intricate embroidery, neon lightning lamp floating above palm, cinematic lighting"],
    ["A majestic dragon soaring through clouds at sunset, fantasy art, ultra detailed"],
    ["Cozy coffee shop interior, rain on windows, warm light, photorealistic"],
    ["Astronaut riding a horse on Mars, cinematic sci-fi concept art"],
    ["Portrait of an old wizard with glowing staff, magical forest"],
]

# ================== THEME ==================
custom_theme = gr.themes.Soft(
    primary_hue="yellow",
    secondary_hue="amber",
    neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter"),
    text_size="lg",
    spacing_size="md",
    radius_size="lg",
)

# ================== UI ==================
with gr.Blocks(fill_height=True, theme=custom_theme) as demo:
    gr.Markdown(
        """
        # 🤖 Burak Image
        **Ultra-fast AI image generation** • CPU / GPU Auto
        """
    )

    with gr.Row():
        with gr.Column(scale=1, min_width=320):
            prompt = gr.Textbox(
                label="✨ Prompt",
                placeholder="Describe the image you want...",
                lines=5,
            )

            with gr.Accordion("⚙️ Advanced Settings", open=False):
                with gr.Row():
                    height = gr.Slider(512, 2048, 1024, step=64, label="Height")
                    width = gr.Slider(512, 2048, 1024, step=64, label="Width")

                num_inference_steps = gr.Slider(
                    1, 20, 9, step=1, label="Inference Steps"
                )

                randomize_seed = gr.Checkbox(
                    label="🎲 Random Seed", value=True
                )
                seed = gr.Number(
                    label="Seed", value=42, precision=0, visible=False
                )

                randomize_seed.change(
                    lambda x: gr.Number(visible=not x),
                    randomize_seed,
                    seed,
                )

            generate_btn = gr.Button(
                "🚀 Generate Image",
                variant="primary",
                size="lg",
            )

            gr.Examples(
                examples=examples,
                inputs=[prompt],
                label="💡 Example Prompts",
            )

        with gr.Column(scale=1, min_width=320):
            output_image = gr.Image(
                label="Generated Image",
                type="pil",
                height=600,
                buttons=["download", "share"],
            )

            used_seed = gr.Number(
                label="🎲 Seed Used",
                interactive=False,
            )

    generate_btn.click(
        generate_image,
        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
        outputs=[output_image, used_seed],
    )

    prompt.submit(
        generate_image,
        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
        outputs=[output_image, used_seed],
    )

# ================== LAUNCH ==================
if __name__ == "__main__":
    demo.launch()