Spaces:
Sleeping
Sleeping
| import os | |
| import spaces | |
| os.environ["HUGGINGFACE_HUB_TOKEN"] = os.getenv("HF_TOKEN") | |
| import torch | |
| import random | |
| import gradio as gr | |
| from diffusers import DiffusionPipeline | |
| import numpy as np | |
| # Device setup | |
| dtype = torch.bfloat16 | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # Load pipelines | |
| print("Loading models...") | |
| pipe_zimage = DiffusionPipeline.from_pretrained( | |
| "Tongyi-MAI/Z-Image-Turbo", torch_dtype=dtype, low_cpu_mem_usage=False | |
| ).to(device) | |
| pipe_schnell = DiffusionPipeline.from_pretrained( | |
| "black-forest-labs/FLUX.1-schnell", | |
| torch_dtype=dtype, | |
| use_auth_token=True | |
| ).to(device) | |
| print("Models loaded!") | |
| MAX_SEED = np.iinfo(np.int32).max | |
| MAX_IMAGE_SIZE = 2048 | |
| def generate(prompt, model_choice, seed=42, randomize_seed=True, width=1024, height=1024, num_inference_steps=8, progress=gr.Progress(track_tqdm=True)): | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| generator = torch.Generator(device).manual_seed(seed) | |
| pipe = pipe_schnell if model_choice == "Schnell" else pipe_zimage | |
| image = pipe( | |
| prompt=prompt, | |
| width=int(width), | |
| height=int(height), | |
| num_inference_steps=int(num_inference_steps), | |
| guidance_scale=0.0, | |
| generator=generator | |
| ).images[0] | |
| return image, seed | |
| # Examples | |
| examples = [ | |
| "Young Chinese woman in red Hanfu, intricate embroidery...", | |
| "A majestic dragon soaring through clouds at sunset...", | |
| "Cozy coffee shop interior, warm lighting, rain on windows...", | |
| "Astronaut riding a horse on Mars, cinematic lighting...", | |
| "Portrait of a wise old wizard with a long white beard..." | |
| ] | |
| # Build Gradio interface | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# 🎨 Multi-Model Image Generator") | |
| with gr.Row(): | |
| with gr.Column(scale=1, min_width=320): | |
| prompt = gr.Textbox(label="✨ Your Prompt", placeholder="Describe the image...", lines=5) | |
| model_choice = gr.Radio(["Z-Image", "Schnell"], label="Select Model", value="Z-Image") | |
| with gr.Accordion("⚙️ Advanced Settings", open=False): | |
| height = gr.Slider(512, MAX_IMAGE_SIZE, value=1024, step=64, label="Height") | |
| width = gr.Slider(512, MAX_IMAGE_SIZE, value=1024, step=64, label="Width") | |
| num_inference_steps = gr.Slider(1, 50, value=8, step=1, label="Inference Steps") | |
| 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") | |
| gr.Examples(examples=examples, inputs=[prompt], label="💡 Try prompts") | |
| with gr.Column(scale=1, min_width=320): | |
| output_image = gr.Image(label="Generated Image", type="pil", height=600, buttons=["download", "share"]) | |
| used_seed = gr.Number(label="🎲 Seed Used", interactive=False) | |
| generate_btn.click( | |
| fn=generate, | |
| inputs=[prompt, model_choice, seed, randomize_seed, width, height, num_inference_steps], | |
| outputs=[output_image, used_seed] | |
| ) | |
| prompt.submit( | |
| fn=generate, | |
| inputs=[prompt, model_choice, seed, randomize_seed, width, height, num_inference_steps], | |
| outputs=[output_image, used_seed] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |