-
Notifications
You must be signed in to change notification settings - Fork 3
/
gradio_app.py
175 lines (148 loc) · 7.9 KB
/
gradio_app.py
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
175
import gradio as gr
import torch
from free_lunch_utils import register_free_upblock2d, register_free_crossattn_upblock2d
from pipeline_freescale import StableDiffusionXLPipeline
from pipeline_freescale_turbo import StableDiffusionXLPipeline_Turbo
dtype = torch.float16
device = "cuda"
model_ckpt = "stabilityai/stable-diffusion-xl-base-1.0"
model_ckpt_turbo = "stabilityai/sdxl-turbo"
pipe = StableDiffusionXLPipeline.from_pretrained(model_ckpt, torch_dtype=dtype).to(device)
pipe_turbo = StableDiffusionXLPipeline_Turbo.from_pretrained(model_ckpt_turbo, torch_dtype=dtype).to(device)
register_free_upblock2d(pipe, b1=1.1, b2=1.2, s1=0.6, s2=0.4)
register_free_crossattn_upblock2d(pipe, b1=1.1, b2=1.2, s1=0.6, s2=0.4)
register_free_upblock2d(pipe_turbo, b1=1.1, b2=1.2, s1=0.6, s2=0.4)
register_free_crossattn_upblock2d(pipe_turbo, b1=1.1, b2=1.2, s1=0.6, s2=0.4)
torch.cuda.empty_cache()
def infer_gpu_part(seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, restart_steps):
generator = torch.Generator(device='cuda')
generator = generator.manual_seed(seed)
result = pipe(prompt, negative_prompt=negative_prompt, generator=generator,
num_inference_steps=ddim_steps, guidance_scale=guidance_scale,
resolutions_list=resolutions_list, fast_mode=fast_mode, cosine_scale=cosine_scale,
restart_steps=restart_steps,
).images[0]
return result
def infer_gpu_part_turbo(seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, restart_steps):
generator = torch.Generator(device='cuda')
generator = generator.manual_seed(seed)
result = pipe_turbo(prompt, negative_prompt=negative_prompt, generator=generator,
num_inference_steps=ddim_steps, guidance_scale=guidance_scale,
resolutions_list=resolutions_list, fast_mode=fast_mode, cosine_scale=cosine_scale,
restart_steps=restart_steps,
).images[0]
return result
def infer(prompt, output_size, ddim_steps, guidance_scale, cosine_scale, seed, options, negative_prompt):
print(prompt)
print(negative_prompt)
disable_turbo = 'Disable Turbo' in options
if disable_turbo:
fast_mode = True
if output_size == "2048 x 2048":
resolutions_list = [[1024, 1024],
[2048, 2048]]
elif output_size == "1024 x 2048":
resolutions_list = [[512, 1024],
[1024, 2048]]
elif output_size == "2048 x 1024":
resolutions_list = [[1024, 512],
[2048, 1024]]
restart_steps = [int(ddim_steps * 0.3)]
result = infer_gpu_part(seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, restart_steps)
else:
fast_mode = False
if output_size == "2048 x 2048":
resolutions_list = [[512, 512],
[1024, 1024],
[2048, 2048]]
elif output_size == "1024 x 2048":
resolutions_list = [[256, 512],
[512, 1024],
[1024, 2048]]
elif output_size == "2048 x 1024":
resolutions_list = [[512, 256],
[1024, 512],
[2048, 1024]]
restart_steps = [int(ddim_steps * 0.5)] * 2
result = infer_gpu_part_turbo(seed, prompt, negative_prompt, ddim_steps, guidance_scale, resolutions_list, fast_mode, cosine_scale, restart_steps)
return result
examples = [
["A cute and adorable fluffy puppy wearing a witch hat in a Halloween autumn evening forest, falling autumn leaves, brown acorns on the ground, Halloween pumpkins spiderwebs, bats, and a witch’s broom.",],
["Brunette pilot girl in a snowstorm, full body, moody lighting, intricate details, depth of field, outdoors, Fujifilm XT3, RAW, 8K UHD, film grain, Unreal Engine 5, ray tracing.",],
["A panda walking and munching bamboo in a bamboo forest.",],
]
css = """
#col-container {max-width: 768px; margin-left: auto; margin-right: auto;}
"""
def mode_update(options):
if 'Disable Turbo' in options:
return [gr.Slider(minimum=5,
maximum=60,
value=50),
gr.Slider(minimum=1.0,
maximum=20.0,
value=7.5),
gr.Row(visible=True)]
else:
return [gr.Slider(minimum=2,
maximum=6,
value=4),
gr.Slider(minimum=0.0,
maximum=1.0,
value=0.0),
gr.Row(visible=False)]
with gr.Blocks(css=css) as demo:
with gr.Column(elem_id="col-container"):
gr.Markdown(
"""
<h1 style="text-align: center;">FreeScale (unleash the resolution of SDXL)</h1>
<p style="text-align: center;">
FreeScale: Unleashing the Resolution of Diffusion Models via Tuning-Free Scale Fusion
</p>
<p style="text-align: center;">
<a href="https://arxiv.org/abs/2412.09626" target="_blank"><b>[arXiv]</b></a>
<a href="http://haonanqiu.com/projects/FreeScale.html" target="_blank"><b>[Project Page]</b></a>
<a href="https://github.com/ali-vilab/FreeScale" target="_blank"><b>[Code]</b></a>
</p>
"""
)
prompt_in = gr.Textbox(label="Prompt", placeholder="A panda walking and munching bamboo in a bamboo forest.")
with gr.Row():
with gr.Accordion('Advanced Settings', open=False):
with gr.Row():
output_size = gr.Dropdown(["2048 x 2048", "1024 x 2048", "2048 x 1024"], value="2048 x 2048", label="Output Size (H x W)", info="Due to GPU constraints, run the demo locally for higher resolutions.")
options = gr.CheckboxGroup(['Disable Turbo'], label="Options", info="Disable Turbo will get better results but cost more time.")
with gr.Row():
ddim_steps = gr.Slider(label='DDIM Steps',
minimum=2,
maximum=6,
step=1,
value=4)
guidance_scale = gr.Slider(label='Guidance Scale (Disabled in Turbo)',
minimum=0.0,
maximum=1.0,
step=0.1,
value=0.0)
with gr.Row():
cosine_scale = gr.Slider(label='Cosine Scale',
minimum=0,
maximum=10,
step=0.1,
value=2.0)
seed = gr.Slider(label='Random Seed',
minimum=0,
maximum=10000,
step=1,
value=111)
with gr.Row() as row_neg:
negative_prompt = gr.Textbox(label='Negative Prompt', value='blurry, ugly, duplicate, poorly drawn, deformed, mosaic', visible=False)
options.change(mode_update, options, [ddim_steps, guidance_scale, row_neg])
submit_btn = gr.Button("Generate", variant='primary')
image_result = gr.Image(label="Image Output")
gr.Examples(examples=examples, inputs=[prompt_in, output_size, ddim_steps, guidance_scale, cosine_scale, seed, options, negative_prompt])
submit_btn.click(fn=infer,
inputs=[prompt_in, output_size, ddim_steps, guidance_scale, cosine_scale, seed, options, negative_prompt],
outputs=[image_result],
api_name="freescalehf")
if __name__ == "__main__":
demo.queue(max_size=8).launch()