Accurate model card (augment run: 1 GPU, coupled env+target augment, DNS-Noise)

README.md

@@ -13,62 +13,71 @@ library_name: transformers
 pipeline_tag: text-to-speech
 ---
 
-# LongCat-AudioDiT Env-TTS — 10000-step fine-tune
+# LongCat-AudioDiT Env-TTS — `augment` (10,000-step fine-tune)
 
-Fine-tune of [meituan-longcat/LongCat-AudioDiT-1B](https://huggingface.co/meituan-longcat/LongCat-AudioDiT-1B) for the
-**three-stream env-tts task**: given a reference environment audio, a reference
-speaker audio, and three text streams (env caption / speaker caption / target
-speech text), generate target speech that places the target text in the
+Fine-tune of [meituan-longcat/LongCat-AudioDiT-1B](https://huggingface.co/meituan-longcat/LongCat-AudioDiT-1B)
+for the **three-stream env-tts task**: given a reference environment audio, a
+reference speaker audio, and three text streams (env caption / speaker caption /
+target speech text), generate target speech that places the target text in the
 referenced environment with the referenced speaker timbre.
 
+This `augment` variant adds **environment-consistent augmentation** so the
+generated target lives in the referenced acoustic scene.
+
 ## Differences from the base model
 
-The transformer adds **six learnable boundary tokens** (three latent-space, three
-text-space):
+Six learnable boundary tokens (three latent-space, three text-space):
 
 ```
 latent sequence : [<boe>  z_env  <bos>  z_spk  <bon>  z_target]
 text sequence   : [<boe_t> env_text_emb <bos_t> spk_text_emb <bon_t> target_text_emb]
 ```
 
-`encode_multistream_text(env, spk, target, drop_env_text=…, drop_spk_text=…,
-drop_target_text=…)` is the new entry-point. `AudioDiTModel.forward(...)` also
-accepts a pre-assembled `prompt_latent` (replaces `prompt_audio`) so the inference
-path can feed the boundary-tokenized three-stream prompt directly.
+`encode_multistream_text(...)` is the entry-point; `AudioDiTModel.forward(...)`
+also accepts a pre-assembled `prompt_latent`.
 
 ## Training summary
 
 | Field | Value |
 |---|---|
-| Steps | 10000 |
-| Effective batch | 16 × grad_accum 2 × 2 GPU = **64 rows / step** |
+| Steps | 10,000 |
+| Hardware | 1× RTX PRO 6000 Blackwell (96 GB), bf16 |
+| Effective batch | 16 × grad_accum 2 × 1 GPU = **32 rows / step** |
 | Learning rate | cosine 5e-5 (warmup 250) |
 | AdamW | β₁=0.9, β₂=0.999, wd=0.01 |
 | EMA | disabled |
 | LoRA | r=32, alpha=32, target = attn + ffn |
-| Full-train | boundary tokens + AdaLN + text_conv + latent_embed + input_embed + output_proj + time_embed |
-| Audio filter | target duration ∈ [3, 45] s |
-| RMS normalize | three-stream independent to **-23 dBFS** (target_rms=0.0708) |
-| Augmentation | noise + RIR on spk_audio (DNS5 64GB) |
+| Full-train | boundary tokens + AdaLN + text_conv + latent / latent_cond / input embeds + output_proj + time_embed |
 | Data | [ChristianYang/Env-TTS-Clean](https://huggingface.co/datasets/ChristianYang/Env-TTS-Clean) |
+| Audio | target ∈ [3, 15] s; three-stream RMS-norm to −23 dBFS; peak-clip at 0.5 |
+
+### Augmentation (the `augment` change)
+
+Noise + RIR are streamed on-demand from
+[ChristianYang/DNS-Noise](https://huggingface.co/datasets/ChristianYang/DNS-Noise)
+(DNS-Challenge `noise_fullband` + `impulse_responses`, republished as 24 kHz mono):
+
+- **Speaker ref** — an independent 50/25/25 draw: clean / noise / noise+RIR, SNR ∈ [−5, 15] dB.
+- **Env + target (coupled)** — a *separate* 50/25/25 draw whose **same** noise clip and
+  **same** RIR are applied to **both** env and target, placing the generated target in one
+  consistent acoustic scene. The RIR tail is kept; env/target are capped to 15 s.
 
 ## How to load
 
-The model uses **custom code** in this repo, so pass `trust_remote_code=True`:
+Uses custom code in this repo, so pass `trust_remote_code=True`:
 
 ```python
 from transformers import AutoModel, AutoTokenizer
 
 model = AutoModel.from_pretrained(
-    "meituan-longcat/LongCat-AudioDiT-Env-TTS-1B-10000Step",
+    "ChristianYang/LongCat-AudioDiT-Env-TTS-1B-augment",
     trust_remote_code=True,
 ).cuda().eval()
 
 tokenizer = AutoTokenizer.from_pretrained(model.config.text_encoder_model)
 ```
 
-For end-to-end env-tts inference (three-stream prompt + ASR fallback for missing
-env/spk text) see the training repo's `tasks/inference.py`.
+See the training repo's `tasks/inference.py` for end-to-end env-tts inference.
 
 ## License