ChenWu98's picture
Update app.py
dae2af9
from diffusers import CycleDiffusionPipeline, DDIMScheduler
import os
import gradio as gr
import torch
from PIL import Image
import utils
import ptp_utils
import seq_aligner
import torch.nn.functional as nnf
from typing import Optional, Union, Tuple, List, Callable, Dict
import abc
LOW_RESOURCE = False
MAX_NUM_WORDS = 77
is_colab = utils.is_google_colab()
colab_instruction = "" if is_colab else """
<p>You can skip the queue using Colab: <a href="https://colab.research.google.com/gist/ChenWu98/0aa4fe7be80f6b45d3d055df9f14353a/copy-of-fine-tuned-diffusion-gradio.ipynb"><img data-canonical-src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab" src="https://colab.research.google.com/assets/colab-badge.svg"></a></p>"""
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
model_id_or_path = "CompVis/stable-diffusion-v1-4"
device_print = "GPU πŸ”₯" if torch.cuda.is_available() else "CPU πŸ₯Ά"
device = "cuda" if torch.cuda.is_available() else "cpu"
if is_colab:
scheduler = DDIMScheduler.from_config(model_id_or_path, subfolder="scheduler")
pipe = CycleDiffusionPipeline.from_pretrained(model_id_or_path, scheduler=scheduler, torch_dtype=torch_dtype)
else:
# import streamlit as st
# scheduler = DDIMScheduler.from_config(model_id_or_path, use_auth_token=st.secrets["USER_TOKEN"], subfolder="scheduler")
# pipe = CycleDiffusionPipeline.from_pretrained(model_id_or_path, use_auth_token=st.secrets["USER_TOKEN"], scheduler=scheduler, torch_dtype=torch_dtype)
scheduler = DDIMScheduler.from_config(model_id_or_path, use_auth_token=os.environ.get("USER_TOKEN"), subfolder="scheduler")
pipe = CycleDiffusionPipeline.from_pretrained(model_id_or_path, use_auth_token=os.environ.get("USER_TOKEN"), scheduler=scheduler, torch_dtype=torch_dtype)
tokenizer = pipe.tokenizer
if torch.cuda.is_available():
pipe = pipe.to("cuda")
class LocalBlend:
def __call__(self, x_t, attention_store):
k = 1
maps = attention_store["down_cross"][2:4] + attention_store["up_cross"][:3]
maps = [item.reshape(self.alpha_layers.shape[0], -1, 1, 16, 16, MAX_NUM_WORDS) for item in maps]
maps = torch.cat(maps, dim=1)
maps = (maps * self.alpha_layers).sum(-1).mean(1)
mask = nnf.max_pool2d(maps, (k * 2 + 1, k * 2 + 1), (1, 1), padding=(k, k))
mask = nnf.interpolate(mask, size=(x_t.shape[2:]))
mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0]
mask = mask.gt(self.threshold)
mask = (mask[:1] + mask[1:]).to(x_t.dtype)
x_t = x_t[:1] + mask * (x_t - x_t[:1])
return x_t
def __init__(self, prompts: List[str], words: [List[List[str]]], threshold=.3):
alpha_layers = torch.zeros(len(prompts), 1, 1, 1, 1, MAX_NUM_WORDS)
for i, (prompt, words_) in enumerate(zip(prompts, words)):
if type(words_) is str:
words_ = [words_]
for word in words_:
ind = ptp_utils.get_word_inds(prompt, word, tokenizer)
alpha_layers[i, :, :, :, :, ind] = 1
self.alpha_layers = alpha_layers.to(device).to(torch_dtype)
self.threshold = threshold
class AttentionControl(abc.ABC):
def step_callback(self, x_t):
return x_t
def between_steps(self):
return
@property
def num_uncond_att_layers(self):
return self.num_att_layers if LOW_RESOURCE else 0
@abc.abstractmethod
def forward(self, attn, is_cross: bool, place_in_unet: str):
raise NotImplementedError
def __call__(self, attn, is_cross: bool, place_in_unet: str):
if self.cur_att_layer >= self.num_uncond_att_layers:
if LOW_RESOURCE:
attn = self.forward(attn, is_cross, place_in_unet)
else:
h = attn.shape[0]
attn[h // 2:] = self.forward(attn[h // 2:], is_cross, place_in_unet)
self.cur_att_layer += 1
if self.cur_att_layer == self.num_att_layers + self.num_uncond_att_layers:
self.cur_att_layer = 0
self.cur_step += 1
self.between_steps()
return attn
def reset(self):
self.cur_step = 0
self.cur_att_layer = 0
def __init__(self):
self.cur_step = 0
self.num_att_layers = -1
self.cur_att_layer = 0
class EmptyControl(AttentionControl):
def forward(self, attn, is_cross: bool, place_in_unet: str):
return attn
class AttentionStore(AttentionControl):
@staticmethod
def get_empty_store():
return {"down_cross": [], "mid_cross": [], "up_cross": [],
"down_self": [], "mid_self": [], "up_self": []}
def forward(self, attn, is_cross: bool, place_in_unet: str):
key = f"{place_in_unet}_{'cross' if is_cross else 'self'}"
if attn.shape[1] <= 32 ** 2: # avoid memory overhead
self.step_store[key].append(attn)
return attn
def between_steps(self):
if len(self.attention_store) == 0:
self.attention_store = self.step_store
else:
for key in self.attention_store:
for i in range(len(self.attention_store[key])):
self.attention_store[key][i] += self.step_store[key][i]
self.step_store = self.get_empty_store()
def get_average_attention(self):
average_attention = {key: [item / self.cur_step for item in self.attention_store[key]] for key in self.attention_store}
return average_attention
def reset(self):
super(AttentionStore, self).reset()
self.step_store = self.get_empty_store()
self.attention_store = {}
def __init__(self):
super(AttentionStore, self).__init__()
self.step_store = self.get_empty_store()
self.attention_store = {}
class AttentionControlEdit(AttentionStore, abc.ABC):
def step_callback(self, x_t):
if self.local_blend is not None:
x_t = self.local_blend(x_t, self.attention_store)
return x_t
def replace_self_attention(self, attn_base, att_replace):
if att_replace.shape[2] <= 16 ** 2:
return attn_base.unsqueeze(0).expand(att_replace.shape[0], *attn_base.shape)
else:
return att_replace
@abc.abstractmethod
def replace_cross_attention(self, attn_base, att_replace):
raise NotImplementedError
def forward(self, attn, is_cross: bool, place_in_unet: str):
super(AttentionControlEdit, self).forward(attn, is_cross, place_in_unet)
if is_cross or (self.num_self_replace[0] <= self.cur_step < self.num_self_replace[1]):
h = attn.shape[0] // self.batch_size
attn = attn.reshape(self.batch_size, h, *attn.shape[1:])
attn_base, attn_repalce = attn[0], attn[1:]
if is_cross:
alpha_words = self.cross_replace_alpha[self.cur_step]
attn_replace_new = self.replace_cross_attention(attn_base, attn_repalce) * alpha_words + (1 - alpha_words) * attn_repalce
attn[1:] = attn_replace_new
else:
attn[1:] = self.replace_self_attention(attn_base, attn_repalce)
attn = attn.reshape(self.batch_size * h, *attn.shape[2:])
return attn
def __init__(self, prompts, num_steps: int,
cross_replace_steps: Union[float, Tuple[float, float], Dict[str, Tuple[float, float]]],
self_replace_steps: Union[float, Tuple[float, float]],
local_blend: Optional[LocalBlend]):
super(AttentionControlEdit, self).__init__()
self.batch_size = len(prompts)
self.cross_replace_alpha = ptp_utils.get_time_words_attention_alpha(prompts, num_steps, cross_replace_steps, tokenizer).to(device).to(torch_dtype)
if type(self_replace_steps) is float:
self_replace_steps = 0, self_replace_steps
self.num_self_replace = int(num_steps * self_replace_steps[0]), int(num_steps * self_replace_steps[1])
self.local_blend = local_blend
class AttentionReplace(AttentionControlEdit):
def replace_cross_attention(self, attn_base, att_replace):
return torch.einsum('hpw,bwn->bhpn', attn_base, self.mapper)
def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float,
local_blend: Optional[LocalBlend] = None):
super(AttentionReplace, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend)
self.mapper = seq_aligner.get_replacement_mapper(prompts, tokenizer).to(device).to(torch_dtype)
class AttentionRefine(AttentionControlEdit):
def replace_cross_attention(self, attn_base, att_replace):
attn_base_replace = attn_base[:, :, self.mapper].permute(2, 0, 1, 3)
attn_replace = attn_base_replace * self.alphas + att_replace * (1 - self.alphas)
return attn_replace
def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float,
local_blend: Optional[LocalBlend] = None):
super(AttentionRefine, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend)
self.mapper, alphas = seq_aligner.get_refinement_mapper(prompts, tokenizer)
self.mapper, alphas = self.mapper.to(device), alphas.to(device).to(torch_dtype)
self.alphas = alphas.reshape(alphas.shape[0], 1, 1, alphas.shape[1])
def get_equalizer(text: str, word_select: Union[int, Tuple[int, ...]], values: Union[List[float], Tuple[float, ...]]):
if type(word_select) is int or type(word_select) is str:
word_select = (word_select,)
equalizer = torch.ones(len(values), 77)
values = torch.tensor(values, dtype=torch_dtype)
for word in word_select:
inds = ptp_utils.get_word_inds(text, word, tokenizer)
equalizer[:, inds] = values
return equalizer
def inference(source_prompt, target_prompt, source_guidance_scale=1, guidance_scale=5, num_inference_steps=100,
width=512, height=512, seed=0, img=None, strength=0.7,
cross_attention_control="None", cross_replace_steps=0.8, self_replace_steps=0.4):
torch.manual_seed(seed)
ratio = min(height / img.height, width / img.width)
img = img.resize((int(img.width * ratio), int(img.height * ratio)))
# create the CAC controller.
if cross_attention_control == "Replace":
controller = AttentionReplace([source_prompt, target_prompt],
num_inference_steps,
cross_replace_steps=cross_replace_steps,
self_replace_steps=self_replace_steps,
)
ptp_utils.register_attention_control(pipe, controller)
elif cross_attention_control == "Refine":
controller = AttentionRefine([source_prompt, target_prompt],
num_inference_steps,
cross_replace_steps=cross_replace_steps,
self_replace_steps=self_replace_steps,
)
ptp_utils.register_attention_control(pipe, controller)
elif cross_attention_control == "None":
controller = EmptyControl()
ptp_utils.register_attention_control(pipe, controller)
else:
raise ValueError("Unknown cross_attention_control: {}".format(cross_attention_control))
results = pipe(prompt=target_prompt,
source_prompt=source_prompt,
init_image=img,
num_inference_steps=num_inference_steps,
eta=0.1,
strength=strength,
guidance_scale=guidance_scale,
source_guidance_scale=source_guidance_scale,
)
return replace_nsfw_images(results)
def replace_nsfw_images(results):
for i in range(len(results.images)):
if results.nsfw_content_detected[i]:
results.images[i] = Image.open("nsfw.png")
return results.images[0]
css = """.cycle-diffusion-div div{display:inline-flex;align-items:center;gap:.8rem;font-size:1.75rem}.cycle-diffusion-div div h1{font-weight:900;margin-bottom:7px}.cycle-diffusion-div p{margin-bottom:10px;font-size:94%}.cycle-diffusion-div p a{text-decoration:underline}.tabs{margin-top:0;margin-bottom:0}#gallery{min-height:20rem}
"""
with gr.Blocks(css=css) as demo:
gr.HTML(
f"""
<div class="cycle-diffusion-div">
<div>
<h1>CycleDiffusion with Stable Diffusion</h1>
</div>
<p>
Demo for CycleDiffusion with Stable Diffusion. <br>
CycleDiffusion (<a href="https://arxiv.org/abs/2210.05559">πŸ“„ Paper link</a> | <a href="https://huggingface.co/docs/diffusers/main/en/api/pipelines/cycle_diffusion">🧨 Pipeline doc</a>) is an image-to-image translation method that supports stochastic samplers for diffusion models. <br>
We also support the combination of CycleDiffusion and Cross Attention Control (CAC | <a href="https://arxiv.org/abs/2208.01626">πŸ“„ Paper link</a>). CAC is a technique to transfer the attention map from the source prompt to the target prompt. <br>
</p>
<p>
<b>Quick start</b>: <br>
1. Click one row of Examples at the end of this page. It will fill all inputs needed. <br>
2. Click the "Run CycleDiffusion" button. <br>
</p>
<p>
{colab_instruction}
Running on <b>{device_print}</b>{(" in a <b>Google Colab</b>." if is_colab else "")}
</p>
</div>
"""
)
with gr.Accordion("See Details", open=False):
gr.HTML(
f"""
<div class="cycle-diffusion-div">
<p>
<b>How to use:</b> <br>
1. Upload an image. <br>
2. Enter the source and target prompts. <br>
3. Select the source guidance scale (for "encoding") and the target guidance scale (for "decoding"). <br>
4. Select the strength (smaller strength means better content preservation). <br>
5 (optional). Configurate Cross Attention Control options (e.g., CAC type, cross replace steps, self replace steps). <br>
6 (optional). Configurate other options (e.g., image size, inference steps, random seed). <br>
7. Click the "Run CycleDiffusion" button. <br>
</p>
<p>
<b>Notes:</b> <br>
1. CycleDiffusion is likely to fail when drastic changes are intended (e.g., changing a large black car to red). <br>
2. The value of strength can be set larger when CAC is used. <br>
3. If CAC type is "Replace", the source and target prompts should differ in only one token; otherwise, an error will be raised. This is why we deliberately make some grammar mistakes in Examples.<br>
4. If CAC type is "Refine", the source prompt be a subsequence of the target prompt; otherwise, an error will be raised. <br>
</p>
<p>
<b>Runtimes:</b> <br>
1. 20s on A10G. <br>
</p>
</div>
"""
)
with gr.Row():
with gr.Column(scale=55):
with gr.Group():
img = gr.Image(label="Input image", height=512, tool="editor", type="pil")
image_out = gr.Image(label="Output image", height=512)
# gallery = gr.Gallery(
# label="Generated images", show_label=False, elem_id="gallery"
# ).style(grid=[1], height="auto")
with gr.Column(scale=45):
with gr.Tab("Edit options"):
with gr.Group():
with gr.Row():
source_prompt = gr.Textbox(label="Source prompt", placeholder="Source prompt describes the input image")
source_guidance_scale = gr.Slider(label="Source guidance scale", value=1, minimum=1, maximum=10)
with gr.Row():
target_prompt = gr.Textbox(label="Target prompt", placeholder="Target prompt describes the output image")
guidance_scale = gr.Slider(label="Target guidance scale", value=5, minimum=1, maximum=10)
with gr.Row():
strength = gr.Slider(label="Strength", value=0.7, minimum=0.5, maximum=1, step=0.01)
with gr.Row():
generate1 = gr.Button(value="Run CycleDiffusion")
with gr.Tab("CAC options"):
with gr.Group():
with gr.Row():
cross_attention_control = gr.Radio(label="CAC type", choices=["None", "Replace", "Refine"], value="None")
with gr.Row():
# If not "None", the following two parameters will be used.
cross_replace_steps = gr.Slider(label="Cross replace steps", value=0.8, minimum=0.0, maximum=1, step=0.01)
self_replace_steps = gr.Slider(label="Self replace steps", value=0.4, minimum=0.0, maximum=1, step=0.01)
with gr.Row():
generate2 = gr.Button(value="Run CycleDiffusion")
with gr.Tab("Other options"):
with gr.Group():
with gr.Row():
num_inference_steps = gr.Slider(label="Inference steps", value=100, minimum=25, maximum=500, step=1)
width = gr.Slider(label="Width", value=512, minimum=512, maximum=1024, step=8)
height = gr.Slider(label="Height", value=512, minimum=512, maximum=1024, step=8)
with gr.Row():
seed = gr.Slider(0, 2147483647, label='Seed', value=0, step=1)
with gr.Row():
generate3 = gr.Button(value="Run CycleDiffusion")
inputs = [source_prompt, target_prompt, source_guidance_scale, guidance_scale, num_inference_steps,
width, height, seed, img, strength,
cross_attention_control, cross_replace_steps, self_replace_steps]
generate1.click(inference, inputs=inputs, outputs=image_out)
generate2.click(inference, inputs=inputs, outputs=image_out)
generate3.click(inference, inputs=inputs, outputs=image_out)
ex = gr.Examples(
[
["An astronaut riding a horse", "An astronaut riding an elephant", 1, 2, 100, 512, 512, 0, "images/astronaut_horse.png", 0.8, "None", 0, 0],
["An astronaut riding a horse", "An astronaut riding a elephant", 1, 2, 100, 512, 512, 0, "images/astronaut_horse.png", 0.9, "Replace", 0.15, 0.10],
["A black colored car.", "A blue colored car.", 1, 3, 100, 512, 512, 0, "images/black_car.png", 0.85, "None", 0, 0],
["A black colored car.", "A blue colored car.", 1, 5, 100, 512, 512, 0, "images/black_car.png", 0.95, "Replace", 0.8, 0.4],
["A black colored car.", "A red colored car.", 1, 5, 100, 512, 512, 0, "images/black_car.png", 1, "Replace", 0.8, 0.4],
["An aerial view of autumn scene.", "An aerial view of winter scene.", 1, 5, 100, 512, 512, 0, "images/mausoleum.png", 0.9, "None", 0, 0],
["An aerial view of autumn scene.", "An aerial view of winter scene.", 1, 5, 100, 512, 512, 0, "images/mausoleum.png", 1, "Replace", 0.8, 0.4],
["A green apple and a black backpack on the floor.", "A red apple and a black backpack on the floor.", 1, 7, 100, 512, 512, 0, "images/apple_bag.png", 0.9, "None", 0, 0],
["A green apple and a black backpack on the floor.", "A red apple and a black backpack on the floor.", 1, 7, 100, 512, 512, 0, "images/apple_bag.png", 0.9, "Replace", 0.8, 0.4],
["A hotel room with red flowers on the bed.", "A hotel room with a cat sitting on the bed.", 1, 4, 100, 512, 512, 0, "images/flower_hotel.png", 0.8, "None", 0, 0],
["A hotel room with red flowers on the bed.", "A hotel room with blue flowers on the bed.", 1, 5, 100, 512, 512, 0, "images/flower_hotel.png", 0.95, "None", 0, 0],
["A green apple and a black backpack on the floor.", "Two green apples and a black backpack on the floor.", 1, 5, 100, 512, 512, 0, "images/apple_bag.png", 0.89, "None", 0, 0],
],
[source_prompt, target_prompt, source_guidance_scale, guidance_scale, num_inference_steps,
width, height, seed, img, strength,
cross_attention_control, cross_replace_steps, self_replace_steps],
image_out, inference, cache_examples=True)
gr.Markdown('''
Space built with Diffusers 🧨 by HuggingFace πŸ€—.
[![Twitter Follow](https://img.shields.io/twitter/follow/ChenHenryWu?style=social)](https://twitter.com/ChenHenryWu)
![visitors](https://visitor-badge.glitch.me/badge?page_id=ChenWu98.CycleDiffusion)
''')
if not is_colab:
demo.queue(concurrency_count=1)
demo.launch(debug=is_colab, share=is_colab)