Spaces:
Running
on
Zero
Running
on
Zero
File size: 5,866 Bytes
8ccf632 0dd76b1 8ccf632 0dd76b1 06f0278 8ccf632 02dee9c 8ccf632 06f0278 8ccf632 76d8871 8267888 76d8871 0dd76b1 6a99664 0dd76b1 8267888 b2a74a5 0dd76b1 b2a74a5 0dd76b1 8267888 0dd76b1 8ccf632 0dd76b1 8267888 0dd76b1 1d2a6dd 8267888 1d2a6dd 8ccf632 06f0278 8ccf632 0dd76b1 8ccf632 e2944a6 8ccf632 0dd76b1 02dee9c 0dd76b1 8ccf632 0dd76b1 8ccf632 0dd76b1 8ccf632 0dd76b1 8ccf632 0dd76b1 8ccf632 b213a9c 0dd76b1 b213a9c ceb48e8 02dee9c b213a9c 0dd76b1 8ccf632 b213a9c 8ccf632 0dd76b1 8ccf632 0dd76b1 1d2a6dd 0dd76b1 8267888 8ccf632 0dd76b1 8267888 0dd76b1 8267888 8ccf632 0dd76b1 |
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 163 164 165 166 167 168 169 170 171 172 173 174 |
import gradio as gr
import numpy as np
import PIL.Image
import spaces
import torch
from diffusers import AutoencoderTiny, DiffusionPipeline
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"
taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1).to(device)
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048
def get_seed(randomize_seed: bool, seed: int) -> int:
"""Determine and return the random seed to use for model generation.
Args:
randomize_seed (bool): If True, a random seed (an integer in [0, MAX_SEED)) is generated using NumPy's default random number generator. If False, the provided seed argument is returned as-is.
seed (int): The seed value to use if randomize_seed is False.
Returns:
int: The selected seed value. If randomize_seed is True, a randomly generated integer; otherwise, the value of the seed argument.
Notes:
- MAX_SEED is the maximum value for a 32-bit integer (np.iinfo(np.int32).max).
- This function is typically used to ensure reproducibility or to introduce randomness in model generation.
"""
rng = np.random.default_rng()
return int(rng.integers(0, MAX_SEED)) if randomize_seed else seed
@spaces.GPU(duration=75)
def infer(
prompt: str,
seed: int,
width: int = 1024,
height: int = 1024,
guidance_scale: float = 3.5,
num_inference_steps: int = 28,
progress: gr.Progress = gr.Progress(track_tqdm=True), # noqa: ARG001, B008
) -> PIL.Image.Image:
"""Generate an image from a text prompt using the FLUX.1 [dev] model.
Note:
- Prompts must be written in English. Other languages are not currently supported.
- Prompts are limited to 77 tokens due to CLIP tokenizer constraints.
Args:
prompt: A text prompt in English to guide the image generation. Limited to 77 tokens.
seed: The seed value used for reproducible image generation.
width: Width of the output image in pixels. Defaults to 1024.
height: Height of the output image in pixels. Defaults to 1024.
guidance_scale: Controls how strongly the model follows the prompt.
Higher values lead to images more closely aligned with the prompt. Defaults to 3.5.
num_inference_steps: Number of denoising steps during generation. Higher values can improve quality. Defaults to 28.
progress: (Internal) Progress tracker for UI integration; should not be manually set by users.
Returns:
A PIL.Image.Image object representing the generated image.
"""
generator = torch.Generator().manual_seed(seed)
return pipe(
prompt=prompt,
width=width,
height=height,
num_inference_steps=num_inference_steps,
generator=generator,
guidance_scale=guidance_scale,
).images[0]
def run_example(prompt: str) -> tuple[PIL.Image.Image, int]:
return infer(prompt, seed=42)
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",
]
css = """
#col-container {
margin: 0 auto;
max-width: 520px;
}
"""
with gr.Blocks(css=css) as demo:
with gr.Column(elem_id="col-container"):
gr.Markdown("""# FLUX.1 [dev]
12B param rectified flow transformer guidance-distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/)
[[non-commercial license](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)] [[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1-dev)]
""")
with gr.Row():
prompt = gr.Text(
label="Prompt",
show_label=False,
max_lines=1,
placeholder="Enter your prompt",
submit_btn=True,
)
result = gr.Image(label="Result", show_label=False)
with gr.Accordion("Advanced Settings", open=False):
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=MAX_SEED,
step=1,
value=0,
)
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():
guidance_scale = gr.Slider(
label="Guidance Scale",
minimum=1,
maximum=15,
step=0.1,
value=3.5,
)
num_inference_steps = gr.Slider(
label="Number of inference steps",
minimum=1,
maximum=50,
step=1,
value=28,
)
gr.Examples(
examples=examples,
fn=run_example,
inputs=prompt,
outputs=result,
)
prompt.submit(
fn=get_seed,
inputs=[randomize_seed, seed],
outputs=seed,
).then(
fn=infer,
inputs=[prompt, seed, width, height, guidance_scale, num_inference_steps],
outputs=result,
)
if __name__ == "__main__":
demo.launch(mcp_server=True)
|