Skip to content

vllm_omni.diffusion.models.soulx_singer.modules.flow_matching

cond_code module-attribute 

cond_code = torch.randn(2, 100, 256).to(device)

device module-attribute 

device = 'cuda'

diff_loss module-attribute 

diff_loss = F.l1_loss(
    flow_pred, flow_gt, reduction="none"
).float() * final_mask.unsqueeze(-1)

diffusion_cond module-attribute 

diffusion_cond = torch.randn(2, 150, 256).to(device)

diffusion_cond_emb module-attribute 

diffusion_cond_emb = model.cond_emb(diffusion_cond)

diffusion_prompt module-attribute 

diffusion_prompt = torch.randn(2, 50, 128).to(device)

final_mask module-attribute 

final_mask = final_mask.squeeze(-1)

flow_gt module-attribute 

flow_gt = x - (1 - 1e-05) * noise

generated module-attribute 

generated = model.reverse_diffusion(
    diffusion_cond_emb,
    diffusion_prompt,
    n_timesteps=n_timesteps,
)

model module-attribute 

model = FlowMatchingTransformer(
    cfg=model_cfg, **model_cfg
).to(device)

model_cfg module-attribute 

model_cfg = {
    "mel_dim": 128,
    "hidden_size": 256,
    "num_layers": 8,
    "num_heads": 8,
    "cfg_drop_prob": 0.2,
    "use_embedding": False,
    "cond_codebook_size": 256,
    "cond_scale_factor": 1,
    "sigma": 1e-05,
    "time_scheduler": "cos",
}

n_timesteps module-attribute 

n_timesteps = 32

outputs module-attribute 

outputs = model(x, x_mask, cond_code)

x module-attribute 

x = torch.randn(2, 100, 128).to(device)

x_mask module-attribute 

x_mask = torch.ones(2, 100).to(device)

FlowMatchingTransformer

Bases: Module

cfg instance-attribute 

cfg = cfg

cfg_drop_prob instance-attribute 

cfg_drop_prob = cfg_drop_prob

cond_codebook_size instance-attribute 

cond_codebook_size = cond_codebook_size

cond_emb instance-attribute 

cond_emb = nn.Embedding(
    cond_codebook_size, self.hidden_size
)

cond_scale_factor instance-attribute 

cond_scale_factor = cond_scale_factor

ctc_layer_index instance-attribute 

ctc_layer_index = None

ctc_mlp_layer instance-attribute 

ctc_mlp_layer = nn.Sequential(
    nn.Linear(hidden_size, hidden_size * 4),
    nn.SiLU(),
    nn.Linear(hidden_size * 4, cfg.ctc.output_dim),
)

diff_estimator instance-attribute 

diff_estimator = DiffLlama(
    mel_dim=mel_dim,
    hidden_size=hidden_size,
    num_heads=num_heads,
    num_layers=num_layers,
    config=llama_config,
)

do_resampling instance-attribute 

do_resampling = True

hidden_size instance-attribute 

hidden_size = hidden_size

mel_dim instance-attribute 

mel_dim = mel_dim

num_heads instance-attribute 

num_heads = num_heads

num_layers instance-attribute 

num_layers = num_layers

repa_layer_index instance-attribute 

repa_layer_index = None

repa_mlp_layer instance-attribute 

repa_mlp_layer = nn.Sequential(
    nn.Linear(hidden_size, hidden_size * 4),
    nn.SiLU(),
    nn.Linear(hidden_size * 4, cfg.repa.output_dim),
)

resampling_layers instance-attribute 

resampling_layers = nn.Sequential(*up_layers)

sigma instance-attribute 

sigma = sigma

time_scheduler instance-attribute 

time_scheduler = time_scheduler

use_ctc instance-attribute 

use_ctc = 'ctc' in cfg

use_repa instance-attribute 

use_repa = 'repa' in cfg

compute_loss 

compute_loss(x, x_mask, cond=None, is_prompt=None)

forward 

forward(
    x: Tensor,
    x_mask: Tensor,
    cond_code: Tensor,
    is_prompt: Tensor | None = None,
)

Parameters:

Name Type Description Default
x Tensor

(B, T, mel_dim)

 required
x_mask Tensor

(B, T)

 required
cond_code Tensor

(B, T), Note that cond_code might be not at 50Hz!

 required

forward_diffusion 

forward_diffusion(x, t, is_prompt=None)

x: (B, T, mel_dim) t: (B,)

loss_t 

loss_t(x, x_mask, t, cond=None, is_prompt=None)

reset_parameters 

reset_parameters()

reverse_diffusion 

reverse_diffusion(
    cond,
    prompt,
    x_mask=None,
    prompt_mask=None,
    n_timesteps=10,
    cfg=1.0,
    rescale_cfg=0.75,
)

reverse_diffusion_v2 

reverse_diffusion_v2(
    cond,
    prompt,
    x_mask=None,
    prompt_mask=None,
    n_timesteps=10,
    cfg=1.0,
    rescale_cfg=0.75,
)