Spaces:
Sleeping
Sleeping
File size: 6,716 Bytes
713ec7d 3fad0f4 713ec7d c331654 55853c4 c331654 01e731e c331654 39a46ec c331654 a9684d6 fac7910 01e731e 39a46ec 01e731e 27a0683 713ec7d b107b62 713ec7d c331654 713ec7d e27aead 713ec7d a9684d6 713ec7d 8506797 713ec7d a9684d6 713ec7d b23661a b8b3bb8 fc06901 60ed05b e09933d fc06901 e09933d 884e72d fc06901 b107b62 9525232 17dc706 fc06901 5b111e3 fc06901 b8b3bb8 fac7910 7389249 7d24c56 f933a72 0e95c6d f933a72 fac7910 f933a72 7d24c56 f933a72 7389249 1e8d7dc 713ec7d 1e8d7dc 713ec7d 3e7f1a1 |
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 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 |
import gradio as gr
from PIL import Image
import torch
from diffusers import StableDiffusionPipeline
from free_lunch_utils import register_free_upblock2d, register_free_crossattn_upblock2d
model_id = "stabilityai/stable-diffusion-2-1"
# model_id = "./stable-diffusion-2-1"
pip_2_1 = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pip_2_1 = pip_2_1.to("cuda")
prompt_prev = None
sd_options_prev = None
seed_prev = None
sd_image_prev = None
def infer(prompt, sd_options, seed, b1, b2, s1, s2):
global prompt_prev
global sd_options_prev
global seed_prev
global sd_image_prev
# if sd_options == 'SD1.5':
# pip = pip_1_5
# elif sd_options == 'SD2.1':
# pip = pip_2_1
# else:
# pip = pip_1_4
pip = pip_2_1
run_baseline = False
if prompt != prompt_prev or sd_options != sd_options_prev or seed != seed_prev:
run_baseline = True
prompt_prev = prompt
sd_options_prev = sd_options
seed_prev = seed
if run_baseline:
register_free_upblock2d(pip, b1=1.0, b2=1.0, s1=1.0, s2=1.0)
register_free_crossattn_upblock2d(pip, b1=1.0, b2=1.0, s1=1.0, s2=1.0)
torch.manual_seed(seed)
print("Generating SD:")
sd_image = pip(prompt, num_inference_steps=25).images[0]
sd_image_prev = sd_image
else:
sd_image = sd_image_prev
register_free_upblock2d(pip, b1=b1, b2=b2, s1=s1, s2=s1)
register_free_crossattn_upblock2d(pip, b1=b1, b2=b2, s1=s1, s2=s1)
torch.manual_seed(seed)
print("Generating FreeU:")
freeu_image = pip(prompt, num_inference_steps=25).images[0]
# First SD, then freeu
images = [sd_image, freeu_image]
return images
examples = [
[
"A drone view of celebration with Christma tree and fireworks, starry sky - background.",
],
[
"happy dog wearing a yellow turtleneck, studio, portrait, facing camera, studio, dark bg"
],
[
"Campfire at night in a snowy forest with starry sky in the background."
],
[
"a fantasy landscape, trending on artstation"
],
[
"Busy freeway at night."
],
[
"An astronaut is riding a horse in the space in a photorealistic style."
],
[
"Turtle swimming in ocean."
],
[
"A storm trooper vacuuming the beach."
],
[
"An astronaut feeding ducks on a sunny afternoon, reflection from the water."
],
[
"Fireworks."
],
[
"A fat rabbit wearing a purple robe walking through a fantasy landscape."
],
[
"A koala bear playing piano in the forest."
],
[
"An astronaut flying in space, 4k, high resolution."
],
[
"Flying through fantasy landscapes, 4k, high resolution."
],
[
"A small cabin on top of a snowy mountain in the style of Disney, artstation",
],
[
"half human half cat, a human cat hybrid",
],
[
"a drone flying over a snowy forest."
],
]
css = """
h1 {
text-align: center;
}
#component-0 {
max-width: 730px;
margin: auto;
}
"""
block = gr.Blocks(css='style.css')
options = ['SD2.1']
with block:
gr.Markdown("# SD 2.1 vs. FreeU")
with gr.Group():
with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):
with gr.Column():
text = gr.Textbox(
label="Enter your prompt",
show_label=False,
max_lines=1,
placeholder="Enter your prompt",
container=False,
)
btn = gr.Button("Generate image", scale=0)
with gr.Row():
sd_options = gr.Dropdown(["SD2.1"], label="SD options", value="SD2.1", visible=False)
with gr.Group():
with gr.Row():
with gr.Accordion('FreeU Parameters (feel free to adjust these parameters based on your prompt): ', open=False):
with gr.Row():
b1 = gr.Slider(label='b1: backbone factor of the first stage block of decoder',
minimum=1,
maximum=1.6,
step=0.01,
value=1.1)
b2 = gr.Slider(label='b2: backbone factor of the second stage block of decoder',
minimum=1,
maximum=1.6,
step=0.01,
value=1.2)
with gr.Row():
s1 = gr.Slider(label='s1: skip factor of the first stage block of decoder',
minimum=0,
maximum=1,
step=0.1,
value=0.2)
s2 = gr.Slider(label='s2: skip factor of the second stage block of decoder',
minimum=0,
maximum=1,
step=0.1,
value=0.2)
seed = gr.Slider(label='seed',
minimum=0,
maximum=1000,
step=1,
value=42)
with gr.Row():
with gr.Group():
# btn = gr.Button("Generate image", scale=0)
with gr.Row():
with gr.Column() as c1:
image_1 = gr.Image(interactive=False)
image_1_label = gr.Markdown("SD")
with gr.Group():
# btn = gr.Button("Generate image", scale=0)
with gr.Row():
with gr.Column() as c2:
image_2 = gr.Image(interactive=False)
image_2_label = gr.Markdown("FreeU")
ex = gr.Examples(examples=examples, fn=infer, inputs=[text, sd_options, seed, b1, b2, s1, s2], outputs=[image_1, image_2], cache_examples=False)
ex.dataset.headers = [""]
text.submit(infer, inputs=[text, sd_options, seed, b1, b2, s1, s2], outputs=[image_1, image_2])
btn.click(infer, inputs=[text, sd_options, seed, b1, b2, s1, s2], outputs=[image_1, image_2])
block.launch()
# block.queue(default_enabled=False).launch(share=False)
|