models.py

import random
import numpy as np
from tqdm import tqdm
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from einops import repeat
import time
from tools.torch_tools import wav_to_fbank, sinusoidal_positional_embedding

from audioldm.audio.stft import TacotronSTFT
from audioldm.variational_autoencoder import AutoencoderKL
from audioldm.utils import default_audioldm_config, get_metadata

from transformers import CLIPTokenizer, AutoTokenizer, T5Tokenizer
from transformers import CLIPTextModel, T5EncoderModel, AutoModel
from transformers import CLIPVisionModelWithProjection, CLIPTextModelWithProjection
from transformers import CLIPProcessor, CLIPModel

import diffusers
from diffusers.utils.torch_utils import randn_tensor
from diffusers import DDPMScheduler, UNet2DConditionModel
from diffusers import AutoencoderKL as DiffuserAutoencoderKL
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize, InterpolationMode, RandomResizedCrop
from diffusers import AudioLDMPipeline

def build_pretrained_models(name):
    checkpoint = torch.load(name, map_location="cpu")
    scale_factor = checkpoint["state_dict"]["scale_factor"].item()

    vae_state_dict = {k[18:]: v for k, v in checkpoint["state_dict"].items() if "first_stage_model." in k}

    config = default_audioldm_config(name)
    vae_config = config["model"]["params"]["first_stage_config"]["params"]
    vae_config["scale_factor"] = scale_factor

    vae = AutoencoderKL(**vae_config)
    vae.load_state_dict(vae_state_dict)

    fn_STFT = TacotronSTFT(
        config["preprocessing"]["stft"]["filter_length"],
        config["preprocessing"]["stft"]["hop_length"],
        config["preprocessing"]["stft"]["win_length"],
        config["preprocessing"]["mel"]["n_mel_channels"],
        config["preprocessing"]["audio"]["sampling_rate"],
        config["preprocessing"]["mel"]["mel_fmin"],
        config["preprocessing"]["mel"]["mel_fmax"],
    )

    vae.eval()
    fn_STFT.eval()
    return vae, fn_STFT


class EffNetb3(nn.Module):
    def __init__(self, pretrained_model_path, embedding_dim=1024, pretrained=True):
        super(EffNetb3, self).__init__()
        self.model_name = 'effnetb3'
        self.pretrained = pretrained
        # Create model
        # self.effnet = torch.hub.load('rwightman/gen-efficientnet-pytorch', 'efficientnet_b3', pretrained=self.pretrained)
        # torch.save(self.effnet, 'model.pth')
        self.effnet = torch.hub.load(pretrained_model_path, 'efficientnet_b3', trust_repo=True, source='local')
        #self.effnet.conv_stem = nn.Conv2d(1, 40, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
        self.embedder = nn.Conv2d(384, embedding_dim, kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        #out = self.effnet(x)
        out = self.effnet.conv_stem(x)
        out = self.effnet.bn1(out)
        out = self.effnet.act1(out)
        for i in range(len(self.effnet.blocks)):
            out = self.effnet.blocks[i](out)
        out = self.embedder(out)
        return out


class EffNetb3_last_layer(nn.Module):
    def __init__(self, pretrained_model_path, embedding_dim=1024, pretrained=True):
        super(EffNetb3_last_layer, self).__init__()
        self.model_name = 'effnetb3'
        self.pretrained = pretrained
        self.effnet = torch.hub.load(pretrained_model_path, 'efficientnet_b3', trust_repo=True, source='local')
        self.effnet.classifier = nn.Linear(1536, embedding_dim)

    def forward(self, x):
        out = self.effnet(x)
        return out.unsqueeze(-1)


class Clip4Video(nn.Module):
    def __init__(self, model, embedding_dim=1024, pretrained=True, pe=False):
        super(Clip4Video, self).__init__()
        self.pretrained = pretrained
        self.clip_vision = CLIPVisionModelWithProjection.from_pretrained(model)
        self.clip_text = CLIPTextModelWithProjection.from_pretrained(model)
        self.tokenizer = AutoTokenizer.from_pretrained(model)

        input_dim = 512 if "clip-vit-base" in model else 768
        self.linear_layer = nn.Linear(input_dim, embedding_dim)
        self.pe = sinusoidal_positional_embedding(30, input_dim) if pe else None
        print("*****PE*****") if pe else print("*****W/O PE*****")

    def forward(self, text=None, image=None, video=None):
        assert text is not None or image is not None or video is not None, "At least one of text, image or video should be provided"
        if text is not None and video is None:
            inputs = self.tokenizer([text], padding=True, truncation=True, return_tensors="pt", max_length=77).to(self.clip_text.device)
            out = self.clip_text(**inputs)
            out = out.text_embeds.repeat(20, 1)
        elif video is not None and text is None:
            out = self.clip_vision(video.to(self.clip_vision.device))          # input video x: t * 3 * w * h
            out = out.image_embeds        # t * 512
            if self.pe is not None:
                out = out + self.pe[:out.shape[0], :].to(self.clip_vision.device)
            # out['last_hidden_state'].shape # t * 50 * 768
            # out['image_embeds'].shape      # t * 512
        elif text is not None and video is not None:
            text_inputs = self.tokenizer([text], padding=True, truncation=True, return_tensors="pt", max_length=77).to(self.clip_text.device)
            video_out = self.clip_vision(video.to(self.clip_vision.device))
            video_out = video_out.image_embeds
            text_out = self.clip_text(**text_inputs)
            text_out = text_out.text_embeds.repeat(video_out.shape[0], 1)
            # out = text_out + video_out
            # concat 
            out = torch.cat([text_out, video_out], dim=0)
        out = self.linear_layer(out)     # t * 1024
        return out


class AudioDiffusion(nn.Module):
    def __init__(
        self,
        fea_encoder_name,
        scheduler_name,
        unet_model_name=None,
        unet_model_config_path=None,
        snr_gamma=None,
        freeze_text_encoder=True,
        uncondition=False,
        img_pretrained_model_path=None,
        task=None,
        embedding_dim=1024,
        pe=False
    ):
        super().__init__()

        assert unet_model_name is not None or unet_model_config_path is not None, "Either UNet pretrain model name or a config file path is required"

        self.fea_encoder_name = fea_encoder_name
        self.scheduler_name = scheduler_name
        self.unet_model_name = unet_model_name
        self.unet_model_config_path = unet_model_config_path
        self.snr_gamma = snr_gamma
        self.freeze_text_encoder = freeze_text_encoder
        self.uncondition = uncondition
        self.task = task
        self.pe = pe

        # https://huggingface.co./docs/diffusers/v0.14.0/en/api/schedulers/overview
        self.noise_scheduler = DDPMScheduler.from_pretrained(self.scheduler_name, subfolder="scheduler")
        self.inference_scheduler = DDPMScheduler.from_pretrained(self.scheduler_name, subfolder="scheduler")

        if unet_model_config_path:
            unet_config = UNet2DConditionModel.load_config(unet_model_config_path)
            print("unet_config", unet_config)
            self.unet = UNet2DConditionModel.from_config(unet_config, subfolder="unet")
            self.set_from = "random"
            print("UNet initialized randomly.")
        else:
            self.unet = UNet2DConditionModel.from_pretrained(unet_model_name, subfolder="unet")
            self.set_from = "pre-trained"
            self.group_in = nn.Sequential(nn.Linear(8, 512), nn.Linear(512, 4))
            self.group_out = nn.Sequential(nn.Linear(4, 512), nn.Linear(512, 8))
            print("UNet initialized from stable diffusion checkpoint.")

        if self.task == "text2audio":
            if "stable-diffusion" in self.fea_encoder_name:
                self.tokenizer = CLIPTokenizer.from_pretrained(self.fea_encoder_name, subfolder="tokenizer")
                self.text_encoder = CLIPTextModel.from_pretrained(self.fea_encoder_name, subfolder="text_encoder")
            elif "t5" in self.fea_encoder_name and "Chinese" not in self.fea_encoder_name:
                self.tokenizer = AutoTokenizer.from_pretrained(self.fea_encoder_name)
                self.text_encoder = T5EncoderModel.from_pretrained(self.fea_encoder_name)
            elif "Chinese" in self.fea_encoder_name:
                self.tokenizer = T5Tokenizer.from_pretrained(self.fea_encoder_name)
                self.text_encoder = T5EncoderModel.from_pretrained(self.fea_encoder_name)
            elif "clap" in self.fea_encoder_name:
                self.tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
                self.CLAP_model = laion_clap.CLAP_Module(enable_fusion=False)
                self.CLAP_model.load_ckpt(self.fea_encoder_name)
            elif "clip-vit" in self.fea_encoder_name:
                # self.CLIP_model = CLIPModel.from_pretrained(self.fea_encoder_name)
                # self.CLIP_processor = CLIPProcessor.from_pretrained(self.fea_encoder_name)
                self.CLIP_model = CLIPTextModelWithProjection.from_pretrained(self.fea_encoder_name)
                self.tokenizer = AutoTokenizer.from_pretrained(self.fea_encoder_name)
                if "base" in self.fea_encoder_name:
                    self.linear_layer = nn.Linear(512, embedding_dim)
                else:
                    self.linear_layer = nn.Linear(768, embedding_dim)
            else:
                self.tokenizer = AutoTokenizer.from_pretrained(self.fea_encoder_name)
                self.text_encoder = AutoModel.from_pretrained(self.fea_encoder_name)
        elif self.task == "image2audio":
            if "clip-vit" in self.fea_encoder_name:
                self.CLIP_model = CLIPModel.from_pretrained(self.fea_encoder_name)
                self.CLIP_processor = CLIPProcessor.from_pretrained(self.fea_encoder_name)
                self.linear_layer = nn.Linear(512, embedding_dim)
            # self.img_fea_extractor = EffNetb3(img_pretrained_model_path)
            else:
                self.img_fea_extractor = EffNetb3_last_layer(img_pretrained_model_path)
        elif self.task == "video2audio":
            self.vid_fea_extractor = Clip4Video(model=self.fea_encoder_name, embedding_dim=embedding_dim, pe=pe)

    def compute_snr(self, timesteps):
        """
        Computes SNR as per https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849
        """
        alphas_cumprod = self.noise_scheduler.alphas_cumprod
        sqrt_alphas_cumprod = alphas_cumprod**0.5
        sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5

        # Expand the tensors.
        # Adapted from https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026
        sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[timesteps].float()
        while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape):
            sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None]
        alpha = sqrt_alphas_cumprod.expand(timesteps.shape)

        sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(device=timesteps.device)[timesteps].float()
        while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape):
            sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None]
        sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape)

        # Compute SNR.
        snr = (alpha / sigma) ** 2
        return snr

    def encode_text(self, prompt):
        device = self.text_encoder.device
        batch = self.tokenizer(
            prompt, max_length=self.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)

        if self.freeze_text_encoder:
            with torch.no_grad():
                encoder_hidden_states = self.text_encoder(
                    input_ids=input_ids, attention_mask=attention_mask
                )[0]
        else:
            encoder_hidden_states = self.text_encoder(
                input_ids=input_ids, attention_mask=attention_mask
            )[0]

        boolean_encoder_mask = (attention_mask == 1).to(device)
        return encoder_hidden_states, boolean_encoder_mask

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

        if self.freeze_text_encoder:
            with torch.no_grad():
                encoder_hidden_states = self.CLAP_model.model.get_text_embedding(prompt)
        else:
            encoder_hidden_states = self.CLAP_model.model.get_text_embedding(prompt)

        boolean_encoder_mask = (attention_mask == 1).to(device)
        return encoder_hidden_states, boolean_encoder_mask

    def encode_image(self, prompt, device):
        if "clip-vit" in self.fea_encoder_name:
            with torch.no_grad():
                inputs = self.CLIP_processor(text=["aaa"], images=prompt, return_tensors="pt", padding=True).to(device)
                encoder_hidden_states = self.CLIP_model(**inputs).image_embeds
            encoder_hidden_states = self.linear_layer(encoder_hidden_states)    # b * 1024
            encoder_hidden_states = encoder_hidden_states.unsqueeze(1).to(device)
        else:
            img_fea = self.img_fea_extractor(prompt) 
            encoder_hidden_states = img_fea.view(img_fea.shape[0], img_fea.shape[1], -1).permute(0, 2, 1)
        boolean_encoder_mask = torch.ones((encoder_hidden_states.shape[0], encoder_hidden_states.shape[1]), dtype=torch.bool)
        boolean_encoder_mask = boolean_encoder_mask.to(device)

        return encoder_hidden_states, boolean_encoder_mask
    
    def encode_video(self, video_batch, text=None, device=None):
        vid_feas = []
        for i, video in enumerate(video_batch):
            if text:
                vid_fea = self.vid_fea_extractor(video=video, text=text[i]) # t * fea_dim
            else:
                vid_fea = self.vid_fea_extractor(video=video)
            vid_feas.append(vid_fea)
        
        padding = 0
        size = max(v.size(0) for v in vid_feas)
        batch_size = len(vid_feas)
        embed_size = vid_feas[0].size(1)
        encoder_hidden_states = vid_feas[0].new(batch_size, size, embed_size).fill_(padding)
        boolean_encoder_mask = torch.ones((batch_size, size), dtype=torch.bool)

        def copy_tensor(src, dst):
            assert dst.numel() == src.numel()
            dst.copy_(src)

        for i, v in enumerate(vid_feas):
            copy_tensor(v, encoder_hidden_states[i][: len(v)])
            boolean_encoder_mask[i, len(v):] = False    
        return encoder_hidden_states.to(device), boolean_encoder_mask.to(device)

    def encode_text_CLIP(self, prompt, device):
        # tmp_image = np.ones((512, 512, 3))
        # with torch.no_grad():
        #     inputs = self.CLIP_processor(text=prompt, images=tmp_image, return_tensors="pt", padding=True, max_length=77, truncation=True).to(device)
        #     encoder_hidden_states = self.CLIP_model(**inputs).text_embeds   # b * 768
        text_inputs = self.tokenizer(prompt, padding=True, truncation=True, return_tensors="pt", max_length=77).to(device)
        encoder_hidden_states = self.CLIP_model(**text_inputs).text_embeds
        encoder_hidden_states = self.linear_layer(encoder_hidden_states)    # b * 1024
        encoder_hidden_states = encoder_hidden_states.unsqueeze(1).to(device)
        boolean_encoder_mask = torch.ones((encoder_hidden_states.shape[0], encoder_hidden_states.shape[1]), dtype=torch.bool)
        boolean_encoder_mask = boolean_encoder_mask.to(device)

        return encoder_hidden_states, boolean_encoder_mask

    def forward(self, latents, text=None, video=None, image=None, validation_mode=False, device=None):
        num_train_timesteps = self.noise_scheduler.num_train_timesteps
        self.noise_scheduler.set_timesteps(num_train_timesteps, device=device)
        # encoder_hidden_states.shape  [b, t, f]
        if self.task == "text2audio":
            if "clip-vit" in self.fea_encoder_name:
                encoder_hidden_states, boolean_encoder_mask = self.encode_text_CLIP(text, device)
            else:
                encoder_hidden_states, boolean_encoder_mask = self.encode_text(text)
            if self.uncondition:
                mask_indices = [k for k in range(len(text)) if random.random() < 0.1]
                # mask_indices = [k for k in range(len(prompt))]
                if len(mask_indices) > 0:
                    encoder_hidden_states[mask_indices] = 0
        elif self.task == "image2audio":
            encoder_hidden_states, boolean_encoder_mask = self.encode_image(image, device=device)
        elif self.task == "video2audio":
            encoder_hidden_states, boolean_encoder_mask = self.encode_video(video, text, device=device)
        
        bsz = latents.shape[0]
        if validation_mode:
            timesteps = (self.noise_scheduler.num_train_timesteps//2) * torch.ones((bsz,), dtype=torch.int64, device=device)
        else:
            # Sample a random timestep for each instance
            timesteps = torch.randint(0, self.noise_scheduler.num_train_timesteps, (bsz,), device=device)
        timesteps = timesteps.long()

        noise = torch.randn_like(latents)
        noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)

        # Get the target for loss depending on the prediction type
        if self.noise_scheduler.config.prediction_type == "epsilon":
            target = noise
        elif self.noise_scheduler.config.prediction_type == "v_prediction":
            target = self.noise_scheduler.get_velocity(latents, noise, timesteps)
        else:
            raise ValueError(f"Unknown prediction type {self.noise_scheduler.config.prediction_type}")

        if self.set_from == "random":
            model_pred = self.unet(
                noisy_latents, timesteps, encoder_hidden_states, 
                encoder_attention_mask=boolean_encoder_mask
            ).sample

        elif self.set_from == "pre-trained":
            compressed_latents = self.group_in(noisy_latents.permute(0, 2, 3, 1).contiguous()).permute(0, 3, 1, 2).contiguous()
            model_pred = self.unet(
                compressed_latents, timesteps, encoder_hidden_states, 
                encoder_attention_mask=boolean_encoder_mask
            ).sample
            model_pred = self.group_out(model_pred.permute(0, 2, 3, 1).contiguous()).permute(0, 3, 1, 2).contiguous()

        if self.snr_gamma is None:
            loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
        else:
            # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.
            # Adaptef from huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py
            snr = self.compute_snr(timesteps)
            mse_loss_weights = (
                torch.stack([snr, self.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr
            )
            loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
            loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights
            loss = loss.mean()

        return loss

    @torch.no_grad()
    def inference(self, inference_scheduler, text=None, video=None, image=None, num_steps=20, guidance_scale=3, num_samples_per_prompt=1, 
                  disable_progress=True, device=None):
        start = time.time()
        classifier_free_guidance = guidance_scale > 1.0

        #print("ldm time 0", time.time()-start, prompt)
        if self.task == "text2audio":
            batch_size = len(text) * num_samples_per_prompt

            if classifier_free_guidance:
                if "clip-vit" in self.fea_encoder_name:
                    encoder_hidden_states, boolean_encoder_mask = self.encode_text_clip_classifier_free(text, num_samples_per_prompt, device=device)
                else:
                    encoder_hidden_states, boolean_encoder_mask = self.encode_text_classifier_free(text, num_samples_per_prompt)
            else:
                encoder_hidden_states, boolean_encoder_mask = self.encode_text(text)
                encoder_hidden_states = encoder_hidden_states.repeat_interleave(num_samples_per_prompt, 0)
                boolean_encoder_mask = boolean_encoder_mask.repeat_interleave(num_samples_per_prompt, 0)
        elif self.task == "image2audio":
            if classifier_free_guidance:
                encoder_hidden_states, boolean_encoder_mask = self.encode_image_classifier_free(image, num_samples_per_prompt, device=device)
            else:
                encoder_hidden_states, boolean_encoder_mask = self.encode_image_no_grad(image, device=device)
                encoder_hidden_states = encoder_hidden_states.repeat_interleave(num_samples_per_prompt, 0)
                boolean_encoder_mask = boolean_encoder_mask.repeat_interleave(num_samples_per_prompt, 0)
        elif self.task == "video2audio":
            batch_size = len(video) * num_samples_per_prompt
            encoder_hidden_states, boolean_encoder_mask = self.encode_video_classifier_free(video, text, num_samples_per_prompt, device=device)
        # import pdb;pdb.set_trace()
        #print("ldm time 1", time.time()-start)
        inference_scheduler.set_timesteps(num_steps, device=device)
        timesteps = inference_scheduler.timesteps

        num_channels_latents = self.unet.in_channels
        latents = self.prepare_latents(batch_size, inference_scheduler, num_channels_latents, encoder_hidden_states.dtype, device)
        num_warmup_steps = len(timesteps) - num_steps * inference_scheduler.order
        progress_bar = tqdm(range(num_steps), disable=disable_progress)

        #print("ldm time 2", time.time()-start, timesteps)
        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)

            #print("ldm emu", i, time.time()-start)
            noise_pred = self.unet(
                latent_model_input, t, encoder_hidden_states=encoder_hidden_states,
                encoder_attention_mask=boolean_encoder_mask
            ).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

            # call the callback, if provided
            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % inference_scheduler.order == 0):
                progress_bar.update(1)

        #print("ldm time 3", time.time()-start)
        if self.set_from == "pre-trained":
            latents = self.group_out(latents.permute(0, 2, 3, 1).contiguous()).permute(0, 3, 1, 2).contiguous()
        return latents

    def prepare_latents(self, 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(self, prompt, num_samples_per_prompt):
        device = self.text_encoder.device
        batch = self.tokenizer(
            prompt, max_length=self.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 = self.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 = self.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 = self.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)

        # import pdb;pdb.set_trace()
        return prompt_embeds, boolean_prompt_mask

    def encode_image_no_grad(self, prompt, device):
        with torch.no_grad():
            img_fea = self.img_fea_extractor(prompt) 
        encoder_hidden_states = img_fea.view(img_fea.shape[0], img_fea.shape[1], -1).permute(0, 2, 1)
        boolean_encoder_mask = torch.ones((encoder_hidden_states.shape[0], encoder_hidden_states.shape[1]), dtype=torch.bool)
        boolean_encoder_mask = boolean_encoder_mask.to(device)

        return encoder_hidden_states, boolean_encoder_mask
    
    def encode_text_clip_classifier_free(self, prompt, num_samples_per_prompt, device):
        # 如果想测试输入文本的效果，就用下面两行
        with torch.no_grad():
            encoder_hidden_states, boolean_encoder_mask = self.encode_text_CLIP(prompt, device)
        # if "clip-vit" in self.fea_encoder_name:
        #     with torch.no_grad():
        #         inputs = self.CLIP_processor(text=['aaa'], images=prompt, return_tensors="pt", padding=True).to(device)
        #         encoder_hidden_states = self.CLIP_model(**inputs).image_embeds   # b * 768
        #         encoder_hidden_states = self.linear_layer(encoder_hidden_states)    # b * 1024
        #         encoder_hidden_states = encoder_hidden_states.unsqueeze(1).to(device)
        #         boolean_encoder_mask = torch.ones((encoder_hidden_states.shape[0], encoder_hidden_states.shape[1]), dtype=torch.bool)
        #         boolean_encoder_mask = boolean_encoder_mask.to(device)

        b, t, n = encoder_hidden_states.shape
        attention_mask = boolean_encoder_mask.to(device)
        prompt_embeds = encoder_hidden_states.repeat_interleave(num_samples_per_prompt, 0)
        attention_mask = attention_mask.repeat_interleave(num_samples_per_prompt, 0)

        negative_prompt_embeds = encoder_hidden_states.new(b, t, n).fill_(0)
        uncond_attention_mask = torch.ones((b, t), dtype=torch.bool).to(device)

        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])
        boolean_prompt_mask = torch.cat([uncond_attention_mask, attention_mask])

        return prompt_embeds.to(device), boolean_prompt_mask.to(device)


    def encode_image_classifier_free(self, prompt, num_samples_per_prompt, device):
        with torch.no_grad():
            if "clip-vit" in self.fea_encoder_name:
                inputs = self.CLIP_processor(text=["aaa"], images=prompt, return_tensors="pt", padding=True).to(device)
                img_fea = self.CLIP_model(**inputs).image_embeds
                img_fea = self.linear_layer(img_fea)
            else:
                img_fea = self.img_fea_extractor(prompt) 
        encoder_hidden_states = img_fea.view(img_fea.shape[0], img_fea.shape[1], -1).permute(0, 2, 1)
        b, t, n = encoder_hidden_states.shape
        boolean_encoder_mask = torch.ones((b, t), dtype=torch.bool)
        attention_mask = boolean_encoder_mask.to(device)
        prompt_embeds = encoder_hidden_states.repeat_interleave(num_samples_per_prompt, 0)
        attention_mask = attention_mask.repeat_interleave(num_samples_per_prompt, 0)

        negative_prompt_embeds = encoder_hidden_states.new(b, t, n).fill_(0)
        uncond_attention_mask = torch.ones((b, t), dtype=torch.bool).to(device)

        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])
        boolean_prompt_mask = torch.cat([uncond_attention_mask, attention_mask])

        return prompt_embeds.to(device), boolean_prompt_mask.to(device)
    
    def encode_video_classifier_free(self, video_batch, text_batch, num_samples_per_prompt, device):
        vid_feas = []
        for i, video in enumerate(video_batch):
            if text_batch:
                vid_fea = self.vid_fea_extractor(video=video.to(device), text=text_batch[i])
            else:
                vid_fea = self.vid_fea_extractor(video=video.to(device))
            vid_feas.append(vid_fea)
        
        padding = 0
        size = max(v.size(0) for v in vid_feas)
        batch_size = len(vid_feas)
        embed_size = vid_feas[0].size(1)
        encoder_hidden_states = vid_feas[0].new(batch_size, size, embed_size).fill_(padding)
        boolean_encoder_mask = torch.ones((batch_size, size), dtype=torch.bool)

        def copy_tensor(src, dst):
            assert dst.numel() == src.numel()
            dst.copy_(src)

        for i, v in enumerate(vid_feas):
            copy_tensor(v, encoder_hidden_states[i][: len(v)])
            boolean_encoder_mask[i, len(v):] = False    

        b, t, n = encoder_hidden_states.shape
        negative_prompt_embeds = encoder_hidden_states.new(b, t, n).fill_(0)
        uncond_attention_mask = torch.ones((b, t), dtype=torch.bool)

        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
        encoder_hidden_states = torch.cat([negative_prompt_embeds, encoder_hidden_states])
        boolean_encoder_mask = torch.cat([uncond_attention_mask, boolean_encoder_mask])

        return encoder_hidden_states.to(device), boolean_encoder_mask.to(device)