| import gradio as gr | |
| import torch | |
| from diffusers import AuraFlowPipeline | |
| import spaces | |
| import numpy as np | |
| pipeline = AuraFlowPipeline.from_pretrained( | |
| "fal/AuraFlow-v0.3", | |
| torch_dtype=torch.float16, | |
| variant="fp16", | |
| use_safetensors=True, | |
| ).to("cuda") | |
| STYLE_PRESETS = { | |
| "None": "", | |
| "Comic": ", in comic book style, bold outlines, vibrant colors, dynamic shading", | |
| "Watercolor": ", in watercolor style, soft edges, translucent colors, delicate brushstrokes", | |
| "Oil Painting": ", in oil painting style, rich textures, bold brushstrokes, deep colors", | |
| "Cyberpunk": ", in cyberpunk style, neon lights, dark atmosphere, futuristic elements", | |
| "Photorealistic": ", in photorealistic style, highly detailed, lifelike textures, realistic lighting" | |
| } | |
| examples = [ | |
| {"prompt": "A rustic village nestled in a golden autumn valley, with rolling hills and a winding river bathed in warm light", "style": "Oil Painting"}, | |
| {"prompt": "A majestic dragon soaring high above a range of snow-capped mountains under a golden sunset sky", "style": "Comic"}, | |
| {"prompt": "A shiba inu on a rocky cliff overlooking a vibrant sunset ocean view", "style": "Photorealistic"}, | |
| {"prompt": "A futuristic city skyline glowing with neon lights, towering skyscrapers, and flying cars under a stormy night", "style": "Cyberpunk"}, | |
| ] | |
| def generate_images( | |
| prompt, | |
| negative_prompt, | |
| style, | |
| width=1024, | |
| height=1024, | |
| steps=20, | |
| guidance=5.0, | |
| seed=1, | |
| num_images=1, | |
| ): | |
| generator = torch.Generator(device="cuda").manual_seed(seed) | |
| styled_prompt = f"{prompt}{STYLE_PRESETS[style]}" | |
| gallery = [] | |
| for i in range(num_images): | |
| image = pipeline( | |
| prompt=styled_prompt, | |
| negative_prompt=negative_prompt, | |
| width=width, | |
| height=height, | |
| num_inference_steps=steps, | |
| guidance_scale=guidance, | |
| generator=generator, | |
| output_type="pil", | |
| ).images[0] | |
| gallery.append((image, "")) | |
| torch.cuda.empty_cache() | |
| return gallery | |
| def interface_fn( | |
| prompt, | |
| negative_prompt, | |
| style, | |
| width, | |
| height, | |
| steps, | |
| guidance, | |
| seed, | |
| num_images, | |
| randomize_seed, | |
| history, | |
| progress: gr.Progress = gr.Progress(track_tqdm=True) | |
| ): | |
| if not prompt: | |
| raise gr.Error("Please enter a prompt!") | |
| if randomize_seed: | |
| seed = np.random.randint(0, 1000000) | |
| gallery = generate_images( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| style=style, | |
| width=width, | |
| height=height, | |
| steps=steps, | |
| guidance=guidance, | |
| seed=seed, | |
| num_images=num_images | |
| ) | |
| updated_history = update_history(gallery, history) | |
| return gallery, seed, updated_history | |
| def update_history(new_images, history): | |
| if history is None: | |
| history = [] | |
| for img in reversed(new_images): | |
| history.insert(0, img[0]) | |
| return history | |
| def clear_result(): | |
| return gr.update(value=[]), gr.update(value=None) | |
| #my custom css for layout | |
| custom_css = """ | |
| .gr-button {margin: 5px;} | |
| .output-image {border-radius: 8px;} | |
| #advanced_options {margin-top: 20px;} | |
| .style-dropdown {width: 100%; max-width: 800px;} | |
| .gr-textbox {width: 100%;} | |
| .example-row {margin-top: 20px;} | |
| .example-button {white-space: normal; height: auto; min-height: 60px;} | |
| /* Center the Generated Images gallery */ | |
| #output-gallery { | |
| display: block; | |
| width: 100%; | |
| text-align: center; | |
| } | |
| #output-gallery .gallery { | |
| display: inline-flex; | |
| justify-content: center; | |
| align-items: center; | |
| flex-wrap: wrap; | |
| margin: 0 auto; | |
| } | |
| #output-gallery .gallery > div { | |
| display: flex; | |
| justify-content: center; | |
| align-items: center; | |
| margin: 5px; | |
| } | |
| #output-gallery img { | |
| display: block; | |
| margin: 0 auto; | |
| } | |
| """ | |
| with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as interface: | |
| gr.Markdown("# AuraFlow v0.3 Image Generator") | |
| gr.Markdown("Enter a prompt and select a style to generate images. Use the advanced settings for more control or try an example below.") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| prompt_input = gr.Textbox( | |
| label="Prompt", | |
| placeholder="Enter your creative prompt here", | |
| lines=3 | |
| ) | |
| neg_prompt_input = gr.Textbox( | |
| label="Negative Prompt", | |
| placeholder="What you don’t want in the image", | |
| lines=2 | |
| ) | |
| style_input = gr.Dropdown( | |
| choices=list(STYLE_PRESETS.keys()), | |
| value="None", | |
| label="Art Style", | |
| elem_classes=["style-dropdown"] | |
| ) | |
| with gr.Accordion("Advanced Settings", open=False, elem_id="advanced_options"): | |
| width_input = gr.Slider(256, 1536, step=256, value=1024, label="Width") | |
| height_input = gr.Slider(256, 1536, step=256, value=1024, label="Height") | |
| steps_input = gr.Slider(1, 50, step=1, value=20, label="Inference Steps") | |
| guidance_input = gr.Slider(0, 10, step=0.5, value=5.0, label="Guidance Scale") | |
| with gr.Row(): | |
| seed_input = gr.Number(value=1, label="Seed", visible=False) | |
| randomize_seed_input = gr.Checkbox(value=True, label="Randomize Seed") | |
| num_images_input = gr.Slider(1, 4, step=1, value=1, label="Number of Images") | |
| with gr.Column(scale=2): | |
| image_output = gr.Gallery( | |
| label="Generated Images", | |
| show_label=True, | |
| preview=True, | |
| elem_id="output-gallery" | |
| ) | |
| with gr.Row(): | |
| clear_btn = gr.Button("Clear", variant="secondary") | |
| generate_btn = gr.Button("Generate", variant="primary") | |
| history_gallery = gr.Gallery( | |
| label="History", | |
| columns=6, | |
| object_fit="contain", | |
| interactive=False | |
| ) | |
| with gr.Row(equal_height=True, elem_classes=["example-row"]): | |
| gr.Markdown("### Try these examples:") | |
| with gr.Row(equal_height=True, elem_classes=["example-row"]): | |
| for ex in examples: | |
| with gr.Column(scale=1, min_width=200): | |
| btn = gr.Button(ex["prompt"], variant="secondary", elem_classes=["example-button"]) | |
| btn.click( | |
| fn=lambda p=ex["prompt"], s=ex["style"]: (gr.update(value=p), gr.update(value=s)), | |
| inputs=[], | |
| outputs=[prompt_input, style_input] | |
| ).then( | |
| fn=interface_fn, | |
| inputs=[prompt_input, neg_prompt_input, style_input, width_input, height_input, | |
| steps_input, guidance_input, seed_input, num_images_input, randomize_seed_input, history_gallery], | |
| outputs=[image_output, seed_input, history_gallery] | |
| ) | |
| generate_btn.click( | |
| fn=lambda: clear_result(), | |
| inputs=[], | |
| outputs=[image_output, seed_input] | |
| ).then( | |
| fn=interface_fn, | |
| inputs=[prompt_input, neg_prompt_input, style_input, width_input, height_input, | |
| steps_input, guidance_input, seed_input, num_images_input, randomize_seed_input, history_gallery], | |
| outputs=[image_output, seed_input, history_gallery] | |
| ) | |
| clear_btn.click( | |
| fn=clear_result, | |
| inputs=[], | |
| outputs=[image_output, seed_input] | |
| ).then( | |
| fn=lambda x: (gr.update(value=[]), gr.update(value=None), x), | |
| inputs=[history_gallery], | |
| outputs=[image_output, seed_input, history_gallery] | |
| ) | |
| randomize_seed_input.change( | |
| fn=lambda randomize: gr.update(visible=not randomize), | |
| inputs=randomize_seed_input, | |
| outputs=seed_input | |
| ) | |
| interface.launch( | |
| show_error=True, | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| share=True | |
| ) |