Spaces:
Running
on
Zero
Running
on
Zero
File size: 5,150 Bytes
55fe803 62594be 55fe803 eee0c8d 55fe803 62594be eee0c8d 55fe803 3176199 55fe803 c96ee5c bfe4467 55fe803 8da5cdc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
import torch
import spaces
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")
# ======== 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!")
@spaces.GPU
def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
"""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"],
]
# Build the Gradio interface
with gr.Blocks(title="Z-Image-Turbo Demo") as demo:
gr.Markdown(
"""
# 🎨 Z-Image-Turbo Demo
Generate high-quality images using the [Tongyi-MAI/Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) model.
This turbo model generates images in just 8 inference steps!
"""
)
with gr.Row():
with gr.Column(scale=1):
prompt = gr.Textbox(
label="Prompt",
placeholder="Enter your image description...",
lines=4,
)
with gr.Row():
height = gr.Slider(
minimum=512,
maximum=2048,
value=1024,
step=64,
label="Height",
)
width = gr.Slider(
minimum=512,
maximum=2048,
value=1024,
step=64,
label="Width",
)
with gr.Row():
num_inference_steps = gr.Slider(
minimum=1,
maximum=20,
value=9,
step=1,
label="Inference Steps",
info="9 steps results in 8 DiT forwards",
)
with gr.Row():
seed = gr.Number(
label="Seed",
value=42,
precision=0,
)
randomize_seed = gr.Checkbox(
label="Randomize Seed",
value=False,
)
generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")
with gr.Column(scale=1):
output_image = gr.Image(
label="Generated Image",
type="pil",
)
used_seed = gr.Number(
label="Seed Used",
interactive=False,
)
gr.Markdown("### 💡 Example Prompts")
gr.Examples(
examples=examples,
inputs=[prompt],
cache_examples=False,
)
gr.Markdown("Demo by [mrfakename](https://x.com/realmrfakename). Model by Alibaba. The model is licensed under Apache 2.0, you can use generated images commercially! Thanks to [multimodalart](https://huggingface.co/multimodalart) for the FA3 + AoTI enhancements/speedups")
# 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(mcp_server=True) |