|
import gradio as gr |
|
import torch |
|
from diffusers import DiffusionPipeline |
|
import time |
|
|
|
|
|
base_model = "black-forest-labs/FLUX.1-dev" |
|
pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16) |
|
|
|
MAX_SEED = 2**32-1 |
|
|
|
class calculateDuration: |
|
def __init__(self, activity_name=""): |
|
self.activity_name = activity_name |
|
|
|
def __enter__(self): |
|
self.start_time = time.time() |
|
return self |
|
|
|
def __exit__(self, exc_type, exc_value, traceback): |
|
self.end_time = time.time() |
|
self.elapsed_time = self.end_time - self.start_time |
|
if self.activity_name: |
|
print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds") |
|
else: |
|
print(f"Elapsed time: {self.elapsed_time:.6f} seconds") |
|
|
|
def generate_image(prompt, steps, seed, cfg_scale, width, height): |
|
pipe.to("cuda") |
|
generator = torch.Generator(device="cuda").manual_seed(seed) |
|
|
|
with calculateDuration("Generating image"): |
|
|
|
image = pipe( |
|
prompt=prompt, |
|
num_inference_steps=steps, |
|
guidance_scale=cfg_scale, |
|
width=width, |
|
height=height, |
|
generator=generator |
|
).images[0] |
|
return image |
|
|
|
def run_model(prompt, cfg_scale, steps, randomize_seed, seed, width, height): |
|
if randomize_seed: |
|
seed = torch.randint(0, MAX_SEED, (1,)).item() |
|
|
|
image = generate_image(prompt, steps, seed, cfg_scale, width, height) |
|
return image, seed |
|
|
|
with gr.Blocks() as app: |
|
with gr.Row(): |
|
with gr.Column(): |
|
prompt = gr.Textbox(label="Prompt", placeholder="Type a prompt here") |
|
generate_button = gr.Button("Generate") |
|
|
|
with gr.Row(): |
|
result = gr.Image(label="Generated Image") |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5) |
|
steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28) |
|
width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024) |
|
height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024) |
|
randomize_seed = gr.Checkbox(True, label="Randomize seed") |
|
seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0) |
|
|
|
gr.Interface( |
|
fn=run_model, |
|
inputs=[prompt, cfg_scale, steps, randomize_seed, seed, width, height], |
|
outputs=[result, seed], |
|
live=True |
|
).launch() |
