Output:
Chunk ID: e380fa
Wall time: 0.0514 seconds
Process exited with code 0
Original token count: 3706
Output:
     1	import os
     2
     3	import random
     4	import hydra
     5	import numpy as np
     6	import librosa
     7	import torch
     8	import torch.nn as nn
     9	import torch.nn.functional as F
    10	import torch.optim as optim
    11	import pytorch_lightning as pl
    12	from vq import CodecEncoder,  CodecDecoderVocos
    13	from module import HiFiGANMultiPeriodDiscriminator, SpecDiscriminator
    14	from criterions import GANLoss, MultiResolutionMelSpectrogramLoss, MultiResolutionSTFTLoss
    15	from common.schedulers import WarmupLR
    16	from transformers import AutoModel
    17	from vq.module import SemanticDecoder,SemanticEncoder
    18	from transformers import AutoFeatureExtractor, Wav2Vec2BertModel
    19	import sys
    20	sys.path.append('./eval_tools/tools/speaker_verification')    # We use wavlm_large_finetune as a vadidation metric during training, https://github.com/microsoft/UniSpeech/tree/main/downstreams/speaker_verification
    21	from  verification import init_model
    22	model_spk = init_model('wavlm_large','/aifs4su/data/zheny/models_fd_ckpt/wavlm_large_finetune.pth')
    23
    24
    25
    26	class CodecLightningModule(pl.LightningModule):
    27	    def __init__(self, cfg):
    28	        super().__init__()
    29	        self.cfg = cfg
    30	        self.ocwd = hydra.utils.get_original_cwd()
    31	        self.construct_model()
    32	        self.construct_criteria()
    33	        self.save_hyperparameters()
    34	        self.automatic_optimization = False
    35
    36	    def construct_model(self):
    37	        # 初始化 Codec Encoder
    38
    39	        enccfg = self.cfg.model.codec_encoder
    40
    41
    42	        self.CodecEnc = CodecEncoder(
    43
    44	            ngf=enccfg.ngf,
    45	            up_ratios=enccfg.up_ratios,
    46	            dilations=enccfg.dilations,
    47	            hidden_dim=enccfg['hidden_dim'],
    48	            depth=enccfg['depth'],
    49	            heads=enccfg['heads'],
    50	            pos_meb_dim=enccfg['pos_meb_dim'],
    51	        )
    52
    53	        # 初始化 Codec Decoder
    54	        deccfg = self.cfg.model.codec_decoder
    55
    56	        self.generator = CodecDecoderVocos(
    57	            hidden_dim=deccfg.hidden_dim,
    58	            depth=deccfg.depth,
    59	            heads=deccfg.heads,
    60	            pos_meb_dim=deccfg.pos_meb_dim,
    61	            hop_length=320,
    62	            vq_num_quantizers=deccfg.vq_num_quantizers,  # VQ 量化器数量
    63	            vq_dim=deccfg.vq_dim,                   # VQ 维度
    64	            vq_commit_weight=deccfg.vq_commit_weight,    # VQ 提交权重
    65	            vq_weight_init=deccfg.vq_weight_init,         # VQ 权重初始化
    66	            vq_full_commit_loss=deccfg.vq_full_commit_loss,  # 是否使用完整的提交损失
    67	            codebook_size=deccfg.codebook_size,            # 码本大小
    68	            codebook_dim=deccfg.codebook_dim ,              # 码本维度
    69	                  # 隐藏层维度
    70	        )
    71
    72
    73
    74	        # 初始化 MultiPeriod Discriminator
    75	        mpdcfg = self.cfg.model.mpd
    76	        self.discriminator = HiFiGANMultiPeriodDiscriminator(
    77	            periods=mpdcfg.periods,
    78	            max_downsample_channels=mpdcfg.max_downsample_channels,
    79	            channels=mpdcfg.channels,
    80	            channel_increasing_factor=mpdcfg.channel_increasing_factor,
    81	        )
    82
    83	        # 初始化 Spectral Discriminator
    84	        mstftcfg = self.cfg.model.mstft
    85	        self.spec_discriminator = SpecDiscriminator(
    86	            stft_params=mstftcfg.stft_params,
    87	            in_channels=mstftcfg.in_channels,
    88	            out_channels=mstftcfg.out_channels,
    89	            kernel_sizes=mstftcfg.kernel_sizes,
    90	            channels=mstftcfg.channels,
    91	            max_downsample_channels=mstftcfg.max_downsample_channels,
    92	            downsample_scales=mstftcfg.downsample_scales,
    93	            use_weight_norm=mstftcfg.use_weight_norm,
    94	        )
    95
    96
    97
    98	        # 单独编译需要优化的子模块
    99	        # self.CodecEnc = torch.compile(self.CodecEnc)
   100	        # self.generator.backbone = torch.compile(self.generator )
   101	        # self.mel_conv = torch.compile(self.mel_conv)
   102
   103	        self.model_spk = model_spk .eval()
   104
   105	        # self.semantic_model = AutoModel.from_pretrained("microsoft/wavlm-large")
   106	        # self.semantic_model.eval()
   107	        # self.semantic_model.requires_grad_(False)
   108
   109
   110	        self.fc_prior = nn.Linear(1024 + 1024, deccfg.vq_dim,   )
   111	        self.fc_post_a = nn.Linear(deccfg.vq_dim,  deccfg.hidden_dim )
   112	        self.fc_post_s = nn.Linear(deccfg.vq_dim,   1024)
   113
   114	        self.SemanticDecoder_module = SemanticDecoder(1024, 1024, 1024)
   115	        self.SemanticEncoder_module = SemanticEncoder(1024, 1024, 1024)
   116	        self.semantic_model = Wav2Vec2BertModel.from_pretrained("facebook/w2v-bert-2.0", output_hidden_states=True)
   117	        self.semantic_model.eval()
   118	        self.semantic_model.requires_grad_(False)
   119	        # self.register_buffer('mel_basis', mel_basis)
   120
   121	        # self.perception_model = AutoModel.from_pretrained("facebook/wav2vec2-large-xlsr-53")
   122	        # self.perception_model.eval()
   123	        # self.perception_model.requires_grad_(False)
   124
   125	    def construct_criteria(self):
   126	        cfg = self.cfg.train
   127	        self.criteria = nn.ModuleDict()
   128	        if cfg.use_mel_loss:
   129	            self.criteria['mel_loss'] = MultiResolutionMelSpectrogramLoss(sample_rate=self.cfg.preprocess.audio.sr)
   130	        if cfg.use_stft_loss:
   131	            self.criteria['stft_loss'] = MultiResolutionSTFTLoss(
   132	                fft_sizes=cfg.stft_loss_params.fft_sizes,
   133	                hop_sizes=cfg.stft_loss_params.hop_sizes,
   134	                win_sizes=cfg.stft_loss_params.win_lengths
   135	            )
   136	        if cfg.use_feat_match_loss:
   137	            self.criteria['fm_loss'] = nn.L1Loss()
   138	        self.criteria['gan_loss'] = GANLoss()
   139	        self.criteria['l1_loss'] = nn.L1Loss()
   140	        self.criteria['l2_loss'] = nn.MSELoss()
   141	        print(self.criteria)
   142
   143
   144
   145
   146	    def forward(self, batch):
   147	        wav = batch['wav']
   148	        feats= batch['feats']
   149
   150	        vq_emb = self.CodecEnc(wav.unsqueeze(1))
   151	        vq_emb = vq_emb.transpose(1, 2)
   152
   153	        with torch.no_grad():
   154	            semantic_target = self.semantic_model(feats[:,0,:,:])
   155
   156	            semantic_target = semantic_target.hidden_states[16]
   157	            semantic_target = semantic_target.detach()
   158
   159	        semantic_target = semantic_target.transpose(1, 2)
   160	        semantic_target_processed = self.SemanticEncoder_module(semantic_target)
   161	        # 拼接语义嵌入和编码器输出
   162	        vq_emb = torch.cat([semantic_target_processed, vq_emb], dim=1)
   163	        vq_emb = self.fc_prior(vq_emb.transpose(1, 2)).transpose(1, 2)
   164
   165	        vq_post_emb, vq_code, vq_loss = self.generator(vq_emb, vq=True)
   166	        semantic_recon = self.fc_post_s(vq_post_emb.transpose(1, 2)).transpose(1, 2)
   167	        semantic_recon = self.SemanticDecoder_module(semantic_recon)
   168
   169
   170	        y_ ,_ = self.generator(
   171	            self.fc_post_a(vq_post_emb.transpose(1, 2)) ,
   172	            vq=False
   173	        )
   174	        y = wav.unsqueeze(1)
   175
   176	        # gt_perceptual = self.perception_model(wav.squeeze(1), output_hidden_states=True) .hidden_states
   177	        # gen_perceptual = self.perception_model(y_.squeeze(1), output_hidden_states=True) .hidden_states
   178
   179	        # gt_perceptual_se = gt_perceptual[10:22]
   180	        # gen_perceptual_se = gen_perceptual[10:22]
   181
   182	        # perceptual_se_loss = [tensor1 - tensor2 for tensor1, tensor2 in zip(gt_perceptual_se, gen_perceptual_se)]
   183
   184	        # # 使用列表推导式逐元素相减
   185	        # perceptual_se_loss_l2 = [F.mse_loss(tensor1.detach(), tensor2) for tensor1, tensor2 in zip(gt_perceptual_se, gen_perceptual_se)]
   186	        # perceptual_se_loss_l2 =torch.stack(perceptual_se_loss_l2).mean()
   187	        output = {
   188	            'gt_wav': y,
   189	            'gen_wav': y_,
   190	            'vq_loss': vq_loss,
   191	            'vq_code': vq_code,
   192	            'semantic_recon_loss': F.mse_loss(semantic_recon, semantic_target),
   193	            # 'perceptual_se_loss_l2': perceptual_se_loss_l2,
   194
   195	        }
   196	        return output
   197
   198	    @torch.inference_mode()
   199	    def inference(self, wav):
   200	        vq_emb = self.CodecEnc(wav.unsqueeze(1))
   201	        vq_post_emb, vq_code, vq_loss = self.generator(vq_emb, vq=True)
   202	        y_ = self.generator(vq_post_emb, vq=False).squeeze(1)  # [B, T]
   203	        return y_
   204
   205	    def compute_disc_loss(self, batch, output):
   206	        y, y_ = output['gt_wav'], output['gen_wav']
   207	        y_ = y_.detach()
   208	        p = self.discriminator(y)
   209	        p_ = self.discriminator(y_)
   210
   211	        real_loss_list, fake_loss_list = [], []
   212	        for i in range(len(p)):
   213	            real_loss, fake_loss = self.criteria['gan_loss'].disc_loss(p[i][-1], p_[i][-1])
   214	            real_loss_list.append(real_loss)
   215	            fake_loss_list.append(fake_loss)
   216
   217	        if hasattr(self, 'spec_discriminator'):
   218	            sd_p = self.spec_discriminator(y)
   219	            sd_p_ = self.spec_discriminator(y_)
   220
   221	            for i in range(len(sd_p)):
   222	                real_loss, fake_loss = self.criteria['gan_loss'].disc_loss(sd_p[i][-1], sd_p_[i][-1])
   223	                real_loss_list.append(real_loss)
   224	                fake_loss_list.append(fake_loss)
   225
   226	        real_loss = sum(real_loss_list)
   227	        fake_loss = sum(fake_loss_list)
   228
   229	        disc_loss = real_loss + fake_loss
   230	        disc_loss = self.cfg.train.lambdas.lambda_disc * disc_loss
   231
   232	        output = {
   233	            'real_loss': real_loss,
   234	            'fake_loss': fake_loss,
   235	            'disc_loss': disc_loss,
   236	        }
   237	        return output
   238
   239	    def compute_gen_loss(self, batch, output):
   240	        y, y_ = output['gt_wav'], output['gen_wav']
   241	        vq_loss, vq_code = output['vq_loss'], output['vq_code']
   242	        semantic_recon_loss = output['semantic_recon_loss']
   243	        # perceptual_se_loss_l2 = output['perceptual_se_loss_l2']
   244	        # x_feat_recon_loss = output['x_feat_recon_loss']
   245	        gen_loss = 0.0
   246	        self.set_discriminator_gradients(False)
   247	        output_dict = {}
   248	        cfg = self.cfg.train
   249
   250	        # Mel spectrogram loss
   251	        if cfg.use_mel_loss:
   252	            mel_loss = self.criteria['mel_loss'](y_.squeeze(1), y.squeeze(1))
   253	            gen_loss += mel_loss * cfg.lambdas.lambda_mel_loss
   254	            output_dict['mel_loss'] = mel_loss
   255
   256	        # GAN loss
   257	        p_ = self.discriminator(y_)
   258	        adv_loss_list = []
   259	        for i in range(len(p_)):
   260	            adv_loss_list.append(self.criteria['gan_loss'].gen_loss(p_[i][-1]))
   261	        if hasattr(self, 'spec_discriminator'):
   262	            sd_p_ = self.spec_discriminator(y_)
   263	            for i in range(len(sd_p_)):
   264	                adv_loss_list.append(self.criteria['gan_loss'].gen_loss(sd_p_[i][-1]))
   265	        adv_loss = sum(adv_loss_list)
   266	        gen_loss += adv_loss * cfg.lambdas.lambda_adv
   267	        output_dict['adv_loss'] = adv_loss
   268
   269	        # Feature Matching loss
   270	        if cfg.use_feat_match_loss:
   271	            fm_loss = 0.0
   272	            with torch.no_grad():
   273	                p = self.discriminator(y)
   274	            for i in range(len(p_)):
   275	                for j in range(len(p_[i]) - 1):
   276	                    fm_loss += self.criteria['fm_loss'](p_[i][j], p[i][j].detach())
   277	            gen_loss += fm_loss * cfg.lambdas.lambda_feat_match_loss
   278	            output_dict['fm_loss'] = fm_loss
   279	            if hasattr(self, 'spec_discriminator'):
   280	                spec_fm_loss = 0.0
   281	                with torch.no_grad():
   282	                    sd_p = self.spec_discriminator(y)
   283	                for i in range(len(sd_p_)):
   284	                    for j in range(len(sd_p_[i]) - 1):
   285	                        spec_fm_loss += self.criteria['fm_loss'](sd_p_[i][j], sd_p[i][j].detach())
   286	                gen_loss += spec_fm_loss * cfg.lambdas.lambda_feat_match_loss
   287	                output_dict['spec_fm_loss'] = spec_fm_loss
   288
   289	        # VQ loss
   290	        if vq_loss is not None:
   291	            vq_loss = sum(vq_loss)
   292	            gen_loss += vq_loss
   293	            output_dict['vq_loss'] = vq_loss
   294
   295	        # Semantic reconstruction loss
   296	        output_dict['semantic_recon_loss'] = semantic_recon_loss
   297	        gen_loss += output_dict['semantic_recon_loss'] * cfg.lambdas.lambda_semantic_loss
   298
   299	        # Perceptual loss
   300	        # output_dict['perceptual_se_loss_l2'] = perceptual_se_loss_l2
   301	        # gen_loss += output_dict['perceptual_se_loss_l2'] * cfg.lambdas.lambda_perceptual_loss
   302
   303	        self.set_discriminator_gradients(True)
   304	        output_dict['gen_loss'] = gen_loss
   305	        return output_dict
   306
   307	    def training_step(self, batch, batch_idx):
   308	        output = self(batch)
   309
   310	        gen_opt, disc_opt = self.optimizers()
   311	        gen_sche, disc_sche = self.lr_schedulers()
   312
   313	        # 训练判别器
   314	        disc_losses = self.compute_disc_loss(batch, output)
   315	        disc_loss = disc_losses['disc_loss']
   316	        disc_opt.zero_grad()
   317	        self.manual_backward(disc_loss)
   318	        self.clip_gradients(
   319	            disc_opt,
   320	            gradient_clip_val=self.cfg.train.disc_grad_clip,


# Plan Mode (Conversational)

You work in 3 phases, and you should *chat your way* to a great plan before finalizing it. A great plan is very detailed—intent- and implementation-wise—so that it can be handed to another engineer or agent to be implemented right away. It must be **decision complete**, where the implementer does not need to make any decisions.

## Mode rules (strict)

You are in **Plan Mode** until a developer message explicitly ends it.

Plan Mode is not changed by user intent, tone, or imperative language. If a user asks for execution while still in Plan Mode, treat it as a request to **plan the execution**, not perform it.

## Plan Mode vs update_plan tool

Plan Mode is a collaboration mode that can involve requesting user input and eventually issuing a `<proposed_plan>` block.

Separately, `update_plan` is a checklist/progress/TODOs tool; it does not enter or exit Plan Mode. Do not confuse it with Plan mode or try to use it while in Plan mode. If you try to use `update_plan` in Plan mode, it will return an error.

## Execution vs. mutation in Plan Mode

You may explore and execute **non-mutating** actions that improve the plan. You must not perform **mutating** actions.

### Allowed (non-mutating, plan-improving)

Actions that gather truth, reduce ambiguity, or validate feasibility without changing repo-tracked state. Examples:

* Reading or searching files, configs, schemas, types, manifests, and docs
* Static analysis, inspection, and repo exploration
* Dry-run style commands when they do not edit repo-tracked files
* Tests, builds, or checks that may write to caches or build artifacts (for example, `target/`, `.cache/`, or snapshots) so long as they do not edit repo-tracked files

### Not allowed (mutating, plan-executing)

Actions that implement the plan or change repo-tracked state. Examples:

* Editing or writing files
* Running formatters or linters that rewrite files
* Applying patches, migrations, or codegen that updates repo-tracked files
* Side-effectful commands whose purpose is to carry out the plan rather than refine it

When in doubt: if the action would reasonably be described as "doing the work" rather than "planning the work," do not do it.

## PHASE 1 — Ground in the environment (explore first, ask second)

Begin by grounding yourself in the actual environment. Eliminate unknowns in the prompt by discovering facts, not by asking the user. Resolve all questions that can be answered through exploration or inspection. Identify missing or ambiguous details only if they cannot be derived from the environment. Silent exploration between turns is allowed and encouraged.

Before asking the user any question, perform at least one targeted non-mutating exploration pass (for example: search relevant files, inspect likely entrypoints/configs, confirm current implementation shape), unless no local environment/repo is available.

Exception: you may ask clarifying questions about the user's prompt before exploring, ONLY if there are obvious ambiguities or contradictions in the prompt itself. However, if ambiguity might be resolved by exploring, always prefer exploring first.

Do not ask questions that can be answered from the repo or system (for example, "where is this struct?" or "which UI component should we use?" when exploration can make it clear). Only ask once you have exhausted reasonable non-mutating exploration.

## PHASE 2 — Intent chat (what they actually want)

* Keep asking until you can clearly state: goal + success criteria, audience, in/out of scope, constraints, current state, and the key preferences/tradeoffs.
* Bias toward questions over guessing: if any high-impact ambiguity remains, do NOT plan yet—ask.

## PHASE 3 — Implementation chat (what/how we’ll build)

* Once intent is stable, keep asking until the spec is decision complete: approach, interfaces (APIs/schemas/I/O), data flow, edge cases/failure modes, testing + acceptance criteria, rollout/monitoring, and any migrations/compat constraints.

## Asking questions

Critical rules:

* Strongly prefer using the `request_user_input` tool to ask any questions.
* Offer only meaningful multiple‑choice options; don’t include filler choices that are obviously wrong or irrelevant.
* In rare cases where an unavoidable, important question can’t be expressed with reasonable multiple‑choice options (due to extreme ambiguity), you may ask it directly without the tool.

You SHOULD ask many questions, but each question must:

* materially change the spec/plan, OR
* confirm/lock an assumption, OR
* choose between meaningful tradeoffs.
* not be answerable by non-mutating commands.

Use the `request_user_input` tool only for decisions that materially change the plan, for confirming important assumptions, or for information that cannot be discovered via non-mutating exploration.

## Two kinds of unknowns (treat differently)

1. **Discoverable facts** (repo/system truth): explore first.

   * Before asking, run targeted searches and check likely sources of truth (configs/manifests/entrypoints/schemas/types/constants).
   * Ask only if: multiple plausible candidates; nothing found but you need a missing identifier/context; or ambiguity is actually product intent.
   * If asking, present concrete candidates (paths/service names) + recommend one.
   * Never ask questions you can answer from your environment (e.g., “where is this struct”).

2. **Preferences/tradeoffs** (not discoverable): ask early.

   * These are intent or implementation preferences that cannot be derived from exploration.
   * Provide 2–4 mutually exclusive options + a recommended default.
   * If unanswered, proceed with the recommended option and record it as an assumption in the final plan.

## Finalization rule

Only output the final plan when it is decision complete and leaves no decisions to the implementer.

When you present the official plan, wrap it in a `<proposed_plan>` block so the client can render it specially:

1) The opening tag must be on its own line.
2) Start the plan content on the next line (no text on the same line as the tag).
3) The closing tag must be on its own line.
4) Use Markdown inside the block.
5) Keep the tags exactly as `<proposed_plan>` and `</proposed_plan>` (do not translate or rename them), even if the plan content is in another language.

Example:

<proposed_plan>
plan content
</proposed_plan>

plan content should be human and agent digestible. The final plan must be plan-only and include:

* A clear title
* A brief summary section
* Important changes or additions to public APIs/interfaces/types
* Test cases and scenarios
* Explicit assumptions and defaults chosen where needed

Do not ask "should I proceed?" in the final output. The user can easily switch out of Plan mode and request implementation if you have included a `<proposed_plan>` block in your response. Alternatively, they can decide to stay in Plan mode and continue refining the plan.

Only produce at most one `<proposed_plan>` block per turn, and only when you are presenting a complete spec.

If the user stays in Plan mode and asks for revisions after a prior `<proposed_plan>`, any new `<proposed_plan>` must be a complete replacement.

Summary: auto

## Skills
A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list of skills that can be used. Each entry includes a name, description, and file path so you can open the source for full instructions when using a specific skill.
### Available skills
- skill-creator: Guide for creating effective skills. This skill should be used when users want to create a new skill (or update an existing skill) that extends Codex's capabilities with specialized knowledge, workflows, or tool integrations. (file: /home/ubuntu/.codex/skills/.system/skill-creator/SKILL.md)
- skill-installer: Install Codex skills into $CODEX_HOME/skills from a curated list or a GitHub repo path. Use when a user asks to list installable skills, install a curated skill, or install a skill from another repo (including private repos). (file: /home/ubuntu/.codex/skills/.system/skill-installer/SKILL.md)
### How to use skills
- Discovery: The list above is the skills available in this session (name + description + file path). Skill bodies live on disk at the listed paths.
- Trigger rules: If the user names a skill (with `$SkillName` or plain text) OR the task clearly matches a skill's description shown above, you must use that skill for that turn. Multiple mentions mean use them all. Do not carry skills across turns unless re-mentioned.
- Missing/blocked: If a named skill isn't in the list or the path can't be read, say so briefly and continue with the best fallback.
- How to use a skill (progressive disclosure):
  1) After deciding to use a skill, open its `SKILL.md`. Read only enough to follow the workflow.
  2) When `SKILL.md` references relative paths (e.g., `scripts/foo.py`), resolve them relative to the skill directory listed above first, and only consider other paths if needed.
  3) If `SKILL.md` points to extra folders such as `references/`, load only the specific files needed for the request; don't bulk-load everything.
  4) If `scripts/` exist, prefer running or patching them instead of retyping large code blocks.
  5) If `assets/` or templates exist, reuse them instead of recreating from scratch.
- Coordination and sequencing:
  - If multiple skills apply, choose the minimal set that covers the request and state the order you'll use them.
  - Announce which skill(s) you're using and why (one short line). If you skip an obvious skill, say why.
- Context hygiene:
  - Keep context small: summarize long sections instead of pasting them; only load extra files when needed.
  - Avoid deep reference-chasing: prefer opening only files directly linked from `SKILL.md` unless you're blocked.
  - When variants exist (frameworks, providers, domains), pick only the relevant reference file(s) and note that choice.
- Safety and fallback: If a skill can't be applied cleanly (missing files, unclear instructions), state the issue, pick the next-best approach, and continue.


**Inspecting decoder extension points**


I verified the train/inference entry points. I’m taking one more code pass over `CodecDecoderVocos` internals to anchor exact extension points (multi-head ISTFT, residual HF branch, gating, and optional post-filter hook) in the implementation plan.


nl -ba /home/ubuntu/xcodec-sr/vq/codec_decoder_vocos.py | sed -n '1,260p' && nl -ba /home/ubuntu/xcodec-sr/vq/codec_decoder_vocos.py | sed -n '260,560p'


Output:
Chunk ID: 9862a6
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 5939
Output:
     1	import sys
     2
     3	import numpy as np
     4	import torch
     5	import torch.nn as nn
     6	from vq.residual_vq import ResidualVQ
     7	from vq.module import WNConv1d, DecoderBlock, ResLSTM
     8	from vq.alias_free_torch import *
     9	from vq  import activations
    10	from typing import Optional
    11	from vq.module   import ConvNeXtBlock,   AdaLayerNorm
    12	from vq.bs_roformer5 import TransformerBlock
    13	# from rotary_embedding_torch import RotaryEmbedding
    14	from torchtune.modules import RotaryPositionalEmbeddings
    15	from vector_quantize_pytorch import ResidualFSQ
    16	from torch.nn import Module, ModuleList
    17	class ISTFT(nn.Module):
    18	    """
    19	    Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with
    20	    windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges.
    21	    See issue: https://github.com/pytorch/pytorch/issues/62323
    22	    Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs.
    23	    The NOLA constraint is met as we trim padded samples anyway.
    24
    25	    Args:
    26	        n_fft (int): Size of Fourier transform.
    27	        hop_length (int): The distance between neighboring sliding window frames.
    28	        win_length (int): The size of window frame and STFT filter.
    29	        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
    30	    """
    31
    32	    def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"):
    33	        super().__init__()
    34	        if padding not in ["center", "same"]:
    35	            raise ValueError("Padding must be 'center' or 'same'.")
    36	        self.padding = padding
    37	        self.n_fft = n_fft
    38	        self.hop_length = hop_length
    39	        self.win_length = win_length
    40	        window = torch.hann_window(win_length)
    41	        self.register_buffer("window", window)
    42
    43	    def forward(self, spec: torch.Tensor) -> torch.Tensor:
    44	        """
    45	        Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram.
    46
    47	        Args:
    48	            spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size,
    49	                            N is the number of frequency bins, and T is the number of time frames.
    50
    51	        Returns:
    52	            Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal.
    53	        """
    54	        if self.padding == "center":
    55	            # Fallback to pytorch native implementation
    56	            return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True)
    57	        elif self.padding == "same":
    58	            pad = (self.win_length - self.hop_length) // 2
    59	        else:
    60	            raise ValueError("Padding must be 'center' or 'same'.")
    61
    62	        assert spec.dim() == 3, "Expected a 3D tensor as input"
    63	        B, N, T = spec.shape
    64
    65	        # Inverse FFT
    66	        ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
    67	        ifft = ifft * self.window[None, :, None]
    68
    69	        # Overlap and Add
    70	        output_size = (T - 1) * self.hop_length + self.win_length
    71	        y = torch.nn.functional.fold(
    72	            ifft, output_size=(1, output_size), kernel_size=(1, self.win_length), stride=(1, self.hop_length),
    73	        )[:, 0, 0, pad:-pad]
    74
    75	        # Window envelope
    76	        window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
    77	        window_envelope = torch.nn.functional.fold(
    78	            window_sq, output_size=(1, output_size), kernel_size=(1, self.win_length), stride=(1, self.hop_length),
    79	        ).squeeze()[pad:-pad]
    80
    81	        # Normalize
    82	        assert (window_envelope > 1e-11).all()
    83	        y = y / window_envelope
    84
    85	        return y
    86
    87
    88
    89	class FourierHead(nn.Module):
    90	    """Base class for inverse fourier modules."""
    91
    92	    def forward(self, x: torch.Tensor) -> torch.Tensor:
    93	        """
    94	        Args:
    95	            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
    96	                        L is the sequence length, and H denotes the model dimension.
    97
    98	        Returns:
    99	            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
   100	        """
   101	        raise NotImplementedError("Subclasses must implement the forward method.")
   102
   103
   104	class ISTFTHead(FourierHead):
   105	    """
   106	    ISTFT Head module for predicting STFT complex coefficients.
   107
   108	    Args:
   109	        dim (int): Hidden dimension of the model.
   110	        n_fft (int): Size of Fourier transform.
   111	        hop_length (int): The distance between neighboring sliding window frames, which should align with
   112	                          the resolution of the input features.
   113	        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
   114	    """
   115
   116	    def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"):
   117	        super().__init__()
   118	        out_dim = n_fft + 2
   119	        self.out = torch.nn.Linear(dim, out_dim)
   120	        self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding)
   121
   122	    def forward(self, x: torch.Tensor) -> torch.Tensor:
   123	        """
   124	        Forward pass of the ISTFTHead module.
   125
   126	        Args:
   127	            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
   128	                        L is the sequence length, and H denotes the model dimension.
   129
   130	        Returns:
   131	            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
   132	        """
   133	        x_pred = self.out(x )
   134	        # x_pred = x
   135	        x_pred = x_pred.transpose(1, 2)
   136	        mag, p = x_pred.chunk(2, dim=1)
   137	        mag = torch.exp(mag)
   138	        mag = torch.clip(mag, max=1e2)  # safeguard to prevent excessively large magnitudes
   139	        # wrapping happens here. These two lines produce real and imaginary value
   140	        x = torch.cos(p)
   141	        y = torch.sin(p)
   142	        # recalculating phase here does not produce anything new
   143	        # only costs time
   144	        # phase = torch.atan2(y, x)
   145	        # S = mag * torch.exp(phase * 1j)
   146	        # better directly produce the complex value
   147	        S = mag * (x + 1j * y)
   148	        audio = self.istft(S)
   149	        return audio.unsqueeze(1),x_pred
   150
   151
   152	def nonlinearity(x):
   153	    # swish
   154	    return x * torch.sigmoid(x)
   155
   156
   157	def Normalize(in_channels, num_groups=32):
   158	    return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
   159
   160
   161	class ResnetBlock(nn.Module):
   162	    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
   163	                 dropout, temb_channels=512):
   164	        super().__init__()
   165	        self.in_channels = in_channels
   166	        out_channels = in_channels if out_channels is None else out_channels
   167	        self.out_channels = out_channels
   168	        self.use_conv_shortcut = conv_shortcut
   169
   170	        self.norm1 = Normalize(in_channels)
   171	        self.conv1 = torch.nn.Conv1d(in_channels,
   172	                                     out_channels,
   173	                                     kernel_size=3,
   174	                                     stride=1,
   175	                                     padding=1)
   176	        if temb_channels > 0:
   177	            self.temb_proj = torch.nn.Linear(temb_channels,
   178	                                             out_channels)
   179	        self.norm2 = Normalize(out_channels)
   180	        self.dropout = torch.nn.Dropout(dropout)
   181	        self.conv2 = torch.nn.Conv1d(out_channels,
   182	                                     out_channels,
   183	                                     kernel_size=3,
   184	                                     stride=1,
   185	                                     padding=1)
   186	        if self.in_channels != self.out_channels:
   187	            if self.use_conv_shortcut:
   188	                self.conv_shortcut = torch.nn.Conv1d(in_channels,
   189	                                                     out_channels,
   190	                                                     kernel_size=3,
   191	                                                     stride=1,
   192	                                                     padding=1)
   193	            else:
   194	                self.nin_shortcut = torch.nn.Conv1d(in_channels,
   195	                                                    out_channels,
   196	                                                    kernel_size=1,
   197	                                                    stride=1,
   198	                                                    padding=0)
   199
   200	    def forward(self, x, temb=None):
   201	        h = x
   202	        h = self.norm1(h)
   203	        h = nonlinearity(h)
   204	        h = self.conv1(h)
   205
   206	        if temb is not None:
   207	            h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
   208
   209	        h = self.norm2(h)
   210	        h = nonlinearity(h)
   211	        h = self.dropout(h)
   212	        h = self.conv2(h)
   213
   214	        if self.in_channels != self.out_channels:
   215	            if self.use_conv_shortcut:
   216	                x = self.conv_shortcut(x)
   217	            else:
   218	                x = self.nin_shortcut(x)
   219
   220	        return x + h
   221
   222	class AttnBlock(nn.Module):
   223	    def __init__(self, in_channels):
   224	        super().__init__()
   225	        self.in_channels = in_channels
   226
   227	        self.norm = Normalize(in_channels)
   228	        self.q = torch.nn.Conv1d(in_channels,
   229	                                 in_channels,
   230	                                 kernel_size=1,
   231	                                 stride=1,
   232	                                 padding=0)
   233	        self.k = torch.nn.Conv1d(in_channels,
   234	                                 in_channels,
   235	                                 kernel_size=1,
   236	                                 stride=1,
   237	                                 padding=0)
   238	        self.v = torch.nn.Conv1d(in_channels,
   239	                                 in_channels,
   240	                                 kernel_size=1,
   241	                                 stride=1,
   242	                                 padding=0)
   243	        self.proj_out = torch.nn.Conv1d(in_channels,
   244	                                        in_channels,
   245	                                        kernel_size=1,
   246	                                        stride=1,
   247	                                        padding=0)
   248
   249	    def forward(self, x):
   250	        h_ = x
   251	        h_ = self.norm(h_)
   252	        q = self.q(h_)
   253	        k = self.k(h_)
   254	        v = self.v(h_)
   255
   256	        # compute attention
   257	        b, c, h = q.shape
   258	        q = q.permute(0, 2, 1)  # b,hw,c
   259	        w_ = torch.bmm(q, k)  # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
   260	        w_ = w_ * (int(c) ** (-0.5))
   260	        w_ = w_ * (int(c) ** (-0.5))
   261	        w_ = torch.nn.functional.softmax(w_, dim=2)
   262
   263	        # attend to values
   264	        w_ = w_.permute(0, 2, 1)  # b,hw,hw (first hw of k, second of q)
   265	        h_ = torch.bmm(v, w_)  # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
   266
   267	        h_ = self.proj_out(h_)
   268
   269	        return x + h_
   270
   271	def make_attn(in_channels, attn_type="vanilla"):
   272	    assert attn_type in ["vanilla", "linear", "none"], f'attn_type {attn_type} unknown'
   273	    print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
   274	    if attn_type == "vanilla":
   275	        return AttnBlock(in_channels)
   276
   277
   278	class Backbone(nn.Module):
   279	    """Base class for the generator's backbone. It preserves the same temporal resolution across all layers."""
   280
   281	    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
   282	        """
   283	        Args:
   284	            x (Tensor): Input tensor of shape (B, C, L), where B is the batch size,
   285	                        C denotes output features, and L is the sequence length.
   286
   287	        Returns:
   288	            Tensor: Output of shape (B, L, H), where B is the batch size, L is the sequence length,
   289	                    and H denotes the model dimension.
   290	        """
   291	        raise NotImplementedError("Subclasses must implement the forward method.")
   292
   293
   294	class VocosBackbone(Backbone):
   295	    """
   296	    Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization
   297
   298	    Args:
   299	        input_channels (int): Number of input features channels.
   300	        dim (int): Hidden dimension of the model.
   301	        intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock.
   302	        num_layers (int): Number of ConvNeXtBlock layers.
   303	        layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`.
   304	        adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
   305	                                                None means non-conditional model. Defaults to None.
   306	    """
   307
   308	    def __init__(
   309	        self,  hidden_dim=1024,depth=12,heads=16,pos_meb_dim=64):
   310	        super().__init__()
   311
   312	        self.embed = nn.Conv1d(hidden_dim, hidden_dim, kernel_size=7, padding=3)
   313
   314
   315
   316	        self.temb_ch = 0
   317	        block_in = hidden_dim
   318	        dropout = 0.1
   319
   320	        prior_net : tp.List[nn.Module] = [
   321	            ResnetBlock(in_channels=block_in,out_channels=block_in,
   322	                        temb_channels=self.temb_ch,dropout=dropout),
   323	            ResnetBlock(in_channels=block_in,out_channels=block_in,
   324	                        temb_channels=self.temb_ch,dropout=dropout),
   325	        ]
   326	        self.prior_net = nn.Sequential(*prior_net)
   327
   328	        depth = depth
   329	        time_rotary_embed = RotaryPositionalEmbeddings(dim=pos_meb_dim)
   330
   331
   332	        transformer_blocks = [
   333	            TransformerBlock(dim=hidden_dim, n_heads=heads, rotary_embed=time_rotary_embed)
   334	            for _ in range(depth)
   335	        ]
   336
   337
   338	        self.transformers = nn.Sequential(*transformer_blocks)
   339	        self.final_layer_norm = nn.LayerNorm(hidden_dim, eps=1e-6)
   340	        post_net : tp.List[nn.Module] = [
   341	            ResnetBlock(in_channels=block_in,out_channels=block_in,
   342	                        temb_channels=self.temb_ch,dropout=dropout),
   343	            ResnetBlock(in_channels=block_in,out_channels=block_in,
   344	                        temb_channels=self.temb_ch,dropout=dropout),
   345	        ]
   346	        self.post_net = nn.Sequential(*post_net)
   347
   348	    def forward(self, x: torch.Tensor ) -> torch.Tensor:
   349	        x = x.transpose(1, 2)
   350	        x = self.embed(x)
   351	        x = self.prior_net(x)
   352	        x = x.transpose(1, 2)
   353	        x= self.transformers(x)
   354	        x = x.transpose(1, 2)
   355	        x = self.post_net(x)
   356	        x = x.transpose(1, 2)
   357	        x = self.final_layer_norm(x)
   358	        return x
   359
   360
   361
   362
   363
   364
   365
   366	def init_weights(m):
   367	    if isinstance(m, nn.Conv1d):
   368	        nn.init.trunc_normal_(m.weight, std=0.02)
   369	        nn.init.constant_(m.bias, 0)
   370
   371	class CodecDecoderVocos(nn.Module):
   372	    def __init__(self,
   373	                 hidden_dim=1024,
   374	                 depth=12,
   375	                 heads=16,
   376	                 pos_meb_dim=64,
   377	                 hop_length=320,
   378	                 vq_num_quantizers=1,
   379	                 vq_dim=2048, #1024 2048
   380	                 vq_commit_weight=0.25,
   381	                 vq_weight_init=False,
   382	                 vq_full_commit_loss=False,
   383	                 codebook_size=16384,
   384	                 codebook_dim=16,
   385	                ):
   386	        super().__init__()
   387	        self.hop_length = hop_length
   388
   389	        self.quantizer = ResidualFSQ(
   390	            dim = vq_dim,
   391	            levels = [4, 4, 4, 4, 4,4,4,4],
   392	            num_quantizers = 1
   393	        )
   394
   395	        # self.quantizer = ResidualVQ(
   396	        #     num_quantizers=vq_num_quantizers,
   397	        #     dim=vq_dim,
   398	        #     codebook_size=codebook_size,
   399	        #     codebook_dim=codebook_dim,
   400	        #     threshold_ema_dead_code=2,
   401	        #     commitment=vq_commit_weight,
   402	        #     weight_init=vq_weight_init,
   403	        #     full_commit_loss=vq_full_commit_loss,
   404	        # )
   405
   406
   407	        self.backbone = VocosBackbone( hidden_dim=hidden_dim,depth=depth,heads=heads,pos_meb_dim=pos_meb_dim)
   408
   409	        self.head = ISTFTHead(dim=hidden_dim, n_fft=self.hop_length*4, hop_length=self.hop_length, padding="same")
   410
   411	        self.reset_parameters()
   412
   413	    def forward(self, x, vq=True):
   414	        if vq is True:
   415	            # x, q, commit_loss = self.quantizer(x)
   416	            x = x.permute(0, 2, 1)
   417	            x, q = self.quantizer(x)
   418	            x = x.permute(0, 2, 1)
   419	            q = q.permute(0, 2, 1)
   420	            return x, q, None
   421	        x = self.backbone(x)
   422	        x,_  = self.head(x)
   423
   424	        return x ,_
   425
   426	    def vq2emb(self, vq):
   427	        self.quantizer = self.quantizer.eval()
   428	        x = self.quantizer.vq2emb(vq)
   429	        return x
   430
   431	    def get_emb(self):
   432	        self.quantizer = self.quantizer.eval()
   433	        embs = self.quantizer.get_emb()
   434	        return embs
   435
   436	    def inference_vq(self, vq):
   437	        x = vq[None,:,:]
   438	        x = self.model(x)
   439	        return x
   440
   441	    def inference_0(self, x):
   442	        x, q, loss, perp = self.quantizer(x)
   443	        x = self.model(x)
   444	        return x, None
   445
   446	    def inference(self, x):
   447	        x = self.model(x)
   448	        return x, None
   449
   450
   451	    def remove_weight_norm(self):
   452	        """Remove weight normalization module from all of the layers."""
   453
   454	        def _remove_weight_norm(m):
   455	            try:
   456	                torch.nn.utils.remove_weight_norm(m)
   457	            except ValueError:  # this module didn't have weight norm
   458	                return
   459
   460	        self.apply(_remove_weight_norm)
   461
   462	    def apply_weight_norm(self):
   463	        """Apply weight normalization module from all of the layers."""
   464
   465	        def _apply_weight_norm(m):
   466	            if isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d):
   467	                torch.nn.utils.weight_norm(m)
   468
   469	        self.apply(_apply_weight_norm)
   470
   471	    def reset_parameters(self):
   472	        self.apply(init_weights)
   473
   474
   475
   476	class CodecDecoderVocos_transpose(nn.Module):
   477	    def __init__(self,
   478	                 hidden_dim=1024,
   479	                 depth=12,
   480	                 heads=16,
   481	                 pos_meb_dim=64,
   482	                 hop_length=320,
   483	                 vq_num_quantizers=1,
   484	                 vq_dim=1024, #1024 2048
   485	                 vq_commit_weight=0.25,
   486	                 vq_weight_init=False,
   487	                 vq_full_commit_loss=False,
   488	                 codebook_size=16384,
   489	                 codebook_dim=16,
   490	                ):
   491	        super().__init__()
   492	        self.hop_length = hop_length
   493
   494
   495	        self.quantizer = ResidualVQ(
   496	            num_quantizers=vq_num_quantizers,
   497	            dim=vq_dim,
   498	            codebook_size=codebook_size,
   499	            codebook_dim=codebook_dim,
   500	            threshold_ema_dead_code=2,
   501	            commitment=vq_commit_weight,
   502	            weight_init=vq_weight_init,
   503	            full_commit_loss=vq_full_commit_loss,
   504	        )
   505
   506
   507	        self.backbone = VocosBackbone( hidden_dim=hidden_dim,depth=depth,heads=heads,pos_meb_dim=pos_meb_dim)
   508
   509	        self.inverse_mel_conv = nn.Sequential(
   510	            nn.GELU(),
   511	            nn.ConvTranspose1d(
   512	                in_channels=hidden_dim,
   513	                out_channels=hidden_dim,
   514	                kernel_size=3,
   515	                stride=2,
   516	                padding=1,
   517	                output_padding=1  # 确保输出长度与编码前匹配
   518	            ),
   519	            nn.GELU(),
   520	            nn.ConvTranspose1d(
   521	                in_channels=hidden_dim,
   522	                out_channels=hidden_dim,
   523	                kernel_size=3,
   524	                padding=1
   525	            )
   526	        )
   527
   528	        self.head = ISTFTHead(dim=hidden_dim, n_fft=self.hop_length*4, hop_length=self.hop_length, padding="same")
   529
   530	        self.reset_parameters()
   531
   532	    def forward(self, x, vq=True):
   533	        if vq is True:
   534	            x, q, commit_loss = self.quantizer(x)
   535	            return x, q, commit_loss
   536	        x = self.backbone(x)
   537	        x,_  = self.head(x)
   538
   539	        return x ,_
   540
   541	    def vq2emb(self, vq):
   542	        self.quantizer = self.quantizer.eval()
   543	        x = self.quantizer.vq2emb(vq)
   544	        return x
   545
   546	    def get_emb(self):
   547	        self.quantizer = self.quantizer.eval()
   548	        embs = self.quantizer.get_emb()
   549	        return embs
   550
   551	    def inference_vq(self, vq):
   552	        x = vq[None,:,:]
   553	        x = self.model(x)
   554	        return x
   555
   556	    def inference_0(self, x):
   557	        x, q, loss, perp = self.quantizer(x)
   558	        x = self.model(x)
   559	        return x, None
   560


Summary: auto


**Inspecting config model**


find /home/ubuntu/xcodec-sr/config -maxdepth 3 -type f | sort | sed -n '1,200p' && echo '---' && rg -n "hop|sr|codec_decoder|stft_params|lambda|hidden_dim|depth|heads" /home/ubuntu/xcodec-sr/config -g '*.yaml'


Output:
Chunk ID: 6c5770
Wall time: 0.0514 seconds
Process exited with code 0
Original token count: 424
Output:
/home/ubuntu/xcodec-sr/config/dataset/default.yaml
/home/ubuntu/xcodec-sr/config/default.yaml
/home/ubuntu/xcodec-sr/config/model/default.yaml
/home/ubuntu/xcodec-sr/config/preprocess/default.yaml
/home/ubuntu/xcodec-sr/config/train/default.yaml
---
/home/ubuntu/xcodec-sr/config/train/default.yaml:14:lambdas:
/home/ubuntu/xcodec-sr/config/train/default.yaml:15:  lambda_disc: 1.0
/home/ubuntu/xcodec-sr/config/train/default.yaml:16:  lambda_feat_match_loss: 1.0
/home/ubuntu/xcodec-sr/config/train/default.yaml:17:  lambda_mel_loss: 15.0
/home/ubuntu/xcodec-sr/config/train/default.yaml:18:  lambda_adv: 1.0
/home/ubuntu/xcodec-sr/config/train/default.yaml:19:  lambda_stft_loss: 1.0
/home/ubuntu/xcodec-sr/config/train/default.yaml:20:  lambda_semantic_loss: 5
/home/ubuntu/xcodec-sr/config/train/default.yaml:21:  lambda_perceptual_loss: 0
/home/ubuntu/xcodec-sr/config/train/default.yaml:28:  hop_sizes: [128,512]
/home/ubuntu/xcodec-sr/config/model/default.yaml:3:  hidden_dim: 1024
/home/ubuntu/xcodec-sr/config/model/default.yaml:4:  depth: 1
/home/ubuntu/xcodec-sr/config/model/default.yaml:5:  heads: 16
/home/ubuntu/xcodec-sr/config/model/default.yaml:10:codec_decoder:
/home/ubuntu/xcodec-sr/config/model/default.yaml:14:  hidden_dim: 1024
/home/ubuntu/xcodec-sr/config/model/default.yaml:15:  depth: 12
/home/ubuntu/xcodec-sr/config/model/default.yaml:16:  heads: 16
/home/ubuntu/xcodec-sr/config/model/default.yaml:33:  stft_params:
/home/ubuntu/xcodec-sr/config/model/default.yaml:35:    hop_sizes: [39, 63, 103, 167, 271,438,709,1148]
/home/ubuntu/xcodec-sr/config/preprocess/default.yaml:10:  sr: 16000
/home/ubuntu/xcodec-sr/config/preprocess/default.yaml:16:  hop_length: 320


**Reviewing config documentation**


I confirmed config entry points for `sr`, `hop_length`, losses, and decoder depth/width; I’ll include concrete config-touchpoints in the phased checklist so implementation is deterministic.


nl -ba /home/ubuntu/xcodec-sr/config/model/default.yaml && echo '---' && nl -ba /home/ubuntu/xcodec-sr/config/preprocess/default.yaml && echo '---' && nl -ba /home/ubuntu/xcodec-sr/config/train/default.yaml


Output:
Chunk ID: e25a13
Wall time: 0.0511 seconds
Process exited with code 0
Original token count: 807
Output:
     1	codec_encoder:
     2	  ngf: 48
     3	  hidden_dim: 1024
     4	  depth: 1
     5	  heads: 16
     6	  pos_meb_dim: 64
     7	  up_ratios: [2, 2, 4, 4, 5]
     8	  dilations: [1, 3, 9]
     9
    10	codec_decoder:
    11	  ngf: 48
    12	  up_ratios: [5, 4, 4, 2, 2] # [2, 5, 5, 4]
    13	  dilations: [1, 3, 9]
    14	  hidden_dim: 1024
    15	  depth: 12
    16	  heads: 16
    17	  pos_meb_dim: 64
    18	  vq_num_quantizers: 1
    19	  vq_dim: 2048 #1024 2048
    20	  vq_commit_weight: 0.25
    21	  vq_weight_init: False
    22	  vq_full_commit_loss: False
    23	  codebook_size: 16384  #8192 16384
    24	  codebook_dim: 16
    25
    26	mpd:
    27	  periods: [2, 3, 5, 7, 11]
    28	  max_downsample_channels: 512
    29	  channels: 16
    30	  channel_increasing_factor: 4
    31
    32	mstft:
    33	  stft_params:
    34	    fft_sizes: [78, 126, 206, 334, 542,876,1418,2296]
    35	    hop_sizes: [39, 63, 103, 167, 271,438,709,1148]
    36	    win_lengths: [78, 126, 206, 334, 542,876,1418,2296]
    37	    window: hann_window
    38	  in_channels: 1
    39	  out_channels: 1
    40	  kernel_sizes: [5, 3]
    41	  channels: 32
    42	  max_downsample_channels: 512
    43	  downsample_scales: [2, 2, 2]
    44	  use_weight_norm: True
---
     1	datasets:
     2	    LibriSpeech:
     3	        root: /aifs4su/data/zheny/data/data_8_21_2/
     4	        trainsets: ['dev', 'test','train']
     5	        testsets: ['test-clean']
     6	view:
     7	    train_filelist: /aifs4su/data/zheny/data/data_8_21_2/mls_all_audio_path_higher_quality.txt
     8	    test_filelist: /aifs4su/data/zheny/data/data_8_21_2/mls_all_audio_path_test.txt
     9	audio:
    10	  sr: 16000
    11	  normalization: False
    12	  max_wav_value: 32768.0
    13	stft:
    14	  n_fft: 1280
    15	  window_length: 1280
    16	  hop_length: 320
---
     1	trainer:
     2	  accelerator: 'gpu'
     3	  num_nodes: 1
     4	  devices: 8
     5	  min_steps: 3000000
     6	  max_steps: 3000000
     7	  precision: 16-mixed
     8	  limit_val_batches: 2560
     9	  val_check_interval: 4000
    10	  num_sanity_val_steps: 0
    11	  accumulate_grad_batches:  1
    12
    13	# loss
    14	lambdas:
    15	  lambda_disc: 1.0
    16	  lambda_feat_match_loss: 1.0
    17	  lambda_mel_loss: 15.0
    18	  lambda_adv: 1.0
    19	  lambda_stft_loss: 1.0
    20	  lambda_semantic_loss: 5
    21	  lambda_perceptual_loss: 0
    22	use_mel_loss: true
    23	use_feat_match_loss: true
    24	use_stft_loss: false
    25
    26	stft_loss_params:
    27	  fft_sizes: [512,2048]
    28	  hop_sizes: [128,512]
    29	  win_lengths: [512,2048]
    30	  window: hann_window
    31
    32	# optimizer
    33	gen_optim_params:
    34	  lr: 1.0
    35	  betas: [0.8, 0.9]
    36	disc_optim_params:
    37	  lr: 1.0
    38	  betas: [0.8, 0.9]
    39	gen_grad_clip: 1.0
    40	disc_grad_clip: 1.0
    41
    42	# scheduler
    43	gen_schedule_params:
    44	  warmup_step: 3000
    45	  down_step: 400000
    46	  min_lr: 2.0e-5 # this should be the final lr
    47	  max_lr: 1.0e-4
    48	disc_schedule_params:
    49	  warmup_step: 3000
    50	  down_step: 400000
    51	  min_lr: 2.0e-5 # this should be the final lr
    52	  max_lr: 1.0e-4


**Finalizing integrated suggestion plan**


**Finalizing refined architecture statuses**


**Finalizing detailed architecture and ADRs**


Section 1: Inputs summary
- Source plans ingested from `/home/ubuntu/xcodec-sr/SRPlans`: `Claude.md`, `GPT.md`, `GPT_PRO.md`, `Gemini.md`, `perplexity.md`.
- Integration constraints from your brief: low-latency voice-agent/TTS deployment, avoid default diffusion bottlenecks, preserve frozen encoder/VQ whenever possible, support controllable `{16,24,32,48}` output, maximize traceability.
- Repo-grounded implementation anchors confirmed: `vq/codec_decoder_vocos.py`, `lightning_module.py`, `inference.py`, `config/model/default.yaml`, `config/preprocess/default.yaml`, `config/train/default.yaml`.
- Traceability map: Plan A=`Claude`, Plan B=`GPT`, Plan C=`GPT_PRO`, Plan D=`Gemini`, Plan E=`perplexity`.

Section 2: Suggestion Inventory (by plan)
**Plan A (Claude)**
- A1: Decoder-only SR with frozen encoder is feasible and should be exploited.
- A2: Start with direct ISTFT-head retarget to 24k.
- A3: Keep compressed-frame modeling; scale hop/n_fft with target SR (50 fps principle).
- A4: Use linear weight interpolation for enlarged ISTFT projection init.
- A5: For 48k, add projection MLP before expanding backbone.
- A6: Use multi-band PQMF (4-band or 3-band reuse variant) with sub-band discriminators.
- A7: Consider MS-Vocos-style trainable synthesis filter for streaming.
- A8: Keep modular post-hoc BWE path available.
- A9: AP-BWE dual-stream amp/phase CNN as primary post-hoc candidate.
- A10: HiFi-GAN+ / MusicHiFi-BWE skip-connected upsampler as alternatives.
- A11: Explicit BWE failure-mode mitigation (HF discriminators, phase modeling, codec-domain training).
- A12: Optional one-step generative tier (FLowHigh/FlashSR/UniverSR-like).
- A13: Structured training recipe (loss stack, discriminator stack, staged schedule, LR partitioning).
- A14: Multi-rate controllability via SR conditioning.
- A15: VQ-token upsampling (prefer non-AR; AR latency concern).

**Plan B (GPT)**
- B1: Treat problem as bandwidth extension, not true signal recovery.
- B2: Highest-ROI baseline is post-BWE after current 16k decoder.
- B3: Ensure streaming/causal operation and fast/HQ mode toggle.
- B4: BBWEXNet-style causal U-Net is a practical low-latency option.
- B5: Prioritize 32k perceptual sweet spot before 48k.
- B6: Train with codec-decoded inputs to avoid train/infer domain gap.
- B7: Use causal sinc/polyphase pre-upsampling.
- B8: Use upper-band-weighted adversarial loss stack.
- B9: Predict only missing high-band residual.
- B10: Implement residual via PQMF or STFT mag/phase residual.
- B11: Native decoder retarget is valid with scaled hop/n_fft and HF-aware losses.
- B12: Optional distilled/one-step diffusion teacher-student and progressive SR conditioning.

**Plan C (GPT_PRO)**
- C1: Prefer decoder-aware BWE over blind terminal upsamplers.
- C2: Execute in ranked order: 24k retarget -> residual HF -> stronger decoder -> optional HQ generative.
- C3: Favor amplitude/phase or complex spectral decoding.
- C4: Add harmonic/source-filter priors for speech/F0 robustness.
- C5: Include lightweight missing-band models (AP-BWE/MS-BWE/BAE-Net/DSP-informed).
- C6: Keep 50Hz token cadence; scale waveform geometry by target SR.
- C7: Replace plain single linear head with dual-stream/harmonic-aware head.
- C8: Residual HF branch with deterministic low-band anchor.
- C9: Condition HF branch on latent + waveform/STFT context + optional F0/UV/energy.
- C10: Hard low-band bypass/identity to prevent regressions.
- C11: Auxiliary F0/UV/band-energy predictor from post-VQ latents.
- C12: APCodec/Wavehax-lite/HiFTNet-lite as stronger decoder replacements.
- C13: Additive branch gating: Base16 + HF24 + HF32 + HF48.
- C14: Band-aware objectives + teacher forcing/scheduled sampling for cascades.
- C15: Strict experiment ladder E0-E5; diffusion/flow as premium only.

**Plan D (Gemini)**
- D1: Multi-band sub-band reconstruction with PQMF.
- D2: Per-band mag/phase outputs from decoder head.
- D3: SnakeBeta + alias-free upsampling blocks in decoder.
- D4: Freeze encoder/semantic path; retrain decoder-side modules only.
- D5: Neural source-filter decomposition (periodic/noise source + filter).
- D6: Shared trunk with multi-head gated outputs by resolution.
- D7: Resolution-dropout training.
- D8: Sub-band discriminators for 8–16 and 16–24 kHz.
- D9: Generic upsamplers fail on FSQ-clean outputs; need token-conditioned generation.

**Plan E (perplexity)**
- E1: Proven decoder-only retrain with frozen encoder/quantizer.
- E2: Reported 24k variant uses hop=960 plus pre-VQ AvgPool to 25Hz.
- E3: Interpolate decoder weights from checkpoint.
- E4: Direct 48k ISTFT retarget (`n_fft=3840`, `hop=960`) is possible.
- E5: Increase backbone width/depth if 48k capacity is insufficient.
- E6: Vocos parameter sensitivity requires deliberate hyperparameter sweeps.
- E7: AP-BWE as post-filter.
- E8: MS-BWE cascade for controllable rates.
- E9: HiFi-GAN+ and UBGAN as alternative post-filters.
- E10: FLowHigh/FlowDec one-step generative options.
- E11: CodecFlow latent-space BWE path.
- E12: MB-iSTFT / 2D frequency upsampling option.
- E13: WaveNeXt / ComplexDec replacement options.
- E14: Multi-head ISTFT gating + resolution embedding + random resolution sampling.
- E15: Phased roadmap (quick 24k -> post-filter -> integrated 48k -> multi-rate gating).
- E16: Reuse pre-encoded data for rapid iteration.

Section 3: Coverage Matrix (table)
**Plan A**
| ID | Status | Reason + implications |
|---|---|---|
| A1 | ACCEPTED | Core premise matches constraints; enables decoder-focused roadmap. |
| A2 | ACCEPTED | Lowest-risk first milestone; fastest path to audible gain. |
| A3 | MODIFIED | Keep 50Hz cadence; use 24k@hop480 baseline, 32/48 progressive geometry next. |
| A4 | ACCEPTED | Reduces cold-start instability for new head dimensions. |
| A5 | MODIFIED | Add projection MLP only if 48k head underfits (triggered by metrics). |
| A6 | MODIFIED | PQMF multiband kept as fallback/parallel experiment, not default first path. |
| A7 | DEFERRED | Trainable synthesis filter postponed until base cascade is stable. |
| A8 | ACCEPTED | External post-filter path included for modular rollout/fallback. |
| A9 | ACCEPTED | AP-BWE design adopted in residual branch and external option. |
| A10 | MODIFIED | HiFi-GAN+/MusicHiFi used as baselines, not primary architecture. |
| A11 | ACCEPTED | Incorporated as mandatory guardrails in training/eval. |
| A12 | MODIFIED | One-step generative kept as optional HQ tier, not realtime default. |
| A13 | ACCEPTED | Adopted as main training protocol with staged optimization. |
| A14 | ACCEPTED | Multi-rate conditioning is part of final architecture. |
| A15 | DEFERRED | Non-AR token upsampling tracked as R&D; high integration cost now. |

**Plan B**
| ID | Status | Reason + implications |
|---|---|---|
| B1 | ACCEPTED | Sets correct objective: plausible HF synthesis, not reconstruction. |
| B2 | MODIFIED | Post-BWE is not sole first step; runs in parallel with native-24 retarget. |
| B3 | ACCEPTED | Runtime supports fast path + HQ path with latency-aware routing. |
| B4 | DEFERRED | BBWEXNet variant retained for benchmark/fallback only. |
| B5 | MODIFIED | Training starts at 24k; product default can move to 32k if SLO permits. |
| B6 | ACCEPTED | Dataset pipeline uses codec-decoded inputs as model inputs. |
| B7 | ACCEPTED | Deterministic causal/polyphase upsampling included pre-HF branch. |
| B8 | ACCEPTED | Upper-band weighted adversarial stack is mandatory. |
| B9 | ACCEPTED | Residual missing-band prediction is central design choice. |
| B10 | ACCEPTED | STFT residual is primary; PQMF variant kept as secondary branch. |
| B11 | ACCEPTED | Native retarget path included with HF-specific objectives. |
| B12 | ACCEPTED | Distilled one-step + progressive conditioning included as optional tiers. |

**Plan C**
| ID | Status | Reason + implications |
|---|---|---|
| C1 | ACCEPTED | Final design is decoder/token-aware, not blind waveform-only. |
| C2 | ACCEPTED | Implemented as canonical phase ordering. |
| C3 | MODIFIED | Dual amp/phase adopted now; full complex-decoder replacement deferred. |
| C4 | ACCEPTED | Harmonic/voicing priors added via conditioning branch. |
| C5 | MODIFIED | AP/MS-BWE integrated now; BAE-Net/DSP variants deferred benchmarks. |
| C6 | ACCEPTED | 50Hz token-rate invariance preserved. |
| C7 | ACCEPTED | Plain head upgraded to dual-stream/harmonic-aware form. |
| C8 | ACCEPTED | Residual HF branch with low-band anchor is primary mechanism. |
| C9 | ACCEPTED | Conditioning inputs include latent, base STFT, optional F0/UV/energy. |
| C10 | ACCEPTED | Hard low-band bypass enforced to prevent regressions. |
| C11 | ACCEPTED | Auxiliary voicing predictor included in training plan. |
| C12 | DEFERRED | Full decoder swaps are escalation options after plateau. |
| C13 | MODIFIED | Plain high-SR Vocos retained only as ablation baseline. |
| C14 | ACCEPTED | Band-aware losses + scheduled sampling adopted for cascades. |
| C15 | ACCEPTED | Experiment ladder E0-E5 is the formal evaluation sequence. |

**Plan D**
| ID | Status | Reason + implications |
|---|---|---|
| D1 | ACCEPTED | PQMF multiband included as selectable 48k strategy. |
| D2 | ACCEPTED | Per-band mag/phase output used in PQMF path design. |
| D3 | DEFERRED | SnakeBeta upsampling blocks tracked as experimental decoder variant. |
| D4 | ACCEPTED | Encoder/semantic frozen in primary experiments. |
| D5 | MODIFIED | Full source-filter decoder deferred; source/voicing priors added now. |
| D6 | ACCEPTED | Shared trunk + gated heads is adopted. |
| D7 | ACCEPTED | Resolution-dropout added for joint multi-rate robustness. |
| D8 | ACCEPTED | Band-specific discriminators are explicit training components. |
| D9 | ACCEPTED | Token-conditioned architecture chosen over generic post-upsamplers. |

**Plan E**
| ID | Status | Reason + implications |
|---|---|---|
| E1 | ACCEPTED | Decoder-only frozen-encoder baseline is foundational. |
| E2 | MODIFIED | Hop960+AvgPool tested as ablation, not default due temporal-rate tradeoff. |
| E3 | ACCEPTED | Weight interpolation is standard initialization step. |
| E4 | ACCEPTED | Direct 48k head path retained as one candidate in Phase 3. |
| E5 | ACCEPTED | Width/depth scaling gated by 48k quality/capacity diagnostics. |
| E6 | ACCEPTED | Hyperparameter sweeps mandated before architectural escalation. |
| E7 | ACCEPTED | AP-BWE supported as external and internalized design influence. |
| E8 | ACCEPTED | MS-BWE-style progressive cascade is integrated. |
| E9 | MODIFIED | HiFi-GAN+/UBGAN become benchmark/fallback deployments. |
| E10 | DEFERRED | FLowHigh/FlowDec in optional HQ lane after realtime path stabilizes. |
| E11 | DEFERRED | CodecFlow latent-flow tracked as long-horizon R&D. |
| E12 | MODIFIED | 2D freq upsample approach included as targeted 48k fallback experiment. |
| E13 | DEFERRED | WaveNeXt/ComplexDec deferred to v2 redesign track. |
| E14 | ACCEPTED | Multi-head gating + resolution conditioning included. |
| E15 | ACCEPTED | Final implementation mirrors this phased roadmap pattern. |
| E16 | ACCEPTED | Pre-encoded data reuse is built into all phases. |

Section 4: Conflicts & Resolutions
1. Conflict: first target rate (`24k first` vs `32k first` vs `direct 48k`).
Resolution: Train native 24k first for stability, then promote 32k as default if quality/latency trigger passes, and 48k as optional tier.
Decision trigger: choose 32k default if MOS/UTMOS gain over 24k is >=0.10 and p95 latency increase <=20%.

2. Conflict: 24k geometry (`hop=480, 50Hz`) vs (`hop=960 + AvgPool to 25Hz`).
Resolution: Default to 50Hz path; run 25Hz variant as controlled ablation only.
Decision trigger: adopt 25Hz only if quality is within 0.02 MOS and compute improves >=25%.

3. Conflict: external post-filter vs integrated decoder enhancement.
Resolution: Hybrid architecture: integrated residual branch is primary; AP-BWE-like external module remains fallback/rapid-ship path.
Decision trigger: fallback to external path if integrated branch misses realtime SLO by >15%.

4. Conflict: single wide ISTFT vs PQMF multiband vs source-filter.
Resolution: Primary = dual-stream STFT residual branch; secondary = PQMF multiband; source-filter priors added via conditioning before full replacement.
Decision trigger: switch to PQMF primary if HF artifact rate or LSD@8–24k remains above threshold after two sweeps.

5. Conflict: plain Vocos-at-higher-SR vs harmonic/phase-aware decoder variants.
Resolution: keep trunk, upgrade head to dual-stream + voicing conditioning first; full APCodec/Wavehax/HiFTNet replacements deferred.
Decision trigger: escalate if E4 plateau persists (<0.05 MOS gain across two major variants).

6. Conflict: diffusion/flow usage.
Resolution: not in default online path; offered as optional HQ async/premium tier.
Decision trigger: enable only if quality gain >=0.10 MOS with predictable cost envelope and no default-path SLO impact.

7. Conflict: controllability mechanism (`multi-head gating` vs `MS-BWE cascade`).
Resolution: combine both: additive cascaded high-band branches with runtime gating and target-rate embedding.
Decision trigger: if memory budget is exceeded, fall back to shared trunk + shallow rate heads.

Section 5: Ensembled Architecture (final)
**Title**
XSR-Ensemble v1: Native-24 Core + Token-Conditioned Residual High-Band Cascade (16/24/32/48)

**Brief summary**
- Keep encoder, semantic branch, and FSQ bottleneck unchanged.
- Upgrade decoder path in two layers: native 24k reconstruction + progressive high-band residual generation for 32/48k.
- Make SR selectable at inference via gating, with optional HQ generative tier off the realtime path.

**Final text diagram**
```text
Input: frozen X-Codec latents (vq_post_emb @ 50Hz) + optional decoder hidden
   |
   +--> Shared Decoder Trunk (existing VocosBackbone, trainable)
           |
           +--> Base16 Head (existing ISTFT head, 0-8k anchor)
           |
           +--> Native24 Dual-Stream Head (amp/phase-aware ISTFT, 0-12k)
           |
           +--> HF Residual Cascade (token-conditioned)
                 - HF24 branch: add 8-12k
                 - HF32 branch: add 12-16k
                 - HF48 branch: add 16-24k
                 - Conditions: vq_post_emb + base STFT + F0/UV/energy + target_sr embedding
                 - Low-band hard bypass
           |
           +--> Band Merger
                 - Default: STFT-band residual merge
                 - Optional: PQMF multi-band merge path
           |
           +--> Output Controller
                 - target_sr in {16,24,32,48}
                 - quality_tier in {realtime,hq}
                 - fast path / HQ path routing
           |
           +--> Optional external module
                 - AP-BWE/UBGAN/HiFi-GAN+ fallback
           |
           +--> Optional HQ async tier
                 - one-step flow/distilled diffusion refinement
```

**Data flow (request -> processing -> storage)**
1. Request arrives with `target_sr`, `quality_tier`, `streaming_mode`, `voice_id`.
2. Retrieve/generate codec latents (`vq_code`, `vq_post_emb`) from cache or upstream TTS stage.
3. Decode through shared trunk and select branch path by gate:
- `16k`: Base16 only.
- `24k`: Native24 head.
- `32k`: Native24 + HF24/HF32 branches.
- `48k`: Native24 + HF24/HF32/HF48 (or PQMF path if selected).
4. Optional external fallback post-filter if realtime SLA pressure or model-selection policy requires.
5. Optional HQ async refinement path for premium/offline.
6. Persist output audio to object storage; persist metadata/metrics to DB; cache intermediates for repeated prompts.

**Key components and responsibilities**
- Shared trunk: global temporal modeling at fixed 50Hz.
- Native24 head: baseline wideband improvement with minimal architecture shift.
- HF residual cascade: generates only missing bands, preserving low-band identity.
- Voicing conditioner: predicts/provides F0/UV/energy features to stabilize HF synthesis.
- Gating controller: enforces quality/latency policy and output-rate selection.
- External post-filter adapter: allows modular deployment and fallback.
- Eval/observability pipeline: enforces quality + SLO gates continuously.

**Important public API/interfaces/types changes**
- `decode(request)` request fields:
  - `target_sr: Literal[16000,24000,32000,48000]`
  - `quality_tier: Literal["realtime","hq"]`
  - `path_override: Optional[Literal["native","cascade","pqmf","external_bwe","hq_refine"]]`
- New model config fields:
  - `model.sr_gating.enabled`
  - `model.sr_gating.target_sr_embeddings`
  - `model.hf_cascade.enabled_branches = ["hf24","hf32","hf48"]`
  - `model.lowband_bypass.enabled`
  - `model.pqmf.enabled`
- New output metadata:
  - `selected_path`, `rtf`, `hf_energy_ratio`, `fallback_reason`.

**Security & privacy considerations**
- Encrypt audio artifacts at rest and in transit.
- Strict retention TTL for raw and generated audio.
- Role-based access for model endpoints and artifact stores.
- PII-safe logs: no raw waveform in logs/traces, only hashed IDs and metrics.

**Performance & cost considerations**
- Realtime default keeps transformer cadence fixed; extra cost concentrated in heads/branches.
- Compute scales with requested SR by gating branch execution.
- Cache reuse of latents/intermediates lowers repeated-utterance cost.
- HQ generative tier isolated to async path to protect latency-critical traffic.

Section 6: ADRs (key decisions)
- ADR-01 Core model strategy: Hybrid native decoder retarget + residual high-band cascade.
Context: reconcile native quality and modular low latency.
Decision: adopt hybrid.
Consequence: best balance of quality, extensibility, and rollout safety.

- ADR-02 Encoder/VQ mutability: keep frozen by default.
Context: existing pretrained bottleneck and cached data.
Decision: no encoder/VQ changes in v1.
Consequence: fast iteration, lower risk, limited latent bandwidth changes.

- ADR-03 Temporal rate policy: preserve 50Hz baseline.
Context: latency and trunk efficiency.
Decision: default 50Hz geometry; test 25Hz ablation only.
Consequence: predictable latency and minimal architectural churn.

- ADR-04 24k-first execution.
Context: strongest risk-adjusted win.
Decision: native 24k is first production milestone.
Consequence: quick quality uplift before 48k complexity.

- ADR-05 HF synthesis design.
Context: avoid low-band regressions.
Decision: residual high-band only with hard low-band bypass.
Consequence: better stability and interpretable errors.

- ADR-06 Controllability mechanism.
Context: must support 16/24/32/48 without model explosion.
Decision: target-rate embedding + additive branch gating.
Consequence: compute scales with requested quality.

- ADR-07 Model serving topology.
Context: realtime agents need strict p95 latency.
Decision: synchronous realtime inference for base/cascade, async queue for HQ tier.
Consequence: latency isolation and operational safety.

- ADR-08 Queue choice.
Context: optional HQ work must not block realtime.
Decision: durable message queue (NATS JetStream/Kafka class) for async refinements.
Consequence: backpressure control and retry semantics.

- ADR-09 Datastore choice.
Context: artifact lifecycle + experiment traceability.
Decision: object storage for audio artifacts + relational metadata store.
Consequence: cheap storage, strong lineage, manageable queryability.

- ADR-10 Caching strategy.
Context: repeated prompts/voices and shared intermediates.
Decision: cache `vq_code`, base decode, and base STFT keyed by request hash.
Consequence: reduced average compute and improved p95.

- ADR-11 Auth/privacy.
Context: voice data sensitivity.
Decision: service JWT/mTLS + per-tenant policy + retention controls.
Consequence: compliance-ready and lower data-exposure risk.

- ADR-12 Observability.
Context: quality regressions are subtle and latency-sensitive.
Decision: OpenTelemetry traces + structured logs + metrics dashboards.
Consequence: rapid diagnosis and safer rollouts.

- ADR-13 Deployment policy.
Context: multiple architecture options and potential regressions.
Decision: canary + shadow A/B + auto-rollback on SLO breach.
Consequence: controlled adoption of new SR paths.

Section 7: Implementation Plan (phased + checklists w/ suggestion IDs)
**Phase 0: Baseline harness and data contracts**
Milestone: reproducible baselines and paired training corpus.
- [ ] Build paired training pipeline `codec_decoded_input -> high_sr_target` from pre-encoded data. [B6, E16, A11, D9]
- [ ] Add deterministic resampler baseline (E0) for 24/32/48 comparisons. [C15, B7]
- [ ] Add objective suite: MR-STFT, upper-band mel, ViSQOL/LSD/FAD proxies, RTF collection. [A13, B8, C14, E6]

**Phase 1: Native 24k decoder retarget**
Milestone: production-candidate 24k model.
- [ ] Implement head-retarget config path and checkpoint interpolation utility. [A2, A4, E1, E3]
- [ ] Set default 24k geometry (50Hz-preserving), plus one ablation for 25Hz AvgPool path. [A3, C6, E2]
- [ ] Train with staged schedule and discriminator/loss stack from baseline recipe. [A13, D4, E15]
- [ ] Add high-band discriminator focus for 8–12k. [B8, D8, C14]
- [ ] Decision gate: promote if quality and latency criteria pass. [A1, B11, C2]

**Phase 2: Residual HF branch to 32k**
Milestone: controllable 32k with low-band fidelity lock.
- [ ] Add residual HF24/HF32 branch modules with hard low-band bypass. [B9, C8, C10]
- [ ] Add conditioning features (`vq_post_emb`, base STFT, F0/UV/energy). [C9, C11, C4]
- [ ] Upgrade head to dual-stream amplitude/phase outputs. [C3, C7, A9]
- [ ] Train with upper-band-weighted objectives + small adv + band discriminator. [B8, C14, D8]
- [ ] Decision gate: 32k default if MOS gain and p95 latency trigger pass. [B5, B3, C2]

**Phase 3: 48k cascade + multirate gating**
Milestone: stable 48k option with runtime gating.
- [ ] Add HF48 branch and scheduled sampling/teacher-forcing in cascaded training. [C13, C14, E8]
- [ ] Implement runtime gate for `{16,24,32,48}` and resolution dropout in training. [A14, D6, D7, E14]
- [ ] Implement PQMF multiband path and per-band mag/phase heads as secondary strategy. [A6, D1, D2, E12]
- [ ] Add capacity-scale options (projection MLP, trunk widen/deepen) behind config flags. [A5, E5]
- [ ] Hyperparameter sweep for Vocos sensitivity before finalizing 48k default. [E6, E4]

**Phase 4: Serving and platform integration**
Milestone: reliable production rollout.
- [ ] Add realtime/HQ routing controller and path metadata emission. [B3, C15, E15]
- [ ] Integrate external AP-BWE fallback and baseline alternatives (HiFi-GAN+/UBGAN). [A8, A10, E7, E9]
- [ ] Implement queue/cache/storage/auth/observability ADR decisions end-to-end. [ADR-07, ADR-08, ADR-09, ADR-10, ADR-11, ADR-12]
- [ ] Canary + shadow A/B rollout with rollback automation. [ADR-13]

**Phase 5: Optional HQ generative tier**
Milestone: premium quality lane without impacting realtime SLO.
- [ ] Integrate one-step flow/distilled diffusion in async path only. [A12, B12, E10, C15]
- [ ] Evaluate latent-space probabilistic extension prototype track. [E11, A15]
- [ ] Escalate to full decoder replacements only if plateau persists. [C12, D3, D5, E13]

Section 8: Robust Evaluation Plan (metrics + tests + SLOs)
**Unit tests**
- ISTFT/head dimension correctness for each `target_sr`.
- Band-mask and low-band bypass invariants (no low-band drift in residual mode).
- Gating logic tests for branch execution and config compatibility.
- Conditioning feature alignment tests (F0/UV frame-sync with 50Hz decoder frames).

**Integration tests**
- End-to-end decode from cached latents for all SR modes.
- Consistency tests between streaming chunked inference and full-utterance inference.
- Fallback path tests (internal cascade -> external BWE -> base path).
- Model-load and mixed-version compatibility tests (old checkpoints + new heads).

**E2E tests**
- Voice-agent pipeline smoke tests under concurrent sessions.
- Multi-tenant auth and data-retention policy verification.
- Canary shadow comparisons with automatic metric collection.

**Load tests + SLOs**
- SLO-1 Realtime latency: p95 chunk latency <=40 ms for realtime tier.
- SLO-2 RTF budgets (GPU reference): 16k<=0.03, 24k<=0.05, 32k<=0.07, 48k<=0.10.
- SLO-3 Reliability: decode failure rate <0.1%.
- SLO-4 Quality uplift: 24k must improve MOS proxy over 16k by >=0.20; 32k/48k must improve HF-LSD and ViSQOL without low-band degradation.
- SLO-5 Cost: GPU-ms per generated second must stay within configured budget envelope per tier.

**Chaos/failure injection**
- Induce queue backlog and verify realtime path isolation.
- Inject missing/invalid conditioning features (F0/UV) and verify graceful fallback.
- Simulate GPU OOM and assert automatic downgrade to lower-cost mode.
- Corrupt cache entries and verify cache bypass correctness.
- Disable one discriminator stream during training and verify alerting on quality regressions.

**Observability signals**
- Metrics: per-path RTF, p50/p95 latency, branch-hit ratio, fallback rate, HF energy ratio, clipping rate, MOS proxy, discriminator losses by band.
- Logs: request hash, target_sr, selected path, fallback reason, model version, output duration.
- Traces: spans for decode trunk, each HF branch, merger, external fallback, HQ async stage.

**Acceptance criteria per component**
- Native24 head: passes quality uplift and no regression in speaker similarity.
- HF residual branch: improves high-band metrics with low-band error below threshold.
- Gating controller: exact branch activation by mode and stable latency scaling.
- PQMF path: artifact rate not worse than STFT-merge baseline.
- External fallback: seamless takeover under SLO pressure.
- HQ tier: quality improvement measurable, realtime path unaffected.

Section 9: Risks, Failure modes, Mitigations
| Risk / failure mode | Likely cause | Mitigation | Fallback |
|---|---|---|---|
| Metallic HF artifacts | phase inconsistency | dual-stream phase losses + phase discriminators | switch to PQMF path |
| Low-band timbre damage | full-band enhancement drift | hard low-band bypass and identity loss | bypass HF branches |
| 48k underfitting | insufficient trunk/head capacity | projection MLP + width/depth scaling sweeps | keep 32k default |
| Latency breach in 48k | too many active branches | runtime gating + model pruning/quantization | degrade to 32k/24k |
| Domain mismatch | clean-LR training inputs | train on codec-decoded inputs | regenerate paired corpus |
| Cascade compounding errors | stage dependency drift | teacher forcing + scheduled sampling | run fewer stages |
| Overfitting to seen voices | narrow training slices | bandwidth-filtered but speaker-diverse sampling | add regularization/data mix |
| Ops instability | mixed model versions | strict model registry + contract tests | auto-rollback canary |
| Security leakage | raw audio in logs/artifacts | PII-safe logging + TTL + encryption | emergency retention purge |
| HQ path impacting realtime | shared resources contention | async queue isolation + quotas | hard-disable HQ tier |

Section 10: Assumptions & Open questions
- Assumption: target output set is `{16,24,32,48}` only.
- Assumption: encoder, semantic model, and FSQ stay frozen in v1.
- Assumption: pre-encoded HiFi-TTS2 assets are accessible with aligned high-rate references.
- Assumption: realtime tier has dedicated GPU serving capacity.
- Open question: final default rate target for production launch (`24k` vs `32k`) depends on Phase 2 gate.
- Open question: exact hardware SLO envelopes (A10/L4/4090/CPU class) must be finalized for hard latency budgets.
- Open question: whether 44.1k support is required later (impacts gating/head configs).
- Open question: external model licensing constraints for AP-BWE/HiFi-GAN+/UBGAN in commercial deployment.
- Open question: required speaker-similarity threshold for acceptance (define explicit metric cutoffs).
- Open question: appetite for v2 clean-slate decoder replacement (APCodec/Wavehax/ComplexDec) after v1 rollout.