Add application file

app.py

@@ -0,0 +1,429 @@
+import os
+import torch
+import random
+import numpy as np
+import gradio as gr
+from glob import glob
+from datetime import datetime
+
+from diffusers import StableDiffusionPipeline,AutoencoderKL
+from diffusers import DDIMScheduler, LCMScheduler, EulerDiscreteScheduler
+
+import torch.nn.functional as F
+from PIL import Image,ImageDraw
+from utils.pipeline import ZePoPipeline
+from utils.attn_control import AttentionStyle
+from torchvision.utils import save_image
+import utils.ptp_utils as ptp_utils
+
+import torchvision.transforms as transforms
+
+try:
+    import xformers
+    is_xformers = True
+except ImportError:
+    is_xformers = False
+
+css = """
+.toolbutton {
+    margin-buttom: 0em 0em 0em 0em;
+    max-width: 2.5em;
+    min-width: 2.5em !important;
+    height: 2.5em;
+}
+"""
+# import sys
+# sys.setrecursionlimit(100000)
+
+
+class GlobalText:
+    def __init__(self):
+        
+        # config dirs
+        self.basedir                = os.getcwd()
+        self.stable_diffusion_dir   = os.path.join(self.basedir, "models", "StableDiffusion")
+        self.personalized_model_dir = './models/Stable-diffusion'
+        self.lora_model_dir         = './models/Lora'
+        self.savedir            = os.path.join(self.basedir, "samples", datetime.now().strftime("Gradio-%Y-%m-%dT%H-%M-%S"))
+        self.savedir_sample         = os.path.join(self.savedir, "sample")
+
+        # self.savedir_mask         = os.path.join(self.savedir, "mask")
+
+        self.stable_diffusion_list   = ["SimianLuo/LCM_Dreamshaper_v7"
+                                        ]
+        self.personalized_model_list = []
+        self.lora_model_list = []
+
+        self.tokenizer             = None
+        self.text_encoder          = None
+        self.vae                   = None
+        self.unet                  = None
+        self.pipeline              = None
+        self.torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+        self.lora_model_state_dict = {}
+        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+    def init_source_image_path(self, source_path):
+        self.source_paths = sorted(glob(os.path.join(source_path, '*')))
+        self.max_source_index = len(self.source_paths) // 12
+        return self.source_paths[0:12]
+    def init_style_image_path(self, style_path):
+        self.style_paths = sorted(glob(os.path.join(style_path, '*')))
+        self.max_style_index = len(self.style_paths) // 12
+        return self.style_paths[0:12]
+    def init_results_image_path(self):
+        results_paths = [os.path.join(self.savedir_sample, file) for file in os.listdir(self.savedir_sample)]
+        self.results_paths = sorted(results_paths, key=os.path.getctime, reverse=True)
+        self.max_results_index = len(self.results_paths) // 12
+        return self.results_paths[0:12]
+    
+    def load_base_pipeline(self, model_path):
+        
+        time_start = datetime.now()
+
+        self.scheduler = 'LCM'
+        scheduler = LCMScheduler.from_pretrained(model_path, subfolder="scheduler")
+        self.pipeline = ZePoPipeline.from_pretrained(model_path,scheduler=scheduler,torch_dtype=torch.float16,).to('cuda')
+        if is_xformers:
+            self.pipeline.enable_xformers_memory_efficient_attention()
+        time_end = datetime.now()
+        print(f'Load {model_path} successful in {time_end-time_start}')
+        return gr.Dropdown()
+    
+    def refresh_stable_diffusion(self,model_path):
+        
+        self.load_base_pipeline(model_path)
+        
+        return self.stable_diffusion_list[0]
+
+    def update_base_model(self, base_model_dropdown):
+        if self.pipeline is None:
+            gr.Info(f"Please select a pretrained model path.")
+            return None
+        else:
+            base_model = self.personalized_model_list[base_model_dropdown]
+            mid_model = StableDiffusionPipeline.from_single_file(base_model)
+            self.pipeline.vae = mid_model.vae
+            self.pipeline.unet = mid_model.unet
+            self.pipeline.text_encoder = mid_model.text_encoder
+            self.pipeline.to(self.device)
+            self.personal_model_loaded = base_model_dropdown.split('.')[0]
+            print(f'load {base_model_dropdown} model success!')
+            return gr.Dropdown()
+
+    
+    def generate(self, source, style, 
+                       num_steps, co_feat_step,strength,
+                       start_ac_layer, end_ac_layer,
+                       sty_guidance,cfg_scale, mix_q_scale,
+                       Scheduler, save_intermediate, seed, de_bug,
+                       target_prompt, negative_prompt_textbox,
+                       width_slider,height_slider,
+                       tome_sx, tome_sy, tome_ratio,tome,
+                       ):
+
+
+        os.makedirs(self.savedir, exist_ok=True)
+        os.makedirs(self.savedir_sample, exist_ok=True)
+        
+        if self.pipeline == None:
+            self.refresh_stable_diffusion(self.stable_diffusion_list[-1])
+        model = self.pipeline
+
+        if Scheduler == 'DDIM':
+            model.scheduler = DDIMScheduler.from_config(model.scheduler.config)
+            print(f"Successful adoption of DDIM scheduler")
+        if Scheduler == 'LCM':
+            model.scheduler = LCMScheduler.from_config(model.scheduler.config)
+            print(f"Successful adoption of LCM scheduler")
+        if Scheduler == 'EulerDiscrete':
+            model.scheduler = EulerDiscreteScheduler.from_config(model.scheduler.config)
+
+        if seed != '-1' and seed != "": torch.manual_seed(int(seed))
+        else: torch.seed()
+        
+        seed = torch.initial_seed()
+        print(f"Seed: {seed}")
+
+        self.sample_count = len(os.listdir(self.savedir_sample))
+        
+
+        prompts = [target_prompt] * 3
+        source = source.resize((width_slider, height_slider))
+        style = style.resize((width_slider, height_slider))
+
+
+        with torch.no_grad():
+
+            controller = AttentionStyle(num_steps, 
+                                         start_ac_layer,
+                                         end_ac_layer,
+                                         style_guidance=sty_guidance,
+                                         mix_q_scale=mix_q_scale,
+                                         de_bug=de_bug,
+                                         )
+                                                        
+            ptp_utils.register_attention_control(model, controller, 
+                                                 tome, 
+                                                 sx=tome_sx,
+                                                 sy=tome_sy,
+                                                 ratio=tome_ratio,
+                                                 de_bug=de_bug,)
+
+            time_begin = datetime.now()
+            generate_image = model(prompt=prompts,
+                                negative_prompt=negative_prompt_textbox,
+                                image=source,
+                                style=style,
+                                num_inference_steps=num_steps,
+                                eta=0.0,
+                                guidance_scale=cfg_scale,
+                                strength=strength,
+                                save_intermediate=save_intermediate,
+                                fix_step_index=co_feat_step,
+                                de_bug = de_bug,
+                                callback = None
+                   ).images
+            time_end = datetime.now()
+            print('generate one image with time {}'.format(time_end-time_begin))
+
+            save_file_name = f"{self.sample_count}_step{num_steps}_sl{start_ac_layer}_el{end_ac_layer}_ST{strength}_CF{co_feat_step}_STG{sty_guidance}_MQ{mix_q_scale}_CFG{cfg_scale}_seed{seed}.jpg"
+
+            
+            save_file_path = os.path.join(self.savedir, save_file_name)
+                    
+            save_image(torch.tensor(generate_image).permute(0, 3, 1, 2), save_file_path, nrow=3, padding=0)
+            save_image(torch.tensor(generate_image[2:]).permute(0, 3, 1, 2), os.path.join(self.savedir_sample, save_file_name), nrow=3, padding=0)
+            self.init_results_image_path()
+        return [
+            generate_image[0],
+            generate_image[1],
+            generate_image[2],
+            self.init_results_image_path()
+            ]
+    
+
+global_text = GlobalText()
+
+
+def ui():
+    with gr.Blocks(css=css) as demo:
+        gr.Markdown(
+            """
+            # [ZePo: Zero-Shot Portrait Stylization with Faster Sampling](https://arxiv.org/abs/2408.05492)
+            Jin Liu, Huaibo Huang, Jie Cao, Ran He<br>
+            [Arxiv](https://arxiv.org/abs/2408.05492) | [Github](https://github.com/liujin112/ZePo)
+            """
+        )
+        with gr.Column(variant="panel"):
+            gr.Markdown(
+                """
+                ### 1. Select a pretrained model.
+                """
+            )
+            with gr.Row():
+                stable_diffusion_dropdown = gr.Dropdown(
+                    label="Pretrained Model Path",
+                    choices=global_text.stable_diffusion_list,
+                    interactive=True,
+                    allow_custom_value=True
+                )
+                stable_diffusion_dropdown.change(fn=global_text.load_base_pipeline, inputs=[stable_diffusion_dropdown], outputs=[stable_diffusion_dropdown])
+                
+                stable_diffusion_refresh_button = gr.Button(value="\U0001F503", elem_classes="toolbutton")
+                def update_stable_diffusion(stable_diffusion_dropdown):
+                    global_text.refresh_stable_diffusion(stable_diffusion_dropdown)
+                    
+                stable_diffusion_refresh_button.click(fn=update_stable_diffusion, inputs=[stable_diffusion_dropdown], outputs=[stable_diffusion_dropdown])
+
+
+        with gr.Column(variant="panel"):
+            gr.Markdown(
+                """
+                ### 2. Configs for ZePo.
+                """
+            )
+            with gr.Tab("Configs"):
+                
+                with gr.Row():
+                    with gr.Column():
+                        with gr.Row():
+                            source_image = gr.Image(label="Source Image",  elem_id="img2maskimg", sources="upload",  type="pil",image_mode="RGB", height=256)
+                            style_image = gr.Image(label="Style Image", elem_id="img2maskimg", sources="upload", type="pil", image_mode="RGB", height=256)
+                        
+                        generate_image = gr.Image(label="Image with PortraitDiff", type="pil", interactive=True, image_mode="RGB", height=512)
+
+
+                        with gr.Row():
+                            recons_content = gr.Image(label="reconstructed content", type="pil", image_mode="RGB", height=256)
+                            recons_style = gr.Image(label="reconstructed style", type="pil", image_mode="RGB", height=256)
+                        prompt_textbox = gr.Textbox(label="Prompt", value='head', lines=1)
+                        negative_prompt_textbox = gr.Textbox(label="Negative prompt", lines=1)
+                    with gr.Row(equal_height=False):
+                        with gr.Column():
+                            with gr.Tab("Resolution"):
+                                width_slider     = gr.Slider(label="Width", value=512, minimum=256, maximum=1024, step=64)
+                                height_slider    = gr.Slider(label="Height", value=512, minimum=256, maximum=1024, step=64)
+                                Scheduler = gr.Dropdown(
+                                            ["DDIM", "LCM", "EulerDiscrete"],
+                                            value="LCM",
+                                            label="Scheduler", info="Select a Scheduler")
+
+
+                            with gr.Tab("Content Gallery"):
+                                
+                                with gr.Row():
+                                    source_path = gr.Textbox(value='./data/content', label="Source Path")
+                                    refresh_source_list_button = gr.Button(value="\U0001F503", elem_classes="toolbutton")
+                                source_gallery_index = gr.Slider(label="Index", value=0, minimum=0, maximum=50, step=1)   
+                                num_gallery_images = 12   
+                                source_image_gallery = gr.Gallery(value=[], columns=4, label="Source Image List")
+                                refresh_source_list_button.click(fn=global_text.init_source_image_path, inputs=[source_path], outputs=[source_image_gallery]) 
+
+                                def update_source_list(index):
+                                    if int(index) < 0:
+                                        index = 0
+                                    if int(index) > global_text.max_source_index:
+                                        index = global_text.max_source_index
+                                    return global_text.source_paths[int(index)*num_gallery_images:(int(index)+1)*num_gallery_images]
+                                
+                                source_gallery_index.change(fn=update_source_list, inputs=[source_gallery_index], outputs=[source_image_gallery])
+
+                            with gr.Tab("Style Gallery"):
+                                
+                                with gr.Row():
+                                    style_path = gr.Textbox(value='./data/style', label="style Path")
+                                    refresh_style_list_button = gr.Button(value="\U0001F503", elem_classes="toolbutton")
+                                style_gallery_index = gr.Slider(label="Index", value=0, minimum=0, maximum=50, step=1)   
+                                num_gallery_images = 12   
+                                style_image_gallery = gr.Gallery(value=[], columns=4, label="style Image List")
+                                refresh_style_list_button.click(fn=global_text.init_style_image_path, inputs=[style_path], outputs=[style_image_gallery]) 
+                                
+
+                                def update_style_list(index):
+                                    if int(index) < 0:
+                                        index = 0
+                                    if int(index) > global_text.max_style_index:
+                                        index = global_text.max_style_index
+                                    return global_text.style_paths[int(index)*num_gallery_images:(int(index)+1)*num_gallery_images]
+                                
+                                style_gallery_index.change(fn=update_style_list, inputs=[style_gallery_index], outputs=[style_image_gallery])
+                            
+                            with gr.Tab("Results Gallery"):
+                                with gr.Row():
+                                    refresh_results_list_button = gr.Button(value="\U0001F503", elem_classes="toolbutton")
+                                    results_gallery_index = gr.Slider(label="Index", value=0, minimum=0, maximum=50, step=1)   
+                                num_gallery_images = 12   
+                                results_image_gallery = gr.Gallery(value=[], columns=4, label="style Image List")
+                                refresh_results_list_button.click(fn=global_text.init_results_image_path, inputs=[], outputs=[results_image_gallery]) 
+                                
+
+                                def update_results_list(index):
+                                    if int(index) < 0:
+                                        index = 0
+                                    if int(index) > global_text.max_results_index:
+                                        index = global_text.max_results_index
+                                    return global_text.results_paths[int(index)*num_gallery_images:(int(index)+1)*num_gallery_images]
+                                
+                                results_gallery_index.change(fn=update_results_list, inputs=[results_gallery_index], outputs=[style_image_gallery])
+                            
+                            
+
+                            with gr.Row():    
+                                generate_button = gr.Button(value="Generate", variant='primary')
+
+                            with gr.Tab('Base Configs'):
+                                num_steps = gr.Slider(label="Total Steps", value=4, minimum=0, maximum=25, step=1)
+                                strength = gr.Slider(label="Noisy Ratio", value=0.5, minimum=0, maximum=1, step=0.01,info="How much noise applied to souce image, 50% for better balance.")
+                                co_feat_step = gr.Slider(label="Consistency Feature Extract Step", value=99, minimum=0, maximum=999, step=1)
+                                
+
+                                with gr.Row():    
+                                    start_ac_layer = gr.Slider(label="Start Layer of AC",
+                                                            minimum=0,
+                                                            maximum=16,
+                                                            value=8,
+                                                            step=1)
+                                    end_ac_layer = gr.Slider(label="End Layer of AC",
+                                                            minimum=0,
+                                                            maximum=16,
+                                                            value=16,
+                                                            step=1)
+                                    
+                                with gr.Row():     
+                                    Style_Guidance = gr.Slider(label="Style Guidance Scale",
+                                                    minimum=-1,
+                                                    maximum=3,
+                                                    value=1.2,
+                                                    step=0.01,
+                                                    )
+                                    mix_q_scale = gr.Slider(label='Query Mix Ratio',
+                                                            minimum=0,
+                                                            maximum=2,
+                                                            step=0.05,
+                                                            value=1.0,
+                                                            )
+                                cfg_scale_slider = gr.Slider(label="CFG Scale", value=2.5, minimum=0, maximum=20, info="Classifier-free guidance scale.")
+                                
+                                with gr.Row():     
+                                    save_intermediate = gr.Checkbox(label="save_intermediate", value=False)
+                                    de_bug = gr.Checkbox(value=False,label='DeBug')
+                            with gr.Tab('ToMe'):
+                                with gr.Row():
+                                    tome = gr.Checkbox(label="Token Merge", value=True)
+                                    
+                                    tome_ratio = gr.Slider(label='ratio: ',
+                                                            minimum=0,
+                                                            maximum=1,
+                                                            step=0.1,
+                                                            value=0.5)
+                                with gr.Row():
+                                    tome_sx = gr.Slider(label='sx:',
+                                                            minimum=0,
+                                                            maximum=64,
+                                                            step=2,
+                                                            value=2)
+                                    tome_sy = gr.Slider(label='sy:',
+                                                            minimum=0,
+                                                            maximum=64,
+                                                            step=2,
+                                                            value=2)
+                            
+                                
+                            with gr.Row():
+                                seed_textbox = gr.Textbox(label="Seed", value=-1)
+                                seed_button  = gr.Button(value="\U0001F3B2", elem_classes="toolbutton")
+                                seed_button.click(fn=lambda: random.randint(1, 1e16), inputs=[], outputs=[seed_textbox])
+        inputs = [
+            source_image, style_image,
+            num_steps,co_feat_step,strength,
+            start_ac_layer, end_ac_layer,
+            Style_Guidance,cfg_scale_slider,mix_q_scale,
+            Scheduler, save_intermediate, seed_textbox, de_bug,  
+            prompt_textbox, negative_prompt_textbox,
+            width_slider,height_slider,
+            tome_sx, tome_sy, tome_ratio, tome,
+        ]
+
+        generate_button.click(
+            fn=global_text.generate,
+            inputs=inputs,
+            outputs=[recons_style,recons_content,generate_image,results_image_gallery]
+        )
+
+        ex = gr.Examples(
+        [
+          ["./data/content/27032.jpg","./data/style/27.jpg",4,0.8,0.5,8427921159605868845],         
+          ["./data/content/29812.jpg","./data/style/47.jpg",4,0.5,0.65,8119359809263726691],
+        ],
+        [source_image, style_image, num_steps,strength, mix_q_scale, seed_textbox],
+        [
+            "Example 1",
+        ],)
+
+
+    return demo
+
+if __name__ == "__main__":
+    demo = ui()
+    demo.launch(server_name='172.18.32.44',show_error=True)

requirements.txt

@@ -0,0 +1,85 @@
+accelerate==0.33.0
+aiofiles==23.2.1
+annotated-types==0.7.0
+anyio==4.4.0
+Brotli @ file:///croot/brotli-split_1714483155106/work
+certifi @ file:///croot/certifi_1720453481653/work/certifi
+charset-normalizer @ file:///croot/charset-normalizer_1721748349566/work
+click==8.1.7
+contourpy==1.2.1
+cycler==0.12.1
+diffusers==0.30.0
+exceptiongroup==1.2.2
+fastapi==0.112.0
+ffmpy==0.4.0
+filelock @ file:///croot/filelock_1700591183607/work
+fonttools==4.53.1
+fsspec==2024.6.1
+gmpy2 @ file:///tmp/build/80754af9/gmpy2_1645438755360/work
+gradio==4.41.0
+gradio_client==1.3.0
+h11==0.14.0
+httpcore==1.0.5
+httpx==0.27.0
+huggingface-hub==0.24.5
+idna @ file:///croot/idna_1714398848350/work
+importlib_metadata==8.2.0
+importlib_resources==6.4.0
+Jinja2 @ file:///croot/jinja2_1716993405101/work
+kiwisolver==1.4.5
+markdown-it-py==3.0.0
+MarkupSafe @ file:///croot/markupsafe_1704205993651/work
+matplotlib==3.9.1.post1
+mdurl==0.1.2
+mkl-fft @ file:///croot/mkl_fft_1695058164594/work
+mkl-random @ file:///croot/mkl_random_1695059800811/work
+mkl-service==2.4.0
+mpmath @ file:///croot/mpmath_1690848262763/work
+mypy-extensions==1.0.0
+networkx @ file:///croot/networkx_1717597493534/work
+numpy @ file:///croot/numpy_and_numpy_base_1708638617955/work/dist/numpy-1.26.4-cp39-cp39-linux_x86_64.whl#sha256=6094eeedd869502faa0fd0a8c5ad3a70c5779be06ddd1feb7627e5c212fac420
+orjson==3.10.7
+packaging==24.1
+pandas==2.2.2
+peft==0.12.0
+pillow @ file:///croot/pillow_1721059439630/work
+psutil==6.0.0
+pydantic==2.8.2
+pydantic_core==2.20.1
+pydub==0.25.1
+Pygments==2.18.0
+pyparsing==3.1.2
+pyre-extensions==0.0.29
+PySocks @ file:///tmp/build/80754af9/pysocks_1605305812635/work
+python-dateutil==2.9.0.post0
+python-multipart==0.0.9
+pytz==2024.1
+PyYAML==6.0.2
+regex==2024.7.24
+requests @ file:///croot/requests_1721410876868/work
+rich==13.7.1
+ruff==0.5.7
+safetensors==0.4.4
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.1
+starlette==0.37.2
+sympy @ file:///croot/sympy_1701397643339/work
+tokenizers==0.19.1
+tomlkit==0.12.0
+torch==2.0.1
+torchaudio==2.0.2
+torchvision==0.15.2
+tqdm==4.66.5
+transformers==4.44.0
+triton==2.0.0
+typer==0.12.3
+typing-inspect==0.9.0
+typing_extensions @ file:///croot/typing_extensions_1715268824938/work
+tzdata==2024.1
+urllib3 @ file:///croot/urllib3_1718912636303/work
+uvicorn==0.30.5
+websockets==12.0
+xformers==0.0.20
+zipp==3.19.2

utils/attn_control.py

@@ -0,0 +1,129 @@
+import torch
+import torch.nn.functional as nnf
+import abc
+import math
+from torchvision.utils import save_image
+
+
+LOW_RESOURCE = False
+MAX_NUM_WORDS = 77
+device = "cuda" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+
+
+class AttentionControl(abc.ABC):
+
+    def step_callback(self, x_t):
+        return x_t
+
+    def between_steps(self):
+        return
+
+    @property
+    def start_att_layers(self):
+        return self.start_ac_layer #if LOW_RESOURCE else 0
+    @property
+    def end_att_layers(self):
+        return self.end_ac_layer
+    
+    
+    
+
+    @abc.abstractmethod
+    def forward(self, q, k, v, num_heads,attn):
+        raise NotImplementedError
+
+    def attn_forward(self, q, k, v, num_heads,attention_probs,attn):
+        if q.shape[0]//num_heads == 3:
+            h_s_re = self.forward(q, k, v, num_heads,attention_probs, attn)
+            
+        else:
+            uq,cq = q.chunk(2)
+            uk,ck = k.chunk(2)
+            uv,cv = v.chunk(2)
+            u_attn, c_attn = attention_probs.chunk(2)
+            
+            u_h_s_re = self.forward(uq, uk, uv, num_heads,u_attn, attn)
+
+            c_h_s_re = self.forward(cq, ck, cv, num_heads,c_attn, attn)
+            h_s_re = (u_h_s_re, c_h_s_re)
+        return h_s_re
+    
+    def __call__(self, q, k, v, num_heads,attention_probs,attn):
+
+        
+        if self.cur_att_layer >= self.start_att_layers and self.cur_att_layer < self.end_att_layers:
+            h_s_re = self.attn_forward(q, k, v, num_heads,attention_probs,attn)
+        else:
+            h_s_re=None
+        
+        
+        self.cur_att_layer += 1
+        
+        if self.cur_att_layer == self.num_att_layers // 2: #+ self.num_uncond_att_layers:
+            self.cur_att_layer = 0 #self.num_uncond_att_layers
+            self.cur_step += 1
+            self.between_steps()
+        return h_s_re
+
+    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
+
+def enhance_tensor(tensor: torch.Tensor, contrast_factor: float = 1.67) -> torch.Tensor:
+    """ Compute the attention map contrasting. """
+    mean_feat = tensor.mean(dim=-1, keepdims=True)
+    adjusted_tensor = (tensor - mean_feat) * contrast_factor + mean_feat
+    return adjusted_tensor
+
+class AttentionStyle(AttentionControl):
+
+    def __init__(self, 
+                 num_steps,
+                 start_ac_layer, end_ac_layer,
+                 style_guidance=0.3,
+                 mix_q_scale=1.0,
+                 de_bug=False,
+                 ):
+        super(AttentionStyle, self).__init__()
+
+
+        self.start_ac_layer = start_ac_layer
+        self.end_ac_layer = end_ac_layer
+        self.num_steps=num_steps
+        self.de_bug = de_bug
+        self.style_guidance = style_guidance
+        self.coef = None
+        self.mix_q_scale = mix_q_scale
+        
+    def forward(self, q, k, v, num_heads, attention_probs, attn):
+
+        
+        if self.de_bug:
+                import pdb; pdb.set_trace()
+
+        if self.mix_q_scale < 1.0:
+            q[num_heads*2:] = q[num_heads*2:] * self.mix_q_scale + (1 - self.mix_q_scale) * q[num_heads*1:num_heads*2]
+        b,n,d = k.shape
+        re_q = q[num_heads*2:] # b,n,d,
+        re_k = torch.cat([k[num_heads*1:num_heads*2],k[num_heads*0:num_heads*1]],dim=1) #b,2n,d
+        v_re = torch.cat([v[num_heads*1:num_heads*2],v[num_heads*0:num_heads*1]],dim=1) #b,2n,d
+        re_sim = torch.bmm(re_q, re_k.transpose(-1, -2)) * attn.scale
+        re_sim[:,:,n:] = re_sim[:,:,n:] * self.style_guidance
+        re_attention_map = re_sim.softmax(-1)
+        h_s_re = torch.bmm(re_attention_map, v_re)
+
+
+
+        return h_s_re
+        
+
+    
+
+

utils/merge.py

@@ -0,0 +1,124 @@
+import torch
+from typing import Tuple, Callable
+
+
+def do_nothing(x: torch.Tensor, mode:str=None):
+    return x
+
+
+def mps_gather_workaround(input, dim, index):
+    if input.shape[-1] == 1:
+        return torch.gather(
+            input.unsqueeze(-1),
+            dim - 1 if dim < 0 else dim,
+            index.unsqueeze(-1)
+        ).squeeze(-1)
+    else:
+        return torch.gather(input, dim, index)
+
+
+def bipartite_soft_matching_random2d(metric: torch.Tensor,
+                                     w: int, h: int, sx: int, sy: int, r: int,
+                                     no_rand: bool = False,
+                                     generator: torch.Generator = None) -> Tuple[Callable, Callable]:
+    """
+    Partitions the tokens into src and dst and merges r tokens from src to dst.
+    Dst tokens are partitioned by choosing one randomy in each (sx, sy) region.
+
+    Args:
+     - metric [B, N, C]: metric to use for similarity
+     - w: image width in tokens
+     - h: image height in tokens
+     - sx: stride in the x dimension for dst, must divide w
+     - sy: stride in the y dimension for dst, must divide h
+     - r: number of tokens to remove (by merging)
+     - no_rand: if true, disable randomness (use top left corner only)
+     - rand_seed: if no_rand is false, and if not None, sets random seed.
+    """
+    B, N, _ = metric.shape
+
+    if r <= 0:
+        return do_nothing, do_nothing
+
+    gather = mps_gather_workaround if metric.device.type == "mps" else torch.gather
+    
+    with torch.no_grad():
+        hsy, wsx = h // sy, w // sx
+
+        # For each sy by sx kernel, randomly assign one token to be dst and the rest src
+        if no_rand:
+            rand_idx = torch.zeros(hsy, wsx, 1, device=metric.device, dtype=torch.int64)
+        else:
+            rand_idx = torch.randint(sy*sx, size=(hsy, wsx, 1), device=generator.device, generator=generator).to(metric.device)
+        
+        # The image might not divide sx and sy, so we need to work on a view of the top left if the idx buffer instead
+        idx_buffer_view = torch.zeros(hsy, wsx, sy*sx, device=metric.device, dtype=torch.int64)
+        idx_buffer_view.scatter_(dim=2, index=rand_idx, src=-torch.ones_like(rand_idx, dtype=rand_idx.dtype))
+        idx_buffer_view = idx_buffer_view.view(hsy, wsx, sy, sx).transpose(1, 2).reshape(hsy * sy, wsx * sx)
+
+        # Image is not divisible by sx or sy so we need to move it into a new buffer
+        if (hsy * sy) < h or (wsx * sx) < w:
+            idx_buffer = torch.zeros(h, w, device=metric.device, dtype=torch.int64)
+            idx_buffer[:(hsy * sy), :(wsx * sx)] = idx_buffer_view
+        else:
+            idx_buffer = idx_buffer_view
+
+        # We set dst tokens to be -1 and src to be 0, so an argsort gives us dst|src indices
+        rand_idx = idx_buffer.reshape(1, -1, 1).argsort(dim=1)
+
+        # We're finished with these
+        del idx_buffer, idx_buffer_view
+
+        # rand_idx is currently dst|src, so split them
+        num_dst = hsy * wsx
+        a_idx = rand_idx[:, num_dst:, :] # src
+        b_idx = rand_idx[:, :num_dst, :] # dst
+
+        def split(x):
+            C = x.shape[-1]
+            src = gather(x, dim=1, index=a_idx.expand(B, N - num_dst, C))
+            dst = gather(x, dim=1, index=b_idx.expand(B, num_dst, C))
+            return src, dst
+
+        # Cosine similarity between A and B
+        metric = metric / metric.norm(dim=-1, keepdim=True)
+        a, b = split(metric)
+        scores = a @ b.transpose(-1, -2)
+
+        # Can't reduce more than the # tokens in src
+        r = min(a.shape[1], r)
+
+        # Find the most similar greedily
+        node_max, node_idx = scores.max(dim=-1)
+        edge_idx = node_max.argsort(dim=-1, descending=True)[..., None]
+
+        unm_idx = edge_idx[..., r:, :]  # Unmerged Tokens
+        src_idx = edge_idx[..., :r, :]  # Merged Tokens
+        dst_idx = gather(node_idx[..., None], dim=-2, index=src_idx)
+
+    def merge(x: torch.Tensor, mode="mean") -> torch.Tensor:
+        src, dst = split(x)
+        n, t1, c = src.shape
+        
+        unm = gather(src, dim=-2, index=unm_idx.expand(n, t1 - r, c))
+        src = gather(src, dim=-2, index=src_idx.expand(n, r, c))
+        dst = dst.scatter_reduce(-2, dst_idx.expand(n, r, c), src, reduce=mode)
+
+        return torch.cat([unm, dst], dim=1)
+
+    def unmerge(x: torch.Tensor) -> torch.Tensor:
+        unm_len = unm_idx.shape[1]
+        unm, dst = x[..., :unm_len, :], x[..., unm_len:, :]
+        _, _, c = unm.shape
+
+        src = gather(dst, dim=-2, index=dst_idx.expand(B, r, c))
+
+        # Combine back to the original shape
+        out = torch.zeros(B, N, c, device=x.device, dtype=x.dtype)
+        out.scatter_(dim=-2, index=b_idx.expand(B, num_dst, c), src=dst)
+        out.scatter_(dim=-2, index=gather(a_idx.expand(B, a_idx.shape[1], 1), dim=1, index=unm_idx).expand(B, unm_len, c), src=unm)
+        out.scatter_(dim=-2, index=gather(a_idx.expand(B, a_idx.shape[1], 1), dim=1, index=src_idx).expand(B, r, c), src=src)
+
+        return out
+
+    return merge, unmerge

utils/pipeline.py

@@ -0,0 +1,434 @@
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+import os
+import numpy as np
+import PIL
+import torch
+from packaging import version
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+
+from diffusers.configuration_utils import FrozenDict
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import LoraLoaderMixin, TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL, UNet2DConditionModel
+from diffusers.schedulers import LCMScheduler
+from diffusers.utils import PIL_INTERPOLATION, deprecate, logging
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+
+logger = logging.get_logger(__name__)  
+
+
+
+class ZePoPipeline(DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin):
+    model_cpu_offload_seq = "text_encoder->unet->vae"
+    _optional_components = ["safety_checker", "feature_extractor"]
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: LCMScheduler,
+        safety_checker: StableDiffusionSafetyChecker,
+        feature_extractor: CLIPImageProcessor,
+        requires_safety_checker: bool = True,
+    ):
+        super().__init__()
+
+        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
+                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
+                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
+                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
+                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
+                " file"
+            )
+            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(scheduler.config)
+            new_config["steps_offset"] = 1
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        if safety_checker is None and requires_safety_checker:
+            logger.warning(
+                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
+                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
+                " results in services or applications open to the public. Both the diffusers team and Hugging Face"
+                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
+                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
+                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
+            )
+
+        if safety_checker is not None and feature_extractor is None:
+            raise ValueError(
+                "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
+                " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
+            )
+        is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(
+            version.parse(unet.config._diffusers_version).base_version
+        ) < version.parse("0.9.0.dev0")
+        is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64
+        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:
+            deprecation_message = (
+                "The configuration file of the unet has set the default `sample_size` to smaller than"
+                " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the"
+                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"
+                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"
+                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"
+                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"
+                " in the config might lead to incorrect results in future versions. If you have downloaded this"
+                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"
+                " the `unet/config.json` file"
+            )
+            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(unet.config)
+            new_config["sample_size"] = 64
+            unet._internal_dict = FrozenDict(new_config)
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            unet=unet,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+            feature_extractor=feature_extractor,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.register_to_config(requires_safety_checker=requires_safety_checker)
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.check_inputs
+    def check_inputs(
+        self, prompt, strength, callback_steps, negative_prompt=None, prompt_embeds=None, negative_prompt_embeds=None
+    ):
+        if strength < 0 or strength > 1:
+            raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
+
+        if (callback_steps is None) or (
+            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
+        ):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}."
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+    def prepare_extra_step_kwargs(self, generator, eta):
+
+
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+   
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents
+    def decode_latents(self, latents):
+        deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"
+        deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)
+
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents, return_dict=False)[0]
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        return image
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps
+    def get_timesteps(self, num_inference_steps, strength, device):
+        # get the original timestep using init_timestep
+        init_timestep = min( int(num_inference_steps * strength), num_inference_steps)
+        init_timestep = max(init_timestep, 1)
+        t_start = max(num_inference_steps - init_timestep, 0)
+        
+        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
+
+        return timesteps, num_inference_steps - t_start
+
+    def prepare_latents(self, image, timestep, device,dtype, denoise_model, generator=None):
+        image = image.to(device=device,dtype=dtype)
+
+        batch_size = image.shape[0]
+
+        if image.shape[1] == 4:
+            init_latents = image
+
+        else:
+            if isinstance(generator, list) and len(generator) != batch_size:
+                raise ValueError(
+                    f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                    f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+                )
+
+            if isinstance(generator, list):
+                init_latents = [
+                    self.vae.encode(image[i : i + 1]).latent_dist.sample(generator[i]) for i in range(batch_size)
+                ]
+                init_latents = torch.cat(init_latents, dim=0)
+            else:
+                init_latents = self.vae.encode(image).latent_dist.sample(generator)
+
+            init_latents = self.vae.config.scaling_factor * init_latents
+
+        
+
+        # add noise to latents using the timestep
+        shape = init_latents.shape
+        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+
+        # get latents
+        clean_latents = init_latents
+        if denoise_model:
+            init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
+            latents = init_latents
+        else:
+            latents = noise
+
+        return latents, clean_latents
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        negative_prompt: Union[str, List[str]]=None,
+        image: PipelineImageInput = None,
+        style: PipelineImageInput = None,
+        strength: float = 0.5,
+        num_inference_steps: Optional[int] = 50,
+        original_inference_steps: Optional[int]  = 50,
+        guidance_scale: Optional[float] = 7.5,
+        eta: Optional[float] = 1.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        denoise_model: Optional[bool] = True,
+        fix_step_index = 0,
+        target_start_step = -1,
+        save_intermediate = False,
+        de_bug=False,
+
+    ):
+        # 1. Check inputs
+        self.check_inputs(prompt, strength, callback_steps)
+        num_inference_steps = int(num_inference_steps * (1/strength))
+        print(f'num_inference_steps {num_inference_steps} is multiple by {int(1/strength)}.')
+        # 2. Define call parameters
+        batch_size = len(prompt) 
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        
+        # text embeddings
+        text_input = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=77,
+            return_tensors="pt"
+        )
+        dtype=self.unet.dtype
+        prompt_embeds = self.text_encoder(text_input.input_ids.to(device))[0]
+        prompt_embeds=prompt_embeds.to(dtype=dtype, device=device)
+        #print("input text embeddings :", prompt_embeds.shape)
+        
+        if guidance_scale > 1.:
+            max_length = text_input.input_ids.shape[-1]
+            if negative_prompt:
+                uc_text = negative_prompt
+            else:
+                uc_text = ""
+            # uc_text = "ugly, tiling, poorly drawn hands, poorly drawn feet, body out of frame, cut off, low contrast, underexposed, distorted face"
+            unconditional_input = self.tokenizer(
+                [uc_text] * batch_size,
+                padding="max_length",
+                max_length=77,
+                return_tensors="pt"
+            )
+            # unconditional_input.input_ids = unconditional_input.input_ids[:, 1:]
+            unconditional_embeddings = self.text_encoder(unconditional_input.input_ids.to(device))[0]
+            unconditional_embeddings=unconditional_embeddings.to(dtype=dtype, device=device)
+            prompt_embeds = torch.cat([unconditional_embeddings, prompt_embeds], dim=0)
+
+        #print("prompt embeds shape: ", prompt_embeds.shape)
+        
+        
+        # 4. Preprocess image
+        image = self.image_processor.preprocess(image)
+        style = self.image_processor.preprocess(style)
+
+        # 5. Prepare timesteps
+        if isinstance(self.scheduler, LCMScheduler):
+            self.scheduler.set_timesteps(
+            num_inference_steps=num_inference_steps, 
+            device=device, 
+            original_inference_steps=original_inference_steps)
+        else:
+            self.scheduler.set_timesteps(
+            num_inference_steps=num_inference_steps, 
+            device=device,)
+        print(f"num_inference_steps is {self.scheduler.timesteps}")
+        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
+        
+        #print(f"All timesteps is : {timesteps}")
+        latent_timestep = torch.tensor([fix_step_index], device=device)
+        
+        assert timesteps != []
+
+        print("The time-steps are: ", timesteps)
+        
+        # 6. Prepare latent variables
+        src_latents, src_clean_latents = self.prepare_latents(
+            image, latent_timestep, device,dtype, denoise_model, generator
+        )
+
+        
+        sty_latents, sty_clean_latents = self.prepare_latents(
+            style, latent_timestep, device,dtype, denoise_model, generator
+        )
+
+        
+        mutual_latents, _ = self.prepare_latents(
+            image, timesteps[:1], device, dtype, denoise_model, generator
+        )
+
+        # mutual_latents = src_latents
+        #latents = torch.cat([sty_t_latents, src_t_latents], dim=0)
+        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+        generator = extra_step_kwargs.pop("generator", None)
+
+        # 8. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if de_bug:
+
+                    import pdb; pdb.set_trace()
+                
+                model_input = torch.cat(
+                        [
+                            sty_latents,
+                            src_latents,
+                            mutual_latents
+
+                        ],
+                        dim=0,
+                    )
+                # predict the noise residual
+                if do_classifier_free_guidance:
+                    concat_latent_model_input = torch.cat([model_input] * 2)
+                    concat_prompt_embeds = prompt_embeds
+                    #raise NotImplementedError("Classifier free guidance is not yet supported")
+                else:
+                    concat_latent_model_input = model_input
+                    concat_prompt_embeds = prompt_embeds
+                    assert len(concat_prompt_embeds) == len(concat_latent_model_input)
+
+                timestep = torch.cat([latent_timestep] * (batch_size-1)+[t[None]], dim=0)
+
+                if do_classifier_free_guidance:
+                    timestep = torch.cat([timestep] * 2)
+
+                
+                concat_noise_pred = self.unet(
+                    concat_latent_model_input,
+                    timestep,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    encoder_hidden_states=concat_prompt_embeds,
+                ).sample
+                # perform guidance
+                if do_classifier_free_guidance:
+
+                    (
+                        noise_pred, 
+                        noise_pred_uncond, 
+                    ) = concat_noise_pred.chunk(2, dim=0)
+                    
+
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
+                    
+                else:
+                    noise_pred = concat_noise_pred
+                
+                (style_noise_pred, source_noise_pred, mutual_noise_pred) = noise_pred.chunk(3, dim=0)
+
+                noise = torch.randn_like(
+                    source_noise_pred
+                )
+
+
+                if isinstance(self.scheduler, LCMScheduler):
+                    mutual_latents, pred_x0_mutual = self.scheduler.step(mutual_noise_pred, t, mutual_latents, return_dict=False)
+                else:
+                    ddim_out = self.scheduler.step(mutual_noise_pred, t, mutual_latents)
+                    mutual_latents, pred_x0_mutual = ddim_out.prev_sample, ddim_out.pred_original_sample
+                
+                
+                pred_x0 = torch.cat([sty_clean_latents,src_clean_latents,pred_x0_mutual ], dim=0)
+      
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                
+        model_input = torch.cat([sty_latents,src_latents,mutual_latents],dim=0,)
+        
+        # 9. Post-processing
+        if not output_type == "latent":
+            image = self.vae.decode(pred_x0 / self.vae.config.scaling_factor, return_dict=False)[0]
+            has_nsfw_concept = None
+        else:
+            image = pred_x0
+            has_nsfw_concept = None
+
+        if has_nsfw_concept is None:
+            do_denormalize = [True] * image.shape[0]
+        else:
+            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
+
+        image = self.image_processor.postprocess(image, output_type='np', do_denormalize=do_denormalize)
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)

utils/ptp_utils.py

@@ -0,0 +1,225 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import torch
+from typing import Optional, Union, Tuple, Dict
+from PIL import Image
+from . import merge
+from .utils import isinstance_str, init_generator
+
+def save_images(images,dest, num_rows=1, offset_ratio=0.02):
+    if type(images) is list:
+        num_empty = len(images) % num_rows
+    elif images.ndim == 4:
+        num_empty = images.shape[0] % num_rows
+    else:
+        images = [images]
+        num_empty = 0
+
+    pil_img = Image.fromarray(images[-1])
+    pil_img.save(dest)
+    # display(pil_img)
+
+
+def save_image(images,dest, num_rows=1, offset_ratio=0.02):
+    print(images.shape)
+    pil_img = Image.fromarray(images[0])
+    pil_img.save(dest)
+
+def register_attention_control(model, controller, tome, ratio, sx, sy, de_bug):
+    class AttnProcessor():
+        def __init__(self,place_in_unet,de_bug):
+            self.place_in_unet = place_in_unet
+            self.de_bug = de_bug
+        def __call__(self,
+            attn,
+            hidden_states,
+            encoder_hidden_states=None,
+            attention_mask=None,
+            temb=None,
+            scale=1.0,):
+            # The `Attention` class can call different attention processors / attention functions
+    
+            residual = hidden_states
+
+            if attn.spatial_norm is not None:
+                hidden_states = attn.spatial_norm(hidden_states, temb)
+
+            input_ndim = hidden_states.ndim
+
+            if input_ndim == 4:
+                batch_size, channel, height, width = hidden_states.shape
+                hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+            h = attn.heads
+            is_cross = encoder_hidden_states is not None
+            if encoder_hidden_states is None:
+                encoder_hidden_states = hidden_states
+            elif attn.norm_cross:
+                encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+            batch_size, sequence_length, _ = (
+                hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+            )
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+
+            
+
+            q = attn.to_q(hidden_states)
+            k = attn.to_k(encoder_hidden_states)
+            v = attn.to_v(encoder_hidden_states)
+            q = attn.head_to_batch_dim(q)
+            k = attn.head_to_batch_dim(k)
+            v = attn.head_to_batch_dim(v)
+
+            
+            # print('unmerge:', q.shape)
+                #pass
+            attention_probs = attn.get_attention_scores(q, k, attention_mask) # bh,n,n
+
+            #
+            
+            if is_cross:
+                pass
+                #attention_probs  = controller(attention_probs , is_cross, self.place_in_unet)
+            x = hidden_states
+            
+            hidden_states = torch.bmm(attention_probs, v)
+            
+            if not is_cross:
+
+                if tome:
+                    
+                    r = int(x.shape[1] * ratio)
+                    H = W = int(np.sqrt(x.shape[1]))
+                    generator = init_generator(x.device)
+                    m, u = merge.bipartite_soft_matching_random2d(x, W, H, sx, sy, r, 
+                                                        no_rand=False, generator=generator)
+                    x = m(x)
+                    m_k = attn.to_k(x)
+                    m_v = attn.to_v(x)
+                    m_k = attn.head_to_batch_dim(m_k)
+                    m_v = attn.head_to_batch_dim(m_v)
+                    # print('merged:', m_q.shape)
+                    # m_k = k
+                    # m_v = v
+                    #m_k, m_v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (m_k, m_v))
+                else:
+                    m_k = k
+                    m_v = v
+                # if self.de_bug:
+                #     import pdb;pdb.set_trace()
+                h_s_re = controller(q, m_k, m_v, attn.heads, attention_probs, attn)
+                
+                if h_s_re != None and hidden_states.shape[0]//attn.heads == 3:
+                    
+                    hidden_states[2*attn.heads:]=h_s_re
+
+                if hidden_states.shape[0]//attn.heads != 3 and h_s_re != None:
+                    (u_h_s_re, c_h_s_re) = h_s_re
+                    if u_h_s_re != None:
+               
+                        hidden_states[2*attn.heads:3*attn.heads] = u_h_s_re
+                        hidden_states[5*attn.heads:] = c_h_s_re
+            hidden_states = attn.batch_to_head_dim(hidden_states)
+
+            # linear proj   
+            hidden_states = attn.to_out[0](hidden_states)
+            # dropout
+            hidden_states = attn.to_out[1](hidden_states)
+
+            if input_ndim == 4:
+                hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+            if attn.residual_connection:
+                hidden_states = hidden_states + residual
+
+            hidden_states = hidden_states / attn.rescale_output_factor
+
+            return hidden_states
+
+
+    def register_recr(net_, count, place_in_unet):
+        for idx, m in enumerate(net_.modules()):
+            # print(m.__class__.__name__)
+            if m.__class__.__name__ == "Attention":
+                count+=1
+                m.processor = AttnProcessor( place_in_unet, de_bug)
+        return count
+
+    cross_att_count = 0
+    sub_nets = model.unet.named_children()
+    for net in sub_nets:
+        if "down" in net[0]:
+            cross_att_count += register_recr(net[1], 0, "down")
+        elif "up" in net[0]:
+            cross_att_count += register_recr(net[1], 0, "up")
+        elif "mid" in net[0]:
+            cross_att_count += register_recr(net[1], 0, "mid")
+    controller.num_att_layers = cross_att_count
+    #print(f'this model have {cross_att_count} attn layer')
+    
+def get_word_inds(text: str, word_place: int, tokenizer):
+    split_text = text.split(" ")
+    if type(word_place) is str:
+        word_place = [i for i, word in enumerate(split_text) if word_place == word]
+    elif type(word_place) is int:
+        word_place = [word_place]
+    out = []
+    if len(word_place) > 0:
+        words_encode = [tokenizer.decode([item]).strip("#") for item in tokenizer.encode(text)][1:-1]
+        cur_len, ptr = 0, 0
+
+        for i in range(len(words_encode)):
+            cur_len += len(words_encode[i])
+            if ptr in word_place:
+                out.append(i + 1)
+            if cur_len >= len(split_text[ptr]):
+                ptr += 1
+                cur_len = 0
+    return np.array(out)
+
+
+def update_alpha_time_word(alpha, bounds: Union[float, Tuple[float, float]], prompt_ind: int, word_inds: Optional[torch.Tensor]=None):
+    if type(bounds) is float:
+        bounds = 0, bounds
+    start, end = int(bounds[0] * alpha.shape[0]), int(bounds[1] * alpha.shape[0])
+    if word_inds is None:
+        word_inds = torch.arange(alpha.shape[2])
+    alpha[: start, prompt_ind, word_inds] = 0
+    alpha[start: end, prompt_ind, word_inds] = 1
+    alpha[end:, prompt_ind, word_inds] = 0
+    return alpha
+
+
+def get_time_words_attention_alpha(prompts, num_steps, cross_replace_steps: Union[float, Tuple[float, float], Dict[str, Tuple[float, float]]],
+                                   tokenizer, max_num_words=77):
+    if type(cross_replace_steps) is not dict:
+        cross_replace_steps = {"default_": cross_replace_steps}
+    if "default_" not in cross_replace_steps:
+        cross_replace_steps["default_"] = (0., 1.)
+    alpha_time_words = torch.zeros(num_steps + 1, len(prompts) - 1, max_num_words)
+    for i in range(len(prompts) - 1):
+        alpha_time_words = update_alpha_time_word(alpha_time_words, cross_replace_steps["default_"],
+                                                  i)
+    for key, item in cross_replace_steps.items():
+        if key != "default_":
+             inds = [get_word_inds(prompts[i], key, tokenizer) for i in range(1, len(prompts))]
+             for i, ind in enumerate(inds):
+                 if len(ind) > 0:
+                    alpha_time_words = update_alpha_time_word(alpha_time_words, item, i, ind)
+    alpha_time_words = alpha_time_words.reshape(num_steps + 1, len(prompts) - 1, 1, 1, max_num_words) # time, batch, heads, pixels, words
+    return alpha_time_words

utils/utils.py

@@ -0,0 +1,32 @@
+import torch
+
+
+def isinstance_str(x: object, cls_name: str):
+    """
+    Checks whether x has any class *named* cls_name in its ancestry.
+    Doesn't require access to the class's implementation.
+    
+    Useful for patching!
+    """
+
+    for _cls in x.__class__.__mro__:
+        if _cls.__name__ == cls_name:
+            return True
+    
+    return False
+
+
+def init_generator(device: torch.device, fallback: torch.Generator=None):
+    """
+    Forks the current default random generator given device.
+    """
+    if device.type == "cpu":
+        return torch.Generator(device="cpu").set_state(torch.get_rng_state())
+    elif device.type == "cuda":
+        return torch.Generator(device=device).set_state(torch.cuda.get_rng_state())
+    else:
+        if fallback is None:
+            return init_generator(torch.device("cpu"))
+        else:
+            return fallback
+