ldm/modules/encoders/modules.py

import os.path

import torch
import torch.nn as nn
from functools import partial
from transformers import T5Tokenizer, T5EncoderModel, AutoTokenizer
from importlib_resources import files
from ldm.modules.encoders.CLAP.utils import read_config_as_args
from ldm.modules.encoders.CLAP.clap import TextEncoder
from ldm.util import count_params
import numpy as np




class Video_Feat_Encoder_NoPosembed(nn.Module):
    """ Transform the video feat encoder"""

    def __init__(self, origin_dim, embed_dim, seq_len=40):
        super().__init__()
        self.embedder = nn.Sequential(nn.Linear(origin_dim, embed_dim))

    def forward(self, x):
        # Revise the shape here:
        x = self.embedder(x)        # B x 117 x C

        return x



class Video_Feat_Encoder_NoPosembed_inpaint(Video_Feat_Encoder_NoPosembed):
    """ Transform the video feat encoder"""

    def forward(self, x):
        # Revise the shape here:
        video, spec = x['mix_video_feat'], x['mix_spec']
        video = self.embedder(video)        # B x 117 x C

        return (video, spec)

class AbstractEncoder(nn.Module):
    def __init__(self):
        super().__init__()

    def encode(self, *args, **kwargs):
        raise NotImplementedError


def disabled_train(self, mode=True):
    """Overwrite model.train with this function to make sure train/eval mode
    does not change anymore."""
    return self

class FrozenFLANEmbedder(AbstractEncoder):
    """Uses the T5 transformer encoder for text"""

    def __init__(self, version="google/flan-t5-large", device="cuda", max_length=77,
                 freeze=True):  # others are google/t5-v1_1-xl and google/t5-v1_1-xxl
        device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        super().__init__()

        self.tokenizer = T5Tokenizer.from_pretrained(version)
        self.transformer = T5EncoderModel.from_pretrained(version)
        self.device = device
        self.max_length = max_length  # TODO: typical value?
        if freeze:
            self.freeze()

    def freeze(self):
        self.transformer = self.transformer.eval()
        # self.train = disabled_train
        for param in self.parameters():
            param.requires_grad = False

    def forward(self, text):
        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
        tokens = batch_encoding["input_ids"].to(self.device)  # tango的flanT5是不定长度的batch，这里做成定长的batch
        outputs = self.transformer(input_ids=tokens)

        z = outputs.last_hidden_state
        return z

    def encode(self, text):
        return self(text)


class FrozenCLAPEmbedder(AbstractEncoder):
    """Uses the CLAP transformer encoder for text from microsoft"""

    def __init__(self, weights_path, freeze=True, device="cuda", max_length=77):  # clip-vit-base-patch32
        device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        super().__init__()
        model_state_dict = torch.load(weights_path, map_location=torch.device('cpu'))['model']
        match_params = dict()
        for key in list(model_state_dict.keys()):
            if 'caption_encoder' in key:
                match_params[key.replace('caption_encoder.', '')] = model_state_dict[key]

        config_as_str = files('ldm').joinpath('modules/encoders/CLAP/config.yml').read_text()
        args = read_config_as_args(config_as_str, is_config_str=True)

        self.tokenizer = AutoTokenizer.from_pretrained(args.text_model)  # args.text_model
        self.caption_encoder = TextEncoder(
            args.d_proj, args.text_model, args.transformer_embed_dim
        )
        
        self.max_length = max_length
        self.device = device
        if freeze: self.freeze()

        print(f"{self.caption_encoder.__class__.__name__} comes with {count_params(self.caption_encoder) * 1.e-6:.2f} M params.")

    def freeze(self):  # only freeze
        self.caption_encoder.base = self.caption_encoder.base.eval()
        for param in self.caption_encoder.base.parameters():
            param.requires_grad = False

    def encode(self, text):
        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
                                        return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
        tokens = batch_encoding["input_ids"].to(self.device)

        outputs = self.caption_encoder.base(input_ids=tokens)
        z = self.caption_encoder.projection(outputs.last_hidden_state)
        return z


class FrozenCLAPFLANEmbedder(AbstractEncoder):
    """Uses the CLAP transformer encoder for text from microsoft"""

    def __init__(self, weights_path, t5version="google/t5-v1_1-large", freeze=True, device="cuda",
                 max_length=77):  # clip-vit-base-patch32
        device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        super().__init__()

        model_state_dict = torch.load(weights_path, map_location=torch.device('cpu'))['model']
        match_params = dict()
        for key in list(model_state_dict.keys()):
            if 'caption_encoder' in key:
                match_params[key.replace('caption_encoder.', '')] = model_state_dict[key]

        config_as_str = files('ldm').joinpath('modules/encoders/CLAP/config.yml').read_text()
        args = read_config_as_args(config_as_str, is_config_str=True)


        self.clap_tokenizer = AutoTokenizer.from_pretrained(args.text_model)  # args.text_model
        self.caption_encoder = TextEncoder(
            args.d_proj, args.text_model, args.transformer_embed_dim
        )

        self.t5_tokenizer = T5Tokenizer.from_pretrained(t5version)
        self.t5_transformer = T5EncoderModel.from_pretrained(t5version)

        self.max_length = max_length
        self.to(device=device)
        if freeze: self.freeze()

        print(
            f"{self.caption_encoder.__class__.__name__} comes with {count_params(self.caption_encoder) * 1.e-6:.2f} M params.")

    def freeze(self):
        self.caption_encoder = self.caption_encoder.eval()
        for param in self.caption_encoder.parameters():
            param.requires_grad = False

    def to(self, device):
        self.t5_transformer.to(device)
        self.caption_encoder.to(device)
        self.device = device

    def encode(self, text):
        ori_caption = text['ori_caption']
        struct_caption = text['struct_caption']
        # print(ori_caption,struct_caption)
        clap_batch_encoding = self.clap_tokenizer(ori_caption, truncation=True, max_length=self.max_length,
                                                  return_length=True,
                                                  return_overflowing_tokens=False, padding="max_length",
                                                  return_tensors="pt")
        ori_tokens = clap_batch_encoding["input_ids"].to(self.device)
        t5_batch_encoding = self.t5_tokenizer(struct_caption, truncation=True, max_length=self.max_length,
                                              return_length=True,
                                              return_overflowing_tokens=False, padding="max_length",
                                              return_tensors="pt")
        struct_tokens = t5_batch_encoding["input_ids"].to(self.device)
        outputs = self.caption_encoder.base(input_ids=ori_tokens)
        z = self.caption_encoder.projection(outputs.last_hidden_state)
        z2 = self.t5_transformer(input_ids=struct_tokens).last_hidden_state
        return torch.concat([z, z2], dim=1)