Spaces:
Running
on
Zero
Running
on
Zero
File size: 6,138 Bytes
40462a0 7aafe2f c82ffd4 2097b43 40462a0 7aafe2f 40462a0 7aafe2f 2097b43 7aafe2f deb928c 7aafe2f 7d90483 f1c2277 7aafe2f c82ffd4 7aafe2f b7bdba8 7aafe2f b7bdba8 7aafe2f b7bdba8 7aafe2f ffe402c 7aafe2f ffe402c 9494d71 5d87e58 7aafe2f a25ba7a a8abbc8 7aafe2f bf1be4a 7aafe2f c9efd38 7aafe2f db97287 7aafe2f db97287 f1c2277 7aafe2f 1c4647b 7aafe2f 5d87e58 7aafe2f 3bad26c 7aafe2f 3bad26c 97d3c4e 7aafe2f 5358c2f 8374546 7aafe2f f1c2277 7aafe2f 23fd89e 7aafe2f f1c2277 05a0589 66892aa f1c2277 8374546 f1c2277 3bad26c f1c2277 3bad26c 8374546 7aafe2f f1c2277 c9c4986 8374546 f1c2277 7aafe2f f1c2277 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
import gradio as gr
import numpy as np
import random
import spaces
import torch
import time
from diffusers import DiffusionPipeline, AutoencoderTiny
from custom_pipeline import HighSpeedFluxPipeline
# Constants
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048
DEFAULT_WIDTH = 1024
DEFAULT_HEIGHT = 1024
DEFAULT_INFERENCE_STEPS = 1
# Device and model setup
dtype = torch.float16
pipe = HighSpeedFluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-schnell", torch_dtype=dtype
).to("cuda")
pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.float16).to("cuda")
torch.cuda.empty_cache()
# Inference function
@spaces.GPU(duration=25)
def generate_image(prompt, seed=42, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, randomize_seed=False, num_inference_steps=2, progress=gr.Progress(track_tqdm=True)):
if randomize_seed:
seed = random.randint(0, MAX_SEED)
generator = torch.Generator().manual_seed(int(float(seed)))
start_time = time.time()
# Only generate the last image in the sequence
img = pipe.generate_images(
prompt=prompt,
width=width,
height=height,
num_inference_steps=num_inference_steps,
generator=generator
)
latency = f"Latency: {(time.time()-start_time):.2f} seconds"
return img, seed, latency
# Example prompts
examples = [
"a tiny astronaut hatching from an egg on the moon",
"a cute white cat holding a sign that says hello world",
"an anime illustration of a wiener schnitzel",
"Create mage of Modern house in minecraft style",
"Imagine steve jobs as Star Wars movie character",
"Lion",
"Photo of a young woman with long, wavy brown hair tied in a bun and glasses. She has a fair complexion and is wearing subtle makeup, emphasizing her eyes and lips. She is dressed in a black top. The background appears to be an urban setting with a building facade, and the sunlight casts a warm glow on her face.",
]
# --- Gradio UI ---
with gr.Blocks() as demo:
with gr.Column(elem_id="app-container"):
gr.Markdown("# 🎨 Realtime FLUX Image Generator")
gr.Markdown("Generate stunning images in real-time with Modified Flux.Schnell pipeline.")
gr.Markdown("<span style='color: red;'>Note: Sometimes it stucks or stops generating images (I don't know why). In that situation just refresh the site.</span>")
with gr.Row():
with gr.Column(scale=2.5):
result = gr.Image(label="Generated Image", show_label=False, interactive=False)
with gr.Column(scale=1):
prompt = gr.Text(
label="Prompt",
placeholder="Describe the image you want to generate...",
lines=3,
show_label=False,
container=False,
)
generateBtn = gr.Button("🖼️ Generate Image")
enhanceBtn = gr.Button("🚀 Enhance Image")
with gr.Column("Advanced Options"):
with gr.Row():
realtime = gr.Checkbox(label="Realtime Toggler", info="If TRUE then uses more GPU but create image in realtime.", value=False)
latency = gr.Text(label="Latency")
with gr.Row():
seed = gr.Number(label="Seed", value=42)
randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
with gr.Row():
width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_WIDTH)
height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_HEIGHT)
num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=4, step=1, value=DEFAULT_INFERENCE_STEPS)
with gr.Row():
gr.Markdown("### 🌟 Inspiration Gallery")
with gr.Row():
gr.Examples(
examples=examples,
fn=generate_image,
inputs=[prompt],
outputs=[result, seed, latency],
cache_examples="lazy"
)
def enhance_image(*args):
gr.Info("Enhancing Image") # currently just runs optimized pipeline for 2 steps. Further implementations later.
return next(generate_image(*args))
enhanceBtn.click(
fn=enhance_image,
inputs=[prompt, seed, width, height],
outputs=[result, seed, latency],
show_progress="full",
queue=False,
concurrency_limit=None
)
generateBtn.click(
fn=generate_image,
inputs=[prompt, seed, width, height, randomize_seed, num_inference_steps],
outputs=[result, seed, latency],
show_progress="full",
api_name="RealtimeFlux",
queue=False
)
def update_ui(realtime_enabled):
return {
prompt: gr.update(interactive=True),
generateBtn: gr.update(visible=not realtime_enabled)
}
realtime.change(
fn=update_ui,
inputs=[realtime],
outputs=[prompt, generateBtn],
queue=False,
concurrency_limit=None
)
def realtime_generation(*args):
if args[0]: # If realtime is enabled
return next(generate_image(*args[1:]))
prompt.submit(
fn=generate_image,
inputs=[prompt, seed, width, height, randomize_seed, num_inference_steps],
outputs=[result, seed, latency],
show_progress="full",
queue=False,
concurrency_limit=None
)
for component in [prompt, width, height, num_inference_steps]:
component.input(
fn=realtime_generation,
inputs=[realtime, prompt, seed, width, height, randomize_seed, num_inference_steps],
outputs=[result, seed, latency],
show_progress="hidden",
trigger_mode="always_last",
queue=False,
concurrency_limit=None
)
# Launch the app
demo.launch()
|