docs/voxcpm_patch_flow.md

# 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]` 한 덩어리를 생성**한다.

```text
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. 전체 파이프라인 다이어그램

```text
                ┌─────────────────────────────────────────────────────────────────────────┐
                │                        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 조건 역할을 한다.

```text
                        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`

```python
# 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`

```python
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`

```python
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`

```python
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`

```python
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`

```python
@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`

```python
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`

```python
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`

```python
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` 로
  통일)을 두는 편이 안전하다.