# Copyright (c) 2026 DeepBeepMeep - part of WanGP, WanGP Community License 2.0 (see LICENSE.txt)
import os
import re
import warnings

import torch

from postprocessing.seedvc_processor import assets as seedvc_assets, component as seedvc_component
from shared.deepy.assets import whisper_download_def, whisper_model_name
from shared.utils import files_locator as fl
from shared.utils.hf import build_hf_url

from .ltx2_handler import _GEMMA_FILENAME, _GEMMA_FOLDER, _GEMMA_QUANTO_FILENAME
from .prompt_enhancer import DRAMABOX_DIALOGUE_PROMPT, DRAMABOX_SPEECH_PROMPT, SCENEMA_DIALOGUE_PROMPT, SCENEMA_SPEECH_PROMPT


SCENEMA_REPO_ID = "DeepBeepMeep/LTX-2"
SCENEMA_ASSET_DIR = ""
SCENEMA_MAIN_FILENAME = "scenema-audio-transformer_bf16.safetensors"
SCENEMA_QUANT_FILENAME = "scenema-audio-transformer_quanto_bf16_int8.safetensors"
LTX23_AUDIO_VAE_FILENAME = "ltx-2.3-22b_audio_vae.safetensors"
LTX23_VOCODER_FILENAME = "ltx-2.3-22b_vocoder.safetensors"
LTX23_TEXT_EMBEDDING_PROJECTION_FILENAME = "ltx-2.3-22b_text_embedding_projection.safetensors"
LTX23_EMBEDDINGS_CONNECTOR_FILENAME = "ltx-2.3-22b_embeddings_connector.safetensors"
SCENEMA_KOKORO_DIR = "kokoro"
SCENEMA_KOKORO_VOICE_DIR = "kokoro/voices"
SCENEMA_KOKORO_FILES = ["config.json", "kokoro-v1_0.pth"]
SCENEMA_KOKORO_VOICE_FILES = ["af_heart.pt"]
SCENEMA_KOKORO_G2P_DIR = "openphonemizer"
SCENEMA_KOKORO_G2P_FILES = ["openphonemizer_fp32.safetensors", "openphonemizer_config.json", "openphonemizer_dict.json"]
SCENEMA_DEFAULT_PACE = 1.5
SCENEMA_DEFAULT_DURATION_SECONDS = 120
SCENEMA_MAX_DURATION_SECONDS = 30 * 60
DRAMABOX_DEFAULT_DURATION_SECONDS = 0
DRAMABOX_MAX_DURATION_SECONDS = 60
DRAMABOX_DEFAULT_NEGATIVE_PROMPT = "worst quality, inconsistent, robotic, distorted, noise, static, muffled, unclear, unnatural, monotone"
SCENEMA_DEFAULT_CUSTOM_SETTINGS = {
    "vc_steps": 25,
    "vc_cfg_rate": 0.5,
    "pace": SCENEMA_DEFAULT_PACE,
}
DRAMABOX_DEFAULT_CUSTOM_SETTINGS = {
    "duration_multiplier": 1.1,
}
SCENEMA_CUSTOM_SETTINGS = [
    {
        "id": "vc_steps",
        "label": "SeedVC Steps (default 25)",
        "name": "SeedVC Steps",
        "type": "int",
        "default": SCENEMA_DEFAULT_CUSTOM_SETTINGS["vc_steps"],
    },
    {
        "id": "vc_cfg_rate",
        "label": "SeedVC CFG Rate (default 0.5)",
        "name": "SeedVC CFG Rate",
        "type": "float",
        "default": SCENEMA_DEFAULT_CUSTOM_SETTINGS["vc_cfg_rate"],
    },
    {
        "id": "pace",
        "label": "Pace (default 1.5)",
        "name": "Pace",
        "type": "float",
        "default": SCENEMA_DEFAULT_CUSTOM_SETTINGS["pace"],
        "min": 0.2,
        "max": 3.0,
    },
]
DRAMABOX_CUSTOM_SETTINGS = [
    {
        "id": "duration_multiplier",
        "label": "Auto Duration Multiplier (default 1.1)",
        "name": "Auto Duration Multiplier",
        "type": "float",
        "default": DRAMABOX_DEFAULT_CUSTOM_SETTINGS["duration_multiplier"],
        "min": 0.5,
        "max": 3.0,
    },
]
LTX_AUDIO_TTS_TOKENIZER_FILES = [
    "added_tokens.json",
    "chat_template.json",
    "config_light.json",
    "generation_config.json",
    "preprocessor_config.json",
    "processor_config.json",
    "special_tokens_map.json",
    "tokenizer.json",
    "tokenizer.model",
    "tokenizer_config.json",
]
SCENEMA_INFOS = """
## WanGP Speech Format
For one speaker, write normal text and put one performance cue in square brackets before each sentence. WanGP converts each `[cue]` into a Scenema `<action>` block.

```text
[Soft, close to the microphone] The lights are already on, so we can start whenever you are ready.
[Brighter, reassuring] Take your time. I am not going anywhere.
```

Plain speech uses only the text and action cues. Use Scenema XML when a single-speaker prompt also needs explicit `voice`, `gender`, `scene`, `shot`, or `language` properties.

## WanGP Dialogue Format

Use `Speaker N:` blocks for dialogue. Every new speaker block becomes its own audio chunk.

```text
Speaker 1{voice="A tired older man, low gravelly voice", gender="male", scene="a quiet kitchen at night"}:
[Trying to stay calm] I left the light on because I knew you would come back.

Speaker 2{voice="A younger woman, controlled but shaken", gender="female"}:
[A guarded breath before speaking] You always say that like it fixes everything.

Speaker 1:
[Softer, almost apologetic] No. But it keeps the room from feeling empty.
```

Properties set in `{...}` are remembered for that speaker and reused when the same speaker appears again. Add a new `{...}` later to override them.

## Supported Properties

- `voice`: the main voice and delivery description. Include age, timbre, accent, energy, microphone distance, or emotional color here.
- `gender`: `male` or `female`. It complements `voice`; for stronger control, also mention the gender in `voice`.
- `scene`: acoustic or cinematic context, such as `a quiet office late at night`.
- `shot`: `closeup`, `wide`, or `scene`. `closeup` is best for speech-first TTS.
- `language`: language code such as `en`, `fr`, or `it`.
- `speaker`: XML-only speaker id, used as `<speak speaker="2">...`.

## Voice References

The voice dropdown uses SeedVC for references. `Speaker 1 reference using SeedVC` applies the first audio reference. `Two Speakers references using SeedVC` applies the first reference to Speaker 1 and the second reference to Speaker 2, and `Three Speakers references using SeedVC` adds the third reference for Speaker 3. Additional speakers are supported, but only the first three can use uploaded reference audio.
"""

DRAMABOX_INFOS = """
## DramaBox Prompt Format

Write a scene prompt with spoken dialogue in double quotes and performance or sound cues outside the quotes.

```text
A woman speaks tenderly, "It has been a long day, my love." She whispers, "Close your eyes. I am right here." She hums quietly, "Mmmm-mmm. Sleep now."
```

DramaBox does not use square-bracket action cues or Scenema XML. Start each segment with a compact speaker voice or delivery description, then quoted dialogue on the same line. Put physical actions, pauses, sighs, and scene reactions after a quoted line or between quoted lines, but do not write standalone action-only, quote-only, or split description/quote segment lines.

## Multiline Speech

For one speaker, every non-empty line is generated as a separate segment. Later lines can reuse the first generated segment as an internal voice reference.

```text
A tired woman speaks close to the microphone, "I waited until the hallway went quiet."
A frightened whisper catches in her throat, "That was when I heard the lock turn."
Her voice breaks into a brittle laugh, "Hahaha, I should have left when I had the chance."
```

## Dialogue

Use `Speaker N:` blocks for dialogue. Any number of speakers is supported. Speakers 1 to 3 can use uploaded voice references; other speakers reuse the last 10 seconds of their first generated segment as their reference for later segments.

```text
Speaker 1:
An older detective speaks in a low tired voice, "You came back after midnight."
His low voice sharpens with quiet urgency, "That means you found something."

Speaker 2:
A young woman replies in a controlled but shaken voice, "I found the room they kept hidden."

Speaker 1:
His voice drops to a tense hush, "Then we do not have much time."
```

Use the voice reference mode to condition on a short reference clip. DramaBox uses a fixed 10 second reference budget internally.

Enable `Remove Unexpected Words` to trim generated words at the beginning of each segment when Whisper alignment can match them against the text inside double quotes. Segments without complete double quotes are left unchanged.
"""


SCENEMA_OVERVIEW_INFOS = """Generate expressive speech from text with performance cues, or from Scenema `<speak>` XML, including dialogues with any number of speakers. Put a cue such as `[Soft, close to the microphone]` before each sentence, and define each speaker's voice with properties such as `Speaker 1{voice="...", gender="female"}:`.

- **Text or <speak> XML:** voices come only from the prompt descriptions.
- **Speaker 1 / Two Speakers / Three Speakers references using SeedVC:** after generation, SeedVC converts the speech of Speakers 1 to 3 to the matching uploaded reference voice (upload order). Other speakers keep their generated voice, stabilized across their turns. **SeedVC Steps** and **SeedVC CFG Rate** trade conversion speed against fidelity to the reference.

**Pace** scales the time planned for each chunk: higher values give slower, more spacious delivery. Max Duration caps the output, and Early Stop returns the chunks already completed. See the prompt help for the full cue, property and XML syntax."""

SCENEMA_DEEPY_INFOS = "Speech or dialogue from `prompt`: `[cue]` before sentences, `Speaker N{voice=\"...\", gender=\"male\"}:` blocks, or Scenema `<speak>` XML. `audio_prompt_type`: empty = voices from descriptions; `A2` / `AB2` / `ABD2` convert Speakers 1 / 1-2 / 1-3 to `audio_guide` / `audio_guide2` / `audio_guide3` with SeedVC. `custom_settings.pace` (default 1.5) slows or speeds delivery. Max duration caps output."

SCENEMA_DEEPY_PROMPT_INFOS = "One performance cue in square brackets before each sentence, e.g. '[Soft, close to the microphone] The lights are on.' Dialogue: 'Speaker 1{voice=\"A tired older man\", gender=\"male\"}:' then cued lines; properties persist for that speaker. Keep spoken words outside brackets."

DRAMABOX_OVERVIEW_INFOS = """Generate dramatic, expressive speech, including whispers, laughter and emotional reactions, from scene-style prompts: a compact voice or delivery description followed by the spoken words in double quotes. Every line becomes its own audio segment, and `Speaker N:` blocks create dialogues with any number of speakers.

- **Text prompt:** voices come from the descriptions; a speaker's later segments reuse their first generated segment as a voice reference.
- **Speaker 1 / 1 and 2 / 1, 2 and 3 voice references:** Speakers 1 to 3 follow the matching uploaded voice (upload order), within a 10-second reference budget.
- **Remove Unexpected Words:** DramaBox sometimes adds invented words at the start of a segment; this option uses Whisper to trim them when the segment's quoted text can be matched.

**Target Duration** of 0 plans each segment from its text; **Auto Duration Multiplier** stretches those estimates. CFG, STG and Guidance Rescale control how strictly the audio follows the prompt, and the negative prompt lists sounds to avoid. Early Stop returns the segments already completed."""

DRAMABOX_DEEPY_INFOS = "Expressive speech from `prompt` lines: voice/delivery description then quoted words; `Speaker N:` blocks for dialogue. `audio_prompt_type`: empty = voices from descriptions; `A` / `AB` / `ABD` condition Speakers 1 / 1-2 / 1-3 on `audio_guide` / `audio_guide2` / `audio_guide3`; add `0` (Remove Unexpected Words) to trim invented words at segment starts. `duration_seconds` 0 = automatic."

DRAMABOX_DEEPY_PROMPT_INFOS = "Each line: compact voice/delivery description, then the words in double quotes, e.g. 'A tired woman speaks close to the microphone, \"I waited until the hallway went quiet.\"' Actions and sounds go outside the quotes, never alone on a line. Dialogue: 'Speaker 1:' block headers."


def _get_scenema_model_def():
    return {
        "audio_only": True,
        "image_outputs": False,
        "sliding_window": False,
        "guidance_max_phases": 0,
        "no_negative_prompt": True,
        "inference_steps": False,
        "temperature": False,
        "lock_inference_steps": True,
        "image_prompt_types_allowed": "",
        "supports_early_stop": True,
        "profiles_dir": ["scenema_audio"],
        "duration_slider": {
            "label": "Max Duration (seconds)",
            "name": "Max Duration",
            "min": 1,
            "max": SCENEMA_MAX_DURATION_SECONDS,
            "increment": 0.5,
            "default": SCENEMA_DEFAULT_DURATION_SECONDS,
        },
        "profile_type": "video",
        "preserve_empty_prompt_lines": True,
        "any_audio_prompt": True,
        "audio_prompt_choices": True,
        "audio_prompt_type_sources": {
            "selection": ["", "A2", "AB2", "ABD2"],
            "labels": {
                "": "Text or <speak> XML",
                "A2": "Speaker 1 reference using SeedVC",
                "AB2": "Two Speakers references using SeedVC",
                "ABD2": "Three Speakers references using SeedVC",
            },
            "letters_filter": "ABD2",
            "custom_flags": {"2": "SeedVC"},
            "default": "",
        },
        "audio_guide_label": "Speaker 1 reference voice (optional for multi-speaker)",
        "audio_guide2_label": "Speaker 2 reference voice (optional)",
        "audio_guide3_label": "Speaker 3 reference voice (optional)",
        "custom_settings": [one.copy() for one in SCENEMA_CUSTOM_SETTINGS],
        "infos": SCENEMA_OVERVIEW_INFOS,
        "prompt_infos": SCENEMA_INFOS,
        "deepy_infos": SCENEMA_DEEPY_INFOS,
        "deepy_prompt_infos": SCENEMA_DEEPY_PROMPT_INFOS,
        "prompt_description": "Speech text or Scenema <speak> XML",
        "text_prompt_enhancer_instructions": SCENEMA_SPEECH_PROMPT,
        "text_prompt_enhancer_instructions1": SCENEMA_DIALOGUE_PROMPT,
        "text_prompt_enhancer_max_tokens": 768,
        "text_prompt_enhancer_max_tokens1": 1024,
        "prompt_enhancer_def": {
            "selection": ["T", "T1"],
            "labels": {
                "T": "A Speech with Action Cues",
                "T1": "A Dialogue with Action Cues",
            },
            "default": "T",
        },
        "prompt_enhancer_button_label": "Write",
        "compile": False,
        "text_encoder_folder": _GEMMA_FOLDER,
        "text_encoder_URLs": [
            build_hf_url("DeepBeepMeep/LTX-2", _GEMMA_FOLDER, _GEMMA_FILENAME),
            build_hf_url("DeepBeepMeep/LTX-2", _GEMMA_FOLDER, _GEMMA_QUANTO_FILENAME),
        ],
        "dtype": "bf16",
    }


def _get_dramabox_model_def():
    return {
        "audio_only": True,
        "image_outputs": False,
        "sliding_window": False,
        "guidance_max_phases": 0,
        "no_negative_prompt": False,
        "inference_steps": True,
        "alt_scale": "Guidance Rescale",
        "temperature": False,
        "image_prompt_types_allowed": "",
        "supports_early_stop": True,
        "profiles_dir": ["dramabox_audio"],
        "duration_slider": {
            "label": "Target Duration (seconds, 0 = auto)",
            "name": "Target Duration",
            "min": 0,
            "max": DRAMABOX_MAX_DURATION_SECONDS,
            "increment": 0.5,
            "default": DRAMABOX_DEFAULT_DURATION_SECONDS,
        },
        "profile_type": "video",
        "preserve_empty_prompt_lines": True,
        "any_audio_prompt": True,
        "audio_prompt_choices": True,
        "audio_prompt_type_sources": {
            "selection": ["", "A", "AB", "ABD"],
            "labels": {
                "": "Text prompt",
                "A": "Speaker 1 voice reference",
                "AB": "Speaker 1 and 2 voice references",
                "ABD": "Speaker 1, 2 and 3 voice references",
            },
            "letters_filter": "ABD",
            "default": "",
        },
        "audio_prompt_type_custom_option": {"label": "Remove Unexpected Words", "flag": "0"},
        "audio_guide_label": "Speaker 1 reference voice (optional)",
        "audio_guide2_label": "Speaker 2 reference voice (optional)",
        "audio_guide3_label": "Speaker 3 reference voice (optional)",
        "custom_settings": [one.copy() for one in DRAMABOX_CUSTOM_SETTINGS],
        "infos": DRAMABOX_OVERVIEW_INFOS,
        "prompt_infos": DRAMABOX_INFOS,
        "deepy_infos": DRAMABOX_DEEPY_INFOS,
        "deepy_prompt_infos": DRAMABOX_DEEPY_PROMPT_INFOS,
        "prompt_description": "DramaBox scene prompt",
        "text_prompt_enhancer_instructions": DRAMABOX_SPEECH_PROMPT,
        "text_prompt_enhancer_instructions1": DRAMABOX_DIALOGUE_PROMPT,
        "text_prompt_enhancer_max_tokens": 768,
        "text_prompt_enhancer_max_tokens1": 1024,
        "prompt_enhancer_def": {
            "selection": ["T", "T1"],
            "labels": {
                "T": "A Speech",
                "T1": "A Dialogue",
            },
            "default": "T",
        },
        "prompt_enhancer_button_label": "Write",
        "compile": False,
        "text_encoder_folder": _GEMMA_FOLDER,
        "text_encoder_URLs": [
            build_hf_url("DeepBeepMeep/LTX-2", _GEMMA_FOLDER, _GEMMA_FILENAME),
            build_hf_url("DeepBeepMeep/LTX-2", _GEMMA_FOLDER, _GEMMA_QUANTO_FILENAME),
        ],
        "dtype": "bf16",
    }


def _get_scenema_download_def():
    return [
        {
            "repoId": SCENEMA_REPO_ID,
            "sourceFolderList": [""],
            "fileList": [[LTX23_AUDIO_VAE_FILENAME, LTX23_VOCODER_FILENAME, LTX23_TEXT_EMBEDDING_PROJECTION_FILENAME, LTX23_EMBEDDINGS_CONNECTOR_FILENAME]],
        },
        {
            "repoId": "DeepBeepMeep/LTX-2",
            "sourceFolderList": [_GEMMA_FOLDER],
            "fileList": [LTX_AUDIO_TTS_TOKENIZER_FILES],
        },
        whisper_download_def(),
        {
            "repoId": SCENEMA_REPO_ID,
            "sourceFolderList": [SCENEMA_KOKORO_DIR, SCENEMA_KOKORO_VOICE_DIR, SCENEMA_KOKORO_G2P_DIR],
            "fileList": [SCENEMA_KOKORO_FILES, SCENEMA_KOKORO_VOICE_FILES, SCENEMA_KOKORO_G2P_FILES],
        },
    ] + seedvc_assets.query_download_def(mode=seedvc_assets.SEEDVC_MODE_SPEECH)


def _get_dramabox_download_def():
    return [
        {
            "repoId": SCENEMA_REPO_ID,
            "sourceFolderList": [""],
            "fileList": [[LTX23_AUDIO_VAE_FILENAME, LTX23_VOCODER_FILENAME, LTX23_TEXT_EMBEDDING_PROJECTION_FILENAME, LTX23_EMBEDDINGS_CONNECTOR_FILENAME]],
        },
        {
            "repoId": "DeepBeepMeep/LTX-2",
            "sourceFolderList": [_GEMMA_FOLDER],
            "fileList": [LTX_AUDIO_TTS_TOKENIZER_FILES],
        },
        whisper_download_def(),
    ]


def _load_alignment_whisper():
    from shared.deepy.transcription import load_whisper

    alignment_whisper = load_whisper(whisper_model_name(), torch.device("cpu"))
    alignment_heads = alignment_whisper.alignment_heads
    del alignment_whisper._buffers["alignment_heads"]
    object.__setattr__(alignment_whisper, "alignment_heads", alignment_heads)
    for module in alignment_whisper.modules():
        if isinstance(module, torch.nn.LayerNorm):
            module._lock_dtype = torch.float32
    alignment_whisper._offload_hooks = ["transcribe"]
    alignment_whisper._model_dtype = torch.float16
    alignment_whisper._budget = 0
    alignment_whisper.eval().requires_grad_(False)
    return alignment_whisper


def _load_kokoro_pipeline():
    from preprocessing.kokoro import KPipeline

    try:
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=FutureWarning, message=r"`torch\.nn\.utils\.weight_norm` is deprecated.*")
            kokoro_pipeline = KPipeline(lang_code="a", device="cpu", repo_id=fl.locate_folder(SCENEMA_KOKORO_DIR))
    except Exception as exc:
        raise RuntimeError(f"Kokoro TTS is required for Scenema Audio duration estimation. Error: {exc}") from exc
    kokoro_model = getattr(kokoro_pipeline, "model", None)
    if kokoro_model is None:
        raise RuntimeError("Kokoro TTS is required for Scenema Audio duration estimation.")
    kokoro_model._model_dtype = torch.float32
    return kokoro_pipeline


def _is_dramabox(model_type: str) -> bool:
    return model_type == "dramabox_audio"


class family_handler:
    @staticmethod
    def query_supported_types():
        return ["scenema_audio", "dramabox_audio"]

    @staticmethod
    def query_family_maps():
        return {}, {}

    @staticmethod
    def query_model_family():
        return "tts"

    @staticmethod
    def query_family_infos():
        return {"tts": (2200, "TTS")}

    @staticmethod
    def get_lora_dir(base_model_type):
        if _is_dramabox(base_model_type):
            return "dramabox_audio"
        return "scenema_audio"

    @staticmethod
    def query_model_def(base_model_type, model_def):
        from models.help_shared import specialities
        return {**(_get_dramabox_model_def() if _is_dramabox(base_model_type) else _get_scenema_model_def()), "device_explicit": True, "specialities": specialities("voice design", "voice cloning", "speech dialogue")}

    @staticmethod
    def query_model_files(computeList, base_model_type, model_def=None):
        return _get_dramabox_download_def() if _is_dramabox(base_model_type) else _get_scenema_download_def()

    @staticmethod
    def validate_generative_settings(base_model_type, model_def, inputs):
        if _is_dramabox(base_model_type):
            inputs.setdefault("num_inference_steps", 30)
            inputs.setdefault("guidance_scale", 2.5)
            inputs.setdefault("audio_guidance_scale", 1.5)
            inputs.setdefault("alt_scale", 0.0)
            custom_settings = inputs.get("custom_settings", None)
            if not isinstance(custom_settings, dict):
                custom_settings = {}
            for key, value in DRAMABOX_DEFAULT_CUSTOM_SETTINGS.items():
                custom_settings.setdefault(key, value)
            try:
                if int(inputs.get("num_inference_steps", 30)) <= 0:
                    return "DramaBox Audio inference steps must be greater than 0."
            except Exception:
                return "DramaBox Audio inference steps must be an integer."
            try:
                if float(inputs.get("guidance_scale", 2.5)) < 1.0:
                    return "DramaBox Audio CFG scale must be at least 1.0."
            except Exception:
                return "DramaBox Audio CFG scale must be a number."
            try:
                if float(inputs.get("audio_guidance_scale", 1.5)) < 0.0:
                    return "DramaBox Audio STG scale must be zero or greater."
            except Exception:
                return "DramaBox Audio STG scale must be a number."
            try:
                if float(custom_settings.get("duration_multiplier", DRAMABOX_DEFAULT_CUSTOM_SETTINGS["duration_multiplier"])) <= 0:
                    return "DramaBox Audio duration multiplier must be greater than 0."
            except Exception:
                return "DramaBox Audio duration multiplier must be a number."
            try:
                alt_scale = float(inputs.get("alt_scale", 0.0))
                if alt_scale < 0 or alt_scale > 1:
                    return "DramaBox Audio guidance rescale must be between 0 and 1."
            except Exception:
                return "DramaBox Audio guidance rescale must be a number."
            audio_prompt_type = str(inputs.get("audio_prompt_type", "") or "").upper()
            if "A" in audio_prompt_type and inputs.get("audio_guide") is None:
                return "DramaBox Audio Speaker 1 reference mode requires a reference audio file."
            if "B" in audio_prompt_type and inputs.get("audio_guide2") is None:
                return "DramaBox Audio Speaker 2 reference mode requires a second reference audio file."
            if "B" in audio_prompt_type and not re.search(r"(?im)^\s*Speaker\s*2\s*(?:\{[^\n{}]*\})?\s*:", str(inputs.get("prompt", "") or "")):
                return "DramaBox Audio two-reference mode requires a Speaker 2: block."
            if "D" in audio_prompt_type and inputs.get("audio_guide3") is None:
                return "DramaBox Audio Speaker 3 reference mode requires a third reference audio file."
            if "D" in audio_prompt_type and not re.search(r"(?im)^\s*Speaker\s*3\s*(?:\{[^\n{}]*\})?\s*:", str(inputs.get("prompt", "") or "")):
                return "DramaBox Audio three-reference mode requires a Speaker 3: block."
            custom_settings = {"duration_multiplier": float(custom_settings.get("duration_multiplier", DRAMABOX_DEFAULT_CUSTOM_SETTINGS["duration_multiplier"]))}
            inputs["alt_scale"] = alt_scale
            inputs["custom_settings"] = custom_settings
            return None

        inputs.update(
            {
                "num_inference_steps": 8,
                "guidance_scale": 1.0,
                "audio_guidance_scale": 1.0,
                "audio_cfg_scale": 1.0,
                "alt_guidance_scale": 1.0,
                "alt_scale": 0.0,
            }
        )
        custom_settings = inputs.get("custom_settings", None)
        if isinstance(custom_settings, dict):
            vc_steps = custom_settings.get("vc_steps", SCENEMA_DEFAULT_CUSTOM_SETTINGS["vc_steps"])
            vc_cfg_rate = custom_settings.get("vc_cfg_rate", SCENEMA_DEFAULT_CUSTOM_SETTINGS["vc_cfg_rate"])
            pace = custom_settings.get("pace", SCENEMA_DEFAULT_CUSTOM_SETTINGS["pace"])
            try:
                if int(vc_steps) <= 0:
                    return "Scenema Audio SeedVC Steps must be greater than 0."
            except Exception:
                return "Scenema Audio SeedVC Steps must be an integer."
            try:
                if float(vc_cfg_rate) <= 0:
                    return "Scenema Audio SeedVC CFG Rate must be greater than 0."
            except Exception:
                return "Scenema Audio SeedVC CFG Rate must be a number."
            try:
                if float(pace) <= 0:
                    return "Scenema Audio pace must be greater than 0."
            except Exception:
                return "Scenema Audio pace must be a number."
            custom_settings["vc_steps"] = int(vc_steps)
            custom_settings["vc_cfg_rate"] = float(vc_cfg_rate)
            custom_settings["pace"] = float(pace)
            inputs["custom_settings"] = custom_settings
        return None

    @staticmethod
    def validate_generative_prompt(base_model_type, model_def, inputs, one_prompt):
        if _is_dramabox(base_model_type):
            if one_prompt is None or len(str(one_prompt).strip()) == 0:
                return "Prompt text cannot be empty for DramaBox Audio."
            audio_prompt_type = str(inputs.get("audio_prompt_type", "") or "").upper()
            if "A" in audio_prompt_type and inputs.get("audio_guide") is None:
                return "DramaBox Audio reference voice mode requires a reference audio file."
            return None

        if one_prompt is None or len(str(one_prompt).strip()) == 0:
            return "Prompt text cannot be empty for Scenema Audio."
        audio_prompt_type = str(inputs.get("audio_prompt_type", "") or "").upper()
        if "A" in audio_prompt_type and "B" not in audio_prompt_type and inputs.get("audio_guide") is None:
            return "Scenema Audio reference voice mode requires a reference audio file."
        if "B" in audio_prompt_type and re.search(r"(?is)(Speaker\s*\d+|<speak[^>]*\bspeaker\s*=)", str(one_prompt)) is None:
            return "Scenema Audio multi-speaker mode requires SpeakerN text or <speak speaker=\"N\"> XML blocks."
        return None

    @staticmethod
    def load_model(
        model_filename,
        model_type,
        base_model_type,
        model_def,
        quantizeTransformer=False,
        text_encoder_quantization=None,
        dtype=torch.bfloat16,
        VAE_dtype=torch.float32,
        mixed_precision_transformer=False,
        save_quantized=False,
        submodel_no_list=None,
        text_encoder_filename=None,
        profile=0,
        **kwargs,
    ):
        if _is_dramabox(base_model_type):
            from .dramabox_audio import DramaBoxAudioPipeline

            weights_path = model_filename[0] if isinstance(model_filename, (list, tuple)) else model_filename
            if not text_encoder_filename:
                raise ValueError("DramaBox Audio requires the LTX2 Gemma text encoder.")
            audio_vae_path = fl.locate_file(LTX23_AUDIO_VAE_FILENAME)
            vocoder_path = fl.locate_file(LTX23_VOCODER_FILENAME)
            text_projection_path = fl.locate_file(LTX23_TEXT_EMBEDDING_PROJECTION_FILENAME)
            text_connector_path = fl.locate_file(LTX23_EMBEDDINGS_CONNECTOR_FILENAME)
            config_path = os.path.join(os.path.dirname(__file__), "configs", "ltx2_22b_config.json")
            pipeline = DramaBoxAudioPipeline(
                model_weights_path=weights_path,
                gemma_path=text_encoder_filename,
                audio_vae_path=audio_vae_path,
                vocoder_path=vocoder_path,
                text_projection_path=text_projection_path,
                text_connector_path=text_connector_path,
                config_path=config_path,
                device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
                dtype=dtype or torch.bfloat16,
            )
            pipe = {
                "pipe": {
                    "transformer": pipeline.model,
                    "text_encoder": pipeline.text_encoder,
                    "text_embedding_projection": pipeline.text_embedding_projection,
                    "text_embeddings_connector": pipeline.text_embeddings_connector,
                    "audio_encoder": pipeline.audio_encoder,
                    "audio_decoder": pipeline.audio_decoder,
                    "vocoder": pipeline.vocoder,
                }
            }
            if save_quantized and weights_path:
                from wgp import save_quantized_model

                quantized_transformer = getattr(pipeline.model, "velocity_model", pipeline.model)
                save_quantized_model(quantized_transformer, model_type, weights_path, dtype or torch.bfloat16, config_path)
            return pipeline, pipe

        from .scenema_audio import ScenemaAudioPipeline
        seedvc_component.ensure_installed()
        from postprocessing import seedvc

        weights_path = model_filename[0] if isinstance(model_filename, (list, tuple)) else model_filename
        if not text_encoder_filename:
            raise ValueError("Scenema Audio requires the LTX2 Gemma text encoder.")
        audio_vae_path = fl.locate_file(LTX23_AUDIO_VAE_FILENAME)
        vocoder_path = fl.locate_file(LTX23_VOCODER_FILENAME)
        text_projection_path = fl.locate_file(LTX23_TEXT_EMBEDDING_PROJECTION_FILENAME)
        text_connector_path = fl.locate_file(LTX23_EMBEDDINGS_CONNECTOR_FILENAME)
        config_path = os.path.join(os.path.dirname(__file__), "configs", "ltx2_22b_config.json")
        alignment_whisper = _load_alignment_whisper()
        kokoro_pipeline = _load_kokoro_pipeline()
        kokoro_model = kokoro_pipeline.model
        seedvc_model = seedvc.get_model(dtype=torch.float16, mode=seedvc_assets.SEEDVC_MODE_SPEECH)
        seedvc_pipe = seedvc.get_pipe(profile_no=profile, model=seedvc_model, mode=seedvc_assets.SEEDVC_MODE_SPEECH)

        pipeline = ScenemaAudioPipeline(
            model_weights_path=weights_path,
            audio_vae_path=audio_vae_path,
            vocoder_path=vocoder_path,
            text_projection_path=text_projection_path,
            text_connector_path=text_connector_path,
            gemma_path=text_encoder_filename,
            config_path=config_path,
            alignment_whisper=alignment_whisper,
            kokoro_pipeline=kokoro_pipeline,
            seedvc=seedvc_model,
            device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
            dtype=dtype or torch.bfloat16,
        )
        pipe = {
            "transformer": pipeline.model,
            "text_encoder": pipeline.text_encoder,
            "text_embedding_projection": pipeline.text_embedding_projection,
            "text_embeddings_connector": pipeline.text_embeddings_connector,
            "audio_encoder": pipeline.audio_encoder,
            "audio_decoder": pipeline.audio_decoder,
            "vocoder": pipeline.vocoder,
            "alignment_whisper": alignment_whisper,
            "kokoro": kokoro_model,
        }
        pipe.update(seedvc_pipe)
        pipe = {"pipe": pipe, "coTenantsMap": seedvc.get_cotenants_map(seedvc_pipe)}

        if save_quantized and weights_path:
            from wgp import save_quantized_model

            quantized_transformer = getattr(pipeline.model, "velocity_model", pipeline.model)
            save_quantized_model(quantized_transformer, model_type, weights_path, dtype or torch.bfloat16, config_path)

        return pipeline, pipe

    @staticmethod
    def fix_settings(base_model_type, settings_version, model_def, ui_defaults):
        if _is_dramabox(base_model_type):
            audio_prompt_type = str(ui_defaults.get("audio_prompt_type", "") or "").upper()
            ui_defaults["audio_prompt_type"] = ("ABD" if all(flag in audio_prompt_type for flag in "ABD") else "AB" if "B" in audio_prompt_type and "A" in audio_prompt_type else "A" if "A" in audio_prompt_type else "") + "".join(flag for flag in "0VN" if flag in audio_prompt_type)  # keep wgp audio options
            ui_defaults["alt_prompt"] = ""
            ui_defaults.setdefault("duration_seconds", DRAMABOX_DEFAULT_DURATION_SECONDS)
            ui_defaults.setdefault("num_inference_steps", 30)
            ui_defaults.setdefault("guidance_scale", 2.5)
            ui_defaults.setdefault("audio_guidance_scale", 1.5)
            ui_defaults.setdefault("negative_prompt", DRAMABOX_DEFAULT_NEGATIVE_PROMPT)
            custom_settings = ui_defaults.get("custom_settings", None)
            if not isinstance(custom_settings, dict):
                custom_settings = {}
            if "alt_scale" not in ui_defaults and "rescale_scale" in custom_settings:
                ui_defaults["alt_scale"] = custom_settings["rescale_scale"]
            ui_defaults.setdefault("alt_scale", 0.0)
            custom_settings = {"duration_multiplier": custom_settings.get("duration_multiplier", DRAMABOX_DEFAULT_CUSTOM_SETTINGS["duration_multiplier"])}
            ui_defaults["custom_settings"] = custom_settings
            return

        audio_prompt_type = str(ui_defaults.get("audio_prompt_type", "") or "").upper()
        ui_defaults["audio_prompt_type"] = ("ABD2" if "2" in audio_prompt_type and "D" in audio_prompt_type else "AB2" if "2" in audio_prompt_type and "B" in audio_prompt_type else "A2" if "2" in audio_prompt_type and "A" in audio_prompt_type else "") + "".join(flag for flag in "VN" if flag in audio_prompt_type)  # keep wgp audio options
        ui_defaults["alt_prompt"] = ""
        ui_defaults.setdefault("duration_seconds", model_def.get("duration_slider", {}).get("default", SCENEMA_DEFAULT_DURATION_SECONDS))
        custom_settings = ui_defaults.get("custom_settings", None)
        if not isinstance(custom_settings, dict):
            custom_settings = {}
        for key, value in SCENEMA_DEFAULT_CUSTOM_SETTINGS.items():
            custom_settings.setdefault(key, value)
        ui_defaults["custom_settings"] = custom_settings

    @staticmethod
    def update_default_settings(base_model_type, model_def, ui_defaults):
        if _is_dramabox(base_model_type):
            ui_defaults.update(
                {
                    "audio_prompt_type": "0",
                    "alt_prompt": "",
                    "repeat_generation": 1,
                    "duration_seconds": DRAMABOX_DEFAULT_DURATION_SECONDS,
                    "video_length": 0,
                    "num_inference_steps": 30,
                    "negative_prompt": DRAMABOX_DEFAULT_NEGATIVE_PROMPT,
                    "guidance_scale": 2.5,
                    "audio_guidance_scale": 1.5,
                    "alt_scale": 0.0,
                    "custom_settings": dict(DRAMABOX_DEFAULT_CUSTOM_SETTINGS),
                    "multi_prompts_gen_type": "FG",
                }
            )
            return

        ui_defaults.update(
            {
                "audio_prompt_type": "",
                "alt_prompt": "",
                "repeat_generation": 1,
                "duration_seconds": model_def.get("duration_slider", {}).get("default", SCENEMA_DEFAULT_DURATION_SECONDS),
                "video_length": 0,
                "num_inference_steps": 8,
                "negative_prompt": "",
                "guidance_scale": 1.0,
                "custom_settings": dict(SCENEMA_DEFAULT_CUSTOM_SETTINGS),
                "multi_prompts_gen_type": "FG",
            }
        )
