FreeU / app.py
ChenyangSi's picture
Update app.py
9b76cbb
raw
history blame
6.54 kB
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 horse galloping on the ocean",
],
[
"a teddy bear walking in the snowstorm"
],
[
"Campfire at night in a snowy forest with starry sky in the background."
],
[
"a fantasy landscape, trending on artstation"
],
[
"An astronaut flying in space, 4k, high resolution."
],
[
"An astronaut is riding a horse in the space in a photorealistic style."
],
[
"Turtle swimming in ocean."
],
[
"A storm trooper vacuuming the beach."
],
[
"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=2.0,
step=0.01,
value=1.4)
b2 = gr.Slider(label='b2: backbone factor of the second stage block of decoder',
minimum=1,
maximum=2.0,
step=0.01,
value=1.6)
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.9)
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)