import spaces import gradio as gr import torch from PIL import Image from diffusers import DiffusionPipeline import random import uuid from typing import Union, List, Optional import numpy as np import time import zipfile # Description for the app DESCRIPTION = """## Qwen Image Generator""" # Helper functions def save_image(img): unique_name = str(uuid.uuid4()) + ".png" img.save(unique_name) return unique_name def randomize_seed_fn(seed: int, randomize_seed: bool) -> int: if randomize_seed: seed = random.randint(0, MAX_SEED) return seed MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = 2048 # Load Qwen/Qwen-Image pipeline dtype = torch.bfloat16 device = "cuda" if torch.cuda.is_available() else "cpu" pipe_qwen = DiffusionPipeline.from_pretrained("Qwen/Qwen-Image", torch_dtype=dtype).to(device) # Aspect ratios aspect_ratios = { "1:1": (1328, 1328), "16:9": (1664, 928), "9:16": (928, 1664), "4:3": (1472, 1140), "3:4": (1140, 1472) } # Generation function for Qwen/Qwen-Image @spaces.GPU def generate_qwen( prompt: str, negative_prompt: str = "", seed: int = 0, width: int = 1024, height: int = 1024, guidance_scale: float = 4.0, randomize_seed: bool = False, num_inference_steps: int = 50, num_images: int = 1, zip_images: bool = False, progress=gr.Progress(track_tqdm=True), ): if randomize_seed: seed = random.randint(0, MAX_SEED) generator = torch.Generator(device).manual_seed(seed) start_time = time.time() images = pipe_qwen( prompt=prompt, negative_prompt=negative_prompt if negative_prompt else None, height=height, width=width, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, num_images_per_prompt=num_images, generator=generator, output_type="pil", ).images end_time = time.time() duration = end_time - start_time image_paths = [save_image(img) for img in images] zip_path = None if zip_images: zip_name = str(uuid.uuid4()) + ".zip" with zipfile.ZipFile(zip_name, 'w') as zipf: for i, img_path in enumerate(image_paths): zipf.write(img_path, arcname=f"Img_{i}.png") zip_path = zip_name return image_paths, seed, f"{duration:.2f}", zip_path # Wrapper function to handle UI logic @spaces.GPU def generate( prompt: str, negative_prompt: str, use_negative_prompt: bool, seed: int, width: int, height: int, guidance_scale: float, randomize_seed: bool, num_inference_steps: int, num_images: int, zip_images: bool, progress=gr.Progress(track_tqdm=True), ): final_negative_prompt = negative_prompt if use_negative_prompt else "" return generate_qwen( prompt=prompt, negative_prompt=final_negative_prompt, seed=seed, width=width, height=height, guidance_scale=guidance_scale, randomize_seed=randomize_seed, num_inference_steps=num_inference_steps, num_images=num_images, zip_images=zip_images, progress=progress, ) # Examples examples = [ "An attractive young woman with blue eyes lying face down on the bed, light white and light amber, timeless beauty, sunrays shine upon it", "Headshot of handsome young man, wearing dark gray sweater, brown hair and short beard, serious look, black background, soft studio lighting", "A medium-angle shot of a young woman with long brown hair, wearing glasses, standing in front of purple and white lights", "High-resolution photograph of a woman, photorealistic, vibrant colors" ] css = ''' .gradio-container { max-width: 590px !important; margin: 0 auto !important; } h1 { text-align: center; } footer { visibility: hidden; } ''' # Gradio interface with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo: gr.Markdown(DESCRIPTION) with gr.Row(): prompt = gr.Text( label="Prompt", show_label=False, max_lines=1, placeholder="Enter your prompt", container=False, ) run_button = gr.Button("Run", scale=0, variant="primary") result = gr.Gallery(label="Result", columns=1, show_label=False, preview=True) with gr.Accordion("Additional Options", open=False): aspect_ratio = gr.Dropdown( label="Aspect Ratio", choices=list(aspect_ratios.keys()), value="1:1", ) use_negative_prompt = gr.Checkbox( label="Use negative prompt", value=False, visible=True ) negative_prompt = gr.Text( label="Negative prompt", max_lines=1, placeholder="Enter a negative prompt", visible=False, ) seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, ) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) with gr.Row(): width = gr.Slider( label="Width", minimum=512, maximum=2048, step=64, value=1024, ) height = gr.Slider( label="Height", minimum=512, maximum=2048, step=64, value=1024, ) guidance_scale = gr.Slider( label="Guidance Scale", minimum=0.0, maximum=20.0, step=0.1, value=4.0, ) num_inference_steps = gr.Slider( label="Number of inference steps", minimum=1, maximum=100, step=1, value=50, ) num_images = gr.Slider( label="Number of images", minimum=1, maximum=5, step=1, value=1, ) zip_images = gr.Checkbox(label="Zip generated images", value=False) gr.Markdown("### Output Information") seed_display = gr.Textbox(label="Seed used", interactive=False) generation_time = gr.Textbox(label="Generation time (seconds)", interactive=False) zip_file = gr.File(label="Download ZIP") # Update aspect ratio def set_dimensions(ar): w, h = aspect_ratios[ar] return gr.update(value=w), gr.update(value=h) aspect_ratio.change( fn=set_dimensions, inputs=aspect_ratio, outputs=[width, height] ) # Negative prompt visibility use_negative_prompt.change( fn=lambda x: gr.update(visible=x), inputs=use_negative_prompt, outputs=negative_prompt ) # Run button and prompt submit gr.on( triggers=[prompt.submit, run_button.click], fn=generate, inputs=[ prompt, negative_prompt, use_negative_prompt, seed, width, height, guidance_scale, randomize_seed, num_inference_steps, num_images, zip_images, ], outputs=[result, seed_display, generation_time, zip_file], api_name="run", ) # Examples gr.Examples( examples=examples, inputs=prompt, outputs=[result, seed_display, generation_time, zip_file], fn=generate, cache_examples=False, ) if __name__ == "__main__": demo.queue(max_size=30).launch(mcp_server=True, ssr_mode=False, show_error=True)