# Copyright 2025 The HuggingFace Team. All rights reserved.

#

# 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.

import math

import numpy as np

import torch

import torch.nn.functional as F

from torch import nn

from ..utils import deprecate

from .activations import FP32SiLU, get_activation

from .attention_processor import Attention

def get_timestep_embedding(

timesteps: torch.Tensor,

embedding_dim: int,

flip_sin_to_cos: bool = False,

downscale_freq_shift: float = 1,

scale: float = 1,

max_period: int = 10000,

) -> torch.Tensor:

"""

This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.

Args

timesteps (torch.Tensor):

a 1-D Tensor of N indices, one per batch element. These may be fractional.

embedding_dim (int):

the dimension of the output.

flip_sin_to_cos (bool):

Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)

downscale_freq_shift (float):

Controls the delta between frequencies between dimensions

scale (float):

Scaling factor applied to the embeddings.

max_period (int):

Controls the maximum frequency of the embeddings

Returns

torch.Tensor: an [N x dim] Tensor of positional embeddings.

"""

assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"

half_dim = embedding_dim // 2

exponent = -math.log(max_period) * torch.arange(

start=0, end=half_dim, dtype=torch.float32, device=timesteps.device

)

exponent = exponent / (half_dim - downscale_freq_shift)

emb = torch.exp(exponent)

emb = timesteps[:, None].float() * emb[None, :]

# scale embeddings

emb = scale * emb

# concat sine and cosine embeddings

emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)

# flip sine and cosine embeddings

if flip_sin_to_cos:

emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)

# zero pad

if embedding_dim % 2 == 1:

emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))

return emb

def get_3d_sincos_pos_embed(

embed_dim: int,

spatial_size: int | tuple[int, int],

temporal_size: int,

spatial_interpolation_scale: float = 1.0,

temporal_interpolation_scale: float = 1.0,

device: torch.device | None = None,

output_type: str = "np",

) -> torch.Tensor:

r"""

Creates 3D sinusoidal positional embeddings.

Args:

embed_dim (`int`):

The embedding dimension of inputs. It must be divisible by 16.

spatial_size (`int` or `tuple[int, int]`):

The spatial dimension of positional embeddings. If an integer is provided, the same size is applied to both

spatial dimensions (height and width).

temporal_size (`int`):

The temporal dimension of positional embeddings (number of frames).

spatial_interpolation_scale (`float`, defaults to 1.0):

Scale factor for spatial grid interpolation.

temporal_interpolation_scale (`float`, defaults to 1.0):

Scale factor for temporal grid interpolation.

Returns:

`torch.Tensor`:

The 3D sinusoidal positional embeddings of shape `[temporal_size, spatial_size[0] * spatial_size[1],

embed_dim]`.

"""

if output_type == "np":

return _get_3d_sincos_pos_embed_np(

embed_dim=embed_dim,

spatial_size=spatial_size,

temporal_size=temporal_size,

spatial_interpolation_scale=spatial_interpolation_scale,

temporal_interpolation_scale=temporal_interpolation_scale,

)

if embed_dim % 4 != 0:

raise ValueError("`embed_dim` must be divisible by 4")

if isinstance(spatial_size, int):

spatial_size = (spatial_size, spatial_size)

embed_dim_spatial = 3 * embed_dim // 4

embed_dim_temporal = embed_dim // 4

# 1. Spatial

grid_h = torch.arange(spatial_size[1], device=device, dtype=torch.float32) / spatial_interpolation_scale

grid_w = torch.arange(spatial_size[0], device=device, dtype=torch.float32) / spatial_interpolation_scale

grid = torch.meshgrid(grid_w, grid_h, indexing="xy") # here w goes first

grid = torch.stack(grid, dim=0)

grid = grid.reshape([2, 1, spatial_size[1], spatial_size[0]])

pos_embed_spatial = get_2d_sincos_pos_embed_from_grid(embed_dim_spatial, grid, output_type="pt")

# 2. Temporal

grid_t = torch.arange(temporal_size, device=device, dtype=torch.float32) / temporal_interpolation_scale

pos_embed_temporal = get_1d_sincos_pos_embed_from_grid(embed_dim_temporal, grid_t, output_type="pt")

# 3. Concat

pos_embed_spatial = pos_embed_spatial[None, :, :]

pos_embed_spatial = pos_embed_spatial.repeat_interleave(

temporal_size, dim=0, output_size=pos_embed_spatial.shape[0] * temporal_size

) # [T, H*W, D // 4 * 3]

pos_embed_temporal = pos_embed_temporal[:, None, :]

pos_embed_temporal = pos_embed_temporal.repeat_interleave(

spatial_size[0] * spatial_size[1], dim=1

) # [T, H*W, D // 4]

pos_embed = torch.concat([pos_embed_temporal, pos_embed_spatial], dim=-1) # [T, H*W, D]

return pos_embed

def _get_3d_sincos_pos_embed_np(

embed_dim: int,

spatial_size: int | tuple[int, int],

temporal_size: int,

spatial_interpolation_scale: float = 1.0,

temporal_interpolation_scale: float = 1.0,

) -> np.ndarray:

r"""

Creates 3D sinusoidal positional embeddings.

Args:

embed_dim (`int`):

The embedding dimension of inputs. It must be divisible by 16.

spatial_size (`int` or `tuple[int, int]`):

The spatial dimension of positional embeddings. If an integer is provided, the same size is applied to both

spatial dimensions (height and width).

temporal_size (`int`):

The temporal dimension of positional embeddings (number of frames).

spatial_interpolation_scale (`float`, defaults to 1.0):

Scale factor for spatial grid interpolation.

temporal_interpolation_scale (`float`, defaults to 1.0):

Scale factor for temporal grid interpolation.

Returns:

`np.ndarray`:

The 3D sinusoidal positional embeddings of shape `[temporal_size, spatial_size[0] * spatial_size[1],

embed_dim]`.

"""

deprecation_message = (

"`get_3d_sincos_pos_embed` uses `torch` and supports `device`."

" `from_numpy` is no longer required."

" Pass `output_type='pt' to use the new version now."

)

deprecate("output_type=='np'", "0.33.0", deprecation_message, standard_warn=False)

if embed_dim % 4 != 0:

raise ValueError("`embed_dim` must be divisible by 4")

if isinstance(spatial_size, int):

spatial_size = (spatial_size, spatial_size)

embed_dim_spatial = 3 * embed_dim // 4

embed_dim_temporal = embed_dim // 4

# 1. Spatial

grid_h = np.arange(spatial_size[1], dtype=np.float32) / spatial_interpolation_scale

grid_w = np.arange(spatial_size[0], dtype=np.float32) / spatial_interpolation_scale

grid = np.meshgrid(grid_w, grid_h) # here w goes first

grid = np.stack(grid, axis=0)

grid = grid.reshape([2, 1, spatial_size[1], spatial_size[0]])

pos_embed_spatial = get_2d_sincos_pos_embed_from_grid(embed_dim_spatial, grid)

# 2. Temporal

grid_t = np.arange(temporal_size, dtype=np.float32) / temporal_interpolation_scale

pos_embed_temporal = get_1d_sincos_pos_embed_from_grid(embed_dim_temporal, grid_t)

# 3. Concat

pos_embed_spatial = pos_embed_spatial[np.newaxis, :, :]

pos_embed_spatial = np.repeat(pos_embed_spatial, temporal_size, axis=0) # [T, H*W, D // 4 * 3]

pos_embed_temporal = pos_embed_temporal[:, np.newaxis, :]

pos_embed_temporal = np.repeat(pos_embed_temporal, spatial_size[0] * spatial_size[1], axis=1) # [T, H*W, D // 4]

pos_embed = np.concatenate([pos_embed_temporal, pos_embed_spatial], axis=-1) # [T, H*W, D]

return pos_embed

def get_2d_sincos_pos_embed(

embed_dim,

grid_size,

cls_token=False,

extra_tokens=0,

interpolation_scale=1.0,

base_size=16,

device: torch.device | None = None,

output_type: str = "np",

):

"""

Creates 2D sinusoidal positional embeddings.

Args:

embed_dim (`int`):

The embedding dimension.

grid_size (`int`):

The size of the grid height and width.

cls_token (`bool`, defaults to `False`):

Whether or not to add a classification token.

extra_tokens (`int`, defaults to `0`):

The number of extra tokens to add.

interpolation_scale (`float`, defaults to `1.0`):

The scale of the interpolation.

Returns:

pos_embed (`torch.Tensor`):

Shape is either `[grid_size * grid_size, embed_dim]` if not using cls_token, or `[1 + grid_size*grid_size,

embed_dim]` if using cls_token

"""

if output_type == "np":

deprecation_message = (

"`get_2d_sincos_pos_embed` uses `torch` and supports `device`."

" `from_numpy` is no longer required."

" Pass `output_type='pt' to use the new version now."

)

deprecate("output_type=='np'", "0.33.0", deprecation_message, standard_warn=False)

return get_2d_sincos_pos_embed_np(

embed_dim=embed_dim,

grid_size=grid_size,

cls_token=cls_token,

extra_tokens=extra_tokens,

interpolation_scale=interpolation_scale,

base_size=base_size,

)

if isinstance(grid_size, int):

grid_size = (grid_size, grid_size)

grid_h = (

torch.arange(grid_size[0], device=device, dtype=torch.float32)

/ (grid_size[0] / base_size)

/ interpolation_scale

)

grid_w = (

torch.arange(grid_size[1], device=device, dtype=torch.float32)

/ (grid_size[1] / base_size)

/ interpolation_scale

)

grid = torch.meshgrid(grid_w, grid_h, indexing="xy") # here w goes first

grid = torch.stack(grid, dim=0)

grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])

pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid, output_type=output_type)

if cls_token and extra_tokens > 0:

pos_embed = torch.concat([torch.zeros([extra_tokens, embed_dim]), pos_embed], dim=0)

return pos_embed

def get_2d_sincos_pos_embed_from_grid(embed_dim, grid, output_type="np"):

r"""

This function generates 2D sinusoidal positional embeddings from a grid.

Args:

embed_dim (`int`): The embedding dimension.

grid (`torch.Tensor`): Grid of positions with shape `(H * W,)`.

Returns:

`torch.Tensor`: The 2D sinusoidal positional embeddings with shape `(H * W, embed_dim)`

"""

if output_type == "np":

deprecation_message = (

"`get_2d_sincos_pos_embed_from_grid` uses `torch` and supports `device`."

" `from_numpy` is no longer required."

" Pass `output_type='pt' to use the new version now."

)

deprecate("output_type=='np'", "0.33.0", deprecation_message, standard_warn=False)

return get_2d_sincos_pos_embed_from_grid_np(

embed_dim=embed_dim,

grid=grid,

)

if embed_dim % 2 != 0:

raise ValueError("embed_dim must be divisible by 2")

# use half of dimensions to encode grid_h

emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0], output_type=output_type) # (H*W, D/2)

emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1], output_type=output_type) # (H*W, D/2)

emb = torch.concat([emb_h, emb_w], dim=1) # (H*W, D)

return emb

def get_1d_sincos_pos_embed_from_grid(embed_dim, pos, output_type="np", flip_sin_to_cos=False, dtype=None):

"""

This function generates 1D positional embeddings from a grid.

Args:

embed_dim (`int`): The embedding dimension `D`

pos (`torch.Tensor`): 1D tensor of positions with shape `(M,)`

output_type (`str`, *optional*, defaults to `"np"`): Output type. Use `"pt"` for PyTorch tensors.

flip_sin_to_cos (`bool`, *optional*, defaults to `False`): Whether to flip sine and cosine embeddings.

dtype (`torch.dtype`, *optional*): Data type for frequency calculations. If `None`, defaults to

`torch.float32` on MPS devices (which don't support `torch.float64`) and `torch.float64` on other devices.

Returns:

`torch.Tensor`: Sinusoidal positional embeddings of shape `(M, D)`.

"""

if output_type == "np":

deprecation_message = (

"`get_1d_sincos_pos_embed_from_grid` uses `torch` and supports `device`."

" `from_numpy` is no longer required."

" Pass `output_type='pt' to use the new version now."

)

deprecate("output_type=='np'", "0.34.0", deprecation_message, standard_warn=False)

return get_1d_sincos_pos_embed_from_grid_np(embed_dim=embed_dim, pos=pos)

if embed_dim % 2 != 0:

raise ValueError("embed_dim must be divisible by 2")

# Auto-detect appropriate dtype if not specified

if dtype is None:

dtype = torch.float32 if pos.device.type == "mps" else torch.float64

omega = torch.arange(embed_dim // 2, device=pos.device, dtype=dtype)

omega /= embed_dim / 2.0

omega = 1.0 / 10000**omega # (D/2,)

pos = pos.reshape(-1) # (M,)

out = torch.outer(pos, omega) # (M, D/2), outer product

emb_sin = torch.sin(out) # (M, D/2)

emb_cos = torch.cos(out) # (M, D/2)

emb = torch.concat([emb_sin, emb_cos], dim=1) # (M, D)

# flip sine and cosine embeddings

if flip_sin_to_cos:

emb = torch.cat([emb[:, embed_dim // 2 :], emb[:, : embed_dim // 2]], dim=1)

return emb

def get_2d_sincos_pos_embed_np(

embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16

):

"""

Creates 2D sinusoidal positional embeddings.

Args:

embed_dim (`int`):

The embedding dimension.

grid_size (`int`):

The size of the grid height and width.

cls_token (`bool`, defaults to `False`):

Whether or not to add a classification token.

extra_tokens (`int`, defaults to `0`):

The number of extra tokens to add.

interpolation_scale (`float`, defaults to `1.0`):

The scale of the interpolation.

Returns:

pos_embed (`np.ndarray`):

Shape is either `[grid_size * grid_size, embed_dim]` if not using cls_token, or `[1 + grid_size*grid_size,

embed_dim]` if using cls_token

"""

if isinstance(grid_size, int):

grid_size = (grid_size, grid_size)

grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size) / interpolation_scale

grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size) / interpolation_scale

grid = np.meshgrid(grid_w, grid_h) # here w goes first

grid = np.stack(grid, axis=0)

grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])

pos_embed = get_2d_sincos_pos_embed_from_grid_np(embed_dim, grid)

if cls_token and extra_tokens > 0:

pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)

return pos_embed

def get_2d_sincos_pos_embed_from_grid_np(embed_dim, grid):

r"""

This function generates 2D sinusoidal positional embeddings from a grid.

Args:

embed_dim (`int`): The embedding dimension.

grid (`np.ndarray`): Grid of positions with shape `(H * W,)`.

Returns:

`np.ndarray`: The 2D sinusoidal positional embeddings with shape `(H * W, embed_dim)`

"""

if embed_dim % 2 != 0:

raise ValueError("embed_dim must be divisible by 2")

# use half of dimensions to encode grid_h

emb_h = get_1d_sincos_pos_embed_from_grid_np(embed_dim // 2, grid[0]) # (H*W, D/2)

emb_w = get_1d_sincos_pos_embed_from_grid_np(embed_dim // 2, grid[1]) # (H*W, D/2)

emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)

return emb

def get_1d_sincos_pos_embed_from_grid_np(embed_dim, pos):

"""

This function generates 1D positional embeddings from a grid.

Args:

embed_dim (`int`): The embedding dimension `D`

pos (`numpy.ndarray`): 1D tensor of positions with shape `(M,)`

Returns:

`numpy.ndarray`: Sinusoidal positional embeddings of shape `(M, D)`.

"""

if embed_dim % 2 != 0:

raise ValueError("embed_dim must be divisible by 2")

omega = np.arange(embed_dim // 2, dtype=np.float64)

omega /= embed_dim / 2.0

omega = 1.0 / 10000**omega # (D/2,)

pos = pos.reshape(-1) # (M,)

out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product

emb_sin = np.sin(out) # (M, D/2)

emb_cos = np.cos(out) # (M, D/2)

emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)

return emb

class PatchEmbed(nn.Module):

"""

2D Image to Patch Embedding with support for SD3 cropping.

Args:

height (`int`, defaults to `224`): The height of the image.

width (`int`, defaults to `224`): The width of the image.

patch_size (`int`, defaults to `16`): The size of the patches.

in_channels (`int`, defaults to `3`): The number of input channels.

embed_dim (`int`, defaults to `768`): The output dimension of the embedding.

layer_norm (`bool`, defaults to `False`): Whether or not to use layer normalization.

flatten (`bool`, defaults to `True`): Whether or not to flatten the output.

bias (`bool`, defaults to `True`): Whether or not to use bias.

interpolation_scale (`float`, defaults to `1`): The scale of the interpolation.

pos_embed_type (`str`, defaults to `"sincos"`): The type of positional embedding.

pos_embed_max_size (`int`, defaults to `None`): The maximum size of the positional embedding.

"""

def __init__(

self,

height=224,

width=224,

patch_size=16,

in_channels=3,

embed_dim=768,

layer_norm=False,

flatten=True,

bias=True,

interpolation_scale=1,

pos_embed_type="sincos",

pos_embed_max_size=None, # For SD3 cropping

):

super().__init__()

num_patches = (height // patch_size) * (width // patch_size)

self.flatten = flatten

self.layer_norm = layer_norm

self.pos_embed_max_size = pos_embed_max_size

self.proj = nn.Conv2d(

in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias

)

if layer_norm:

self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6)

else:

self.norm = None

self.patch_size = patch_size

self.height, self.width = height // patch_size, width // patch_size

self.base_size = height // patch_size

self.interpolation_scale = interpolation_scale

# Calculate positional embeddings based on max size or default

if pos_embed_max_size:

grid_size = pos_embed_max_size

else:

grid_size = int(num_patches**0.5)

if pos_embed_type is None:

self.pos_embed = None

elif pos_embed_type == "sincos":

pos_embed = get_2d_sincos_pos_embed(

embed_dim,

grid_size,

base_size=self.base_size,

interpolation_scale=self.interpolation_scale,

output_type="pt",

)

persistent = True if pos_embed_max_size else False

self.register_buffer("pos_embed", pos_embed.float().unsqueeze(0), persistent=persistent)

else:

raise ValueError(f"Unsupported pos_embed_type: {pos_embed_type}")

def cropped_pos_embed(self, height, width):

"""Crops positional embeddings for SD3 compatibility."""

if self.pos_embed_max_size is None:

raise ValueError("`pos_embed_max_size` must be set for cropping.")

height = height // self.patch_size

width = width // self.patch_size

if height > self.pos_embed_max_size:

raise ValueError(

f"Height ({height}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}."

)

if width > self.pos_embed_max_size:

raise ValueError(

f"Width ({width}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}."

)

top = (self.pos_embed_max_size - height) // 2

left = (self.pos_embed_max_size - width) // 2

spatial_pos_embed = self.pos_embed.reshape(1, self.pos_embed_max_size, self.pos_embed_max_size, -1)

spatial_pos_embed = spatial_pos_embed[:, top : top + height, left : left + width, :]

spatial_pos_embed = spatial_pos_embed.reshape(1, -1, spatial_pos_embed.shape[-1])

return spatial_pos_embed

def forward(self, latent):

if self.pos_embed_max_size is not None:

height, width = latent.shape[-2:]

else:

height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size

latent = self.proj(latent)

if self.flatten:

latent = latent.flatten(2).transpose(1, 2) # BCHW -> BNC

if self.layer_norm:

latent = self.norm(latent)

if self.pos_embed is None:

return latent.to(latent.dtype)

# Interpolate or crop positional embeddings as needed

if self.pos_embed_max_size:

pos_embed = self.cropped_pos_embed(height, width)

else:

if self.height != height or self.width != width:

pos_embed = get_2d_sincos_pos_embed(

embed_dim=self.pos_embed.shape[-1],

grid_size=(height, width),

base_size=self.base_size,

interpolation_scale=self.interpolation_scale,

device=latent.device,

output_type="pt",

)

pos_embed = pos_embed.float().unsqueeze(0)

else:

pos_embed = self.pos_embed

return (latent + pos_embed).to(latent.dtype)

class LuminaPatchEmbed(nn.Module):

"""

2D Image to Patch Embedding with support for Lumina-T2X

Args:

patch_size (`int`, defaults to `2`): The size of the patches.

in_channels (`int`, defaults to `4`): The number of input channels.

embed_dim (`int`, defaults to `768`): The output dimension of the embedding.

bias (`bool`, defaults to `True`): Whether or not to use bias.

"""

def __init__(self, patch_size=2, in_channels=4, embed_dim=768, bias=True):

super().__init__()

self.patch_size = patch_size

self.proj = nn.Linear(

in_features=patch_size * patch_size * in_channels,

out_features=embed_dim,

bias=bias,

)

def forward(self, x, freqs_cis):

"""

Patchifies and embeds the input tensor(s).

Args:

x (list[torch.Tensor] | torch.Tensor): The input tensor(s) to be patchified and embedded.

Returns:

tuple[torch.Tensor, torch.Tensor, list[tuple[int, int]], torch.Tensor]: A tuple containing the patchified

and embedded tensor(s), the mask indicating the valid patches, the original image size(s), and the

frequency tensor(s).

"""

freqs_cis = freqs_cis.to(x[0].device)

patch_height = patch_width = self.patch_size

batch_size, channel, height, width = x.size()

height_tokens, width_tokens = height // patch_height, width // patch_width

x = x.view(batch_size, channel, height_tokens, patch_height, width_tokens, patch_width).permute(

0, 2, 4, 1, 3, 5

)

x = x.flatten(3)

x = self.proj(x)

x = x.flatten(1, 2)

mask = torch.ones(x.shape[0], x.shape[1], dtype=torch.int32, device=x.device)

return (

x,

mask,

[(height, width)] * batch_size,

freqs_cis[:height_tokens, :width_tokens].flatten(0, 1).unsqueeze(0),

)

class CogVideoXPatchEmbed(nn.Module):

def __init__(

self,

patch_size: int = 2,

patch_size_t: int | None = None,

in_channels: int = 16,

embed_dim: int = 1920,

text_embed_dim: int = 4096,

bias: bool = True,

sample_width: int = 90,

sample_height: int = 60,

sample_frames: int = 49,

temporal_compression_ratio: int = 4,

max_text_seq_length: int = 226,

spatial_interpolation_scale: float = 1.875,

temporal_interpolation_scale: float = 1.0,

use_positional_embeddings: bool = True,

use_learned_positional_embeddings: bool = True,

) -> None:

super().__init__()

self.patch_size = patch_size

self.patch_size_t = patch_size_t

self.embed_dim = embed_dim

self.sample_height = sample_height

self.sample_width = sample_width

self.sample_frames = sample_frames

self.temporal_compression_ratio = temporal_compression_ratio

self.max_text_seq_length = max_text_seq_length

self.spatial_interpolation_scale = spatial_interpolation_scale

self.temporal_interpolation_scale = temporal_interpolation_scale

self.use_positional_embeddings = use_positional_embeddings

self.use_learned_positional_embeddings = use_learned_positional_embeddings

if patch_size_t is None:

# CogVideoX 1.0 checkpoints

self.proj = nn.Conv2d(

in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias

)

else:

# CogVideoX 1.5 checkpoints

self.proj = nn.Linear(in_channels * patch_size * patch_size * patch_size_t, embed_dim)

self.text_proj = nn.Linear(text_embed_dim, embed_dim)

if use_positional_embeddings or use_learned_positional_embeddings:

persistent = use_learned_positional_embeddings

pos_embedding = self._get_positional_embeddings(sample_height, sample_width, sample_frames)

self.register_buffer("pos_embedding", pos_embedding, persistent=persistent)

def _get_positional_embeddings(

self, sample_height: int, sample_width: int, sample_frames: int, device: torch.device | None = None

) -> torch.Tensor:

post_patch_height = sample_height // self.patch_size

post_patch_width = sample_width // self.patch_size

post_time_compression_frames = (sample_frames - 1) // self.temporal_compression_ratio + 1

num_patches = post_patch_height * post_patch_width * post_time_compression_frames

pos_embedding = get_3d_sincos_pos_embed(

self.embed_dim,

(post_patch_width, post_patch_height),

post_time_compression_frames,

self.spatial_interpolation_scale,

self.temporal_interpolation_scale,

device=device,

output_type="pt",

)

pos_embedding = pos_embedding.flatten(0, 1)

joint_pos_embedding = pos_embedding.new_zeros(

1, self.max_text_seq_length + num_patches, self.embed_dim, requires_grad=False

)

joint_pos_embedding.data[:, self.max_text_seq_length :].copy_(pos_embedding)

return joint_pos_embedding

def forward(self, text_embeds: torch.Tensor, image_embeds: torch.Tensor):

r"""

Args:

text_embeds (`torch.Tensor`):

Input text embeddings. Expected shape: (batch_size, seq_length, embedding_dim).

image_embeds (`torch.Tensor`):

Input image embeddings. Expected shape: (batch_size, num_frames, channels, height, width).

"""

text_embeds = self.text_proj(text_embeds)

batch_size, num_frames, channels, height, width = image_embeds.shape

if self.patch_size_t is None:

image_embeds = image_embeds.reshape(-1, channels, height, width)

image_embeds = self.proj(image_embeds)

image_embeds = image_embeds.view(batch_size, num_frames, *image_embeds.shape[1:])

image_embeds = image_embeds.flatten(3).transpose(2, 3) # [batch, num_frames, height x width, channels]

image_embeds = image_embeds.flatten(1, 2) # [batch, num_frames x height x width, channels]

else:

p = self.patch_size

p_t = self.patch_size_t

image_embeds = image_embeds.permute(0, 1, 3, 4, 2)

image_embeds = image_embeds.reshape(

batch_size, num_frames // p_t, p_t, height // p, p, width // p, p, channels

)

image_embeds = image_embeds.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(4, 7).flatten(1, 3)

image_embeds = self.proj(image_embeds)

embeds = torch.cat(

[text_embeds, image_embeds], dim=1

).contiguous() # [batch, seq_length + num_frames x height x width, channels]

if self.use_positional_embeddings or self.use_learned_positional_embeddings:

if self.use_learned_positional_embeddings and (self.sample_width != width or self.sample_height != height):

raise ValueError(

"It is currently not possible to generate videos at a different resolution that the defaults. This should only be the case with 'THUDM/CogVideoX-5b-I2V'."

"If you think this is incorrect, please open an issue at https://github.com/huggingface/diffusers/issues."

)

pre_time_compression_frames = (num_frames - 1) * self.temporal_compression_ratio + 1

if (

self.sample_height != height

or self.sample_width != width

or self.sample_frames != pre_time_compression_frames

):

pos_embedding = self._get_positional_embeddings(

height, width, pre_time_compression_frames, device=embeds.device

)

else:

pos_embedding = self.pos_embedding

pos_embedding = pos_embedding.to(dtype=embeds.dtype)

embeds = embeds + pos_embedding

return embeds

class CogView3PlusPatchEmbed(nn.Module):

def __init__(

self,

in_channels: int = 16,

hidden_size: int = 2560,

patch_size: int = 2,

text_hidden_size: int = 4096,

pos_embed_max_size: int = 128,

):

super().__init__()

self.in_channels = in_channels

self.hidden_size = hidden_size

self.patch_size = patch_size

self.text_hidden_size = text_hidden_size

self.pos_embed_max_size = pos_embed_max_size

# Linear projection for image patches

self.proj = nn.Linear(in_channels * patch_size**2, hidden_size)

# Linear projection for text embeddings

self.text_proj = nn.Linear(text_hidden_size, hidden_size)

pos_embed = get_2d_sincos_pos_embed(

hidden_size, pos_embed_max_size, base_size=pos_embed_max_size, output_type="pt"

)

pos_embed = pos_embed.reshape(pos_embed_max_size, pos_embed_max_size, hidden_size)

self.register_buffer("pos_embed", pos_embed.float(), persistent=False)

def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor) -> torch.Tensor:

batch_size, channel, height, width = hidden_states.shape

if height % self.patch_size != 0 or width % self.patch_size != 0:

raise ValueError("Height and width must be divisible by patch size")

height = height // self.patch_size

width = width // self.patch_size

hidden_states = hidden_states.view(batch_size, channel, height, self.patch_size, width, self.patch_size)

hidden_states = hidden_states.permute(0, 2, 4, 1, 3, 5).contiguous()

hidden_states = hidden_states.view(batch_size, height * width, channel * self.patch_size * self.patch_size)

# Project the patches

hidden_states = self.proj(hidden_states)

encoder_hidden_states = self.text_proj(encoder_hidden_states)

hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)

# Calculate text_length

text_length = encoder_hidden_states.shape[1]

image_pos_embed = self.pos_embed[:height, :width].reshape(height * width, -1)

text_pos_embed = torch.zeros(

(text_length, self.hidden_size), dtype=image_pos_embed.dtype, device=image_pos_embed.device

)

pos_embed = torch.cat([text_pos_embed, image_pos_embed], dim=0)[None, ...]

return (hidden_states + pos_embed).to(hidden_states.dtype)

def get_3d_rotary_pos_embed(

embed_dim,

crops_coords,

grid_size,

temporal_size,

theta: int = 10000,

use_real: bool = True,

grid_type: str = "linspace",

max_size: tuple[int, int] | None = None,

device: torch.device | None = None,

) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

"""

RoPE for video tokens with 3D structure.

Args:

embed_dim: (`int`):

The embedding dimension size, corresponding to hidden_size_head.

crops_coords (`tuple[int]`):

The top-left and bottom-right coordinates of the crop.

grid_size (`tuple[int]`):

The grid size of the spatial positional embedding (height, width).

temporal_size (`int`):

The size of the temporal dimension.

theta (`float`):

Scaling factor for frequency computation.

grid_type (`str`):

Whether to use "linspace" or "slice" to compute grids.

Returns:

`torch.Tensor`: positional embedding with shape `(temporal_size * grid_size[0] * grid_size[1], embed_dim/2)`.

"""

if use_real is not True:

raise ValueError(" `use_real = False` is not currently supported for get_3d_rotary_pos_embed")

if grid_type == "linspace":

start, stop = crops_coords

grid_size_h, grid_size_w = grid_size

grid_h = torch.linspace(

start[0], stop[0] * (grid_size_h - 1) / grid_size_h, grid_size_h, device=device, dtype=torch.float32

)

grid_w = torch.linspace(

start[1], stop[1] * (grid_size_w - 1) / grid_size_w, grid_size_w, device=device, dtype=torch.float32

)

grid_t = torch.arange(temporal_size, device=device, dtype=torch.float32)

grid_t = torch.linspace(

0, temporal_size * (temporal_size - 1) / temporal_size, temporal_size, device=device, dtype=torch.float32

)

elif grid_type == "slice":

max_h, max_w = max_size

grid_size_h, grid_size_w = grid_size

grid_h = torch.arange(max_h, device=device, dtype=torch.float32)

grid_w = torch.arange(max_w, device=device, dtype=torch.float32)

grid_t = torch.arange(temporal_size, device=device, dtype=torch.float32)

else:

raise ValueError("Invalid value passed for `grid_type`.")

# Compute dimensions for each axis

dim_t = embed_dim // 4

dim_h = embed_dim // 8 * 3

dim_w = embed_dim // 8 * 3

# Temporal frequencies

freqs_t = get_1d_rotary_pos_embed(dim_t, grid_t, theta=theta, use_real=True)

# Spatial frequencies for height and width

freqs_h = get_1d_rotary_pos_embed(dim_h, grid_h, theta=theta, use_real=True)

freqs_w = get_1d_rotary_pos_embed(dim_w, grid_w, theta=theta, use_real=True)

# BroadCast and concatenate temporal and spaial frequencie (height and width) into a 3d tensor

def combine_time_height_width(freqs_t, freqs_h, freqs_w):

freqs_t = freqs_t[:, None, None, :].expand(

-1, grid_size_h, grid_size_w, -1

) # temporal_size, grid_size_h, grid_size_w, dim_t

freqs_h = freqs_h[None, :, None, :].expand(

temporal_size, -1, grid_size_w, -1

) # temporal_size, grid_size_h, grid_size_2, dim_h

freqs_w = freqs_w[None, None, :, :].expand(

temporal_size, grid_size_h, -1, -1

) # temporal_size, grid_size_h, grid_size_2, dim_w

freqs = torch.cat(

[freqs_t, freqs_h, freqs_w], dim=-1

) # temporal_size, grid_size_h, grid_size_w, (dim_t + dim_h + dim_w)

freqs = freqs.view(

temporal_size * grid_size_h * grid_size_w, -1

) # (temporal_size * grid_size_h * grid_size_w), (dim_t + dim_h + dim_w)

return freqs

t_cos, t_sin = freqs_t # both t_cos and t_sin has shape: temporal_size, dim_t

h_cos, h_sin = freqs_h # both h_cos and h_sin has shape: grid_size_h, dim_h

w_cos, w_sin = freqs_w # both w_cos and w_sin has shape: grid_size_w, dim_w

if grid_type == "slice":

t_cos, t_sin = t_cos[:temporal_size], t_sin[:temporal_size]

h_cos, h_sin = h_cos[:grid_size_h], h_sin[:grid_size_h]

w_cos, w_sin = w_cos[:grid_size_w], w_sin[:grid_size_w]

cos = combine_time_height_width(t_cos, h_cos, w_cos)

sin = combine_time_height_width(t_sin, h_sin, w_sin)

return cos, sin

def get_3d_rotary_pos_embed_allegro(

embed_dim,

crops_coords,

grid_size,

temporal_size,

interpolation_scale: tuple[float, float, float] = (1.0, 1.0, 1.0),

theta: int = 10000,

device: torch.device | None = None,

) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

# TODO(aryan): docs

start, stop = crops_coords

grid_size_h, grid_size_w = grid_size

interpolation_scale_t, interpolation_scale_h, interpolation_scale_w = interpolation_scale

grid_t = torch.linspace(

0, temporal_size * (temporal_size - 1) / temporal_size, temporal_size, device=device, dtype=torch.float32

)

grid_h = torch.linspace(

start[0], stop[0] * (grid_size_h - 1) / grid_size_h, grid_size_h, device=device, dtype=torch.float32

)

grid_w = torch.linspace(

start[1], stop[1] * (grid_size_w - 1) / grid_size_w, grid_size_w, device=device, dtype=torch.float32

)

# Compute dimensions for each axis

dim_t = embed_dim // 3

dim_h = embed_dim // 3

dim_w = embed_dim // 3

# Temporal frequencies

freqs_t = get_1d_rotary_pos_embed(

dim_t, grid_t / interpolation_scale_t, theta=theta, use_real=True, repeat_interleave_real=False

)

# Spatial frequencies for height and width

freqs_h = get_1d_rotary_pos_embed(

dim_h, grid_h / interpolation_scale_h, theta=theta, use_real=True, repeat_interleave_real=False

)

freqs_w = get_1d_rotary_pos_embed(

dim_w, grid_w / interpolation_scale_w, theta=theta, use_real=True, repeat_interleave_real=False

)

return freqs_t, freqs_h, freqs_w, grid_t, grid_h, grid_w

def get_2d_rotary_pos_embed(

embed_dim, crops_coords, grid_size, use_real=True, device: torch.device | None = None, output_type: str = "np"

):

"""

RoPE for image tokens with 2d structure.

Args:

embed_dim: (`int`):

The embedding dimension size

crops_coords (`tuple[int]`)

The top-left and bottom-right coordinates of the crop.

grid_size (`tuple[int]`):

The grid size of the positional embedding.

use_real (`bool`):

If True, return real part and imaginary part separately. Otherwise, return complex numbers.

device: (`torch.device`, **optional**):

The device used to create tensors.

Returns:

`torch.Tensor`: positional embedding with shape `( grid_size * grid_size, embed_dim/2)`.

"""

if output_type == "np":

deprecation_message = (

"`get_2d_sincos_pos_embed` uses `torch` and supports `device`."

" `from_numpy` is no longer required."

" Pass `output_type='pt' to use the new version now."

)