File size: 2,598 Bytes
466c880 0a68e6e 466c880 0a68e6e c3a9e27 0a68e6e c3a9e27 0a68e6e c3a9e27 |
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 |
import gradio as gr
import torch
from diffusers import DiffusionPipeline
import time
# Initialize the base model
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"):
# Generate 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() |