Spaces:
Running
on
Zero
Running
on
Zero
import random | |
import gradio as gr | |
import numpy as np | |
import spaces | |
import torch | |
from dataset_viber import CollectorInterface | |
from diffusers import DiffusionPipeline | |
dtype = torch.bfloat16 | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).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, progress=gr.Progress(track_tqdm=True)): | |
if randomize_seed: | |
seed = random.randint(0, MAX_SEED) | |
generator = torch.Generator().manual_seed(seed) | |
image = pipe( | |
prompt=prompt, | |
width=width, | |
height=height, | |
num_inference_steps=num_inference_steps, | |
generator=generator, | |
guidance_scale=0.0 | |
).images[0] | |
return image | |
examples = [ | |
["a tiny astronaut hatching from an egg on the moon", 0, True, 1024, 1024, 4], | |
["a cat holding a sign that says hello world", 0, True, 1024, 1024, 4], | |
["an anime illustration of a wiener schnitzel", 0, True, 1024, 1024, 4], | |
] | |
css = """ | |
#col-container { | |
margin: 0 auto; | |
max-width: 520px; | |
} | |
""" | |
description = """# FLUX.1 [schnell] | |
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-schnell)] | |
""" | |
interface = CollectorInterface( | |
fn=infer, | |
inputs=[ | |
gr.Textbox(label="Prompt", placeholder="Enter your prompt") | |
], | |
outputs=[ | |
gr.Image(label="Result"), | |
], | |
additional_inputs=[ | |
gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0), | |
gr.Checkbox(label="Randomize seed", value=True), | |
gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024), | |
gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024), | |
gr.Slider(label="Number of inference steps", minimum=1, maximum=50, step=1, value=4), | |
], | |
title="FLUX.1 [schnell] - with Dataset Viber data collection", | |
description=description, | |
examples=examples, | |
css=css, | |
dataset_name="image-generation-flux1-schnell" | |
) | |
interface.launch() |