|
import gradio as gr |
|
import numpy as np |
|
import random |
|
import torch |
|
import time |
|
from diffusers import DiffusionPipeline |
|
|
|
|
|
dtype = torch.float16 |
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
|
|
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype) |
|
pipe = pipe.to(device) |
|
|
|
MAX_SEED = np.iinfo(np.int32).max |
|
MAX_IMAGE_SIZE = 2048 |
|
|
|
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, guidance_scale=7.5, progress=gr.Progress(track_tqdm=True)): |
|
start_time = time.time() |
|
|
|
if width > MAX_IMAGE_SIZE or height > MAX_IMAGE_SIZE: |
|
raise ValueError("Image size exceeds the maximum allowed dimensions.") |
|
|
|
if randomize_seed: |
|
seed = random.randint(0, MAX_SEED) |
|
generator = torch.Generator(device=device).manual_seed(seed) |
|
|
|
try: |
|
image = pipe( |
|
prompt=prompt, |
|
width=width, |
|
height=height, |
|
num_inference_steps=num_inference_steps, |
|
generator=generator, |
|
guidance_scale=guidance_scale |
|
).images[0] |
|
except Exception as e: |
|
print(f"Error generating image: {e}") |
|
return None, seed, f"Error: {str(e)}" |
|
|
|
if time.time() - start_time > 60: |
|
return None, seed, "Image generation took too long and was cancelled." |
|
|
|
return image, seed, None |
|
|
|
examples = [ |
|
["a tiny astronaut hatching from an egg on the moon"], |
|
["a cat holding a sign that says hello world"], |
|
["an anime illustration of a wiener schnitzel"], |
|
] |
|
|
|
with gr.Blocks(theme=gr.themes.Soft()) as demo: |
|
gr.Markdown(""" |
|
# Custom Image Creator |
|
12B param rectified flow transformer distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/) for 4 step generation |
|
[[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1)] |
|
""") |
|
|
|
with gr.Row(): |
|
with gr.Column(scale=2): |
|
prompt = gr.Textbox( |
|
label="Prompt", |
|
placeholder="Enter your prompt", |
|
lines=3 |
|
) |
|
run_button = gr.Button("Generate Image", variant="primary") |
|
|
|
with gr.Column(scale=2): |
|
result = gr.Image(label="Generated Image") |
|
seed_output = gr.Number(label="Seed Used") |
|
|
|
with gr.Accordion("Advanced Settings", open=False): |
|
seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42) |
|
randomize_seed = gr.Checkbox(label="Randomize seed", value=True) |
|
|
|
with gr.Row(): |
|
width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024) |
|
height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024) |
|
|
|
with gr.Row(): |
|
num_inference_steps = gr.Slider(label="Number of inference steps", minimum=1, maximum=50, step=1, value=4) |
|
guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=20.0, step=0.5, value=7.5) |
|
|
|
gr.Examples( |
|
examples=examples, |
|
inputs=[prompt], |
|
outputs=[result, seed_output], |
|
fn=infer, |
|
cache_examples=True |
|
) |
|
|
|
run_button.click( |
|
fn=infer, |
|
inputs=[prompt, seed, randomize_seed, width, height, num_inference_steps, guidance_scale], |
|
outputs=[result, seed_output] |
|
) |
|
|
|
gr.Markdown(""" |
|
## Save Your Image |
|
Right-click on the generated image and select 'Save image as' to download it. |
|
""") |
|
|
|
if __name__ == "__main__": |
|
demo.launch() |
