vllm_omni.diffusion.models.soulx_singer.modules ¶

model = FlowMatchingTransformer(cfg=config, **config)

act = nn.GELU()

dwconv = nn.Conv1d(
    dim,
    dim,
    kernel_size=7,
    padding=padding,
    groups=dim,
    dilation=dilation,
)

grn = GRN(intermediate_dim)

norm = nn.LayerNorm(dim, eps=1e-06)

pwconv1 = nn.Linear(dim, intermediate_dim)

pwconv2 = nn.Linear(intermediate_dim, dim)

cond_mlp = nn.Sequential(
    nn.Linear(hidden_size, hidden_size * 4),
    nn.SiLU(),
    nn.Linear(hidden_size * 4, hidden_size),
)

diff_step_embedding = SinusoidalPosEmb(hidden_size)

diff_step_mlp = nn.Sequential(
    nn.Linear(hidden_size, hidden_size * 4),
    nn.SiLU(),
    nn.Linear(hidden_size * 4, hidden_size),
)

embed_tokens = None

layers = nn.ModuleList(
    [
        (LlamaNARDecoderLayer(layer_config, layer_idx=i))
        for i in (range(num_layers))
    ]
)

mel_mlp = nn.Sequential(
    nn.Linear(mel_dim, hidden_size * 4),
    nn.SiLU(),
    nn.Linear(hidden_size * 4, hidden_size),
)

mel_out_mlp = nn.Sequential(
    nn.Linear(hidden_size, hidden_size * 4),
    nn.SiLU(),
    nn.Linear(hidden_size * 4, mel_dim),
)

norm = LlamaAdaptiveRMSNorm(
    hidden_size, dim_cond=hidden_size
)

cfg = cfg

cfg_drop_prob = cfg_drop_prob

cond_codebook_size = cond_codebook_size

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

cond_scale_factor = cond_scale_factor

ctc_layer_index = None

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 = DiffLlama(
    mel_dim=mel_dim,
    hidden_size=hidden_size,
    num_heads=num_heads,
    num_layers=num_layers,
    config=llama_config,
)

do_resampling = True

hidden_size = hidden_size

mel_dim = mel_dim

num_heads = num_heads

num_layers = num_layers

repa_layer_index = None

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 = nn.Sequential(*up_layers)

sigma = sigma

time_scheduler = time_scheduler

use_ctc = 'ctc' in cfg

use_repa = 'repa' in cfg

mel_mean = audio_config.get('mel_mean', -4.92)

mel_var = audio_config.get('mel_var', 8.14)

model = load_mel_spectrogram_from_cfg(audio_config)

model = load_vocos_model(
    ckpt_path=ckpt_path, config=model_cfg
)

device = torch.device(
    device
    or ("cuda" if torch.cuda.is_available() else "cpu")
)

dtype: dtype

fe = WhisperFeatureExtractor.from_pretrained(
    "openai/whisper-base"
)

model = self.model.to(
    device=self.device, dtype=torch.float32
)