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.

from matplotlib import pyplot as plt
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont
import cv2
from typing import Optional, Union, Tuple, List, Callable, Dict
from IPython.display import display
from tqdm import tqdm
import math
from diffusers.utils import randn_tensor
def text_under_image(image: np.ndarray, text: str, text_color: Tuple[int, int, int] = (0, 0, 0)):
    h, w, c = image.shape
    offset = int(h * .2)
    img = np.ones((h + offset, w, c), dtype=np.uint8) * 255
    font = cv2.FONT_HERSHEY_SIMPLEX
    # font = ImageFont.truetype("/usr/share/fonts/truetype/noto/NotoMono-Regular.ttf", font_size)
    img[:h] = image
    textsize = cv2.getTextSize(text, font, 1, 2)[0]
    text_x, text_y = (w - textsize[0]) // 2, h + offset - textsize[1] // 2
    cv2.putText(img, text, (text_x, text_y ), font, 1, text_color, 2)
    return img

def save_plot(mel, num_cols=2, save_path='mels.png', names='mel-spectrogram'):  ###[B, n_mels, time_frames] cpu numpy
    if mel.ndim == 2:
        if type(mel) is torch.Tensor:
            mel = mel.cpu().numpy()
        mel = np.expand_dims(mel, axis=0)
    if type(names) is str:
        names = [names] * mel.shape[0]
    assert mel.shape[0] == len(names)

    batch_size = mel.shape[0]
    # 计算行数（自动计算）
    num_rows = (batch_size + num_cols - 1) // num_cols

    # 创建一个大图形
    fig, axes = plt.subplots(max(2, num_rows), max(2, num_cols), figsize=(15, 8)) ###最小2行2列

    # 循环绘制每张子图
    for i in range(batch_size):
        row = i // num_cols  # 计算行索引
        col = i % num_cols   # 计算列索引
        
        ax = axes[row, col]  # 获取当前子图对象
        
        # 绘制Mel谱数据（在这里使用示例数据）
        imm = ax.imshow(mel[i], origin='lower', aspect='auto')
        colorbar = plt.colorbar(imm)
        # 添加标题
        ax.set_title(names[i], fontsize=14)

    # 去除空白的子图
    for i in range(batch_size, num_rows * num_cols):
        row = i // num_cols
        col = i % num_cols
        fig.delaxes(axes[row, col])

    # 调整子图的布局，确保它们适应大图的大小
    plt.tight_layout()

    # 显示图形
    plt.show()
    plt.savefig(save_path)
    return

def view_images(images, num_rows=1, offset_ratio=0.02, save_path=None): ###[B, H, W, C] cpu tensor
    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

    empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255
    images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty
    num_items = len(images)

    h, w, c = images[0].shape
    offset = int(h * offset_ratio)
    num_cols = num_items // num_rows
    image_ = np.ones((h * num_rows + offset * (num_rows - 1),
                      w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255
    for i in range(num_rows):
        for j in range(num_cols):
            image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images[
                i * num_cols + j]

    pil_img = Image.fromarray(image_)
    if save_path:
        pil_img.save(save_path)  # 保存图像到指定路径
        

def save_mel(tensor, savepath):
    plt.style.use('default')
    fig, ax = plt.subplots(figsize=(12, 3))
    im = ax.imshow(tensor, aspect="auto", origin="lower", interpolation='none')
    plt.colorbar(im, ax=ax)
    plt.tight_layout()
    fig.canvas.draw()
    plt.savefig(savepath)
    plt.close()
    return

def diffusion_step(model, controller, latents, context, boolean_prompt_mask, t, guidance_scale, low_resource=False, inference_scheduler=None):
    if low_resource:
        noise_pred_uncond = model.unet(latents, t, encoder_hidden_states=context[0])["sample"]
        noise_prediction_text = model.unet(latents, t, encoder_hidden_states=context[1])["sample"]
    else:
        # expand the latents if we are doing classifier free guidance
        latent_model_input = torch.cat([latents] * 2) 
        latent_model_input = inference_scheduler.scale_model_input(latent_model_input, t)

        noise_pred = self.unet(
            latent_model_input, t, encoder_hidden_states=prompt_embeds,
            encoder_attention_mask=boolean_prompt_mask
        ).sample

            # perform guidance
        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

            # compute the previous noisy sample x_t -> x_t-1
        latents = inference_scheduler.step(noise_pred, t, latents).prev_sample

            # call the callback, if provided
        if t == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % inference_scheduler.order == 0):
            progress_bar.update(1)
    noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond)
    latents = model.inference_scheduler.step(noise_pred, t, latents)["prev_sample"]
    latents = controller.step_callback(latents)
    return latents


def latent2image(vae, latents):
    latents = 1 / 0.18215 * latents
    image = vae.decode(latents)['sample']
    image = (image / 2 + 0.5).clamp(0, 1)
    image = image.cpu().permute(0, 2, 3, 1).numpy()
    image = (image * 255).astype(np.uint8)
    return image


def init_latent(latent, model, height, width, generator, batch_size, device, dtype):
    if latent is None:
        shape = (1, model.unet.in_channels, height // 4, width // 4)
        latent = randn_tensor(shape, generator=None, device=device, dtype=dtype)
        
    latents = latent.expand(batch_size,  model.unet.in_channels, height // 4, width // 4).to(device)
    return latent, latents



def prepare_latents(batch_size, inference_scheduler, num_channels_latents, dtype, device):
    shape = (batch_size, num_channels_latents, 256, 16)
    latents = randn_tensor(shape, generator=None, device=device, dtype=dtype)
    # scale the initial noise by the standard deviation required by the scheduler
    latents = latents * inference_scheduler.init_noise_sigma
    return latents

def encode_text_classifier_free(model, prompt, num_samples_per_prompt):
    device = model.text_encoder.device
    batch = model.tokenizer(
            prompt, max_length=model.tokenizer.model_max_length, padding=True, truncation=True, return_tensors="pt"
        )
    input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device)

    with torch.no_grad():
        prompt_embeds = model.text_encoder(
            input_ids=input_ids, attention_mask=attention_mask
        )[0]
                
    prompt_embeds = prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
    attention_mask = attention_mask.repeat_interleave(num_samples_per_prompt, 0)

        # get unconditional embeddings for classifier free guidance
    uncond_tokens = [""] * len(prompt)

    max_length = prompt_embeds.shape[1]
    uncond_batch = model.tokenizer(
        uncond_tokens, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt",
    )
    uncond_input_ids = uncond_batch.input_ids.to(device)
    uncond_attention_mask = uncond_batch.attention_mask.to(device)

    with torch.no_grad():
        negative_prompt_embeds = model.text_encoder(
            input_ids=uncond_input_ids, attention_mask=uncond_attention_mask
        )[0]
                
    negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
    uncond_attention_mask = uncond_attention_mask.repeat_interleave(num_samples_per_prompt, 0)

        # For classifier free guidance, we need to do two forward passes.
        # We concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes
    prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
    prompt_mask = torch.cat([uncond_attention_mask, attention_mask])
    boolean_prompt_mask = (prompt_mask == 1).to(device)

    return prompt_embeds, boolean_prompt_mask

@torch.no_grad()
def text2mel_ldm_stable(
    tango,          ##给定tango
    prompt: List[str],
    controller,
    num_inference_steps: int = 200,
    guidance_scale: float = 3,
    generator: Optional[torch.Generator] = None,
    latent: Optional[torch.FloatTensor] = None,
    low_resource: bool = False,
):
    register_attention_control(tango.model, controller)
    height = 1024
    width = 64
    batch_size = len(prompt)
    device = tango.model.unet.device
    classifier_free_guidance = guidance_scale > 1.0
    batch_size = len(prompt) 

    if classifier_free_guidance:
        context, prompt_mask = encode_text_classifier_free(tango.model, prompt, num_samples_per_prompt=1)
    inference_scheduler = tango.scheduler
    inference_scheduler.set_timesteps(num_inference_steps, device=device)
    timesteps = inference_scheduler.timesteps

    num_channels_latents = tango.model.unet.config.in_channels
    latents = prepare_latents(batch_size, inference_scheduler, num_channels_latents, context.dtype, device)


    # set timesteps
    tango.scheduler.set_timesteps(num_inference_steps)
    timesteps = tango.scheduler.timesteps
    num_warmup_steps = len(timesteps) - num_inference_steps * tango.scheduler.order  ##1000-200*1
    progress_bar = tqdm(range(num_inference_steps), disable=False)
    for i, t in enumerate(timesteps):
        # expand the latents if we are doing classifier free guidance
        latent_model_input = torch.cat([latents] * 2) if classifier_free_guidance else latents
        latent_model_input = inference_scheduler.scale_model_input(latent_model_input, t)

        noise_pred, _, _, _ = tango.model.unet(
            latent_model_input, t, encoder_hidden_states=context,
            encoder_attention_mask=prompt_mask
        )
        noise_pred = noise_pred.sample

            # perform guidance
        if classifier_free_guidance:
            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

            # compute the previous noisy sample x_t -> x_t-1
        latents = inference_scheduler.step(noise_pred, t, latents).prev_sample    
        if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % tango.scheduler.order == 0):
            progress_bar.update(1)
        latents = controller.step_callback(latents)
    mel = tango.vae.decode_first_stage(latents)
    wave = tango.vae.decode_to_waveform(mel)
    return mel, latent, wave[0]



def register_attention_control(model, controller):
    def ca_forward(self, place_in_unet):
        to_out = self.to_out
        if type(to_out) is torch.nn.modules.container.ModuleList:
            to_out = self.to_out[0]
        else:
            to_out = self.to_out
        
        def forward(hidden_states, encoder_hidden_states=None, attention_mask=None,temb=None,):
            is_cross = encoder_hidden_states is not None
            
            batch_size, sequence_length, _ = (
                hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
            )
            attention_mask = self.prepare_attention_mask(attention_mask, sequence_length, batch_size)
            query = self.to_q(hidden_states)

            if encoder_hidden_states is None:
                encoder_hidden_states = hidden_states
            elif self.cross_attention_norm:
                encoder_hidden_states = self.norm_cross(encoder_hidden_states)

            key = self.to_k(encoder_hidden_states)
            value = self.to_v(encoder_hidden_states)

            query = self.head_to_batch_dim(query)
            key = self.head_to_batch_dim(key)
            value = self.head_to_batch_dim(value)

        
            attention_probs = self.get_attention_scores(query, key, attention_mask)
            attention_probs = controller(attention_probs, is_cross, place_in_unet)

            hidden_states = torch.bmm(attention_probs, value)
            hidden_states = self.batch_to_head_dim(hidden_states)

            # linear proj
            hidden_states = self.to_out[0](hidden_states)
            # dropout
            hidden_states = self.to_out[1](hidden_states)
            return hidden_states, attention_probs
        return forward

    class DummyController:

        def __call__(self, *args):
            return args[0]

        def __init__(self):
            self.num_att_layers = 0

    if controller is None:
        controller = DummyController()

    def register_recr(net_, count, place_in_unet): ###递归登记attion模块
        if net_.__class__.__name__ == 'Attention':
            net_.forward = ca_forward(net_, place_in_unet)
            return count + 1
        elif hasattr(net_, 'children'):
            for net__ in net_.children():
                count = register_recr(net__, count, place_in_unet)
        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

    
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, 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)
    return alpha_time_words