app.py

import torch
import spaces
import gradio as gr
from diffusers import ZImagePipeline, ZImageTransformer2DModel
import numpy as np
import random

# Load the pipeline once at startup
print("Loading Z-Image-Turbo pipeline...")
transformer = ZImageTransformer2DModel.from_pretrained(
    "linoyts/beyond-reality-z-image-diffusers",
    torch_dtype=torch.bfloat16,
)

# Load pipeline with custom transformer
pipe = ZImagePipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    transformer=transformer,
    torch_dtype=torch.bfloat16,
)
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("Pipeline loaded!")
MAX_SEED = np.iinfo(np.int32).max



@spaces.GPU
def generate_image(prompt, height, width, num_inference_steps=10, seed=42, randomize_seed=True, progress=gr.Progress(track_tqdm=True)):
    """Generate an image from the given prompt."""
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    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 = [
    ["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"],
]


css = """
#col-container { max-width: 1000px; margin: 0 auto; }
.dark .progress-text { color: white !important; }
#examples { max-width: 1000px; margin: 0 auto; }
.dark #examples button,
.dark #examples .example,
.dark #examples span {
    color: white !important;
}
"""

# Build the Gradio interface
with gr.Blocks(fill_height=True) as demo:
    # Header
    gr.Markdown(
        """
        # Z-Image-Turbo BEYOND REALITY
        **Ultra-fast AI image generation** with [Nurburgring/BEYOND_REALITY_Z_IMAGE](https://huggingface.co/Nurburgring/BEYOND_REALITY_Z_IMAGE) - Z Image Turbo fine-tuned for enhanced texture fidelity, analog photography alignment & improved balance of realism & aesthetics.
        """,
        elem_classes="header-text"
    )
    
    with gr.Row(equal_height=False):
        # Left column - Input controls
        with gr.Column(scale=1, min_width=320,elem_id="col-container"):
            prompt = gr.Textbox(
                label="✨ Your Prompt",
                placeholder="Describe the image you want to create...",
                lines=5,
                max_lines=10,
                autofocus=True,
            )
            
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                with gr.Row():
                    height = gr.Slider(
                        minimum=512,
                        maximum=2048,
                        value=1024,
                        step=64,
                        label="Height",
                        info="Image height in pixels"
                    )
                    width = gr.Slider(
                        minimum=512,
                        maximum=2048,
                        value=1024,
                        step=64,
                        label="Width",
                        info="Image width in pixels"
                    )
                
                num_inference_steps = gr.Slider(
                    minimum=1,
                    maximum=20,
                    value=10,
                    step=1,
                    label="Inference Steps",
                    info="9 steps = 8 DiT forwards (recommended)"
                )
                
                with gr.Row():
                    randomize_seed = gr.Checkbox(
                        label="🎲 Random Seed",
                        value=True,
                    )
                    seed = gr.Number(
                        label="Seed",
                        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(
                "🚀 Generate Image",
                variant="primary",
                size="lg",
                scale=1
            )
            
            # Example prompts
            gr.Examples(
                examples=examples,
                inputs=[prompt],
                label="💡 Try these prompts",
                examples_per_page=5,
            )
        
        # Right column - Output
        with gr.Column(scale=1, min_width=320):
            output_image = gr.Image(
                label="Generated Image",
                type="pil",
                show_label=False,
                height=600,
                buttons=["download", "share"],
            )
            
            used_seed = gr.Number(
                label="🎲 Seed Used",
                interactive=False,
                container=True,
            )

    
    # Connect the generate button
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
        outputs=[output_image, used_seed],
    )
    
    # 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, used_seed],
    )

if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Citrus(),
        css=css,
        footer_links=[
            "api",
            "gradio"
        ]
    )