Cosmos

[a-r-r-o-w]

The cosmos is within us. We are made of star-stuff. We are a way for the universe to know itself.

Models

• https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-7B-Text2World
• https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-7B-Video2World
• https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-14B-Text2World
• https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-14B-Video2World
• https://huggingface.co/nvidia/Cosmos-1.0-Tokenizer-CV8x8x82

Transformer

test attention

fromtypingimportOptionalfromeinopsimportrearrangeimporttorchimporttorch.nnasnnclassRMSNorm(torch.nn.Module):
def__init__(
self, dim: int, elementwise_affine: bool=False, eps: float=1e-6, device=None, dtype=None
):
super().__init__()
self.eps=epsself.learnable_scale=elementwise_affineifself.learnable_scale:
self.weight=nn.Parameter(torch.empty(dim, device=device, dtype=dtype))
else:
self.register_parameter("weight", None)
defforward(self, x):
r=x*torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) +self.eps)
ifself.weightisNone:
returnrelse:
returnr*self.weight.to(dtype=x.dtype, device=x.device)
defget_normalization(name: str, channels: int):
ifname=="I":
returnnn.Identity()
elifname=="R":
# return te.pytorch.RMSNorm(channels, eps=1e-6)returnRMSNorm(channels, eps=1e-6)
else:
raiseValueError(f"Normalization {name} not found")
classAttention(nn.Module):
def__init__(
self,
query_dim: int,
context_dim=None,
heads=8,
dim_head=64,
dropout=0.0,
qkv_bias: bool=False,
out_bias: bool=False,
qkv_norm: str="SSI",
qkv_norm_mode: str="per_head",
backend: str="transformer_engine",
qkv_format: str="bshd",
) ->None:
super().__init__()
self.is_selfattn=context_dimisNone# self attentioninner_dim=dim_head*headscontext_dim=query_dimifcontext_dimisNoneelsecontext_dimself.heads=headsself.dim_head=dim_headself.qkv_norm_mode=qkv_norm_modeself.qkv_format=qkv_formatifself.qkv_norm_mode=="per_head":
norm_dim=dim_headelse:
raiseValueError(f"Normalization mode {self.qkv_norm_mode} not found, only support 'per_head'")
self.backend=backendself.to_q=nn.Sequential(
nn.Linear(query_dim, inner_dim, bias=qkv_bias),
get_normalization(qkv_norm[0], norm_dim),
)
self.to_k=nn.Sequential(
nn.Linear(context_dim, inner_dim, bias=qkv_bias),
get_normalization(qkv_norm[1], norm_dim),
)
self.to_v=nn.Sequential(
nn.Linear(context_dim, inner_dim, bias=qkv_bias),
get_normalization(qkv_norm[2], norm_dim),
)
self.to_out=nn.Sequential(
nn.Linear(inner_dim, query_dim, bias=out_bias),
nn.Dropout(dropout),
)
defcal_qkv(
self, x, context=None, mask=None, rope_emb=None, **kwargs
) ->tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
q=self.to_q[0](x)
context=xifcontextisNoneelsecontextk=self.to_k[0](context)
v=self.to_v[0](context)
q, k, v=map(
# lambda t: rearrange(t, "b ... (n c) -> b ... n c", n=self.heads, c=self.dim_head),lambdat: rearrange(t, "s b (n c) -> b n s c", n=self.heads, c=self.dim_head),
(q, k, v),
)
q=self.to_q[1](q)
k=self.to_k[1](k)
v=self.to_v[1](v)
ifself.is_selfattnandrope_embisnotNone: # only apply to self-attention!print("here")
# q = apply_rotary_pos_emb(q, rope_emb, tensor_format=self.qkv_format, fused=True)# k = apply_rotary_pos_emb(k, rope_emb, tensor_format=self.qkv_format, fused=True)# apply_rotary_pos_emb inlinedq_shape=q.shapeq=q.reshape(*q.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2)
q=torch.cat([rope_emb[..., 0] *q[..., 0], rope_emb[..., 1] *q[..., 1]], dim=-1)
# q = rope_emb[..., 0] * q[..., 0] + rope_emb[..., 1] * q[..., 1]q=q.movedim(-1, -2).reshape(*q_shape).to(x.dtype)
# apply_rotary_pos_emb inlinedk_shape=k.shapek=k.reshape(*k.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2)
k=torch.cat([rope_emb[..., 0] *k[..., 0], rope_emb[..., 1] *k[..., 1]], dim=-1)
# k = rope_emb[..., 0] * k[..., 0] + rope_emb[..., 1] * k[..., 1]k=k.movedim(-1, -2).reshape(*k_shape).to(x.dtype)
returnq, k, vdefcal_attn(self, q, k, v, mask=None):
out=torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask)
out=rearrange(out, "b n s c -> s b (n c)")
out=self.to_out(out)
returnoutdefforward(
self,
x,
context=None,
mask=None,
rope_emb=None,
**kwargs,
):
""" Args: x (Tensor): The query tensor of shape [B, Mq, K] context (Optional[Tensor]): The key tensor of shape [B, Mk, K] or use x as context [self attention] if None """q, k, v=self.cal_qkv(x, context, mask, rope_emb=rope_emb, **kwargs)
returnself.cal_attn(q, k, v, mask)
@torch.no_grad()defmatch_rms_norm():
fromdiffusers.models.normalizationimportRMSNormasDiffusersRMSNormtheirs_rmsnorm=RMSNorm(128, elementwise_affine=True, eps=1e-6)
ours_rmsnorm=DiffusersRMSNorm(128, eps=1e-6, elementwise_affine=True)
ours_rmsnorm.weight.data.copy_(theirs_rmsnorm.weight.data)
input=torch.randn(1, 128)
theirs_output=theirs_rmsnorm(input)
ours_output=ours_rmsnorm(input)
print(sum(p.numel() forpintheirs_rmsnorm.parameters()))
print(sum(p.numel() forpinours_rmsnorm.parameters()))
print(torch.allclose(theirs_output, ours_output))
@torch.no_grad()defmatch_attention():
fromdiffusers.models.attentionimportAttentionasDiffusersAttentiontheirs_attention=Attention(128, 128, heads=8, dim_head=16, qkv_bias=False, out_bias=False, qkv_norm="RRI")
ours_attention=DiffusersAttention(128, 128, heads=8, dim_head=16, qk_norm="rms_norm", out_bias=False, elementwise_affine=False)
ours_attention.to_q.weight.data.copy_(theirs_attention.to_q[0].weight.data)
ours_attention.to_k.weight.data.copy_(theirs_attention.to_k[0].weight.data)
ours_attention.to_v.weight.data.copy_(theirs_attention.to_v[0].weight.data)
ours_attention.to_out[0].weight.data.copy_(theirs_attention.to_out[0].weight.data)
input=torch.randn(1, 42, 128)
theirs_output=rearrange(theirs_attention(rearrange(input, "b s c -> s b c")), "s b c -> b s c")
ours_output=ours_attention(input)
print(sum(p.numel() forpintheirs_attention.parameters()))
print(sum(p.numel() forpinours_attention.parameters()))
print(torch.allclose(theirs_output, ours_output, atol=1e-3))
match_rms_norm()
match_attention()

test ff

fromtypingimportOptionalimporttorchimporttorch.nnasnnfromtorch.utils.checkpointimportcheckpointclassFeedForward(nn.Module):
def__init__(
self,
d_model: int,
d_ff: int,
dropout: float=0.1,
activation=nn.ReLU(),
is_gated: bool=False,
bias: bool=False,
) ->None:
super().__init__()
self.layer1=nn.Linear(d_model, d_ff, bias=bias)
self.layer2=nn.Linear(d_ff, d_model, bias=bias)
self.dropout=nn.Dropout(dropout)
self.activation=activationself.is_gated=is_gatedifis_gated:
self.linear_gate=nn.Linear(d_model, d_ff, bias=False)
defforward(self, x: torch.Tensor):
g=self.activation(self.layer1(x))
ifself.is_gated:
x=g*self.linear_gate(x)
else:
x=gassertself.dropout.p==0.0, "we skip dropout"returnself.layer2(x)
classGPT2FeedForward(FeedForward):
def__init__(self, d_model: int, d_ff: int, dropout: float=0.1, bias: bool=False):
super().__init__(
d_model=d_model,
d_ff=d_ff,
dropout=dropout,
activation=nn.GELU(),
is_gated=False,
bias=bias,
)
defforward(self, x: torch.Tensor):
assertself.dropout.p==0.0, "we skip dropout"x=self.layer1(x)
defactivation_layer2_forward(x):
x=self.activation(x)
x=self.layer2(x)
returnxx=checkpoint(activation_layer2_forward, x, use_reentrant=False)
returnx@torch.no_grad()defmatch_ff():
fromdiffusers.models.attentionimportFeedForwardasDiffusersFeedForwardtheirs_ff=FeedForward(128, 512, 0.0, activation=nn.GELU(), is_gated=True, bias=False)
ours_ff=DiffusersFeedForward(128, mult=4, dropout=0.0, activation_fn="geglu", bias=False)
ours_ff.net[0].proj.weight.data[:512, :].copy_(theirs_ff.linear_gate.weight.data)
ours_ff.net[0].proj.weight.data[512:, :].copy_(theirs_ff.layer1.weight.data)
ours_ff.net[2].weight.data.copy_(theirs_ff.layer2.weight.data)
input=torch.randn(1, 128)
theirs_output=theirs_ff(input)
ours_output=ours_ff(input)
print(sum(p.numel() forpintheirs_ff.parameters()))
print(sum(p.numel() forpinours_ff.parameters()))
print(torch.allclose(theirs_output, ours_output))
match_ff()

test timesteps

importitertoolsimportmathimporttorchimporttorch.nnasnnclassTimesteps(nn.Module):
def__init__(self, num_channels):
super().__init__()
self.num_channels=num_channelsdefforward(self, timesteps):
in_dype=timesteps.dtypehalf_dim=self.num_channels//2exponent=-math.log(10000) *torch.arange(half_dim, dtype=torch.float32, device=timesteps.device)
exponent=exponent/ (half_dim-0.0)
emb=torch.exp(exponent)
emb=timesteps[:, None].float() *emb[None, :]
sin_emb=torch.sin(emb)
cos_emb=torch.cos(emb)
emb=torch.cat([cos_emb, sin_emb], dim=-1)
# return emb.to(in_dype)returnembclassTimestepEmbedding(nn.Module):
def__init__(self, in_features: int, out_features: int, use_adaln_lora: bool=False):
super().__init__()
self.linear_1=nn.Linear(in_features, out_features, bias=notuse_adaln_lora)
self.activation=nn.SiLU()
self.use_adaln_lora=use_adaln_loraifuse_adaln_lora:
self.linear_2=nn.Linear(out_features, 3*out_features, bias=False)
else:
self.linear_2=nn.Linear(out_features, out_features, bias=True)
defforward(self, sample: torch.Tensor) ->torch.Tensor:
sample=sample.to(self.linear_1.weight.dtype)
emb=self.linear_1(sample)
emb=self.activation(emb)
emb=self.linear_2(emb)
ifself.use_adaln_lora:
emb_B_D=sampleadaln_lora_B_3D=embelse:
emb_B_D=embadaln_lora_B_3D=Nonereturnemb_B_D, adaln_lora_B_3DclassCosmosTimestepEmbedding(nn.Module):
def__init__(self, in_features: int, out_features: int) ->None:
super().__init__()
self.linear_1=nn.Linear(in_features, out_features, bias=False)
self.activation=nn.SiLU()
self.linear_2=nn.Linear(out_features, 3*out_features, bias=False)
defforward(self, hidden_states: torch.Tensor) ->torch.Tensor:
emb=self.linear_1(hidden_states)
emb=self.activation(emb)
emb=self.linear_2(emb)
returnhidden_states, emb@torch.no_grad()defmatch_timestep():
fromdiffusers.models.embeddingsimportTimestepsasDiffusersTimestepstheirs_timesteps=Timesteps(256)
ours_timesteps=DiffusersTimesteps(256, flip_sin_to_cos=True, downscale_freq_shift=0.0)
input=torch.tensor([1000.0], dtype=torch.float32)
theirs_output=theirs_timesteps(input)
ours_output=ours_timesteps(input)
print(torch.allclose(theirs_output, ours_output))
@torch.no_grad()defmatch_timestep_embedding():
theirs_temb=TimestepEmbedding(256, 256, use_adaln_lora=True)
ours_temb=CosmosTimestepEmbedding(256, 256)
ours_temb.linear_1.weight.data.copy_(theirs_temb.linear_1.weight.data)
ours_temb.linear_2.weight.data.copy_(theirs_temb.linear_2.weight.data)
input=torch.randn(1, 256)
theirs_output=theirs_temb(input)
ours_output=ours_temb(input)
print(sum(p.numel() forpintheirs_temb.parameters()))
print(sum(p.numel() forpinours_temb.parameters()))
print(torch.allclose(theirs_output[0], ours_output[0]))
print(torch.allclose(theirs_output[1], ours_output[1]))
@torch.no_grad()defmatch_timestep_embedding_2():
fromdiffusers.models.transformers.transformer_cosmosimportCosmosTimestepEmbeddingtheirs_temb=TimestepEmbedding(256, 256, use_adaln_lora=True)
ours_temb=CosmosTimestepEmbedding(256, 256)
ours_temb.linear_1.weight.data.copy_(theirs_temb.linear_1.weight.data)
ours_temb.linear_2.weight.data.copy_(theirs_temb.linear_2.weight.data)
input=torch.randn(1, 256)
theirs_output=theirs_temb(input)
ours_output=ours_temb(input)
print(sum(p.numel() forpintheirs_temb.parameters()))
print(sum(p.numel() forpinours_temb.parameters()))
print(torch.allclose(theirs_output[1], ours_output))
@torch.no_grad()defmatch_timestep_prepare_embedding():
fromdiffusers.models.transformers.transformer_cosmosimportCosmosEmbeddingfromdiffusers.models.normalizationimportRMSNormtheirs_t_embedder=nn.Sequential(
Timesteps(4096),
TimestepEmbedding(4096, 4096, use_adaln_lora=True),
)
theirs_norm=RMSNorm(4096, 1e-6, True)
ours_t_embedder=CosmosEmbedding(4096, 4096)
ours_t_embedder.t_embedder.linear_1.weight.data.copy_(theirs_t_embedder[1].linear_1.weight.data)
ours_t_embedder.t_embedder.linear_2.weight.data.copy_(theirs_t_embedder[1].linear_2.weight.data)
ours_t_embedder.norm.weight.data.copy_(theirs_norm.weight.data)
hidden_states=torch.randn(1, 1, 4096)
input=torch.randint(0, 1000, (1,)).long()
theirs_output=theirs_t_embedder(input)
ours_output=ours_t_embedder(hidden_states, input)
print(sum(p.numel() forpinitertools.chain(theirs_t_embedder.parameters(), theirs_norm.parameters())))
print(sum(p.numel() forpinours_t_embedder.parameters()))
print(torch.allclose(theirs_output[1], ours_output[0]))
print(torch.allclose(theirs_norm(theirs_output[0]), ours_output[1]))
match_timestep()
print()
match_timestep_embedding()
print()
match_timestep_embedding_2()
print()
match_timestep_prepare_embedding()
print()

test patch embed

importtorchimporttorch.nnasnnfromeinops.layers.torchimportRearrangeclassPatchEmbed(nn.Module):
def__init__(
self,
spatial_patch_size,
temporal_patch_size,
in_channels=3,
out_channels=768,
bias=True,
):
super().__init__()
self.spatial_patch_size=spatial_patch_sizeself.temporal_patch_size=temporal_patch_sizeself.proj=nn.Sequential(
Rearrange(
"b c (t r) (h m) (w n) -> b t h w (c r m n)",
r=temporal_patch_size,
m=spatial_patch_size,
n=spatial_patch_size,
),
nn.Linear(
in_channels*spatial_patch_size*spatial_patch_size*temporal_patch_size, out_channels, bias=bias
),
)
self.out=nn.Identity()
defforward(self, x):
assertx.dim() ==5_, _, T, H, W=x.shapeassertH%self.spatial_patch_size==0andW%self.spatial_patch_size==0assertT%self.temporal_patch_size==0x=self.proj(x)
returnself.out(x)
@torch.no_grad()defmatch_patch_embed():
fromdiffusers.models.transformers.transformer_cosmosimportCosmosPatchEmbedtheirs_patch_embed=PatchEmbed(2, 1, 16, 4096, bias=False)
ours_patch_embed=CosmosPatchEmbed(16, 4096, (1, 2, 2), bias=False)
ours_patch_embed.proj.weight.data.copy_(theirs_patch_embed.proj[1].weight.data)
input=torch.randn(1, 16, 128, 240, 240)
theirs_output=theirs_patch_embed(input)
ours_output=ours_patch_embed(input)
print(torch.allclose(theirs_output, ours_output))
match_patch_embed()

test positional embed

importmathfromtypingimportOptional, Listimportnumpyasnpimporttorchfromeinopsimportrearrange, repeatdefnormalize(x: torch.Tensor, dim: Optional[List[int]] =None, eps: float=0) ->torch.Tensor:
ifdimisNone:
dim=list(range(1, x.ndim))
norm=torch.linalg.vector_norm(x, dim=dim, keepdim=True, dtype=torch.float32)
norm=torch.add(eps, norm, alpha=np.sqrt(norm.numel() /x.numel()))
returnx/norm.to(x.dtype)
def_no_grad_trunc_normal_(tensor, mean, std, a, b):
# Cut & paste from PyTorch official master until it's in a few official releases - RW# Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdfdefnorm_cdf(x):
# Computes standard normal cumulative distribution functionreturn (1.0+math.erf(x/math.sqrt(2.0))) /2.0withtorch.no_grad():
# Values are generated by using a truncated uniform distribution and# then using the inverse CDF for the normal distribution.# Get upper and lower cdf valuesl=norm_cdf((a-mean) /std)
u=norm_cdf((b-mean) /std)
# Uniformly fill tensor with values from [l, u], then translate to# [2l-1, 2u-1].tensor.uniform_(2*l-1, 2*u-1)
# Use inverse cdf transform for normal distribution to get truncated# standard normaltensor.erfinv_()
# Transform to proper mean, stdtensor.mul_(std*math.sqrt(2.0))
tensor.add_(mean)
# Clamp to ensure it's in the proper rangetensor.clamp_(min=a, max=b)
returntensordeftrunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
return_no_grad_trunc_normal_(tensor, mean, std, a, b)
classVideoPositionEmb(torch.nn.Module):
defforward(self, x_B_T_H_W_C: torch.Tensor, fps=Optional[torch.Tensor]) ->torch.Tensor:
""" It delegates the embedding generation to generate_embeddings function. """B_T_H_W_C=x_B_T_H_W_C.shapeembeddings=self.generate_embeddings(B_T_H_W_C, fps=fps)
returnembeddingsdefgenerate_embeddings(self, B_T_H_W_C: torch.Size, fps=Optional[torch.Tensor]):
raiseNotImplementedErrorclassVideoRopePosition3DEmb(VideoPositionEmb):
def__init__(
self,
*, # enforce keyword argumentshead_dim: int,
len_h: int,
len_w: int,
len_t: int,
base_fps: int=24,
h_extrapolation_ratio: float=1.0,
w_extrapolation_ratio: float=1.0,
t_extrapolation_ratio: float=1.0,
**kwargs, # used for compatibility with other positional embeddings; unused in this class
):
delkwargssuper().__init__()
self.register_buffer("seq", torch.arange(max(len_h, len_w, len_t), dtype=torch.float))
self.base_fps=base_fpsself.max_h=len_hself.max_w=len_wdim=head_dimdim_h=dim//6*2dim_w=dim_hdim_t=dim-2*dim_hassertdim==dim_h+dim_w+dim_t, f"bad dim: {dim} != {dim_h} + {dim_w} + {dim_t}"self.register_buffer(
"dim_spatial_range",
# torch.arange(0, dim_h, 2)[: (dim_h // 2)].float().cuda() / dim_h,torch.arange(0, dim_h, 2)[: (dim_h//2)].float() /dim_h,
persistent=False,
)
self.register_buffer(
"dim_temporal_range",
# torch.arange(0, dim_t, 2)[: (dim_t // 2)].float().cuda() / dim_t,torch.arange(0, dim_t, 2)[: (dim_t//2)].float() /dim_t,
persistent=False,
)
self.h_ntk_factor=h_extrapolation_ratio** (dim_h/ (dim_h-2))
self.w_ntk_factor=w_extrapolation_ratio** (dim_w/ (dim_w-2))
self.t_ntk_factor=t_extrapolation_ratio** (dim_t/ (dim_t-2))
defgenerate_embeddings(
self,
B_T_H_W_C: torch.Size,
fps: Optional[torch.Tensor] =None,
h_ntk_factor: Optional[float] =None,
w_ntk_factor: Optional[float] =None,
t_ntk_factor: Optional[float] =None,
):
""" Generate embeddings for the given input size. Args: B_T_H_W_C (torch.Size): Input tensor size (Batch, Time, Height, Width, Channels). fps (Optional[torch.Tensor], optional): Frames per second. Defaults to None. h_ntk_factor (Optional[float], optional): Height NTK factor. If None, uses self.h_ntk_factor. w_ntk_factor (Optional[float], optional): Width NTK factor. If None, uses self.w_ntk_factor. t_ntk_factor (Optional[float], optional): Time NTK factor. If None, uses self.t_ntk_factor. Returns: Not specified in the original code snippet. """h_ntk_factor=h_ntk_factorifh_ntk_factorisnotNoneelseself.h_ntk_factorw_ntk_factor=w_ntk_factorifw_ntk_factorisnotNoneelseself.w_ntk_factort_ntk_factor=t_ntk_factorift_ntk_factorisnotNoneelseself.t_ntk_factorh_theta=10000.0*h_ntk_factorw_theta=10000.0*w_ntk_factort_theta=10000.0*t_ntk_factorh_spatial_freqs=1.0/ (h_theta**self.dim_spatial_range)
w_spatial_freqs=1.0/ (w_theta**self.dim_spatial_range)
temporal_freqs=1.0/ (t_theta**self.dim_temporal_range)
B, T, H, W, _=B_T_H_W_Cuniform_fps= (fpsisNone) or (fps.min() ==fps.max())
assert (
uniform_fpsorB==1orT==1
), "For video batch, batch size should be 1 for non-uniform fps. For image batch, T should be 1"assert (
H<=self.max_handW<=self.max_w
), f"Input dimensions (H={H}, W={W}) exceed the maximum dimensions (max_h={self.max_h}, max_w={self.max_w})"half_emb_h=torch.outer(self.seq[:H], h_spatial_freqs)
half_emb_w=torch.outer(self.seq[:W], w_spatial_freqs)
# apply sequence scaling in temporal dimensioniffpsisNone: # image caseassertT==1, "T should be 1 for image batch."half_emb_t=torch.outer(self.seq[:T], temporal_freqs)
else:
half_emb_t=torch.outer(self.seq[:T] /fps[:1] *self.base_fps, temporal_freqs)
em_T_H_W_D=torch.cat(
[
repeat(half_emb_t, "t d -> t h w d", h=H, w=W),
repeat(half_emb_h, "h d -> t h w d", t=T, w=W),
repeat(half_emb_w, "w d -> t h w d", t=T, h=H),
]
*2,
dim=-1,
)
returnrearrange(em_T_H_W_D, "t h w d -> (t h w) 1 1 d").float()
classLearnablePosEmbAxis(VideoPositionEmb):
def__init__(
self,
*, # enforce keyword argumentsinterpolation: str,
model_channels: int,
len_h: int,
len_w: int,
len_t: int,
**kwargs,
):
""" Args: interpolation (str): we curretly only support "crop", ideally when we need extrapolation capacity, we should adjust frequency or other more advanced methods. they are not implemented yet. """delkwargs# unusedsuper().__init__()
self.interpolation=interpolationassertself.interpolationin ["crop"], f"Unknown interpolation method {self.interpolation}"self.pos_emb_h=torch.nn.Parameter(torch.zeros(len_h, model_channels))
self.pos_emb_w=torch.nn.Parameter(torch.zeros(len_w, model_channels))
self.pos_emb_t=torch.nn.Parameter(torch.zeros(len_t, model_channels))
trunc_normal_(self.pos_emb_h, std=0.02)
trunc_normal_(self.pos_emb_w, std=0.02)
trunc_normal_(self.pos_emb_t, std=0.02)
defgenerate_embeddings(self, B_T_H_W_C: torch.Size, fps=Optional[torch.Tensor]) ->torch.Tensor:
B, T, H, W, _=B_T_H_W_Cifself.interpolation=="crop":
emb_h_H=self.pos_emb_h[:H]
emb_w_W=self.pos_emb_w[:W]
emb_t_T=self.pos_emb_t[:T]
emb= (
repeat(emb_t_T, "t d-> b t h w d", b=B, h=H, w=W)
+repeat(emb_h_H, "h d-> b t h w d", b=B, t=T, w=W)
+repeat(emb_w_W, "w d-> b t h w d", b=B, t=T, h=H)
)
assertlist(emb.shape)[:4] == [B, T, H, W], f"bad shape: {list(emb.shape)[:4]} != {B, T, H, W}"else:
raiseValueError(f"Unknown interpolation method {self.interpolation}")
returnnormalize(emb, dim=-1, eps=1e-6)
@torch.no_grad()defmatch_rope():
fromdiffusers.models.transformers.transformer_cosmosimportCosmosRotaryPosEmbedtheirs_rope=VideoRopePosition3DEmb(head_dim=128, len_h=240//2, len_w=240//2, len_t=128//1, base_fps=24, h_extrapolation_ratio=1.0, w_extrapolation_ratio=1.0, t_extrapolation_ratio=2.0)
ours_rope=CosmosRotaryPosEmbed(hidden_size=128, max_size=(128, 240, 240), patch_size=(1, 2, 2), base_fps=24, rope_scale=(2.0, 1.0, 1.0))
hidden_states=torch.randn(2, 2, 32, 32, 16)
fps=30theirs_output=theirs_rope(hidden_states[:, :, :16, :16, :], fps=torch.tensor([fps])) # the input slicing is to replicate patchification operationours_output=ours_rope(hidden_states.permute(0, 4, 1, 2, 3), fps=fps)
theirs_cos, theirs_sin=torch.cos(theirs_output), torch.sin(theirs_output)
print(torch.allclose(ours_output[0][:, None, None, :], theirs_cos))
print(torch.allclose(ours_output[1][:, None, None, :], theirs_sin))
@torch.no_grad()defmatch_learnable_pe():
fromdiffusers.models.transformers.transformer_cosmosimportCosmosLearnablePositionalEmbedtheirs_pe=LearnablePosEmbAxis(interpolation="crop", model_channels=4096, len_h=240//2, len_w=240//2, len_t=128//1)
ours_pe=CosmosLearnablePositionalEmbed(4096, max_size=(128, 240, 240), patch_size=(1, 2, 2), eps=1e-6)
ours_pe.pos_emb_t.data.copy_(theirs_pe.pos_emb_t.data)
ours_pe.pos_emb_h.data.copy_(theirs_pe.pos_emb_h.data)
ours_pe.pos_emb_w.data.copy_(theirs_pe.pos_emb_w.data)
hidden_states=torch.randn(2, 2, 32, 32, 16)
theirs_output=theirs_pe(hidden_states[:, :, :16, :16, :])
ours_output=ours_pe(hidden_states.permute(0, 4, 1, 2, 3))
theirs_output=theirs_output.flatten(1, 3)
print(torch.allclose(ours_output, theirs_output))
# match_rope()match_learnable_pe()

test transformer block

importsyssys.path.append("/raid/aryan/cosmos-code/")
importtorchfromcosmos1.models.diffusion.module.blocksimportGeneralDITTransformerBlock@torch.no_grad()defmatch_transformer_block():
fromdiffusers.models.transformers.transformer_cosmosimportCosmosTransformerBlocktheirs_transformer_block=GeneralDITTransformerBlock(
x_dim=4096,
context_dim=1024,
num_heads=32,
block_config="FA-CA-MLP",
mlp_ratio=4.0,
x_format="BTHWD",
use_adaln_lora=True,
adaln_lora_dim=256,
)
ours_transformer_block=CosmosTransformerBlock(
num_attention_heads=32,
attention_head_dim=128,
cross_attention_dim=1024,
mlp_ratio=4,
adaln_lora_dim=256,
qk_norm="rms_norm",
out_bias=False,
)
ours_transformer_block.norm1.linear_1.weight.data.copy_(theirs_transformer_block.blocks[0].adaLN_modulation[1].weight.data)
ours_transformer_block.norm1.linear_2.weight.data.copy_(theirs_transformer_block.blocks[0].adaLN_modulation[2].weight.data)
ours_transformer_block.attn1.to_q.weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_q[0].weight.data)
ours_transformer_block.attn1.to_k.weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_k[0].weight.data)
ours_transformer_block.attn1.to_v.weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_v[0].weight.data)
ours_transformer_block.attn1.to_out[0].weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_out[0].weight.data)
ours_transformer_block.attn1.norm_q.weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_q[1].weight.data)
ours_transformer_block.attn1.norm_k.weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_k[1].weight.data)
ours_transformer_block.norm2.linear_1.weight.data.copy_(theirs_transformer_block.blocks[1].adaLN_modulation[1].weight.data)
ours_transformer_block.norm2.linear_2.weight.data.copy_(theirs_transformer_block.blocks[1].adaLN_modulation[2].weight.data)
ours_transformer_block.attn2.to_q.weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_q[0].weight.data)
ours_transformer_block.attn2.to_k.weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_k[0].weight.data)
ours_transformer_block.attn2.to_v.weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_v[0].weight.data)
ours_transformer_block.attn2.to_out[0].weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_out[0].weight.data)
ours_transformer_block.attn2.norm_q.weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_q[1].weight.data)
ours_transformer_block.attn2.norm_k.weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_k[1].weight.data)
ours_transformer_block.norm3.linear_1.weight.data.copy_(theirs_transformer_block.blocks[2].adaLN_modulation[1].weight.data)
ours_transformer_block.norm3.linear_2.weight.data.copy_(theirs_transformer_block.blocks[2].adaLN_modulation[2].weight.data)
ours_transformer_block.ff.net[0].proj.weight.data.copy_(theirs_transformer_block.blocks[2].block.layer1.weight.data)
ours_transformer_block.ff.net[2].weight.data.copy_(theirs_transformer_block.blocks[2].block.layer2.weight.data)
# ============batch_size=1latent_num_frames=2latent_height=16latent_width=16embedding_dim=4096encoder_seq_length=64encoder_dim=1024hidden_states=torch.randn(batch_size, latent_num_frames, latent_height, latent_width, embedding_dim)
temb=torch.randn(batch_size, embedding_dim)
encoder_hidden_states=torch.randn(batch_size, encoder_seq_length, encoder_dim)
attention_mask=Nonefreqs=torch.randn(1, 1, latent_num_frames*latent_height*latent_width, 128)
embedded_timestep=torch.randn(batch_size, 3*embedding_dim)
extra_per_block_emb=torch.randn(batch_size, latent_num_frames, latent_height, latent_width, embedding_dim)
theirs_output=theirs_transformer_block(
x=hidden_states.flatten(1, 3).permute(1, 0, 2),
emb_B_D=temb,
crossattn_emb=encoder_hidden_states.permute(1, 0, 2),
crossattn_mask=attention_mask,
rope_emb_L_1_1_D=freqs.permute(2, 0, 1, 3),
adaln_lora_B_3D=embedded_timestep,
extra_per_block_pos_emb=extra_per_block_emb.flatten(1, 3).permute(1, 0, 2),
)
ours_output=ours_transformer_block(
hidden_states=hidden_states.flatten(1, 3),
encoder_hidden_states=encoder_hidden_states,
temb=temb,
embedded_timestep=embedded_timestep,
image_rotary_emb=(torch.cos(freqs.flatten(0, 2)), torch.sin(freqs.flatten(0, 2))),
extra_pos_emb=extra_per_block_emb.flatten(1, 3),
attention_mask=attention_mask,
)
theirs_output=theirs_output.flatten(0, 2).permute(1, 0, 2)
print(sum(p.numel() forpintheirs_transformer_block.parameters()))
print(sum(p.numel() forpinours_transformer_block.parameters()))
print(torch.allclose(theirs_output.flatten(), ours_output.flatten(), atol=1e-4))
match_transformer_block()
# GeneralDITTransformerBlock(# (blocks): ModuleList(# (0): DITBuildingBlock(# (block): VideoAttn(# (attn): Attention(# (to_q): Sequential(# (0): Linear(in_features=4096, out_features=4096, bias=False)# (1): RMSNorm()# )# (to_k): Sequential(# (0): Linear(in_features=4096, out_features=4096, bias=False)# (1): RMSNorm()# )# (to_v): Sequential(# (0): Linear(in_features=4096, out_features=4096, bias=False)# (1): Identity()# )# (to_out): Sequential(# (0): Linear(in_features=4096, out_features=4096, bias=False)# (1): Dropout(p=0.0, inplace=False)# )# )# )# (norm_state): LayerNorm((4096,), eps=1e-06, elementwise_affine=False)# (adaLN_modulation): Sequential(# (0): SiLU()# (1): Linear(in_features=4096, out_features=256, bias=False)# (2): Linear(in_features=256, out_features=12288, bias=False)# )# )# (1): DITBuildingBlock(# (block): VideoAttn(# (attn): Attention(# (to_q): Sequential(# (0): Linear(in_features=4096, out_features=4096, bias=False)# (1): RMSNorm()# )# (to_k): Sequential(# (0): Linear(in_features=1204, out_features=4096, bias=False)# (1): RMSNorm()# )# (to_v): Sequential(# (0): Linear(in_features=1204, out_features=4096, bias=False)# (1): Identity()# )# (to_out): Sequential(# (0): Linear(in_features=4096, out_features=4096, bias=False)# (1): Dropout(p=0.0, inplace=False)# )# )# )# (norm_state): LayerNorm((4096,), eps=1e-06, elementwise_affine=False)# (adaLN_modulation): Sequential(# (0): SiLU()# (1): Linear(in_features=4096, out_features=256, bias=False)# (2): Linear(in_features=256, out_features=12288, bias=False)# )# )# (2): DITBuildingBlock(# (block): GPT2FeedForward(# (layer1): Linear(in_features=4096, out_features=16384, bias=False)# (layer2): Linear(in_features=16384, out_features=4096, bias=False)# (dropout): Dropout(p=0.0, inplace=False)# (activation): GELU(approximate='none')# )# (norm_state): LayerNorm((4096,), eps=1e-06, elementwise_affine=False)# (adaLN_modulation): Sequential(# (0): SiLU()# (1): Linear(in_features=4096, out_features=256, bias=False)# (2): Linear(in_features=256, out_features=12288, bias=False)# )# )# )# )# CosmosTransformerBlock(# (norm1): CosmosAdaLayerNormZero(# (norm): LayerNorm((4096,), eps=1e-06, elementwise_affine=False)# (activation): SiLU()# (linear_1): Linear(in_features=4096, out_features=256, bias=False)# (linear_2): Linear(in_features=256, out_features=12288, bias=False)# )# (attn1): Attention(# (norm_q): RMSNorm()# (norm_k): RMSNorm()# (to_q): Linear(in_features=4096, out_features=4096, bias=False)# (to_k): Linear(in_features=4096, out_features=4096, bias=False)# (to_v): Linear(in_features=4096, out_features=4096, bias=False)# (to_out): ModuleList(# (0): Linear(in_features=4096, out_features=4096, bias=False)# (1): Dropout(p=0.0, inplace=False)# )# )# (norm2): CosmosAdaLayerNormZero(# (norm): LayerNorm((4096,), eps=1e-06, elementwise_affine=False)# (activation): SiLU()# (linear_1): Linear(in_features=4096, out_features=256, bias=False)# (linear_2): Linear(in_features=256, out_features=12288, bias=False)# )# (attn2): Attention(# (norm_q): RMSNorm()# (norm_k): RMSNorm()# (to_q): Linear(in_features=4096, out_features=4096, bias=False)# (to_k): Linear(in_features=1024, out_features=4096, bias=False)# (to_v): Linear(in_features=1024, out_features=4096, bias=False)# (to_out): ModuleList(# (0): Linear(in_features=4096, out_features=4096, bias=False)# (1): Dropout(p=0.0, inplace=False)# )# )# (norm3): CosmosAdaLayerNormZero(# (norm): LayerNorm((4096,), eps=1e-06, elementwise_affine=False)# (activation): SiLU()# (linear_1): Linear(in_features=4096, out_features=256, bias=False)# (linear_2): Linear(in_features=256, out_features=12288, bias=False)# )# (ff): FeedForward(# (net): ModuleList(# (0): GELU(# (proj): Linear(in_features=4096, out_features=16384, bias=False)# )# (1): Dropout(p=0.0, inplace=False)# (2): Linear(in_features=16384, out_features=4096, bias=False)# )# )# )

test transformer

importsyssys.path.append("/raid/aryan/cosmos-code/")
importtorchfromcosmos1.models.diffusion.networks.general_ditimportGeneralDIT@torch.no_grad()defmatch_transformer():
fromdiffusers.models.transformers.transformer_cosmosimportCosmosTransformer3DModeltheirs_transformer=GeneralDIT(
max_img_h=240,
max_img_w=240,
max_frames=128,
in_channels=16,
out_channels=16,
patch_spatial=2,
patch_temporal=1,
concat_padding_mask=True,
block_config="FA-CA-MLP",
model_channels=4096,
num_blocks=2,
num_heads=32,
mlp_ratio=4,
block_x_format="THWBD",
crossattn_emb_channels=1024,
use_cross_attn_mask=False,
pos_emb_cls="rope3d",
pos_emb_learnable=True,
pos_emb_interpolation="crop",
affline_emb_norm=True,
use_adaln_lora=True,
adaln_lora_dim=256,
rope_h_extrapolation_ratio=1.0,
rope_w_extrapolation_ratio=1.0,
rope_t_extrapolation_ratio=2.0,
extra_per_block_abs_pos_emb=True,
extra_per_block_abs_pos_emb_type="learnable",
)
ours_transformer=CosmosTransformer3DModel(
in_channels=16,
out_channels=16,
num_attention_heads=32,
attention_head_dim=128,
num_layers=2,
mlp_ratio=4,
text_embed_dim=1024,
adaln_lora_dim=256,
max_size=(128, 240, 240),
patch_size=(1, 2, 2),
rope_scale=(2.0, 1.0, 1.0),
concat_padding_mask=True,
extra_pos_embed_type="learnable",
)
# Patch embeddingours_transformer.patch_embed.proj.weight.data.copy_(theirs_transformer.x_embedder.proj[1].weight.data)
# Timestep embeddingours_t_embedder=ours_transformer.time_embedtheirs_t_embedder=theirs_transformer.t_embeddertheirs_norm=theirs_transformer.affline_normours_t_embedder.t_embedder.linear_1.weight.data.copy_(theirs_t_embedder[1].linear_1.weight.data)
ours_t_embedder.t_embedder.linear_2.weight.data.copy_(theirs_t_embedder[1].linear_2.weight.data)
ours_t_embedder.norm.weight.data.copy_(theirs_norm.weight.data)
# Learnable position embeddingours_pe=ours_transformer.learnable_pos_embedtheirs_pe=theirs_transformer.extra_pos_embedderours_pe.pos_emb_t.data.copy_(theirs_pe.pos_emb_t.data)
ours_pe.pos_emb_h.data.copy_(theirs_pe.pos_emb_h.data)
ours_pe.pos_emb_w.data.copy_(theirs_pe.pos_emb_w.data)
# Transformer blocksforiinrange(2):
ours_transformer_block=ours_transformer.transformer_blocks[i]
theirs_transformer_block=theirs_transformer.blocks[f"block{i}"]
ours_transformer_block.norm1.linear_1.weight.data.copy_(theirs_transformer_block.blocks[0].adaLN_modulation[1].weight.data)
ours_transformer_block.norm1.linear_2.weight.data.copy_(theirs_transformer_block.blocks[0].adaLN_modulation[2].weight.data)
ours_transformer_block.attn1.to_q.weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_q[0].weight.data)
ours_transformer_block.attn1.to_k.weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_k[0].weight.data)
ours_transformer_block.attn1.to_v.weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_v[0].weight.data)
ours_transformer_block.attn1.to_out[0].weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_out[0].weight.data)
ours_transformer_block.attn1.norm_q.weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_q[1].weight.data)
ours_transformer_block.attn1.norm_k.weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_k[1].weight.data)
ours_transformer_block.norm2.linear_1.weight.data.copy_(theirs_transformer_block.blocks[1].adaLN_modulation[1].weight.data)
ours_transformer_block.norm2.linear_2.weight.data.copy_(theirs_transformer_block.blocks[1].adaLN_modulation[2].weight.data)
ours_transformer_block.attn2.to_q.weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_q[0].weight.data)
ours_transformer_block.attn2.to_k.weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_k[0].weight.data)
ours_transformer_block.attn2.to_v.weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_v[0].weight.data)
ours_transformer_block.attn2.to_out[0].weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_out[0].weight.data)
ours_transformer_block.attn2.norm_q.weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_q[1].weight.data)
ours_transformer_block.attn2.norm_k.weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_k[1].weight.data)
ours_transformer_block.norm3.linear_1.weight.data.copy_(theirs_transformer_block.blocks[2].adaLN_modulation[1].weight.data)
ours_transformer_block.norm3.linear_2.weight.data.copy_(theirs_transformer_block.blocks[2].adaLN_modulation[2].weight.data)
ours_transformer_block.ff.net[0].proj.weight.data.copy_(theirs_transformer_block.blocks[2].block.layer1.weight.data)
ours_transformer_block.ff.net[2].weight.data.copy_(theirs_transformer_block.blocks[2].block.layer2.weight.data)
# Output layersours_transformer.norm_out.linear_1.weight.data.copy_(theirs_transformer.final_layer.adaLN_modulation[1].weight.data)
ours_transformer.norm_out.linear_2.weight.data.copy_(theirs_transformer.final_layer.adaLN_modulation[2].weight.data)
ours_transformer.proj_out.weight.data.copy_(theirs_transformer.final_layer.linear.weight.data)
forname, paramintheirs_transformer.named_parameters():
if"bias"inname:
print(name, param.shape)
forname, paraminours_transformer.named_parameters():
if"bias"inname:
print(name, param.shape)
# ============batch_size=1latent_num_frames=2latent_height=16latent_width=16encoder_seq_length=64encoder_dim=1024fps=30.0hidden_states=torch.randn(batch_size, latent_num_frames, latent_height, latent_width, 16)
timestep=torch.randint(0, 1000, (batch_size,)).float()
encoder_hidden_states=torch.randn(batch_size, encoder_seq_length, encoder_dim)
attention_mask=Nonepadding_mask=torch.zeros((1, 1, latent_height*8, latent_width*8))
theirs_output=theirs_transformer(
x=hidden_states.permute(0, 4, 1, 2, 3),
timesteps=timestep,
crossattn_emb=encoder_hidden_states,
crossattn_mask=attention_mask,
fps=torch.tensor([fps]),
padding_mask=padding_mask,
)
print()
ours_output=ours_transformer(
hidden_states=hidden_states.permute(0, 4, 1, 2, 3),
timestep=timestep.long(),
encoder_hidden_states=encoder_hidden_states,
attention_mask=attention_mask,
fps=fps,
padding_mask=padding_mask,
)[0]
print(torch.allclose(theirs_output, ours_output, atol=1e-4))
match_transformer()

test transformer video

importsyssys.path.append("/raid/aryan/cosmos-code/")
importtorchfromcosmos1.models.diffusion.networks.general_dit_video_conditionedimportVideoExtendGeneralDIT@torch.no_grad()defmatch_transformer():
fromdiffusers.models.transformers.transformer_cosmosimportCosmosTransformer3DModeltheirs_transformer=VideoExtendGeneralDIT(
max_img_h=240,
max_img_w=240,
max_frames=128,
in_channels=16+1,
out_channels=16,
patch_spatial=2,
patch_temporal=1,
concat_padding_mask=True,
block_config="FA-CA-MLP",
model_channels=4096,
num_blocks=2,
num_heads=32,
mlp_ratio=4,
block_x_format="THWBD",
crossattn_emb_channels=1024,
use_cross_attn_mask=False,
pos_emb_cls="rope3d",
pos_emb_learnable=True,
pos_emb_interpolation="crop",
affline_emb_norm=True,
use_adaln_lora=True,
adaln_lora_dim=256,
rope_h_extrapolation_ratio=1.0,
rope_w_extrapolation_ratio=1.0,
rope_t_extrapolation_ratio=2.0,
extra_per_block_abs_pos_emb=True,
extra_per_block_abs_pos_emb_type="learnable",
)
ours_transformer=CosmosTransformer3DModel(
in_channels=16+1,
out_channels=16,
num_attention_heads=32,
attention_head_dim=128,
num_layers=2,
mlp_ratio=4,
text_embed_dim=1024,
adaln_lora_dim=256,
max_size=(128, 240, 240),
patch_size=(1, 2, 2),
rope_scale=(2.0, 1.0, 1.0),
concat_padding_mask=True,
extra_pos_embed_type="learnable",
)
# Patch embeddingours_transformer.patch_embed.proj.weight.data.copy_(theirs_transformer.x_embedder.proj[1].weight.data)
# Timestep embeddingours_t_embedder=ours_transformer.time_embedtheirs_t_embedder=theirs_transformer.t_embeddertheirs_norm=theirs_transformer.affline_normours_t_embedder.t_embedder.linear_1.weight.data.copy_(theirs_t_embedder[1].linear_1.weight.data)
ours_t_embedder.t_embedder.linear_2.weight.data.copy_(theirs_t_embedder[1].linear_2.weight.data)
ours_t_embedder.norm.weight.data.copy_(theirs_norm.weight.data)
# Learnable position embeddingours_pe=ours_transformer.learnable_pos_embedtheirs_pe=theirs_transformer.extra_pos_embedderours_pe.pos_emb_t.data.copy_(theirs_pe.pos_emb_t.data)
ours_pe.pos_emb_h.data.copy_(theirs_pe.pos_emb_h.data)
ours_pe.pos_emb_w.data.copy_(theirs_pe.pos_emb_w.data)
# Transformer blocksforiinrange(2):
ours_transformer_block=ours_transformer.transformer_blocks[i]
theirs_transformer_block=theirs_transformer.blocks[f"block{i}"]
ours_transformer_block.norm1.linear_1.weight.data.copy_(theirs_transformer_block.blocks[0].adaLN_modulation[1].weight.data)
ours_transformer_block.norm1.linear_2.weight.data.copy_(theirs_transformer_block.blocks[0].adaLN_modulation[2].weight.data)
ours_transformer_block.attn1.to_q.weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_q[0].weight.data)
ours_transformer_block.attn1.to_k.weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_k[0].weight.data)
ours_transformer_block.attn1.to_v.weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_v[0].weight.data)
ours_transformer_block.attn1.to_out[0].weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_out[0].weight.data)
ours_transformer_block.attn1.norm_q.weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_q[1].weight.data)
ours_transformer_block.attn1.norm_k.weight.data.copy_(theirs_transformer_block.blocks[0].block.attn.to_k[1].weight.data)
ours_transformer_block.norm2.linear_1.weight.data.copy_(theirs_transformer_block.blocks[1].adaLN_modulation[1].weight.data)
ours_transformer_block.norm2.linear_2.weight.data.copy_(theirs_transformer_block.blocks[1].adaLN_modulation[2].weight.data)
ours_transformer_block.attn2.to_q.weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_q[0].weight.data)
ours_transformer_block.attn2.to_k.weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_k[0].weight.data)
ours_transformer_block.attn2.to_v.weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_v[0].weight.data)
ours_transformer_block.attn2.to_out[0].weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_out[0].weight.data)
ours_transformer_block.attn2.norm_q.weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_q[1].weight.data)
ours_transformer_block.attn2.norm_k.weight.data.copy_(theirs_transformer_block.blocks[1].block.attn.to_k[1].weight.data)
ours_transformer_block.norm3.linear_1.weight.data.copy_(theirs_transformer_block.blocks[2].adaLN_modulation[1].weight.data)
ours_transformer_block.norm3.linear_2.weight.data.copy_(theirs_transformer_block.blocks[2].adaLN_modulation[2].weight.data)
ours_transformer_block.ff.net[0].proj.weight.data.copy_(theirs_transformer_block.blocks[2].block.layer1.weight.data)
ours_transformer_block.ff.net[2].weight.data.copy_(theirs_transformer_block.blocks[2].block.layer2.weight.data)
# Output layersours_transformer.norm_out.linear_1.weight.data.copy_(theirs_transformer.final_layer.adaLN_modulation[1].weight.data)
ours_transformer.norm_out.linear_2.weight.data.copy_(theirs_transformer.final_layer.adaLN_modulation[2].weight.data)
ours_transformer.proj_out.weight.data.copy_(theirs_transformer.final_layer.linear.weight.data)
forname, paramintheirs_transformer.named_parameters():
if"bias"inname:
print(name, param.shape)
forname, paraminours_transformer.named_parameters():
if"bias"inname:
print(name, param.shape)
# ============batch_size=1latent_num_frames=2latent_height=16latent_width=16encoder_seq_length=64encoder_dim=1024fps=30.0hidden_states=torch.randn(batch_size, latent_num_frames, latent_height, latent_width, 16)
timestep=torch.randint(0, 1000, (batch_size,)).float()
encoder_hidden_states=torch.randn(batch_size, encoder_seq_length, encoder_dim)
attention_mask=Nonecondition_mask=torch.ones(batch_size, 1, latent_num_frames, latent_height, latent_width)
padding_mask=torch.zeros((1, 1, latent_height*8, latent_width*8))
theirs_output=theirs_transformer(
x=hidden_states.permute(0, 4, 1, 2, 3),
timesteps=timestep,
crossattn_emb=encoder_hidden_states,
crossattn_mask=attention_mask,
fps=torch.tensor([fps]),
condition_video_input_mask=condition_mask,
padding_mask=padding_mask,
)
print()
ours_output=ours_transformer(
hidden_states=hidden_states.permute(0, 4, 1, 2, 3),
timestep=timestep.long(),
encoder_hidden_states=encoder_hidden_states,
attention_mask=attention_mask,
fps=fps,
condition_mask=condition_mask,
padding_mask=padding_mask,
)[0]
print(torch.allclose(theirs_output, ours_output, atol=1e-4))
match_transformer()

VAE

test vae attention

importsyssys.path.append("/raid/aryan/cosmos-tokenizer-code/")
fromcosmos_tokenizer.modules.layers3dimportCausalAttnBlock, CausalTemporalAttnBlockfromtypingimportUnion, Optionalimporttorchimporttorch.nnasnnimporttorch.nn.functionalasFfromdiffusers.models.autoencoders.autoencoder_kl_cosmosimportCosmosCausalGroupNorm, CosmosCausalConv3dclassCosmosCausalAttention(nn.Module):
def__init__(self, num_attention_heads: int, attention_head_dim: int, num_groups: int=1, dropout: float=0.0, processor: Union["CosmosSpatialAttentionProcessor2_0", "CosmosTemporalAttentionProcessor2_0"] =None) ->None:
super().__init__()
self.num_attention_heads=num_attention_headsself.norm=CosmosCausalGroupNorm(attention_head_dim, num_groups=num_groups)
self.to_q=CosmosCausalConv3d(attention_head_dim, attention_head_dim, kernel_size=1, stride=1, padding=0)
self.to_k=CosmosCausalConv3d(attention_head_dim, attention_head_dim, kernel_size=1, stride=1, padding=0)
self.to_v=CosmosCausalConv3d(attention_head_dim, attention_head_dim, kernel_size=1, stride=1, padding=0)
self.to_out=nn.ModuleList([])
self.to_out.append(CosmosCausalConv3d(attention_head_dim, attention_head_dim, kernel_size=1, stride=1, padding=0))
self.to_out.append(nn.Dropout(dropout))
self.processor=processorifself.processorisNone:
raiseValueError("CosmosCausalAttention requires a processor.")
defforward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] =None) ->torch.Tensor:
returnself.processor(self, hidden_states=hidden_states, attention_mask=attention_mask)
classCosmosSpatialAttentionProcessor2_0:
def__init__(self):
ifnothasattr(F, "scaled_dot_product_attention"):
raiseImportError("CosmosSpatialAttentionProcessor2_0 requires PyTorch 2.0 or higher. To use it, please upgrade PyTorch.")
def__call__(self, attn: CosmosCausalAttention, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] =None) ->torch.Tensor:
batch_size, num_channels, num_frames, height, width=hidden_states.shaperesidual=hidden_stateshidden_states=attn.norm(hidden_states)
query=attn.to_q(hidden_states)
key=attn.to_k(hidden_states)
value=attn.to_v(hidden_states)
# [B, C, T, H, W] -> [B * T, H * W, C]query=query.permute(0, 2, 3, 4, 1).flatten(2, 3).flatten(0, 1)
key=key.permute(0, 2, 3, 4, 1).flatten(2, 3).flatten(0, 1)
value=value.permute(0, 2, 3, 4, 1).flatten(2, 3).flatten(0, 1)
# [B * T, H * W, C] -> [B * T, N, H * W, C // N]query=query.unflatten(2, (attn.num_attention_heads, -1)).transpose(1, 2)
key=key.unflatten(2, (attn.num_attention_heads, -1)).transpose(1, 2)
value=value.unflatten(2, (attn.num_attention_heads, -1)).transpose(1, 2)
hidden_states=F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask)
hidden_states=hidden_states.transpose(1, 2).flatten(2, 3).type_as(query)
hidden_states=hidden_states.unflatten(1, (height, width)).unflatten(0, (batch_size, num_frames))
hidden_states=hidden_states.permute(0, 4, 1, 2, 3)
hidden_states=attn.to_out[0](hidden_states)
hidden_states=attn.to_out[1](hidden_states)
returnhidden_states+residualclassCosmosTemporalAttentionProcessor2_0:
def__init__(self):
ifnothasattr(F, "scaled_dot_product_attention"):
raiseImportError("CosmosSpatialAttentionProcessor2_0 requires PyTorch 2.0 or higher. To use it, please upgrade PyTorch.")
def__call__(self, attn: CosmosCausalAttention, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] =None) ->torch.Tensor:
batch_size, num_channels, num_frames, height, width=hidden_states.shaperesidual=hidden_stateshidden_states=attn.norm(hidden_states)
query=attn.to_q(hidden_states)
key=attn.to_k(hidden_states)
value=attn.to_v(hidden_states)
# [B, C, T, H, W] -> [B * T, H * W, C]query=query.permute(0, 3, 4, 2, 1).flatten(0, 2)
key=key.permute(0, 3, 4, 2, 1).flatten(0, 2)
value=value.permute(0, 3, 4, 2, 1).flatten(0, 2)
# [B * T, H * W, C] -> [B * T, N, H * W, C // N]query=query.unflatten(2, (attn.num_attention_heads, -1)).transpose(1, 2)
key=key.unflatten(2, (attn.num_attention_heads, -1)).transpose(1, 2)
value=value.unflatten(2, (attn.num_attention_heads, -1)).transpose(1, 2)
hidden_states=F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask)
hidden_states=hidden_states.transpose(1, 2).flatten(2, 3).type_as(query)
hidden_states=hidden_states.unflatten(0, (batch_size, height, width))
hidden_states=hidden_states.permute(0, 4, 3, 1, 2)
hidden_states=attn.to_out[0](hidden_states)
hidden_states=attn.to_out[1](hidden_states)
returnhidden_states+residual@torch.no_grad()deftest_causal_attn_block_spatial():
in_channels=128torch.manual_seed(0)
theirs_attn=CausalAttnBlock(in_channels, num_groups=1)
ours_attn=CosmosCausalAttention(num_attention_heads=1, attention_head_dim=in_channels, num_groups=1, dropout=0.0, processor=CosmosSpatialAttentionProcessor2_0())
ours_attn.to_q.conv.weight.data.copy_(theirs_attn.q.conv3d.weight.data)
ours_attn.to_k.conv.weight.data.copy_(theirs_attn.k.conv3d.weight.data)
ours_attn.to_v.conv.weight.data.copy_(theirs_attn.v.conv3d.weight.data)
ours_attn.to_out[0].conv.weight.data.copy_(theirs_attn.proj_out.conv3d.weight.data)
ours_attn.to_q.conv.bias.data.copy_(theirs_attn.q.conv3d.bias.data)
ours_attn.to_k.conv.bias.data.copy_(theirs_attn.k.conv3d.bias.data)
ours_attn.to_v.conv.bias.data.copy_(theirs_attn.v.conv3d.bias.data)
ours_attn.to_out[0].conv.bias.data.copy_(theirs_attn.proj_out.conv3d.bias.data)
ours_attn.norm.norm.weight.data.copy_(theirs_attn.norm.norm.weight.data)
ours_attn.norm.norm.bias.data.copy_(theirs_attn.norm.norm.bias.data)
batch_size=2num_frames=16height=8width=8hidden_states=torch.randn(batch_size, in_channels, num_frames, height, width)
theirs_output=theirs_attn(hidden_states)
ours_output=ours_attn(hidden_states)
diff=theirs_output-ours_outputprint(f"absmax diff: {diff.abs().max()}")
print(f"absmean diff: {diff.abs().mean()}")
@torch.no_grad()deftest_causal_attn_block_temporal():
in_channels=128torch.manual_seed(0)
theirs_attn=CausalTemporalAttnBlock(in_channels, num_groups=1)
ours_attn=CosmosCausalAttention(num_attention_heads=1, attention_head_dim=in_channels, num_groups=1, dropout=0.0, processor=CosmosTemporalAttentionProcessor2_0())
ours_attn.to_q.conv.weight.data.copy_(theirs_attn.q.conv3d.weight.data)
ours_attn.to_k.conv.weight.data.copy_(theirs_attn.k.conv3d.weight.data)
ours_attn.to_v.conv.weight.data.copy_(theirs_attn.v.conv3d.weight.data)
ours_attn.to_out[0].conv.weight.data.copy_(theirs_attn.proj_out.conv3d.weight.data)
ours_attn.to_q.conv.bias.data.copy_(theirs_attn.q.conv3d.bias.data)
ours_attn.to_k.conv.bias.data.copy_(theirs_attn.k.conv3d.bias.data)
ours_attn.to_v.conv.bias.data.copy_(theirs_attn.v.conv3d.bias.data)
ours_attn.to_out[0].conv.bias.data.copy_(theirs_attn.proj_out.conv3d.bias.data)
ours_attn.norm.norm.weight.data.copy_(theirs_attn.norm.norm.weight.data)
ours_attn.norm.norm.bias.data.copy_(theirs_attn.norm.norm.bias.data)
batch_size=2num_frames=16height=8width=8hidden_states=torch.randn(batch_size, in_channels, num_frames, height, width)
theirs_output=theirs_attn(hidden_states)
attn_mask=torch.tril(hidden_states.new_ones(num_frames, num_frames)).bool()
ours_output=ours_attn(hidden_states, attn_mask)
diff=theirs_output-ours_outputprint(f"absmax diff: {diff.abs().max()}")
print(f"absmean diff: {diff.abs().mean()}")
test_causal_attn_block_temporal()
test_causal_attn_block_spatial()

test vae

importsyssys.path.append("/raid/aryan/cosmos-tokenizer-code/")
fromcosmos_tokenizer.modulesimport (
ContinuousFormulation,
Encoder3DType,
Decoder3DType,
)
fromcosmos_tokenizer.networks.continuous_videoimportCausalContinuousVideoTokenizerimporttorchfromaccelerateimportinit_empty_weightsfromtypingimportDict, Anydefremove_keys_(key: str, state_dict: Dict[str, Any]):
state_dict.pop(key)
defupdate_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) ->Dict[str, Any]:
state_dict[new_key] =state_dict.pop(old_key)
VAE_KEYS_RENAME_DICT= {
"down.0": "down_blocks.0",
"down.1": "down_blocks.1",
"down.2": "down_blocks.2",
"up.0": "up_blocks.2",
"up.1": "up_blocks.1",
"up.2": "up_blocks.0",
".block.": ".resnets.",
"downsample": "downsamplers.0",
"upsample": "upsamplers.0",
"mid.block_1": "mid_block.resnets.0",
"mid.attn_1.0": "mid_block.attentions.0",
"mid.attn_1.1": "mid_block.temp_attentions.0",
"mid.block_2": "mid_block.resnets.1",
".q.conv3d": ".to_q",
".k.conv3d": ".to_k",
".v.conv3d": ".to_v",
".proj_out.conv3d": ".to_out.0",
".0.conv3d": ".conv_s",
".1.conv3d": ".conv_t",
"conv1.conv3d": "conv1",
"conv2.conv3d": "conv2",
"conv3.conv3d": "conv3",
"nin_shortcut.conv3d": "conv_shortcut",
"quant_conv.conv3d": "quant_conv",
"post_quant_conv.conv3d": "post_quant_conv",
}
VAE_SPECIAL_KEYS_REMAP= {}
@torch.no_grad()deftest_vae():
fromdiffusersimportAutoencoderKLCosmostorch.manual_seed(0)
theirs_config=dict(
attn_resolutions=[32],
channels=128,
channels_mult=[2, 4, 4],
dropout=0.0,
in_channels=3,
num_res_blocks=2,
out_channels=3,
resolution=1024,
patch_size=4,
patch_method="haar",
latent_channels=16,
z_channels=16,
z_factor=1,
num_groups=1,
legacy_mode=False,
spatial_compression=8,
temporal_compression=8,
formulation=ContinuousFormulation.AE.name,
encoder=Encoder3DType.FACTORIZED.name,
decoder=Decoder3DType.FACTORIZED.name,
name="CV",
)
theirs_model=CausalContinuousVideoTokenizer(**theirs_config)
ours_model=AutoencoderKLCosmos()
# print(theirs_model.decoder)# print()# print()# print()# print()# print(ours_model.decoder)theirs_num_params=sum(p.numel() forpintheirs_model.parameters())
ours_num_params=sum(p.numel() forpinours_model.parameters())
print(f"theirs_num_params: {theirs_num_params}")
print(f"ours_num_params: {ours_num_params}")
PREFIX_KEY=""original_state_dict=theirs_model.state_dict()
forkeyinlist(original_state_dict.keys()):
new_key=key[:]
ifnew_key.startswith(PREFIX_KEY):
new_key=new_key.removeprefix(PREFIX_KEY)
forreplace_key, rename_keyinVAE_KEYS_RENAME_DICT.items():
new_key=new_key.replace(replace_key, rename_key)
update_state_dict_(original_state_dict, key, new_key)
forkeyinlist(original_state_dict.keys()):
forspecial_key, handler_fn_inplaceinVAE_SPECIAL_KEYS_REMAP.items():
ifspecial_keynotinkey:
continuehandler_fn_inplace(key, original_state_dict)
ours_model.load_state_dict(original_state_dict, strict=True, assign=True)
batch_size=2num_channels=3num_frames=49height=256width=256hidden_states=torch.randn(batch_size, num_channels, num_frames, height, width)
theirs_output=theirs_model(hidden_states)["reconstructions"]
ours_output=ours_model(hidden_states)[0]
# torch.Size([2, 3, 49, 256, 256]) torch.Size([2, 3, 97, 512, 512])print(theirs_output.shape, ours_output.shape)
diff=theirs_output-ours_outputprint(f"absmax diff: {diff.abs().max()}")
print(f"absmean diff: {diff.abs().mean()}")
test_vae()

Text-to-World:

importtorchfromdiffusersimportCosmosTextToWorldPipelinefromdiffusers.utilsimportexport_to_videomodel_id="nvidia/Cosmos-1.0-Diffusion-7B-Text2World"pipe=CosmosTextToWorldPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
pipe.to("cuda")
prompt="A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves. The robot's metallic body gleams under the bright, even lighting, highlighting its futuristic design and intricate joints. A glowing blue light emanates from its chest, adding a touch of advanced technology. The background is dominated by rows of boxes, suggesting a highly organized storage system. The floor is lined with wooden pallets, enhancing the industrial setting. The camera remains static, capturing the robot's poised stance amidst the orderly environment, with a shallow depth of field that keeps the focus on the robot while subtly blurring the background for a cinematic effect."output=pipe(prompt=prompt).frames[0]
export_to_video(output, "output.mp4", fps=30)

Video-to-World (image-conditioning):

importtorchfromdiffusersimportCosmosVideoToWorldPipelinefromdiffusers.utilsimportexport_to_video, load_imagemodel_id="nvidia/Cosmos-1.0-Diffusion-7B-Video2World"pipe=CosmosVideoToWorldPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
pipe.to("cuda")
prompt="The video depicts a long, straight highway stretching into the distance, flanked by metal guardrails. The road is divided into multiple lanes, with a few vehicles visible in the far distance. The surrounding landscape features dry, grassy fields on one side and rolling hills on the other. The sky is mostly clear with a few scattered clouds, suggesting a bright, sunny day."image=load_image(
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cosmos/cosmos-video2world-input.jpg"
)
video=pipe(image=image, prompt=prompt).frames[0]
export_to_video(video, "output.mp4", fps=30)

Video-to-World (video-conditioning):

importtorchfromdiffusersimportCosmosVideoToWorldPipelinefromdiffusers.utilsimportexport_to_video, load_videomodel_id="nvidia/Cosmos-1.0-Diffusion-7B-Video2World"pipe=CosmosVideoToWorldPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
pipe.transformer=torch.compile(pipe.transformer)
pipe.to("cuda")
prompt="The video depicts a winding mountain road covered in snow, with a single vehicle traveling along it. The road is flanked by steep, rocky cliffs and sparse vegetation. The landscape is characterized by rugged terrain and a river visible in the distance. The scene captures the solitude and beauty of a winter drive through a mountainous region."video=load_video(
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cosmos/cosmos-video2world-input-vid.mp4"
)[:21] # This example uses only the first 21 framesvideo=pipe(video=video, prompt=prompt).frames[0]
export_to_video(video, "output.mp4", fps=30)

Note that the model repos are not yet compatible with Diffusers-loading. I'll open PRs for weights once nvidia team gives the thumbs up.

Inference code (old)

importosfromtypingimportAny, DictimporttorchfromdiffusersimportCosmosTransformer3DModel, CosmosPipeline, EDMEulerScheduler, EDMDPMSolverMultistepSchedulerfromdiffusers.utilsimportexport_to_videofromtransformersimportT5EncoderModel, T5TokenizerFastdefremove_keys_(key: str, state_dict: Dict[str, Any]):
state_dict.pop(key)
defupdate_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) ->Dict[str, Any]:
state_dict[new_key] =state_dict.pop(old_key)
defrename_transformer_blocks_(key: str, state_dict: Dict[str, Any]):
block_index=int(key.split(".")[1].removeprefix("block"))
new_key=keyold_prefix=f"blocks.block{block_index}"new_prefix=f"transformer_blocks.{block_index}"new_key=new_prefix+new_key.removeprefix(old_prefix)
state_dict[new_key] =state_dict.pop(key)
TRANSFORMER_KEYS_RENAME_DICT= {
"t_embedder.1": "time_embed.t_embedder",
"affline_norm": "time_embed.norm",
".blocks.0.block.attn": ".attn1",
".blocks.1.block.attn": ".attn2",
".blocks.2.block": ".ff",
".blocks.0.adaLN_modulation.1": ".norm1.linear_1",
".blocks.0.adaLN_modulation.2": ".norm1.linear_2",
".blocks.1.adaLN_modulation.1": ".norm2.linear_1",
".blocks.1.adaLN_modulation.2": ".norm2.linear_2",
".blocks.2.adaLN_modulation.1": ".norm3.linear_1",
".blocks.2.adaLN_modulation.2": ".norm3.linear_2",
"to_q.0": "to_q",
"to_q.1": "norm_q",
"to_k.0": "to_k",
"to_k.1": "norm_k",
"to_v.0": "to_v",
"layer1": "net.0.proj",
"layer2": "net.2",
"proj.1": "proj",
"x_embedder": "patch_embed",
"extra_pos_embedder": "learnable_pos_embed",
"final_layer.adaLN_modulation.1": "norm_out.linear_1",
"final_layer.adaLN_modulation.2": "norm_out.linear_2",
"final_layer.linear": "proj_out",
}
TRANSFORMER_SPECIAL_KEYS_REMAP= {
"blocks.block": rename_transformer_blocks_,
"logvar.0.freqs": remove_keys_,
"logvar.0.phases": remove_keys_,
"logvar.1.weight": remove_keys_,
"pos_embedder.seq": remove_keys_,
}
defconvert_transformer(state_dict):
PREFIX_KEY="net."forkeyinlist(state_dict.keys()):
new_key=key[:]
ifnew_key.startswith(PREFIX_KEY):
new_key=key[len(PREFIX_KEY) :]
forreplace_key, rename_keyinTRANSFORMER_KEYS_RENAME_DICT.items():
new_key=new_key.replace(replace_key, rename_key)
update_state_dict_inplace(state_dict, key, new_key)
forkeyinlist(state_dict.keys()):
forspecial_key, handler_fn_inplaceinTRANSFORMER_SPECIAL_KEYS_REMAP.items():
ifspecial_keynotinkey:
continuehandler_fn_inplace(key, state_dict)
returnstate_dicttorch.manual_seed(0)
device="cuda"dtype=torch.bfloat16withtorch.no_grad():
withtorch.device("meta"):
transformer=CosmosTransformer3DModel()
num_parameters=sum(p.numel() forpintransformer.parameters())
print(f"{num_parameters=}")
checkpoint_file="/raid/aryan/cosmos-code/checkpoints/Cosmos-1.0-Diffusion-7B-Text2World/model.pt"checkpoint=torch.load(checkpoint_file, map_location="cpu", weights_only=True)
checkpoint=convert_transformer(checkpoint)
transformer.load_state_dict(checkpoint, strict=True, assign=True)
text_encoder=T5EncoderModel.from_pretrained("google-t5/t5-11b", torch_dtype=dtype, cache_dir="/raid/aryan/cosmos-code/checkpoints")
tokenizer=T5TokenizerFast.from_pretrained("google-t5/t5-11b", cache_dir="/raid/aryan/cosmos-code/checkpoints")
vae_dir="/raid/aryan/cosmos-code/checkpoints/Cosmos-1.0-Tokenizer-CV8x8x8"decoder=torch.jit.load(os.path.join(vae_dir, "decoder.jit")).to(device=device, dtype=dtype)
latent_mean, latent_std=torch.load(os.path.join(vae_dir, "mean_std.pt"), weights_only=True)
scheduler=EDMEulerScheduler(final_sigmas_type="sigma_min")
pipe=CosmosPipeline(text_encoder, tokenizer, transformer, vae=None, scheduler=scheduler)
pipe.to(device, dtype=dtype)
prompt="A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves. The robot's metallic body gleams under the bright, even lighting, highlighting its futuristic design and intricate joints. A glowing blue light emanates from its chest, adding a touch of advanced technology. The background is dominated by rows of boxes, suggesting a highly organized storage system. The floor is lined with wooden pallets, enhancing the industrial setting. The camera remains static, capturing the robot's poised stance amidst the orderly environment, with a shallow depth of field that keeps the focus on the robot while subtly blurring the background for a cinematic effect."negative_prompt="The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality."latents=pipe(
prompt=prompt,
negative_prompt=negative_prompt,
height=704,
width=960,
# width=1280,num_frames=121,
num_inference_steps=36,
output_type="latent",
).framestorch.save(latents, "latents.pt")
latent_mean=latent_mean.to(device=device).reshape(1, 16, 16, 1, 1).float()[:, :, :latents.size(2)]
latent_std=latent_std.to(device=device).reshape(1, 16, 16, 1, 1).float()[:, :, :latents.size(2)]
sigma_data=0.5latents=latents/sigma_datalatents= (latents.float() *latent_std+latent_mean).type_as(latents)
output=decoder(latents.to(device=device, dtype=dtype))
video=pipe.video_processor.postprocess_video(output, output_type="pil")[0]
export_to_video(video, "output.mp4", fps=30)

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[a-r-r-o-w]

To match our sigmas to original exactly, without any rounding errors, I had to use torch.float64. This change is maybe not required since the values are nearly the same, but just something to keep in mind.

# theirs: [80.0, 68.32506, 58.14207, 49.28863, 41.61683, 34.99219, 29.29279, 24.40834, 20.23932, 16.69618, 13.69857, 11.17463, 9.06026, 7.29851, 5.83895, 4.63707, 3.65381, 2.85496, 2.21074, 1.69537, 1.28661, 0.96542, 0.71556, 0.52331, 0.37715, 0.26748, 0.18636, 0.12731, 0.08509, 0.05548, 0.03519, 0.02162, 0.01281, 0.00727, 0.00393, 0.002]
# ours_original: [79.99998, 68.00508, 57.58597, 48.56622, 40.78557, 34.09878, 28.37458, 23.49461, 19.35245, 15.8527, 12.91008, 10.44864, 8.40094, 6.70731, 5.31519, 4.17847, 3.25682, 2.51522, 1.92334, 1.4551, 1.08817, 0.80359, 0.58535, 0.42002, 0.29644, 0.20544, 0.13952, 0.09262, 0.05995, 0.03769, 0.02293, 0.01343, 0.00753, 0.004, 0.002, 0.0]
# ours_modified: [80.0, 68.32506, 58.14206, 49.28863, 41.61682, 34.99219, 29.29279, 24.40834, 20.23932, 16.69618, 13.69858, 11.17463, 9.06026, 7.29851, 5.83895, 4.63708, 3.65381, 2.85496, 2.21074, 1.69537, 1.28661, 0.96542, 0.71556, 0.52331, 0.37715, 0.26748, 0.18636, 0.12731, 0.08509, 0.05548, 0.03519, 0.02162, 0.01281, 0.00727, 0.00393, 0.002, 0.0]

Also, we only match the sigmas if we set our our_num_inference_steps=their_num_inference_steps + 1. This is because they do an extra inference step without scheduler step (effective same as setting our final_sigmas_type="sigma_min"

[yiyixuxu]

@a-r-r-o-w
i made the package here if you want to to take a look and test it out https://github.com/yiyixuxu/cosmos-guardrail
I can publish later once it's ok

[a-r-r-o-w]

Thanks @yiyixuxu! I'll take a look and update our implementation accordingly tomorrow

[a-r-r-o-w]

@yiyixuxu I've updated the code to use the package. Could you take a look again? If everything looks good, let's try to get the 7B model weights merged and I'll open the 14B model weight PRs soon

[yiyixuxu]

thanks! I left some questions but the PR looks good to me

[yiyixuxu]

is this just to make code more organized?

[a-r-r-o-w]

Nope, it's called from the cosmos video-to-video pipeline for rescaling the latents with augment sigmas

[yiyixuxu]

was this float64 before? just curious

[a-r-r-o-w]

It was float32 before. But, to match the sigmas value from the original implementation (which is correct and ours was a little off^*), you need the arange and division by num_train_timesteps to happen in float64. This was done to match the final outputs perfectly (if not, the diff is on the order of 1e-4 to 1e-6, so this change is not particularly required but I think we should keep it)

_{^* If done in float32 (ours), the sigmas start from 79.998 IIRC due to precision issues. If done in float64 and then converted to float32 (current), the sigmas start as expected at 80.0}

[a-r-r-o-w]

@pjannaty@amolfasale@asfiyab-nvidia Hey, could you take a look at the following PRs and let me know if the changes look alright? If all looks good, I can open PRs to the other Cosmos repos with the similar README and weight updates

• https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-7B-Text2World/discussions/9
• https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-7B-Video2World/discussions/2

Additionally, in a follow up PR, we will add support for loading the original format weights directly too

[pjannaty]

@pjannaty@amolfasale@asfiyab-nvidia Hey, could you take a look at the following PRs and let me know if the changes look alright? If all looks good, I can open PRs to the other Cosmos repos with the similar README and weight updates

• https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-7B-Text2World/discussions/9
• https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-7B-Video2World/discussions/2

Additionally, in a follow up PR, we will add support for loading the original format weights directly too

@a-r-r-o-w

Changes look good to me. What a major lift! Let's merge!

[a-r-r-o-w]

@pjannaty@asfiyab-nvidia@amolfasale

Here's the list of all the weight PRs:

• https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-7B-Text2World/discussions/9
• https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-7B-Video2World/discussions/2
• https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-14B-Text2World/discussions/1
• https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-14B-Video2World/discussions/1
• https://huggingface.co/nvidia/Cosmos-1.0-Tokenizer-CV8x8x8/discussions/2

We should be good to merge this code PR already, but users will not be able to download or be able to use the example code snippets until the weight PRs are merged (unless they add revision="refs/pr/<weight_pr_number>" to the model loading code)

[pjannaty]

Thank you for the major lift, team! Let's merge!

[a-r-r-o-w]

@pjannaty The weight PRs cannot be merged by us since we don't have access to the nvidia org. If you or someone with access to the repositories could take a look and merge those, it'd be great

[a-r-r-o-w]

Thank you for merging the PRs @pjannaty!