hao-ai-lab-FastVideo/fastvideo/v1/models/text_encoder.py

from dataclasses import dataclass
from typing import Any, Optional

import torch
import torch.nn as nn
from transformers.utils import ModelOutput

from fastvideo.v1.forward_context import set_forward_context
from fastvideo.v1.logger import init_logger

logger = init_logger(__name__)


def use_default(value, default) -> Any:
    return value if value is not None else default


@dataclass
class TextEncoderModelOutput(ModelOutput):
    """
    Base class for model's outputs that also contains a pooling of the last hidden states.

    Args:
        hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
            Sequence of hidden-states at the output of the last layer of the model.

        attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
        hidden_states_list (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed):
            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
        text_outputs (`list`, *optional*, returned when `return_texts=True` is passed):
            List of decoded texts.
    """

    hidden_state: torch.FloatTensor = None
    attention_mask: Optional[torch.LongTensor] = None
    text_outputs: Optional[list] = None


class TextEncoder(nn.Module):

    def __init__(
        self,
        text_encoder,
        tokenizer,
        max_length: int,
        text_encoder_precision: Optional[str] = None,
        text_encoder_path: Optional[str] = None,
        output_key: Optional[str] = None,
        use_attention_mask: bool = True,
        prompt_template: Optional[dict] = None,
        prompt_template_video: Optional[dict] = None,
        hidden_state_skip_layer: Optional[int] = None,
        apply_final_norm: bool = False,
        device=None,
    ):
        super().__init__()
        # TODO(will): check if there's a cleaner way to do this
        self.text_encoder_type = text_encoder.config.architectures[0]
        self.max_length = max_length
        self.precision = text_encoder_precision
        self.model_path = text_encoder_path
        self.use_attention_mask = use_attention_mask
        if prompt_template_video is not None:
            assert (use_attention_mask is True
                    ), "Attention mask is True required when training videos."
        self.prompt_template = prompt_template
        self.prompt_template_video = prompt_template_video
        self.hidden_state_skip_layer = hidden_state_skip_layer
        self.apply_final_norm = apply_final_norm

        if "T5" in self.text_encoder_type:
            self.output_key = output_key or "last_hidden_state"
        elif "CLIPTextModel" in self.text_encoder_type:
            self.output_key = output_key or "pooler_output"
        elif "LlamaModel" in self.text_encoder_type or "glm" in self.text_encoder_type:
            self.output_key = output_key or "last_hidden_state"
        else:
            raise ValueError(
                f"Unsupported text encoder type: {self.text_encoder_type}")

        self.model = text_encoder
        # self.dtype = self.model.dtype
        self.device = device

        self.tokenizer = tokenizer

    def __repr__(self):
        return f"{self.text_encoder_type} ({self.precision} - {self.model_path})"

    @staticmethod
    def apply_text_to_template(text, template, prevent_empty_text=True) -> str:
        """
        Apply text to template.

        Args:
            text (str): Input text.
            template (str or list): Template string or list of chat conversation.
            prevent_empty_text (bool): If True, we will prevent the user text from being empty
                by adding a space. Defaults to True.
        """
        if isinstance(template, str):
            # Will send string to tokenizer. Used for llm
            return template.format(text)
        else:
            raise TypeError(f"Unsupported template type: {type(template)}")

    def text2tokens(self, text) -> dict:
        """
        Tokenize the input text.

        Args:
            text (str or list): Input text.
        """
        if self.prompt_template_video is not None:
            prompt_template = self.prompt_template_video["template"]

            text = self.apply_text_to_template(text, prompt_template)

        kwargs = dict(
            truncation=True,
            max_length=self.max_length,
            return_tensors="pt",
        )
        batch_encoding: dict = self.tokenizer(
            text,
            return_length=False,
            return_overflowing_tokens=False,
            return_attention_mask=True,
            **kwargs,
        )
        return batch_encoding

    def encode(
        self,
        batch_encoding,
        use_attention_mask=None,
        hidden_state_skip_layer=None,
        device=None,
    ) -> TextEncoderModelOutput:
        """
        Args:
            batch_encoding (dict): Batch encoding from tokenizer.
            use_attention_mask (bool): Whether to use attention mask. If None, use self.use_attention_mask.
                Defaults to None.
            output_hidden_states (bool): Whether to output hidden states. If False, return the value of
                self.output_key. If True, return the entire output. If set self.hidden_state_skip_layer,
                output_hidden_states will be set True. Defaults to False.
            hidden_state_skip_layer (int): Number of hidden states to hidden_state_skip_layer. 0 means the last layer.
                If None, self.output_key will be used. Defaults to None.
            return_texts (bool): Whether to return the decoded texts. Defaults to False.
        """
        device = self.model.device if device is None else device
        use_attention_mask = use_default(use_attention_mask,
                                         self.use_attention_mask)
        hidden_state_skip_layer = use_default(hidden_state_skip_layer,
                                              self.hidden_state_skip_layer)

        # note: clip will need attention mask
        # TODO(will): unify interface with dit
        # TODO (peiyuan): why clip need attention mask?
        with set_forward_context(current_timestep=0, attn_metadata=None):
            outputs = self.model(
                input_ids=batch_encoding["input_ids"].to(device),
                output_hidden_states=hidden_state_skip_layer is not None,
            )
        if hidden_state_skip_layer is not None:
            last_hidden_state = outputs.hidden_states[-(
                hidden_state_skip_layer + 1)]
            # Real last hidden state already has layer norm applied. So here we only apply it
            # for intermediate layers.
            if hidden_state_skip_layer > 0 and self.apply_final_norm:
                last_hidden_state = self.model.final_layer_norm(
                    last_hidden_state)
        else:
            last_hidden_state = outputs[self.output_key]

        # Remove hidden states of instruction tokens, only keep prompt tokens.
        if self.prompt_template_video is not None:

            crop_start = self.prompt_template_video.get("crop_start", -1)

            last_hidden_state = last_hidden_state[:, crop_start:]

        return TextEncoderModelOutput(last_hidden_state)

    def forward(
        self,
        text,
        use_attention_mask=None,
        output_hidden_states=False,
        hidden_state_skip_layer=None,
        return_texts=False,
    ):
        batch_encoding = self.text2tokens(text)
        return self.encode(
            batch_encoding,
            use_attention_mask=use_attention_mask,
            hidden_state_skip_layer=hidden_state_skip_layer,
        )