File size: 4,332 Bytes
5c6edfb
a60b83b
1528648
5c6edfb
1528648
 
5c6edfb
65a7901
 
4daeaec
65a7901
a60b83b
 
5c6edfb
 
0b1e883
a60b83b
 
4daeaec
8fa2633
 
 
 
 
 
 
 
a60b83b
 
 
 
 
 
 
 
 
 
 
 
 
8fa2633
a60b83b
 
1e94087
5c6edfb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fab03a8
5c6edfb
 
 
9c908f8
5c6edfb
 
 
0b1e883
60c1752
5c6edfb
 
d13a106
5c6edfb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ec78c92
 
 
 
 
 
 
5c6edfb
 
 
 
0b1e883
1e94087
5c6edfb
 
 
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
import gradio as gr
from gradio_client import Client, handle_file
import numpy as np

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


flux_1_schell_spaces = ["https://black-forest-labs-flux-1-schnell.hf.space", "ChristianHappy/FLUX.1-schnell", "innoai/FLUX.1-schnell", "tuan2308/FLUX.1-schnell", "FiditeNemini/FLUX.1-schnell"]
# flux_1_schnell_space = "https://black-forest-labs-flux-1-schnell.hf.space"

client = None
job = None


def infer(selected_space, prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
    global job
    global client

    if client is None:
        try:
            client = Client(selected_space)
            print(f"Loaded custom model from {selected_space}")
        except ValueError as e:
            client = None
            print(f"Failed to load custom model from {selected_space}: {e}")
            raise gr.Error("Failed to load client after trying all spaces.")
        
    try:
        job = client.submit(
            prompt=prompt,
            seed=seed,
            randomize_seed=randomize_seed,
            width=width,
            height=height,
            num_inference_steps=num_inference_steps,
            api_name="/infer"
        )
        result = job.result()
    except ValueError as e:
        client = None
        raise gr.Error(e)
    
    return result
 
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:
    selected_space_index = gr.State(0);
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# FLUX.1 [schnell]
[black-forest-labs/FLUX.1-schnell](https://huggingface.co/spaces/black-forest-labs/FLUX.1-schnell)
12B param rectified flow transformer distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/) for 4 step generation
[[blog](https://blackforestlabs.ai/2024/07/31/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1-schnell)]
        """)
        space = gr.Radio(flux_1_schell_spaces, label="HF Space", value=flux_1_schell_spaces[0])
     
        
        with gr.Row():
            prompt = gr.Textbox(
                label="Prompt",
                show_label=False,
                placeholder="Enter your prompt",
                container=False,
            )
            
            run_button = gr.Button("Run", scale=0)
        
        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():
                
  
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=4,
                )
        
        # gr.Examples(
        #     examples = examples,
        #     fn = infer,
        #     inputs = [selected_space_index, prompt],
        #     outputs = [selected_space_index, space, result, seed],
        #     cache_examples="lazy"
        # )

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

demo.launch()