timesteps: torch.Tensor,

embedding_dim: int,

flip_sin_to_cos: bool = False,

downscale_freq_shift: float = 1,

scale: float = 1,

max_period: int = 10000,

):

"""

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))

embed_dim: int,

spatial_size: Union[int, Tuple[int, int]],

temporal_size: int,

spatial_interpolation_scale: float = 1.0,

temporal_interpolation_scale: float = 1.0,

) -> np.ndarray:

r"""

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]

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

):

"""

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(embed_dim, grid)

if cls_token and extra_tokens > 0:

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

r"""

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]) # (H*W, D/2)

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

emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, 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)

"""

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

)

persistent = True if pos_embed_max_size else False

self.register_buffer("pos_embed", torch.from_numpy(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,

)

pos_embed = torch.from_numpy(pos_embed).float().unsqueeze(0).to(latent.device)

else:

pos_embed = self.pos_embed

"""

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):

"""

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),

def __init__(

self,

patch_size: int = 2,

patch_size_t: Optional[int] = 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) -> 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,

)

pos_embedding = torch.from_numpy(pos_embedding).flatten(0, 1)

joint_pos_embedding = torch.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"""

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)

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

else:

pos_embedding = self.pos_embedding

embeds = embeds + pos_embedding

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)

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

self.register_buffer("pos_embed", torch.from_numpy(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, ...]

embed_dim,

crops_coords,

grid_size,

temporal_size,

theta: int = 10000,

use_real: bool = True,

grid_type: str = "linspace",

max_size: Optional[Tuple[int, int]] = None,

device: Optional[torch.device] = None,

) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:

"""

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)

embed_dim,

crops_coords,

grid_size,

temporal_size,

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

theta: int = 10000,

device: Optional[torch.device] = None,

) -> Union[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

)

"""

start, stop = crops_coords

grid_h = np.linspace(start[0], stop[0], grid_size[0], endpoint=False, dtype=np.float32)

grid_w = np.linspace(start[1], stop[1], grid_size[1], endpoint=False, dtype=np.float32)

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

grid = np.stack(grid, axis=0) # [2, W, H]

grid = grid.reshape([2, 1, *grid.shape[1:]])

pos_embed = get_2d_rotary_pos_embed_from_grid(embed_dim, grid, use_real=use_real)

"""

assert embed_dim % 4 == 0

# use half of dimensions to encode grid_h

emb_h = get_1d_rotary_pos_embed(

embed_dim // 2, grid[0].reshape(-1), use_real=use_real

) # (H*W, D/2) if use_real else (H*W, D/4)

emb_w = get_1d_rotary_pos_embed(

embed_dim // 2, grid[1].reshape(-1), use_real=use_real

) # (H*W, D/2) if use_real else (H*W, D/4)

if use_real:

cos = torch.cat([emb_h[0], emb_w[0]], dim=1) # (H*W, D)

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

return cos, sin

else:

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

"""

assert embed_dim % 4 == 0

emb_h = get_1d_rotary_pos_embed(

embed_dim // 2, len_h, linear_factor=linear_factor, ntk_factor=ntk_factor

) # (H, D/4)

emb_w = get_1d_rotary_pos_embed(

embed_dim // 2, len_w, linear_factor=linear_factor, ntk_factor=ntk_factor

) # (W, D/4)

emb_h = emb_h.view(len_h, 1, embed_dim // 4, 1).repeat(1, len_w, 1, 1) # (H, W, D/4, 1)

emb_w = emb_w.view(1, len_w, embed_dim // 4, 1).repeat(len_h, 1, 1, 1) # (H, W, D/4, 1)

emb = torch.cat([emb_h, emb_w], dim=-1).flatten(2) # (H, W, D/2)

dim: int,

pos: Union[np.ndarray, int],

theta: float = 10000.0,

use_real=False,

linear_factor=1.0,

ntk_factor=1.0,

repeat_interleave_real=True,

freqs_dtype=torch.float32, # torch.float32, torch.float64 (flux)

):

"""

assert dim % 2 == 0

if isinstance(pos, int):

pos = torch.arange(pos)

if isinstance(pos, np.ndarray):

pos = torch.from_numpy(pos) # type: ignore # [S]

theta = theta * ntk_factor

freqs = (

1.0

/ (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype, device=pos.device)[: (dim // 2)] / dim))

/ linear_factor

) # [D/2]

freqs = torch.outer(pos, freqs) # type: ignore # [S, D/2]

if use_real and repeat_interleave_real:

# flux, hunyuan-dit, cogvideox

freqs_cos = freqs.cos().repeat_interleave(2, dim=1).float() # [S, D]

freqs_sin = freqs.sin().repeat_interleave(2, dim=1).float() # [S, D]

return freqs_cos, freqs_sin

elif use_real:

# stable audio, allegro

freqs_cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).float() # [S, D]

freqs_sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).float() # [S, D]

return freqs_cos, freqs_sin

else:

# lumina

freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 # [S, D/2]

x: torch.Tensor,

freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],

use_real: bool = True,

use_real_unbind_dim: int = -1,

) -> Tuple[torch.Tensor, torch.Tensor]:

"""

if use_real:

cos, sin = freqs_cis # [S, D]

cos = cos[None, None]

sin = sin[None, None]

cos, sin = cos.to(x.device), sin.to(x.device)

if use_real_unbind_dim == -1:

# Used for flux, cogvideox, hunyuan-dit

x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2]

x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)

elif use_real_unbind_dim == -2:

# Used for Stable Audio

x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) # [B, S, H, D//2]

x_rotated = torch.cat([-x_imag, x_real], dim=-1)

else:

raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")

out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)

return out

else:

# used for lumina

x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))

freqs_cis = freqs_cis.unsqueeze(2)

x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)

# TODO(aryan): rewrite

def apply_1d_rope(tokens, pos, cos, sin):

cos = F.embedding(pos, cos)[:, None, :, :]

sin = F.embedding(pos, sin)[:, None, :, :]

x1, x2 = tokens[..., : tokens.shape[-1] // 2], tokens[..., tokens.shape[-1] // 2 :]

tokens_rotated = torch.cat((-x2, x1), dim=-1)

return (tokens.float() * cos + tokens_rotated.float() * sin).to(tokens.dtype)

(t_cos, t_sin), (h_cos, h_sin), (w_cos, w_sin) = freqs_cis

t, h, w = x.chunk(3, dim=-1)

t = apply_1d_rope(t, positions[0], t_cos, t_sin)

h = apply_1d_rope(h, positions[1], h_cos, h_sin)

w = apply_1d_rope(w, positions[2], w_cos, w_sin)

x = torch.cat([t, h, w], dim=-1)

# modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11

def __init__(self, theta: int, axes_dim: List[int]):

super().__init__()

self.theta = theta

self.axes_dim = axes_dim

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

n_axes = ids.shape[-1]

cos_out = []

sin_out = []

pos = ids.float()

is_mps = ids.device.type == "mps"

freqs_dtype = torch.float32 if is_mps else torch.float64

for i in range(n_axes):

cos, sin = get_1d_rotary_pos_embed(

self.axes_dim[i],

pos[:, i],

theta=self.theta,

repeat_interleave_real=True,

use_real=True,

freqs_dtype=freqs_dtype,

)

cos_out.append(cos)

sin_out.append(sin)

freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)

freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)

def __init__(

self,

in_channels: int,

time_embed_dim: int,

act_fn: str = "silu",

out_dim: int = None,

post_act_fn: Optional[str] = None,

cond_proj_dim=None,

sample_proj_bias=True,

):