utils/ptp_utils.py

# 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