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)