app.py

"""
Hugging Face Image Generator - ZeroGPU Version
Simple version with Z-Image-Turbo
"""

import spaces
import torch
from typing import Tuple
from PIL import Image

import gradio as gr
from diffusers import ZImagePipeline

MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
pipe = None


@spaces.GPU(duration=180)
def generate(prompt, negative_prompt, steps, width, height, seed, guidance_scale):
    global pipe
    
    if not prompt:
        return None, "Please enter a prompt"
    
    try:
        if pipe is None:
            print("Loading model...")
            pipe = ZImagePipeline.from_pretrained(
                MODEL_ID,
                torch_dtype=torch.bfloat16,
                low_cpu_mem_usage=False,
            )
            pipe.to("cuda")
            print("Model loaded!")
        
        # Parse parameters
        try:
            steps_int = int(steps) if steps else 9
        except:
            steps_int = 9
        
        try:
            seed_int = int(seed) if seed else 0
        except:
            seed_int = 0
            
        try:
            guidance_float = float(guidance_scale) if guidance_scale else 0.0
        except:
            guidance_float = 0.0
        
        # Seed
        generator = None
        if seed_int > 0:
            generator = torch.Generator("cuda").manual_seed(seed_int)
        
        # Generate
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt if negative_prompt else None,
            num_inference_steps=steps_int,
            guidance_scale=guidance_float,
            width=width,
            height=height,
            generator=generator,
        )
        
        image = result.images[0]
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        return image, "Success!"
        
    except Exception as e:
        err = str(e)
        return None, "Error: " + err


with gr.Blocks(title="AI Image Generator") as demo:
    gr.Markdown("# 🎨 AI Image Generator")
    gr.Markdown("Z-Image-Turbo - 1024x1024")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", placeholder="A cat...", lines=3)
            negative_prompt = gr.Textbox(label="Negative Prompt", lines=2)
    
    with gr.Row():
        with gr.Column():
            steps = gr.Slider(1, 30, value=9, step=1, label="Steps")
            guidance_scale = gr.Slider(0, 10, value=0, step=0.5, label="Guidance Scale")
        with gr.Column():
            width = gr.Dropdown([512, 768, 1024], value=1024, label="Width")
            height = gr.Dropdown([512, 768, 1024], value=1024, label="Height")
            seed = gr.Number(value=0, label="Seed")
    
    generate_btn = gr.Button("Generate", variant="primary")
    output = gr.Image(label="Result")
    status = gr.Textbox(label="Status")
    
    generate_btn.click(generate, inputs=[prompt, negative_prompt, steps, width, height, seed, guidance_scale], outputs=[output, status])

demo.launch()