import torch
import gradio as gr
from diffusers import DiffusionPipeline

# Load the pipeline once at startup
print("🚀 Loading Z-Image-Turbo pipeline...")
pipe = DiffusionPipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=False,
)
pipe.to("cuda")

print("✅ Pipeline loaded!")

def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed):
    """Generate an image from the given prompt."""
    if randomize_seed:
        seed = torch.randint(0, 2**32 - 1, (1,)).item()
    
    generator = torch.Generator("cuda").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,  # Guidance should be 0 for Turbo models
        generator=generator,
    ).images[0]
    
    return image, seed

# Example prompts
examples = [
    ["Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp, bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda, blurred colorful distant lights."],
    ["A majestic dragon soaring through clouds at sunset, scales shimmering with iridescent colors, detailed fantasy art style"],
    ["Cozy coffee shop interior, warm lighting, rain on windows, plants on shelves, vintage aesthetic, photorealistic"],
    ["Astronaut riding a horse on Mars, cinematic lighting, sci-fi concept art, highly detailed"],
    ["Portrait of a wise old wizard with a long white beard, holding a glowing crystal staff, magical forest background"],
]

# Modern Gradio 6 Blocks with minimal design
with gr.Blocks(
    title="Z-Image-Turbo",
    fill_height=True,
    footer_links=[{"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}]
) as demo:
    
    # Header with clean, modern design
    with gr.Row():
        gr.Markdown("""
        # 🎨 Z-Image-Turbo
        
        Generate high-quality images in just 8 inference steps using the [Tongyi-MAI/Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) model.
        """)
    
    # Main content area
    with gr.Row():
        with gr.Column(scale=1, min_width=320):
        # Prompt input - clean and focused
        prompt = gr.Textbox(
            label="✨ Image Description",
            placeholder="Describe the image you want to generate...",
            lines=3,
            max_lines=6,
            autofocus=True,
        )
        
        # Compact settings in an accordion
        with gr.Accordion("⚙️ Advanced Settings", open=False):
            with gr.Row():
                height = gr.Slider(
                    label="📏 Height",
                    minimum=512,
                    maximum=2048,
                    value=1024,
            )
            with gr.Row():
                width = gr.Slider(
                    label="📐 Width",
                    minimum=512,
                    maximum=2048,
                    value=1024,
                )
            
            with gr.Row():
                num_inference_steps = gr.Slider(
                    label="⚡ Inference Steps",
                    minimum=1,
                    maximum=20,
                    value=9,
                )
            
            with gr.Row():
                seed = gr.Number(
                    label="🌱 Seed",
                    value=42,
                    precision=0,
                )
                randomize_seed = gr.Checkbox(
                    label="🎲 Randomize Seed",
                    value=False,
                )
        
        # Generate button - prominent and clear
        generate_btn = gr.Button(
            "🚀 Generate Image",
            variant="primary",
            size="lg",
            scale=1,
        )
        
        with gr.Column(scale=1, min_width=320):
            output_image = gr.Image(
                label="🎨 Generated Image",
                type="pil",
            )
    
    # Examples section
    with gr.Row():
        gr.Markdown("### 💡 Try These Examples")
        gr.Examples(
            examples=examples,
            inputs=prompt,
            cache_examples=False,
        )
    
    # Connect the generate button
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
        outputs=[output_image, seed],
        api_visibility="public",
    )
    
    # Also allow generating by pressing Enter in the prompt box
    prompt.submit(
        fn=generate_image,
        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
        outputs=[output_image, seed],
        api_visibility="public",
    )

if __name__ == "__main__":
    demo.launch(
        mcp_server=True,
        show_error=True,
        share=False,
    )