Spaces:
Running on Zero
Running on Zero
| import spaces | |
| import random | |
| import re | |
| import torch | |
| import gradio as gr | |
| from diffusers import ZImagePipeline | |
| # ==================== Configuration ==================== | |
| MODEL_PATH = "Tongyi-MAI/Z-Image" # Full model with CFG support | |
| RESOLUTIONS = [ | |
| # Square | |
| "1024x1024 ( 1:1 )", | |
| # Landscape | |
| "1344x576 ( 21:9 )", | |
| "1280x720 ( 16:9 )", | |
| "1248x832 ( 3:2 )", | |
| "1152x864 ( 4:3 )", | |
| # Portrait | |
| "576x1344 ( 9:21 )", | |
| "720x1280 ( 9:16 )", | |
| "832x1248 ( 2:3 )", | |
| "864x1152 ( 3:4 )", | |
| ] | |
| EXAMPLE_PROMPTS = [ | |
| ["A golden retriever puppy sitting in a field of wildflowers at sunset, warm golden hour lighting, shallow depth of field, photorealistic."], | |
| ["A cozy coffee shop interior with exposed brick walls, vintage Edison bulbs hanging from the ceiling, wooden tables with steaming cups of latte art, rain visible through large windows, warm ambient lighting."], | |
| ["An ancient Japanese temple in autumn, surrounded by brilliant red and orange maple trees, morning mist rising from a koi pond in the foreground, traditional architecture with curved roofs, serene and peaceful atmosphere."], | |
| ["A futuristic cityscape at night, towering glass skyscrapers with holographic advertisements, flying cars leaving light trails, neon reflections on wet streets, cyberpunk aesthetic, highly detailed."], | |
| ["A macro photograph of a mechanical watch movement, intricate gears and springs in gold and silver, precise engineering details visible, dramatic lighting highlighting the craftsmanship, dark background, ultra high detail."], | |
| ["An astronaut floating in space with Earth visible in the background, stars and nebulae in the distance, realistic spacesuit with reflective visor showing Earth's reflection, cinematic lighting, awe-inspiring scale."], | |
| ] | |
| def get_resolution(resolution: str) -> tuple[int, int]: | |
| match = re.search(r"(\d+)\s*[×x]\s*(\d+)", resolution) | |
| if match: | |
| return int(match.group(1)), int(match.group(2)) | |
| return 1024, 1024 | |
| # ==================== Load Model ==================== | |
| print(f"Loading Z-Image pipeline from {MODEL_PATH}...") | |
| pipe = ZImagePipeline.from_pretrained( | |
| MODEL_PATH, | |
| torch_dtype=torch.bfloat16, | |
| ) | |
| pipe.to("cuda") | |
| print("Pipeline loaded successfully!") | |
| # ==================== Generation Function ==================== | |
| def generate( | |
| prompt: str, | |
| negative_prompt: str = "", | |
| resolution: str = "1024x1024 ( 1:1 )", | |
| seed: int = 42, | |
| steps: int = 28, | |
| guidance_scale: float = 4.0, | |
| cfg_normalization: bool = False, | |
| random_seed: bool = True, | |
| progress: gr.Progress = gr.Progress(track_tqdm=True), | |
| ): | |
| """ | |
| Generate an image using the Z-Image model. | |
| Args: | |
| prompt: Text prompt describing the desired image content. | |
| negative_prompt: What to avoid in the image (e.g., "blurry, low quality"). | |
| resolution: Output resolution in format "WIDTHxHEIGHT ( RATIO )". | |
| seed: Seed for reproducible generation. Ignored if random_seed is True. | |
| steps: Number of inference steps (28-50 recommended). | |
| guidance_scale: CFG scale controlling prompt adherence (3.0-5.0 recommended). | |
| cfg_normalization: Enable CFG normalization for more realistic outputs. | |
| random_seed: If True, a new random seed will be sampled. | |
| Returns: | |
| tuple: (image, seed_str, seed_int) | |
| """ | |
| if not prompt.strip(): | |
| raise gr.Error("Please enter a prompt.") | |
| if random_seed: | |
| new_seed = random.randint(1, 1000000) | |
| else: | |
| new_seed = seed if seed != -1 else random.randint(1, 1000000) | |
| width, height = get_resolution(resolution) | |
| generator = torch.Generator("cuda").manual_seed(new_seed) | |
| image = pipe( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt if negative_prompt.strip() else None, | |
| height=height, | |
| width=width, | |
| cfg_normalization=cfg_normalization, | |
| num_inference_steps=steps, | |
| guidance_scale=guidance_scale, | |
| generator=generator, | |
| ).images[0] | |
| return image, str(new_seed), int(new_seed) | |
| # ==================== Gradio UI ==================== | |
| with gr.Blocks(title="Z-Image Generation MCP") as demo: | |
| gr.Markdown( | |
| """<div align="center"> | |
| # Z-Image Generation MCP | |
| <a href="https://huggingface.co/settings/mcp?add=victor/Z-Image-MCP" target="_blank" | |
| style="padding: 10px 24px; background: #22c55e; color: white; text-decoration: none; border-radius: 9999px; font-weight: 600; font-size: 14px;"> | |
| Use via MCP | |
| </a> | |
| <a href="https://victor-z-image-mcp.hf.space/gradio_api/openapi.json" target="_blank" | |
| style="padding: 10px 24px; background: #3b82f6; margin-left: 12px; color: white; text-decoration: none; border-radius: 9999px; font-weight: 600; font-size: 14px;"> | |
| View API Schema | |
| </a> | |
| *An Efficient Image Generation Foundation Model with Single-Stream Diffusion Transformer* | |
| </div>""" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| prompt_input = gr.Textbox( | |
| label="Prompt", lines=3, placeholder="Enter your prompt here..." | |
| ) | |
| negative_prompt = gr.Textbox( | |
| label="Negative Prompt", | |
| lines=2, | |
| placeholder="What to avoid in the image (e.g., blurry, low quality)...", | |
| ) | |
| resolution = gr.Dropdown( | |
| value=RESOLUTIONS[0], | |
| choices=RESOLUTIONS, | |
| label="Resolution", | |
| ) | |
| with gr.Row(): | |
| seed = gr.Number(label="Seed", value=42, precision=0) | |
| random_seed = gr.Checkbox(label="Random Seed", value=True) | |
| with gr.Row(): | |
| steps = gr.Slider( | |
| label="Steps (28-50 recommended)", | |
| minimum=28, | |
| maximum=50, | |
| value=28, | |
| step=1, | |
| ) | |
| with gr.Row(): | |
| guidance_scale = gr.Slider( | |
| label="Guidance Scale (3-5 recommended)", | |
| minimum=3.0, | |
| maximum=5.0, | |
| value=4.0, | |
| step=0.5, | |
| ) | |
| cfg_normalization = gr.Checkbox( | |
| label="CFG Normalization", | |
| value=False, | |
| ) | |
| generate_btn = gr.Button("Generate", variant="primary") | |
| gr.Markdown("### Example Prompts") | |
| gr.Examples(examples=EXAMPLE_PROMPTS, inputs=prompt_input, label=None) | |
| with gr.Column(scale=1): | |
| output_image = gr.Image( | |
| label="Generated Image", | |
| format="png", | |
| height=600, | |
| interactive=False, | |
| ) | |
| used_seed = gr.Textbox(label="Seed Used", interactive=False) | |
| generate_btn.click( | |
| generate, | |
| inputs=[ | |
| prompt_input, | |
| negative_prompt, | |
| resolution, | |
| seed, | |
| steps, | |
| guidance_scale, | |
| cfg_normalization, | |
| random_seed, | |
| ], | |
| outputs=[output_image, used_seed, seed], | |
| api_name="generate", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(mcp_server=True) | |