Add DramaBox LoRA training workflow

- Add the ready-to-use DramaBox LoRA training workflow and cover
- List the workflow in the README examples
- Keep DramaBox fine-tuning documentation general and capability-focused

CHANGELOG.md

@@ -5,6 +5,51 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [5.6.5] - 2026-08-03
+
+### Fixed
+
+- Fix MOSS-TTS training settings in saved workflows
+- Fix existing MOSS Dataset Prep workflows loading values into the wrong fields
+- Fix invalid validation split and preparation batch size errors after updating
+- Fix MOSS training tensor shape errors caused by shifted codec settings
+## [5.6.4] - 2026-08-03
+
+### Added
+
+- Add MOSS-TTS training dataset folder support
+- Add direct loading of matching audio and transcript files from a folder
+- Support WAV, FLAC, MP3, OGG, and M4A training clips
+- Add optional recursive scanning for datasets organized into subfolders
+- Preserve existing JSONL manifest workflows
+## [5.6.3] - 2026-08-01
+
+### Changed
+
+- Improve runtime availability checks so package startup code is not executed during installation
+
+### Fixed
+
+- Fix TTS Audio Suite installer validation failures
+- Fix ComfyUI Desktop installation failing on supported PyTorch and TorchAudio combinations
+## [5.6.2] - 2026-07-30
+
+### Changed
+
+- Improve F5-TTS fallback so the standard PyTorch attention backend continues working
+
+### Fixed
+
+- Fix F5-TTS failing to load with incomplete FlashAttention installations
+- Fix F5-TTS startup crashes when optional FlashAttention components are missing
+## [5.6.1] - 2026-07-30
+
+### Fixed
+
+- Fix Fish Audio S2 installation in headless Linux environments
+- Fix missing optional audio libraries preventing Fish Audio S2 setup
+- Improve Linux and macOS dependency warnings so core TTS installation continues
+- Correct Fedora package installation guidance
 ## [5.6.0] - 2026-07-25
 
 ### Added

README.md

@@ -7,7 +7,7 @@
 [![Dynamic TOML Badge][version-shield]][version-url]
 [![Ko-Fi](https://img.shields.io/badge/Ko--fi-F16061?style=for-the-badge&logo=ko-fi&logoColor=white)](https://ko-fi.com/diogogo)
 
-# TTS Audio Suite v5.6.0
+# TTS Audio Suite v5.6.5
 
 [![ko-fi](https://ko-fi.com/img/githubbutton_sm.svg)](https://ko-fi.com/diogogo)
 
@@ -277,6 +277,9 @@ This matters because the suite now has a clearer split:
   FP8-cast transformer storage, and optional `torch.compile`
 * **Generation diagnostics**: conservative near-silence detection in console
   output, TTS generation information, and SRT timing reports
+* **LoRA training**: official DramaBox audio-branch IC-LoRA training through
+  the unified training nodes, with normalized manifest/index input and managed
+  adapter export
 
 **Important limitations:**
 
@@ -288,6 +291,8 @@ This matters because the suite now has a clearer split:
 
 See the **[DramaBox Prompting Guide](docs/DRAMABOX_PROMPTING_GUIDE.md)** for
 prompt syntax, controls, memory modes, duration behavior, and examples.
+See the **[DramaBox LoRA Training Guide](docs/DRAMABOX_LORA_GUIDE.md)** for
+dataset formats, training workflow, adapter loading, and CPU-safe preflight.
 
 </details>
 
@@ -1268,19 +1273,19 @@ This section provides a detailed guide for installing TTS Audio Suite, covering
 
 * Python 3.12 or higher
 
-* **System libraries** (Linux only):
+* **Optional system libraries** (Linux only):
   
   ```bash
-  # Ubuntu/Debian - Required for audio processing
+  # Ubuntu/Debian - Optional audio features
   sudo apt-get install portaudio19-dev libsamplerate0-dev
   
   # Fedora/RHEL
   sudo dnf install portaudio-devel libsamplerate-devel
   ```
   
-  > **📋 Why needed?** `libsamplerate0-dev` provides audio resampling libraries for packages like `resampy` and `soxr`. `portaudio19-dev` enables voice recording features.
+  > **📋 Optional:** `libsamplerate0-dev` provides additional audio-resampling support. `portaudio19-dev` enables voice recording. Missing either package no longer blocks installation of the TTS engines.
 
-* **macOS dependencies**:
+* **Optional macOS dependencies**:
   
   ```bash
   brew install portaudio
@@ -1375,17 +1380,17 @@ If you have a direct installation with a virtual environment (venv), follow thes
 
 ### Troubleshooting Dependency Issues
 
-#### System Dependencies (Linux)
+#### Optional System Dependencies (Linux)
 
-**Our install script automatically detects missing system libraries** and will display helpful error messages like:
+**Our install script automatically detects missing optional system libraries** and will display feature warnings like:
 
 ```
-[!] Missing system dependencies detected!
+[!] Optional system dependencies are missing
 ============================================================
-SYSTEM DEPENDENCIES REQUIRED
+OPTIONAL LINUX SYSTEM DEPENDENCIES
 ============================================================
-• libsamplerate0-dev (for audio resampling)  
-• portaudio19-dev (for voice recording)
+• libsamplerate0-dev (optional additional audio-resampling support)
+• portaudio19-dev (optional voice recording)
 
 Please install with:
 # Ubuntu/Debian:
@@ -1394,7 +1399,7 @@ sudo apt-get install libsamplerate0-dev portaudio19-dev
 # Fedora/RHEL:
 sudo dnf install libsamplerate-devel portaudio-devel
 ============================================================
-Then run this install script again.
+Core TTS installation will continue; only the listed features may be unavailable.
 ```
 
 #### Python Environment Issues
@@ -1527,7 +1532,7 @@ For offline/manual setup:
 | Granite ASR | `ComfyUI/models/TTS/granite_asr/` | ✅ | Granite ASR models; plus adds native diarization/timestamps, optional Qwen forced aligner reused lazily for timestamps/SRT fallback |
 | Echo-TTS | `ComfyUI/models/TTS/echo-tts-base/` | ✅ | ~7.1GB total (base + dac); CC-BY-NC-SA |
 | Dots TTS | `ComfyUI/models/TTS/dots_tts/` | ✅ | Official base / soar / mf checkpoints with tokenizer, vocoder, speaker encoder |
-| DramaBox | `ComfyUI/models/TTS/dramabox/DramaBox/` | ✅ | ~16.4GB download; fast mode roughly 24GB VRAM; experimental FP8 peaks on RTX 4090: staged ~15.1GB allocated, sequential ~11.7GB allocated / ~12.4GB reserved; conditional LTX-2 Community License |
+| DramaBox | `ComfyUI/models/TTS/dramabox/DramaBox/` | ✅ | ~16.4GB download; fast mode roughly 24GB VRAM; experimental FP8, staged, and sequential options can reduce VRAM, but no minimum GPU size is guaranteed; conditional LTX-2 Community License |
 | Fish Audio S2 Pro | `ComfyUI/models/TTS/fish_audio_s2_pro/` | ✅ | Official BF16 or optional community FP8 checkpoint; the official checkpoint can be quantized on load with BNB INT8/NF4; main T5 environment with process teardown for Clear VRAM; Fish Audio Research License |
 | OmniVoice | `ComfyUI/models/TTS/omnivoice/` | ✅ | Official OmniVoice model. Voice cloning in this suite requires explicit reference text. |
 
@@ -1568,6 +1573,7 @@ Your support helps maintain and improve this project for the entire community!
 | Workflow                                       | Description                                                | Status               | Files                                                                                                               |
 | ---------------------------------------------- | ---------------------------------------------------------- | -------------------- | ------------------------------------------------------------------------------------------------------------------- |
 | **🤐 Voice Cleaning**                          | Audio restoration & cleanup with dual tool pipeline        | ✅ **New in v4.13**   | [📁 JSON](example_workflows/Voice%20Cleaning%20-%20🤐%20Noise%20or%20Vocal%20Removal%20+%20🤐%20Voice%20Fixer.json) |
+| **DramaBox LoRA 🎓 Model Training**           | DramaBox IC-LoRA training workflow from staged speech clips | ✅ **New**           | [📁 JSON](example_workflows/DramaBox%20LoRA%20🎓%20Model%20Training.json)                                          |
 | **MOSS LoRA 🎓 Model Training**               | Initial MOSS LoRA training workflow from clipped speech dataset | ✅ **New in v4.27** | [📁 JSON](example_workflows/MOSS%20LoRA%20🎓%20Model%20Training.json)                                               |
 | **RVC 🎓 Model Training**                     | RVC voice model training workflow                          | ✅ **New in v4.25**   | [📁 JSON](example_workflows/RVC%20🎓%20Model%20Training.json)                                                       |
 | **🎨 Step Audio EditX - Audio Editor**         | Step Audio EditX audio editing with inline edit tags       | ✅ **New in v4.14**   | [📁 JSON](example_workflows/🎨%20Step%20Audio%20EditX%20-%20Audio%20Editor%20+%20Inline%20Edit%20Tags.json)        |

docs/DRAMABOX_LORA_GUIDE.md

@@ -0,0 +1,99 @@
+# DramaBox LoRA training
+
+TTS Audio Suite exposes the official DramaBox audio-branch IC-LoRA trainer
+through the unified `🎓 Model Training` flow. The bundled scripts are pinned to
+the same upstream DramaBox revision as the inference implementation.
+
+See the official DramaBox
+[LoRA training guide](https://github.com/resemble-ai/DramaBox#training-a-lora-on-top-of-dramabox)
+for the upstream dataset format and training behavior.
+
+## Workflow
+
+1. Build a `⚙️ DramaBox Engine`.
+2. Create the dataset either externally or entirely inside ComfyUI:
+   `🎞️ Training Clip Staging` → `🧾 DramaBox Dataset Rows`.
+3. Connect the resulting manifest to `📦 DramaBox Dataset Prep` and keep
+   `dataset_type` set to `manifest`.
+4. Provide at least two clips per speaker.
+5. Connect the dataset to `🎛️ DramaBox Training Config` and then to `🎓 Model Training`.
+6. Select the resulting adapter in the DramaBox engine, or enter its path in the
+   advanced LoRA override field.
+
+The dataset node accepts:
+
+- JSONL/JSON manifests with `audio_filepath` (or `audio_path`) and `text` (or
+  `transcript`)
+- TSV rows with audio path and text
+- the official `gemini_synthetic` and `libriheavy` index formats
+
+Manifest rows may include `speaker`, `speaker_id`, `language`, and `duration`.
+If `speaker` is omitted, rows are grouped as `speaker_1`. Duration and audio
+metadata are measured without loading the waveform into the GPU. The suite
+converts all accepted formats into the `~`-delimited speaker index required by
+the upstream training loop. Clips are restricted to 2–20 seconds by default.
+
+For an all-ComfyUI dataset, connect one or more `AUDIO` sources to
+`🎞️ Training Clip Staging`, then enter one transcript per clip in
+`🧾 DramaBox Dataset Rows`. Speaker and language lines are optional; shared
+defaults are used when those lines are blank.
+
+### Transcripts and scene descriptions
+
+The official trainer accepts either plain spoken transcripts or the same
+scene-style prompt format used for inference. For example, both of these are
+valid training text:
+
+```text
+This is the spoken sentence.
+A woman speaks warmly, "This is the spoken sentence."
+```
+
+Use scene descriptions only when they accurately describe the clip. Plain
+transcripts remain valid and are the safer choice when no reliable style or
+scene annotation is available.
+
+## What training does
+
+The first preprocessing pass uses Gemma and the DramaBox audio VAE to create
+cached conditions and audio latents. The training process then attaches a LoRA
+to the audio transformer branch. It saves periodic checkpoints and exports the
+selected adapter to:
+
+```text
+ComfyUI/models/TTS/dramabox/loras/<adapter_name>/
+```
+
+The job directory, normalized index, preprocessing cache, progress file, and
+logs are stored under:
+
+```text
+ComfyUI/output/tts_audio_suite_training/dramabox/
+```
+
+`continue_from` is a warm start from an existing LoRA checkpoint; it is not an
+exact optimizer-state resume. Use saved checkpoints to compare quality rather
+than assuming the last step is best. Optional upstream validation can be
+enabled with a `val_config` YAML path, but it launches full DramaBox inference
+at each save step. It requires a second GPU: set `validation_gpu` to that
+physical CUDA device index. The suite rejects validation on the training GPU
+instead of allowing both full model processes to compete for the same VRAM.
+
+DramaBox LoRA inference supports normal transformer precision, `fp8_cast`, and
+the optional `torch.compile` path. With normal precision the live adapter is
+reversibly merged for fast inference. With FP8 storage the BF16 adapter remains
+unmerged above the immutable FP8 base weights, avoiding unsafe mixed-dtype
+weight fusion while retaining the main FP8 memory saving.
+
+The base DramaBox runtime is reused when the selected adapter or LoRA strength
+changes. Strength updates are applied directly to the live PEFT adapter, while
+the generated-audio cache still treats adapter path, file revision, and strength
+as distinct generation settings. Replacing an adapter with a different rank may
+retrace compiled transformer blocks, but does not reload the base checkpoint.
+
+## CPU-safe preflight
+
+Training and Gemma/VAE preprocessing are GPU workloads. For development or
+validation without touching CUDA, enable `dry_run` in the training config and
+the dataset node's `dry_run`/`preprocess_now` controls. This writes the
+normalized index and official command/config without loading DramaBox weights.

docs/DRAMABOX_PROMPTING_GUIDE.md

@@ -140,9 +140,8 @@ The segment override ends at the next character tag.
   VRAM. It additionally keeps the diffusion transformer in system RAM while
   another major stage uses CUDA. It transfers the transformer for every
   generated segment or long-form chunk and is therefore substantially slower.
-  With `fp8_cast`,
-  this measured about 11.7GB peak allocated and 12.4GB peak reserved VRAM on
-  an RTX 4090; leave additional headroom for ComfyUI and other loaded models.
+  Actual peak usage varies with the environment, generation settings, and
+  other loaded components; no minimum GPU size is guaranteed.
   System RAM must hold the offloaded transformer (about 3.4GB with FP8 or
   6.6GB without it).
 - `fp8_cast` uses the official LTX FP8 transformer weight-storage policy and

docs/Dev reports/tts_audio_suite_engines.yaml

@@ -1351,7 +1351,7 @@ engines:
       srt: true
       vc: false
       asr: false
-      training: false
+      training: true
 
     special_features:
       - "Expressive scene prompting and stage directions"
@@ -1362,6 +1362,7 @@ engines:
       - "Explicit generation/reference durations, CFG rescale control, and optional Perth watermark"
       - "Experimental staged and sequential strategies for lowering peak VRAM; official LTX FP8-cast storage"
       - "Optional official torch.compile path"
+      - "Official audio-branch IC-LoRA training workflow"
 
     model_sources:
       - component: "DramaBox DiT + audio components"
@@ -1415,8 +1416,8 @@ engines:
       emotion_control: { supported: true, notes: "Natural-language scene prompt and stage directions" }
       native_long_form: { supported: true, notes: "Official duration-aware quote-group chunking; ~37s target / 45s cap" }
       native_srt_duration_targeting: { supported: true, notes: "Subtitle duration is passed as gen_duration before the selected SRT timing mode applies final correction" }
-      community_finetunes: { supported: false, notes: "Not integrated" }
-      vram_efficient: { supported: true, notes: "Fast mode is ~24GB; experimental FP8 peaks on RTX 4090: staged ~15.1GB allocated, sequential ~11.7GB allocated / ~12.4GB reserved" }
+      community_finetunes: { supported: true, notes: "Official audio-branch IC-LoRA adapters can be trained and loaded" }
+      vram_efficient: { supported: true, notes: "Fast mode is ~24GB; experimental FP8, staged, and sequential options can reduce VRAM, but practical peak usage varies by environment and loaded components" }
       speed_performance: { supported: "partial", notes: "Fast favors repeated generation; staged reloads temporary components; sequential also transfers the transformer; optional block-level torch.compile accelerates repeat denoising" }
       reference_free_tts: { supported: true, notes: "Voice reference is optional" }
 
@@ -1983,7 +1984,7 @@ readme_model_download_table:
   - engine: "DramaBox"
     primary_model_path: "ComfyUI/models/TTS/dramabox/DramaBox/"
     auto_download: "✅"
-    notes: "~16.4GB download; fast mode roughly 24GB VRAM; experimental FP8 peaks on RTX 4090: staged ~15.1GB allocated, sequential ~11.7GB allocated / ~12.4GB reserved; conditional LTX-2 Community License"
+    notes: "~16.4GB download; fast mode roughly 24GB VRAM; experimental FP8, staged, and sequential options can reduce VRAM, but no minimum GPU size is guaranteed; conditional LTX-2 Community License"
   - engine: "Fish Audio S2 Pro"
     primary_model_path: "ComfyUI/models/TTS/fish_audio_s2_pro/"
     auto_download: "✅"
@@ -2223,7 +2224,7 @@ model_layouts_markdown: |
 
   ```text
   ComfyUI/models/TTS/dramabox/
-  └── DramaBox/
+  ├── DramaBox/
       ├── dramabox-dit-v1.safetensors
       ├── dramabox-audio-components.safetensors
       ├── assets/
@@ -2234,6 +2235,10 @@ model_layouts_markdown: |
           ├── model-00002-of-00002.safetensors
           ├── model.safetensors.index.json
           └── tokenizer and processor files...
+  └── loras/
+      └── <adapter_name>/
+          ├── adapter_config.json
+          └── adapter_model.safetensors
   ```
 
   Notes:
@@ -2241,6 +2246,8 @@ model_layouts_markdown: |
   - Both repositories download directly into the organized suite folder.
   - Transformers is forced into local-only loading after download.
   - Requires NVIDIA CUDA. Fast mode targets approximately 24GB VRAM; experimental staged/sequential modes can run with less.
+  - Integrated training exports LoRA adapters into `dramabox/loras/<adapter_name>/`.
+  - Training jobs, normalized datasets, preprocessing caches, checkpoints, and logs are stored under `ComfyUI/output/tts_audio_suite_training/dramabox/`.
   - The LTX-2 Community License requires a paid license for entities with at
     least USD 10 million in annual revenue.
 

docs/ENGINE_COMPARISON.md

@@ -18,7 +18,7 @@
 | **Echo-TTS**       | Main      | echo-tts-base + fish-s1-dac-min           | ~5.3GB + ~1.8GB | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | CC-BY-NC-SA-4.0          | Diffusion-based (~30s best), Force Speaker KV (speaker drift control)                    | 1                                                                                        |
 | **Fish Audio S2 Pro** | Main      | S2 Pro 4B / FP8                           | ~10.3GB / ~8.0GB | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | Fish Audio Research License | Free-form sub-word emotion/prosody tags, Native multi-speaker and multi-turn dialogue with dynamic speaker references, Zero-shot voice cloning, Optional per-segment custom character switching, Configurable 4K-32K native context with reduced KV-cache VRAM, Optional community FP8 weight-only checkpoint with BF16 activations, Optional on-the-fly BitsAndBytes INT8/NF4 for the official checkpoint | 80+ languages                                                                            |
 | **Dots TTS**       | Main      | dots.tts-base, dots.tts-soar, dots.tts-mf | ~6GB         | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | Apache-2.0               | Official auto language detect / language control, SOAR and MeanFlow distilled variants   | 19                                                                                       |
-| **DramaBox**       | Main      | DramaBox 3.3B                             | ~16.4GB      | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | LTX-2 Community License  | Expressive scene prompting and stage directions, Native and SRT-aware duration targeting, Official duration-aware long-form chunking with scene-prefix preservation, Optional 10-second zero-shot voice reference, CFG negative prompt with per-segment switching, Explicit generation/reference durations, CFG rescale control, and optional Perth watermark, Experimental staged and sequential strategies for lowering peak VRAM; official LTX FP8-cast storage, Optional official torch.compile path | 1                                                                                        |
+| **DramaBox**       | Main      | DramaBox 3.3B                             | ~16.4GB      | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | LTX-2 Community License  | Expressive scene prompting and stage directions, Native and SRT-aware duration targeting, Official duration-aware long-form chunking with scene-prefix preservation, Optional 10-second zero-shot voice reference, CFG negative prompt with per-segment switching, Explicit generation/reference durations, CFG rescale control, and optional Perth watermark, Experimental staged and sequential strategies for lowering peak VRAM; official LTX FP8-cast storage, Optional official torch.compile path, Official audio-branch IC-LoRA training workflow | 1                                                                                        |
 | **OmniVoice**      | Main      | OmniVoice                                 | ~3.7GB       | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | Apache-2.0               | Inline non-verbal tags and pronunciation overrides, Reference-free voice design, 600+ language support, Upstream long-form chunk orchestration | 600+                                                                                     |
 | **MOSS-TTS**       | Main      | Local 1.7B, Delay 8B v1.5/1.0, VoiceGenerator 1.7B, SoundEffect 8B v1, TTSD 8B | ~8.5GB tokenizer + ~6.1GB/17GB/18GB model | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | Apache-2.0               | Reference-free voice design with MOSS-VoiceGenerator, Native 1-5 speaker TTSD dialogue, 31-language generation with MOSS-TTS-v1.5, Prompt-only sound-effect generation with MOSS-SoundEffect v1, Long-form generation (TTSD/Delay), Duration token hint, Local/Delay/TTSD variants, Initial integrated LoRA training workflow (Delay 8B) | 24                                                                                       |
 | **MOSS-SoundEffect v2** | Main      | MOSS-SoundEffect-v2.0                     | ~11.2GB      | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | Apache-2.0               | Durations up to 30 seconds, Native negative prompting, CFG, flow shift, and diffusion-step controls, Prompt-only text-to-sound generation, 48 kHz mono output, Seeded generation | 2                                                                                        |

docs/FEATURE_COMPARISON.md

@@ -9,14 +9,14 @@
 | **Voice Conversion**         | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ |
 | **ASR (Transcribe)**         | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **Sound Effects**            | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
-| **Training**                 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ |
+| **Training**                 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ✅ |
 | **Voice Cloning**            | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (Base model) | ❌ | ✅ | ✅ | ✅ Reference audio plus exact transcript | ✅ | ✅ Optional reference audio; configurable 3-30 second window from the beginning of the source, default 10 seconds | ✅ | ✅ | ❌ | ⚠️ (needs training) |
 | **Reference Transcript†**    | **Required** | Not used | Not used | Not used | Optional | Optional | Not used | Conditional | Conditional | N/A | **Required** | Not used | **Required** | Optional | Not used | **Required** | Conditional | Not used | N/A |
 | **Native Multi-Speaker**     | ❌ | ❌ | ❌ | ✅ (Base only, Kugel uses fallback) | ✅ | ❌ | ❌ | ❌ | ❌ | ⚠️ (Plus variant speaker attribution / diarization) | ❌ | ❌ | ✅ Native dialogue is default; optional custom mode generates each [Character] segment independently as local speaker 0 | ❌ | ❌ | ❌ | ✅ (TTSD v1.0; 1-5 speakers) | ❌ | ❌ |
 | **Emotion Control**          | ❌ | ❌ | ⚠️ (v2 tags - doesn't work) | ❌ | ⚠️ (via prompt) | ✅ (native inline tags) | ✅ (8 emotions) | ⚠️ (via instruct) | ⚠️ (via instruct) | ❌ | ✅ (14 emotions) | ❌ | ✅ Free-form inline natural-language tags | ❌ | ✅ Natural-language scene prompt and stage directions | ⚠️ (voice-design instruct + inline non-verbal tags) | ✅ (MOSS-VoiceGenerator instruction-conditioned voice design) | ❌ | ❌ |
 | **Native Long-form**         | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ Configurable 4K-32K native context; suite text chunking is bypassed | ❌ | ✅ Official duration-aware quote-group chunking; ~37s target / 45s cap | ✅ (uses upstream audio_chunk_duration / audio_chunk_threshold orchestration; bypasses suite char-based chunk splitting) | ✅ (TTSD/Delay long-form; use chunk orchestration for very long inputs) | ❌ | N/A |
-| **Community Finetunes**      | ✅ | ✅ | ✅ | ✅ KugelAudio, Hindi | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ (LoRA adapter inference supported; initial integrated LoRA training support added for MOSS-TTS Delay 8B) | ❌ | ✅ |
-| **VRAM Efficient**           | ✅ | ✅ | ✅ | ⚠️ (5-18GB) | ⚠️ (9GB) | ⚠️ (~8-10GB) | ⚠️ (9-12GB) | ✅ (5.4GB) | ✅ (3-6GB) | ✅ (~4.6GB) | ⚠️ (7GB) | ⚠️ (~7GB total) | ⚠️ 8K context measured at ~15.2GB BF16, ~11.2GB FP8, ~11.2GB BNB INT8, or ~8.9GB BNB NF4; BF16 codec and activations; BNB is a load-time option for the official checkpoint | ⚠️ (main env works; 2B-class model is not lightweight) | ✅ Fast mode is ~24GB; experimental FP8 peaks on RTX 4090: staged ~15.1GB allocated, sequential ~11.7GB allocated / ~12.4GB reserved | ⚠️ (main env works; 2B-class model is not lightweight) | ⚠️ (Local 1.7B smaller; tokenizer is large) | ⚠️ (runs in the main ComfyUI environment but remains GPU-heavy) | ✅ |
+| **Community Finetunes**      | ✅ | ✅ | ✅ | ✅ KugelAudio, Hindi | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ Official audio-branch IC-LoRA adapters can be trained and loaded | ❌ | ✅ (LoRA adapter inference supported; initial integrated LoRA training support added for MOSS-TTS Delay 8B) | ❌ | ✅ |
+| **VRAM Efficient**           | ✅ | ✅ | ✅ | ⚠️ (5-18GB) | ⚠️ (9GB) | ⚠️ (~8-10GB) | ⚠️ (9-12GB) | ✅ (5.4GB) | ✅ (3-6GB) | ✅ (~4.6GB) | ⚠️ (7GB) | ⚠️ (~7GB total) | ⚠️ 8K context measured at ~15.2GB BF16, ~11.2GB FP8, ~11.2GB BNB INT8, or ~8.9GB BNB NF4; BF16 codec and activations; BNB is a load-time option for the official checkpoint | ⚠️ (main env works; 2B-class model is not lightweight) | ✅ Fast mode is ~24GB; experimental FP8, staged, and sequential options can reduce VRAM, but practical peak usage varies by environment and loaded components | ⚠️ (main env works; 2B-class model is not lightweight) | ⚠️ (Local 1.7B smaller; tokenizer is large) | ⚠️ (runs in the main ComfyUI environment but remains GPU-heavy) | ✅ |
 | **Speed/Performance**        | ✅ Very Fast | ✅ Fast | ✅ Fast | ⚠️ | ⚠️ | ⚠️ CUDA recommended | ⚠️ | ✅ Fast | ⚠️ | ⚠️ Moderate | ⚠️ | ✅ Fast (diffusion, realtime-capable) | ✅ Main-environment subprocess with reliable teardown; local compile measurements: ~40 it/s BF16, ~11.8 it/s NF4 at ~8.9GB VRAM, and ~3.7 it/s INT8 at ~11.2GB VRAM; quality comparison pending | ⚠️ Moderate; mf variant is faster | ⚠️ Fast favors repeated generation; staged reloads temporary components; sequential also transfers the transformer; optional block-level torch.compile accelerates repeat denoising | ✅ Fast; upstream reports sub-realtime RTF | ✅ Fast with CUDA/FlashAttention | ⚠️ (100 diffusion steps by default) | ✅ Fast |
 | **No Narrator Required**     | ❌ | ✅ (default speaker) | ✅ (default speaker) | ✅ (zero-shot / default speaker) | ✅ (basic TTS if no narrator/reference is provided) | ✅ (zero-shot) | ❌ | ✅ (cross-lingual or instruct mode) | ✅ (Base default voice or CustomVoice presets) | N/A | ❌ | ❌ | ✅ Reference audio is optional | ✅ (default speaker) | ✅ Voice reference is optional | ✅ (native default voice; instruct also works without narrator) | ✅ (direct TTS and prompt-only generation) | ❌ | N/A |
 

docs/MODEL_LAYOUTS.md

@@ -227,7 +227,7 @@ Notes:
 
 ```text
 ComfyUI/models/TTS/dramabox/
-└── DramaBox/
+├── DramaBox/
     ├── dramabox-dit-v1.safetensors
     ├── dramabox-audio-components.safetensors
     ├── assets/
@@ -238,6 +238,10 @@ ComfyUI/models/TTS/dramabox/
         ├── model-00002-of-00002.safetensors
         ├── model.safetensors.index.json
         └── tokenizer and processor files...
+└── loras/
+    └── <adapter_name>/
+        ├── adapter_config.json
+        └── adapter_model.safetensors
 ```
 
 Notes:
@@ -245,6 +249,8 @@ Notes:
 - Both repositories download directly into the organized suite folder.
 - Transformers is forced into local-only loading after download.
 - Requires NVIDIA CUDA. Fast mode targets approximately 24GB VRAM; experimental staged/sequential modes can run with less.
+- Integrated training exports LoRA adapters into `dramabox/loras/<adapter_name>/`.
+- Training jobs, normalized datasets, preprocessing caches, checkpoints, and logs are stored under `ComfyUI/output/tts_audio_suite_training/dramabox/`.
 - The LTX-2 Community License requires a paid license for entities with at
   least USD 10 million in annual revenue.
 

docs/MOSS_TRAINING_GUIDE.md

@@ -23,12 +23,29 @@ It does **not** currently support:
 
 Current ComfyUI flow:
 
-1. `🎞️ MOSS Clip Staging`
+1. `🎞️ Training Clip Staging`
 2. `🧾 MOSS Dataset Rows`
 3. `📦 MOSS Dataset Prep`
 4. `🎛️ MOSS Training Config`
 5. `🎓 Model Training`
 
+If clips and transcripts are already prepared on disk, you can skip the first two
+nodes. Set `dataset_source` on `📦 MOSS Dataset Prep` to a folder containing
+same-name audio and text pairs:
+
+```text
+my_dataset/
+├── clip001.wav
+├── clip001.txt
+├── clip002.flac
+└── clip002.txt
+```
+
+Each `.txt` file must contain the transcript spoken in its matching audio file.
+Folder scanning supports WAV, FLAC, MP3, OGG, and M4A. Subfolders are ignored
+unless `recursive_folder_scan` is enabled. Existing JSONL manifest paths continue
+to work unchanged.
+
 ## The Important Fields
 
 ### `text_lines`
@@ -229,7 +246,7 @@ If you do not have a separate validation manifest:
 
 If you want the least confusing starting point:
 
-- use `🎞️ MOSS Clip Staging`
+- use `🎞️ Training Clip Staging`
 - use `🧾 MOSS Dataset Rows`
 - fill only `text_lines`
 - leave `reference_clip_lines` blank

engines/adapters/dramabox_adapter.py

@@ -26,11 +26,36 @@ class DramaBoxEngineAdapter:
         self.audio_cache = get_audio_cache()
         self._last_config: Optional[ModelLoadConfig] = None
         self._load_signature = None
+        self._lora_signature = None
         self.last_generation_status: Dict[str, Any] = {"near_silent": False}
 
     def update_config(self, new_config: Dict[str, Any]):
         self.config = new_config.copy() if new_config else {}
 
+    @staticmethod
+    def _lora_revision(path: Any) -> str:
+        """Return a cheap cache token that changes when a managed adapter is replaced."""
+        value = str(path or "").strip()
+        if not value:
+            return ""
+        try:
+            candidate = os.path.abspath(os.path.expanduser(value))
+            if os.path.isfile(candidate):
+                stat = os.stat(candidate)
+                return f"{candidate}:{stat.st_size}:{stat.st_mtime_ns}"
+            if os.path.isdir(candidate):
+                entries = []
+                for item in os.listdir(candidate):
+                    if not item.endswith(".safetensors"):
+                        continue
+                    item_path = os.path.join(candidate, item)
+                    stat = os.stat(item_path)
+                    entries.append(f"{item}:{stat.st_size}:{stat.st_mtime_ns}")
+                return f"{candidate}|{'|'.join(sorted(entries))}"
+        except OSError:
+            pass
+        return value
+
     @classmethod
     def _warn_if_near_silent(
         cls,
@@ -76,6 +101,7 @@ class DramaBoxEngineAdapter:
         }
 
     def _build_load_signature(self) -> Tuple[Any, ...]:
+        """Identity of the expensive base runtime, excluding live LoRA state."""
         return (
             self.config.get("model_name", "DramaBox"),
             self.config.get("device", "auto"),
@@ -85,35 +111,50 @@ class DramaBoxEngineAdapter:
             bool(self.config.get("compile_model", False)),
         )
 
+    def _build_lora_signature(self) -> Tuple[Any, ...]:
+        path = self.config.get("lora_path", "")
+        return (
+            str(path or "").strip(),
+            self._lora_revision(path),
+            float(self.config.get("lora_strength", 1.0)),
+        )
+
     def _ensure_model_loaded(self):
         signature = self._build_load_signature()
-        if signature == self._load_signature and self._last_config is not None:
-            return
-
-        self._last_config = ModelLoadConfig(
-            engine_name="dramabox",
-            model_type="tts",
-            model_name=self.config.get("model_name", "DramaBox"),
-            device=self.config.get("device", "auto"),
-            additional_params={
-                "precision": self.config.get("precision", "auto"),
-                "memory_mode": self.config.get("memory_mode", "fast"),
-                "transformer_quantization": self.config.get(
-                    "transformer_quantization", "none"
-                ),
-                "compile_model": bool(self.config.get("compile_model", False)),
-            },
-        )
         from utils.models.unified_model_interface import unified_model_interface
 
-        unified_model_interface.load_model(self._last_config)
-        self._load_signature = signature
+        if signature != self._load_signature or self._last_config is None:
+            self._last_config = ModelLoadConfig(
+                engine_name="dramabox",
+                model_type="tts",
+                model_name=self.config.get("model_name", "DramaBox"),
+                device=self.config.get("device", "auto"),
+                additional_params={
+                    "precision": self.config.get("precision", "auto"),
+                    "memory_mode": self.config.get("memory_mode", "fast"),
+                    "transformer_quantization": self.config.get(
+                        "transformer_quantization", "none"
+                    ),
+                    "compile_model": bool(self.config.get("compile_model", False)),
+                },
+            )
+            self._load_signature = signature
+            self._lora_signature = None
+
+        engine = unified_model_interface.load_model(self._last_config)
+        lora_signature = self._build_lora_signature()
+        if lora_signature != self._lora_signature:
+            lora_path, lora_revision, lora_strength = lora_signature
+            engine.set_lora(
+                lora_path=lora_path,
+                strength=lora_strength,
+                revision=lora_revision,
+            )
+            self._lora_signature = lora_signature
+        return engine
 
     def _get_engine(self):
-        self._ensure_model_loaded()
-        from utils.models.unified_model_interface import unified_model_interface
-
-        return unified_model_interface.load_model(self._last_config)
+        return self._ensure_model_loaded()
 
     def _extract_voice_reference(
         self, voice_ref: Optional[Dict[str, Any]]
@@ -195,6 +236,9 @@ class DramaBoxEngineAdapter:
                 ),
                 memory_mode=self.config.get("memory_mode", "fast"),
                 compile_model=bool(self.config.get("compile_model", False)),
+                lora_path=self.config.get("lora_path", ""),
+                lora_strength=float(self.config.get("lora_strength", 1.0)),
+                lora_revision=self._lora_revision(self.config.get("lora_path", "")),
                 seed=int(seed),
                 character=character_name or "narrator",
             )

engines/dramabox/dramabox_engine.py

@@ -28,6 +28,8 @@ class DramaBoxEngine:
         memory_mode: str = "fast",
         transformer_quantization: str = "none",
         compile_model: bool = False,
+        lora_path: str = "",
+        lora_strength: float = 1.0,
     ):
         self.model_name = model_name
         self.device = resolve_torch_device(device)
@@ -36,6 +38,8 @@ class DramaBoxEngine:
         self.memory_mode = str(memory_mode)
         self.transformer_quantization = str(transformer_quantization)
         self.compile_model = bool(compile_model)
+        self.lora_path = str(lora_path or "").strip()
+        self.lora_strength = float(lora_strength)
         self._server = None
         self._server_module = None
 
@@ -104,6 +108,8 @@ class DramaBoxEngine:
             bnb_4bit=True,
             memory_mode=self.memory_mode,
             transformer_quantization=self.transformer_quantization,
+            lora_path=self.lora_path,
+            lora_strength=self.lora_strength,
         )
         print("✅ DramaBox runtime ready")
 
@@ -169,6 +175,17 @@ class DramaBoxEngine:
             "sample_rate": int(sample_rate),
         }
 
+    def set_lora(self, lora_path: str = "", strength: float = 1.0, revision: str = ""):
+        """Update the live adapter without rebuilding the base DramaBox runtime."""
+        self.lora_path = str(lora_path or "").strip()
+        self.lora_strength = float(strength)
+        if self._server is not None:
+            self._server.configure_lora(
+                self.lora_path,
+                self.lora_strength,
+                revision=str(revision or ""),
+            )
+
     def parameters(self) -> Iterator[torch.nn.Parameter]:
         """Expose loaded submodule parameters for ComfyUI memory accounting."""
         if self._server is None:

engines/dramabox/training/__init__.py

@@ -0,0 +1,5 @@
+"""DramaBox LoRA dataset and training integration."""
+
+from .handler import DramaBoxTrainingHandler
+
+__all__ = ["DramaBoxTrainingHandler"]

engines/dramabox/training/dataset.py

@@ -0,0 +1,458 @@
+"""Dataset normalization for the official DramaBox IC-LoRA trainer.
+
+The upstream preprocessor accepts JSONL and TSV, but the upstream training
+loop builds its speaker map from ``~``-delimited index rows.  This module keeps
+that conversion in the suite so a manifest that is valid for preprocessing is
+also valid for training.
+"""
+
+from __future__ import annotations
+
+import csv
+import hashlib
+import json
+import os
+import re
+import wave
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+import folder_paths
+
+
+AUDIO_EXTENSIONS = {".wav", ".flac", ".mp3", ".ogg", ".m4a", ".aac"}
+PREPROCESSED_SAMPLE_PATTERN = re.compile(r"sample_(\d+)\.pt$")
+
+
+def slugify(value: Any) -> str:
+    safe = "".join(
+        ch if ch.isalnum() or ch in ("-", "_") else "_"
+        for ch in str(value or "").strip()
+    )
+    safe = safe.strip("_")
+    return safe or "dramabox_lora"
+
+
+def get_dramabox_training_root() -> str:
+    root = os.path.join(
+        folder_paths.get_output_directory(), "tts_audio_suite_training", "dramabox"
+    )
+    os.makedirs(root, exist_ok=True)
+    return root
+
+
+def _resolve_source_path(value: str) -> Path:
+    raw = os.path.expanduser(str(value or "").strip())
+    if not raw:
+        raise ValueError("dataset_source is required")
+
+    candidates = [Path(raw)]
+    input_root = Path(folder_paths.get_input_directory())
+    candidates.extend((input_root / raw, input_root / "datasets" / raw))
+    for candidate in candidates:
+        if candidate.is_file():
+            return candidate.resolve()
+    raise FileNotFoundError(f"DramaBox dataset source not found: {value}")
+
+
+def _resolve_audio_path(raw_path: Any, *, source_path: Path, audio_dir: str) -> Path:
+    value = os.path.expanduser(str(raw_path or "").strip())
+    if not value:
+        raise ValueError("Dataset row is missing audio_filepath/audio_path")
+
+    candidates: List[Path] = []
+    if os.path.isabs(value):
+        candidates.append(Path(value))
+    else:
+        if audio_dir:
+            candidates.append(Path(os.path.expanduser(audio_dir)) / value)
+        candidates.append(source_path.parent / value)
+        candidates.append(Path(value))
+
+    for candidate in candidates:
+        if candidate.is_file():
+            return candidate.resolve()
+    raise FileNotFoundError(f"DramaBox audio file not found: {raw_path}")
+
+
+def _clean_text(value: Any) -> str:
+    return re.sub(r"\s+", " ", str(value or "").replace("\x00", "")).strip()
+
+
+def _speaker_value(row: Dict[str, Any], default: str = "speaker_1") -> str:
+    value = (
+        row.get("speaker")
+        or row.get("speaker_id")
+        or row.get("voice")
+        or row.get("character")
+        or default
+    )
+    return _clean_text(value).replace("~", "_") or default
+
+
+def _language_value(row: Dict[str, Any]) -> str:
+    return _clean_text(row.get("language") or row.get("lang") or "en").replace("~", "_") or "en"
+
+
+def _coerce_float(value: Any, default: float = 0.0) -> float:
+    try:
+        parsed = float(value)
+    except (TypeError, ValueError):
+        return float(default)
+    return parsed if parsed > 0 else float(default)
+
+
+def _probe_audio(path: Path) -> Tuple[int, int, float]:
+    """Return sample rate, frame count, and duration without loading audio."""
+    try:
+        import torchaudio
+
+        info = torchaudio.info(str(path))
+        sample_rate = int(getattr(info, "sample_rate", 0) or 0)
+        frames = int(getattr(info, "num_frames", 0) or 0)
+        if sample_rate > 0 and frames > 0:
+            return sample_rate, frames, frames / sample_rate
+    except Exception:
+        pass
+
+    if path.suffix.lower() == ".wav":
+        with wave.open(str(path), "rb") as handle:
+            sample_rate = int(handle.getframerate())
+            frames = int(handle.getnframes())
+        if sample_rate > 0 and frames > 0:
+            return sample_rate, frames, frames / sample_rate
+
+    raise RuntimeError(
+        f"Could not inspect audio duration for '{path}'. Add a positive duration "
+        "field to the manifest or install a Torchaudio-compatible decoder."
+    )
+
+
+def _parse_manifest(source_path: Path, audio_dir: str) -> Iterable[Dict[str, Any]]:
+    text = source_path.read_text(encoding="utf-8-sig")
+    stripped = text.lstrip()
+    if stripped.startswith("["):
+        raw_rows = json.loads(text)
+    else:
+        raw_rows = [json.loads(line) for line in text.splitlines() if line.strip()]
+    for row in raw_rows:
+        if not isinstance(row, dict):
+            continue
+        yield {
+            "audio": _resolve_audio_path(
+                row.get("audio_filepath", row.get("audio_path", row.get("audio"))),
+                source_path=source_path,
+                audio_dir=audio_dir,
+            ),
+            "text": _clean_text(row.get("text", row.get("transcript", ""))),
+            "duration": _coerce_float(row.get("duration")),
+            "sample_rate": int(_coerce_float(row.get("sample_rate"))),
+            "samples": int(_coerce_float(row.get("samples", row.get("num_frames")))),
+            "speaker": _speaker_value(row),
+            "language": _language_value(row),
+        }
+
+
+def _parse_tsv(source_path: Path, audio_dir: str) -> Iterable[Dict[str, Any]]:
+    with source_path.open("r", encoding="utf-8-sig", newline="") as handle:
+        for row_number, row in enumerate(csv.reader(handle, delimiter="\t"), start=1):
+            if len(row) < 2:
+                continue
+            yield {
+                "audio": _resolve_audio_path(row[0], source_path=source_path, audio_dir=audio_dir),
+                "text": _clean_text(row[1]),
+                "duration": _coerce_float(row[2]) if len(row) > 2 else 0.0,
+                "sample_rate": 0,
+                "samples": 0,
+                "speaker": _clean_text(row[3]).replace("~", "_") if len(row) > 3 else "speaker_1",
+                "language": _clean_text(row[4]).replace("~", "_") if len(row) > 4 else "en",
+                "row_number": row_number,
+            }
+
+
+def _parse_gemini(source_path: Path, audio_dir: str) -> Iterable[Dict[str, Any]]:
+    for line in source_path.read_text(encoding="utf-8-sig").splitlines():
+        parts = line.strip().split("~")
+        if len(parts) < 8:
+            continue
+        file_id, speaker, language = parts[:3]
+        sample_rate = int(_coerce_float(parts[3], 24000))
+        samples = int(_coerce_float(parts[4]))
+        duration = _coerce_float(parts[5])
+        text = _clean_text(parts[-1])
+        yield {
+            "audio": _resolve_audio_path(file_id, source_path=source_path, audio_dir=audio_dir),
+            "text": text,
+            "duration": duration,
+            "sample_rate": sample_rate,
+            "samples": samples,
+            "speaker": _clean_text(speaker).replace("~", "_") or "speaker_1",
+            "language": _clean_text(language).replace("~", "_") or "en",
+        }
+
+
+def _parse_libriheavy(source_path: Path, audio_dir: str) -> Iterable[Dict[str, Any]]:
+    for line in source_path.read_text(encoding="utf-8-sig").splitlines():
+        parts = line.strip().split("~")
+        if len(parts) < 7:
+            continue
+        file_id, speaker, language = parts[:3]
+        # Format: id~speaker~lang~samples~duration_ms~phonemes~text.
+        sample_rate = 24000
+        samples = int(_coerce_float(parts[3]))
+        duration = _coerce_float(parts[4]) / 1000.0 if len(parts) >= 5 else 0.0
+        yield {
+            "audio": _resolve_audio_path(file_id, source_path=source_path, audio_dir=audio_dir),
+            "text": _clean_text(parts[-1]),
+            "duration": duration,
+            "sample_rate": sample_rate,
+            "samples": samples,
+            "speaker": _clean_text(speaker).replace("~", "_") or "speaker_1",
+            "language": _clean_text(language).replace("~", "_") or "en",
+        }
+
+
+def _raw_rows(source_path: Path, dataset_type: str, audio_dir: str) -> Iterable[Dict[str, Any]]:
+    parsers = {
+        "manifest": _parse_manifest,
+        "tsv": _parse_tsv,
+        "gemini_synthetic": _parse_gemini,
+        "libriheavy": _parse_libriheavy,
+    }
+    try:
+        parser = parsers[str(dataset_type)]
+    except KeyError as exc:
+        raise ValueError(f"Unsupported DramaBox dataset type: {dataset_type}") from exc
+    return parser(source_path, audio_dir)
+
+
+def _fingerprint(source_path: Path, *, dataset_type: str, audio_dir: str, min_duration: float, max_duration: float) -> str:
+    stat = source_path.stat()
+    raw = f"{source_path}|{stat.st_size}|{stat.st_mtime_ns}|{dataset_type}|{audio_dir}|{min_duration}|{max_duration}"
+    return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16]
+
+
+def _normalize_rows(
+    source_path: Path,
+    *,
+    dataset_type: str,
+    audio_dir: str,
+    min_duration: float,
+    max_duration: float,
+) -> List[Dict[str, Any]]:
+    records: List[Dict[str, Any]] = []
+    for row_index, row in enumerate(_raw_rows(source_path, dataset_type, audio_dir)):
+        text = _clean_text(row.get("text"))
+        if not text:
+            continue
+
+        audio = Path(row["audio"]).resolve()
+        sample_rate = int(row.get("sample_rate") or 0)
+        samples = int(row.get("samples") or 0)
+        duration = _coerce_float(row.get("duration"))
+        if not sample_rate or not samples or not duration:
+            try:
+                probed_rate, probed_samples, probed_duration = _probe_audio(audio)
+                sample_rate = sample_rate or probed_rate
+                samples = samples or probed_samples
+                duration = duration or probed_duration
+            except RuntimeError:
+                if duration <= 0:
+                    raise
+                sample_rate = sample_rate or 24000
+                samples = samples or max(1, round(duration * sample_rate))
+
+        if duration < float(min_duration) or duration > float(max_duration):
+            continue
+        records.append(
+            {
+                "id": f"sample_{row_index:06d}",
+                "audio": str(audio),
+                "text": text,
+                "duration": float(duration),
+                "sample_rate": int(sample_rate),
+                "samples": int(samples),
+                "speaker": _speaker_value(row),
+                "language": _language_value(row),
+            }
+        )
+
+    if not records:
+        raise ValueError(
+            "DramaBox dataset preparation produced no usable rows. Check the audio paths, "
+            "transcripts, and the min/max duration filters."
+        )
+
+    speaker_counts: Dict[str, int] = {}
+    for record in records:
+        speaker_counts[record["speaker"]] = speaker_counts.get(record["speaker"], 0) + 1
+    unusable = sorted(name for name, count in speaker_counts.items() if count < 2)
+    if unusable:
+        raise ValueError(
+            "DramaBox LoRA training needs at least two clips per speaker so the official "
+            f"trainer can choose a reference clip. Speakers with fewer than two clips: {', '.join(unusable)}."
+        )
+    return records
+
+
+def _write_index(records: List[Dict[str, Any]], index_path: Path) -> None:
+    index_path.parent.mkdir(parents=True, exist_ok=True)
+    with index_path.open("w", encoding="utf-8") as handle:
+        for record in records:
+            text = str(record["text"]).replace("\r", " ").replace("\n", " ")
+            handle.write(
+                "~".join(
+                    (
+                        str(Path(record["audio"]).resolve()),
+                        str(record["speaker"]),
+                        str(record["language"]),
+                        str(int(record["sample_rate"])),
+                        str(int(record["samples"])),
+                        f"{float(record['duration']):.6f}",
+                        "_",
+                        text,
+                    )
+                )
+                + "\n"
+            )
+
+
+def _preprocessed_indices(directory: Path) -> set[int]:
+    indices: set[int] = set()
+    if not directory.is_dir():
+        return indices
+    for path in directory.glob("sample_*.pt"):
+        match = PREPROCESSED_SAMPLE_PATTERN.fullmatch(path.name)
+        if match:
+            indices.add(int(match.group(1)))
+    return indices
+
+
+def validate_preprocessed_dataset(
+    records: List[Dict[str, Any]],
+    preprocessed_dir: str | Path,
+    *,
+    raise_on_missing: bool = False,
+) -> bool:
+    """Require matching text conditions and audio latents for every index row."""
+    root = Path(preprocessed_dir)
+    expected = set(range(len(records)))
+    available = _preprocessed_indices(root / "conditions") & _preprocessed_indices(
+        root / "audio_latents"
+    )
+    missing = sorted(expected - available)
+    complete = bool(expected) and not missing
+    if raise_on_missing and not complete:
+        preview = ", ".join(str(index) for index in missing[:10]) or "all"
+        suffix = "..." if len(missing) > 10 else ""
+        raise RuntimeError(
+            "DramaBox preprocessing did not produce matching condition/audio-latent "
+            f"files for {len(missing) or len(expected)} sample(s) (indices: {preview}{suffix}). "
+            "Fix the reported source-audio errors and run Dataset Prep again."
+        )
+    return complete
+
+
+def prepare_dramabox_dataset(
+    shared_settings: Dict[str, Any],
+    *,
+    dataset_source: str,
+    model_name: str,
+    dataset_type: str = "manifest",
+    audio_dir: str = "",
+    min_duration: float = 2.0,
+    max_duration: float = 20.0,
+    reuse_existing: bool = True,
+    preprocess_now: bool = True,
+    dry_run: bool = False,
+) -> Dict[str, Any]:
+    source_path = _resolve_source_path(dataset_source)
+    fingerprint = _fingerprint(
+        source_path,
+        dataset_type=dataset_type,
+        audio_dir=audio_dir,
+        min_duration=min_duration,
+        max_duration=max_duration,
+    )
+    safe_name = slugify(model_name)
+    dataset_root = Path(get_dramabox_training_root()) / "datasets" / f"{safe_name}_{fingerprint}"
+    index_path = dataset_root / "speaker_index.txt"
+    metadata_path = dataset_root / "dataset.json"
+    preprocessed_dir = dataset_root / "preprocessed"
+
+    if reuse_existing and metadata_path.is_file() and index_path.is_file():
+        try:
+            metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+            records = metadata.get("records") or []
+        except Exception:
+            records = []
+    else:
+        records = []
+
+    if not records:
+        records = _normalize_rows(
+            source_path,
+            dataset_type=dataset_type,
+            audio_dir=audio_dir,
+            min_duration=float(min_duration),
+            max_duration=float(max_duration),
+        )
+        dataset_root.mkdir(parents=True, exist_ok=True)
+        _write_index(records, index_path)
+        metadata_path.write_text(
+            json.dumps(
+                {
+                    "type": "dramabox_dataset",
+                    "source_path": str(source_path),
+                    "dataset_type": dataset_type,
+                    "audio_dir": audio_dir,
+                    "min_duration": float(min_duration),
+                    "max_duration": float(max_duration),
+                    "records": records,
+                },
+                indent=2,
+                ensure_ascii=False,
+            ),
+            encoding="utf-8",
+        )
+
+    # Rewrite cached indexes as well so datasets prepared by older suite
+    # builds migrate from synthetic sample ids to resolvable audio paths.
+    _write_index(records, index_path)
+
+    dataset: Dict[str, Any] = {
+        "type": "training_dataset",
+        "engine_type": "dramabox",
+        "training_mode": "audio_lora",
+        "model_name": model_name,
+        "dataset_type": dataset_type,
+        "source_path": str(source_path),
+        "index_path": str(index_path),
+        "speaker_index": str(index_path),
+        "data_dir": [str(preprocessed_dir)],
+        "preprocessed_dir": str(preprocessed_dir),
+        "min_duration": float(min_duration),
+        "max_duration": float(max_duration),
+        "records": records,
+        "train_records": len(records),
+        "speakers": sorted({str(record["speaker"]) for record in records}),
+        "preprocessed": validate_preprocessed_dataset(records, preprocessed_dir),
+        "dry_run": bool(dry_run),
+        "shared_settings": dict(shared_settings or {}),
+    }
+
+    if preprocess_now and not dry_run and not dataset["preprocessed"]:
+        from .trainer import run_dramabox_preprocess
+
+        run_dramabox_preprocess(dataset, shared_settings, batch_size=8)
+        dataset["preprocessed"] = True
+
+    return dataset
+
+
+__all__ = [
+    "get_dramabox_training_root",
+    "prepare_dramabox_dataset",
+    "slugify",
+    "validate_preprocessed_dataset",
+]

engines/dramabox/training/handler.py

@@ -0,0 +1,82 @@
+"""DramaBox backend for the unified model-training node."""
+
+from __future__ import annotations
+
+from typing import Any, Dict
+
+from engines.training.base_handler import BaseTrainingHandler
+from engines.training.registry import register_training_handler
+
+
+class DramaBoxTrainingHandler(BaseTrainingHandler):
+    engine_type = "dramabox"
+    artifact_type = "lora_adapter"
+
+    def _shared_settings(self, tts_engine: Any) -> Dict[str, Any]:
+        config = self.ensure_engine_type(tts_engine)
+        return {
+            "model_name": config.get("model_name", "DramaBox"),
+            "device": str(config.get("device", "auto")),
+            "precision": str(config.get("precision", "auto")),
+        }
+
+    def build_default_training_config(self, tts_engine: Any) -> Dict[str, Any]:
+        self._shared_settings(tts_engine)
+        return {
+            "type": "training_config",
+            "engine_type": "dramabox",
+            "training_mode": "audio_lora",
+            "base_model": "dev",
+            "steps": 10000,
+            "learning_rate": 1e-4,
+            "lr_scheduler": "cosine",
+            "warmup_steps": 500,
+            "batch_size": 1,
+            "grad_accum": 4,
+            "max_grad_norm": 1.0,
+            "save_every": 500,
+            "log_every": 10,
+            "seed": 42,
+            "lora_rank": 128,
+            "lora_alpha": 128,
+            "lora_dropout": 0.1,
+            "ref_ratio": 0.3,
+            "max_ref_tokens": 200,
+            "text_dropout": 0.4,
+            "preprocess_batch_size": 8,
+            "validation_config": "",
+            "validation_gpu": "",
+            "dry_run": False,
+        }
+
+    def prepare_dataset(self, tts_engine: Any, **kwargs) -> Dict[str, Any]:
+        from .dataset import prepare_dramabox_dataset
+
+        return prepare_dramabox_dataset(self._shared_settings(tts_engine), **kwargs)
+
+    def train(
+        self,
+        tts_engine: Any,
+        training_dataset: Dict[str, Any],
+        training_config: Dict[str, Any],
+        output_name: str = "",
+        resume: bool = False,
+        overwrite: bool = False,
+        continue_from: Any = None,
+        node_id: str = "",
+    ) -> Dict[str, Any]:
+        from .trainer import run_dramabox_training_job
+
+        return run_dramabox_training_job(
+            shared_settings=self._shared_settings(tts_engine),
+            dataset_info=training_dataset,
+            training_config=training_config,
+            output_name=output_name,
+            resume=resume,
+            overwrite=overwrite,
+            continue_from=continue_from,
+            node_id=node_id,
+        )
+
+
+register_training_handler("dramabox", DramaBoxTrainingHandler)

engines/dramabox/training/trainer.py

@@ -0,0 +1,687 @@
+"""Process runner for the official DramaBox IC-LoRA trainer."""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import shutil
+import subprocess
+import sys
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, Iterable, Optional
+
+import folder_paths
+
+from engines.dramabox.dramabox_downloader import DramaBoxDownloader
+from engines.training.progress_io import write_json_progress_file
+from engines.training.progress_registry import (
+    finalize_training_job,
+    register_training_job,
+    update_training_job,
+)
+
+from .dataset import (
+    get_dramabox_training_root,
+    slugify,
+    validate_preprocessed_dataset,
+)
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[3]
+VENDOR_ROOT = PROJECT_ROOT / "engines" / "dramabox" / "vendor"
+PREPROCESS_SCRIPT = VENDOR_ROOT / "src" / "preprocess.py"
+TRAIN_SCRIPT = VENDOR_ROOT / "src" / "train.py"
+
+
+def _write_progress(progress_file: str, *, status: str, phase: str, **updates: Any) -> None:
+    payload: Dict[str, Any] = {}
+    if progress_file and os.path.isfile(progress_file):
+        try:
+            with open(progress_file, "r", encoding="utf-8") as handle:
+                existing = json.load(handle)
+                if isinstance(existing, dict):
+                    payload.update(existing)
+        except Exception:
+            pass
+    payload.update(updates)
+    payload["status"] = status
+    payload["phase"] = phase
+    payload["updated_at"] = datetime.now().isoformat()
+    if progress_file:
+        write_json_progress_file(progress_file, payload, default=str)
+
+
+def _interrupt_requested() -> bool:
+    try:
+        import comfy.model_management as model_management
+    except Exception:
+        return False
+    try:
+        return bool(model_management.processing_interrupted())
+    except Exception:
+        return bool(getattr(model_management, "interrupt_processing", False))
+
+
+def _device_environment(shared_settings: Dict[str, Any]) -> Dict[str, str]:
+    env = os.environ.copy()
+    device = str(shared_settings.get("device", "auto") or "auto").strip().lower()
+    if device.startswith("cpu"):
+        # CPU mode is explicit. This also prevents a CUDA-enabled torch build
+        # from silently taking the user's GPU during preprocessing.
+        env["CUDA_VISIBLE_DEVICES"] = ""
+    elif device.startswith("cuda:"):
+        env["CUDA_VISIBLE_DEVICES"] = device.split(":", 1)[1]
+    return env
+
+
+def _run_process(
+    command: Iterable[str],
+    *,
+    cwd: Path,
+    env: Dict[str, str],
+    phase: str,
+    progress_file: str = "",
+    node_id: str = "",
+    total_steps: int = 0,
+) -> None:
+    command = [str(value) for value in command]
+    print(f"🎓 DramaBox {phase} command: {' '.join(command)}")
+    process = subprocess.Popen(
+        command,
+        cwd=str(cwd),
+        env=env,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        encoding="utf-8",
+        errors="replace",
+        bufsize=1,
+    )
+    tail: list[str] = []
+    recent_loss_trace: list[Dict[str, Any]] = []
+    best_loss: Optional[float] = None
+    try:
+        assert process.stdout is not None
+        for raw_line in process.stdout:
+            line = raw_line.rstrip()
+            if line:
+                telemetry_match = re.fullmatch(
+                    r"TTS_SUITE_PROGRESS\s+step=(\d+)\s+total=(\d+)", line
+                )
+                if telemetry_match is None:
+                    print(f"[DramaBox {phase}] {line}")
+                    tail.append(line)
+                    del tail[:-30]
+
+                if progress_file:
+                    match = telemetry_match or re.search(
+                        r"(?:Step|step)\s+(\d+)(?:/(\d+))?", line
+                    )
+                    if match:
+                        step = int(match.group(1))
+                        parsed_total = int(match.group(2) or total_steps or 0)
+                        overall_progress = (step / parsed_total) if parsed_total else 0.0
+                        progress_updates: Dict[str, Any] = {
+                            "step": step,
+                            "total_steps": parsed_total,
+                            "overall_progress": overall_progress,
+                            "latest_log": line,
+                        }
+                        loss_match = re.search(
+                            r"\bloss=([-+0-9.eE]+)", line, re.IGNORECASE
+                        )
+                        if loss_match:
+                            loss_value = float(loss_match.group(1))
+                            lr_match = re.search(
+                                r"\blr=([-+0-9.eE]+)", line, re.IGNORECASE
+                            )
+                            learning_rate = (
+                                float(lr_match.group(1)) if lr_match else None
+                            )
+                            recent_loss_trace.append(
+                                {"step": step, "total_loss": loss_value}
+                            )
+                            recent_loss_trace = recent_loss_trace[-120:]
+                            best_loss = (
+                                loss_value
+                                if best_loss is None
+                                else min(best_loss, loss_value)
+                            )
+                            progress_updates.update(
+                                latest_loss=loss_value,
+                                best_gen_loss=best_loss,
+                                recent_loss_trace=recent_loss_trace,
+                                current_metrics={
+                                    "loss_gen_all": loss_value,
+                                    "loss_disc_all": 0.0,
+                                    "loss_mel": 0.0,
+                                    "loss_kl": 0.0,
+                                    "loss_fm": 0.0,
+                                    "learning_rate": learning_rate,
+                                },
+                            )
+                        _write_progress(
+                            progress_file,
+                            status="running",
+                            phase=phase,
+                            **progress_updates,
+                        )
+                        update_training_job(
+                            node_id,
+                            status="running",
+                            phase=phase,
+                            **progress_updates,
+                        )
+                    elif "encoding:" in line.lower():
+                        match = re.search(r"(\d+)\s*/\s*(\d+)", line)
+                        if match:
+                            step = int(match.group(1))
+                            parsed_total = int(match.group(2))
+                            overall_progress = step / max(parsed_total, 1)
+                            _write_progress(
+                                progress_file,
+                                status="running",
+                                phase=phase,
+                                step=step,
+                                total_steps=parsed_total,
+                                overall_progress=overall_progress,
+                                latest_log=line,
+                            )
+                            update_training_job(
+                                node_id,
+                                status="running",
+                                phase=phase,
+                                step=step,
+                                total_steps=parsed_total,
+                                overall_progress=overall_progress,
+                                latest_log=line,
+                            )
+
+            if _interrupt_requested():
+                process.terminate()
+                try:
+                    process.wait(timeout=10)
+                except subprocess.TimeoutExpired:
+                    process.kill()
+                raise InterruptedError(f"DramaBox {phase} interrupted by user")
+
+        return_code = process.wait()
+    except BaseException:
+        if process.poll() is None:
+            process.terminate()
+        raise
+
+    if return_code != 0:
+        details = "\n".join(tail[-10:])
+        raise RuntimeError(
+            f"DramaBox {phase} process failed with exit code {return_code}."
+            + (f"\nLast output:\n{details}" if details else "")
+        )
+
+
+def _resolve_model_paths(shared_settings: Dict[str, Any]) -> Dict[str, str]:
+    model_name = str(shared_settings.get("model_name", "DramaBox") or "DramaBox")
+    return DramaBoxDownloader().resolve_model_path(model_name)
+
+
+def build_preprocess_command(
+    dataset_info: Dict[str, Any],
+    shared_settings: Dict[str, Any],
+    *,
+    batch_size: int = 8,
+    skip_existing: bool = True,
+) -> list[str]:
+    paths = _resolve_model_paths(shared_settings)
+    command = [
+        sys.executable,
+        str(PREPROCESS_SCRIPT),
+        "--dataset-type",
+        "gemini_synthetic",
+        "--index",
+        str(dataset_info["index_path"]),
+        "--output-dir",
+        str(dataset_info["preprocessed_dir"]),
+        "--checkpoint",
+        paths["audio_components"],
+        "--audio-only-ckpt",
+        paths["audio_components"],
+        "--gemma-root",
+        paths["gemma_root"],
+        "--max-duration",
+        str(float(dataset_info.get("max_duration", 20.0))),
+        "--min-duration",
+        str(float(dataset_info.get("min_duration", 2.0))),
+        "--batch-size",
+        str(max(1, int(batch_size))),
+    ]
+    if skip_existing:
+        command.append("--skip-existing")
+    return command
+
+
+def run_dramabox_preprocess(
+    dataset_info: Dict[str, Any],
+    shared_settings: Dict[str, Any],
+    *,
+    batch_size: int = 8,
+    progress_file: str = "",
+    node_id: str = "",
+) -> Dict[str, Any]:
+    command = build_preprocess_command(
+        dataset_info,
+        shared_settings,
+        batch_size=batch_size,
+        skip_existing=True,
+    )
+    _run_process(
+        command,
+        cwd=VENDOR_ROOT,
+        env=_device_environment(shared_settings),
+        phase="preprocess",
+        progress_file=progress_file,
+        node_id=node_id,
+    )
+    validate_preprocessed_dataset(
+        dataset_info.get("records") or [],
+        dataset_info["preprocessed_dir"],
+        raise_on_missing=True,
+    )
+    dataset_info["preprocessed"] = True
+    return dataset_info
+
+
+def _resolve_validation_config(value: str) -> str:
+    raw = os.path.expanduser(str(value or "").strip())
+    if not raw:
+        return ""
+    candidates = [Path(raw)]
+    if not os.path.isabs(raw):
+        candidates.extend(
+            (
+                Path(folder_paths.get_input_directory()) / raw,
+                VENDOR_ROOT / raw,
+            )
+        )
+    for candidate in candidates:
+        if candidate.is_file():
+            return str(candidate.resolve())
+    raise FileNotFoundError(f"DramaBox validation config not found: {value}")
+
+
+def _validation_gpu(training_device: str, requested_gpu: Any) -> str:
+    value = str(requested_gpu or "").strip()
+    if not value:
+        raise ValueError(
+            "DramaBox validation_config requires validation_gpu because official validation "
+            "runs a second full model process. Reserve a GPU different from the training GPU."
+        )
+    if not value.isdigit():
+        raise ValueError("DramaBox validation_gpu must be a non-negative CUDA device index")
+    device = str(training_device or "auto").strip().lower()
+    training_gpu = device.split(":", 1)[1] if device.startswith("cuda:") else "0"
+    if value == training_gpu:
+        raise ValueError(
+            f"DramaBox validation_gpu ({value}) must differ from the training GPU ({training_gpu})"
+        )
+    return value
+
+
+def _resolve_continue_lora(continue_from: Any) -> str:
+    if continue_from is None:
+        return ""
+    if isinstance(continue_from, str):
+        value = os.path.abspath(os.path.expanduser(continue_from.strip()))
+    elif isinstance(continue_from, dict):
+        if str(continue_from.get("engine_type", "") or "").strip().lower() not in {"", "dramabox"}:
+            raise ValueError("continue_from TRAINING_ARTIFACTS must come from a DramaBox training run")
+        value = str(
+            continue_from.get("lora_path")
+            or continue_from.get("model_path")
+            or (continue_from.get("lora_adapter") or {}).get("adapter_path", "")
+        ).strip()
+        value = os.path.abspath(os.path.expanduser(value)) if value else ""
+    else:
+        raise ValueError("Unsupported DramaBox continue_from input")
+
+    if not value:
+        return ""
+    if os.path.isdir(value):
+        candidates = sorted(Path(value).glob("lora_step_*.safetensors"))
+        candidates += [Path(value) / "adapter_model.safetensors"]
+        for candidate in reversed(candidates):
+            if candidate.is_file():
+                return str(candidate)
+        raise FileNotFoundError(f"No DramaBox LoRA weights found in '{value}'")
+    if not os.path.isfile(value):
+        raise FileNotFoundError(f"DramaBox LoRA checkpoint not found: {value}")
+    return value
+
+
+def _managed_lora_root() -> Path:
+    try:
+        from utils.models.extra_paths import get_all_tts_model_paths
+
+        for base_path in get_all_tts_model_paths("TTS"):
+            root = Path(base_path) / "dramabox" / "loras"
+            root.mkdir(parents=True, exist_ok=True)
+            return root
+    except Exception:
+        pass
+    root = Path(folder_paths.models_dir) / "TTS" / "dramabox" / "loras"
+    root.mkdir(parents=True, exist_ok=True)
+    return root
+
+
+def _next_managed_lora_dir(name: str, *, overwrite: bool) -> Path:
+    target = _managed_lora_root() / slugify(name)
+    if overwrite or not target.exists():
+        return target
+    counter = 2
+    while True:
+        candidate = target.parent / f"{target.name}_{counter}"
+        if not candidate.exists():
+            return candidate
+        counter += 1
+
+
+def _latest_lora_file(output_dir: Path) -> Optional[Path]:
+    candidates = sorted(
+        output_dir.glob("lora_step_*.safetensors"),
+        key=lambda path: int(re.search(r"(\d+)", path.stem).group(1))
+        if re.search(r"(\d+)", path.stem)
+        else -1,
+    )
+    if candidates:
+        return candidates[-1]
+    candidate = output_dir / "adapter_model.safetensors"
+    return candidate if candidate.is_file() else None
+
+
+def _build_train_config(
+    dataset_info: Dict[str, Any],
+    shared_settings: Dict[str, Any],
+    training_config: Dict[str, Any],
+    *,
+    output_dir: Path,
+    continue_lora: str,
+    resolve_paths: bool = True,
+) -> Dict[str, Any]:
+    if shared_settings.get("model_paths"):
+        paths = dict(shared_settings["model_paths"])
+    elif resolve_paths:
+        paths = _resolve_model_paths(shared_settings)
+    else:
+        paths = {
+            "transformer": "<dramabox-transformer.safetensors>",
+            "audio_components": "<dramabox-audio-components.safetensors>",
+        }
+    config: Dict[str, Any] = {
+        "data_dir": [str(dataset_info["preprocessed_dir"])],
+        "speaker_index": [str(dataset_info["index_path"])],
+        "output_dir": str(output_dir),
+        "checkpoint": paths["transformer"],
+        "full_checkpoint": paths["audio_components"],
+        "base_model": str(training_config.get("base_model", "dev")),
+        "lora_rank": int(training_config.get("lora_rank", 128)),
+        "lora_alpha": int(training_config.get("lora_alpha", 128)),
+        "lora_dropout": float(training_config.get("lora_dropout", 0.1)),
+        "ref_ratio": float(training_config.get("ref_ratio", 0.3)),
+        "max_ref_tokens": int(training_config.get("max_ref_tokens", 200)),
+        "text_dropout": float(training_config.get("text_dropout", 0.4)),
+        "steps": int(training_config.get("steps", 10000)),
+        "lr": float(training_config.get("learning_rate", 1e-4)),
+        "lr_scheduler": str(training_config.get("lr_scheduler", "cosine")),
+        "warmup_steps": int(training_config.get("warmup_steps", 500)),
+        "batch_size": int(training_config.get("batch_size", 1)),
+        "grad_accum": int(training_config.get("grad_accum", 4)),
+        "max_grad_norm": float(training_config.get("max_grad_norm", 1.0)),
+        "save_every": max(1, int(training_config.get("save_every", 500))),
+        "log_every": int(training_config.get("log_every", 10)),
+        "seed": int(training_config.get("seed", 42)),
+    }
+    if continue_lora:
+        config["resume_lora"] = continue_lora
+    validation_config = _resolve_validation_config(
+        training_config.get("validation_config", "")
+    )
+    if validation_config:
+        config["val_config"] = validation_config
+    return config
+
+
+def _accelerate_command() -> list[str]:
+    executable = shutil.which("accelerate")
+    if executable:
+        return [executable, "launch", "--num_processes", "1"]
+    return [sys.executable, "-m", "accelerate.commands.launch", "--num_processes", "1"]
+
+
+def run_dramabox_training_job(
+    shared_settings: Dict[str, Any],
+    dataset_info: Dict[str, Any],
+    training_config: Dict[str, Any],
+    *,
+    output_name: str = "",
+    resume: bool = False,
+    overwrite: bool = False,
+    continue_from: Any = None,
+    node_id: str = "",
+) -> Dict[str, Any]:
+    if str(dataset_info.get("engine_type", "") or "").strip().lower() != "dramabox":
+        raise ValueError("DramaBox training requires a DramaBox TRAINING_DATASET payload")
+    if str(training_config.get("training_mode", "audio_lora") or "").strip().lower() != "audio_lora":
+        raise ValueError("DramaBox training currently supports audio_lora mode only")
+    if resume:
+        raise RuntimeError(
+            "DramaBox does not support exact optimizer-state resume. Use continue_from with a saved LoRA checkpoint for a warm start."
+        )
+    if str(shared_settings.get("device", "auto") or "auto").strip().lower().startswith("cpu") and not bool(
+        training_config.get("dry_run", False)
+    ):
+        raise RuntimeError(
+            "DramaBox model training requires CUDA. Use dry_run for CPU-only validation; "
+            "no model weights or CUDA process will be started in that mode."
+        )
+    requested_validation = str(
+        training_config.get("validation_config", "") or ""
+    ).strip()
+    if requested_validation:
+        _resolve_validation_config(requested_validation)
+        _validation_gpu(
+            shared_settings.get("device", "auto"),
+            training_config.get("validation_gpu", ""),
+        )
+
+    safe_name = slugify(output_name or dataset_info.get("model_name") or "dramabox_lora")
+    root = Path(get_dramabox_training_root()) / "jobs"
+    root.mkdir(parents=True, exist_ok=True)
+    fingerprint = f"{safe_name}|{dataset_info.get('index_path')}|{training_config}"
+    job_hash = __import__("hashlib").sha256(fingerprint.encode("utf-8")).hexdigest()[:12]
+    job_dir = root / f"{safe_name}_{job_hash}"
+    if job_dir.exists() and not overwrite:
+        job_dir = root / f"{safe_name}_{job_hash}_{int(time.time())}"
+    if overwrite and job_dir.exists():
+        shutil.rmtree(job_dir)
+    job_dir.mkdir(parents=True, exist_ok=True)
+    train_output_dir = job_dir / "lora"
+    progress_file = str(job_dir / "progress.json")
+    managed_dir = _next_managed_lora_dir(safe_name, overwrite=overwrite)
+    continue_lora = _resolve_continue_lora(continue_from)
+
+    register_training_job(
+        node_id,
+        engine_type="dramabox",
+        progress_file=progress_file,
+        job_dir=str(job_dir),
+        model_name=safe_name,
+        sample_rate="48k",
+        total_epochs=1,
+    )
+    try:
+        _write_progress(
+            progress_file,
+            status="starting",
+            phase="setup",
+            engine_type="dramabox",
+            model_name=safe_name,
+            dataset_records=int(dataset_info.get("train_records", 0)),
+            speakers=dataset_info.get("speakers", []),
+            started_at=time.time(),
+        )
+
+        if not bool(dataset_info.get("preprocessed")):
+            if bool(training_config.get("dry_run", False)):
+                print("🧪 DramaBox dry-run: skipping GPU dataset preprocessing")
+            else:
+                _write_progress(progress_file, status="running", phase="preprocess")
+                run_dramabox_preprocess(
+                    dataset_info,
+                    shared_settings,
+                    batch_size=int(training_config.get("preprocess_batch_size", 8)),
+                    progress_file=progress_file,
+                    node_id=node_id,
+                )
+
+        train_config = _build_train_config(
+            dataset_info,
+            shared_settings,
+            training_config,
+            output_dir=train_output_dir,
+            continue_lora=continue_lora,
+            resolve_paths=not bool(training_config.get("dry_run", False)),
+        )
+        config_path = job_dir / "training_config.yaml"
+        import yaml
+
+        config_path.write_text(yaml.safe_dump(train_config, sort_keys=False), encoding="utf-8")
+        (job_dir / "resolved_training_config.json").write_text(
+            json.dumps(
+                {
+                    "dataset": dataset_info,
+                    "shared_settings": shared_settings,
+                    "training_config": training_config,
+                    "official_config": train_config,
+                    "continue_from": continue_lora,
+                },
+                indent=2,
+                ensure_ascii=False,
+                default=str,
+            ),
+            encoding="utf-8",
+        )
+
+        command = [*_accelerate_command(), str(TRAIN_SCRIPT), "--config", str(config_path)]
+        if bool(training_config.get("dry_run", False)):
+            summary = (
+                f"DramaBox dry-run ready: {safe_name} | {dataset_info.get('train_records', 0)} rows | "
+                f"official command prepared without loading CUDA or model weights"
+            )
+            _write_progress(
+                progress_file,
+                status="completed",
+                phase="dry_run",
+                overall_progress=1.0,
+                summary=summary,
+                command=command,
+            )
+            finalize_training_job(node_id, status="completed", summary=summary, dry_run=True)
+            return {
+                "type": "training_artifacts",
+                "engine_type": "dramabox",
+                "training_mode": "audio_lora",
+                "dry_run": True,
+                "job_dir": str(job_dir),
+                "training_config": str(config_path),
+                "summary": summary,
+                "command": command,
+            }
+
+        _write_progress(progress_file, status="running", phase="train", total_steps=int(train_config["steps"]))
+        train_env = _device_environment(shared_settings)
+        if train_config.get("val_config"):
+            paths = _resolve_model_paths(shared_settings)
+            train_env["LTX_CHECKPOINT"] = paths["transformer"]
+            train_env["LTX_FULL_CHECKPOINT"] = paths["audio_components"]
+            train_env["GEMMA_ROOT"] = paths["gemma_root"]
+            train_env["TRAIN_VAL_GPU"] = _validation_gpu(
+                shared_settings.get("device", "auto"),
+                training_config.get("validation_gpu", ""),
+            )
+        _run_process(
+            command,
+            cwd=VENDOR_ROOT,
+            env=train_env,
+            phase="train",
+            progress_file=progress_file,
+            node_id=node_id,
+            total_steps=int(train_config["steps"]),
+        )
+
+        selected_lora = _latest_lora_file(train_output_dir)
+        if selected_lora is None:
+            raise RuntimeError(
+                f"DramaBox training exited successfully but produced no LoRA file in '{train_output_dir}'."
+            )
+        if managed_dir.exists():
+            shutil.rmtree(managed_dir)
+        managed_dir.mkdir(parents=True, exist_ok=True)
+        managed_lora = managed_dir / selected_lora.name
+        shutil.copy2(selected_lora, managed_lora)
+        if selected_lora.name != "adapter_model.safetensors":
+            shutil.copy2(selected_lora, managed_dir / "adapter_model.safetensors")
+        adapter_config = train_output_dir / "adapter_config.json"
+        if adapter_config.is_file():
+            shutil.copy2(adapter_config, managed_dir / adapter_config.name)
+        shutil.copy2(config_path, managed_dir / "training_config.yaml")
+
+        summary = (
+            f"DramaBox audio LoRA training complete: {safe_name} | "
+            f"steps={train_config['steps']} | adapter={managed_lora}"
+        )
+        _write_progress(
+            progress_file,
+            status="completed",
+            phase="done",
+            overall_progress=1.0,
+            output_adapter=str(managed_lora),
+            output_dir=str(managed_dir),
+            summary=summary,
+        )
+        finalize_training_job(
+            node_id,
+            status="completed",
+            output_adapter=str(managed_lora),
+            summary=summary,
+        )
+        return {
+            "type": "training_artifacts",
+            "engine_type": "dramabox",
+            "training_mode": "audio_lora",
+            "model_path": str(managed_dir),
+            "lora_path": str(managed_lora),
+            "job_dir": str(job_dir),
+            "summary": summary,
+            "lora_adapter": {
+                "type": "dramabox_lora",
+                "adapter_path": str(managed_lora),
+                "adapter_dir": str(managed_dir),
+            },
+        }
+    except InterruptedError as error:
+        _write_progress(progress_file, status="cancelled", phase="cancelled", error=str(error))
+        finalize_training_job(node_id, status="cancelled", error=str(error))
+        raise
+    except Exception as error:
+        _write_progress(progress_file, status="error", phase="error", error=str(error))
+        finalize_training_job(node_id, status="error", error=str(error))
+        raise
+
+
+__all__ = [
+    "build_preprocess_command",
+    "run_dramabox_preprocess",
+    "run_dramabox_training_job",
+]

engines/dramabox/vendor/UPSTREAM.md

@@ -1,4 +1,4 @@
-# Bundled DramaBox inference source
+# Bundled DramaBox inference and training source
 
 This directory contains the inference-critical source copied unchanged from:
 
@@ -15,4 +15,27 @@ The bundled-code changes are marked inline:
 - `src/inference_server.py`: ComfyUI cancellation exceptions are allowed to
   propagate from progress callbacks instead of being swallowed; the official
   negative-prompt, FP8-cast, compile, and staged-memory controls are exposed to
-  the suite wrapper.
+  the suite wrapper; suite-managed PEFT LoRA loading is added for trained
+  DramaBox audio adapters.
+- `src/validate.py`: validation accepts the suite's separately organized
+  DramaBox transformer and audio-components checkpoints.
+- `src/preprocess.py`: suite-distributed pre-quantized Gemma checkpoints use
+  the same bitsandbytes-aware prompt-encoder loader as DramaBox inference.
+- `src/train.py`: the batch collator lives at module scope so Windows
+  spawn-based DataLoader workers can serialize it; lightweight per-step
+  telemetry keeps the suite's training dashboard current between normal logs.
+- `ltx2/ltx_core/text_encoders/gemma/encoders/encoder_configurator.py`:
+  supports both wrapped and direct SigLIP vision-tower layouts for the suite's
+  newer Transformers runtime.
+
+The official training entry points are also bundled at this pin:
+
+- `src/preprocess.py`
+- `src/train.py`
+- `src/validate.py`
+- `configs/training_args.example.yaml`
+- `configs/val_config.example.yaml`
+
+The suite invokes these scripts through the unified training backend. Apart
+from the documented compatibility patches, training behavior stays upstream;
+dataset normalization, job lifecycle, and UI wiring remain suite-side.

engines/dramabox/vendor/configs/training_args.example.yaml

@@ -0,0 +1,63 @@
+# DramaBox IC-LoRA training config — values become the defaults for
+# `accelerate launch src/train.py --config configs/training_args.example.yaml`.
+# Any flag explicitly passed on the CLI overrides the YAML.
+
+# ── Data ───────────────────────────────────────────────────────────────────
+# One entry per preprocessed dataset (output dirs from src/preprocess.py).
+data_dir:
+  - /path/to/preprocessed_dataset_a/
+  - /path/to/preprocessed_dataset_b/
+
+# One index file per data_dir entry. Each line follows the format you fed to
+# preprocess.py — see README "Prepare your index file".
+speaker_index:
+  - /path/to/preprocessed_dataset_a/index.txt
+  - /path/to/preprocessed_dataset_b/index.txt
+
+# Output directory for LoRA shards + logs (relative paths resolve against the
+# repo root).
+output_dir: tts_iclora_v1
+
+# ── Base model ─────────────────────────────────────────────────────────────
+# Train your LoRA on top of DramaBox itself (recommended) — the trimmed audio
+# components are enough; no need to ship the raw LTX-2.3 base.
+checkpoint: dramabox-dit-v1.safetensors
+full_checkpoint: dramabox-audio-components.safetensors
+base_model: dev          # 'dev' = ShiftedLogitNormal sampler; 'distilled' = DistilledTimestepSampler
+
+# ── LoRA hyperparams (rank == alpha → scale = 1.0) ─────────────────────────
+lora_rank: 128
+lora_alpha: 128
+lora_dropout: 0.1        # ~0.1 helps regularize on small datasets
+
+# Resume an existing LoRA — step number parsed from the filename
+# (e.g. lora_step_05000.safetensors → starts at step 5000).
+# resume_lora: tts_iclora_v0/lora_step_05000.safetensors
+
+# ── Voice-cloning reference tokens ─────────────────────────────────────────
+ref_ratio: 0.3           # fraction of training samples that get a ref-token tail
+max_ref_tokens: 200      # cap on appended ref tokens after patchification
+
+# CFG training: probability of zeroing the text condition (forces reliance on
+# the voice ref / unconditional path).
+text_dropout: 0.4
+
+# ── Schedule ───────────────────────────────────────────────────────────────
+# Cosine + 1e-4 = from-scratch fine-tune.
+# Constant + 1e-5 = polish on top of an existing LoRA (use with `resume_lora`).
+steps: 10000
+lr: 1.0e-04
+lr_scheduler: cosine
+warmup_steps: 500
+
+batch_size: 1
+grad_accum: 4
+max_grad_norm: 1.0
+
+save_every: 500
+log_every: 50
+seed: 53
+
+# Optional per-save-step validation pass. Generates a sample for every speaker
+# in the val_config so you can A/B listen during training.
+# val_config: configs/val_config.example.yaml

engines/dramabox/vendor/configs/val_config.example.yaml

@@ -0,0 +1,25 @@
+# Validation prompts run by src/validate.py at every --save-every checkpoint.
+# Each entry produces one .wav under <output_dir>/val_step_<N>/<name>.wav.
+#
+# Fields:
+#   name      — short tag used as the output filename
+#   prompt    — full DramaBox-style scene prompt
+#   reference — (optional) absolute path to a 10+ s voice reference clip;
+#               omit for prompt-only generation
+
+speakers:
+  - name: villain_growl
+    prompt: 'A shadowy villain speaks with cold menace, "You have entered my domain, mortal." He chuckles darkly, "Such arrogance will be your undoing."'
+    reference: /path/to/voice_refs/male_villain.wav
+
+  - name: tender_whisper
+    prompt: 'A woman speaks tenderly, "It has been a long day, my love." She whispers, "Close your eyes. I am right here."'
+    reference: /path/to/voice_refs/female_warm.wav
+
+  - name: catgirl_giggle
+    prompt: 'A playful girl already mid-giggle, "Hehehe, oh my gosh you should see your face!" She gasps, "Oh my, hehe, I cannot stop!"'
+    # No `reference:` here — pure prompt-driven generation.
+
+  - name: announcer_smug
+    prompt: 'A confident announcer speaks proudly, "And now, the moment you have all been waiting for." He chuckles knowingly, "Heheh."'
+    reference: /path/to/voice_refs/male_announcer.wav

engines/dramabox/vendor/ltx2/ltx_core/text_encoders/gemma/encoders/encoder_configurator.py

@@ -154,22 +154,48 @@ VIDEO_ONLY_EMBEDDINGS_PROCESSOR_KEY_OPS = (
 
 def create_and_populate(module: GemmaTextEncoder) -> GemmaTextEncoder:
     model = module.model
-    v_model = model.model.vision_tower.vision_model
+    vision_tower = model.model.vision_tower
+    # TTS Audio Suite patch: Transformers 5 exposes SiglipVisionModel
+    # directly, while the upstream-pinned layout wraps it in `.vision_model`.
+    v_model = getattr(vision_tower, "vision_model", vision_tower)
     l_model = model.model.language_model
 
     config = model.config.text_config
     dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
-    base = config.rope_local_base_freq
+    if hasattr(config, "rope_local_base_freq"):
+        base = config.rope_local_base_freq
+        rope_type = config.rope_scaling["rope_type"]
+        rope_kwargs = {}
+    else:
+        # TTS Audio Suite patch: Transformers 5 migrates Gemma 3's local and
+        # full-attention RoPE settings into named `rope_parameters` entries.
+        rope_parameters = config.rope_parameters
+        base = rope_parameters["sliding_attention"]["rope_theta"]
+        rope_type = rope_parameters["full_attention"]["rope_type"]
+        rope_kwargs = {"layer_type": "full_attention"}
     local_rope_freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(dtype=torch.float) / dim))
-    inv_freqs, _ = ROPE_INIT_FUNCTIONS[config.rope_scaling["rope_type"]](config)
+    inv_freqs, _ = ROPE_INIT_FUNCTIONS[rope_type](config, **rope_kwargs)
 
     positions_length = len(v_model.embeddings.position_ids[0])
     position_ids = torch.arange(positions_length, dtype=torch.long, device="cpu").unsqueeze(0)
     v_model.embeddings.register_buffer("position_ids", position_ids)
     embed_scale = torch.tensor(model.config.text_config.hidden_size**0.5, device="cpu")
     l_model.embed_tokens.register_buffer("embed_scale", embed_scale)
-    l_model.rotary_emb_local.register_buffer("inv_freq", local_rope_freqs)
-    l_model.rotary_emb.register_buffer("inv_freq", inv_freqs)
+    if hasattr(l_model, "rotary_emb_local"):
+        l_model.rotary_emb_local.register_buffer("inv_freq", local_rope_freqs)
+        l_model.rotary_emb.register_buffer("inv_freq", inv_freqs)
+    else:
+        # TTS Audio Suite patch: Transformers 5 consolidates both attention
+        # variants into one rotary module with separately named buffers.
+        rotary_emb = l_model.rotary_emb
+        rotary_emb.register_buffer("sliding_attention_inv_freq", local_rope_freqs)
+        rotary_emb.register_buffer(
+            "sliding_attention_original_inv_freq", local_rope_freqs.clone()
+        )
+        rotary_emb.register_buffer("full_attention_inv_freq", inv_freqs)
+        rotary_emb.register_buffer(
+            "full_attention_original_inv_freq", inv_freqs.clone()
+        )
 
     return module
 

engines/dramabox/vendor/src/inference_server.py

@@ -125,7 +125,8 @@ def auto_rescale_for_cfg(cfg: float) -> float:
 class TTSServer:
     def __init__(self, checkpoint=None, full_checkpoint=None, gemma_root=None,
                  device="cuda", dtype="bf16", compile_model=True, bnb_4bit=True,
-                 memory_mode="fast", transformer_quantization="none"):
+                 memory_mode="fast", transformer_quantization="none",
+                 lora_path="", lora_strength=1.0):
         MODELS = APP_DIR / "models"
         self.checkpoint = checkpoint or str(MODELS / "ltx-2.3-22b-dev-audio-only-v13-merged.safetensors")
         self.full_checkpoint = full_checkpoint or os.environ.get(
@@ -140,6 +141,14 @@ class TTSServer:
         self.bnb_4bit = bnb_4bit
         self.memory_mode = str(memory_mode)
         self.transformer_quantization = str(transformer_quantization)
+        # TTS Audio Suite patch: accept a trained DramaBox audio LoRA at
+        # runtime so training artifacts can be used without a second CLI.
+        self.lora_path = str(lora_path or "").strip()
+        self.lora_strength = float(lora_strength)
+        self._active_lora_revision = ""
+        self._active_lora_file = ""
+        self._applied_lora_strength = 0.0
+        self._unmerged_lora_weight_scale = 1.0
         if self.memory_mode not in {"fast", "staged", "sequential"}:
             raise ValueError(f"Unknown DramaBox memory mode: {self.memory_mode}")
         if self.transformer_quantization not in {"none", "fp8_cast"}:
@@ -262,6 +271,8 @@ class TTSServer:
         self._velocity_model = builder.build(
             device=self.device, dtype=build_dtype
         ).to(self.device).eval()
+        if self.lora_path and self.lora_strength != 0.0:
+            self.configure_lora(self.lora_path, self.lora_strength)
         n_params = sum(p.numel() for p in self._velocity_model.parameters()) / 1e9
         vram_gb = sum(p.numel() * p.element_size() for p in self._velocity_model.parameters()) / 1e9
         logging.info(f"  Transformer: {time.time()-t0:.1f}s ({n_params:.1f}B params, {vram_gb:.1f}GB VRAM, {self.dtype})")
@@ -287,6 +298,260 @@ class TTSServer:
         )
         logging.info(f"  AudioDecoder (warm): {time.time()-t0:.1f}s")
 
+    @staticmethod
+    def _resolve_lora_file(lora_path: str) -> Path:
+        path = Path(os.path.expanduser(str(lora_path or "").strip()))
+        if path.is_dir():
+            candidates = sorted(path.glob("lora_step_*.safetensors"))
+            candidates += [path / "adapter_model.safetensors"]
+            for candidate in reversed(candidates):
+                if candidate.is_file():
+                    return candidate
+        if path.is_file():
+            return path
+        raise FileNotFoundError(f"DramaBox LoRA file not found: {lora_path}")
+
+    # TTS Audio Suite patch: keep PEFT state attached and reversibly merge it
+    # so strength changes avoid both a base reload and per-step LoRA matmuls.
+    @staticmethod
+    def _set_lora_strength(model, strength: float) -> None:
+        """Re-merge the live adapter at a new strength without reloading the base."""
+        try:
+            from peft.tuners.lora.layer import LoraLayer
+        except ImportError as exc:
+            raise RuntimeError("DramaBox LoRA inference requires peft.") from exc
+
+        updated = 0
+        if any(
+            isinstance(module, LoraLayer) and bool(module.merged)
+            for module in model.modules()
+        ):
+            model.unmerge_adapter()
+
+        for module in model.modules():
+            if isinstance(module, LoraLayer) and "default" in module.lora_A:
+                module.set_scale("default", float(strength))
+                updated += 1
+        if updated <= 0:
+            raise RuntimeError("DramaBox LoRA modules are missing from the live model.")
+        if hasattr(model, "set_adapter"):
+            model.set_adapter("default")
+        if float(strength) == 0.0:
+            model.disable_adapter_layers()
+        else:
+            model.enable_adapter_layers()
+            model.merge_adapter(adapter_names=["default"])
+
+    @staticmethod
+    def _set_unmerged_lora_strength(
+        model, strength: float, current_weight_scale: float
+    ) -> float:
+        """Scale a BF16 PEFT branch over an immutable FP8 base in place."""
+        try:
+            from peft.tuners.lora.layer import LoraLayer
+        except ImportError as exc:
+            raise RuntimeError("DramaBox LoRA inference requires peft.") from exc
+
+        if float(strength) == 0.0:
+            model.disable_adapter_layers()
+            return float(current_weight_scale)
+
+        model.enable_adapter_layers()
+        if hasattr(model, "set_adapter"):
+            model.set_adapter("default")
+        ratio = float(strength) / float(current_weight_scale)
+        updated = 0
+        for module in model.modules():
+            if isinstance(module, LoraLayer) and "default" in module.lora_A:
+                if ratio != 1.0:
+                    with torch.no_grad():
+                        module.lora_B["default"].weight.mul_(ratio)
+                updated += 1
+        if updated <= 0:
+            raise RuntimeError("DramaBox LoRA modules are missing from the live model.")
+        return float(strength)
+
+    def _prepare_unmerged_lora(self, model) -> None:
+        """Keep PEFT matrices in the activation dtype used above FP8 storage."""
+        from peft.tuners.lora.layer import LoraLayer
+
+        for module in model.modules():
+            if isinstance(module, LoraLayer) and "default" in module.lora_A:
+                module.lora_A["default"].to(device=self.device, dtype=self.dtype)
+                module.lora_B["default"].to(device=self.device, dtype=self.dtype)
+
+    # TTS Audio Suite patch: replace only the live adapter modules while
+    # preserving the already-loaded official DramaBox transformer weights.
+    @classmethod
+    def _attach_lora(cls, model, lora_path: str, strength: float):
+        """Attach or replace the official PEFT-compatible audio LoRA in place."""
+        try:
+            from peft import LoraConfig, PeftModel, get_peft_model
+            from safetensors.torch import load_file
+        except ImportError as exc:
+            raise RuntimeError(
+                "DramaBox LoRA inference requires peft and safetensors."
+            ) from exc
+
+        lora_file = cls._resolve_lora_file(lora_path)
+        adapter_config_path = lora_file.parent / "adapter_config.json"
+        rank = 128
+        alpha = 128
+        if adapter_config_path.is_file():
+            try:
+                metadata = json.loads(adapter_config_path.read_text(encoding="utf-8"))
+                rank = int(metadata.get("r", rank))
+                alpha = int(metadata.get("lora_alpha", alpha))
+            except Exception as exc:
+                logging.warning("Could not read DramaBox LoRA adapter_config.json: %s", exc)
+
+        lora_state = load_file(str(lora_file))
+        # Standalone upstream checkpoints may omit adapter_config.json. Infer
+        # the rank from the first LoRA-A tensor so those files remain usable.
+        if not adapter_config_path.is_file():
+            for key, value in lora_state.items():
+                if ".lora_A." in key or key.endswith(".lora_A.weight"):
+                    rank = int(value.shape[0])
+                    alpha = rank
+                    break
+        lora_config = LoraConfig(
+            r=rank,
+            lora_alpha=alpha,
+            lora_dropout=0.0,
+            bias="none",
+            target_modules=[
+                "audio_attn1.to_k",
+                "audio_attn1.to_q",
+                "audio_attn1.to_v",
+                "audio_attn1.to_out.0",
+                "audio_attn2.to_k",
+                "audio_attn2.to_q",
+                "audio_attn2.to_v",
+                "audio_attn2.to_out.0",
+                "audio_ff.net.0.proj",
+                "audio_ff.net.2",
+            ],
+        )
+        if isinstance(model, PeftModel):
+            if any(
+                hasattr(module, "merged") and bool(module.merged)
+                for module in model.modules()
+            ):
+                model.unmerge_adapter()
+            if "default" in model.peft_config:
+                model.delete_adapter("default")
+            model.add_adapter("default", lora_config)
+            adapted = model
+        else:
+            adapted = get_peft_model(model, lora_config)
+        mapped = {}
+        is_peft_format = any("base_model.model." in key for key in lora_state)
+        is_original_format = any("diffusion_model." in key for key in lora_state)
+        compiled_blocks = any("._orig_mod." in key for key in adapted.state_dict())
+        for key, value in lora_state.items():
+            if is_peft_format:
+                new_key = key
+            elif is_original_format:
+                new_key = key.replace("diffusion_model.", "base_model.model.")
+            else:
+                continue
+            new_key = new_key.replace(".lora_A.weight", ".lora_A.default.weight")
+            new_key = new_key.replace(".lora_B.weight", ".lora_B.default.weight")
+            if compiled_blocks and "._orig_mod." not in new_key:
+                # TTS Audio Suite patch: torch.compile wraps every official
+                # transformer block in OptimizedModule and inserts `_orig_mod`
+                # into its state-dict path before PEFT attaches the adapter.
+                new_key = re.sub(
+                    r"(transformer_blocks\.\d+)\.",
+                    r"\1._orig_mod.",
+                    new_key,
+                    count=1,
+                )
+            mapped[new_key] = value
+        if not mapped:
+            raise RuntimeError(
+                f"DramaBox LoRA '{lora_file}' is not in a recognized PEFT/ID-LoRA format."
+            )
+
+        missing, unexpected = adapted.load_state_dict(mapped, strict=False)
+        loaded = len(mapped) - len(unexpected)
+        if loaded <= 0:
+            raise RuntimeError(
+                f"DramaBox LoRA '{lora_file}' did not match the audio transformer modules."
+            )
+        logging.info(
+            "DramaBox LoRA loaded: %s (%d tensors, strength %.2f)",
+            lora_file,
+            loaded,
+            float(strength),
+        )
+        return adapted.eval(), str(lora_file.resolve())
+
+    # TTS Audio Suite patch: split mutable adapter identity from the expensive
+    # base-model cache identity used by the suite's ComfyUI model wrapper.
+    def configure_lora(self, lora_path: str, strength: float, revision: str = "") -> None:
+        """Hot-swap a DramaBox adapter or update only its runtime strength."""
+        path = str(lora_path or "").strip()
+        strength = float(strength)
+        revision = str(revision or "")
+        use_unmerged_fp8 = self.transformer_quantization == "fp8_cast"
+
+        if not path:
+            if self._active_lora_file:
+                if use_unmerged_fp8:
+                    self._unmerged_lora_weight_scale = self._set_unmerged_lora_strength(
+                        self._velocity_model, 0.0, self._unmerged_lora_weight_scale
+                    )
+                else:
+                    self._set_lora_strength(self._velocity_model, 0.0)
+                logging.info("DramaBox LoRA disabled without reloading the base model")
+            self.lora_path = ""
+            self.lora_strength = strength
+            self._applied_lora_strength = 0.0
+            return
+
+        lora_file = str(self._resolve_lora_file(path).resolve())
+        same_adapter = (
+            lora_file == self._active_lora_file
+            and revision == self._active_lora_revision
+        )
+        if same_adapter:
+            if strength != self._applied_lora_strength:
+                if use_unmerged_fp8:
+                    self._unmerged_lora_weight_scale = self._set_unmerged_lora_strength(
+                        self._velocity_model,
+                        strength,
+                        self._unmerged_lora_weight_scale,
+                    )
+                else:
+                    self._set_lora_strength(self._velocity_model, strength)
+                logging.info(
+                    "DramaBox LoRA strength updated in place: %.2f", strength
+                )
+        else:
+            self._velocity_model, lora_file = self._attach_lora(
+                self._velocity_model, path, strength
+            )
+            self._active_lora_file = lora_file
+            self._active_lora_revision = revision
+            if use_unmerged_fp8:
+                self._prepare_unmerged_lora(self._velocity_model)
+                self._unmerged_lora_weight_scale = 1.0
+                self._unmerged_lora_weight_scale = self._set_unmerged_lora_strength(
+                    self._velocity_model,
+                    strength,
+                    self._unmerged_lora_weight_scale,
+                )
+                logging.info(
+                    "DramaBox FP8 base: using an unmerged BF16 LoRA branch"
+                )
+            else:
+                self._set_lora_strength(self._velocity_model, strength)
+
+        self.lora_path = path
+        self.lora_strength = strength
+        self._applied_lora_strength = strength
+
     def _move_velocity_model(self, target: torch.device) -> None:
         """Move the persistent DiT between CUDA and RAM for staged inference."""
         target = torch.device(target)

engines/dramabox/vendor/src/preprocess.py

@@ -0,0 +1,384 @@
+#!/usr/bin/env python3
+"""
+Preprocess TTS datasets for LTX-2.3 audio-only LoRA fine-tuning.
+
+Takes paired (audio, transcript) data and produces the format expected by
+the LTX trainer:
+    .precomputed/
+    ├── latents/sample_N.pt         # Dummy video latents (minimal)
+    ├── conditions/sample_N.pt      # Text embeddings from Gemma
+    └── audio_latents/sample_N.pt   # Audio VAE-encoded latents
+
+Supports multiple dataset formats:
+  - gemini_synthetic: index.txt with ~-separated fields (id~speaker~lang~sr~samples~dur~phonemes~text)
+  - libriheavy: index_ft.txt with ~-separated fields (id~speaker~lang~samples~dur~phonemes~text)
+  - manifest: JSON/JSONL with {"audio_filepath": ..., "text": ...}
+  - tsv: TSV file with audio_path<TAB>text columns
+
+Usage:
+    python preprocess_tts_data.py \
+        --dataset-type gemini_synthetic \
+        --index /path/to/dataset/index.txt \
+        --audio-dir /path/to/dataset/wavs \
+        --output-dir /path/to/output/tts_training_data \
+        --max-samples 10000 \
+        --max-duration 20.0 \
+        --min-duration 3.0
+"""
+
+import argparse
+import json
+import logging
+import os
+import sys
+from pathlib import Path
+
+import torch
+import torchaudio
+
+REPO_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ltx2"))
+# ltx-pipelines on path via ltx2/
+
+MODEL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+GEMMA_DIR = os.environ.get("GEMMA_DIR", "gemma-3-12b-it-qat-q4_0-unquantized")
+
+
+def parse_args():
+    p = argparse.ArgumentParser(description="Preprocess TTS data for LTX-2.3 fine-tuning")
+    p.add_argument("--dataset-type", required=True,
+                   choices=["gemini_synthetic", "libriheavy", "manifest", "tsv"],
+                   help="Dataset format type")
+    p.add_argument("--index", required=True, help="Path to index/manifest file")
+    p.add_argument("--audio-dir", default=None,
+                   help="Base directory for audio files (if paths in index are relative)")
+    p.add_argument("--output-dir", required=True, help="Output directory for preprocessed data")
+    p.add_argument("--checkpoint", default=os.path.join(MODEL_DIR, "ltx-2.3-22b-distilled.safetensors"))
+    p.add_argument("--gemma-root", default=GEMMA_DIR)
+    p.add_argument("--max-samples", type=int, default=0, help="Max samples to process (0=all)")
+    p.add_argument("--max-duration", type=float, default=20.0, help="Max audio duration in seconds")
+    p.add_argument("--min-duration", type=float, default=2.0, help="Min audio duration in seconds")
+    p.add_argument("--batch-size", type=int, default=8, help="Batch size for text encoding")
+    p.add_argument("--skip-existing", action="store_true", help="Skip already processed samples")
+    p.add_argument("--audio-only-ckpt", default=None,
+                   help="Audio-only checkpoint for VAE encoding (optional, uses full ckpt if not set)")
+    p.add_argument("--shard", type=int, default=0, help="Shard index (for parallel processing)")
+    p.add_argument("--num-shards", type=int, default=1, help="Total number of shards")
+    p.add_argument("--gpu", type=int, default=None, help="GPU device index to use")
+    return p.parse_args()
+
+
+def parse_gemini_synthetic(index_path: str, audio_dir: str | None) -> list[dict]:
+    """Parse gemini_synthetic format: id~speaker~lang~sr~samples~dur~phonemes~text"""
+    samples = []
+    with open(index_path) as f:
+        for line in f:
+            parts = line.strip().split("~")
+            if len(parts) < 7:
+                continue
+            file_id = parts[0]
+            text = parts[-1]  # Last field is always the text
+            sr = int(parts[3])
+            n_samples = int(parts[4])
+            duration = n_samples / sr
+
+            # Find audio file
+            if audio_dir:
+                # Try common extensions
+                for ext in [".flac", ".wav", ".mp3"]:
+                    audio_path = os.path.join(audio_dir, file_id + ext)
+                    if os.path.exists(audio_path):
+                        break
+                else:
+                    continue
+            else:
+                audio_path = file_id
+
+            samples.append({
+                "id": file_id,
+                "audio_path": audio_path,
+                "text": text,
+                "duration": duration,
+            })
+    return samples
+
+
+def parse_libriheavy(index_path: str, audio_dir: str | None) -> list[dict]:
+    """Parse libriheavy format: id~speaker~lang~samples~dur~phonemes~text"""
+    samples = []
+    with open(index_path) as f:
+        for line in f:
+            parts = line.strip().split("~")
+            if len(parts) < 7:
+                continue
+            file_id = parts[0]
+            text = parts[-1]
+            n_samples = int(parts[3])
+            duration = int(parts[4]) / 1000.0  # milliseconds to seconds
+
+            if audio_dir:
+                for ext in [".flac", ".wav", ".mp3"]:
+                    audio_path = os.path.join(audio_dir, file_id + ext)
+                    if os.path.exists(audio_path):
+                        break
+                else:
+                    continue
+            else:
+                audio_path = file_id
+
+            samples.append({
+                "id": file_id,
+                "audio_path": audio_path,
+                "text": text,
+                "duration": duration,
+            })
+    return samples
+
+
+def parse_manifest(index_path: str, audio_dir: str | None) -> list[dict]:
+    """Parse JSON/JSONL manifest with audio_filepath and text fields."""
+    samples = []
+    with open(index_path) as f:
+        for line in f:
+            entry = json.loads(line.strip())
+            audio_path = entry.get("audio_filepath", entry.get("audio_path", ""))
+            text = entry.get("text", entry.get("transcript", ""))
+            duration = entry.get("duration", 0.0)
+
+            if audio_dir and not os.path.isabs(audio_path):
+                audio_path = os.path.join(audio_dir, audio_path)
+
+            if os.path.exists(audio_path) and text:
+                samples.append({
+                    "id": Path(audio_path).stem,
+                    "audio_path": audio_path,
+                    "text": text,
+                    "duration": duration,
+                })
+    return samples
+
+
+def parse_tsv(index_path: str, audio_dir: str | None) -> list[dict]:
+    """Parse TSV file with audio_path<TAB>text."""
+    samples = []
+    with open(index_path) as f:
+        for line in f:
+            parts = line.strip().split("\t")
+            if len(parts) < 2:
+                continue
+            audio_path, text = parts[0], parts[1]
+            if audio_dir and not os.path.isabs(audio_path):
+                audio_path = os.path.join(audio_dir, audio_path)
+            if os.path.exists(audio_path):
+                samples.append({
+                    "id": Path(audio_path).stem,
+                    "audio_path": audio_path,
+                    "text": text,
+                    "duration": 0.0,
+                })
+    return samples
+
+
+PARSERS = {
+    "gemini_synthetic": parse_gemini_synthetic,
+    "libriheavy": parse_libriheavy,
+    "manifest": parse_manifest,
+    "tsv": parse_tsv,
+}
+
+
+@torch.inference_mode()
+def main():
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+    args = parse_args()
+
+    from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
+    from ltx_core.types import Audio
+    from ltx_pipelines.utils.blocks import AudioConditioner, PromptEncoder
+    from ltx_pipelines.utils.media_io import decode_audio_from_file
+    from ltx_trainer.model_loader import load_text_encoder, load_embeddings_processor
+
+    if args.gpu is not None:
+        os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = torch.bfloat16
+
+    # Create output directories
+    out = Path(args.output_dir)
+    (out / "latents").mkdir(parents=True, exist_ok=True)
+    (out / "conditions").mkdir(parents=True, exist_ok=True)
+    (out / "audio_latents").mkdir(parents=True, exist_ok=True)
+
+    # Parse dataset
+    logging.info(f"Parsing {args.dataset_type} dataset from {args.index}...")
+    samples = PARSERS[args.dataset_type](args.index, args.audio_dir)
+    logging.info(f"Found {len(samples)} samples")
+
+    # Filter by duration
+    before = len(samples)
+    samples = [s for s in samples if args.min_duration <= s["duration"] <= args.max_duration]
+    logging.info(f"After duration filter [{args.min_duration}s, {args.max_duration}s]: {len(samples)} (dropped {before - len(samples)})")
+
+    if args.max_samples > 0:
+        samples = samples[:args.max_samples]
+        logging.info(f"Limiting to {len(samples)} samples")
+
+    # Assign global indices before sharding
+    for i, s in enumerate(samples):
+        s["global_idx"] = i
+
+    # Shard the data for parallel processing
+    if args.num_shards > 1:
+        total = len(samples)
+        samples = samples[args.shard::args.num_shards]
+        logging.info(f"Shard {args.shard}/{args.num_shards}: {len(samples)} samples (of {total} total)")
+
+    # ── Step 1: Encode text with Gemma (Blocks 1+2 only) ──
+    # The trainer runs Block 3 (embeddings processor/connectors) during training,
+    # so we only precompute Blocks 1+2 here (Gemma LLM + feature extractor).
+    logging.info("Loading text encoder (Gemma + feature extractor)...")
+    gemma_config_path = Path(args.gemma_root) / "config.json"
+    gemma_config = {}
+    if gemma_config_path.is_file():
+        try:
+            gemma_config = json.loads(gemma_config_path.read_text(encoding="utf-8"))
+        except Exception:
+            pass
+
+    prompt_encoder = None
+    if "quantization_config" in gemma_config:
+        # TTS Audio Suite patch: the suite distributes a pre-quantized BNB
+        # Gemma checkpoint. The official tensor builder treats its packed
+        # weights as dense matrices, producing thousands of shape mismatches.
+        # Reuse the inference loader that already understands this format.
+        logging.info("Loading pre-quantized Gemma through the BNB prompt encoder...")
+        prompt_encoder = PromptEncoder(
+            checkpoint_path=args.checkpoint,
+            gemma_root=args.gemma_root,
+            dtype=dtype,
+            device=device,
+            warm=True,
+            use_bnb_4bit=True,
+            audio_only=True,
+        )
+        text_encoder = prompt_encoder._warm_text_encoder
+        embeddings_processor = prompt_encoder._warm_embeddings_processor
+        text_encoder.feature_extractor = embeddings_processor.feature_extractor
+        prompt_encoder._warm_text_encoder = None
+        prompt_encoder._warm_embeddings_processor = None
+    else:
+        text_encoder = load_text_encoder(args.gemma_root, device=device, dtype=dtype)
+
+        # Load feature extractor on CPU first to save GPU memory, then move to device
+        logging.info("Loading feature extractor (on CPU first to save GPU memory)...")
+        emb_proc = load_embeddings_processor(args.checkpoint, device="cpu", dtype=dtype)
+        text_encoder.feature_extractor = emb_proc.feature_extractor.to(device)
+        del emb_proc
+    torch.cuda.empty_cache()
+
+    logging.info("Encoding text prompts (Blocks 1+2: Gemma + feature extractor)...")
+    for i, sample in enumerate(samples):
+        gidx = sample["global_idx"]
+        cond_path = out / "conditions" / f"sample_{gidx:06d}.pt"
+        if args.skip_existing and cond_path.exists():
+            continue
+
+        text = sample["text"]
+        # Run Blocks 1+2: Gemma LLM → feature extractor
+        hidden_states, attention_mask = text_encoder.encode(text)
+        video_feats, audio_feats = text_encoder.feature_extractor(
+            hidden_states, attention_mask, "left"
+        )
+
+        torch.save({
+            "video_prompt_embeds": video_feats.squeeze(0).cpu(),
+            "audio_prompt_embeds": audio_feats.squeeze(0).cpu() if audio_feats is not None else video_feats.squeeze(0).cpu(),
+            "prompt_attention_mask": attention_mask.squeeze(0).bool().cpu(),
+        }, cond_path)
+
+        if i % 100 == 0:
+            logging.info(f"  Text encoding: {i}/{len(samples)}")
+
+    del text_encoder
+    if prompt_encoder is not None:
+        del embeddings_processor
+        del prompt_encoder
+    torch.cuda.empty_cache()
+
+    # ── Step 2: Encode audio with Audio VAE ──
+    ckpt_for_vae = args.audio_only_ckpt or args.checkpoint
+    logging.info(f"Loading audio VAE from {ckpt_for_vae}...")
+
+    ac = AudioConditioner(checkpoint_path=ckpt_for_vae, dtype=dtype, device=device)
+
+    logging.info("Encoding audio samples...")
+    for idx, sample in enumerate(samples):
+        gidx = sample["global_idx"]
+        audio_path = out / "audio_latents" / f"sample_{gidx:06d}.pt"
+        if args.skip_existing and audio_path.exists():
+            continue
+
+        try:
+            # Load audio
+            voice = decode_audio_from_file(sample["audio_path"], device, 0.0, args.max_duration)
+            if voice is None:
+                logging.warning(f"  Skipping {sample['id']}: no audio")
+                continue
+
+            w = voice.waveform
+            if w.dim() == 2:
+                if w.shape[0] == 1:
+                    w = w.repeat(2, 1)
+                w = w.unsqueeze(0)
+            elif w.dim() == 3 and w.shape[1] == 1:
+                w = w.repeat(1, 2, 1)
+            voice = Audio(waveform=w, sampling_rate=voice.sampling_rate)
+
+            # Encode through Audio VAE
+            audio_latent = ac(lambda enc: vae_encode_audio(voice, enc, None))
+
+            # Save audio latent
+            torch.save({
+                "latents": audio_latent.squeeze(0).cpu(),  # [C=8, T, F=16]
+                "sample_rate": 16000,
+            }, audio_path)
+
+        except Exception as e:
+            logging.warning(f"  Skipping {sample['id']}: {e}")
+            continue
+
+        if idx % 100 == 0:
+            logging.info(f"  Audio encoding: {idx}/{len(samples)}")
+
+    del ac
+    torch.cuda.empty_cache()
+
+    # ── Step 3: Create dummy video latents ──
+    logging.info("Creating dummy video latents...")
+    # Minimal video: 1 frame, 64x64 = 2x2 in latent space
+    dummy_video = {
+        "latents": torch.zeros(128, 1, 2, 2),
+        "num_frames": 1,
+        "height": 2,
+        "width": 2,
+        "fps": 24.0,
+    }
+    for idx, sample in enumerate(samples):
+        gidx = sample["global_idx"]
+        latent_path = out / "latents" / f"sample_{gidx:06d}.pt"
+        if args.skip_existing and latent_path.exists():
+            continue
+        torch.save(dummy_video, latent_path)
+
+    # ── Summary ──
+    n_audio = len(list((out / "audio_latents").glob("*.pt")))
+    n_cond = len(list((out / "conditions").glob("*.pt")))
+    n_lat = len(list((out / "latents").glob("*.pt")))
+    logging.info(f"\nDone! Output: {args.output_dir}")
+    logging.info(f"  audio_latents: {n_audio} files")
+    logging.info(f"  conditions:    {n_cond} files")
+    logging.info(f"  latents:       {n_lat} files")
+
+
+if __name__ == "__main__":
+    main()

engines/dramabox/vendor/src/train.py

@@ -0,0 +1,900 @@
+#!/usr/bin/env python3
+"""
+Audio-Only IC-LoRA Training for Voice Cloning on LTX-2.3.
+
+Uses the IC-LoRA pattern: reference audio tokens are APPENDED to the end of
+the target sequence using AudioConditionByReferenceLatent.  Loss is computed
+only on target tokens; reference tokens remain clean (denoise_mask=0).
+
+This follows the official video-to-video IC-LoRA strategy closely, but adapted
+for the audio-only modality path.
+
+Usage (single GPU):
+    CUDA_VISIBLE_DEVICES=0 python train_audio_iclora.py --data-dir ... --speaker-index ...
+
+Usage (multi-GPU with accelerate):
+    CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch --num_processes=4 train_audio_iclora.py ...
+"""
+
+import argparse
+import logging
+import math
+import os
+import random
+import shutil
+import sys
+import time
+from collections import defaultdict
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+
+REPO_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ltx2"))
+# ltx-pipelines already on path via ltx2/
+
+MODEL_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+# Import audio conditioning item from our module
+sys.path.insert(0, MODEL_DIR)
+from audio_conditioning import AudioConditionByReferenceLatent
+
+
+# ─── Timestep Sampling ───
+
+class DistilledTimestepSampler:
+    """Sample timesteps from the distilled sigma schedule.
+
+    The distilled model was trained to denoise at these specific sigma values.
+    We sample uniformly from the intervals between consecutive sigmas,
+    matching the distribution the model actually operates on.
+    """
+
+    # Distilled 8-step sigma values (boundaries of denoising intervals)
+    SIGMAS = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0]
+
+    def __init__(self, jitter: float = 0.02):
+        self.jitter = jitter
+
+    def sample(self, batch_size: int, seq_length: int = None, device: torch.device = None) -> torch.Tensor:
+        n_intervals = len(self.SIGMAS) - 1
+        interval_idx = torch.randint(0, n_intervals, (batch_size,), device=device)
+        t = torch.rand(batch_size, device=device)
+        sigma_high = torch.tensor([self.SIGMAS[i] for i in interval_idx], device=device)
+        sigma_low = torch.tensor([self.SIGMAS[i + 1] for i in interval_idx], device=device)
+        sigma = sigma_low + t * (sigma_high - sigma_low)
+        return sigma.clamp(0.01, 0.99)
+
+
+class ShiftedLogitNormalTimestepSampler:
+    """Shifted logit-normal distribution, shift depends on sequence length."""
+
+    def __init__(self, std: float = 1.0, eps: float = 1e-3, uniform_prob: float = 0.1):
+        self.std = std
+        self.eps = eps
+        self.uniform_prob = uniform_prob
+        self.normal_999_percentile = 3.0902 * std
+        self.normal_005_percentile = -2.5758 * std
+
+    def sample(self, batch_size: int, seq_length: int, device: torch.device = None) -> torch.Tensor:
+        mu = self._get_shift(seq_length)
+        normal = torch.randn(batch_size, device=device) * self.std + mu
+        logitnormal = torch.sigmoid(normal)
+
+        p999 = torch.sigmoid(torch.tensor(mu + self.normal_999_percentile, device=device))
+        p005 = torch.sigmoid(torch.tensor(mu + self.normal_005_percentile, device=device))
+        stretched = (logitnormal - p005) / (p999 - p005)
+        stretched = torch.where(stretched >= self.eps, stretched, 2 * self.eps - stretched)
+        stretched = stretched.clamp(0, 1)
+
+        uniform = (1 - self.eps) * torch.rand(batch_size, device=device) + self.eps
+        prob = torch.rand(batch_size, device=device)
+        return torch.where(prob > self.uniform_prob, stretched, uniform)
+
+    @staticmethod
+    def _get_shift(seq_length, min_tok=1024, max_tok=4096, min_s=0.95, max_s=2.05):
+        m = (max_s - min_s) / (max_tok - min_tok)
+        return m * seq_length + (min_s - m * min_tok)
+
+
+# ─── Dataset ───
+
+def build_speaker_map(index_paths, data_dirs):
+    """Map speaker → [(data_dir, sample_idx)] from index file(s).
+
+    The sample index comes from field 0 of the `~`-delimited row when it
+    parses as int (allows subset indexes that keep original sample numbers),
+    otherwise we fall back to the row's line number (legacy behaviour for
+    string-keyed indexes like tts_training_data_podcast).
+    """
+    speaker_to_samples = defaultdict(list)
+    for index_path, data_dir in zip(index_paths, data_dirs):
+        with open(index_path) as f:
+            for line_num, line in enumerate(f):
+                parts = line.strip().split("~")
+                if len(parts) < 7:
+                    continue
+                try:
+                    idx = int(parts[0])
+                except ValueError:
+                    idx = line_num
+                speaker_id = parts[1]
+                speaker_to_samples[speaker_id].append((data_dir, idx))
+    return {k: v for k, v in speaker_to_samples.items() if len(v) >= 2}
+
+
+class IDLoRADataset(Dataset):
+    # Silence-latent reference loaded once, used to detect and strip any
+    # leading silence frames baked into the preprocessed audio_latents. The
+    # training loop ALREADY prepends 0-25 random silence frames, so we don't
+    # want accidental silence in the source data compounding on top.
+    _silence_ref = None
+
+    @classmethod
+    def _load_silence_ref(cls):
+        if cls._silence_ref is None:
+            p = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+                             "assets", "silence_latent_frame.pt")
+            if os.path.exists(p):
+                cls._silence_ref = torch.load(p, weights_only=True).float().squeeze()  # [C, F]
+        return cls._silence_ref
+
+    def __init__(self, speaker_map):
+        self.samples = []
+        self.speaker_map = {}
+        for speaker, entries in speaker_map.items():
+            valid = []
+            for data_dir, idx in entries:
+                audio_path = Path(data_dir) / "audio_latents" / f"sample_{idx:06d}.pt"
+                cond_path = Path(data_dir) / "conditions" / f"sample_{idx:06d}.pt"
+                if audio_path.exists() and cond_path.exists():
+                    valid.append((data_dir, idx))
+            if len(valid) >= 2:
+                self.speaker_map[speaker] = valid
+        for speaker, entries in self.speaker_map.items():
+            for entry in entries:
+                self.samples.append((entry, speaker))
+        IDLoRADataset._load_silence_ref()
+
+    def __len__(self):
+        return len(self.samples)
+
+    def _load_sample(self, data_dir, idx):
+        base = Path(data_dir)
+        audio = torch.load(base / "audio_latents" / f"sample_{idx:06d}.pt", weights_only=False)
+        # Prefer prefix-stripped text embeddings if they exist (re-encoded with
+        # just the quoted dialogue, dropping the "A woman says, " / "A man
+        # speaks with X accent, " scene-description prefix).
+        stripped = base / "conditions_stripped" / f"sample_{idx:06d}.pt"
+        cond_path = stripped if stripped.exists() else base / "conditions" / f"sample_{idx:06d}.pt"
+        cond = torch.load(cond_path, weights_only=False)
+        if isinstance(audio, dict):
+            audio = audio.get("audio_latent", audio.get("latent", list(audio.values())[0]))
+        if audio.dim() == 2:
+            audio = audio.unsqueeze(0)
+        audio_feats = cond.get("audio_prompt_embeds", cond.get("prompt_embeds"))
+        attn_mask = cond.get("prompt_attention_mask")
+        # The audio_connector has num_learnable_registers=128 and asserts the
+        # input sequence length is divisible by 128. Our new preprocessing
+        # saved trimmed conditions (dropping left-padding to save disk), which
+        # produces short/irregular sequence lengths. Left-pad back to the next
+        # multiple of 128 with zeros (matching the tokenizer's left-padding
+        # convention) so this assertion holds.
+        REG = 128
+        L = audio_feats.shape[0]
+        target_L = ((L + REG - 1) // REG) * REG
+        if target_L != L:
+            pad_len = target_L - L
+            pad_emb = torch.zeros(pad_len, audio_feats.shape[1],
+                                  dtype=audio_feats.dtype)
+            pad_mask = torch.zeros(pad_len, dtype=attn_mask.dtype)
+            audio_feats = torch.cat([pad_emb, audio_feats], dim=0)
+            attn_mask = torch.cat([pad_mask, attn_mask], dim=0)
+        return audio, audio_feats, attn_mask
+
+    def __getitem__(self, idx):
+        (data_dir, tgt_idx), speaker = self.samples[idx]
+        tgt_latent, audio_feats, attn_mask = self._load_sample(data_dir, tgt_idx)
+
+        # Drop the reference entirely for non-voice-cloning categories:
+        #   - SFX samples (speaker starts with "sfx_"): descriptive sound events,
+        #     no speaker identity to clone.
+        #   - Song/music samples (suno dataset): prompts describe the music style,
+        #     reference audio doesn't transfer anything useful.
+        # Return a zero-length ref so the model trains target-only for these.
+        drop_ref = speaker.startswith("sfx_") or "preprocessed_ltx_suno" in str(data_dir)
+        if drop_ref:
+            C, F_dim = tgt_latent.shape[0], tgt_latent.shape[2]
+            ref_latent = torch.zeros(C, 0, F_dim, dtype=tgt_latent.dtype)
+        else:
+            entries = self.speaker_map[speaker]
+            ref_entry = random.choice([e for e in entries if e[1] != tgt_idx])
+            ref_latent, _, _ = self._load_sample(*ref_entry)
+
+        return {
+            "tgt_latent": tgt_latent,
+            "ref_latent": ref_latent,
+            "audio_features": audio_feats,
+            "attention_mask": attn_mask,
+        }
+
+
+# ─── Model building ───
+
+def build_audio_only_model(checkpoint_path, device, dtype):
+    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
+    from ltx_core.loader.registry import DummyRegistry
+    from ltx_core.loader.sd_ops import SDOps
+    from ltx_core.model.transformer.model import LTXModel, LTXModelType
+    from ltx_core.model.model_protocol import ModelConfigurator
+    from ltx_core.model.transformer.attention import AttentionFunction
+    from ltx_core.model.transformer.rope import LTXRopeType
+
+    sd_ops = SDOps("AO").with_matching(prefix="model.diffusion_model.").with_replacement("model.diffusion_model.", "")
+
+    class Cfg(ModelConfigurator[LTXModel]):
+        @classmethod
+        def from_config(cls, config):
+            t = config.get("transformer", {})
+            cp = None
+            if not t.get("caption_proj_before_connector", False):
+                from ltx_core.model.transformer.text_projection import create_caption_projection
+                with torch.device("meta"):
+                    cp = create_caption_projection(t, audio=True)
+            return LTXModel(
+                model_type=LTXModelType.AudioOnly,
+                audio_num_attention_heads=t.get("audio_num_attention_heads", 32),
+                audio_attention_head_dim=t.get("audio_attention_head_dim", 64),
+                audio_in_channels=t.get("audio_in_channels", 128),
+                audio_out_channels=t.get("audio_out_channels", 128),
+                num_layers=t.get("num_layers", 48),
+                audio_cross_attention_dim=t.get("audio_cross_attention_dim", 2048),
+                norm_eps=t.get("norm_eps", 1e-6),
+                attention_type=AttentionFunction(t.get("attention_type", "default")),
+                positional_embedding_theta=t.get("positional_embedding_theta", 10000.0),
+                audio_positional_embedding_max_pos=t.get("audio_positional_embedding_max_pos", [20]),
+                timestep_scale_multiplier=t.get("timestep_scale_multiplier", 1000),
+                use_middle_indices_grid=t.get("use_middle_indices_grid", True),
+                rope_type=LTXRopeType(t.get("rope_type", "interleaved")),
+                double_precision_rope=t.get("frequencies_precision", False) == "float64",
+                apply_gated_attention=t.get("apply_gated_attention", False),
+                audio_caption_projection=cp,
+                cross_attention_adaln=t.get("cross_attention_adaln", False),
+            )
+
+    builder = Builder(model_path=checkpoint_path, model_class_configurator=Cfg,
+                      model_sd_ops=sd_ops, registry=DummyRegistry())
+    return builder.build(device=device, dtype=dtype)
+
+
+def load_audio_connector(checkpoint_path, device, dtype):
+    # ltx-trainer already on path via ltx2/
+    from ltx_trainer.model_loader import load_embeddings_processor
+    emb_proc = load_embeddings_processor(checkpoint_path, device=device, dtype=dtype)
+    connector = emb_proc.audio_connector
+    del emb_proc
+    return connector
+
+
+def apply_lora(model, rank, alpha, dropout=0.0):
+    from peft import LoraConfig, get_peft_model
+    config = LoraConfig(
+        r=rank, lora_alpha=alpha, lora_dropout=dropout, bias="none",
+        target_modules=[
+            # Self-attention over audio tokens (voice-transfer pathway via ref).
+            "audio_attn1.to_k", "audio_attn1.to_q", "audio_attn1.to_v", "audio_attn1.to_out.0",
+            # Cross-attention (audio ↔ text context) NOT adapted — keep base
+            # model's prompt→audio behaviour intact and rely on dataset balance
+            # to drive expressiveness. (v15c tried this with adaLN unfreeze,
+            # that proved too destructive; v16 tries it adaLN-frozen.)
+            # FFN — non-linear capacity for style/phonetic adaptation.
+            "audio_ff.net.0.proj", "audio_ff.net.2",
+        ],
+    )
+    model = get_peft_model(model, config)
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total = sum(p.numel() for p in model.parameters())
+    logging.info(f"LoRA: {trainable:,} trainable / {total:,} total ({100*trainable/total:.1f}%)")
+    return model
+
+
+@torch.no_grad()
+def prepare_audio_context(audio_connector, audio_features, attention_mask, device, dtype):
+    from ltx_core.text_encoders.gemma.embeddings_processor import convert_to_additive_mask
+    audio_features = audio_features.to(device=device, dtype=dtype)
+    attention_mask = attention_mask.to(device=device)
+    if audio_features.shape[0] > 1:
+        results = []
+        for i in range(audio_features.shape[0]):
+            feat_i = audio_features[i:i+1]
+            mask_i = attention_mask[i:i+1]
+            additive = convert_to_additive_mask(mask_i, feat_i.dtype)
+            enc_i, _ = audio_connector(feat_i, additive)
+            results.append(enc_i)
+        return torch.cat(results, dim=0)
+    additive_mask = convert_to_additive_mask(attention_mask, audio_features.dtype)
+    audio_encoded, _ = audio_connector(audio_features, additive_mask)
+    return audio_encoded
+
+
+# ─── Validation ───
+
+def _unwrap_model_safe(model):
+    """Strip DDP / peft wrappers without going through accelerate.unwrap_model,
+    which imports deepspeed — broken in our env (torch API drift)."""
+    while hasattr(model, "module"):
+        model = model.module
+    return model
+
+
+def run_validation(lora_path, val_config_path, output_dir, step, lora_rank=128):
+    """Call validate.py in a subprocess. It loads TTSServer (the same stack
+    the warm server / Gradio app uses), attaches our LoRA, then iterates every
+    entry in val_config with the same inference settings the user tests with.
+    Single subprocess amortises the model-load cost across all val entries.
+
+    Forces validation onto VAL_GPU (default "0") because training already
+    occupies the rest. Override via TRAIN_VAL_GPU env var.
+    """
+    import subprocess
+    val_dir = os.path.join(output_dir, "validation", f"step_{step:05d}")
+    os.makedirs(val_dir, exist_ok=True)
+    script = os.path.join(os.path.dirname(__file__), "validate.py")
+    cmd = [
+        sys.executable, script,
+        "--val-config", val_config_path,
+        "--output-dir", val_dir,
+        "--lora", lora_path,
+        "--lora-rank", str(lora_rank),
+        # Use raw estimator output (no +10% buffer) so we can hear
+        # whether the model needs more/less duration at current quality.
+        "--duration-multiplier", "1.0",
+    ]
+    log_path = os.path.join(val_dir, "validate.log")
+    env = os.environ.copy()
+    # Validation needs its OWN GPU (training fills the others).
+    env["CUDA_VISIBLE_DEVICES"] = os.environ.get("TRAIN_VAL_GPU", "0")
+    try:
+        with open(log_path, "w") as logf:
+            result = subprocess.run(
+                cmd, stdout=logf, stderr=subprocess.STDOUT, timeout=1800, env=env,
+            )
+        if result.returncode == 0:
+            logging.info(f"  Validation step {step}: OK → {val_dir}")
+        else:
+            logging.warning(f"  Validation step {step} FAILED (see {log_path})")
+    except subprocess.TimeoutExpired:
+        logging.warning(f"  Validation step {step} TIMEOUT (>30min)")
+
+
+# ─── Args ───
+
+def collate_audio_batch(batch):
+    """Pad variable-length audio and track real lengths for loss masking."""
+    # TTS Audio Suite patch: keep this callable at module scope so Windows
+    # spawn-based DataLoader workers can pickle it.
+    max_tgt_T = max(item["tgt_latent"].shape[1] for item in batch)
+    max_ref_T = max(item["ref_latent"].shape[1] for item in batch)
+    channels = batch[0]["tgt_latent"].shape[0]
+    feature_dim = batch[0]["tgt_latent"].shape[2]
+
+    tgt_list, ref_list, feat_list, mask_list = [], [], [], []
+    tgt_lengths, ref_lengths = [], []
+    for item in batch:
+        tgt = item["tgt_latent"]
+        ref = item["ref_latent"]
+        tgt_lengths.append(tgt.shape[1])
+        ref_lengths.append(ref.shape[1])
+
+        if tgt.shape[1] < max_tgt_T:
+            pad = torch.zeros(
+                channels,
+                max_tgt_T - tgt.shape[1],
+                feature_dim,
+                dtype=tgt.dtype,
+            )
+            tgt = torch.cat([tgt, pad], dim=1)
+        tgt_list.append(tgt)
+
+        if ref.shape[1] < max_ref_T:
+            pad = torch.zeros(
+                channels,
+                max_ref_T - ref.shape[1],
+                feature_dim,
+                dtype=ref.dtype,
+            )
+            ref = torch.cat([ref, pad], dim=1)
+        ref_list.append(ref)
+        feat_list.append(item["audio_features"])
+        mask_list.append(item["attention_mask"])
+
+    return {
+        "tgt_latent": torch.stack(tgt_list),
+        "ref_latent": torch.stack(ref_list),
+        "audio_features": torch.stack(feat_list),
+        "attention_mask": torch.stack(mask_list),
+        "tgt_lengths": torch.tensor(tgt_lengths),
+        "ref_lengths": torch.tensor(ref_lengths),
+    }
+
+
+def parse_args():
+    # First pass: pull out --config so its values can become argparse defaults.
+    cfg_parser = argparse.ArgumentParser(add_help=False)
+    cfg_parser.add_argument("--config", default=None,
+                            help="YAML file with default values for any of the flags below. "
+                                 "Explicit CLI flags still override the YAML.")
+    cfg_args, remaining = cfg_parser.parse_known_args()
+    yaml_defaults: dict = {}
+    if cfg_args.config:
+        import yaml as _yaml
+        with open(cfg_args.config) as f:
+            yaml_defaults = _yaml.safe_load(f) or {}
+        # YAML keys are dashes-or-underscores → normalize to argparse dest (underscore).
+        yaml_defaults = {k.replace("-", "_"): v for k, v in yaml_defaults.items()}
+
+    def _yaml(name, fallback):
+        return yaml_defaults.get(name, fallback)
+
+    p = argparse.ArgumentParser(
+        parents=[cfg_parser],
+        description="Audio-Only IC-LoRA Training for Voice Cloning",
+    )
+    p.add_argument("--data-dir", required="data_dir" not in yaml_defaults,
+                   nargs="+", default=_yaml("data_dir", None))
+    p.add_argument("--speaker-index", required="speaker_index" not in yaml_defaults,
+                   nargs="+", default=_yaml("speaker_index", None))
+    p.add_argument("--output-dir", default=_yaml("output_dir", os.path.join(MODEL_DIR, "tts_iclora_v1")))
+    p.add_argument("--checkpoint", default=_yaml("checkpoint", os.path.join(MODEL_DIR, "dramabox-dit-v1.safetensors")))
+    p.add_argument("--full-checkpoint", default=_yaml("full_checkpoint", os.path.join(MODEL_DIR, "dramabox-audio-components.safetensors")))
+    p.add_argument("--base-model", choices=["distilled", "dev"], default=_yaml("base_model", "dev"),
+                   help="Base model type: distilled uses DistilledTimestepSampler, dev uses ShiftedLogitNormal")
+    p.add_argument("--lora-rank", type=int, default=_yaml("lora_rank", 128))
+    p.add_argument("--lora-alpha", type=int, default=_yaml("lora_alpha", 128))
+    p.add_argument("--lora-dropout", type=float, default=_yaml("lora_dropout", 0.0),
+                   help="Dropout applied to LoRA A/B matrices during training. "
+                        "Recommended ~0.1 for small datasets to regularize.")
+    p.add_argument("--resume-lora", default=_yaml("resume_lora", None))
+    p.add_argument("--resume-step-offset", type=int, default=_yaml("resume_step_offset", None),
+                   help="Step to add when naming saved checkpoints. If None, inferred "
+                        "from --resume-lora filename (e.g. lora_step_10000.safetensors → 10000). "
+                        "Set to 0 to start numbering at 0 regardless.")
+    p.add_argument("--ref-ratio", type=float, default=_yaml("ref_ratio", 0.3),
+                   help="Fraction of target length to use as reference (default 0.3)")
+    p.add_argument("--max-ref-tokens", type=int, default=_yaml("max_ref_tokens", 200),
+                   help="Maximum reference tokens after patchification (default 200)")
+    p.add_argument("--text-dropout", type=float, default=_yaml("text_dropout", 0.0),
+                   help="Probability of dropping text conditioning (forces reliance on voice ref)")
+    p.add_argument("--steps", type=int, default=_yaml("steps", 30000))
+    p.add_argument("--lr", type=float, default=_yaml("lr", 3e-5))
+    p.add_argument("--lr-scheduler", choices=["cosine", "linear", "constant"], default=_yaml("lr_scheduler", "cosine"))
+    p.add_argument("--batch-size", type=int, default=_yaml("batch_size", 1))
+    p.add_argument("--grad-accum", type=int, default=_yaml("grad_accum", 4))
+    p.add_argument("--max-grad-norm", type=float, default=_yaml("max_grad_norm", 1.0))
+    p.add_argument("--save-every", type=int, default=_yaml("save_every", 1000))
+    p.add_argument("--log-every", type=int, default=_yaml("log_every", 50))
+    p.add_argument("--seed", type=int, default=_yaml("seed", 42))
+    p.add_argument("--warmup-steps", type=int, default=_yaml("warmup_steps", 100))
+    p.add_argument("--val-config", default=_yaml("val_config", None))
+    return p.parse_args(remaining)
+
+
+# ─── Main ───
+
+def main():
+    from accelerate import Accelerator
+    from accelerate.utils import set_seed
+
+    args = parse_args()
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.grad_accum,
+        mixed_precision="bf16",
+    )
+
+    is_main = accelerator.is_main_process
+    if is_main:
+        logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+    else:
+        logging.basicConfig(level=logging.WARNING)
+
+    set_seed(args.seed)
+    device = accelerator.device
+    dtype = torch.bfloat16
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Save training args
+    if is_main:
+        import yaml
+        args_dict = vars(args).copy()
+        args_dict["_meta"] = {
+            "world_size": accelerator.num_processes,
+            "dtype": str(dtype),
+            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+            "script": "train_audio_iclora.py",
+            "pattern": "IC-LoRA (ref appended to end)",
+        }
+        with open(os.path.join(args.output_dir, "training_args.yaml"), "w") as f:
+            yaml.dump(args_dict, f, default_flow_style=False, sort_keys=False)
+
+    from ltx_core.components.patchifiers import AudioPatchifier
+    from ltx_core.model.transformer.modality import Modality
+    from ltx_core.guidance.perturbations import BatchedPerturbationConfig
+    from ltx_core.tools import AudioLatentTools
+    from ltx_core.types import AudioLatentShape, LatentState
+    from ltx_pipelines.utils.helpers import modality_from_latent_state, timesteps_from_mask
+
+    # Build speaker map
+    if is_main:
+        logging.info("Building speaker map...")
+    speaker_map = build_speaker_map(args.speaker_index, args.data_dir)
+    if is_main:
+        logging.info(f"Speaker map: {len(speaker_map)} speakers, "
+                     f"{sum(len(v) for v in speaker_map.values())} samples")
+
+    # Load model
+    if is_main:
+        logging.info("Loading audio-only model...")
+    model = build_audio_only_model(args.checkpoint, device, dtype)
+
+    if is_main:
+        logging.info("Loading audio connector...")
+    audio_connector = load_audio_connector(args.full_checkpoint, device, dtype)
+    audio_connector.eval()
+    for p in audio_connector.parameters():
+        p.requires_grad = False
+
+    if is_main:
+        logging.info(f"Applying LoRA (rank={args.lora_rank}, alpha={args.lora_alpha})...")
+    model = apply_lora(model, args.lora_rank, args.lora_alpha, args.lora_dropout)
+
+    # Resume from checkpoint
+    if args.resume_lora:
+        from safetensors.torch import load_file as st_load
+        if is_main:
+            logging.info(f"Resuming from: {args.resume_lora}")
+        lora_sd = st_load(args.resume_lora)
+        mapped = {}
+        for k, v in lora_sd.items():
+            nk = k.replace(".lora_A.weight", ".lora_A.default.weight").replace(
+                ".lora_B.weight", ".lora_B.default.weight")
+            mapped[nk] = v
+        model.load_state_dict(mapped, strict=False)
+
+    # Determine step offset for save filenames. Without this, resuming a run
+    # restarts step numbering at 0 and would overwrite earlier phase-1
+    # checkpoints with the same save_every cadence.
+    if args.resume_step_offset is None:
+        resume_offset = 0
+        if args.resume_lora:
+            import re as _re
+            m = _re.search(r"lora_step_(\d+)", os.path.basename(args.resume_lora))
+            if m:
+                resume_offset = int(m.group(1))
+        args.resume_step_offset = resume_offset
+    if is_main and args.resume_step_offset:
+        logging.info(f"Save-step offset: +{args.resume_step_offset}")
+
+    model.train()
+    model.base_model.model.set_gradient_checkpointing(True)
+
+    # Dataset & DataLoader
+    dataset = IDLoRADataset(speaker_map)
+    if is_main:
+        logging.info(f"Dataset: {len(dataset)} samples, {len(dataset.speaker_map)} speakers")
+
+    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=2,
+                            pin_memory=True, drop_last=True, collate_fn=collate_audio_batch)
+
+    # Optimizer & Scheduler
+    optimizer = torch.optim.AdamW(
+        [p for p in model.parameters() if p.requires_grad],
+        lr=args.lr, betas=(0.9, 0.999), weight_decay=0.01,
+    )
+
+    from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR, ConstantLR
+    warmup = LinearLR(optimizer, start_factor=0.01, end_factor=1.0, total_iters=args.warmup_steps)
+    remaining = args.steps - args.warmup_steps
+    if args.lr_scheduler == "cosine":
+        # Warmup -> constant hold (20% of remaining) -> cosine decay
+        hold_steps = max(remaining // 5, 0)
+        decay_steps = max(remaining - hold_steps, 1)
+        hold_sched = ConstantLR(optimizer, factor=1.0, total_iters=hold_steps)
+        decay_sched = CosineAnnealingLR(optimizer, T_max=decay_steps, eta_min=1e-6)
+        scheduler = SequentialLR(
+            optimizer,
+            [warmup, hold_sched, decay_sched],
+            milestones=[args.warmup_steps, args.warmup_steps + hold_steps],
+        )
+    elif args.lr_scheduler == "linear":
+        main_sched = LinearLR(optimizer, start_factor=1.0, end_factor=0.01, total_iters=max(remaining, 1))
+        scheduler = SequentialLR(optimizer, [warmup, main_sched], milestones=[args.warmup_steps])
+    else:
+        main_sched = ConstantLR(optimizer, factor=1.0, total_iters=max(remaining, 1))
+        scheduler = SequentialLR(optimizer, [warmup, main_sched], milestones=[args.warmup_steps])
+
+    # Prepare with Accelerate — but NOT the scheduler. AcceleratedScheduler
+    # calls the underlying scheduler.step() `num_processes` times per sync,
+    # which silently scales down our warmup/cosine spans by that factor.
+    # We call scheduler.step() ourselves, gated on sync_gradients → exactly
+    # one advance per optimizer step, as the yaml spec intends.
+    model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
+
+    patchifier = AudioPatchifier(patch_size=1)
+
+    # Select timestep sampler based on base model type
+    if args.base_model == "distilled":
+        timestep_sampler = DistilledTimestepSampler()
+        if is_main:
+            logging.info("Using DistilledTimestepSampler (matching distilled model sigmas)")
+    else:
+        timestep_sampler = ShiftedLogitNormalTimestepSampler()
+        if is_main:
+            logging.info("Using ShiftedLogitNormalTimestepSampler (dev model)")
+
+    # Training loop
+    if is_main:
+        logging.info(f"Training: {args.steps} steps, lr={args.lr}, scheduler={args.lr_scheduler}, "
+                     f"batch={args.batch_size}, grad_accum={args.grad_accum}, "
+                     f"world_size={accelerator.num_processes}, "
+                     f"ref_ratio={args.ref_ratio}, max_ref_tokens={args.max_ref_tokens}")
+        logging.info("IC-LoRA pattern: ref tokens APPENDED to target, loss on target only")
+
+    data_iter = iter(dataloader)
+    step = 0
+    accum_loss = 0.0
+    best_loss = float("inf")
+    best_step = 0
+    t0 = time.time()
+
+    total_micro_steps = args.steps * args.grad_accum
+
+    for micro_step in range(total_micro_steps):
+        try:
+            batch = next(data_iter)
+        except StopIteration:
+            data_iter = iter(dataloader)
+            batch = next(data_iter)
+
+        is_opt_step = (micro_step + 1) % args.grad_accum == 0
+        if is_opt_step:
+            step += 1
+            if is_main:
+                # TTS Audio Suite patch: provide lightweight per-step telemetry
+                # to the parent process even when human logs use log_every=50.
+                print(
+                    f"TTS_SUITE_PROGRESS step={step} total={args.steps}",
+                    flush=True,
+                )
+
+        with accelerator.accumulate(model):
+            tgt_latent = batch["tgt_latent"].to(dtype=dtype)  # [B, C, max_tgt_T, F]
+            ref_latent = batch["ref_latent"].to(dtype=dtype)  # [B, C, max_ref_T, F]
+            tgt_lengths = batch["tgt_lengths"].to(device=device)  # [B]
+            B = tgt_latent.shape[0]
+
+            # ── Random silence padding (0-1s) ── ltx_audio_tts baseline.
+            # User observed reference-audio leak at end of generations when this
+            # was reduced to 5 (v14) or 10 frames (v16/v17) — the model seemed
+            # to use the extra target budget to regurgitate ref content. Full
+            # 25 frames (0-1s avg 500ms) was apparently load-bearing for
+            # regularising the boundary and reducing hallucinations.
+            # Uses the real silence latent (not zeros) so the VAE decodes it as
+            # true silence instead of static noise.
+            max_pad_frames = 25  # ~1s at 25 latent frames/sec
+            pad_frames = random.randint(0, max_pad_frames)
+            if pad_frames > 0:
+                C, F_dim = tgt_latent.shape[1], tgt_latent.shape[3]
+                if not hasattr(args, '_silence_frame') or args._silence_frame is None:
+                    _sf_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "assets", "silence_latent_frame.pt")
+                    if os.path.exists(_sf_path):
+                        args._silence_frame = torch.load(_sf_path, weights_only=True)  # [C, 1, F]
+                        if is_main:
+                            logging.info(f"Loaded silence latent from {_sf_path}")
+                    else:
+                        args._silence_frame = False  # fallback to zeros
+                        if is_main:
+                            logging.warning(f"silence_latent_frame.pt not found, using zeros")
+                if args._silence_frame is not False:
+                    sf = args._silence_frame.to(dtype=dtype, device=device)  # [C, 1, F]
+                    silence_pad = sf.unsqueeze(0).expand(B, -1, pad_frames, -1)  # [B, C, pad, F]
+                else:
+                    silence_pad = torch.zeros(B, C, pad_frames, F_dim, dtype=dtype, device=device)
+                tgt_latent = torch.cat([silence_pad, tgt_latent], dim=2)
+
+            # Cap reference to max_ref_tokens (in latent frames, before patchification)
+            # After patchification, ref_T tokens = ref frames (patch_size=1)
+            ref_T_frames = min(ref_latent.shape[2], args.max_ref_tokens)
+            ref_latent = ref_latent[:, :, :ref_T_frames, :]
+
+            tgt_T_frames = tgt_latent.shape[2]  # max (padded) target frames
+
+            # ── Step 1: Create target AudioLatentShape and AudioLatentTools ──
+            tgt_shape = AudioLatentShape(
+                batch=B,
+                channels=tgt_latent.shape[1],  # 8
+                frames=tgt_T_frames,
+                mel_bins=tgt_latent.shape[3],   # 16
+            )
+
+            audio_tools = AudioLatentTools(
+                patchifier=patchifier,
+                target_shape=tgt_shape,
+            )
+
+            # ── Step 2: Create initial state from target latent ──
+            # create_initial_state patchifies: [B, C, T, F] -> [B, T, C*F]
+            # Also creates denoise_mask=1 (all target tokens will be denoised)
+            # and computes temporal positions
+            state = audio_tools.create_initial_state(
+                device=device,
+                dtype=dtype,
+                initial_latent=tgt_latent,
+            )
+            # state.latent: [B, tgt_T, 128], state.denoise_mask: [B, tgt_T, 1]
+            # state.positions: [B, 1, tgt_T, 2]
+
+            tgt_T = audio_tools.target_shape.token_count()  # = tgt_T_frames
+
+            # ── Step 3: Apply flow-matching noise to target BEFORE appending ref ──
+            # Sample sigma
+            total_tokens = tgt_T + ref_T_frames
+            sigma = timestep_sampler.sample(B, total_tokens, device=device)
+            sigma_exp = sigma.view(-1, 1, 1)  # [B, 1, 1]
+
+            noise = torch.randn_like(state.latent)  # [B, tgt_T, 128]
+            noisy_tgt = (1 - sigma_exp) * state.latent + sigma_exp * noise
+
+            # Replace the latent in state with the noisy version
+            # (clean_latent stays clean for post_process_latent pattern)
+            state = LatentState(
+                latent=noisy_tgt,
+                denoise_mask=state.denoise_mask,
+                positions=state.positions,
+                clean_latent=state.clean_latent,
+                attention_mask=state.attention_mask,
+            )
+
+            # ── Step 4: Append reference tokens using AudioConditionByReferenceLatent ──
+            # This appends ref tokens to the END with denoise_mask=0 (frozen/clean)
+            # Skip entirely when ref_T=0 (SFX / song samples): the model trains
+            # target-only for those categories since there's no voice to clone.
+            if ref_T_frames > 0:
+                ref_conditioning = AudioConditionByReferenceLatent(
+                    latent=ref_latent,
+                    strength=1.0,  # 1.0 = ref fully clean (denoise_mask=0)
+                )
+                state = ref_conditioning.apply_to(
+                    latent_state=state,
+                    latent_tools=audio_tools,
+                )
+            # state.latent: [B, tgt_T + ref_T, 128]
+            # state.denoise_mask: [B, tgt_T + ref_T, 1]
+            #   target tokens: 1.0 (denoise), ref tokens: 0.0 (frozen)
+            # state.positions: [B, 1, tgt_T + ref_T, 2]
+
+            # ── Step 5: Build loss mask for target tokens (excluding padding) ──
+            # loss_mask: 1 for real target tokens, 0 for padding and ref tokens
+            loss_mask = torch.zeros(B, tgt_T, device=device)
+            for b_idx in range(B):
+                real_len = min(tgt_lengths[b_idx].item(), tgt_T)
+                loss_mask[b_idx, :real_len] = 1.0
+
+            # ── Step 6: Prepare text context ──
+            # Text conditioning dropout: randomly zero out text context to force
+            # the model to rely on the voice reference for identity/style.
+            with torch.no_grad():
+                audio_context = prepare_audio_context(
+                    audio_connector, batch["audio_features"],
+                    batch["attention_mask"], device, dtype)
+                if args.text_dropout > 0 and random.random() < args.text_dropout:
+                    audio_context = torch.zeros_like(audio_context)
+
+            # ── Step 7: Build Modality using modality_from_latent_state ──
+            # timesteps = sigma * denoise_mask (ref gets 0, target gets sigma)
+            audio_mod = modality_from_latent_state(
+                state=state,
+                context=audio_context,
+                sigma=sigma,
+                enabled=True,
+            )
+
+            # ── Step 8: Forward pass ──
+            perturbations = BatchedPerturbationConfig.empty(B)
+            with torch.autocast(device_type="cuda", dtype=dtype):
+                _, velocity_pred = model(video=None, audio=audio_mod, perturbations=perturbations)
+
+            # ── Step 9: Compute loss (IC-LoRA pattern) ──
+            # Target is at the FRONT (indices 0..tgt_T), ref at the END
+            # velocity target = noise - clean
+            tgt_patchified = audio_tools.patchifier.patchify(tgt_latent)  # [B, tgt_T, 128]
+            target_velocity = noise - tgt_patchified
+
+            # Extract target portion of prediction
+            pred_tgt = velocity_pred[:, :tgt_T]  # [B, tgt_T, 128]
+
+            # MSE loss with mask: only on real target tokens (not padding or ref)
+            per_token_mse = (pred_tgt - target_velocity).pow(2).mean(dim=-1)  # [B, tgt_T]
+            loss = per_token_mse.mul(loss_mask).div(loss_mask.mean().clamp(min=1e-6)).mean()
+
+            accelerator.backward(loss)
+
+            if accelerator.sync_gradients and args.max_grad_norm > 0:
+                accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+
+            optimizer.step()
+            optimizer.zero_grad()
+            # Only advance the LR scheduler once per OPTIMIZER step (not per
+            # micro-step). Mirrors AcceleratedOptimizer.step() which is
+            # internally gated on sync_gradients.
+            if accelerator.sync_gradients:
+                scheduler.step()
+
+        accum_loss += loss.item()
+
+        # Logging & saving on optimization steps only
+        if is_opt_step and step % args.log_every == 0 and is_main:
+            avg_loss = accum_loss / (args.log_every * args.grad_accum)
+            lr = optimizer.param_groups[0]["lr"]
+            elapsed = time.time() - t0
+            sps = step / elapsed if elapsed > 0 else 0
+            eta = (args.steps - step) / sps if sps > 0 else 0
+            logging.info(
+                f"Step {step}/{args.steps} | loss={avg_loss:.4f} | lr={lr:.2e} | "
+                f"tgt_T={tgt_T} ref_T={ref_T_frames} total={tgt_T + ref_T_frames} | "
+                f"{sps:.1f} steps/s | ETA {eta/60:.0f}min"
+            )
+
+            # Save best whenever loss improves — no warmup gate, so we can
+            # observe best checkpoints during warmup too.
+            if avg_loss < best_loss:
+                best_loss = avg_loss
+                old_best = os.path.join(args.output_dir, f"best_step_{best_step:05d}.safetensors")
+                best_step = step + args.resume_step_offset
+                new_best = os.path.join(args.output_dir, f"best_step_{best_step:05d}.safetensors")
+                unwrapped = _unwrap_model_safe(model)
+                unwrapped.save_pretrained(args.output_dir)
+                adapter = os.path.join(args.output_dir, "adapter_model.safetensors")
+                if os.path.exists(adapter):
+                    shutil.copy(adapter, new_best)
+                if old_best != new_best and os.path.exists(old_best):
+                    os.remove(old_best)
+                logging.info(f"New best: loss={best_loss:.4f} at step {best_step}")
+
+            accum_loss = 0.0
+
+        if is_opt_step and step % args.save_every == 0 and is_main:
+            global_step = step + args.resume_step_offset
+            save_path = os.path.join(args.output_dir, f"lora_step_{global_step:05d}.safetensors")
+            logging.info(f"Saving: {save_path}")
+            unwrapped = _unwrap_model_safe(model)
+            unwrapped.save_pretrained(args.output_dir)
+            adapter = os.path.join(args.output_dir, "adapter_model.safetensors")
+            if os.path.exists(adapter):
+                shutil.copy(adapter, save_path)
+
+            if args.val_config:
+                logging.info(f"Running validation at step {global_step}...")
+                model.eval()
+                run_validation(save_path, args.val_config, args.output_dir, global_step,
+                               lora_rank=args.lora_rank)
+                model.train()
+
+    # Final save
+    if is_main:
+        unwrapped = _unwrap_model_safe(model)
+        unwrapped.save_pretrained(args.output_dir)
+        adapter = os.path.join(args.output_dir, "adapter_model.safetensors")
+        global_step = step + args.resume_step_offset
+        save_path = os.path.join(args.output_dir, f"lora_step_{global_step:05d}.safetensors")
+        if os.path.exists(adapter):
+            shutil.copy(adapter, save_path)
+        logging.info(f"Training complete! {step} steps in {time.time()-t0:.0f}s")
+        logging.info(f"Best loss: {best_loss:.4f} at step {best_step}")
+
+
+if __name__ == "__main__":
+    main()

engines/dramabox/vendor/src/validate.py

@@ -0,0 +1,370 @@
+#!/usr/bin/env python3
+"""Warm validation runner — loads base dev + LoRA + all aux models ONCE,
+then iterates every speaker in val_config generating each output.
+
+Matches the same generation path as inference.py but keeps Gemma / audio VAE
+/ velocity model / audio decoder resident across entries. Inference
+settings default to the Gradio warm-server values (cfg=2.5, stg=1.5,
+modality=1.0, rescale=0, 30 steps, fps=25) — use --inference-params to
+override.
+"""
+import argparse
+import logging
+import os
+import sys
+import time
+import traceback
+
+import torch
+import torchaudio
+
+REPO_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+MODEL_DIR = REPO_DIR
+sys.path.insert(0, os.path.join(REPO_DIR, "ltx2"))
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+DEV_FULL_CKPT = os.environ.get(
+    "LTX_FULL_CHECKPOINT",
+    os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ltx-2.3-22b-dev.safetensors"),
+)
+# TTS Audio Suite patch: organized DramaBox installs keep the transformer and
+# audio components in separate checkpoints, unlike the upstream full checkpoint.
+DRAMABOX_TRANSFORMER_CKPT = os.environ.get("LTX_CHECKPOINT", DEV_FULL_CKPT)
+GEMMA_ROOT = os.environ.get(
+    "GEMMA_ROOT",
+    os.path.expanduser("~/.cache/dramabox/gemma-3-12b-it-bnb-4bit"),
+)
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--val-config", required=True)
+    p.add_argument("--output-dir", required=True)
+    p.add_argument("--lora", default=None)
+    p.add_argument("--lora-rank", type=int, default=128)
+    p.add_argument("--checkpoint", default=DRAMABOX_TRANSFORMER_CKPT)
+    p.add_argument("--full-checkpoint", default=DEV_FULL_CKPT)
+    p.add_argument("--gemma-root", default=GEMMA_ROOT)
+    p.add_argument("--cfg-scale", type=float, default=2.5)
+    p.add_argument("--stg-scale", type=float, default=1.5)
+    p.add_argument("--rescale-scale", type=float, default=0.0)
+    p.add_argument("--modality-scale", type=float, default=1.0)
+    p.add_argument("--steps", type=int, default=30)
+    p.add_argument("--fps", type=float, default=25.0)
+    p.add_argument("--stg-block", type=int, default=29)
+    p.add_argument("--cfg-clamp", type=float, default=0.0)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--duration-multiplier", type=float, default=1.1)
+    # Match Gradio / inference_server.py DEFAULT_NEG exactly
+    p.add_argument("--negative-prompt", default=(
+        "worst quality, inconsistent, robotic, distorted, noise, static, "
+        "muffled, unclear, unnatural, monotone"
+    ))
+    return p.parse_args()
+
+
+def estimate_speech_duration(prompt: str, speed: float = 1.0) -> float:
+    import re
+    quoted = re.findall(r'"([^"]*)"', prompt) or re.findall(r"'([^']*)'", prompt)
+    text = " ".join(quoted) if quoted else prompt
+    duration = len(text) * 0.065 / max(speed, 0.1) + 1.5
+    return max(3.0, round(duration, 1))
+
+
+class WarmValidator:
+    def __init__(self, checkpoint, full_checkpoint, gemma_root, lora_path=None, lora_rank=128,
+                 device="cuda", dtype=torch.bfloat16):
+        from audio_conditioning import AudioConditionByReferenceLatent  # noqa: F401 (imported by inference.py)
+        from ltx_core.components.patchifiers import AudioPatchifier
+        from ltx_pipelines.utils.blocks import PromptEncoder, AudioConditioner, AudioDecoder
+
+        self.device = torch.device(device)
+        self.dtype = dtype
+        self.full_checkpoint = full_checkpoint
+        self.gemma_root = gemma_root
+        self.patchifier = AudioPatchifier(patch_size=1)
+
+        logging.info("Loading PromptEncoder (Gemma + embeddings_processor)...")
+        t0 = time.time()
+        self.prompt_encoder = PromptEncoder(
+            checkpoint_path=full_checkpoint, gemma_root=gemma_root,
+            dtype=dtype, device=self.device, warm=True, audio_only=True,
+        )
+        logging.info(f"  PromptEncoder ready in {time.time()-t0:.1f}s")
+
+        logging.info("Loading AudioConditioner (audio VAE encoder)...")
+        t0 = time.time()
+        self.audio_conditioner = AudioConditioner(
+            checkpoint_path=full_checkpoint, dtype=dtype, device=self.device, warm=True,
+        )
+        logging.info(f"  AudioConditioner ready in {time.time()-t0:.1f}s")
+
+        logging.info("Loading AudioDecoder...")
+        t0 = time.time()
+        self.audio_decoder = AudioDecoder(
+            checkpoint_path=full_checkpoint, dtype=dtype, device=self.device, warm=True,
+        )
+        logging.info(f"  AudioDecoder ready in {time.time()-t0:.1f}s")
+
+        logging.info("Building velocity model (audio-only from base dev)...")
+        t0 = time.time()
+        # TTS Audio Suite patch: build the DiT from the dedicated transformer
+        # checkpoint while keeping the audio connector/VAE/decoder checkpoint separate.
+        self.velocity_model = self._build_velocity_model(checkpoint, lora_path, lora_rank)
+        logging.info(f"  Velocity model ready in {time.time()-t0:.1f}s "
+                     f"({sum(p.numel() for p in self.velocity_model.parameters()) / 1e9:.1f}B params)")
+
+    def _build_velocity_model(self, checkpoint_path, lora_path, lora_rank):
+        from ltx_core.loader.registry import DummyRegistry
+        from ltx_core.loader.sd_ops import SDOps
+        from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
+        from ltx_core.model.model_protocol import ModelConfigurator
+        from ltx_core.model.transformer.attention import AttentionFunction
+        from ltx_core.model.transformer.model import LTXModel, LTXModelType
+        from ltx_core.model.transformer.rope import LTXRopeType
+
+        sd_ops = (
+            SDOps("AO")
+            .with_matching(prefix="model.diffusion_model.")
+            .with_replacement("model.diffusion_model.", "")
+        )
+
+        class Cfg(ModelConfigurator[LTXModel]):
+            @classmethod
+            def from_config(cls, config):
+                t = config.get("transformer", {})
+                cp = None
+                if not t.get("caption_proj_before_connector", False):
+                    from ltx_core.model.transformer.text_projection import create_caption_projection
+                    with torch.device("meta"):
+                        cp = create_caption_projection(t, audio=True)
+                return LTXModel(
+                    model_type=LTXModelType.AudioOnly,
+                    audio_num_attention_heads=t.get("audio_num_attention_heads", 32),
+                    audio_attention_head_dim=t.get("audio_attention_head_dim", 64),
+                    audio_in_channels=t.get("audio_in_channels", 128),
+                    audio_out_channels=t.get("audio_out_channels", 128),
+                    num_layers=t.get("num_layers", 48),
+                    audio_cross_attention_dim=t.get("audio_cross_attention_dim", 2048),
+                    norm_eps=t.get("norm_eps", 1e-6),
+                    attention_type=AttentionFunction(t.get("attention_type", "default")),
+                    positional_embedding_theta=10000.0,
+                    audio_positional_embedding_max_pos=[20.0],
+                    timestep_scale_multiplier=t.get("timestep_scale_multiplier", 1000),
+                    use_middle_indices_grid=t.get("use_middle_indices_grid", True),
+                    rope_type=LTXRopeType(t.get("rope_type", "interleaved")),
+                    double_precision_rope=t.get("frequencies_precision", False) == "float64",
+                    apply_gated_attention=t.get("apply_gated_attention", False),
+                    audio_caption_projection=cp,
+                    cross_attention_adaln=t.get("cross_attention_adaln", False),
+                )
+
+        builder = Builder(
+            model_path=checkpoint_path, model_class_configurator=Cfg,
+            model_sd_ops=sd_ops, registry=DummyRegistry(),
+        )
+        velocity = builder.build(device=self.device, dtype=self.dtype).to(self.device).eval()
+
+        if lora_path and os.path.exists(lora_path):
+            from peft import LoraConfig, get_peft_model
+            from safetensors.torch import load_file as st_load
+            logging.info(f"Attaching LoRA: {lora_path}")
+            lora_sd = st_load(lora_path)
+            is_peft = any("base_model.model." in k for k in lora_sd.keys())
+            is_iclora = any("diffusion_model." in k for k in lora_sd.keys())
+            cfg = LoraConfig(
+                r=lora_rank, lora_alpha=lora_rank, lora_dropout=0.0, bias="none",
+                target_modules=[
+                    "audio_attn1.to_k", "audio_attn1.to_q",
+                    "audio_attn1.to_v", "audio_attn1.to_out.0",
+                    "audio_attn2.to_k", "audio_attn2.to_q",
+                    "audio_attn2.to_v", "audio_attn2.to_out.0",
+                    "audio_ff.net.0.proj", "audio_ff.net.2",
+                ],
+            )
+            velocity = get_peft_model(velocity, cfg)
+
+            if is_peft:
+                mapped = {}
+                for k, v in lora_sd.items():
+                    nk = k
+                    if ".lora_A.weight" in k and ".lora_A.default.weight" not in k:
+                        nk = k.replace(".lora_A.weight", ".lora_A.default.weight")
+                    if ".lora_B.weight" in k and ".lora_B.default.weight" not in k:
+                        nk = k.replace(".lora_B.weight", ".lora_B.default.weight")
+                    mapped[nk] = v
+                _, unexpected = velocity.load_state_dict(mapped, strict=False)
+                logging.info(f"  Loaded {len(mapped) - len(unexpected)} LoRA weights (peft)")
+            elif is_iclora:
+                audio_keys = {k: v for k, v in lora_sd.items()
+                              if "audio_attn1" in k or "audio_attn2" in k or "audio_ff" in k}
+                mapped = {}
+                for k, v in audio_keys.items():
+                    nk = k.replace("diffusion_model.", "base_model.model.")
+                    nk = nk.replace(".lora_A.weight", ".lora_A.default.weight")
+                    nk = nk.replace(".lora_B.weight", ".lora_B.default.weight")
+                    mapped[nk] = v
+                _, unexpected = velocity.load_state_dict(mapped, strict=False)
+                logging.info(f"  Loaded {len(mapped) - len(unexpected)} LoRA weights (iclora)")
+
+            velocity = velocity.merge_and_unload()
+            logging.info("  Merged LoRA into base weights")
+
+        return velocity
+
+    @torch.inference_mode()
+    def generate(self, prompt, output_path, voice_ref=None, args=None):
+        from audio_conditioning import AudioConditionByReferenceLatent
+        from ltx_core.batch_split import BatchSplitAdapter
+        from ltx_core.components.diffusion_steps import EulerDiffusionStep
+        from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
+        from ltx_core.components.noisers import GaussianNoiser
+        from ltx_core.components.schedulers import LTX2Scheduler
+        from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
+        from ltx_core.model.transformer.model import X0Model
+        from ltx_core.tools import AudioLatentTools
+        from ltx_core.types import Audio, AudioLatentShape, VideoPixelShape
+        from ltx_pipelines.utils.denoisers import GuidedDenoiser, SimpleDenoiser
+        from ltx_pipelines.utils.gpu_model import gpu_model
+        from ltx_pipelines.utils.media_io import decode_audio_from_file
+        from ltx_pipelines.utils.samplers import euler_denoising_loop
+
+        t_total = time.time()
+
+        # ---- Duration + shape ----
+        gen_dur = estimate_speech_duration(prompt) * args.duration_multiplier
+        raw_frames = int(round(gen_dur * args.fps)) + 1
+        num_frames = ((raw_frames - 1 + 4) // 8) * 8 + 1
+        pixel_shape = VideoPixelShape(batch=1, frames=num_frames, height=64, width=64, fps=args.fps)
+        tgt_shape = AudioLatentShape.from_video_pixel_shape(pixel_shape)
+        audio_tools = AudioLatentTools(patchifier=self.patchifier, target_shape=tgt_shape)
+
+        state = audio_tools.create_initial_state(self.device, self.dtype)
+
+        # ---- Voice reference ----
+        if voice_ref and os.path.exists(voice_ref):
+            voice = decode_audio_from_file(voice_ref, self.device, 0.0, 10.0)
+            if voice is not None:
+                w = voice.waveform
+                if w.dim() == 2:
+                    if w.shape[0] == 1:
+                        w = w.repeat(2, 1)
+                    w = w.unsqueeze(0)
+                elif w.dim() == 3 and w.shape[1] == 1:
+                    w = w.repeat(1, 2, 1)
+                target_samples = int(10.0 * voice.sampling_rate)
+                if w.shape[-1] < target_samples:
+                    w = w.repeat(1, 1, (target_samples // w.shape[-1]) + 1)
+                w = w[..., :target_samples]
+                peak = w.abs().max()
+                if peak > 0:
+                    w = w * (10 ** (-4.0 / 20) / peak)
+                voice = Audio(waveform=w, sampling_rate=voice.sampling_rate)
+                ref_latent = self.audio_conditioner(lambda enc: vae_encode_audio(voice, enc, None))
+                cond = AudioConditionByReferenceLatent(
+                    latent=ref_latent.to(self.device, self.dtype), strength=1.0,
+                )
+                state = cond.apply_to(latent_state=state, latent_tools=audio_tools)
+
+        # ---- Noise ----
+        gen = torch.Generator(device=self.device).manual_seed(args.seed)
+        noiser = GaussianNoiser(generator=gen)
+        state = noiser(state, noise_scale=1.0)
+
+        # ---- Prompt encode ----
+        use_cfg = args.cfg_scale > 1.0
+        prompts = [prompt, args.negative_prompt] if use_cfg else [prompt]
+        ctx = self.prompt_encoder(prompts, streaming_prefetch_count=None)
+        a_ctx = ctx[0].audio_encoding
+        a_ctx_neg = ctx[1].audio_encoding if use_cfg else None
+
+        # ---- Denoiser ----
+        needs_guidance = args.cfg_scale > 1.0 or args.stg_scale > 0.0 or args.modality_scale > 1.0
+        if needs_guidance:
+            guider = MultiModalGuider(
+                params=MultiModalGuiderParams(
+                    cfg_scale=args.cfg_scale, stg_scale=args.stg_scale,
+                    stg_blocks=[args.stg_block] if args.stg_scale > 0 else [],
+                    rescale_scale=args.rescale_scale,
+                    modality_scale=args.modality_scale,
+                    cfg_clamp_scale=args.cfg_clamp,
+                ),
+                negative_context=a_ctx_neg,
+            )
+            denoiser = GuidedDenoiser(
+                v_context=None, a_context=a_ctx,
+                video_guider=None, audio_guider=guider,
+            )
+        else:
+            denoiser = SimpleDenoiser(v_context=None, a_context=a_ctx)
+
+        sigmas = LTX2Scheduler().execute(steps=args.steps, latent=state.latent).to(self.device)
+
+        # ---- Denoise ----
+        # NOTE: don't wrap in gpu_model() — that context manager moves the
+        # model back off GPU on exit, which breaks subsequent iterations of
+        # our warm validator. We keep the velocity model resident.
+        x0 = X0Model(self.velocity_model)
+        batched = BatchSplitAdapter(x0, max_batch_size=1)
+        _, audio_state = euler_denoising_loop(
+            sigmas=sigmas, video_state=None, audio_state=state,
+            stepper=EulerDiffusionStep(), transformer=batched, denoiser=denoiser,
+        )
+
+        audio_state = audio_tools.clear_conditioning(audio_state)
+        audio_state = audio_tools.unpatchify(audio_state)
+        decoded = self.audio_decoder(audio_state.latent)
+
+        wav = decoded.waveform
+        if wav.dim() == 1:
+            wav = wav.unsqueeze(0)
+        os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
+        torchaudio.save(output_path, wav.float().cpu(), decoded.sampling_rate)
+        logging.info(f"  -> {output_path} ({wav.shape[-1]/decoded.sampling_rate:.1f}s, "
+                     f"{time.time()-t_total:.1f}s)")
+
+
+def main():
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+    args = parse_args()
+    import yaml
+    with open(args.val_config) as f:
+        val_cfg = yaml.safe_load(f)
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Build validator once (models warm for all entries).
+    validator = WarmValidator(
+        checkpoint=args.checkpoint,
+        full_checkpoint=args.full_checkpoint,
+        gemma_root=args.gemma_root,
+        lora_path=args.lora,
+        lora_rank=args.lora_rank,
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        dtype=torch.bfloat16,
+    )
+
+    n_ok = n_fail = 0
+    t0 = time.time()
+    for entry in val_cfg.get("speakers", []):
+        name = entry["name"]
+        out_path = os.path.join(args.output_dir, f"{name}.wav")
+        try:
+            validator.generate(
+                prompt=entry["prompt"],
+                output_path=out_path,
+                voice_ref=entry.get("reference"),
+                args=args,
+            )
+            n_ok += 1
+            logging.info(f"  [{name}] OK")
+        except Exception as e:
+            n_fail += 1
+            logging.warning(f"  [{name}] FAILED: {e}")
+            traceback.print_exc()
+
+    logging.info(f"Validation done: ok={n_ok} fail={n_fail} in {(time.time()-t0)/60:.1f}min "
+                 f"at {args.output_dir}")
+
+
+if __name__ == "__main__":
+    main()

engines/f5_tts/model/modules.py

@@ -434,9 +434,19 @@ class Attention(nn.Module):
 
 # Attention processor
 
+# TTS Audio Suite patch: FlashAttention is optional and F5 defaults to the
+# PyTorch backend. Treat partial/incompatible installations as unavailable
+# instead of crashing module import when submodules or compiled symbols are
+# missing.
+FLASH_ATTN_AVAILABLE = False
 if is_package_available("flash_attn"):
-    from flash_attn.bert_padding import pad_input, unpad_input
-    from flash_attn import flash_attn_varlen_func, flash_attn_func
+    try:
+        from flash_attn.bert_padding import pad_input, unpad_input
+        from flash_attn import flash_attn_varlen_func, flash_attn_func
+
+        FLASH_ATTN_AVAILABLE = True
+    except (ImportError, OSError, AttributeError):
+        pass
 
 
 class AttnProcessor:
@@ -447,7 +457,7 @@ class AttnProcessor:
         attn_mask_enabled: bool = True,
     ):
         if attn_backend == "flash_attn":
-            assert is_package_available("flash_attn"), "Please install flash-attn first."
+            assert FLASH_ATTN_AVAILABLE, "Please install a compatible flash-attn build first."
 
         self.pe_attn_head = pe_attn_head
         self.attn_backend = attn_backend

engines/moss_tts/training/common.py

@@ -27,6 +27,7 @@ FRIENDLY_VARIANT_MAP = {
 }
 
 SUPPORTED_DELAY_TRAINING_VARIANTS = {"MOSS-TTS", "MOSS-TTS-v1.5"}
+MOSS_DATASET_AUDIO_EXTENSIONS = {".wav", ".flac", ".mp3", ".ogg", ".m4a"}
 
 
 def slugify(value: str) -> str:
@@ -65,6 +66,71 @@ def resolve_manifest_path(dataset_source: str) -> str:
     raise FileNotFoundError(f"MOSS training manifest not found: {dataset_source}")
 
 
+def _resolve_dataset_source_path(dataset_source: str) -> Path:
+    raw = os.path.expanduser(str(dataset_source or "").strip())
+    if not raw:
+        raise ValueError("dataset_source is required")
+
+    input_dir = folder_paths.get_input_directory()
+    candidates = [Path(raw), Path(input_dir, raw), Path(input_dir, "datasets", raw)]
+    for candidate in candidates:
+        if candidate.is_file() or candidate.is_dir():
+            return candidate.resolve()
+    raise FileNotFoundError(f"MOSS dataset source not found: {dataset_source}")
+
+
+def _build_manifest_from_audio_folder(dataset_dir: Path, recursive: bool) -> str:
+    iterator = dataset_dir.rglob("*") if recursive else dataset_dir.iterdir()
+    audio_paths = sorted(
+        (path for path in iterator if path.is_file() and path.suffix.lower() in MOSS_DATASET_AUDIO_EXTENSIONS),
+        key=lambda path: str(path.relative_to(dataset_dir)).lower(),
+    )
+    if not audio_paths:
+        scope = "recursively" if recursive else ""
+        raise ValueError(f"No supported audio files found {scope} in MOSS dataset folder: {dataset_dir}")
+
+    records: List[Dict[str, str]] = []
+    missing_transcripts: List[str] = []
+    source_paths: List[str] = []
+    for audio_path in audio_paths:
+        transcript_path = audio_path.with_suffix(".txt")
+        if not transcript_path.is_file():
+            missing_transcripts.append(str(audio_path.relative_to(dataset_dir)))
+            continue
+        transcript = transcript_path.read_text(encoding="utf-8-sig").strip()
+        if not transcript:
+            missing_transcripts.append(str(audio_path.relative_to(dataset_dir)))
+            continue
+        records.append({"audio": str(audio_path.resolve()), "text": transcript})
+        source_paths.extend((str(audio_path), str(transcript_path)))
+
+    if missing_transcripts:
+        preview = ", ".join(missing_transcripts[:10])
+        remainder = len(missing_transcripts) - 10
+        if remainder > 0:
+            preview += f", and {remainder} more"
+        raise ValueError(
+            "Every MOSS dataset audio file needs a non-empty .txt transcript with the same basename. "
+            f"Missing or empty transcripts for: {preview}"
+        )
+
+    source_hash = fingerprint_paths(source_paths)
+    manifest_dir = Path(get_moss_training_root(), "imported_manifests")
+    manifest_path = manifest_dir / f"{slugify(dataset_dir.name)}_{source_hash[:12]}.jsonl"
+    if not manifest_path.is_file():
+        dump_jsonl(records, manifest_path)
+    print(f"MOSS dataset folder imported: {dataset_dir} | {len(records)} clips")
+    return str(manifest_path)
+
+
+def resolve_moss_dataset_source(dataset_source: str, recursive: bool = False) -> str:
+    """Resolve an existing JSONL manifest or import a folder of audio/.txt pairs."""
+    source_path = _resolve_dataset_source_path(dataset_source)
+    if source_path.is_file():
+        return str(source_path)
+    return _build_manifest_from_audio_folder(source_path, recursive=bool(recursive))
+
+
 def fingerprint_paths(paths: Sequence[str]) -> str:
     digest = hashlib.md5()
     for path in paths:

engines/moss_tts/training/dataset.py

@@ -25,6 +25,7 @@ from engines.moss_tts.training.common import (
     load_jsonl,
     resolve_codec_path,
     resolve_delay_training_variant,
+    resolve_moss_dataset_source,
     resolve_model_path,
     split_train_val,
     slugify,
@@ -215,18 +216,19 @@ def prepare_moss_training_dataset(
     n_vq: int = 0,
     encode_reference_audio: bool = True,
     reuse_existing: bool = True,
+    recursive_folder_scan: bool = False,
 ) -> Dict[str, Any]:
     # Node UI uses prep_batch_size; keep batch_size for compatibility with older callers.
     effective_batch_size = int(prep_batch_size) if int(prep_batch_size or 0) > 0 else int(batch_size)
 
     variant = resolve_delay_training_variant(shared_settings)
 
-    train_manifest_path = os.path.abspath(dataset_source)
-    if not os.path.isfile(train_manifest_path):
-        raise FileNotFoundError(f"MOSS training manifest not found: {dataset_source}")
-    val_manifest_path = os.path.abspath(validation_source) if str(validation_source or "").strip() else ""
-    if val_manifest_path and not os.path.isfile(val_manifest_path):
-        raise FileNotFoundError(f"MOSS validation manifest not found: {validation_source}")
+    train_manifest_path = resolve_moss_dataset_source(dataset_source, recursive=recursive_folder_scan)
+    val_manifest_path = (
+        resolve_moss_dataset_source(validation_source, recursive=recursive_folder_scan)
+        if str(validation_source or "").strip()
+        else ""
+    )
 
     fingerprint_inputs = [train_manifest_path]
     if val_manifest_path:

engines/training/registry.py

@@ -14,6 +14,7 @@ _HANDLERS: Dict[str, Type[BaseTrainingHandler]] = {}
 _HANDLER_MODULES = {
     "rvc": "engines.rvc.training.handler",
     "moss_tts": "engines.moss_tts.training.handler",
+    "dramabox": "engines.dramabox.training.handler",
 }
 
 

install.py

@@ -15,6 +15,7 @@ import subprocess
 import sys
 import os
 import platform
+import importlib.machinery
 import importlib.util
 import hashlib
 import json
@@ -177,8 +178,8 @@ class TTSAudioInstaller:
         else:
             self.log("All requirements.txt dependencies already satisfied", "SUCCESS")
 
-    def check_system_dependencies(self):
-        """Check for required system libraries and provide helpful error messages"""
+    def check_system_dependencies(self):
+        """Check optional system libraries and provide helpful feature warnings."""
         if self.is_windows:
             return True  # Windows packages come pre-compiled
             
@@ -210,8 +211,8 @@ class TTSAudioInstaller:
 
         return False
 
-    def check_macos_dependencies(self):
-        """Check for required system libraries on macOS"""
+    def check_macos_dependencies(self):
+        """Check for optional system libraries on macOS."""
         self.log("Checking macOS system dependencies...", "INFO")
         missing_deps = []
 
@@ -230,12 +231,12 @@ class TTSAudioInstaller:
             pass
         
         if missing_deps:
-            self.log("Missing system dependencies detected!", "WARNING")
-            print("\n" + "="*60)
-            print("MACOS SYSTEM DEPENDENCIES REQUIRED")
-            print("="*60)
-            for dep, purpose in missing_deps:
-                print(f"• {dep} (for {purpose})")
+            self.log("Optional system dependencies are missing", "WARNING")
+            print("\n" + "="*60)
+            print("OPTIONAL MACOS SYSTEM DEPENDENCIES")
+            print("="*60)
+            for dep, purpose in missing_deps:
+                print(f"• {dep} (for {purpose})")
             
             print("\nPlease install with Homebrew:")
             deps_list = " ".join([dep for dep, _ in missing_deps])
@@ -249,15 +250,19 @@ class TTSAudioInstaller:
                 print("Check with: python -c \"import platform; print(platform.machine())\"")
                 print("Should show: arm64 (not x86_64)")
             
-            print("="*60)
-            print("Then run this install script again.\n")
-            return False
-        
-        self.log("macOS system dependencies check passed", "SUCCESS")
-        return True
-    
-    def check_linux_dependencies(self):
-        """Check for required system libraries on Linux"""
+            print("="*60)
+            print("Core TTS installation will continue; only the listed features may be unavailable.\n")
+            self.add_install_warning(
+                "Optional macOS audio libraries",
+                "Install the listed Homebrew packages to enable every audio feature.",
+            )
+
+        if not missing_deps:
+            self.log("macOS system dependencies check passed", "SUCCESS")
+        return True
+
+    def check_linux_dependencies(self):
+        """Check for optional system libraries on Linux."""
         self.log("Checking Linux system dependencies...", "INFO")
         missing_deps = []
         
@@ -279,26 +284,34 @@ class TTSAudioInstaller:
             pass
         
         if missing_deps:
-            self.log("Missing system dependencies detected!", "WARNING")
-            print("\n" + "="*60)
-            print("LINUX SYSTEM DEPENDENCIES REQUIRED")
-            print("="*60)
-            for dep, purpose in missing_deps:
-                print(f"• {dep} (for {purpose})")
+            self.log("Optional system dependencies are missing", "WARNING")
+            print("\n" + "="*60)
+            print("OPTIONAL LINUX SYSTEM DEPENDENCIES")
+            print("="*60)
+            for dep, purpose in missing_deps:
+                print(f"• {dep} (for {purpose})")
             
             print("\nPlease install with:")
             print("# Ubuntu/Debian:")
             deps_list = " ".join([dep for dep, _ in missing_deps])
             print(f"sudo apt-get install {deps_list}")
             print("\n# Fedora/RHEL:")
-            fedora_deps = deps_list.replace('-dev', '-devel').replace('19', '')
-            print(f"sudo dnf install {fedora_deps}")
-            print("="*60)
-            print("Then run this install script again.\n")
-            return False
-        
-        self.log("Linux system dependencies check passed", "SUCCESS")
-        return True
+            fedora_packages = {
+                "libsamplerate0-dev": "libsamplerate-devel",
+                "portaudio19-dev": "portaudio-devel",
+            }
+            fedora_deps = " ".join(fedora_packages.get(dep, dep) for dep, _ in missing_deps)
+            print(f"sudo dnf install {fedora_deps}")
+            print("="*60)
+            print("Core TTS installation will continue; only the listed features may be unavailable.\n")
+            self.add_install_warning(
+                "Optional Linux audio libraries",
+                "Install the listed system packages to enable every audio feature.",
+            )
+
+        if not missing_deps:
+            self.log("Linux system dependencies check passed", "SUCCESS")
+        return True
 
     def run_pip_command(self, args: List[str], description: str, ignore_errors: bool = False) -> bool:
         """Execute pip command with error handling and Windows UTF-8 support"""
@@ -507,7 +520,25 @@ class TTSAudioInstaller:
     def module_available(self, module_name: str) -> bool:
         """Check module presence without starting another Python process or importing it."""
         try:
-            return importlib.util.find_spec(module_name) is not None
+            parts = module_name.split(".")
+            spec = importlib.util.find_spec(parts[0])
+            if spec is None:
+                return False
+
+            # util.find_spec() imports the parent when given a dotted name.
+            # Walk the package paths directly so presence checks stay side-effect free.
+            for index in range(1, len(parts)):
+                search_locations = spec.submodule_search_locations
+                if search_locations is None:
+                    return False
+                qualified_name = ".".join(parts[: index + 1])
+                spec = importlib.machinery.PathFinder.find_spec(
+                    qualified_name,
+                    search_locations,
+                )
+                if spec is None:
+                    return False
+            return True
         except (ImportError, ModuleNotFoundError, AttributeError, ValueError):
             return False
 
@@ -1907,10 +1938,8 @@ def main():
         else:
             installer.ensure_requirements_installed()  # Ensure requirements.txt is installed first
 
-            # Check system dependencies (Linux only)
-            if not installer.check_system_dependencies():
-                installer.log("System dependency check failed - aborting installation", "ERROR")
-                sys.exit(1)
+            # Report optional system dependencies without blocking engine setup.
+            installer.check_system_dependencies()
 
             # Install in correct order to prevent conflicts
             installer.install_pytorch_with_cuda()  # Install PyTorch first with proper CUDA detection

nodes.py

@@ -11,7 +11,7 @@ except ImportError:
     pass
 
 # Version and constants
-VERSION = "5.6.0"
+VERSION = "5.6.5"
 IS_DEV = False  # Set to False for release builds
 VERSION_DISPLAY = f"v{VERSION}" + (" (dev)" if IS_DEV else "")
 SEPARATOR = "=" * 70
@@ -500,6 +500,30 @@ except Exception as e:
     print(f"❌ MOSS Dataset Rows failed: {e}")
     MOSS_DATASET_ROWS_AVAILABLE = False
 
+try:
+    dramabox_dataset_prep_module = load_node_module("dramabox_dataset_prep_node", "training/dramabox_dataset_prep_node.py")
+    DramaBoxDatasetPrepNode = dramabox_dataset_prep_module.DramaBoxDatasetPrepNode
+    DRAMABOX_DATASET_PREP_AVAILABLE = True
+except Exception as e:
+    print(f"❌ DramaBox Dataset Prep failed: {e}")
+    DRAMABOX_DATASET_PREP_AVAILABLE = False
+
+try:
+    dramabox_dataset_rows_module = load_node_module("dramabox_dataset_rows_node", "training/dramabox_dataset_rows_node.py")
+    DramaBoxDatasetRowsNode = dramabox_dataset_rows_module.DramaBoxDatasetRowsNode
+    DRAMABOX_DATASET_ROWS_AVAILABLE = True
+except Exception as e:
+    print(f"❌ DramaBox Dataset Rows failed: {e}")
+    DRAMABOX_DATASET_ROWS_AVAILABLE = False
+
+try:
+    dramabox_training_config_module = load_node_module("dramabox_training_config_node", "training/dramabox_training_config_node.py")
+    DramaBoxTrainingConfigNode = dramabox_training_config_module.DramaBoxTrainingConfigNode
+    DRAMABOX_TRAINING_CONFIG_AVAILABLE = True
+except Exception as e:
+    print(f"❌ DramaBox Training Config failed: {e}")
+    DRAMABOX_TRAINING_CONFIG_AVAILABLE = False
+
 try:
     phoneme_text_normalizer_module = load_node_module("phoneme_text_normalizer_node", "text/phoneme_text_normalizer_node.py")
     PhonemeTextNormalizer = phoneme_text_normalizer_module.PhonemeTextNormalizer
@@ -839,12 +863,24 @@ if MOSS_TRAINING_CONFIG_AVAILABLE:
 
 if MOSS_CLIP_STAGING_AVAILABLE:
     NODE_CLASS_MAPPINGS["MossClipStagingNode"] = MossClipStagingNode
-    NODE_DISPLAY_NAME_MAPPINGS["MossClipStagingNode"] = "🎞️ MOSS Clip Staging"
+    NODE_DISPLAY_NAME_MAPPINGS["MossClipStagingNode"] = "🎞️ Training Clip Staging"
 
 if MOSS_DATASET_ROWS_AVAILABLE:
     NODE_CLASS_MAPPINGS["MossDatasetRowsNode"] = MossDatasetRowsNode
     NODE_DISPLAY_NAME_MAPPINGS["MossDatasetRowsNode"] = "🧾 MOSS Dataset Rows"
 
+if DRAMABOX_DATASET_PREP_AVAILABLE:
+    NODE_CLASS_MAPPINGS["DramaBoxDatasetPrepNode"] = DramaBoxDatasetPrepNode
+    NODE_DISPLAY_NAME_MAPPINGS["DramaBoxDatasetPrepNode"] = "📦 DramaBox Dataset Prep"
+
+if DRAMABOX_DATASET_ROWS_AVAILABLE:
+    NODE_CLASS_MAPPINGS["DramaBoxDatasetRowsNode"] = DramaBoxDatasetRowsNode
+    NODE_DISPLAY_NAME_MAPPINGS["DramaBoxDatasetRowsNode"] = "🧾 DramaBox Dataset Rows"
+
+if DRAMABOX_TRAINING_CONFIG_AVAILABLE:
+    NODE_CLASS_MAPPINGS["DramaBoxTrainingConfigNode"] = DramaBoxTrainingConfigNode
+    NODE_DISPLAY_NAME_MAPPINGS["DramaBoxTrainingConfigNode"] = "🎛️ DramaBox Training Config"
+
 # Register text processing nodes
 if PHONEME_TEXT_NORMALIZER_AVAILABLE:
     NODE_CLASS_MAPPINGS["PhonemeTextNormalizer"] = PhonemeTextNormalizer

nodes/engines/dramabox_engine_node.py

@@ -3,6 +3,7 @@
 import importlib.util
 import os
 import sys
+from typing import List
 
 current_dir = os.path.dirname(__file__)
 nodes_dir = os.path.dirname(current_dir)
@@ -18,6 +19,7 @@ base_spec.loader.exec_module(base_module)
 BaseTTSNode = base_module.BaseTTSNode
 
 from engines.dramabox.dramabox_downloader import DramaBoxDownloader
+from utils.models.extra_paths import get_all_tts_model_paths
 
 
 class DramaBoxEngineNode(BaseTTSNode):
@@ -144,7 +146,8 @@ class DramaBoxEngineNode(BaseTTSNode):
                     "default": "none",
                     "tooltip": (
                         "Official LTX FP8 weight-storage policy for the diffusion transformer. "
-                        "fp8_cast lowers VRAM but upcasts each linear layer during inference."
+                        "fp8_cast lowers VRAM but upcasts each linear layer during inference. "
+                        "DramaBox LoRAs remain as an unmerged BF16 branch over the FP8 base."
                     ),
                 }),
                 "compile_model": ("BOOLEAN", {
@@ -154,6 +157,27 @@ class DramaBoxEngineNode(BaseTTSNode):
                         "slower and may reserve more VRAM; later denoising can be faster."
                     ),
                 }),
+                "local_lora_adapter": (cls._get_ui_lora_options(), {
+                    "default": "None",
+                    "tooltip": (
+                        "Optional DramaBox audio LoRA discovered under models/TTS/dramabox/loras. "
+                        "Training outputs are copied there when a run completes."
+                    ),
+                }),
+                "lora_adapter_override": ("STRING", {
+                    "default": "",
+                    "tooltip": (
+                        "Advanced local path to a DramaBox LoRA file or adapter folder. "
+                        "If filled, this overrides the local adapter dropdown."
+                    ),
+                }),
+                "lora_strength": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.0,
+                    "max": 2.0,
+                    "step": 0.05,
+                    "tooltip": "Scale applied to the trained DramaBox LoRA. 1.0 uses the adapter's trained strength; 0 disables it.",
+                }),
             },
         }
 
@@ -179,7 +203,11 @@ class DramaBoxEngineNode(BaseTTSNode):
         memory_mode: str = "fast",
         transformer_quantization: str = "none",
         compile_model: bool = False,
+        local_lora_adapter: str = "None",
+        lora_adapter_override: str = "",
+        lora_strength: float = 1.0,
     ) -> tuple:
+        lora_path = self._resolve_lora_adapter(local_lora_adapter, lora_adapter_override)
         config = {
             "engine_type": "dramabox",
             "model_name": model_name,
@@ -197,6 +225,8 @@ class DramaBoxEngineNode(BaseTTSNode):
             "memory_mode": str(memory_mode),
             "transformer_quantization": str(transformer_quantization),
             "compile_model": bool(compile_model),
+            "lora_path": lora_path,
+            "lora_strength": float(lora_strength),
         }
         print(f"⚙️ DramaBox: {model_name} on {device} ({precision})")
         print(
@@ -206,6 +236,8 @@ class DramaBoxEngineNode(BaseTTSNode):
             f"watermark={watermark}, memory_mode={memory_mode}, "
             f"transformer_quantization={transformer_quantization}, compile={compile_model}"
         )
+        if lora_path:
+            print(f"   LoRA: {lora_path} (strength={float(lora_strength):.2f})")
         print("   Prompt: dialogue in quotes; expressive stage directions outside quotes")
         return ({
             "engine_type": "dramabox",
@@ -213,6 +245,51 @@ class DramaBoxEngineNode(BaseTTSNode):
             "capabilities": ["tts"],
         },)
 
+    @classmethod
+    def _discover_local_loras(cls) -> List[str]:
+        discovered: List[str] = []
+        seen = set()
+        try:
+            for base_path in get_all_tts_model_paths("TTS"):
+                root = os.path.join(base_path, "dramabox", "loras")
+                if not os.path.isdir(root):
+                    continue
+                for name in sorted(os.listdir(root)):
+                    candidate = os.path.join(root, name)
+                    if os.path.isdir(candidate):
+                        has_weights = any(
+                            filename.endswith(".safetensors")
+                            for filename in os.listdir(candidate)
+                        )
+                    else:
+                        has_weights = os.path.isfile(candidate) and candidate.endswith(".safetensors")
+                    if has_weights and f"local:{name}" not in seen:
+                        seen.add(f"local:{name}")
+                        discovered.append(f"local:{name}")
+        except Exception:
+            pass
+        return discovered
+
+    @classmethod
+    def _get_ui_lora_options(cls) -> List[str]:
+        return ["None"] + cls._discover_local_loras()
+
+    @classmethod
+    def _resolve_lora_adapter(cls, local_value: str, override: str) -> str:
+        manual = str(override or "").strip()
+        if manual:
+            return os.path.abspath(os.path.expanduser(manual))
+        selected = str(local_value or "").strip()
+        if not selected or selected == "None":
+            return ""
+        if selected.startswith("local:"):
+            name = selected.split(":", 1)[1]
+            for base_path in get_all_tts_model_paths("TTS"):
+                candidate = os.path.join(base_path, "dramabox", "loras", name)
+                if os.path.exists(candidate):
+                    return candidate
+        return os.path.abspath(os.path.expanduser(selected))
+
     @staticmethod
     def _validate_rescale_scale(value: str):
         text = str(value).strip().lower()

nodes/training/dramabox_dataset_prep_node.py

@@ -0,0 +1,111 @@
+"""DramaBox dataset normalization and official-preprocessor node."""
+
+import importlib.util
+import os
+import sys
+
+from engines.training.registry import get_training_handler
+
+current_dir = os.path.dirname(__file__)
+nodes_dir = os.path.dirname(current_dir)
+project_root = os.path.dirname(nodes_dir)
+if project_root not in sys.path:
+    sys.path.insert(0, project_root)
+
+base_node_path = os.path.join(nodes_dir, "base", "base_node.py")
+base_spec = importlib.util.spec_from_file_location("base_node_module", base_node_path)
+base_module = importlib.util.module_from_spec(base_spec)
+sys.modules["base_node_module"] = base_module
+base_spec.loader.exec_module(base_module)
+BaseTTSNode = base_module.BaseTTSNode
+
+
+class DramaBoxDatasetPrepNode(BaseTTSNode):
+    @classmethod
+    def NAME(cls):
+        return "📦 DramaBox Dataset Prep"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "TTS_engine": ("TTS_ENGINE", {
+                    "tooltip": "Connect a DramaBox engine. Its selected model supplies the official transformer, audio components, and Gemma paths.",
+                }),
+                "model_name": ("STRING", {
+                    "default": "MyDramaBoxLoRA",
+                    "tooltip": "Name used for the prepared dataset and eventual managed LoRA adapter.",
+                }),
+                "dataset_source": ("STRING", {
+                    "default": "",
+                    "tooltip": "JSONL/JSON manifest, TSV, gemini_synthetic index, or libriheavy index. Manifest rows should contain audio_filepath/audio_path and text/transcript.",
+                }),
+                "dataset_type": (["manifest", "tsv", "gemini_synthetic", "libriheavy"], {
+                    "default": "manifest",
+                    "tooltip": "Input format. The suite converts every format into the official ~-delimited speaker index used by the trainer.",
+                }),
+            },
+            "optional": {
+                "audio_dir": ("STRING", {
+                    "default": "",
+                    "tooltip": "Base folder for relative audio paths. Blank resolves paths relative to the dataset file.",
+                }),
+                "min_duration": ("FLOAT", {
+                    "default": 2.0,
+                    "min": 0.1,
+                    "max": 60.0,
+                    "step": 0.1,
+                    "tooltip": "Minimum clip duration passed to the official preprocessor.",
+                }),
+                "max_duration": ("FLOAT", {
+                    "default": 20.0,
+                    "min": 0.5,
+                    "max": 120.0,
+                    "step": 0.5,
+                    "tooltip": "Maximum clip duration passed to the official preprocessor.",
+                }),
+                "reuse_existing": ("BOOLEAN", {
+                    "default": True,
+                    "tooltip": "Reuse a matching normalized index and already-preprocessed cache when available.",
+                }),
+                "preprocess_now": ("BOOLEAN", {
+                    "default": True,
+                    "tooltip": "Run the official Gemma/audio-VAE preprocessing now. Turn this off to prepare only the CPU-side index and let Model Training preprocess later.",
+                }),
+                "dry_run": ("BOOLEAN", {
+                    "default": False,
+                    "tooltip": "CPU-safe index-only mode. No model download, Gemma load, or CUDA preprocessing is started.",
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("TRAINING_DATASET", "STRING")
+    RETURN_NAMES = ("training_dataset", "dataset_info")
+    FUNCTION = "prepare_dataset"
+    CATEGORY = "TTS Audio Suite/🎓 Training"
+
+    def prepare_dataset(self, TTS_engine, model_name, dataset_source, dataset_type, **kwargs):
+        handler = get_training_handler("dramabox")
+        if handler is None:
+            raise RuntimeError("DramaBox training backend is not available")
+        dataset = handler.prepare_dataset(
+            TTS_engine,
+            dataset_source=dataset_source,
+            model_name=model_name,
+            dataset_type=dataset_type,
+            **kwargs,
+        )
+        info = (
+            f"DramaBox dataset ready: {dataset['model_name']} | "
+            f"{dataset['train_records']} clips | "
+            f"{len(dataset['speakers'])} speaker(s) | "
+            f"preprocessed={dataset.get('preprocessed', False)}"
+        )
+        print(f"📦 {info}")
+        return dataset, info
+
+
+NODE_CLASS_MAPPINGS = {"DramaBoxDatasetPrepNode": DramaBoxDatasetPrepNode}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "DramaBoxDatasetPrepNode": "📦 DramaBox Dataset Prep"
+}

nodes/training/dramabox_dataset_rows_node.py

@@ -0,0 +1,197 @@
+"""Build a DramaBox training manifest from engine-neutral staged clips."""
+
+import importlib.util
+import json
+import os
+import sys
+from typing import List
+
+import folder_paths
+
+current_dir = os.path.dirname(__file__)
+nodes_dir = os.path.dirname(current_dir)
+project_root = os.path.dirname(nodes_dir)
+if project_root not in sys.path:
+    sys.path.insert(0, project_root)
+
+base_node_path = os.path.join(nodes_dir, "base", "base_node.py")
+base_spec = importlib.util.spec_from_file_location("base_node_module", base_node_path)
+base_module = importlib.util.module_from_spec(base_spec)
+sys.modules["base_node_module"] = base_module
+base_spec.loader.exec_module(base_module)
+BaseTTSNode = base_module.BaseTTSNode
+
+
+def _required_lines(raw_text: str, expected_count: int) -> List[str]:
+    lines = str(raw_text or "").splitlines()
+    if len(lines) != expected_count:
+        raise ValueError(
+            "DramaBox transcript line count mismatch: "
+            f"expected {expected_count} line(s), got {len(lines)}. "
+            "Enter exactly one transcript per staged clip."
+        )
+    return [line.strip() for line in lines]
+
+
+def _optional_lines(raw_text: str, expected_count: int, field_name: str) -> List[str]:
+    lines = str(raw_text or "").splitlines()
+    if len(lines) > expected_count:
+        raise ValueError(
+            f"DramaBox {field_name} line count mismatch: expected at most "
+            f"{expected_count} line(s), got {len(lines)}."
+        )
+    return [line.strip() for line in lines] + [""] * (expected_count - len(lines))
+
+
+class DramaBoxDatasetRowsNode(BaseTTSNode):
+    @classmethod
+    def NAME(cls):
+        return "🧾 DramaBox Dataset Rows"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "clip_dataset": ("TRAINING_CLIP_DATASET", {
+                    "tooltip": "Staged audio from Training Clip Staging.",
+                }),
+                "manifest_name": ("STRING", {
+                    "default": "dramabox_train.jsonl",
+                    "tooltip": "Output manifest filename. .jsonl is appended when missing.",
+                }),
+                "transcript_lines": ("STRING", {
+                    "default": "Hello there, this is a training sample.\nThis is the second sample from the same speaker.",
+                    "multiline": True,
+                    "tooltip": "Exactly one line per staged clip, in clip order. Blank lines skip the corresponding clip.",
+                }),
+            },
+            "optional": {
+                "speaker_lines": ("STRING", {
+                    "default": "",
+                    "multiline": True,
+                    "tooltip": "Optional speaker name per clip. Blank lines use default_speaker.",
+                }),
+                "language_lines": ("STRING", {
+                    "default": "",
+                    "multiline": True,
+                    "tooltip": "Optional language code per clip. Blank lines use default_language.",
+                }),
+                "default_speaker": ("STRING", {
+                    "default": "speaker_1",
+                    "tooltip": "Speaker assigned when the corresponding speaker line is blank. Each DramaBox speaker needs at least two clips.",
+                }),
+                "default_language": ("STRING", {
+                    "default": "en",
+                    "tooltip": "Language code assigned when the corresponding language line is blank.",
+                }),
+                "output_subdir": ("STRING", {
+                    "default": "tts_audio_suite_training/dramabox/manifests",
+                    "tooltip": "Subdirectory inside ComfyUI input/ for the generated manifest.",
+                }),
+                "overwrite": ("BOOLEAN", {
+                    "default": True,
+                    "tooltip": "Overwrite an existing manifest with the same name.",
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("STRING", "STRING")
+    RETURN_NAMES = ("manifest_path", "manifest_info")
+    FUNCTION = "build_rows"
+    CATEGORY = "TTS Audio Suite/🎓 Training"
+
+    def build_rows(
+        self,
+        clip_dataset,
+        manifest_name: str,
+        transcript_lines: str,
+        speaker_lines: str = "",
+        language_lines: str = "",
+        default_speaker: str = "speaker_1",
+        default_language: str = "en",
+        output_subdir: str = "",
+        overwrite: bool = True,
+    ):
+        if not isinstance(clip_dataset, dict) or clip_dataset.get("type") not in {
+            "training_clip_dataset",
+            "moss_clip_dataset",
+        }:
+            raise ValueError("clip_dataset must come from Training Clip Staging")
+
+        clips = clip_dataset.get("clips") or []
+        if not clips:
+            raise ValueError("clip_dataset contains no clips")
+
+        clip_count = len(clips)
+        transcripts = _required_lines(transcript_lines, clip_count)
+        speakers = _optional_lines(speaker_lines, clip_count, "speaker_lines")
+        languages = _optional_lines(language_lines, clip_count, "language_lines")
+        fallback_speaker = str(default_speaker or "").strip() or "speaker_1"
+        fallback_language = str(default_language or "").strip() or "en"
+
+        records = []
+        speaker_counts = {}
+        skipped_rows = 0
+        for index, clip in enumerate(clips):
+            if not transcripts[index]:
+                skipped_rows += 1
+                continue
+            speaker = speakers[index] or fallback_speaker
+            language = languages[index] or fallback_language
+            speaker_counts[speaker] = speaker_counts.get(speaker, 0) + 1
+            records.append({
+                "audio_filepath": str(clip["audio"]),
+                "text": transcripts[index],
+                "speaker": speaker,
+                "language": language,
+                "duration": float(clip["duration_seconds"]),
+                "sample_rate": int(clip["sample_rate"]),
+                "samples": round(
+                    float(clip["duration_seconds"]) * int(clip["sample_rate"])
+                ),
+            })
+
+        if not records:
+            raise RuntimeError(
+                "DramaBox Dataset Rows produced no records. Add at least two "
+                "non-empty transcripts for one speaker."
+            )
+
+        short_speakers = sorted(
+            speaker for speaker, count in speaker_counts.items() if count < 2
+        )
+        if short_speakers:
+            raise ValueError(
+                "DramaBox needs at least two clips per speaker. Speakers with only "
+                "one staged clip: " + ", ".join(short_speakers)
+            )
+
+        filename = str(manifest_name or "").strip() or "dramabox_train.jsonl"
+        if not filename.lower().endswith(".jsonl"):
+            filename += ".jsonl"
+        input_root = folder_paths.get_input_directory()
+        subdir = str(output_subdir or "").strip().strip("/\\")
+        output_dir = os.path.join(input_root, subdir) if subdir else input_root
+        os.makedirs(output_dir, exist_ok=True)
+        manifest_path = os.path.join(output_dir, filename)
+        if os.path.exists(manifest_path) and not overwrite:
+            raise FileExistsError(f"DramaBox manifest already exists: {manifest_path}")
+
+        with open(manifest_path, "w", encoding="utf-8") as handle:
+            for record in records:
+                handle.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+        info = (
+            f"DramaBox manifest ready: {os.path.basename(manifest_path)} | "
+            f"{len(records)} clips | {len(speaker_counts)} speaker(s)"
+        )
+        if skipped_rows:
+            info += f" | skipped {skipped_rows} blank transcript row(s)"
+        print(f"🧾 {info}")
+        return manifest_path, info
+
+
+NODE_CLASS_MAPPINGS = {"DramaBoxDatasetRowsNode": DramaBoxDatasetRowsNode}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "DramaBoxDatasetRowsNode": "🧾 DramaBox Dataset Rows"
+}

nodes/training/dramabox_training_config_node.py

@@ -0,0 +1,194 @@
+"""DramaBox IC-LoRA training configuration node."""
+
+import importlib.util
+import os
+import sys
+
+current_dir = os.path.dirname(__file__)
+nodes_dir = os.path.dirname(current_dir)
+project_root = os.path.dirname(nodes_dir)
+if project_root not in sys.path:
+    sys.path.insert(0, project_root)
+
+base_node_path = os.path.join(nodes_dir, "base", "base_node.py")
+base_spec = importlib.util.spec_from_file_location("base_node_module", base_node_path)
+base_module = importlib.util.module_from_spec(base_spec)
+sys.modules["base_node_module"] = base_module
+base_spec.loader.exec_module(base_module)
+BaseTTSNode = base_module.BaseTTSNode
+
+
+class DramaBoxTrainingConfigNode(BaseTTSNode):
+    @classmethod
+    def NAME(cls):
+        return "🎛️ DramaBox Training Config"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "training_mode": (["Audio LoRA (IC-LoRA)"], {
+                    "default": "Audio LoRA (IC-LoRA)",
+                    "tooltip": "Official DramaBox audio-branch IC-LoRA training mode.",
+                }),
+                "base_model": (["dev", "distilled"], {
+                    "default": "dev",
+                    "tooltip": "Official timestep schedule. dev is the normal DramaBox fine-tuning choice; distilled is experimental.",
+                }),
+                "steps": ("INT", {
+                    "default": 10000,
+                    "min": 1,
+                    "max": 1000000,
+                    "step": 100,
+                    "tooltip": "Optimizer steps. The upstream example uses 10,000; listen to saved checkpoints instead of assuming the final step is best.",
+                }),
+                "learning_rate": ("FLOAT", {
+                    "default": 1e-4,
+                    "min": 1e-8,
+                    "max": 1.0,
+                    "step": 1e-6,
+                    "tooltip": "LoRA learning rate. The official example uses 1e-4 for a fresh adapter.",
+                }),
+                "batch_size": ("INT", {
+                    "default": 1,
+                    "min": 1,
+                    "max": 32,
+                    "step": 1,
+                    "tooltip": "Per-device batch size. Keep this at 1 unless the dataset and GPU have room.",
+                }),
+                "grad_accum": ("INT", {
+                    "default": 4,
+                    "min": 1,
+                    "max": 256,
+                    "step": 1,
+                    "tooltip": "Gradient accumulation steps. This increases effective batch size without loading more samples at once.",
+                }),
+                "lora_rank": ("INT", {
+                    "default": 128,
+                    "min": 1,
+                    "max": 512,
+                    "step": 1,
+                    "tooltip": "LoRA rank. The official DramaBox example uses 128.",
+                }),
+                "lora_alpha": ("INT", {
+                    "default": 128,
+                    "min": 1,
+                    "max": 1024,
+                    "step": 1,
+                    "tooltip": "LoRA alpha. Keeping alpha equal to rank gives a 1.0 adapter scale.",
+                }),
+                "lora_dropout": ("FLOAT", {
+                    "default": 0.1,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.01,
+                    "tooltip": "LoRA dropout. The official small-dataset example uses 0.1.",
+                }),
+            },
+            "optional": {
+                "lr_scheduler": (["cosine", "linear", "constant"], {
+                    "default": "cosine",
+                    "tooltip": "Learning-rate schedule passed to the official trainer.",
+                }),
+                "warmup_steps": ("INT", {
+                    "default": 500,
+                    "min": 0,
+                    "max": 100000,
+                    "step": 10,
+                    "tooltip": "Warmup steps before the selected schedule. The official example uses 500.",
+                }),
+                "max_grad_norm": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.0,
+                    "max": 10.0,
+                    "step": 0.1,
+                    "tooltip": "Gradient clipping threshold.",
+                }),
+                "ref_ratio": ("FLOAT", {
+                    "default": 0.3,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.05,
+                    "tooltip": "Fraction of a training target used as the appended voice-reference tail.",
+                }),
+                "max_ref_tokens": ("INT", {
+                    "default": 200,
+                    "min": 0,
+                    "max": 4096,
+                    "step": 1,
+                    "tooltip": "Maximum reference tokens after audio patchification.",
+                }),
+                "text_dropout": ("FLOAT", {
+                    "default": 0.4,
+                    "min": 0.0,
+                    "max": 1.0,
+                    "step": 0.05,
+                    "tooltip": "Probability of dropping text conditioning so the adapter learns to use the reference voice path.",
+                }),
+                "save_every": ("INT", {
+                    "default": 500,
+                    "min": 1,
+                    "max": 100000,
+                    "step": 10,
+                    "tooltip": "Checkpoint cadence. The official trainer requires a value of at least 1.",
+                }),
+                "log_every": ("INT", {
+                    "default": 10,
+                    "min": 1,
+                    "max": 10000,
+                    "step": 1,
+                    "tooltip": "Human-readable console update cadence. The training panel receives quieter per-step updates.",
+                }),
+                "seed": ("INT", {
+                    "default": 42,
+                    "min": 0,
+                    "max": 2**31 - 1,
+                    "step": 1,
+                    "tooltip": "Training random seed.",
+                }),
+                "preprocess_batch_size": ("INT", {
+                    "default": 8,
+                    "min": 1,
+                    "max": 64,
+                    "step": 1,
+                    "tooltip": "Audio/text preprocessing batch size. Lower this if preprocessing runs out of memory.",
+                }),
+                "validation_config": ("STRING", {
+                    "default": "",
+                    "tooltip": "Optional path to the official val_config YAML. Validation launches another full inference process at each save step and requires a separate GPU.",
+                }),
+                "validation_gpu": ("STRING", {
+                    "default": "",
+                    "tooltip": "Physical CUDA device index reserved for validation, for example 1. Required when validation_config is set and must differ from the training GPU.",
+                }),
+                "dry_run": ("BOOLEAN", {
+                    "default": False,
+                    "tooltip": "CPU-safe preflight only: writes the normalized official config and command without loading DramaBox weights or starting CUDA training.",
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("TRAINING_CONFIG", "STRING")
+    RETURN_NAMES = ("training_config", "config_info")
+    FUNCTION = "create_config"
+    CATEGORY = "TTS Audio Suite/🎓 Training"
+
+    def create_config(self, **kwargs):
+        kwargs["training_mode"] = "audio_lora"
+        config = {
+            "type": "training_config",
+            "engine_type": "dramabox",
+            **kwargs,
+        }
+        info = (
+            f"DramaBox audio LoRA config: {config['base_model']} | "
+            f"{config['steps']} steps | batch {config['batch_size']} | "
+            f"rank {config['lora_rank']} | lr {config['learning_rate']}"
+        )
+        return config, info
+
+
+NODE_CLASS_MAPPINGS = {"DramaBoxTrainingConfigNode": DramaBoxTrainingConfigNode}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "DramaBoxTrainingConfigNode": "🎛️ DramaBox Training Config"
+}

nodes/training/moss_clip_staging_node.py

@@ -1,6 +1,4 @@
-"""
-MOSS clip staging node for unified training workflows.
-"""
+"""Engine-neutral audio clip staging for training workflows."""
 
 import os
 import re
@@ -59,7 +57,7 @@ class DynamicAudioOptionalInputs(dict):
 def _slugify(value: str) -> str:
     safe = "".join(ch if ch.isalnum() or ch in ("-", "_") else "_" for ch in str(value).strip())
     safe = safe.strip("_")
-    return safe or "moss_dataset"
+    return safe or "training_dataset"
 
 
 def _iter_audio_batches(waveform):
@@ -74,7 +72,7 @@ def _iter_audio_batches(waveform):
             yield clip[None, :]
         return
     if waveform.ndim != 3:
-        raise ValueError(f"Unsupported audio tensor shape for MOSS clip staging: {tuple(waveform.shape)}")
+        raise ValueError(f"Unsupported audio tensor shape for clip staging: {tuple(waveform.shape)}")
     for clip in waveform:
         if clip.ndim == 1:
             yield clip[None, :]
@@ -96,17 +94,19 @@ def _write_audio_clip(audio_tensor, sample_rate: int, output_path: str):
 
 
 class MossClipStagingNode(BaseTTSNode):
+    """Legacy class id retained so existing MOSS workflows keep loading."""
+
     @classmethod
     def NAME(cls):
-        return "🎞️ MOSS Clip Staging"
+        return "🎞️ Training Clip Staging"
 
     @classmethod
     def INPUT_TYPES(cls):
         optional_inputs = DynamicAudioOptionalInputs(
             {
                 "output_subdir": ("STRING", {
-                    "default": "tts_audio_suite_training/moss_tts/staged_audio",
-                    "tooltip": "Subdirectory inside ComfyUI input/ where staged MOSS training clips will be written."
+                    "default": "tts_audio_suite_training/staged_audio",
+                    "tooltip": "Subdirectory inside ComfyUI input/ where reusable training clips will be written."
                 }),
                 "overwrite": ("BOOLEAN", {
                     "default": True,
@@ -124,14 +124,14 @@ class MossClipStagingNode(BaseTTSNode):
         return {
             "required": {
                 "dataset_name": ("STRING", {
-                    "default": "MyMossDataset",
-                    "tooltip": "Base name for the staged clip set."
+                    "default": "MyTrainingDataset",
+                    "tooltip": "Base name for the staged clip set. The output can feed engine-specific Dataset Rows nodes."
                 }),
             },
             "optional": optional_inputs,
         }
 
-    RETURN_TYPES = ("MOSS_CLIP_DATASET", "STRING")
+    RETURN_TYPES = ("TRAINING_CLIP_DATASET", "STRING")
     RETURN_NAMES = ("clip_dataset", "dataset_info")
     FUNCTION = "stage_clips"
     CATEGORY = "TTS Audio Suite/🎓 Training"
@@ -159,7 +159,7 @@ class MossClipStagingNode(BaseTTSNode):
     ):
         audio_inputs = self._collect_audio_inputs(opt_audio1=opt_audio1, **kwargs)
         if not audio_inputs:
-            raise ValueError("MOSS Clip Staging requires at least one connected AUDIO input")
+            raise ValueError("Training Clip Staging requires at least one connected AUDIO input")
 
         dataset_slug = _slugify(dataset_name)
         input_root = folder_paths.get_input_directory()
@@ -172,7 +172,7 @@ class MossClipStagingNode(BaseTTSNode):
                 import shutil
                 shutil.rmtree(dataset_dir)
             else:
-                raise FileExistsError(f"MOSS staged clip folder already exists: {dataset_dir}")
+                raise FileExistsError(f"Staged clip folder already exists: {dataset_dir}")
         os.makedirs(dataset_dir, exist_ok=True)
 
         clips: List[Dict[str, object]] = []
@@ -201,18 +201,18 @@ class MossClipStagingNode(BaseTTSNode):
                 })
 
         if not clips:
-            raise RuntimeError("MOSS Clip Staging produced no clips")
+            raise RuntimeError("Training Clip Staging produced no clips")
 
         dataset = {
-            "type": "moss_clip_dataset",
+            "type": "training_clip_dataset",
             "dataset_name": dataset_name,
             "dataset_dir": dataset_dir,
             "clips": clips,
         }
-        info = f"MOSS clip dataset ready: {dataset_name} | {len(clips)} clips"
+        info = f"Training clip dataset ready: {dataset_name} | {len(clips)} clips"
         print(f"🎞️ {info}")
         return dataset, info
 
 
 NODE_CLASS_MAPPINGS = {"MossClipStagingNode": MossClipStagingNode}
-NODE_DISPLAY_NAME_MAPPINGS = {"MossClipStagingNode": "🎞️ MOSS Clip Staging"}
+NODE_DISPLAY_NAME_MAPPINGS = {"MossClipStagingNode": "🎞️ Training Clip Staging"}

nodes/training/moss_dataset_prep_node.py

@@ -41,9 +41,9 @@ class MossDatasetPrepNode(BaseTTSNode):
                 "dataset_source": ("STRING", {
                     "default": "",
                     "tooltip": (
-                        "Path to the main MOSS manifest JSONL.\n"
-                        "This is your training set manifest: one JSON row per clip.\n"
-                        "In the normal workflow, connect the manifest path produced by MOSS Dataset Rows here."
+                        "Path to a MOSS manifest JSONL or a folder of paired audio and transcript files.\n"
+                        "For folders, use matching names such as clip001.wav + clip001.txt.\n"
+                        "In the node workflow, connect the manifest path produced by MOSS Dataset Rows here."
                     )
                 }),
             },
@@ -104,6 +104,13 @@ class MossDatasetPrepNode(BaseTTSNode):
                     "default": True,
                     "tooltip": "Reuse a matching prepared dataset cache instead of re-encoding audio codes every run."
                 }),
+                "recursive_folder_scan": ("BOOLEAN", {
+                    "default": False,
+                    "tooltip": (
+                        "When dataset_source or validation_source is a folder, also scan its subfolders. "
+                        "Disabled by default; direct files in the selected folder are always scanned."
+                    )
+                }),
             },
         }
 

nodes/training/moss_dataset_rows_node.py

@@ -58,8 +58,8 @@ class MossDatasetRowsNode(BaseTTSNode):
     def INPUT_TYPES(cls):
         return {
             "required": {
-                "clip_dataset": ("MOSS_CLIP_DATASET", {
-                    "tooltip": "Staged clip dataset from MOSS Clip Staging."
+                "clip_dataset": ("TRAINING_CLIP_DATASET", {
+                    "tooltip": "Staged clip dataset from Training Clip Staging."
                 }),
                 "manifest_name": ("STRING", {
                     "default": "moss_train.jsonl",
@@ -224,8 +224,11 @@ class MossDatasetRowsNode(BaseTTSNode):
         output_subdir: str = "",
         overwrite: bool = True,
     ):
-        if not isinstance(clip_dataset, dict) or clip_dataset.get("type") != "moss_clip_dataset":
-            raise ValueError("clip_dataset must be a MOSS_CLIP_DATASET payload from MOSS Clip Staging")
+        if not isinstance(clip_dataset, dict) or clip_dataset.get("type") not in {
+            "training_clip_dataset",
+            "moss_clip_dataset",
+        }:
+            raise ValueError("clip_dataset must come from Training Clip Staging")
 
         clips = clip_dataset.get("clips") or []
         if not clips:

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "tts_audio_suite"
 description = "TTS Audio Suite - Universal multi-engine TTS extension for ComfyUI with unified architecture supporting IndexTTS-2, ChatterBox, Chatterbox Multilingual TTS (Official 23-Lang), F5-TTS, Higgs Audio 2, VibeVoice, and RVC engines. It has character voice management, SRT subtitle TTS support, and audio processing capabilities."
-version = "5.6.0"
+version = "5.6.5"
 license = {file = "LICENSE"}
 
 [project.urls]

tests/unit/test_installer_runtime_repairs.py

@@ -18,6 +18,25 @@ from engines.dots_tts.dots_tts_engine import DotsTTSEngine
 from utils.runtimes.profiles import get_runtime_profile
 
 
+@pytest.mark.unit
+def test_nested_module_check_does_not_import_parent(tmp_path, monkeypatch):
+    package = tmp_path / "explosive_parent"
+    package.mkdir()
+    (package / "__init__.py").write_text(
+        "raise RuntimeError('parent package must not be imported')\n",
+        encoding="utf-8",
+    )
+    (package / "runtime.py").write_text("", encoding="utf-8")
+    monkeypatch.syspath_prepend(str(tmp_path))
+    monkeypatch.delitem(sys.modules, "explosive_parent", raising=False)
+
+    installer = INSTALL_MODULE.TTSAudioInstaller()
+
+    assert installer.module_available("explosive_parent.runtime") is True
+    assert installer.module_available("explosive_parent.missing") is False
+    assert "explosive_parent" not in sys.modules
+
+
 @pytest.mark.unit
 def test_fish_source_restore_accepts_namespace_package(tmp_path, monkeypatch):
     source_root = tmp_path / "source"

utils/audio/cache.py

@@ -598,6 +598,9 @@ class DramaBoxCacheKeyGenerator(CacheKeyGenerator):
             'transformer_quantization': params.get('transformer_quantization', 'none'),
             'memory_mode': params.get('memory_mode', 'fast'),
             'compile_model': bool(params.get('compile_model', False)),
+            'lora_path': params.get('lora_path', ''),
+            'lora_strength': round(float(params.get('lora_strength', 1.0)), 4),
+            'lora_revision': params.get('lora_revision', ''),
             'seed': params.get('seed', 42),
             'character': params.get('character', 'narrator'),
             'engine': 'dramabox',

utils/models/engine_registry.py

@@ -162,6 +162,8 @@ ENGINE_REGISTRY: Dict[str, EngineCapabilities] = {
         multilingual_model_switching=False,
         can_corrupt_on_reload=False,
         fallback_languages=["English"],
+        supports_training=True,
+        training_modes=["audio_lora"],
         supports_runtime_isolation=False,
     ),