app.py

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

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

# ======== AoTI compilation + FA3 ========
# pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
# spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")

print("Neural core online. Synaptic links established.")

@spaces.GPU
def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
    """Jack in. Generate."""
    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,
        generator=generator,
    ).images[0]
    
    return image, seed

# Example prompts
examples = [
    ["A highly detailed cyberpunk street scene in Tokyo at night after rain. Pink and cyan neon signs reflecting in wet asphalt puddles. A flying food stall in the foreground with rising steam. Cinematic look, 8k resolution, soft bokeh background, hyper-realistic."],
    ["A floating island in a sea of clouds at sunset. A large, glowing ancient tree with golden leaves stands on the island. Small waterfalls cascading from the cliffs into the void. Soft brushstroke digital art style, magical lighting, pastel colors, dreamy atmosphere."],
    ["Macro photography of a single water droplet on a deep green fern leaf. An entire forest is reflected inside the droplet. Extremely shallow depth of field, intricate leaf texture, natural sunlight filtering through the canopy, 100mm macro lens aesthetic."],
    ["Isometric 3D miniature of a cozy reading room. A small fireplace burning, a comfortable armchair nearby, and walls lined with bookshelves. Octane render style, soft clay textures, warm volumetric lighting, minimalist and cute."],
    ["A dramatic oil painting in the style of Caravaggio. A lone astronomer sitting at a wooden table, studying an ancient star map by candlelight. Strong chiaroscuro, visible brushstrokes, rich textures of parchment and heavy velvet, deep gold and brown tones."],
]


CYBERPUNK_CSS = """
/* ============================================
   BRAINDANCE FACTORY — NEURAL INTERFACE v2.0.77
   ============================================ */

@import url('https://fonts.googleapis.com/css2?family=Share+Tech+Mono&family=Orbitron:wght@400;700;900&family=Rajdhani:wght@300;400;600&display=swap');

/* Root & Scanlines */
.gradio-container {
    max-width: 1500px !important;
    margin: 0 auto !important;
    background: #050811 !important;
    position: relative;
}

.gradio-container::before {
    content: '';
    position: fixed;
    top: 0; left: 0; right: 0; bottom: 0;
    background: repeating-linear-gradient(
        0deg,
        transparent,
        transparent 2px,
        rgba(0, 245, 255, 0.015) 2px,
        rgba(0, 245, 255, 0.015) 4px
    );
    pointer-events: none;
    z-index: 9999;
}

/* Header */
.bd-header {
    font-family: 'Orbitron', monospace;
    text-align: center;
    padding: 2.5rem 1rem 1.5rem;
    position: relative;
    border-bottom: 1px solid #00f5ff18;
    margin-bottom: 1.5rem;
    overflow: hidden;
}

.bd-header::after {
    content: '';
    position: absolute;
    bottom: 0; left: 10%; right: 10%;
    height: 1px;
    background: linear-gradient(90deg, transparent, #00f5ff, #ff00c8, #00f5ff, transparent);
    animation: scanline-sweep 4s linear infinite;
}

@keyframes scanline-sweep {
    0% { opacity: 0.3; } 50% { opacity: 1; } 100% { opacity: 0.3; }
}

.bd-header h1 {
    font-size: clamp(2rem, 5vw, 3.8rem) !important;
    font-weight: 900 !important;
    letter-spacing: 0.12em;
    margin: 0 0 0.3rem !important;
    color: #00f5ff !important;
    text-shadow:
        0 0 10px #00f5ff,
        0 0 30px #00f5ffaa,
        0 0 60px #00f5ff44;
    -webkit-text-fill-color: unset !important;
    background: none !important;
    animation: flicker 6s infinite;
}

@keyframes flicker {
    0%, 95%, 100% { opacity: 1; }
    96% { opacity: 0.7; }
    97% { opacity: 1; }
    98% { opacity: 0.4; }
    99% { opacity: 1; }
}

.bd-header .glitch-sub {
    font-family: 'Rajdhani', sans-serif;
    font-size: 1rem;
    letter-spacing: 0.35em;
    text-transform: uppercase;
    color: #ff00c8 !important;
    opacity: 0.85;
    text-shadow: 0 0 8px #ff00c8aa;
}

.bd-header .glitch-sub span {
    color: #00f5ff88;
    margin: 0 0.5em;
}

/* Tagline badge */
.bd-tagline {
    display: inline-block;
    font-family: 'Share Tech Mono', monospace;
    font-size: 0.72rem;
    color: #00f5ff66;
    border: 1px solid #00f5ff22;
    padding: 0.2em 0.8em;
    margin-top: 0.8rem;
    letter-spacing: 0.15em;
    background: #00f5ff08;
}

/* Panel blocks */
.block, .gr-block, .gr-panel {
    border: 1px solid #00f5ff1a !important;
    border-radius: 2px !important;
    background: #0a0f1e !important;
    box-shadow: 0 0 20px #00f5ff08, inset 0 0 20px #00000033 !important;
}

/* Corner decorators on panels */
.block::before, .block::after {
    content: '';
    position: absolute;
    width: 8px; height: 8px;
    border-color: #00f5ff55;
    border-style: solid;
}
.block::before {
    top: -1px; left: -1px;
    border-width: 1px 0 0 1px;
}
.block::after {
    bottom: -1px; right: -1px;
    border-width: 0 1px 1px 0;
}

/* Generate button — the big CTA */
.generate-btn button, button[class*="primary"] {
    font-family: 'Orbitron', monospace !important;
    font-size: 0.95rem !important;
    font-weight: 700 !important;
    letter-spacing: 0.2em !important;
    text-transform: uppercase !important;
    background: linear-gradient(135deg, #00f5ff, #00b8cc) !important;
    color: #020710 !important;
    border: none !important;
    border-radius: 2px !important;
    padding: 0.9rem 1.5rem !important;
    position: relative;
    overflow: hidden;
    transition: all 0.25s ease !important;
    box-shadow: 0 0 20px #00f5ff55, 0 0 40px #00f5ff22 !important;
    cursor: pointer;
}

.generate-btn button::before, button[class*="primary"]::before {
    content: '';
    position: absolute;
    top: 0; left: -100%;
    width: 100%; height: 100%;
    background: linear-gradient(90deg, transparent, #ffffff33, transparent);
    transition: left 0.4s ease;
}

.generate-btn button:hover::before, button[class*="primary"]:hover::before {
    left: 100%;
}

.generate-btn button:hover, button[class*="primary"]:hover {
    background: linear-gradient(135deg, #ff00c8, #cc0099) !important;
    box-shadow: 0 0 25px #ff00c888, 0 0 50px #ff00c833 !important;
    transform: translateY(-2px) !important;
    color: #fff !important;
}

/* Input fields */
textarea, input[type="text"], input[type="number"] {
    font-family: 'Rajdhani', sans-serif !important;
    font-size: 1.05rem !important;
    color: #c8e0ef !important;
    background: #0d1526 !important;
    border: 1px solid #00f5ff22 !important;
    border-radius: 2px !important;
    caret-color: #00f5ff;
    transition: border-color 0.2s, box-shadow 0.2s !important;
}

textarea:focus, input:focus {
    border-color: #00f5ff88 !important;
    box-shadow: 0 0 12px #00f5ff22 !important;
    outline: none !important;
}

textarea::placeholder {
    color: #2a4055 !important;
    font-style: italic;
}

/* Labels */
label span, .gr-block label {
    font-family: 'Share Tech Mono', monospace !important;
    font-size: 0.78rem !important;
    letter-spacing: 0.12em !important;
    text-transform: uppercase !important;
    color: #00f5ffaa !important;
}

/* Sliders */
input[type="range"] {
    accent-color: #00f5ff !important;
}

input[type="range"]::-webkit-slider-thumb {
    background: #00f5ff !important;
    box-shadow: 0 0 8px #00f5ff !important;
}

input[type="range"]::-webkit-slider-runnable-track {
    background: linear-gradient(to right, #00f5ff, #0d1526) !important;
}

/* Accordion */
.gr-accordion > .label-wrap {
    border-bottom: 1px solid #00f5ff18 !important;
}

.gr-accordion {
    border: 1px solid #00f5ff1a !important;
    background: #070c1a !important;
    border-radius: 2px !important;
}

/* Image output */
.gr-image img, .output-image img {
    border: 1px solid #00f5ff33 !important;
    border-radius: 2px !important;
    box-shadow: 0 0 30px #00f5ff18, 0 0 60px #00000066 !important;
    transition: box-shadow 0.4s ease;
}

.gr-image img:hover, .output-image img:hover {
    box-shadow: 0 0 40px #00f5ff44, 0 0 80px #ff00c822 !important;
}

/* Examples — catch-all for every Gradio version */
.examples, .examples *,
.svelte-examples, .svelte-examples *,
[class*="examples"], [class*="examples"] *,
[class*="example-"] button, [class*="example-"] td,
#examples-block, #examples-block *,
.wrap > .examples_table,
.wrap > .examples_table * {
    color: #b8d8e8 !important;
    background-color: #0a1220 !important;
    font-family: 'Rajdhani', sans-serif !important;
    font-weight: 600 !important;
}

[class*="examples"] button:hover,
[class*="example-"] button:hover {
    background-color: #0d1e30 !important;
    color: #00f5ff !important;
    text-shadow: 0 0 8px #00f5ff77 !important;
}

/* Also catch bare table cells anywhere in the UI */
td {
    color: #b8d8e8 !important;
    background-color: #0a1220 !important;
    font-family: 'Rajdhani', sans-serif !important;
    font-weight: 600 !important;
    font-size: 0.9rem !important;
    border-color: #00f5ff15 !important;
    padding: 0.4em 0.7em !important;
}

tr:hover td {
    background-color: #0d1e30 !important;
    color: #00f5ff !important;
}

/* Footer */
.bd-footer {
    font-family: 'Share Tech Mono', monospace;
    text-align: center;
    padding: 1.5rem 0 0.5rem;
    border-top: 1px solid #00f5ff12;
    margin-top: 1rem;
    font-size: 0.75rem;
    color: #2a4055;
    letter-spacing: 0.08em;
}

.bd-footer a {
    color: #00f5ff66 !important;
    text-decoration: none !important;
    transition: color 0.2s, text-shadow 0.2s;
}

.bd-footer a:hover {
    color: #ff00c8 !important;
    text-shadow: 0 0 8px #ff00c888 !important;
}

.bd-footer .sep {
    color: #00f5ff22;
    margin: 0 0.6em;
}

/* Status dot */
.status-dot {
    display: inline-block;
    width: 6px; height: 6px;
    border-radius: 50%;
    background: #00f5ff;
    box-shadow: 0 0 6px #00f5ff;
    animation: pulse-dot 1.8s infinite;
    vertical-align: middle;
    margin-right: 0.4em;
}

@keyframes pulse-dot {
    0%, 100% { opacity: 1; box-shadow: 0 0 6px #00f5ff; }
    50% { opacity: 0.4; box-shadow: 0 0 2px #00f5ff; }
}

/* Scrollbar */
::-webkit-scrollbar { width: 4px; }
::-webkit-scrollbar-track { background: #050811; }
::-webkit-scrollbar-thumb { background: #00f5ff33; border-radius: 2px; }
::-webkit-scrollbar-thumb:hover { background: #00f5ff66; }

/* Mobile */
@media (max-width: 768px) {
    .bd-header h1 { font-size: 1.8rem !important; }
}
"""

# ─── Build the Gradio Interface ───────────────────────────────────────────────
with gr.Blocks() as demo:

    # ── Header ──
    gr.HTML("""
        <div class="bd-header">
            <h1>BRAINDANCE FACTORY</h1>
            <div class="glitch-sub">
                Neural Image Synthesis
                <span>///</span>
                Z-Image-Turbo Core
                <span>///</span>
                8-Step Protocol
            </div>
            <div class="bd-tagline">
                <span class="status-dot"></span>
                SYNAPTIC UPLINK ESTABLISHED — JACK IN TO BEGIN
            </div>
        </div>
    """)

    with gr.Row(equal_height=False):
        # ── Left: Controls ──
        with gr.Column(scale=1, min_width=320):

            prompt = gr.Textbox(
                label="⬡ NEURAL PROMPT INPUT",
                placeholder="Inject your vision into the braindance matrix...",
                lines=5,
                max_lines=10,
                autofocus=True,
            )

            with gr.Accordion("⬡ SYSTEM PARAMETERS", open=False):
                with gr.Row():
                    height = gr.Slider(
                        minimum=512, maximum=2048, value=1024, step=64,
                        label="HEIGHT  [px]",
                        info="Vertical resolution"
                    )
                    width = gr.Slider(
                        minimum=512, maximum=2048, value=1024, step=64,
                        label="WIDTH  [px]",
                        info="Horizontal resolution"
                    )

                num_inference_steps = gr.Slider(
                    minimum=1, maximum=20, value=9, step=1,
                    label="INFERENCE CYCLES",
                    info="9 steps = 8 DiT forwards (optimal)"
                )

                with gr.Row():
                    randomize_seed = gr.Checkbox(
                        label="⬡ RANDOMIZE SEED",
                        value=True,
                    )
                    seed = gr.Number(
                        label="SEED NODE",
                        value=42,
                        precision=0,
                        visible=False,
                    )

                def toggle_seed(randomize):
                    return gr.Number(visible=not randomize)

                randomize_seed.change(
                    toggle_seed,
                    inputs=[randomize_seed],
                    outputs=[seed]
                )

            generate_btn = gr.Button(
                "⬡ INITIATE BRAINDANCE",
                variant="primary",
                size="lg",
                elem_classes=["generate-btn"],
            )

            gr.Examples(
                examples=examples,
                inputs=[prompt],
                label="⬡ MEMORY FRAGMENTS — SAMPLE INPUTS",
                examples_per_page=5,
                elem_id="examples-block",
            )

        # ── Right: Output ──
        with gr.Column(scale=1, min_width=320):
            output_image = gr.Image(
                label="SYNTHESIZED BRAINDANCE OUTPUT",
                type="pil",
                format="png",
                show_label=True,
                height=600,
            )
            used_seed = gr.Number(
                label="⬡ SEED NODE USED",
                interactive=False,
                container=True,
            )

    # ── Footer ──
    gr.HTML("""
        <div class="bd-footer">
            <span class="status-dot"></span>
            MODEL: <a href="https://huggingface.co/Tongyi-MAI/Z-Image-Turbo" target="_blank">TONGYI-MAI / Z-IMAGE-TURBO</a>
            <span class="sep">//</span>
            APACHE 2.0 LICENSE
            <span class="sep">//</span>
            DEMO: <a href="https://x.com/realmrfakename" target="_blank">@MRFAKENAME</a>
            <span class="sep">//</span>
            OPTIMIZATIONS: <a href="https://huggingface.co/multimodalart" target="_blank">@MULTIMODALART</a>
            <span class="sep">//</span>
            BRAINDANCE SKIN: ANYCODER
            <br><br>
            &gt;&gt; ALL MEMORIES ARE PROPERTY OF THE FACTORY. UNAUTHORIZED DREAMING IS ENCOURAGED. &lt;&lt;
        </div>
    """)

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

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

if __name__ == "__main__":
    demo.launch(
        css=CYBERPUNK_CSS,
        mcp_server=True,
        footer_links=["api", "gradio"],
    )