VoxCPM-BACKUP_2 / docs /voxcpm_patch_flow.md

Add files using upload-large-folder tool

9cd59fc verified about 1 month ago

22.8 kB

VoxCPM `patch_size` 흐름과 DiT 아키텍처

본 문서는 /data/mm-llm-backbone_890/personal/sirius/VoxCPM/ 코드에서 확인한 사실을 기반으로, patch_size 가 AudioVAE latent → LM → LocDiT → AudioVAE.decode 까지 어떻게 흐르는지를 정리한 것이다. SimWhisper-Codec 자체가 아니라 VoxCPM 의 latent-AR + flow-matching diffusion 구조에 대한 노트이며, 본 SimWhisper-Codec 레포의 audiocodec/ 와는 직접적인 코드 공유는 없다.

0. 한 줄 요약

patch_size(이하 P)는 DiT 가 한 step 에 예측하는 AudioVAE latent frame 개수다. LM/Residual LM 은 latent frame 단위가 아니라 patch 토큰 단위(T_seq = ceil(T_vae / P))로 돈다. DiT 는 전체 오디오 시퀀스를 한 번에 보지 않고, 각 LM step 마다 local patch [P, D] 한 덩어리를 생성한다.

wav
 -> AudioVAE.encode
 -> latent [B, D, T_vae]
 -> patching
 -> audio_feats [B, T_seq, P, D]
 -> feat_encoder -> [B, T_seq, lm_dim]
 -> base_lm  + residual_lm
 -> dit_hidden [B, T_seq, dit_dim]
 -> flatten   [(B*T_seq), dit_dim]    (μ for DiT)
 -> LocDiT (flow-matching, in_channels=D, sequence length=P)
 -> feat_pred_seq [(B*T_seq), D, P]
 -> unpatch [B, D, T_seq*P] = [B, D, T_vae]
 -> AudioVAE.decode
 -> wav

1. 텐서 shape cheat-sheet

단계	shape	위치
wav	`[B, 1, T_wav]`	input
latent (VAE)	`[B, D=1280, T_vae]` (50 Hz)	`audio_whisper_vae.py:174`
patched feats	`[B, T_seq, P, D]`	`packers.py:51-68`, `voxcpm2.py:421-422`
feat_encoder 출력	`[B, T_seq, lm_dim]`	`local_encoder.py:7-30`
base_lm hidden	`[B, T_seq, lm_dim]` (shift-right)	`voxcpm2.py:325-331`
residual_lm hidden	`[B, T_seq, lm_dim]` (shift-right)	`voxcpm2.py:334-342`
dit_hidden (μ)	`[(B·T_seq), dit_dim]`	`voxcpm2.py:344-345`
feat_gt / feat_cond	`[(B·T_seq), D, P]`	`voxcpm2.py:348-352`
DiT noise z	`[(B·T_seq), D, P]`	`unified_cfm.py:65`
DiT 출력	`[(B·T_seq), D, P]`	`local_dit_v2.py:108-114`
unpatch 결과	`[B, D, T_seq·P] = [B, D, T_vae]`	`voxcpm2.py:382`
recon wav	`[B, 1, T_wav]`	`audio_whisper_vae.py:189`

2. 전체 파이프라인 다이어그램

                ┌─────────────────────────────────────────────────────────────────────────┐
                │                        TRAINING / INFERENCE FORWARD                      │
                └─────────────────────────────────────────────────────────────────────────┘

   wav  [B, 1, T_wav]                                       AudioVAE = WhisperAudioVAE
     │                                                          sample_rate = 16000
     │                                                          hop_length  = 320  ── 50 Hz latent
     │   pad to multiple of patch_len                           chunk_size  = 320  ★ hop과 동일
     │     train: patch_len = audio_vae.hop_length * P           latent_dim  = 1280
     │              packers.py:22
     │     infer: patch_len = self.patch_size * self.chunk_size
     │              voxcpm2.py:416
     ▼
   AudioVAE.encode(wav, sr=16000)                          audio_whisper_vae.py:132-174
     │   feats = Whisper.encoder(mel).transpose(1,2)
     ▼
   latent  z : [B, D=1280, T_vae]   (50 Hz)
     │
     │   ── PATCHING ──
     │   train (packers + collate):
     │     feat = z.transpose(1,2)                # [B, T_vae, D]
     │     view → [B, T_seq, P, D]
     │   infer (_encode_wav):                     voxcpm2.py:421-422
     │     feat.view(D, -1, P).permute(1,2,0)     # [T_seq, P, D]
     ▼
   audio_feats : [B, T_seq, P, D]                          T_seq = ceil(T_vae / P)

   ┌───────────────────────────────────────────────────────────────────────────────────┐
   │                          LM-side  (patch-token sequence, T_seq tokens)             │
   ├───────────────────────────────────────────────────────────────────────────────────┤
   │                                                                                    │
   │   audio_feats [B, T, P, D]                                                         │
   │       │                                                                            │
   │       │  feat_encoder = VoxCPMLocEnc                local_encoder.py:7-30          │
   │       │   • in_proj  : Linear(D=64, H_lm)                                          │
   │       │   • prepend learnable [CLS] token                                          │
   │       │   • MiniCPMModel (bi-dir, is_causal=False)                                 │
   │       │   • take CLS:  [(B*T), P+1, H] → [(B*T), H]                                │
   │       ▼                                                                            │
   │   feat_embed [B, T, H_lm]   →  enc_to_lm_proj  →  [B, T, lm_hidden]                │
   │       │                                                                            │
   │       │  combined_embed = text_mask·text_embed + audio_mask·feat_embed             │
   │       ▼                                                                            │
   │   base_lm (MiniCPM, is_causal=True)              voxcpm2.py:325-326                │
   │       │ enc_outputs [B, T, lm_dim]                                                 │
   │       │   ↳ fsq_layer on audio positions       voxcpm2.py:327                      │
   │       │   ↳ shift-right ⇒ lm_hidden                                                │
   │       ▼                                                                            │
   │   residual path:                                voxcpm2.py:334-342                 │
   │     residual_inputs = fusion_concat_proj([enc_outputs, audio_mask·feat_embed])     │
   │     residual_lm(... is_causal=True) ⇒ residual_outputs [B, T, lm_dim]              │
   │     shift-right ⇒ residual_hidden                                                  │
   │       │                                                                            │
   │       ▼                                                                            │
   │   dit_hidden = cat( lm_to_dit_proj(lm_hidden),                                     │
   │                     res_to_dit_proj(residual_hidden) )   [B, T, dit_dim]           │
   │                                                                                    │
   └───────────────────────────────────────────────────────────────────────────────────┘
                                              │
                       ── FLATTEN ──         voxcpm2.py:345-352
                       dit_hidden : [B, T, dit_dim] ──► [(B*T), dit_dim]            (mu)
                       feat_gt    : [B, T, P, D]   ──► [(B*T), P, D]  →ᵀ→ [(B*T), D, P]
                       feat_cond  : shift-right of audio_feats, same reshape   [(B*T), D, P]
                                              │
                                              ▼
   ┌───────────────────────────────────────────────────────────────────────────────────┐
   │                       DiT / Flow-matching  (local, per-patch)                     │
   ├───────────────────────────────────────────────────────────────────────────────────┤
   │                                                                                    │
   │   UnifiedCFM.forward(mu, n_steps, P, cond, …)        unified_cfm.py:54-79         │
   │     z = randn((b, in_channels=D, t=P))                                            │
   │     iterate t_span via solve_euler                                                │
   │                                                                                    │
   │   estimator = VoxCPMLocDiT                            local_dit_v2.py:50-116      │
   │     x    [N, D, P]    cond [N, D, P]    mu [N, H]    t,dt [N]                     │
   │       │                                                                            │
   │       │ in_proj  : Linear(D=64 → H)   on xᵀ  ⇒  [N, P, H]                          │
   │       │ cond_proj: Linear(D=64 → H)   on condᵀ ⇒ [N, P, H]                         │
   │       │ time emb : SinPosEmb(t) + SinPosEmb(dt) → time_mlp  ⇒ [N, H]               │
   │       │ mu       : view to [N, 1, H]                                              │
   │       │ seq      : cat([mu, t_token, cond, x], dim=1)   length 1+1+P+P            │
   │       │ MiniCPMModel(seq, is_causal=False)                                        │
   │       │ slice last P tokens → out_proj → [N, P, D] →ᵀ→ [N, D, P]                  │
   │       ▼                                                                            │
   │   feat_pred_seq : [(B*T), D, P]                                                   │
   │                                                                                    │
   │   Training loss:  UnifiedCFM.compute_loss(x1=feat_gt, mu, cond)                   │
   │                                                       unified_cfm.py:174-248      │
   │     y = (1-r)·z + r·x1   v = x1 - z                                               │
   │     mse(model(y,r,t), v) , masked by tgt_mask                                     │
   │                                                                                    │
   └───────────────────────────────────────────────────────────────────────────────────┘
                                              │
                       ── UNPATCH ──          voxcpm2.py:382
                       feat_pred_seq [(B*T), D, P]
                          → transpose to [(B*T), P, D]
                          → rearrange "(b t) p d -> b d (t p)"
                          ⇒ feat_pred : [B, D, T*P]   (= [B, D, T_vae])
                                              │
                                              ▼
                       AudioVAE.decode(feat_pred)                audio_whisper_vae.py:176
                          feats = z.transpose(1,2).unsqueeze(2)   # [B, T, 1, D]
                          recon = generator.inference(feats)      # [B, 1, T_wav]
                                              │
                                              ▼
                                        ŵav  [B, 1, T_wav]

3. LocDiT 한 step 내부 (token 단위)

mu 는 한 개 토큰 으로, timestep 도 별도의 한 개 토큰 으로 prepend 된다. 추가로 이전 patch (feat_cond) 가 P 토큰의 prefix 로 항상 함께 들어가서 in-context 조건 역할을 한다.

                        Local sequence (length = 1 + 1 + P + P = 2 + 2P)
                   ┌────────┬────────┬─────────────────┬─────────────────┐
   input tokens ── │   μ    │   t̂    │  cond[0..P-1]   │   x[0..P-1]    │
                   │ [1,H]  │ [1,H]  │     [P,H]       │     [P,H]      │
                   └────────┴────────┴─────────────────┴─────────────────┘
                        │       │             │                 │
                LM hidden    sinusoidal     prefix          noisy / current
                (== μ)       time emb       (prev patch)    sample x_r
                             + Δt emb

                           │  bi-directional MiniCPM transformer  │
                           ▼
                   ┌────────┬────────┬─────────────────┬─────────────────┐
   output slice ── │   ·    │   ·    │       ·         │  out[0..P-1]   │ → out_proj
                   └────────┴────────┴─────────────────┴─────────────────┘
                                                              │
                                                              ▼
                                                       [N, P, D] →ᵀ→ [N, D, P]
                                                                     = predicted velocity

4. 핵심 코드 인용

4.1 Patching (latent → audio_feats)

Training packer — VoxCPM/src/voxcpm/training/packers.py:22, 51-68

# packers.py:22
self.patch_len = audio_vae.hop_length * self.patch_size

# packers.py:51-68
def encode_audio(self, wav: torch.Tensor):
    wav = wav.unsqueeze(0).unsqueeze(1)              # [1, 1, T]
    if wav.size(-1) % self.patch_len != 0:
        padding_size = self.patch_len - wav.size(-1) % self.patch_len
        wav = torch.nn.functional.pad(wav, (0, padding_size))
    with torch.no_grad():
        z = self.audio_vae.encode(wav, self.audio_vae.sample_rate)  # [1, D, T']
        feat = z.transpose(1, 2)                                    # [1, T', D]
    return feat

Inference path — VoxCPM/src/voxcpm/model/voxcpm2.py:416-422

patch_len = self.patch_size * self.chunk_size
if audio.size(1) % patch_len != 0:
    padding_size = patch_len - audio.size(1) % patch_len
    pad = (padding_size, 0) if padding_mode == "left" else (0, padding_size)
    audio = torch.nn.functional.pad(audio, pad)
feat = self.audio_vae.encode(audio.to(self.device), self._encode_sample_rate).cpu()
return feat.view(self.audio_vae.latent_dim, -1, self.patch_size).permute(1, 2, 0)  # [T_seq, P, D]

4.2 feat_encoder (`VoxCPMLocEnc`)

VoxCPM/src/voxcpm/modules/locenc/local_encoder.py:7-30

class VoxCPMLocEnc(nn.Module):
    def __init__(self, config: MiniCPM4Config, input_dim: int = 64):
        super().__init__()
        self.special_token = nn.Parameter(torch.randn(1, 1, 1, config.hidden_size))
        self.in_proj = nn.Linear(input_dim, config.hidden_size, bias=True)
        self.encoder = MiniCPMModel(config)

    def forward(self, x):  # x: [B, T, P, D]
        B, T, P, D = x.shape
        x = self.in_proj(x)                                  # [B, T, P, H]
        special_tokens = self.special_token.expand(B, T, 1, -1)
        x = torch.cat([special_tokens, x], dim=2)
        x = rearrange(x, "b t p c -> (b t) p c")
        outputs, _ = self.encoder(x, is_causal=False)
        cls_output = outputs[:, 0, :]
        return rearrange(cls_output, "(b t) c -> b t c", b=B)

→ 단순 Linear 가 아니라 mini-transformer + [CLS] 다.

4.3 LM + Residual LM

VoxCPM/src/voxcpm/model/voxcpm2.py:319-345

B, T, P, D = audio_feats.shape
feat_embed = self.feat_encoder(audio_feats)                  # [B, T, H_enc]
feat_embed = self.enc_to_lm_proj(feat_embed)                 # [B, T, lm_dim]

text_embed = self.base_lm.embed_tokens(text_tokens) * scale_emb
combined_embed = text_mask.unsqueeze(-1) * text_embed + audio_mask.unsqueeze(-1) * feat_embed
enc_outputs, _ = self.base_lm(inputs_embeds=combined_embed, is_causal=True)
enc_outputs = self.fsq_layer(enc_outputs) * audio_mask.unsqueeze(-1) \
              + enc_outputs * text_mask.unsqueeze(-1)
lm_hidden = torch.cat((torch.zeros_like(enc_outputs[:, 0:1, :]), enc_outputs[:, :-1, :]), dim=1)

residual_inputs = self.fusion_concat_proj(
    torch.cat((enc_outputs, audio_mask.unsqueeze(-1) * feat_embed), dim=-1)
)
residual_outputs, _ = self.residual_lm(inputs_embeds=residual_inputs, is_causal=True)
residual_hidden = torch.cat(
    (torch.zeros_like(residual_outputs[:, 0:1, :]), residual_outputs[:, :-1, :]), dim=1,
)

dit_hidden = torch.cat(
    (self.lm_to_dit_proj(lm_hidden), self.res_to_dit_proj(residual_hidden)), dim=-1
)
dit_hidden = rearrange(dit_hidden, "b t c -> (b t) c")        # [(B*T), dit_dim]

4.4 DiT 입력 만들기

VoxCPM/src/voxcpm/model/voxcpm2.py:348-358

feat_gt = rearrange(audio_feats.to(target_dtype), "b t p d -> (b t) p d")
feat_cond = torch.cat(
    (torch.zeros_like(audio_feats[:, 0:1, ...]), audio_feats[:, :-1, ...]), dim=1,
)
feat_cond = rearrange(feat_cond.to(target_dtype), "b t p d -> (b t) p d")

loss_seq_mask = loss_mask.unsqueeze(-1).repeat(1, 1, self.patch_size)
loss_seq_mask = rearrange(loss_seq_mask, "b t p -> (b t) p 1").to(target_dtype)

4.5 UnifiedCFM — noise/sampling

VoxCPM/src/voxcpm/modules/locdit/unified_cfm.py:54-79

@torch.inference_mode()
def forward(self, mu, n_timesteps, patch_size, cond, ...):
    b, _ = mu.shape
    t = patch_size
    z = torch.randn((b, self.in_channels, t),
                    device=mu.device, dtype=mu.dtype) * temperature
    t_span = torch.linspace(1, 0, n_timesteps + 1, device=mu.device, dtype=mu.dtype)
    t_span = t_span + sway_sampling_coef * (torch.cos(torch.pi / 2 * t_span) - 1 + t_span)
    return self.solve_euler(x=z, t_span=t_span, mu=mu, cond=cond, ...)

학습 손실 — unified_cfm.py:174-248 의 compute_loss(x1, mu, cond, tgt_mask, progress).

4.6 LocDiT (`VoxCPMLocDiT`)

VoxCPM/src/voxcpm/modules/locdit/local_dit_v2.py:50-116

def forward(self, x, mu, t, cond, dt):
    """
    x:    [N, C, P]   noisy / current sample
    mu:   [N, C]      LM hidden
    t:    [N]         diffusion timestep
    cond: [N, C, P]   previous-patch condition
    dt:   [N]         delta for mean velocity
    """
    x    = self.in_proj(x.transpose(1, 2).contiguous())     # [N, P, H]
    cond = self.cond_proj(cond.transpose(1, 2).contiguous())# [N, P, H]
    prefix = cond.size(1)

    t  = self.time_mlp(self.time_embeddings(t).to(x.dtype))
    dt = self.delta_time_mlp(self.time_embeddings(dt).to(x.dtype))
    t  = t + dt

    mu = mu.view(x.size(0), -1, x.size(-1))                  # [N, 1, H]
    x  = torch.cat([mu, t.unsqueeze(1), cond, x], dim=1)     # [N, 1+1+P+P, H]
    hidden, _ = self.decoder(x, is_causal=False)
    hidden = hidden[:, prefix + mu.size(1) + 1:, :]
    hidden = self.out_proj(hidden)                           # [N, P, D]
    return hidden.transpose(1, 2).contiguous()               # [N, D, P]

4.7 Unpatch

VoxCPM/src/voxcpm/model/voxcpm2.py:382

feat_pred = rearrange(
    feat_pred_seq.transpose(1, 2),
    "(b t) d p -> b d (t p)",
    b=B, p=self.patch_size,
)   # [B, D, T_vae]

스트리밍 inference 시에는 한 step 단위 [B, 1, P, D] → [B, D, P] (voxcpm2.py:1083), non-streaming 시에는 T_seq 만큼 모은 뒤 한 번에 (b t) p d -> b d (t p) (voxcpm2.py:1106).

4.8 WhisperAudioVAE

VoxCPM/src/voxcpm/modules/audiovae/audio_whisper_vae.py:28-87

self.sample_rate = 16000
self.hop_length = 320              # 16000 / 50 Hz Whisper latent
self.chunk_size = 320              # ★ VoxCPM 의 patch_len 계산에서 hop_length 와 동일해야 함
self.latent_dim = 1280             # Whisper-large hidden size
self.max_audio_samples = self.sample_rate * 30   # SimWhisper extractor 30s 한계

encode 는 [B, D=1280, T_lat], decode 는 [B, 1, T_wav] 를 돌려준다.

5. `chunk_size` vs `hop_length` 정리

위치	수식	사용 attribute
Training packer (`packers.py:22`)	`patch_len = audio_vae.hop_length * patch_size`	`hop_length`
Inference `_encode_wav` (`voxcpm2.py:416`)	`patch_len = self.patch_size * self.chunk_size`	`chunk_size`

→ 두 경로가 같은 의미("VAE 한 frame 당 raw sample 수 × P")여야 하므로 chunk_size == hop_length 여야 함. WhisperAudioVAE 에서는 둘 다 320 으로 하드코드되어 있어 실효적으로 일치 (audio_whisper_vae.py:84-85).

30 초 제한은 chunk_size 가 아니라 별도 max_audio_samples = sample_rate * 30 이다. 의미가 다른 두 값을 하나에 묶지 않도록 주의.

6. 사용자 요약 vs 실제 코드 — 짚어둘 점

feat_encoder 는 Linear 가 아니라 mini-transformer + CLS 이다 (VoxCPMLocEnc). patch 내부 P frame 을 양방향 self-attention 으로 섞은 뒤 [CLS] 토큰만 빼서 [B, T, lm_dim] 을 만든다.
Residual LM 은 base LM 과 별개 모듈 (voxcpm2.py:334-342). base LM 출력 + audio embed 를 fusion_concat_proj 으로 합쳐 다시 causal LM 한 번 더 돌려서 dit_hidden 의 절반을 만들고, 나머지 절반은 base LM 출력. 두 hidden 을 채널 축으로 concat 하여 최종 dit_dim 을 채운다.
DiT prefix 는 [μ, t̂, cond, x] 로 4 종류 토큰 이 prepend 된다. μ 와 timestep 은 합쳐진 단일 토큰이 아니라 별개의 두 토큰. 추가로 이전 patch(feat_cond) 가 P 개 prefix 토큰으로 들어가서 in-context conditioning 을 수행.
노이즈 z 는 randn((b, in_channels=D, P)) — DiT 의 "local time" 길이는 정확히 patch_size 이다.
chunk_size == hop_length 동치는 WhisperAudioVAE 의 현재 하드코드 에서만 성립. 새 AudioVAE 를 붙일 때는 두 값이 분리되지 않도록 단언문(또는 chunk_size = hop_length 로 통일)을 두는 편이 안전하다.

VoxCPM patch_size 흐름과 DiT 아키텍처