Publish pinned visual components with exact source payloads and native numerical validation

b46df1b verified 13 days ago

16.3 kB

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from .modules import modulate, RMSNorm

from timm.layers.helpers import to_2tuple
from transformers import AutoTokenizer


def next_token_prediction(logits, labels, vocab_szie):
    return F.cross_entropy(logits[:, :-1].contiguous().view(-1, vocab_szie), labels[:, 1:].contiguous().view(-1),
                           ignore_index=-100)


def velocity_prediction(latents, labels, mask=None):
    if mask is not None:
        loss = F.mse_loss(latents, labels, reduction='none')[mask.bool()]
        return loss.mean()
    else:
        return F.mse_loss(latents, labels)


def prepare_gen_input(prompts, text_tokenizer, num_image_tokens, bos_id, eos_id, boi_id, eoi_id, pad_id, img_pad_id,
                      max_text_len, device):
    batch_text_tokens = []
    batch_modality_positions = []
    batch_text_tokens_null = []
    batch_modality_positions_null = []
    for prompt in prompts:
        text_tokens = text_tokenizer(prompt, add_special_tokens=False)['input_ids'][:(max_text_len)]

        modality_positions = torch.tensor([len(text_tokens) + 1 + 1, num_image_tokens]).unsqueeze(0)
        text_tokens = [bos_id] + text_tokens + [boi_id] + [img_pad_id] * num_image_tokens + \
                      [eoi_id] + [eos_id] + [pad_id] * (max_text_len - len(text_tokens))

        batch_text_tokens.append(torch.tensor(text_tokens))
        batch_modality_positions.append(modality_positions)

        text_tokens_null = []
        modality_positions_null = torch.tensor([len(text_tokens_null) + 1 + 1, num_image_tokens]).unsqueeze(0)
        text_tokens_null = [bos_id] + text_tokens_null + [boi_id] + [img_pad_id] * num_image_tokens + \
                           [eoi_id] + [eos_id] + [pad_id] * (max_text_len - len(text_tokens_null))

        batch_text_tokens_null.append(torch.tensor(text_tokens_null))
        batch_modality_positions_null.append(modality_positions_null)

    batch_text_tokens = torch.stack(batch_text_tokens, dim=0).to(device)
    batch_modality_positions = torch.stack(batch_modality_positions, dim=0).to(device)

    batch_text_tokens_null = torch.stack(batch_text_tokens_null, dim=0).to(device)
    batch_modality_positions_null = torch.stack(batch_modality_positions_null, dim=0).to(device)

    return batch_text_tokens, batch_text_tokens_null, batch_modality_positions, batch_modality_positions_null


def prepare_mixed_modal_gen_input(prompts, nulls, text_tokenizer, num_image_tokens, bos_id, boi_id, eoi_id, pad_id, img_pad_id, device):
    batch_text_tokens = []
    batch_modality_positions = []
    batch_text_tokens_null = []
    batch_modality_positions_null = []
    for prompt, null in zip(prompts, nulls):
        text_tokens = text_tokenizer(prompt, add_special_tokens=False).input_ids
        modality_positions = torch.tensor([len(text_tokens) + 1 + 1, num_image_tokens]).unsqueeze(0)
        text_tokens = [bos_id] + text_tokens + [boi_id] + [img_pad_id] * num_image_tokens + [eoi_id]

        text_tokens_null = text_tokenizer(null, add_special_tokens=False).input_ids
        modality_positions_null = torch.tensor([len(text_tokens_null) + 1 + 1, num_image_tokens]).unsqueeze(0)
        text_tokens_null = [bos_id] + text_tokens_null + [boi_id] + [img_pad_id] * num_image_tokens + [eoi_id]

        len_a = len(text_tokens)
        len_b = len(text_tokens_null)

        max_len = max(len_a, len_b)

        if max_len % 128 != 0:
            max_len = (max_len // 128 + 1) * 128

        num_pads_a = max_len - len_a
        num_pads_b = max_len - len_b

        text_tokens = text_tokens + [pad_id] * num_pads_a
        text_tokens_null = text_tokens_null + [pad_id] * num_pads_b

        batch_text_tokens.append(torch.tensor(text_tokens))
        batch_modality_positions.append(modality_positions)

        batch_text_tokens_null.append(torch.tensor(text_tokens_null))
        batch_modality_positions_null.append(modality_positions_null)

    batch_text_tokens = torch.stack(batch_text_tokens, dim=0).to(device)
    batch_modality_positions = torch.stack(batch_modality_positions, dim=0).to(device)

    batch_text_tokens_null = torch.stack(batch_text_tokens_null, dim=0).to(device)
    batch_modality_positions_null = torch.stack(batch_modality_positions_null, dim=0).to(device)

    return batch_text_tokens, batch_text_tokens_null, batch_modality_positions, batch_modality_positions_null


# def prepare_mixed_modal_gen_input(prompt, text_tokenizer, num_image_tokens, boi_id, eoi_id, img_pad_id, pad_id, device):
#     text_tokens = text_tokenizer(prompt, add_special_tokens=False).input_ids
#
#     modality_positions = torch.Tensor([[len(text_tokens), num_image_tokens]]).long().unsqueeze(0)
#     text_tokens = text_tokens + [img_pad_id] * num_image_tokens + [eoi_id]
#
#     modality_positions_null = torch.Tensor([[2, num_image_tokens]]).long().unsqueeze(0)
#     text_tokens_null = [text_tokens[0]] + [boi_id] + [img_pad_id] * num_image_tokens + [eoi_id]
#
#     len_a = len(text_tokens)
#     len_b = len(text_tokens_null)
#     num_pads_a = max(len_a, len_b) - len_a
#     num_pads_b = max(len_a, len_b) - len_b
#
#     text_tokens += [pad_id] * num_pads_a
#     text_tokens_null += [pad_id] * num_pads_b
#
#     text_tokens = torch.tensor(text_tokens).unsqueeze(0)
#     text_tokens_null = torch.tensor(text_tokens_null).unsqueeze(0)
#
#     return text_tokens.to(device), text_tokens_null.to(device), \
#            modality_positions.to(device), modality_positions_null.to(device)


class PatchEmbed(nn.Module):
    """2D Image to Patch Embedding"""

    def __init__(
            self,
            patch_size=16,
            in_chans=3,
            embed_dim=768,
            kernel_size=None,
            padding=0,
            norm_layer=None,
            flatten=True,
            bias=True,
    ):
        super().__init__()
        kernel_size = kernel_size or patch_size
        patch_size = to_2tuple(patch_size)
        self.patch_size = patch_size
        self.flatten = flatten
        self.proj = nn.Conv2d(
            in_chans, embed_dim, kernel_size=kernel_size, stride=patch_size, bias=bias
        )
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

    def forward(self, x):
        x = self.proj(x)
        if self.flatten:
            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
        x = self.norm(x)
        return x


class TimestepEmbedder(nn.Module):
    """
    Embeds scalar timesteps into vector representations.
    """

    def __init__(self, hidden_size, frequency_embedding_size=256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size, bias=True),
        )
        self.frequency_embedding_size = frequency_embedding_size

    @staticmethod
    def timestep_embedding(t, dim, max_period=10000):
        """
        Create sinusoidal timestep embeddings.
        :param t: a 1-D Tensor of N indices, one per batch element.
                          These may be fractional.
        :param dim: the dimension of the output.
        :param max_period: controls the minimum frequency of the embeddings.
        :return: an (N, D) Tensor of positional embeddings.
        """
        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
        half = dim // 2
        freqs = torch.exp(
            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
        ).to(device=t.device)
        args = t[:, None].float() * freqs[None]
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        if dim % 2:
            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
        return embedding

    def forward(self, t, dtype):
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(dtype)
        t_emb = self.mlp(t_freq)
        return t_emb


class FinalLayer(nn.Module):
    """
    The final layer of DiT.
    """

    def __init__(self, hidden_size, patch_size, out_channels):
        super().__init__()
        self.norm_final = RMSNorm(hidden_size)
        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 2 * hidden_size, bias=True)
        )

    def forward(self, x, adaln_input):
        shift, scale = self.adaLN_modulation(adaln_input).chunk(2, dim=1)
        x = modulate(self.norm_final(x), shift, scale)
        x = self.linear(x)
        return x


class UpdatedVisionTransformer(nn.Module):
    def __init__(self, model, del_last_layer=True):
        super().__init__()
        self.model = model
        if del_last_layer:
            del self.model.transformer.resblocks[-1]

    def forward(self, x: torch.Tensor):
        x = self.model.conv1(x)  # shape = [*, width, grid, grid]
        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
        x = torch.cat([self.model.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype,
                                                                            device=x.device), x],
                      dim=1)  # shape = [*, grid ** 2 + 1, width]
        x = x + self.model.positional_embedding.to(x.dtype)
        x = self.model.ln_pre(x)

        x = x.permute(1, 0, 2)  # NLD -> LND
        x = self.model.transformer(x)
        x = x.permute(1, 0, 2)[:, 1:]  # LND -> NLD

        return x


class CLIPVisionEncoder(nn.Module):
    def __init__(self, model, del_last_layer=False):
        super().__init__()
        self.model = model
        if del_last_layer:
            del self.model.transformer.resblocks[-1]

    def forward(self, x: torch.Tensor):
        x = self.model.conv1(x)  # shape = [*, width, grid, grid]
        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
        x = torch.cat([self.model.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype,
                                                                            device=x.device), x],
                      dim=1)  # shape = [*, grid ** 2 + 1, width]
        x = x + self.model.positional_embedding.to(x.dtype)
        x = self.model.ln_pre(x)

        x = x.permute(1, 0, 2)  # NLD -> LND
        x = self.model.transformer(x)
        x = x.permute(1, 0, 2)[:, 1:]  # LND -> NLD

        return x


class SigLipVisionEncoder(nn.Module):
    def __init__(self, model, del_last_layer=True):
        """
        A wrapper for extracting features from the penultimate layer of a vision transformer model.
        Args:
            model: The pre-trained model (e.g., CLIP or SigLIP).
            del_last_layer (bool): Whether to delete the last layer of the vision transformer.
        """
        super().__init__()
        self.model = model

        # Remove the text model (if not needed)
        if hasattr(self.model, "text_model"):
            del self.model.text_model

        # Remove the last layer of the vision transformer
        if del_last_layer and hasattr(self.model.vision_model, "encoder"):
            del self.model.vision_model.encoder.layers[-1]

        # Replace the classification head (if it exists) with an identity layer
        if hasattr(self.model.vision_model, "head"):
            self.model.vision_model.head = nn.Identity()
        if hasattr(self.model.vision_model, "post_layernorm"):
            self.model.vision_model.post_layernorm = nn.Identity()

    def forward(self, x):
        """
        Forward pass to extract features from the penultimate layer.
        Args:
            x: Input image tensor (pixel values).
        Returns:
            Tensor: Features from the penultimate layer.
        """
        return self.model.get_image_features(pixel_values=x)

from transformers.utils import torch_int
def interpolate_pos_encoding(dim: int, position_embedding: torch.nn.Embedding, height: int, width: int,
                             patch_size: int) -> torch.Tensor:
    """
    This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution
    images. This method is also adapted to support torch.jit tracing and no class embeddings.
    Adapted from:
    - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
    - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
    """

    num_positions = position_embedding.weight.shape[0]
    patch_pos_embed = position_embedding.weight.unsqueeze(0)

    new_height = height // patch_size
    new_width = width // patch_size

    sqrt_num_positions = torch_int(num_positions**0.5)
    patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)
    patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)

    patch_pos_embed = nn.functional.interpolate(
        patch_pos_embed,
        size=(new_height, new_width),
        mode="bicubic",
        align_corners=False,
    )

    patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
    return patch_pos_embed


def get_text_tokenizer(model_path, add_showo_tokens=True, return_showo_token_ids=False, llm_name="qwen2_5"):
    text_tokenizer = AutoTokenizer.from_pretrained(model_path)
    text_tokenizer.add_special_tokens({'pad_token': '[PAD]'})
    if add_showo_tokens:
        if llm_name == "llama3":
            text_tokenizer.add_tokens('<|img_start|>')
            text_tokenizer.add_tokens('<|img_end|>')
            text_tokenizer.add_tokens('<|image_pad|>')
            text_tokenizer.add_tokens('<|video_pad|>')
            text_tokenizer.add_tokens('<|vid_start|>')
            text_tokenizer.add_tokens('<|vid_end|>')
            text_tokenizer.add_tokens('<image>')
        elif llm_name == "qwen2_5":
            text_tokenizer.add_tokens('<image>')
            text_tokenizer.add_tokens('<|vid_start|>')
            text_tokenizer.add_tokens('<|vid_end|>')
        else:
            raise NotImplementedError

    if return_showo_token_ids:
        if llm_name == "llama3":
            showo_token_ids = {
                "bos_id": text_tokenizer.get_vocab()["<|begin_of_text|>"],
                "eos_id": text_tokenizer.eos_token_id,
                "boi_id": text_tokenizer.get_vocab()["<|img_start|>"],
                "eoi_id": text_tokenizer.get_vocab()["<|img_end|>"],
                "bov_id": text_tokenizer.get_vocab()["<|vid_start|>"],
                "eov_id": text_tokenizer.get_vocab()["<|vid_end|>"],
                "img_pad_id": text_tokenizer.get_vocab()["<|image_pad|>"],
                "vid_pad_id": text_tokenizer.get_vocab()["<|video_pad|>"],
                "img_id": text_tokenizer.get_vocab()["<image>"],
            }
        elif llm_name == "qwen2_5":
            showo_token_ids = {
                "bos_id": text_tokenizer.get_vocab()["<|im_start|>"],
                "eos_id": text_tokenizer.eos_token_id,
                "boi_id": text_tokenizer.get_vocab()["<|vision_start|>"],
                "eoi_id": text_tokenizer.get_vocab()["<|vision_end|>"],
                "bov_id": text_tokenizer.get_vocab()["<|vid_start|>"],
                "eov_id": text_tokenizer.get_vocab()["<|vid_end|>"],
                "img_pad_id": text_tokenizer.get_vocab()["<|image_pad|>"],
                "vid_pad_id": text_tokenizer.get_vocab()["<|video_pad|>"],
                "img_id": text_tokenizer.get_vocab()["<image>"],
            }
        else:
            raise NotImplementedError

        return text_tokenizer, showo_token_ids

    return text_tokenizer


def get_weight_type(config):
    if config.training.mixed_precision == 'bf16':
        weight_type = torch.bfloat16
    elif config.training.mixed_precision == 'float16':
        weight_type = torch.float16
    else:
        weight_type = torch.float32
    return weight_type

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	import math
	from .modules import modulate, RMSNorm

	from timm.layers.helpers import to_2tuple
	from transformers import AutoTokenizer


	def next_token_prediction(logits, labels, vocab_szie):
	return F.cross_entropy(logits[:, :-1].contiguous().view(-1, vocab_szie), labels[:, 1:].contiguous().view(-1),
	ignore_index=-100)


	def velocity_prediction(latents, labels, mask=None):
	if mask is not None:
	loss = F.mse_loss(latents, labels, reduction='none')[mask.bool()]
	return loss.mean()
	else:
	return F.mse_loss(latents, labels)


	def prepare_gen_input(prompts, text_tokenizer, num_image_tokens, bos_id, eos_id, boi_id, eoi_id, pad_id, img_pad_id,
	max_text_len, device):
	batch_text_tokens = []
	batch_modality_positions = []
	batch_text_tokens_null = []
	batch_modality_positions_null = []
	for prompt in prompts:
	text_tokens = text_tokenizer(prompt, add_special_tokens=False)['input_ids'][:(max_text_len)]

	modality_positions = torch.tensor([len(text_tokens) + 1 + 1, num_image_tokens]).unsqueeze(0)
	text_tokens = [bos_id] + text_tokens + [boi_id] + [img_pad_id] * num_image_tokens + \
	[eoi_id] + [eos_id] + [pad_id] * (max_text_len - len(text_tokens))

	batch_text_tokens.append(torch.tensor(text_tokens))
	batch_modality_positions.append(modality_positions)

	text_tokens_null = []
	modality_positions_null = torch.tensor([len(text_tokens_null) + 1 + 1, num_image_tokens]).unsqueeze(0)
	text_tokens_null = [bos_id] + text_tokens_null + [boi_id] + [img_pad_id] * num_image_tokens + \
	[eoi_id] + [eos_id] + [pad_id] * (max_text_len - len(text_tokens_null))

	batch_text_tokens_null.append(torch.tensor(text_tokens_null))
	batch_modality_positions_null.append(modality_positions_null)

	batch_text_tokens = torch.stack(batch_text_tokens, dim=0).to(device)
	batch_modality_positions = torch.stack(batch_modality_positions, dim=0).to(device)

	batch_text_tokens_null = torch.stack(batch_text_tokens_null, dim=0).to(device)
	batch_modality_positions_null = torch.stack(batch_modality_positions_null, dim=0).to(device)

	return batch_text_tokens, batch_text_tokens_null, batch_modality_positions, batch_modality_positions_null


	def prepare_mixed_modal_gen_input(prompts, nulls, text_tokenizer, num_image_tokens, bos_id, boi_id, eoi_id, pad_id, img_pad_id, device):
	batch_text_tokens = []
	batch_modality_positions = []
	batch_text_tokens_null = []
	batch_modality_positions_null = []
	for prompt, null in zip(prompts, nulls):
	text_tokens = text_tokenizer(prompt, add_special_tokens=False).input_ids
	modality_positions = torch.tensor([len(text_tokens) + 1 + 1, num_image_tokens]).unsqueeze(0)
	text_tokens = [bos_id] + text_tokens + [boi_id] + [img_pad_id] * num_image_tokens + [eoi_id]

	text_tokens_null = text_tokenizer(null, add_special_tokens=False).input_ids
	modality_positions_null = torch.tensor([len(text_tokens_null) + 1 + 1, num_image_tokens]).unsqueeze(0)
	text_tokens_null = [bos_id] + text_tokens_null + [boi_id] + [img_pad_id] * num_image_tokens + [eoi_id]

	len_a = len(text_tokens)
	len_b = len(text_tokens_null)

	max_len = max(len_a, len_b)

	if max_len % 128 != 0:
	max_len = (max_len // 128 + 1) * 128

	num_pads_a = max_len - len_a
	num_pads_b = max_len - len_b

	text_tokens = text_tokens + [pad_id] * num_pads_a
	text_tokens_null = text_tokens_null + [pad_id] * num_pads_b

	batch_text_tokens.append(torch.tensor(text_tokens))
	batch_modality_positions.append(modality_positions)

	batch_text_tokens_null.append(torch.tensor(text_tokens_null))
	batch_modality_positions_null.append(modality_positions_null)

	batch_text_tokens = torch.stack(batch_text_tokens, dim=0).to(device)
	batch_modality_positions = torch.stack(batch_modality_positions, dim=0).to(device)

	batch_text_tokens_null = torch.stack(batch_text_tokens_null, dim=0).to(device)
	batch_modality_positions_null = torch.stack(batch_modality_positions_null, dim=0).to(device)

	return batch_text_tokens, batch_text_tokens_null, batch_modality_positions, batch_modality_positions_null


	# def prepare_mixed_modal_gen_input(prompt, text_tokenizer, num_image_tokens, boi_id, eoi_id, img_pad_id, pad_id, device):
	# text_tokens = text_tokenizer(prompt, add_special_tokens=False).input_ids
	#
	# modality_positions = torch.Tensor([[len(text_tokens), num_image_tokens]]).long().unsqueeze(0)
	# text_tokens = text_tokens + [img_pad_id] * num_image_tokens + [eoi_id]
	#
	# modality_positions_null = torch.Tensor([[2, num_image_tokens]]).long().unsqueeze(0)
	# text_tokens_null = [text_tokens[0]] + [boi_id] + [img_pad_id] * num_image_tokens + [eoi_id]
	#
	# len_a = len(text_tokens)
	# len_b = len(text_tokens_null)
	# num_pads_a = max(len_a, len_b) - len_a
	# num_pads_b = max(len_a, len_b) - len_b
	#
	# text_tokens += [pad_id] * num_pads_a
	# text_tokens_null += [pad_id] * num_pads_b
	#
	# text_tokens = torch.tensor(text_tokens).unsqueeze(0)
	# text_tokens_null = torch.tensor(text_tokens_null).unsqueeze(0)
	#
	# return text_tokens.to(device), text_tokens_null.to(device), \
	# modality_positions.to(device), modality_positions_null.to(device)


	class PatchEmbed(nn.Module):
	"""2D Image to Patch Embedding"""

	def __init__(
	self,
	patch_size=16,
	in_chans=3,
	embed_dim=768,
	kernel_size=None,
	padding=0,
	norm_layer=None,
	flatten=True,
	bias=True,
	):
	super().__init__()
	kernel_size = kernel_size or patch_size
	patch_size = to_2tuple(patch_size)
	self.patch_size = patch_size
	self.flatten = flatten
	self.proj = nn.Conv2d(
	in_chans, embed_dim, kernel_size=kernel_size, stride=patch_size, bias=bias
	)
	self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

	def forward(self, x):
	x = self.proj(x)
	if self.flatten:
	x = x.flatten(2).transpose(1, 2) # BCHW -> BNC
	x = self.norm(x)
	return x


	class TimestepEmbedder(nn.Module):
	"""
	Embeds scalar timesteps into vector representations.
	"""

	def __init__(self, hidden_size, frequency_embedding_size=256):
	super().__init__()
	self.mlp = nn.Sequential(
	nn.Linear(frequency_embedding_size, hidden_size, bias=True),
	nn.SiLU(),
	nn.Linear(hidden_size, hidden_size, bias=True),
	)
	self.frequency_embedding_size = frequency_embedding_size

	@staticmethod
	def timestep_embedding(t, dim, max_period=10000):
	"""
	Create sinusoidal timestep embeddings.
	:param t: a 1-D Tensor of N indices, one per batch element.
	These may be fractional.
	:param dim: the dimension of the output.
	:param max_period: controls the minimum frequency of the embeddings.
	:return: an (N, D) Tensor of positional embeddings.
	"""
	# https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
	half = dim // 2
	freqs = torch.exp(
	-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
	).to(device=t.device)
	args = t[:, None].float() * freqs[None]
	embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
	if dim % 2:
	embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
	return embedding

	def forward(self, t, dtype):
	t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(dtype)
	t_emb = self.mlp(t_freq)
	return t_emb


	class FinalLayer(nn.Module):
	"""
	The final layer of DiT.
	"""

	def __init__(self, hidden_size, patch_size, out_channels):
	super().__init__()
	self.norm_final = RMSNorm(hidden_size)
	self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
	self.adaLN_modulation = nn.Sequential(
	nn.SiLU(),
	nn.Linear(hidden_size, 2 * hidden_size, bias=True)
	)

	def forward(self, x, adaln_input):
	shift, scale = self.adaLN_modulation(adaln_input).chunk(2, dim=1)
	x = modulate(self.norm_final(x), shift, scale)
	x = self.linear(x)
	return x


	class UpdatedVisionTransformer(nn.Module):
	def __init__(self, model, del_last_layer=True):
	super().__init__()
	self.model = model
	if del_last_layer:
	del self.model.transformer.resblocks[-1]

	def forward(self, x: torch.Tensor):
	x = self.model.conv1(x) # shape = [*, width, grid, grid]
	x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [, width, grid * 2]
	x = x.permute(0, 2, 1) # shape = [, grid * 2, width]
	x = torch.cat([self.model.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype,
	device=x.device), x],
	dim=1) # shape = [, grid * 2 + 1, width]
	x = x + self.model.positional_embedding.to(x.dtype)
	x = self.model.ln_pre(x)

	x = x.permute(1, 0, 2) # NLD -> LND
	x = self.model.transformer(x)
	x = x.permute(1, 0, 2)[:, 1:] # LND -> NLD

	return x


	class CLIPVisionEncoder(nn.Module):
	def __init__(self, model, del_last_layer=False):
	super().__init__()
	self.model = model
	if del_last_layer:
	del self.model.transformer.resblocks[-1]

	def forward(self, x: torch.Tensor):
	x = self.model.conv1(x) # shape = [*, width, grid, grid]
	x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [, width, grid * 2]
	x = x.permute(0, 2, 1) # shape = [, grid * 2, width]
	x = torch.cat([self.model.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype,
	device=x.device), x],
	dim=1) # shape = [, grid * 2 + 1, width]
	x = x + self.model.positional_embedding.to(x.dtype)
	x = self.model.ln_pre(x)

	x = x.permute(1, 0, 2) # NLD -> LND
	x = self.model.transformer(x)
	x = x.permute(1, 0, 2)[:, 1:] # LND -> NLD

	return x


	class SigLipVisionEncoder(nn.Module):
	def __init__(self, model, del_last_layer=True):
	"""
	A wrapper for extracting features from the penultimate layer of a vision transformer model.

	Args:
	model: The pre-trained model (e.g., CLIP or SigLIP).
	del_last_layer (bool): Whether to delete the last layer of the vision transformer.
	"""
	super().__init__()
	self.model = model

	# Remove the text model (if not needed)
	if hasattr(self.model, "text_model"):
	del self.model.text_model

	# Remove the last layer of the vision transformer
	if del_last_layer and hasattr(self.model.vision_model, "encoder"):
	del self.model.vision_model.encoder.layers[-1]

	# Replace the classification head (if it exists) with an identity layer
	if hasattr(self.model.vision_model, "head"):
	self.model.vision_model.head = nn.Identity()
	if hasattr(self.model.vision_model, "post_layernorm"):
	self.model.vision_model.post_layernorm = nn.Identity()

	def forward(self, x):
	"""
	Forward pass to extract features from the penultimate layer.

	Args:
	x: Input image tensor (pixel values).

	Returns:
	Tensor: Features from the penultimate layer.
	"""
	return self.model.get_image_features(pixel_values=x)

	from transformers.utils import torch_int
	def interpolate_pos_encoding(dim: int, position_embedding: torch.nn.Embedding, height: int, width: int,
	patch_size: int) -> torch.Tensor:
	"""
	This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution
	images. This method is also adapted to support torch.jit tracing and no class embeddings.

	Adapted from:
	- https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
	- https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
	"""

	num_positions = position_embedding.weight.shape[0]
	patch_pos_embed = position_embedding.weight.unsqueeze(0)

	new_height = height // patch_size
	new_width = width // patch_size

	sqrt_num_positions = torch_int(num_positions**0.5)
	patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)
	patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)

	patch_pos_embed = nn.functional.interpolate(
	patch_pos_embed,
	size=(new_height, new_width),
	mode="bicubic",
	align_corners=False,
	)

	patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
	return patch_pos_embed


	def get_text_tokenizer(model_path, add_showo_tokens=True, return_showo_token_ids=False, llm_name="qwen2_5"):
	text_tokenizer = AutoTokenizer.from_pretrained(model_path)
	text_tokenizer.add_special_tokens({'pad_token': '[PAD]'})
	if add_showo_tokens:
	if llm_name == "llama3":
	text_tokenizer.add_tokens('<\|img_start\|>')
	text_tokenizer.add_tokens('<\|img_end\|>')
	text_tokenizer.add_tokens('<\|image_pad\|>')
	text_tokenizer.add_tokens('<\|video_pad\|>')
	text_tokenizer.add_tokens('<\|vid_start\|>')
	text_tokenizer.add_tokens('<\|vid_end\|>')
	text_tokenizer.add_tokens('<image>')
	elif llm_name == "qwen2_5":
	text_tokenizer.add_tokens('<image>')
	text_tokenizer.add_tokens('<\|vid_start\|>')
	text_tokenizer.add_tokens('<\|vid_end\|>')
	else:
	raise NotImplementedError

	if return_showo_token_ids:
	if llm_name == "llama3":
	showo_token_ids = {
	"bos_id": text_tokenizer.get_vocab()["<\|begin_of_text\|>"],
	"eos_id": text_tokenizer.eos_token_id,
	"boi_id": text_tokenizer.get_vocab()["<\|img_start\|>"],
	"eoi_id": text_tokenizer.get_vocab()["<\|img_end\|>"],
	"bov_id": text_tokenizer.get_vocab()["<\|vid_start\|>"],
	"eov_id": text_tokenizer.get_vocab()["<\|vid_end\|>"],
	"img_pad_id": text_tokenizer.get_vocab()["<\|image_pad\|>"],
	"vid_pad_id": text_tokenizer.get_vocab()["<\|video_pad\|>"],
	"img_id": text_tokenizer.get_vocab()["<image>"],
	}
	elif llm_name == "qwen2_5":
	showo_token_ids = {
	"bos_id": text_tokenizer.get_vocab()["<\|im_start\|>"],
	"eos_id": text_tokenizer.eos_token_id,
	"boi_id": text_tokenizer.get_vocab()["<\|vision_start\|>"],
	"eoi_id": text_tokenizer.get_vocab()["<\|vision_end\|>"],
	"bov_id": text_tokenizer.get_vocab()["<\|vid_start\|>"],
	"eov_id": text_tokenizer.get_vocab()["<\|vid_end\|>"],
	"img_pad_id": text_tokenizer.get_vocab()["<\|image_pad\|>"],
	"vid_pad_id": text_tokenizer.get_vocab()["<\|video_pad\|>"],
	"img_id": text_tokenizer.get_vocab()["<image>"],
	}
	else:
	raise NotImplementedError

	return text_tokenizer, showo_token_ids

	return text_tokenizer


	def get_weight_type(config):
	if config.training.mixed_precision == 'bf16':
	weight_type = torch.bfloat16
	elif config.training.mixed_precision == 'float16':
	weight_type = torch.float16
	else:
	weight_type = torch.float32
	return weight_type