Blame for nodes.py

"""OGG TTS as a self-contained ComfyUI pack: load the model and the Fish codec, encode a voice, generate codes, decode or
edit them, and transcribe a reference with Qwen3-ASR. Nothing needs another custom-node pack.

Values passed between nodes are plain dicts:
  OGG_VOICE = {"transcript": str, "codes": LongTensor [10, frames]}
  OGG_CODES = {"codes": LongTensor [10, frames], "lines": [{speaker, text, chunk, start, end}], "takes": [frames, ...]}
Frames are the codec's: 21.53 per second, 10 codebooks (1 semantic of 4096 + 9 residual of 1024).

The codec runs before and after generation only, so it is never resident next to the model while frames are sampled.
"""
import hashlib
import json
import logging
import os
import random
import re
import wave
import zlib

import torch
from aiohttp import web
from safetensors.torch import save_file

import comfy.audio
import comfy.model_management
import comfy.model_patcher
import comfy.utils
import folder_paths
from comfy_api.latest import ComfyExtension, io
from server import PromptServer

from .ogg import generate as gen
from .ogg import model as of
from .ogg.audio import tidy_clip, trim_pauses
from .ogg.codec import DECODE_CONTEXT, DECODE_FRAME_BYTES, DECODE_WINDOW, FRAME_LENGTH, SAMPLE_RATE, load_codec
from .ogg.editor import DEFAULTS, analyze, parse_edits, plan, render
from .ogg.qwen3_asr.asr import LANGUAGES, load_qwen3_asr, transcribe

OggModel = io.Custom("OGG_MODEL")
OggCodec = io.Custom("OGG_CODEC")
OggVoice = io.Custom("OGG_VOICE")
OggCodes = io.Custom("OGG_CODES")
OggASR = io.Custom("OGG_ASR_MODEL")

# OmniVoice's config.json and tokenizer files (not its weights) live in the pack, next to this file
OMNIVOICE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "omnivoice")
OMNIVOICE_REPO = "k2-fsa/OmniVoice"
# seconds of silence Decode puts between separately generated takes, the Editor's shortest gap between lines
TAKE_GAP = 0.3
# Fish Audio's Drama 3 Space shortens longer pauses to this many seconds
MAX_PAUSE = 0.75
SPEAKER_TAG = re.compile(r"^\s*<\|speaker:\d+\|>\s*")
# how many neighbouring lines a line hears as context (the model was trained with up to 2 each side)
CONTEXT_LINES = 2
BLOCK = 16  # frames per block, as trained (0.74 s)

VOICES_FOLDER = os.path.join(folder_paths.get_output_directory(), "ogg_voices")
CODES_FOLDER = os.path.join(folder_paths.get_output_directory(), "ogg_codes")


def register_folder(name, path):
    """`path` after any paths extra_model_paths.yaml gives, listed with the usual model file extensions."""
    folder_paths.add_model_folder_path(name, path)
    folder_paths.folder_names_and_paths[name] = (folder_paths.folder_names_and_paths[name][0], folder_paths.supported_pt_extensions)


register_folder("asr", os.path.join(folder_paths.models_dir, "asr"))
register_folder("ogg_voices", VOICES_FOLDER)
register_folder("ogg_codes", CODES_FOLDER)


def choices(folder, keyword):
    """The files of a model folder, those with `keyword` in the name first."""
    return sorted(folder_paths.get_filename_list(folder), key=lambda name: (keyword not in name.lower(), name.lower()))


# ----------------------------------------------------------------------------------------------------------------- model


class Ogg:
    """A loaded OGG TTS: the module (managed by a ComfyUI patcher) and its tokenizer."""

    def __init__(self, model, tok, patcher):
        self.model, self.tok, self.patcher = model, tok, patcher


def ensure_omnivoice_files():
    """config.json and the tokenizer files of OmniVoice in the pack's omnivoice folder; fetched once if they are missing."""
    if os.path.isfile(os.path.join(OMNIVOICE_DIR, "config.json")):
        return
    logging.info("OGG TTS: %s has no config.json; fetching OmniVoice's config and tokenizer files (no weights) from %s",
                 OMNIVOICE_DIR, OMNIVOICE_REPO)
    try:
        from huggingface_hub import snapshot_download
        snapshot_download(OMNIVOICE_REPO, allow_patterns=["*.json", "*.txt"], local_dir=OMNIVOICE_DIR)
    except Exception as error:
        raise FileNotFoundError(
            f"OGG TTS needs OmniVoice's config.json and tokenizer files in {OMNIVOICE_DIR}, and fetching them failed ({error}). "
            f'Download them by hand (not the weights): hf download {OMNIVOICE_REPO} --include "*.json" "*.txt" --local-dir "{OMNIVOICE_DIR}"') from error
    if not os.path.isfile(os.path.join(OMNIVOICE_DIR, "config.json")):
        raise FileNotFoundError(f"OGG TTS: {OMNIVOICE_REPO} has no config.json for {OMNIVOICE_DIR}")


def load_ogg(path, dtype):
    ensure_omnivoice_files()
    state = comfy.utils.load_torch_file(path)
    if "audio_embeddings.weight" not in state or "semantic_head.weight" not in state:
        raise ValueError(f"{os.path.basename(path)} is not an OGG TTS file (FashionFlora/ogg-tts model.safetensors)")
    # the released file may or may not carry the prompt cross-attention of the MTP (see omnifish.AxialMTP)
    prompt_x = any(".px." in key for key in state)
    model, tok = of.OmniFish.from_config_dir(OMNIVOICE_DIR, dtype=dtype, prompt_x=prompt_x)
    fresh = model.load_weights(state, tok)
    del state
    tok = tok if fresh else of.IdTokenizer(tok, model.text_vocab)
    model.eval()
    load_device = comfy.model_management.get_torch_device()
    offload_device = comfy.model_management.unet_offload_device()
    return Ogg(model, tok, comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=offload_device))


class OggLoader(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="OggLoader",
            display_name="Load OGG TTS",
            category="audio/ogg",
            description="OGG TTS (FashionFlora/ogg-tts model.safetensors) from models/diffusion_models. OmniVoice's config.json and "
                        "tokenizer files are read from the omnivoice folder of this pack (fetched on first use if missing).",
            inputs=[
                io.Combo.Input("model_name", options=choices("diffusion_models", "ogg"),
                               tooltip="A file of models/diffusion_models; names with 'ogg' come first."),
                io.Combo.Input("dtype", options=["float32", "bfloat16"], default="float32",
                               tooltip="How the weights are held. float32 is what the authors run (with bfloat16 autocast); bfloat16 "
                                       "halves the memory (~3 GB -> ~1.5 GB) and is the thing to try on a small card."),
            ],
            outputs=[OggModel.Output()],
        )

    @classmethod
    def execute(cls, model_name, dtype="float32") -> io.NodeOutput:
        path = folder_paths.get_full_path_or_raise("diffusion_models", model_name)
        return io.NodeOutput(load_ogg(path, torch.float32 if dtype == "float32" else torch.bfloat16))


class OggCodecLoader(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="OggCodecLoader",
            display_name="Load OGG Codec",
            category="audio/ogg",
            description="The Fish S2-Pro codec OGG TTS speaks in (drama3_codec.safetensors in models/vae): 44.1 kHz, 21.5 frames/s.",
            inputs=[io.Combo.Input("codec_name", options=choices("vae", "codec"))],
            outputs=[OggCodec.Output()],
        )

    @classmethod
    def execute(cls, codec_name) -> io.NodeOutput:
        return io.NodeOutput(load_codec(folder_paths.get_full_path_or_raise("vae", codec_name)))


# ----------------------------------------------------------------------------------------------------------------- voices


class OggEncodeVoice(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="OggEncodeVoice",
            display_name="OGG Encode Voice",
            category="audio/ogg",
            description="A voice to clone: a clip and its exact transcript. A few seconds of clean speech work best; the clip "
                        "loses its silent edges and is peak-normalized before it is encoded.",
            inputs=[
                OggCodec.Input("codec"),
                io.Audio.Input("audio"),
                io.String.Input("transcript", multiline=True,
                                tooltip="What the clip says, word for word (Qwen3-ASR Transcribe can write it)."),
            ],
            outputs=[OggVoice.Output()],
        )

    @classmethod
    def execute(cls, codec, audio, transcript) -> io.NodeOutput:
        comfy.model_management.load_models_gpu([codec.patcher])
        device = codec.patcher.load_device
        waveform = audio["waveform"][0].float().mean(0, keepdim=True)
        if audio["sample_rate"] != SAMPLE_RATE:
            waveform = comfy.audio.resample(waveform, audio["sample_rate"], SAMPLE_RATE)
        waveform = tidy_clip(waveform)
        with comfy.model_management.cuda_device_context(device):
            codes = codec.model.encode(waveform.to(device))[0].cpu()
        return io.NodeOutput({"transcript": transcript, "codes": codes})


class OggSaveVoice(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="OggSaveVoice",
            display_name="Save OGG Voice",
            category="audio/ogg",
            description="Save an encoded voice (its codes and transcript, a few KB) to output/ogg_voices, for Load OGG Voice.",
            inputs=[
                OggVoice.Input("voice"),
                io.String.Input("name", default="ogg/voice", tooltip="File name prefix; a counter is added. Subfolders are allowed."),
            ],
            is_output_node=True,
        )

    @classmethod
    def execute(cls, voice, name) -> io.NodeOutput:
        folder, filename, counter, _, _ = folder_paths.get_save_image_path(name, VOICES_FOLDER)
        save_file({"codes": voice["codes"].to(torch.int16).contiguous()}, os.path.join(folder, f"{filename}_{counter:05}_.safetensors"),
                  metadata={"transcript": voice["transcript"]})
        return io.NodeOutput()


class OggLoadVoice(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="OggLoadVoice",
            display_name="Load OGG Voice",
            category="audio/ogg",
            description="A voice saved by Save OGG Voice (output/ogg_voices). Voice files of ComfyUI-Drama3 have the same layout and load too.",
            inputs=[io.Combo.Input("voice", options=folder_paths.get_filename_list("ogg_voices"))],
            outputs=[OggVoice.Output()],
        )

    @classmethod
    def execute(cls, voice) -> io.NodeOutput:
        tensors, metadata = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("ogg_voices", voice), return_metadata=True)
        if not metadata or "transcript" not in metadata or "codes" not in tensors:
            raise ValueError(f"{voice} is not a voice file")
        return io.NodeOutput({"transcript": metadata["transcript"], "codes": tensors["codes"].long()})


# -------------------------------------------------------------------------------------------------------------- generate


# the model's trained identities (ogg/id_vocab.json): 3157 anonymous speakers spk_N, 599 recordings qN (a book or a game each),
# and one "unknown speaker" per language (en_x, ru_x ...)
SPEAKER_IDS = [name for name in of.ID_VOCAB if name.startswith("spk_")]
TAG_ANY = re.compile(r"\[[^\]]*\]")


def resolve_speaker(value, seed=0, name=None):
    """A speaker token for the model from what a node holds: '' -> none (or, with `name`, a token fixed by that name, so a
    speaker with no voice keeps one identity), 'random' -> one of the trained speakers picked by `seed`, else the token
    itself, which has to be one the model knows."""
    value = (value or "").strip()
    if not value:
        return SPEAKER_IDS[zlib.crc32(name.strip().casefold().encode()) % len(SPEAKER_IDS)] if name else None
    if value.lower() == "random":
        return random.Random(seed).choice(SPEAKER_IDS)
    if value not in of.ID_INDEX or value.startswith("q"):
        raise ValueError(f"OGG TTS: '{value}' is not a speaker token of the model. Use spk_0 ... spk_{len(SPEAKER_IDS) - 1}, an unknown-speaker "
                         "token such as ru_x, or random (see ogg/id_vocab.json).")
    return value


def check_quality(value):
    value = (value or "").strip()
    if value and (value not in of.ID_INDEX or not value.startswith("q")):
        raise ValueError(f"OGG TTS: '{value}' is not a recording token of the model: use q0 ... q598 (see ogg/id_vocab.json).")
    return value or None


def split_text(text, split):
    lines = [line.strip() for line in text.splitlines() if line.strip()]
    return lines if split else ([" ".join(lines)] if lines else [])


def auto_cap(prompt_frames, prompt_chars, text):
    """The reference script's safety cap on frames: twice the prompt's frames per character times the text, plus two blocks."""
    rate = prompt_frames / prompt_chars if prompt_chars else 1.6
    return int(2 * rate * len(text)) + 2 * BLOCK


class OggGenerate(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="OggGenerate",
            display_name="OGG TTS Generate",
            category="audio/ogg",
            description="Text to codec codes (for OGG Decode or OGG Editor). Optional emotion and vocal tags go inline: [angry] "
                        "[frustrated] [excited] [tired] [playful] [sad] [scared] [shouting] [whispering] [softly] [voice breaking] "
                        "[singing] [pause] [slowly] [laughing] [sighing] [gasping] [panting] [coughing] [clears throat] [sniffling] "
                        "[scoffing] [yawning] [humming]. With a voice the model clones it; without one it speaks in a voice of its "
                        "own (or the speaker token's).",
            inputs=[
                OggModel.Input("model"),
                io.String.Input("text", multiline=True, tooltip="What to say. With split_lines, each line is its own take."),
                io.Combo.Input("language", options=list(of.LANGUAGES), default="en"),
                io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True),
                io.Int.Input("steps", default=6, min=1, max=32, tooltip="Unmasking iterations per block of 16 frames (0.74 s)."),
                io.Float.Input("cfg", default=1.5, min=0.0, max=6.0, step=0.1,
                               tooltip="Guidance. 1-1.5 keeps the voice's natural tempo; 2 speaks faster and more pointedly; 0 turns it "
                                       "off (and halves the work)."),
                io.Float.Input("temperature", default=0.8, min=0.0, max=2.0, step=0.05),
                io.Boolean.Input("split_lines", default=True,
                                 tooltip="One take per line of text. Each line hears the two lines before and after it as context (not "
                                         "spoken), so the delivery fits the scene. Off: the whole text is one take."),
                io.Int.Input("max_frames", default=0, min=0, max=8192,
                             tooltip="Safety cap on frames per take (21.5 per second); 0 derives it from the text's length. Generation "
                                     "ends earlier by itself, at the model's own end-of-speech."),
                io.String.Input("speaker", default="", advanced=True,
                                tooltip="A voice the model knows, with no clip: spk_0 ... spk_3156 (3157 anonymous trained speakers), an "
                                        "unknown-speaker token such as ru_x, or random (one picked by the seed; the console says which). "
                                        "Empty: none. With a voice connected it is only an extra hint. There is no way to describe a voice in words."),
                io.String.Input("quality", default="", advanced=True,
                                tooltip="A trained recording, q0 ... q598 (each is a book or a game the model heard), or empty. It carries the "
                                        "recording's room and, to a degree, its manner of delivery."),
                io.Float.Input("mtp_temperature", default=0.0, min=0.0, max=2.0, step=0.05, advanced=True,
                               tooltip="Sampling temperature of the acoustic codes; 0 takes the most likely."),
                io.Int.Input("mtp_top_k", default=0, min=0, max=1024, advanced=True),
                io.Int.Input("mtp_rounds", default=4, min=1, max=16, advanced=True,
                             tooltip="Unmasking rounds per codebook and block of the acoustic decoder."),
                OggVoice.Input("voice", optional=True,
                               tooltip="A voice from OGG Encode Voice (a clip and its exact transcript) or Load OGG Voice."),
            ],
            outputs=[OggCodes.Output()],
        )

    @classmethod
    def execute(cls, model, text, language="en", seed=0, steps=6, cfg=1.5, temperature=0.8, split_lines=True, max_frames=0,
                speaker="", quality="", mtp_temperature=0.0, mtp_top_k=0, mtp_rounds=4, voice=None) -> io.NodeOutput:
        texts = split_text(text, split_lines)
        if not texts:
            raise ValueError("OGG TTS: there is no text to speak")
        prompt_codes, prompt_text = None, ""
        if voice is not None:
            prompt_codes = voice["codes"].long()
            prompt_text = SPEAKER_TAG.sub("", voice["transcript"]).strip()
        speaker, quality = resolve_speaker(speaker, seed), check_quality(quality)
        if speaker:
            logging.info("OGG TTS: speaker token %s", speaker)
        comfy.model_management.load_models_gpu([model.patcher])
        caps = [max_frames or auto_cap(prompt_codes.shape[1] if prompt_codes is not None else 0, len(prompt_text), line) for line in texts]
        progress = comfy.utils.ProgressBar(sum(caps))
        done, pieces, lines, takes = 0, [], [], []
        for index, line in enumerate(texts):
            before = texts[max(0, index - CONTEXT_LINES):index] if split_lines else []
            after = texts[index + 1:index + 1 + CONTEXT_LINES] if split_lines else []
            last = None
            for clean, codes_so_far in gen.blocks_cached(
                    model.model, model.tok, prompt_codes, prompt_text, line, language, block=BLOCK, steps=steps, cfg=cfg,
                    temperature=temperature, max_frames=caps[index], seed=seed + index, speaker=speaker, quality=quality,
                    mtp_temperature=mtp_temperature, mtp_top_k=mtp_top_k, mtp_rounds=mtp_rounds,
                    semantic_left=tuple(before), after=tuple(after)):
                comfy.model_management.throw_exception_if_processing_interrupted()
                progress.update_absolute(done + min(clean.shape[0], caps[index]))
                last = (clean, codes_so_far)
            done += caps[index]
            progress.update_absolute(done)
            if last is None:
                continue
            clean, codes_so_far = last
            n = gen.speech_frames(clean)
            if n == 0:
                logging.warning("OGG TTS: line %d came out empty: %r", index, line)
                continue
            start = sum(takes)
            lines.append({"speaker": 0, "text": line, "chunk": index, "start": start, "end": start + n})
            takes.append(n)
            pieces.append(codes_so_far[:, :n].detach().cpu().long().contiguous())
        if not pieces:
            raise RuntimeError("OGG TTS produced no speech; try another seed or a higher cfg")
        return io.NodeOutput({"codes": torch.cat(pieces, -1), "lines": lines, "takes": takes})


# ------------------------------------------------------------------------------------------------------------- dialogue

SLOTS = 6  # speakers a dialogue can name and give a voice
# a line's text before its first colon that looks like a name (one or two words), for catching a misspelled speaker
PLAIN_NAME = re.compile(r"[^\W\d_]+(?: [^\W\d_]+)?")


def parse_script(script, names):
    """A script of lines `Name: text` -> [(slot, text)], `slot` the index of the name in `names` (compared without case).
    A line that starts with no name goes on with the speaker before it; a line `Name:` alone gives the turn to that
    speaker for the lines that follow. A word before a colon that is no known name raises, to catch a misspelled one."""
    slots = {name.strip().casefold(): slot for slot, name in enumerate(names) if name.strip()}
    turns, slot = [], None
    for raw in script.splitlines():
        line = raw.strip()
        if not line:
            continue
        head, colon, rest = line.partition(":")
        key = head.strip().casefold()
        if colon and key in slots:
            slot, line = slots[key], rest.strip()
            if not line:
                continue
        elif colon and PLAIN_NAME.fullmatch(head.strip()):
            raise ValueError(f"OGG Dialogue: '{head.strip()}' before a colon is not one of the speaker names ({', '.join(n.strip() for n in names if n.strip())}). "
                             "Add it as a name, or start the line with a speaker's name if the colon is part of the text.")
        if slot is None:
            raise ValueError("OGG Dialogue: the script must begin with a speaker's name and a colon, e.g. 'Jean: Hello.'")
        turns.append((slot, line))
    return turns


class OggDialogue(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        name_inputs = [io.String.Input(f"name_{i + 1}", default=("Jean", "Maria")[i] if i < 2 else "",
                                       tooltip="The name of this speaker as it is written in the script (without case). Leave empty if unused.")
                       for i in range(SLOTS)]
        speaker_inputs = [io.String.Input(f"speaker_{i + 1}", default="", optional=True, advanced=True,
                                          tooltip=f"For the speaker of name_{i + 1} without a voice: a trained speaker token (spk_N, ru_x ...) or random. "
                                                  "Empty: one is picked from the name, so the same name always gets the same identity (the console says which).")
                          for i in range(SLOTS)]
        voice_inputs = [OggVoice.Input(f"voice_{i + 1}", optional=True,
                                       tooltip=f"The voice of the speaker named in name_{i + 1} (OGG Encode Voice or Load OGG Voice).")
                        for i in range(SLOTS)]
        return io.Schema(
            node_id="OggDialogue",
            display_name="OGG TTS Dialogue",
            category="audio/ogg",
            description="A script of lines `Name: [tag] text`, each speaker with a voice of its own. Name the speakers in name_1..name_6 and "
                        "connect a voice to the voice_N beside each (a speaker with no voice gets a trained identity of the model instead). A line with no name goes on with the speaker before it. Output: codes "
                        "for OGG Decode or OGG Editor. Tags as in OGG TTS Generate.",
            inputs=[
                OggModel.Input("model"),
                io.String.Input("script", multiline=True, default="Jean: [excited] Hello, my name is Jean, and I am an adventurer.\n"
                                "Maria: [sighing] You say that to everyone.\nJean: [laughing] Only to those who ask.",
                                tooltip="One turn per line: `Name: text`."),
                io.Combo.Input("language", options=list(of.LANGUAGES), default="en"),
                io.Combo.Input("mode", options=["line by line", "one pass"], default="line by line",
                               tooltip="line by line: every line is generated on its own from its speaker's voice, hearing the two lines "
                                       "before and after it as context (not spoken), as the authors' script does for dialogue. one pass: the "
                                       "whole scene in one generation with all the voices at once, the model's dialogue mode; it needs a "
                                       "voice for every speaker, gives one take, and the lines' timing is not known to the editor. Experimental: "
                                       "the authors' README documents only line by line, and in their training code the dialogue scenes are "
                                       "switched on only when a data-check file exists, so the released weights may not have learned this mode."),
                io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True),
                io.Int.Input("steps", default=6, min=1, max=32, tooltip="Unmasking iterations per block of 16 frames (0.74 s)."),
                io.Float.Input("cfg", default=1.5, min=0.0, max=6.0, step=0.1,
                               tooltip="Guidance. 1-1.5 keeps the voices' natural tempo; 2 speaks faster and more pointedly."),
                io.Float.Input("temperature", default=0.8, min=0.0, max=2.0, step=0.05),
                io.Int.Input("max_frames", default=0, min=0, max=16384,
                             tooltip="Safety cap on frames per take (21.5 per second); 0 derives it from the text's length."),
                *name_inputs,
                io.String.Input("quality", default="", advanced=True,
                                tooltip="A trained recording id such as q35 (see ogg/id_vocab.json), or empty."),
                io.Float.Input("mtp_temperature", default=0.0, min=0.0, max=2.0, step=0.05, advanced=True,
                               tooltip="Sampling temperature of the acoustic codes; 0 takes the most likely."),
                io.Int.Input("mtp_top_k", default=0, min=0, max=1024, advanced=True),
                io.Int.Input("mtp_rounds", default=4, min=1, max=16, advanced=True,
                             tooltip="Unmasking rounds per codebook and block of the acoustic decoder."),
                io.Combo.Input("dialogue_layout", options=["as trained", "reference script"], default="as trained", advanced=True, optional=True,
                               tooltip="Only for one pass. as trained: the context laid out as the authors' training code builds it (the whole "
                                       "text with each voice's clip text first, then the clips one after another). reference script: as "
                                       "voice_context of their inference script does it (each clip beside its text, the scene last). The "
                                       "two differ in their repository; if one pass gives no speech, try the other."),
                *speaker_inputs,
                *voice_inputs,
            ],
            outputs=[OggCodes.Output()],
        )

    @classmethod
    def execute(cls, model, script, language="en", mode="line by line", seed=0, steps=6, cfg=1.5, temperature=0.8, max_frames=0,
                name_1="", name_2="", name_3="", name_4="", name_5="", name_6="", quality="", mtp_temperature=0.0, mtp_top_k=0,
                mtp_rounds=4, dialogue_layout="as trained", speaker_1="", speaker_2="", speaker_3="", speaker_4="", speaker_5="", speaker_6="",
                voice_1=None, voice_2=None, voice_3=None, voice_4=None, voice_5=None, voice_6=None) -> io.NodeOutput:
        names = [name_1, name_2, name_3, name_4, name_5, name_6]
        voices = [voice_1, voice_2, voice_3, voice_4, voice_5, voice_6]
        turns = parse_script(script, names)
        if not turns:
            raise ValueError("OGG Dialogue: there is no text to speak")
        order = list(dict.fromkeys(slot for slot, _ in turns))  # speakers in order of first appearance
        speakers = [speaker_1, speaker_2, speaker_3, speaker_4, speaker_5, speaker_6]
        if mode == "one pass":
            for slot in order:
                if voices[slot] is None:
                    raise ValueError(f"OGG Dialogue: one pass needs a voice for every speaker; '{names[slot].strip()}' has none (connect one to voice_{slot + 1})")
        # a speaker with no voice speaks in a trained identity (its own token, or one fixed by its name)
        tokens_ = {}
        for slot in order:
            tokens_[slot] = resolve_speaker(speakers[slot], seed, None if voices[slot] is not None else names[slot])
            if tokens_[slot]:
                logging.info("OGG Dialogue: '%s' -> speaker token %s%s", names[slot].strip(), tokens_[slot], "" if voices[slot] is None else " (with a voice)")
        quality = check_quality(quality)
        sampling = dict(block=BLOCK, steps=steps, cfg=cfg, temperature=temperature, quality=quality, mtp_temperature=mtp_temperature,
                        mtp_top_k=mtp_top_k, mtp_rounds=mtp_rounds)
        prompts = {slot: ((voices[slot]["codes"].long(), SPEAKER_TAG.sub("", voices[slot]["transcript"]).strip()) if voices[slot] is not None
                          else (None, "")) for slot in order}
        comfy.model_management.load_models_gpu([model.patcher])
        pieces, lines, takes = [], [], []

        if mode == "one pass":
            # the voices are named S1, S2, ... in order of first appearance, as in training (train_omnifish.dialogue_spans)
            tokens = {slot: f"S{i + 1}" for i, slot in enumerate(order)}
            text = of.dialogue_text([(tokens[slot], line) for slot, line in turns])
            frames = sum(prompts[slot][0].shape[1] for slot in order)
            chars = sum(len(prompts[slot][1]) for slot in order)
            cap = max_frames or auto_cap(frames, chars, text)
            progress = comfy.utils.ProgressBar(cap)
            if dialogue_layout == "reference script":
                args = (None, "", text)
                extra = dict(voices=[(tokens[slot], prompts[slot][1], prompts[slot][0]) for slot in order])
            else:
                # training: [style "dialogue"][all text: each voice's clip text, then the scene][the clips one after another][target]
                clip_text = " ".join(f"<|quad_start|>{tokens[slot]}<|quad_end|>{prompts[slot][1]}" for slot in order)
                args = (torch.cat([prompts[slot][0] for slot in order], 1), clip_text, text)
                extra = dict(speaker="dialogue")
            last = None
            for clean, codes_so_far in gen.blocks_cached(model.model, model.tok, args[0], args[1], args[2], language, max_frames=cap,
                                                         seed=seed, **extra, **sampling):
                comfy.model_management.throw_exception_if_processing_interrupted()
                progress.update_absolute(min(clean.shape[0], cap))
                last = (clean, codes_so_far)
            n = gen.speech_frames(last[0]) if last else 0
            if n == 0:
                raise RuntimeError("OGG Dialogue (one pass) ended before the first frame: the model found no speech to say. Try another seed, "
                                   "the other dialogue_layout, or the line by line mode.")
            pieces.append(last[1][:, :n].detach().cpu().long().contiguous())
            lines.append({"speaker": 0, "text": " ".join(line for _, line in turns), "chunk": 0, "start": 0, "end": n})
            takes.append(n)
            return io.NodeOutput({"codes": pieces[0], "lines": lines, "takes": takes})

        texts = [line for _, line in turns]
        caps = [max_frames or auto_cap(0 if prompts[slot][0] is None else prompts[slot][0].shape[1], len(prompts[slot][1]), line) for slot, line in turns]
        progress = comfy.utils.ProgressBar(sum(caps))
        done = 0
        for index, (slot, line) in enumerate(turns):
            prompt_codes, prompt_text = prompts[slot]
            before = texts[max(0, index - CONTEXT_LINES):index]
            after = texts[index + 1:index + 1 + CONTEXT_LINES]
            last = None
            for clean, codes_so_far in gen.blocks_cached(
                    model.model, model.tok, prompt_codes, prompt_text, line, language, max_frames=caps[index], seed=seed + index,
                    speaker=tokens_[slot], semantic_left=tuple(before), after=tuple(after), **sampling):
                comfy.model_management.throw_exception_if_processing_interrupted()
                progress.update_absolute(done + min(clean.shape[0], caps[index]))
                last = (clean, codes_so_far)
            done += caps[index]
            progress.update_absolute(done)
            n = gen.speech_frames(last[0]) if last else 0
            if n == 0:
                logging.warning("OGG Dialogue: line %d (%s) came out empty: %r", index, names[slot].strip(), line)
                continue
            start = sum(takes)
            lines.append({"speaker": order.index(slot), "text": line, "chunk": index, "start": start, "end": start + n})
            takes.append(n)
            pieces.append(codes_so_far[:, :n].detach().cpu().long().contiguous())
        if not pieces:
            raise RuntimeError("OGG Dialogue produced no speech; try another seed or a higher cfg")
        return io.NodeOutput({"codes": torch.cat(pieces, -1), "lines": lines, "takes": takes})


class OggVoiceFromTake(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="OggVoiceFromTake",
            display_name="OGG Voice from Take",
            category="audio/ogg",
            description="Turn a generated take into a voice to clone: its codes and the words it spoke. This fixes a voice found by "
                        "chance (speaker = random, a quality token, a lucky seed) so that other lines, or a dialogue, can use it.",
            inputs=[
                OggCodes.Input("codes"),
                io.Int.Input("take", default=0, min=0, max=999, tooltip="Which line of the take (0 is the first) to use as the voice."),
                io.Boolean.Input("strip_tags", default=True,
                                 tooltip="Take [tags] out of the voice's transcript, so a neutral clip does not carry the delivery of its tags."),
            ],
            outputs=[OggVoice.Output()],
        )

    @classmethod
    def execute(cls, codes, take=0, strip_tags=True) -> io.NodeOutput:
        lines = codes["lines"]
        if take >= len(lines):
            raise ValueError(f"OGG Voice from Take: the take has {len(lines)} line(s); take {take} does not exist")
        line = lines[take]
        text = line["text"]
        if strip_tags:
            text = re.sub(r"\s{2,}", " ", TAG_ANY.sub("", text)).strip()
        return io.NodeOutput({"transcript": text, "codes": codes["codes"][:, line["start"]:line["end"]].clone()})


# ----------------------------------------------------------------------------------------------------- save / load codes


class OggSaveCodes(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="OggSaveCodes",
            display_name="Save OGG Codes",
            category="audio/ogg",
            description="Save a generated take (its codes and line timing, a few hundred KB) to output/ogg_codes, to decode or edit "
                        "it again with Load OGG Codes without generating it again.",
            inputs=[
                OggCodes.Input("codes"),
                io.String.Input("name", default="ogg/take", tooltip="File name prefix; a counter is added. Subfolders are allowed."),
            ],
            is_output_node=True,
        )

    @classmethod
    def execute(cls, codes, name) -> io.NodeOutput:
        folder, filename, counter, _, _ = folder_paths.get_save_image_path(name, CODES_FOLDER)
        save_file({"codes": codes["codes"].to(torch.int16).contiguous()}, os.path.join(folder, f"{filename}_{counter:05}_.safetensors"),
                  metadata={"lines": json.dumps(codes["lines"]), "takes": json.dumps(codes["takes"])})
        return io.NodeOutput()


class OggLoadCodes(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="OggLoadCodes",
            display_name="Load OGG Codes",
            category="audio/ogg",
            description="A take saved by Save OGG Codes (output/ogg_codes), for Decode or the Editor.",
            inputs=[io.Combo.Input("codes", options=folder_paths.get_filename_list("ogg_codes"))],
            outputs=[OggCodes.Output()],
        )

    @classmethod
    def execute(cls, codes) -> io.NodeOutput:
        tensors, metadata = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("ogg_codes", codes), return_metadata=True)
        if not metadata or "lines" not in metadata or "takes" not in metadata or "codes" not in tensors:
            raise ValueError(f"{codes} is not a codes file")
        return io.NodeOutput({"codes": tensors["codes"].long(), "lines": json.loads(metadata["lines"]), "takes": json.loads(metadata["takes"])})


# ------------------------------------------------------------------------------------------------------ decode / editor


def decode(codec, codes, gap=0.0):
    """Audio [1, samples] of Generate's codes: each take decoded on its own, `gap` seconds of silence between takes, with
    room made in VRAM for one decode window."""
    takes = [take for take in codes["codes"].split(codes["takes"], -1) if take.shape[-1]]
    if not takes:
        return torch.zeros((1, FRAME_LENGTH))
    window = min(max(take.shape[-1] for take in takes), DECODE_CONTEXT + DECODE_WINDOW)
    comfy.model_management.load_models_gpu([codec.patcher], memory_required=window * DECODE_FRAME_BYTES)
    device = codec.patcher.load_device
    with comfy.model_management.cuda_device_context(device):
        audio = [codec.model.decode(take[None].to(device), torch.float32).float().cpu() for take in takes]
    silence = torch.zeros((1, round(gap * SAMPLE_RATE)))
    return torch.cat([piece for take in audio for piece in (silence, take)][1:], -1)


class OggDecode(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="OggDecode",
            display_name="OGG Decode",
            category="audio/ogg",
            description="Codes to 44.1 kHz audio. Takes (one per line of text with split_lines) are decoded one by one and joined "
                        "with 0.3 s of silence between them.",
            inputs=[
                OggCodec.Input("codec"),
                OggCodes.Input("codes"),
                io.Boolean.Input("shorten_long_pauses", default=False,
                                 tooltip="Shorten pauses longer than 0.75 s to 0.75 s, cut the silence before the first sound and "
                                         "keep 0.15 s after the last, as Fish Audio's Drama 3 Space does. Off (default) keeps the "
                                         "take as generated, [pause] tags included."),
            ],
            outputs=[io.Audio.Output()],
        )

    @classmethod
    def execute(cls, codec, codes, shorten_long_pauses=False) -> io.NodeOutput:
        waveform = decode(codec, codes, TAKE_GAP)
        if shorten_long_pauses and codes["codes"].shape[-1]:
            waveform = trim_pauses(waveform, SAMPLE_RATE, MAX_PAUSE)
        return io.NodeOutput({"waveform": waveform[None], "sample_rate": SAMPLE_RATE})


class OggEditor(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="OggEditor",
            display_name="OGG Editor",
            category="audio/ogg",
            is_output_node=True,
            search_aliases=["ogg timeline", "edit take", "remove silence", "shorten pauses"],
            description="Decode a take and edit it: captions follow the voice word by word, one track per voice. Edit opens "
                        "the editor over ComfyUI: choose how the silences sound, check the lines it flags, cut, add silence, "
                        "name voices and set their volume. Use these edits runs the node again with them; they stay in the "
                        "workflow.",
            inputs=[
                OggCodec.Input("codec"),
                OggCodes.Input("codes"),
                OggASR.Input("aligner", optional=True,
                             tooltip="Qwen3-ForcedAligner from Load Qwen3-ASR: words light up as they are spoken. Without "
                                     "it, whole lines do."),
                io.String.Input("edits", default="", optional=True, socketless=True, extra_dict={"hidden": True},
                                tooltip="Settings and edits made in the editor (managed by the editor)."),
            ],
            outputs=[io.Audio.Output()],
        )

    @classmethod
    def execute(cls, codec, codes, aligner=None, edits="") -> io.NodeOutput:
        waveform = decode(codec, codes)[0]
        analysis = analyze(waveform, SAMPLE_RATE, codes["lines"], aligner)
        settings = parse_edits(edits)
        audio = render(waveform, analysis, plan(analysis, settings)["pieces"], settings)
        return io.NodeOutput({"waveform": audio[None, None], "sample_rate": SAMPLE_RATE}, ui={"ogg_editor": [save_preview(waveform, analysis)]})


def save_preview(waveform, analysis):
    """The unedited take (16-bit WAV) and its analysis (JSON) in the temp folder, for the editor to play and draw; the
    workflow keeps only their names."""
    pcm = (waveform.clamp(-1, 1) * 32767).round().to(torch.int16).numpy().tobytes()
    name = f"{hashlib.sha1(pcm).hexdigest()[:16]}{'_aligned' if analysis['aligned'] else ''}"
    folder = os.path.join(folder_paths.get_temp_directory(), "ogg_editor")
    os.makedirs(folder, exist_ok=True)
    with wave.open(os.path.join(folder, f"{name}.wav"), "wb") as file:
        file.setnchannels(1)
        file.setsampwidth(2)
        file.setframerate(SAMPLE_RATE)
        file.writeframes(pcm)
    with open(os.path.join(folder, f"{name}.json"), "w", encoding="utf-8") as file:
        json.dump(analysis, file)
    return {"audio": f"{name}.wav", "analysis": f"{name}.json", "subfolder": "ogg_editor", "type": "temp"}


if getattr(PromptServer, "instance", None) is not None:
    @PromptServer.instance.routes.post("/ogg/editor/plan")
    async def editor_plan(request):
        """The edit list the editor previews: the node's analysis and the editor's current settings and edits."""
        body = await request.json()
        return web.json_response(plan(body["analysis"], {**DEFAULTS, **body.get("edits", {})}))


# ------------------------------------------------------------------------------------------------------------ Qwen3-ASR


class OggASRLoader(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="OggQwen3ASRLoader",
            display_name="Load Qwen3-ASR (OGG)",
            category="audio/ogg",
            description="Qwen3-ASR speech recognition (1.7B or 0.6B) or Qwen3-ForcedAligner, converted files from models/asr.",
            inputs=[io.Combo.Input("model_name", options=choices("asr", "qwen3"))],
            outputs=[OggASR.Output()],
        )

    @classmethod
    def execute(cls, model_name) -> io.NodeOutput:
        return io.NodeOutput(load_qwen3_asr(folder_paths.get_full_path_or_raise("asr", model_name)))


class OggASRTranscribe(io.ComfyNode):
    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="OggQwen3ASRTranscribe",
            display_name="Qwen3-ASR Transcribe (OGG)",
            category="audio/ogg",
            description="Speech to text in 30 languages, singing included. Its text can feed OGG Encode Voice's transcript.",
            inputs=[
                OggASR.Input("model"),
                io.Audio.Input("audio"),
                io.Combo.Input("language", options=["auto", *LANGUAGES], default="auto",
                               tooltip="auto detects the language; naming it writes the text without detecting."),
                io.String.Input("context", multiline=True, default="",
                                tooltip="Optional text that helps recognition: names, terms, the topic."),
                io.Int.Input("max_new_tokens", default=4096, min=1, max=32768, advanced=True,
                             tooltip="At most this many tokens per 20 minutes of audio."),
            ],
            outputs=[io.String.Output("text", display_name="text"), io.String.Output("language", display_name="language")],
        )

    @classmethod
    def execute(cls, model, audio, language="auto", context="", max_new_tokens=4096) -> io.NodeOutput:
        results = [transcribe(model, waveform, audio["sample_rate"], None if language == "auto" else language, context, max_new_tokens)
                   for waveform in audio["waveform"]]
        return io.NodeOutput("\n".join(text for text, _ in results), ",".join(dict.fromkeys(found for _, found in results if found)))


class OggExtension(ComfyExtension):
    async def get_node_list(self) -> list[type[io.ComfyNode]]:
        return [OggLoader, OggCodecLoader, OggEncodeVoice, OggSaveVoice, OggLoadVoice, OggGenerate, OggDialogue, OggVoiceFromTake, OggSaveCodes, OggLoadCodes,
                OggDecode, OggEditor, OggASRLoader, OggASRTranscribe]


async def comfy_entrypoint() -> OggExtension:
    return OggExtension()