from diffusers import CycleDiffusionPipeline, DDIMScheduler 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() model_id_or_path = "CompVis/stable-diffusion-v1-4" if is_colab: scheduler = DDIMScheduler.from_config(model_id_or_path, subfolder="scheduler") pipe = CycleDiffusionPipeline.from_pretrained(model_id_or_path, scheduler=scheduler) 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) tokenizer = pipe.tokenizer if torch.cuda.is_available(): pipe = pipe.to("cuda") device_print = "GPU ๐Ÿ”ฅ" if torch.cuda.is_available() else "CPU ๐Ÿฅถ" device = "cuda" if torch.cuda.is_available() else "cpu" 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:]).float() 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) 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) 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) 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) 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.float32) 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"""

CycleDiffusion with Stable Diffusion

Demo for CycleDiffusion with Stable Diffusion.
CycleDiffusion (๐Ÿ“„ Paper link | ๐Ÿงจ Pipeline doc) is an image-to-image translation method that supports stochastic samplers for diffusion models.
We also support the combination of CycleDiffusion and Cross Attention Control (CAC | ๐Ÿ“„ Paper link). CAC is a technique to transfer the attention map from the source prompt to the target prompt.

Quick start:
1. Click one row of Examples at the end of this page. It will fill all inputs needed.
2. Click the "Edit" button.

How to use:
1. Upload an image.
2. Enter the source and target prompts.
3. Select the source guidance scale (for "encoding") and the target guidance scale (for "decoding").
4. Select the strength (smaller strength means better content preservation).
5 (optional). Configurate Cross Attention Control options (e.g., CAC type, cross replace steps, self replace steps).
6 (optional). Configurate other options (e.g., image size, inference steps, random seed).
7. Click the "Edit" button.

Notes:
1. CycleDiffusion is likely to fail when drastic changes are intended (e.g., changing a large black car to red).
2. The value of strength can be set larger when CAC is used.
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.
4. If CAC type is "Refine", the source prompt be a subsequence of the target prompt; otherwise, an error will be raised.

You can skip the queue using Colab: Open In Colab

Running on {device_print}{(" in a Google Colab." if is_colab else "")}

""" ) 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="Edit") 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="Edit") with gr.Tab("Other options"): with gr.Group(): with gr.Row(): num_inference_steps = gr.Slider(label="Number of inference steps", value=100, minimum=25, maximum=500, step=1) width = gr.Slider(label="Width", value=512, minimum=64, maximum=1024, step=8) height = gr.Slider(label="Height", value=512, minimum=64, 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="Edit") 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, "images/astronaut_horse.png", 0.8, "None", 0, 0], ["An astronaut riding a horse", "An astronaut riding a elephant", 1, 2, 100, "images/astronaut_horse.png", 0.9, "Replace", 0.15, 0.10], ["A black colored car.", "A blue colored car.", 1, 3, 100, "images/black_car.png", 0.85, "None", 0, 0], ["A black colored car.", "A blue colored car.", 1, 5, 100, "images/black_car.png", 0.95, "Replace", 0.8, 0.4], ["A black colored car.", "A red colored car.", 1, 5, 100, "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, "images/mausoleum.png", 0.9, "None", 0, 0], ["An aerial view of autumn scene.", "An aerial view of winter scene.", 1, 5, 100, "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, "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, "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, "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, "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, "images/apple_bag.png", 0.89, "None", 0, 0], ], [source_prompt, target_prompt, source_guidance_scale, guidance_scale, num_inference_steps, img, strength, cross_attention_control, cross_replace_steps, self_replace_steps], image_out, inference, cache_examples=False) 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)