app.py

"""
Z-Image-Turbo: Ultra-fast AI image generation powered by diffusion models.

A Gradio-based web application for generating images using the Tongyi-MAI/Z-Image-Turbo
model. Features a modern dark UI with purple/violet accent theme.
"""

import torch
import spaces
import gradio as gr
from diffusers import DiffusionPipeline


# =============================================================================
# CONFIGURATION CONSTANTS
# =============================================================================

ASPECT_RATIOS = {
    "1:1 Square": (1024, 1024),
    "16:9 Landscape": (1344, 768),
    "9:16 Portrait": (768, 1344),
    "4:3 Standard": (1152, 896),
    "3:4 Portrait": (896, 1152),
    "21:9 Cinematic": (1536, 640),
    "Custom": None,
}

EXAMPLE_PROMPTS = [
    ["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"],
    ["Cyberpunk street scene at night, neon signs reflecting in rain puddles, flying cars overhead, blade runner aesthetic"],
    ["Serene Japanese zen garden with cherry blossoms, koi pond, wooden bridge, morning mist, traditional painting style"],
    ["Steampunk mechanical owl with brass gears and glowing amber eyes, perched on old leather-bound books, warm candlelight"],
]

DEFAULT_SETTINGS = {
    "height": 1024,
    "width": 1024,
    "inference_steps": 9,
    "guidance_scale": 0.0,
    "seed": 42,
}


# =============================================================================
# THEME CONFIGURATION
# =============================================================================

def create_theme():
    """Create the modern dark theme with purple/violet accents."""
    return gr.themes.Base(
        primary_hue="violet",
        secondary_hue="purple",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter"),
        text_size="lg",
        spacing_size="md",
        radius_size="lg"
    ).set(
        # Dark mode colors
        body_background_fill="*neutral_950",
        body_background_fill_dark="*neutral_950",
        background_fill_primary="*neutral_900",
        background_fill_primary_dark="*neutral_900",
        background_fill_secondary="*neutral_800",
        background_fill_secondary_dark="*neutral_800",
        # Text colors
        body_text_color="*neutral_100",
        body_text_color_dark="*neutral_100",
        body_text_color_subdued="*neutral_400",
        body_text_color_subdued_dark="*neutral_400",
        # Border styling
        border_color_primary="*neutral_700",
        border_color_primary_dark="*neutral_700",
        block_border_width="1px",
        # Button styling
        button_primary_background_fill="linear-gradient(135deg, *primary_500 0%, *secondary_600 100%)",
        button_primary_background_fill_hover="linear-gradient(135deg, *primary_400 0%, *secondary_500 100%)",
        button_primary_text_color="white",
        button_primary_border_color="transparent",
        # Input styling
        input_background_fill="*neutral_800",
        input_background_fill_dark="*neutral_800",
        input_border_color="*neutral_600",
        input_border_color_dark="*neutral_600",
        input_border_color_focus="*primary_500",
        input_border_color_focus_dark="*primary_500",
        # Block styling
        block_background_fill="*neutral_900",
        block_background_fill_dark="*neutral_900",
        block_label_background_fill="*neutral_800",
        block_label_text_color="*neutral_200",
        block_title_text_weight="600",
        block_label_text_weight="500",
        # Shadow and depth
        shadow_drop="0 4px 20px rgba(0, 0, 0, 0.3)",
        shadow_drop_lg="0 8px 40px rgba(0, 0, 0, 0.4)",
    )


# =============================================================================
# CUSTOM CSS
# =============================================================================

CUSTOM_CSS = """
/* ===== ROOT VARIABLES ===== */
:root {
    --accent-primary: #8b5cf6;
    --accent-secondary: #a855f7;
    --accent-glow: rgba(139, 92, 246, 0.4);
    --bg-dark: #0a0a0f;
    --bg-card: #111118;
    --bg-elevated: #1a1a24;
    --text-primary: #f1f5f9;
    --text-secondary: #94a3b8;
    --text-muted: #64748b;
    --border-subtle: rgba(148, 163, 184, 0.1);
    --border-hover: rgba(139, 92, 246, 0.3);
    --gradient-primary: linear-gradient(135deg, #8b5cf6 0%, #a855f7 50%, #d946ef 100%);
    --gradient-hero: linear-gradient(135deg, #0f0f1a 0%, #1a1025 50%, #0f0f1a 100%);
}

/* ===== BASE STYLES ===== */
.gradio-container {
    max-width: 1400px !important;
    margin: 0 auto !important;
    background: var(--bg-dark) !important;
}

.dark {
    --background-fill-primary: var(--bg-dark) !important;
}

/* ===== HERO SECTION ===== */
.hero-section {
    position: relative;
    padding: 3rem 2rem;
    text-align: center;
    overflow: hidden;
    border-radius: 20px;
    margin-bottom: 2rem;
    background: var(--gradient-hero);
}

.hero-background {
    position: absolute;
    top: 0;
    left: 0;
    right: 0;
    bottom: 0;
    background:
        radial-gradient(ellipse at 20% 50%, rgba(139, 92, 246, 0.15) 0%, transparent 50%),
        radial-gradient(ellipse at 80% 50%, rgba(168, 85, 247, 0.15) 0%, transparent 50%),
        radial-gradient(ellipse at 50% 100%, rgba(217, 70, 239, 0.1) 0%, transparent 50%);
    animation: pulse-glow 8s ease-in-out infinite;
}

@keyframes pulse-glow {
    0%, 100% { opacity: 0.7; }
    50% { opacity: 1; }
}

.hero-content {
    position: relative;
    z-index: 1;
}

.logo-badge {
    display: inline-flex;
    align-items: center;
    justify-content: center;
    width: 64px;
    height: 64px;
    background: var(--gradient-primary);
    border-radius: 16px;
    margin-bottom: 1rem;
    box-shadow: 0 8px 32px var(--accent-glow);
}

.logo-icon {
    font-size: 2rem;
    font-weight: 800;
    color: white;
    font-family: 'Inter', sans-serif;
}

.hero-title {
    font-size: 3rem;
    font-weight: 800;
    margin: 0.5rem 0;
    background: var(--gradient-primary);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    letter-spacing: -0.02em;
}

.hero-subtitle {
    font-size: 1.1rem;
    color: var(--text-secondary);
    margin: 0.5rem 0 1.5rem;
    font-weight: 400;
}

.hero-stats {
    display: inline-flex;
    align-items: center;
    gap: 1.5rem;
    padding: 1rem 2rem;
    background: rgba(255, 255, 255, 0.03);
    border: 1px solid var(--border-subtle);
    border-radius: 100px;
    backdrop-filter: blur(10px);
}

.stat-item {
    display: flex;
    flex-direction: column;
    align-items: center;
    gap: 0.25rem;
}

.stat-value {
    font-size: 1.25rem;
    font-weight: 700;
    color: var(--accent-primary);
}

.stat-label {
    font-size: 0.75rem;
    color: var(--text-muted);
    text-transform: uppercase;
    letter-spacing: 0.05em;
}

.stat-divider {
    width: 1px;
    height: 32px;
    background: var(--border-subtle);
}

/* ===== MAIN CONTAINER ===== */
.main-container {
    gap: 2rem !important;
}

/* ===== SECTION HEADERS ===== */
.section-header {
    display: flex;
    align-items: center;
    gap: 0.5rem;
    font-size: 0.875rem;
    font-weight: 600;
    color: var(--text-secondary);
    text-transform: uppercase;
    letter-spacing: 0.08em;
    margin-bottom: 0.75rem;
    padding-left: 0.25rem;
}

.section-icon {
    font-size: 1rem;
    color: var(--accent-primary);
}

/* ===== CONTROLS COLUMN ===== */
.controls-column {
    padding: 0 !important;
}

/* ===== PROMPT INPUT ===== */
.prompt-input textarea {
    background: var(--bg-card) !important;
    border: 2px solid var(--border-subtle) !important;
    border-radius: 16px !important;
    padding: 1rem !important;
    font-size: 1rem !important;
    color: var(--text-primary) !important;
    transition: all 0.3s ease !important;
    min-height: 120px !important;
}

.prompt-input textarea:focus {
    border-color: var(--accent-primary) !important;
    box-shadow: 0 0 0 4px var(--accent-glow), 0 4px 20px rgba(0, 0, 0, 0.3) !important;
}

.prompt-input textarea::placeholder {
    color: var(--text-muted) !important;
}

/* ===== QUICK ACTIONS ===== */
.quick-actions {
    gap: 0.5rem !important;
    margin-top: 0.5rem !important;
    margin-bottom: 1.5rem !important;
}

.action-btn {
    flex: 1;
}

.action-btn button {
    width: 100% !important;
    background: var(--bg-elevated) !important;
    border: 1px solid var(--border-subtle) !important;
    color: var(--text-secondary) !important;
    border-radius: 10px !important;
    font-size: 0.875rem !important;
    padding: 0.5rem 1rem !important;
    transition: all 0.2s ease !important;
}

.action-btn button:hover {
    background: var(--bg-card) !important;
    border-color: var(--border-hover) !important;
    color: var(--text-primary) !important;
}

/* ===== ASPECT RATIO SELECTOR ===== */
.aspect-radio {
    margin-bottom: 1rem !important;
}

.aspect-radio .wrap {
    display: grid !important;
    grid-template-columns: repeat(4, 1fr) !important;
    gap: 0.5rem !important;
}

.aspect-radio label {
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    padding: 0.625rem 0.5rem !important;
    background: var(--bg-card) !important;
    border: 1px solid var(--border-subtle) !important;
    border-radius: 10px !important;
    cursor: pointer !important;
    transition: all 0.2s ease !important;
    font-size: 0.75rem !important;
    color: var(--text-secondary) !important;
    text-align: center !important;
}

.aspect-radio label:hover {
    border-color: var(--border-hover) !important;
    background: var(--bg-elevated) !important;
}

.aspect-radio label.selected,
.aspect-radio input:checked + label {
    background: var(--accent-primary) !important;
    border-color: var(--accent-primary) !important;
    color: white !important;
}

/* ===== DIMENSION SLIDERS ===== */
.dimension-row {
    gap: 1rem !important;
    margin-bottom: 1rem !important;
}

.dimension-slider {
    flex: 1;
}

.dimension-slider input[type="range"] {
    accent-color: var(--accent-primary) !important;
}

/* ===== ADVANCED ACCORDION ===== */
.advanced-accordion {
    margin-bottom: 1.5rem !important;
}

.advanced-accordion > .label-wrap {
    background: var(--bg-card) !important;
    border: 1px solid var(--border-subtle) !important;
    border-radius: 12px !important;
    padding: 0.75rem 1rem !important;
}

.advanced-accordion .icon {
    color: var(--accent-primary) !important;
}

/* ===== GENERATE BUTTON ===== */
.generate-btn button {
    width: 100% !important;
    background: var(--gradient-primary) !important;
    border: none !important;
    border-radius: 14px !important;
    padding: 1rem 2rem !important;
    font-size: 1.1rem !important;
    font-weight: 600 !important;
    color: white !important;
    cursor: pointer !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 20px var(--accent-glow) !important;
    position: relative;
    overflow: hidden;
}

.generate-btn button::before {
    content: '';
    position: absolute;
    top: 0;
    left: -100%;
    width: 100%;
    height: 100%;
    background: linear-gradient(90deg, transparent, rgba(255,255,255,0.2), transparent);
    transition: left 0.5s ease;
}

.generate-btn button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 32px var(--accent-glow) !important;
}

.generate-btn button:hover::before {
    left: 100%;
}

.generate-btn button:active {
    transform: translateY(0) !important;
}

/* ===== EXAMPLES GALLERY ===== */
#examples-gallery {
    margin-top: 0.5rem;
}

#examples-gallery .gallery {
    display: grid !important;
    grid-template-columns: repeat(2, 1fr) !important;
    gap: 0.5rem !important;
}

#examples-gallery .gallery button {
    background: var(--bg-card) !important;
    border: 1px solid var(--border-subtle) !important;
    border-radius: 10px !important;
    padding: 0.75rem !important;
    text-align: left !important;
    font-size: 0.75rem !important;
    color: var(--text-secondary) !important;
    line-height: 1.4 !important;
    transition: all 0.2s ease !important;
    overflow: hidden;
    text-overflow: ellipsis;
    display: -webkit-box;
    -webkit-line-clamp: 2;
    -webkit-box-orient: vertical;
}

#examples-gallery .gallery button:hover {
    border-color: var(--border-hover) !important;
    color: var(--text-primary) !important;
    background: var(--bg-elevated) !important;
}

/* ===== OUTPUT COLUMN ===== */
.output-column {
    background: var(--bg-card) !important;
    border-radius: 20px !important;
    padding: 1.5rem !important;
    border: 1px solid var(--border-subtle) !important;
}

.output-image {
    border-radius: 16px !important;
    overflow: hidden !important;
    background: var(--bg-elevated) !important;
    border: 1px solid var(--border-subtle) !important;
}

.output-image img {
    border-radius: 16px !important;
}

.output-image .icon-buttons {
    background: rgba(0, 0, 0, 0.6) !important;
    backdrop-filter: blur(10px) !important;
    border-radius: 10px !important;
}

/* ===== IMAGE INFO BAR ===== */
.image-info-bar {
    margin-top: 1rem !important;
    padding-top: 1rem !important;
    border-top: 1px solid var(--border-subtle) !important;
}

.seed-display input {
    background: var(--bg-elevated) !important;
    border: 1px solid var(--border-subtle) !important;
    border-radius: 8px !important;
    color: var(--text-secondary) !important;
    font-family: 'JetBrains Mono', monospace !important;
}

/* ===== FOOTER ===== */
.app-footer {
    margin-top: 2rem;
    padding: 1.5rem;
    border-top: 1px solid var(--border-subtle);
}

.footer-content {
    display: flex;
    justify-content: space-between;
    align-items: center;
    flex-wrap: wrap;
    gap: 1rem;
}

.footer-links {
    display: flex;
    align-items: center;
    gap: 1rem;
    flex-wrap: wrap;
}

.footer-link {
    display: inline-flex;
    align-items: center;
    gap: 0.375rem;
    color: var(--text-secondary);
    text-decoration: none;
    font-size: 0.875rem;
    transition: color 0.2s ease;
}

.footer-link:hover {
    color: var(--accent-primary);
}

.link-icon {
    font-size: 1rem;
}

.footer-divider {
    color: var(--text-muted);
}

.footer-text {
    color: var(--text-muted);
    font-size: 0.875rem;
}

.footer-badge {
    display: inline-block;
    padding: 0.375rem 0.75rem;
    background: var(--bg-elevated);
    border: 1px solid var(--border-subtle);
    border-radius: 100px;
    font-size: 0.75rem;
    color: var(--text-muted);
}

/* ===== RESPONSIVE DESIGN ===== */
@media (max-width: 768px) {
    .hero-section {
        padding: 2rem 1rem;
    }

    .hero-title {
        font-size: 2rem;
    }

    .hero-subtitle {
        font-size: 0.95rem;
    }

    .hero-stats {
        flex-wrap: wrap;
        padding: 0.75rem 1rem;
        gap: 1rem;
    }

    .stat-divider {
        display: none;
    }

    .main-container {
        flex-direction: column !important;
    }

    .output-column {
        order: -1;
    }

    .aspect-radio .wrap {
        grid-template-columns: repeat(2, 1fr) !important;
    }

    #examples-gallery .gallery {
        grid-template-columns: 1fr !important;
    }

    .footer-content {
        flex-direction: column;
        text-align: center;
    }
}

/* ===== LOADING STATE ===== */
.generating .output-image {
    position: relative;
}

.generating .output-image::after {
    content: '';
    position: absolute;
    top: 0;
    left: 0;
    right: 0;
    bottom: 0;
    background: linear-gradient(
        90deg,
        transparent,
        rgba(139, 92, 246, 0.1),
        transparent
    );
    animation: shimmer 1.5s infinite;
}

@keyframes shimmer {
    0% { transform: translateX(-100%); }
    100% { transform: translateX(100%); }
}

/* ===== SCROLLBAR ===== */
::-webkit-scrollbar {
    width: 8px;
    height: 8px;
}

::-webkit-scrollbar-track {
    background: var(--bg-dark);
}

::-webkit-scrollbar-thumb {
    background: var(--border-subtle);
    border-radius: 4px;
}

::-webkit-scrollbar-thumb:hover {
    background: var(--text-muted);
}

/* ===== ANIMATIONS ===== */
@keyframes fadeIn {
    from { opacity: 0; transform: translateY(10px); }
    to { opacity: 1; transform: translateY(0); }
}

.controls-column, .output-column {
    animation: fadeIn 0.5s ease-out;
}

.output-column {
    animation-delay: 0.1s;
}
"""


# =============================================================================
# HTML TEMPLATES
# =============================================================================

HERO_HTML = """
<div class="hero-section">
    <div class="hero-background"></div>
    <div class="hero-content">
        <div class="logo-badge">
            <span class="logo-icon">Z</span>
        </div>
        <h1 class="hero-title">Z-Image-Turbo</h1>
        <p class="hero-subtitle">Ultra-fast AI image generation powered by state-of-the-art diffusion</p>
        <div class="hero-stats">
            <div class="stat-item">
                <span class="stat-value">8</span>
                <span class="stat-label">DiT Steps</span>
            </div>
            <div class="stat-divider"></div>
            <div class="stat-item">
                <span class="stat-value">2K</span>
                <span class="stat-label">Max Resolution</span>
            </div>
            <div class="stat-divider"></div>
            <div class="stat-item">
                <span class="stat-value">&lt;3s</span>
                <span class="stat-label">Generation</span>
            </div>
        </div>
    </div>
</div>
"""

FOOTER_HTML = """
<footer class="app-footer">
    <div class="footer-content">
        <div class="footer-links">
            <a href="https://huggingface.co/Tongyi-MAI/Z-Image-Turbo" target="_blank" class="footer-link">
                <span class="link-icon">&#129303;</span> Model
            </a>
            <span class="footer-divider">|</span>
            <a href="https://x.com/realmrfakename" target="_blank" class="footer-link">
                <span class="link-icon">&#128104;&#8205;&#128187;</span> @mrfakename
            </a>
            <span class="footer-divider">|</span>
            <span class="footer-text">Redesigned with care</span>
        </div>
        <div class="footer-badge">
            Apache 2.0 License
        </div>
    </div>
</footer>
"""


# =============================================================================
# MODEL INITIALIZATION
# =============================================================================

def load_pipeline():
    """Load and configure the diffusion pipeline."""
    print("Loading Z-Image-Turbo pipeline...")
    pipeline = DiffusionPipeline.from_pretrained(
        "Tongyi-MAI/Z-Image-Turbo",
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=False,
    )
    pipeline.to("cuda")
    print("Pipeline loaded!")
    return pipeline


# Initialize pipeline at module level
pipe = load_pipeline()


# =============================================================================
# CORE FUNCTIONS
# =============================================================================

@spaces.GPU
def generate_image(
    prompt: str,
    height: int,
    width: int,
    num_inference_steps: int,
    seed: int,
    randomize_seed: bool,
    progress=gr.Progress(track_tqdm=True)
):
    """Generate an image from the given prompt.

    Args:
        prompt: Text description of the image to generate
        height: Output image height in pixels
        width: Output image width in pixels
        num_inference_steps: Number of denoising steps
        seed: Random seed for reproducibility
        randomize_seed: Whether to use a random seed
        progress: Gradio progress tracker

    Returns:
        Tuple of (generated image, seed used)
    """
    if not prompt or prompt.strip() == "":
        gr.Warning("Please enter a prompt first!")
        return None, None

    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=DEFAULT_SETTINGS["guidance_scale"],
        generator=generator,
    ).images[0]

    return image, seed


def update_dimensions(aspect_ratio: str):
    """Update height and width sliders based on selected aspect ratio.

    Args:
        aspect_ratio: Selected aspect ratio key from ASPECT_RATIOS

    Returns:
        Tuple of (height slider, width slider) with updated values
    """
    if aspect_ratio == "Custom":
        return gr.Slider(interactive=True), gr.Slider(interactive=True)

    dims = ASPECT_RATIOS.get(aspect_ratio)
    if dims:
        return (
            gr.Slider(value=dims[1], interactive=False),
            gr.Slider(value=dims[0], interactive=False)
        )
    return gr.Slider(interactive=True), gr.Slider(interactive=True)


def toggle_seed_visibility(randomize: bool):
    """Toggle seed input visibility based on randomize checkbox."""
    return gr.Number(visible=not randomize)


def clear_prompt():
    """Clear the prompt textbox."""
    return ""


def enhance_prompt(prompt_text: str) -> str:
    """Add quality enhancement keywords to a prompt.

    Args:
        prompt_text: Original prompt text

    Returns:
        Enhanced prompt with quality keywords added
    """
    if not prompt_text:
        return prompt_text

    enhancements = ", highly detailed, professional photography, 8k resolution, cinematic lighting, masterpiece"
    enhancement_keywords = ["detailed", "8k", "cinematic", "masterpiece"]

    if not any(keyword in prompt_text.lower() for keyword in enhancement_keywords):
        return prompt_text + enhancements
    return prompt_text


# =============================================================================
# UI BUILDER
# =============================================================================

def build_controls_column():
    """Build the left controls column with prompt input and settings."""
    with gr.Column(scale=4, min_width=340, elem_classes="controls-column"):
        # Prompt Section
        gr.HTML('<div class="section-header"><span class="section-icon">&#9998;</span> Describe Your Vision</div>')

        prompt = gr.Textbox(
            label="",
            placeholder="A mystical forest at twilight, bioluminescent mushrooms glowing softly, ethereal fog weaving between ancient trees...",
            lines=4,
            max_lines=8,
            autofocus=True,
            elem_classes="prompt-input",
            show_label=False,
        )

        # Quick Actions Row
        with gr.Row(elem_classes="quick-actions"):
            clear_btn = gr.Button("Clear", size="sm", variant="secondary", elem_classes="action-btn")
            enhance_btn = gr.Button("Enhance Prompt", size="sm", variant="secondary", elem_classes="action-btn")

        # Aspect Ratio Selection
        gr.HTML('<div class="section-header"><span class="section-icon">&#9634;</span> Aspect Ratio</div>')

        with gr.Row(elem_classes="aspect-ratio-grid"):
            aspect_ratio = gr.Radio(
                choices=list(ASPECT_RATIOS.keys()),
                value="1:1 Square",
                label="",
                show_label=False,
                elem_classes="aspect-radio",
            )

        # Dimension Controls
        with gr.Row(visible=True, elem_classes="dimension-row"):
            height = gr.Slider(
                minimum=512,
                maximum=2048,
                value=DEFAULT_SETTINGS["height"],
                step=64,
                label="Height",
                interactive=False,
                elem_classes="dimension-slider"
            )
            width = gr.Slider(
                minimum=512,
                maximum=2048,
                value=DEFAULT_SETTINGS["width"],
                step=64,
                label="Width",
                interactive=False,
                elem_classes="dimension-slider"
            )

        # Advanced Settings
        with gr.Accordion("Advanced Settings", open=False, elem_classes="advanced-accordion"):
            num_inference_steps = gr.Slider(
                minimum=1,
                maximum=20,
                value=DEFAULT_SETTINGS["inference_steps"],
                step=1,
                label="Quality Steps",
                info="Higher = better quality, slower generation (9 recommended)"
            )

            with gr.Row():
                randomize_seed = gr.Checkbox(
                    label="Random Seed",
                    value=True,
                    elem_classes="seed-checkbox"
                )
                seed = gr.Number(
                    label="Seed Value",
                    value=DEFAULT_SETTINGS["seed"],
                    precision=0,
                    visible=False,
                    elem_classes="seed-input"
                )

        # Generate Button
        generate_btn = gr.Button(
            "Generate Image",
            variant="primary",
            size="lg",
            elem_classes="generate-btn",
        )

        # Examples Section
        gr.HTML('<div class="section-header"><span class="section-icon">&#10024;</span> Inspiration Gallery</div>')

        gr.Examples(
            examples=EXAMPLE_PROMPTS,
            inputs=[prompt],
            label="",
            examples_per_page=4,
            elem_id="examples-gallery",
        )

    return {
        "prompt": prompt,
        "clear_btn": clear_btn,
        "enhance_btn": enhance_btn,
        "aspect_ratio": aspect_ratio,
        "height": height,
        "width": width,
        "num_inference_steps": num_inference_steps,
        "randomize_seed": randomize_seed,
        "seed": seed,
        "generate_btn": generate_btn,
    }


def build_output_column():
    """Build the right output column with generated image display."""
    with gr.Column(scale=5, min_width=400, elem_classes="output-column"):
        gr.HTML('<div class="section-header"><span class="section-icon">&#128444;</span> Generated Artwork</div>')

        output_image = gr.Image(
            label="",
            type="pil",
            show_label=False,
            height=580,
            elem_classes="output-image",
            show_download_button=True,
            show_share_button=True,
        )

        # Image Info Bar
        with gr.Row(elem_classes="image-info-bar"):
            used_seed = gr.Number(
                label="Seed",
                interactive=False,
                elem_classes="seed-display",
                scale=1,
            )

    return {
        "output_image": output_image,
        "used_seed": used_seed,
    }


def create_app():
    """Create and configure the Gradio application."""
    theme = create_theme()

    with gr.Blocks(fill_height=True, title="Z-Image-Turbo | AI Image Generator") as demo:
        # Hero Header
        gr.HTML(HERO_HTML)

        with gr.Row(equal_height=False, elem_classes="main-container"):
            # Build UI sections
            controls = build_controls_column()
            outputs = build_output_column()

        # Footer
        gr.HTML(FOOTER_HTML)

        # Wire up event handlers
        controls["randomize_seed"].change(
            toggle_seed_visibility,
            inputs=[controls["randomize_seed"]],
            outputs=[controls["seed"]]
        )

        controls["aspect_ratio"].change(
            update_dimensions,
            inputs=[controls["aspect_ratio"]],
            outputs=[controls["height"], controls["width"]]
        )

        controls["clear_btn"].click(
            clear_prompt,
            outputs=[controls["prompt"]]
        )

        controls["enhance_btn"].click(
            enhance_prompt,
            inputs=[controls["prompt"]],
            outputs=[controls["prompt"]]
        )

        # Generation triggers
        generation_inputs = [
            controls["prompt"],
            controls["height"],
            controls["width"],
            controls["num_inference_steps"],
            controls["seed"],
            controls["randomize_seed"],
        ]
        generation_outputs = [outputs["output_image"], outputs["used_seed"]]

        controls["generate_btn"].click(
            fn=generate_image,
            inputs=generation_inputs,
            outputs=generation_outputs,
        )

        controls["prompt"].submit(
            fn=generate_image,
            inputs=generation_inputs,
            outputs=generation_outputs,
        )

    return demo, theme


# =============================================================================
# MAIN ENTRY POINT
# =============================================================================

# Create the application
demo, custom_theme = create_app()

if __name__ == "__main__":
    demo.launch(
        theme=custom_theme,
        css=CUSTOM_CSS,
    )