File size: 4,410 Bytes
0ac8bb5
134c1b2
 
 
345aaf6
134c1b2
96b59b6
 
134c1b2
 
cf9fce4
134c1b2
 
832c70d
4c2de11
134c1b2
832c70d
134c1b2
4c2de11
137f187
fb3ffbe
8e8990a
134c1b2
ebdce90
134c1b2
 
 
 
 
56c7207
134c1b2
 
4c2de11
 
 
 
 
 
134c1b2
 
 
 
 
 
 
56c7207
d01359b
56c7207
 
 
 
 
 
 
 
 
134c1b2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e4a91d8
134c1b2
 
 
 
 
 
 
 
 
 
 
 
3d86fda
134c1b2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
56c7207
134c1b2
 
 
 
 
 
 
56c7207
134c1b2
 
 
 
 
 
 
 
56c7207
134c1b2
 
 
 
 
 
 
56c7207
134c1b2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import spaces
import gradio as gr
import numpy as np
import random
import os
import torch
from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images

device = "cuda" if torch.cuda.is_available() else "cpu"
hf_token = os.getenv('HF_TOKEN')

if torch.cuda.is_available():
    dtype = torch.float16
    torch.cuda.empty_cache()
else:
    dtype = torch.float32

taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
pipe = pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", token=hf_token, torch_dtype=dtype)
pipe.load_lora_weights('aleksa-codes/flux-ghibsky-illustration', weight_name='lora.safetensors')
good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
pipe = pipe.to(device)
pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)

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


@spaces.GPU
def infer(
    prompt,
    seed=42,
    randomize_seed=True,
    width=1024,
    height=1024,
    guidance_scale=3.5,
    num_inference_steps=28,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator().manual_seed(seed)

    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
            prompt="GHIBSKY style, " + prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator,
            output_type="pil",
            good_vae=good_vae,
        ):
            yield img, seed


examples = [
    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
    "An astronaut riding a green horse",
    "A delicious ceviche cheesecake slice",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(" # Flux Ghibsky Illustration: Create Serene and Enchanting Landscapes")

        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, variant="primary")

        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=0.0,
                    maximum=10.0,
                    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, inputs=[prompt])
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, seed],
    )

if __name__ == "__main__":
    demo.launch()