mstar.model.cosyvoice3.components.dit

mstar.model.cosyvoice3.components.dit#

CosyVoice3’s flow-matching velocity estimator: an F5-TTS style DiT.

x (noisy mel), cond (prompt mel, zero elsewhere), mu (token condition) and the projected x-vector are concatenated per frame, projected to the model width and given a causal convolutional position embedding; 22 AdaLN-zero blocks follow, then a modulated output norm and a projection back to 80 mel bins.

Two quirks of the reference are kept on purpose:

  • rotary positions are applied to the whole to_q/to_k projection before it is split into heads, rotating only its first dim_head (64) channels, i.e. only head 0 sees positions (x-transformers’ partial rotary, interleaved pairs, computed in float32);

  • attention takes a mask, not a length: a block-causal chunk mask when streaming (chunk_attention_mask), the key padding otherwise.

Ported from cosyvoice/flow/DiT/{dit,modules}.py and x-transformers’ RotaryEmbedding / apply_rotary_pos_emb. Parameter paths match flow.pt under decoder.estimator..

Functions

apply_partial_rotary(t, freqs)

Rotate the first freqs.shape[-1] channels of t [B, T, C] in float32 (interleaved pairs); the rest pass through.

chunk_attention_mask(valid, total, chunk, ...)

[B, 1, T, T] bool attention mask (True = attend).

Classes

AdaLayerNormZero(dim)

AdaLayerNormZeroFinal(dim)

Attention(dim, heads, dim_head)

CausalConvPositionEmbedding(dim, ...)

Two left-padded grouped convolutions with Mish.

DiT(config)

DiTBlock(config)

InputEmbedding(config)

TimestepEmbedding(dim, freq_dim)

class mstar.model.cosyvoice3.components.dit.AdaLayerNormZero(dim)[source]#

Bases: Module

Parameters:

dim (int)

forward(x, emb)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:
class mstar.model.cosyvoice3.components.dit.AdaLayerNormZeroFinal(dim)[source]#

Bases: Module

Parameters:

dim (int)

forward(x, emb)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:
Return type:

Tensor

class mstar.model.cosyvoice3.components.dit.Attention(dim, heads, dim_head)[source]#

Bases: Module

Parameters:
forward(x, attn_mask, row_keep, freqs)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:
class mstar.model.cosyvoice3.components.dit.CausalConvPositionEmbedding(dim, kernel_size, groups)[source]#

Bases: Module

Two left-padded grouped convolutions with Mish.

Parameters:
forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Tensor)

Return type:

Tensor

class mstar.model.cosyvoice3.components.dit.DiT(config)[source]#

Bases: Module

Parameters:

config (DiTConfig)

forward(x, mask, mu, t, spks, cond)[source]#

x, mu, cond: [B, 80, T]; mask: [B, 1, T, T] attention mask (chunk_attention_mask; bool, or 0/1 in any float dtype); t: [B]; spks: [B, 80].

Runs in the weights’ dtype and returns the inputs’ dtype; t keeps its own precision through the sinusoid, so a bfloat16 estimator is fed the solve’s float32 time grid, not a rounded one.

Parameters:
Return type:

Tensor

rotary_freqs(seq_len)[source]#

[1, T, dim_head]: each frequency repeated for its pair.

Parameters:

seq_len (int)

Return type:

Tensor

class mstar.model.cosyvoice3.components.dit.DiTBlock(config)[source]#

Bases: Module

Parameters:

config (DiTConfig)

forward(x, t, attn_mask, row_keep, freqs)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class mstar.model.cosyvoice3.components.dit.InputEmbedding(config)[source]#

Bases: Module

Parameters:

config (DiTConfig)

forward(x, cond, mu, spks)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

class mstar.model.cosyvoice3.components.dit.TimestepEmbedding(dim, freq_dim)[source]#

Bases: Module

Parameters:
forward(t)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

t (Tensor)

Return type:

Tensor

mstar.model.cosyvoice3.components.dit.apply_partial_rotary(t, freqs)[source]#

Rotate the first freqs.shape[-1] channels of t [B, T, C] in float32 (interleaved pairs); the rest pass through.

Parameters:
Return type:

Tensor

mstar.model.cosyvoice3.components.dit.chunk_attention_mask(valid, total, chunk, streaming)[source]#

[B, 1, T, T] bool attention mask (True = attend).

Keys past a row’s valid frames are masked. A streaming row also masks every key past the end of the query’s chunk, chunks counted from frame 0 (the reference’s add_optional_chunk_mask with a static chunk size and all left chunks).

Parameters:
Return type:

Tensor