| """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_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "omnivoice") |
| OMNIVOICE_REPO = "k2-fsa/OmniVoice" |
| |
| TAKE_GAP = 0.3 |
| |
| MAX_PAUSE = 0.75 |
| SPEAKER_TAG = re.compile(r"^\s*<\|speaker:\d+\|>\s*") |
| |
| CONTEXT_LINES = 2 |
| BLOCK = 16 |
|
|
| 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())) |
|
|
|
|
| |
|
|
|
|
| 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)") |
| |
| 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))) |
|
|
|
|
| |
|
|
|
|
| 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()}) |
|
|
|
|
| |
|
|
|
|
| |
| |
| 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}) |
|
|
|
|
| |
|
|
| SLOTS = 6 |
| |
| 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 = [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})") |
| |
| 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": |
| |
| 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: |
| |
| 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()}) |
|
|
|
|
| |
|
|
|
|
| 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"])}) |
|
|
|
|
| |
|
|
|
|
| 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", {})})) |
|
|
|
|
| |
|
|
|
|
| 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() |
|
|
