File size: 4,089 Bytes
6f49453
2f01d97
 
 
 
046370b
11419c0
 
6f49453
2f01d97
 
 
11419c0
5258884
11419c0
 
f59f8a1
2f01d97
 
 
 
 
fdfe2cf
2f01d97
 
 
fdfe2cf
 
 
 
 
 
 
 
 
 
 
 
 
 
2f01d97
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fdfe2cf
393347c
 
 
 
 
 
 
 
 
2f01d97
393347c
 
 
2f01d97
393347c
 
 
 
 
 
2f01d97
 
393347c
 
 
 
 
 
 
 
 
2f01d97
393347c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2f01d97
 
393347c
 
 
 
2f01d97
 
 
 
 
393347c
fdfe2cf
393347c
2f01d97
 
11419c0
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
import gradio as gr
import numpy as np
import random
import spaces
import torch
import os
from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

# Include your Hugging Face access token
hf_token = os.getenv("waffles")

# Load the diffusion pipeline with the access token
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, token=hf_token).to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

@spaces.GPU(duration=190)
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=5.0, num_inference_steps=28, num_images=1, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    
    images = []
    for _ in range(num_images):
        image = pipe(
            prompt=prompt, 
            width=width,
            height=height,
            num_inference_steps=num_inference_steps, 
            generator=generator,
            guidance_scale=guidance_scale
        ).images[0]
        images.append(image)
    
    return images, seed
 
examples = [
    "a white husky knocking everything down in a living room",
    "a tuxedo cat with a waffle in her mouth",
    "an anime Chiweenie Dog wearing a hoodie",
]

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(f"""# 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",
                container=False,
            )
            
            run_button = gr.Button("Run", scale=0)
        
        result = gr.Gallery(label="Result", show_label=False)

        # Display all settings directly without the Accordion
        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=5.0,
            )

            num_inference_steps = gr.Slider(
                label="Number of inference steps",
                minimum=1,
                maximum=50,
                step=1,
                value=28,
            )
        
        gr.Examples(
            examples=examples,
            fn=infer,
            inputs=[prompt],
            outputs=[result, seed],
            cache_examples="lazy"
        )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images],
        outputs=[result, seed]
    )

demo.launch()