"""JoyAI-Echo ComfyUI node implementations. Six nodes faithful to the official inference.py: 1. JoyEcho_ModelLoader — load text encoder + DiT + VAEs (bf16) 2. JoyEcho_TextEncode — encode prompts, auto-release text encoder 3. JoyEcho_Generate — multi-shot denoise + decode with memory bank 4. JoyEcho_SingleShotGenerate — single-shot with per-shot text box and memory chaining 5. JoyEcho_PromptFormat — get system prompt for LLM-based prompt enhancement 6. JoyEcho_LLMEnhance — call LLM API to generate shot prompts from a story idea """ from __future__ import annotations import gc import json import os from pathlib import Path from typing import Any import torch DENOISING_SIGMAS = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0] def _empty_cache(): if torch.cuda.is_available(): torch.cuda.empty_cache() def _move(module, device): if module is not None: module.to(device) class SequentialOffloader: """Layer-by-layer GPU offloading for the DiT transformer blocks. Hooks into each transformer block so that only the currently-executing block resides on GPU. All other blocks stay on CPU/pinned memory. Peak VRAM for the generator drops from ~30GB to ~2-3GB (1 block + activations). """ def __init__(self, generator, device: torch.device, pin_memory: bool = True, resident_blocks: int = 0): self._generator = generator self._device = device self._hooks: list[torch.utils.hooks.RemovableHook] = [] self._pin_memory = pin_memory self._installed = False # First N transformer blocks stay permanently on GPU (no hooks, no # streaming). Each streamed block costs a PCIe round-trip per denoise # step; pinning K of 48 cuts that traffic by K/48 at K x per-block # VRAM (bf16 ~0.9GB, fp8-resident ~0.45GB per block). self._resident_blocks = max(0, int(resident_blocks)) def install(self): """Install forward hooks on transformer blocks and move them to CPU.""" if self._installed: return self._installed = True velocity_model = self._generator.model.velocity_model blocks = velocity_model.transformer_blocks # Keep pre/post processing layers on GPU (small footprint) for name, param in velocity_model.named_parameters(): if "transformer_blocks" not in name: param.data = param.data.to(self._device) for name, buf in velocity_model.named_buffers(): if "transformer_blocks" not in name: buf.data = buf.data.to(self._device) n_res = min(self._resident_blocks, len(blocks)) resident = list(blocks)[:n_res] streamed = list(blocks)[n_res:] # Resident blocks live on the GPU permanently. for block in resident: block.to(self._device) # Move streamed blocks to CPU (optionally pinned). Pinning is a # transfer-speed optimization (enables async H2D copies), NOT a # correctness requirement — so it must never be fatal. cudaHostAlloc # exhaustion surfaces as "CUDA error: out of memory" even though it # is HOST page-locked memory that ran out (hit on BEAST 2026-07-19: # the refine's resident_blocks=0 tried to pin all 48 blocks after # the shot passes had pinned only 36 — the last ~11GB of pinning # pushed past what the host could lock). On the first failure we # stop pinning entirely (the pool is exhausted; per-param retries # just burn time) and stream the rest unpinned — the non_blocking # copies silently become synchronous, slower but correct. # Already-pinned tensors are a no-op for pin_memory(), so re-installs # keep whatever pinning already succeeded. _pin = self._pin_memory and torch.cuda.is_available() for block in streamed: block.to("cpu") if _pin: try: for param in block.parameters(): param.data = param.data.pin_memory() for buf in block.buffers(): buf.data = buf.data.pin_memory() except Exception as e: _pin = False print(f"[JoyEcho] WARNING: pinned-memory allocation failed " f"({e}); streaming remaining blocks UNPINNED (slower " f"PCIe transfers, otherwise identical).", flush=True) # Also keep the wrapper's patchifiers and X0Model's non-block params on GPU for name, param in self._generator.named_parameters(): if "velocity_model.transformer_blocks" not in name and "velocity_model" not in name: param.data = param.data.to(self._device) for name, buf in self._generator.named_buffers(): if "velocity_model.transformer_blocks" not in name and "velocity_model" not in name: buf.data = buf.data.to(self._device) def make_pre_hook(block_module): def hook(module, args): block_module.to(self._device, non_blocking=True) if torch.cuda.is_available(): torch.cuda.current_stream().synchronize() return hook def make_post_hook(block_module): def hook(module, args, output): block_module.to("cpu", non_blocking=True) return hook for block in streamed: h1 = block.register_forward_pre_hook(make_pre_hook(block)) h2 = block.register_forward_hook(make_post_hook(block)) self._hooks.extend([h1, h2]) if n_res: print(f"[JoyEcho] Sequential offloading installed: {len(blocks)} blocks " f"({n_res} resident on GPU, {len(streamed)} streamed)", flush=True) else: print(f"[JoyEcho] Sequential offloading installed: {len(blocks)} blocks", flush=True) def remove(self): """Remove all hooks and move entire generator back to CPU.""" for h in self._hooks: h.remove() self._hooks.clear() self._installed = False self._generator.to("cpu") _MODEL_FILE_MANUAL = "(use checkpoint_path)" _MODEL_FILE_CATS = ("checkpoints", "diffusion_models", "unet") _LORA_FILE_MANUAL = "(use lora_path / none)" _LORA_FILE_CATS = ("loras",) _GEMMA_FILE_MANUAL = "(use gemma_path field)" _GEMMA_FILE_CATS = ("text_encoders", "clip") def _list_cat_files(cats, sentinel, exts=("*.safetensors", "*.gguf")) -> list: """Every matching file under the given ComfyUI model-dir categories, as 'category: relative/path' combo entries. Dirs shared between categories (unet is an alias of diffusion_models on newer ComfyUI) are deduped.""" try: import folder_paths except ImportError: return [sentinel] out, seen_dirs, seen = [], set(), set() for cat in cats: try: roots = folder_paths.get_folder_paths(cat) except Exception: continue for root in roots: try: rp = Path(root).resolve() except OSError: continue if not rp.is_dir() or rp in seen_dirs: continue seen_dirs.add(rp) for ext in exts: for f in rp.rglob(ext): label = f"{cat}: {f.relative_to(rp).as_posix()}" if label not in seen: seen.add(label) out.append(label) return [sentinel] + sorted(out) def _resolve_cat_file(choice: str, cats, widget: str, sentinel: str) -> str: import folder_paths cat, _, rel = choice.partition(": ") if cat in cats and rel: for root in folder_paths.get_folder_paths(cat): p = Path(root) / rel if p.is_file(): return str(p) raise FileNotFoundError( f"{widget} {choice!r} no longer exists on disk. Refresh the node " f"list (R) and re-pick, or use {sentinel}.") def _list_model_files() -> list: return _list_cat_files(_MODEL_FILE_CATS, _MODEL_FILE_MANUAL) def _resolve_model_file(choice: str) -> str: return _resolve_cat_file(choice, _MODEL_FILE_CATS, "model_file", _MODEL_FILE_MANUAL) def _list_lora_files() -> list: return _list_cat_files(_LORA_FILE_CATS, _LORA_FILE_MANUAL, exts=("*.safetensors",)) def _resolve_lora_file(choice: str) -> str: return _resolve_cat_file(choice, _LORA_FILE_CATS, "lora_file", _LORA_FILE_MANUAL) def _parse_lora_entries(raw: str, default_strength: float) -> list: """Multi-LoRA (v1.4): lora_path accepts SEVERAL entries separated by commas or newlines. Each entry is "path", "path@strength" or "path:strength" - the suffix counts as a strength only when it parses as a float, so Windows drive letters (F:\\x.safetensors) can never be mistaken for one. Entries without a suffix use the lora_strength widget. The fusion engine (apply_loras) has always summed a LIST of LoRA deltas; only this plumbing was single.""" import re as _re out = [] for part in _re.split(r"[,\n]+", raw or ""): part = part.strip().strip('"').strip("'") if not part: continue strength = float(default_strength) for sep in ("@", ":"): head, s, tail = part.rpartition(sep) if s and head: try: strength = float(tail) part = head.strip() break except ValueError: continue out.append((part, strength)) return out def _list_gemma_files() -> list: # Gemma can be a single-file .safetensors OR a .gguf, and lives in either # models/text_encoders or models/clip depending on how the user filed it. return _list_cat_files(_GEMMA_FILE_CATS, _GEMMA_FILE_MANUAL) def _resolve_gemma_file(choice: str) -> str: return _resolve_cat_file(choice, _GEMMA_FILE_CATS, "gemma_file", _GEMMA_FILE_MANUAL) class JoyEcho_LoraStacker: """Chainable LoRA stack - the familiar dropdown+strength UI. Each node adds up to three LoRAs; wire lora_stack to another stacker to chain more, and the final one into the Model Loader's lora_stack input. Every entry is fused into the DiT at load (safetensors DiT path only).""" @classmethod def INPUT_TYPES(cls): lora_list = ["(none)"] + [x for x in _list_lora_files() if x != _LORA_FILE_MANUAL] opt = {"lora_stack": ("JOYECHO_LORA_STACK", { "tooltip": "Chain from another LoRA Stack node to add more slots."})} for i in (1, 2, 3): opt[f"lora_{i}"] = (lora_list, { "default": "(none)", "tooltip": "LoRA from models/loras. (none) = slot unused."}) opt[f"strength_{i}"] = ("FLOAT", { "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05}) return {"required": {}, "optional": opt} RETURN_TYPES = ("JOYECHO_LORA_STACK",) RETURN_NAMES = ("lora_stack",) FUNCTION = "stack" CATEGORY = "JoyAI-Echo" def stack(self, lora_stack=None, **kw): out = list(lora_stack) if lora_stack else [] for i in (1, 2, 3): choice = kw.get(f"lora_{i}") or "(none)" if choice != "(none)" and ": " in str(choice): out.append((_resolve_lora_file(choice), float(kw.get(f"strength_{i}", 1.0)))) return (out,) class JoyEcho_ModelLoader: """Load JoyAI-Echo model components: text encoder, DiT generator, and VAEs.""" @classmethod def INPUT_TYPES(cls): return { # All inputs are optional: pick from the dropdowns for the common # case, or fall back to the manual *_path fields for a GGUF's VAE # source, an HF gemma DIRECTORY, or a file outside the model tree. "required": {}, "optional": { # --- DiT: pick a full/GGUF model, or type a full checkpoint --- "model_file": (_list_model_files(), { "default": _MODEL_FILE_MANUAL, "tooltip": "Pick the model instead of typing checkpoint_path. " "A .safetensors = FULL checkpoint (replaces checkpoint_path " "entirely: DiT + VAEs + vocoder + text connectors from that " "file). A .gguf = DiT ONLY - checkpoint_path must still point " "at a full safetensors (e.g. the JoyAI release) to supply the " "VAEs/vocoder/connectors. Refresh the node list (R) after " "adding files.", }), "checkpoint_path": ("STRING", { "default": "", "tooltip": "Manual fallback / GGUF VAE source. A full safetensors " "checkpoint supplying the VAEs, vocoder and text connectors. " "REQUIRED when model_file is a .gguf (DiT only); leave empty " "when model_file is a full .safetensors.", }), # --- text encoder: pick a single file, or type a path/dir --- "gemma_file": (_list_gemma_files(), { "default": _GEMMA_FILE_MANUAL, "tooltip": "Pick the Gemma text encoder from models/text_encoders or " "models/clip instead of typing gemma_path. Single-file " ".safetensors or .gguf only - for an HF gemma-3-12b-it " "DIRECTORY, leave this on the sentinel and type the folder in " "gemma_path. Refresh the node list (R) after adding files.", }), "gemma_path": ("STRING", { "default": "", "tooltip": "Manual fallback for the text encoder. Use for an HF " "gemma-3-12b-it DIRECTORY (dropdowns list files, not folders), " "or an encoder outside models/text_encoders and models/clip. " "Leave empty when gemma_file is set.", }), # --- LoRA: pick from the loras tree, or type a path --- "lora_file": (_list_lora_files(), { "default": _LORA_FILE_MANUAL, "tooltip": "Pick a LoRA from the models/loras tree instead of typing " "lora_path. Applied at lora_strength on the safetensors DiT " "path (ignored when a GGUF DiT is selected). Refresh the node " "list (R) after adding files.", }), "lora_stack": ("JOYECHO_LORA_STACK", { "tooltip": "Wire a JoyEcho LoRA Stack node here for the " "dropdown+strength multi-LoRA UI. Stacks with " "lora_file and lora_path entries.", }), "lora_path": ("STRING", { "default": "", "multiline": True, "tooltip": "One or MORE LoRAs, separated by commas or newlines. Each " "entry is a path, optionally with its own strength: " "'a.safetensors@0.7, b.safetensors@0.5'. Entries without a " "strength use lora_strength. Stacks WITH the lora_file pick. " "Safetensors DiT path only (ignored on a GGUF DiT).", }), "lora_strength": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05, }), # --- quantization toggles (DiT, then encoder) --- "fp8_transformer": ("BOOLEAN", { "default": False, "tooltip": "Quantize the DiT's attention/FF linear weights to " "float8_e4m3fn at load (upcast per-layer during inference). " "Roughly halves transformer weight memory - works from the " "normal bf16 checkpoint, keeping JoyAI's memory training and " "projection tensors intact. Slight quality cost; VAEs, text " "encoder and non-linear layers stay bf16. Ignored when a " "GGUF is picked in model_file (already quantized).", }), "fp8_scaled_mm": ("BOOLEAN", { "default": False, "tooltip": "Store DiT linears as float8_e4m3fn AND compute the matmuls " "natively in fp8 via torch._scaled_mm (RTX 40/50-series). " "Unlike fp8_transformer there is NO per-layer upcast tax - " "and at ~22GB resident the DiT can run with " "sequential_offload OFF at moderate resolutions. Overrides " "fp8_transformer. Ignored for GGUF DiTs.", }), "encoder_fp8": ("BOOLEAN", { "default": False, "tooltip": "Store the Gemma text encoder's linear weights as " "float8_e4m3fn (upcast per-layer at encode). Roughly halves " "the encoder's ~24GB footprint - it fits the GPU for the " "encode pass (seconds per shot instead of ~12s on CPU) and " "frees ~11GB system RAM. Encode runs once per queue item, so " "the upcast tax is irrelevant here. Slight embedding shift - " "voice is the canary; A/B before adopting.", }), # --- misc --- "low_vram": ("BOOLEAN", { "default": False, "tooltip": "Load text encoder on CPU for 24GB GPUs. " "Encoding will be slower but uses no GPU memory.", }), }, } RETURN_TYPES = ("JOYECHO_MODEL",) RETURN_NAMES = ("model",) FUNCTION = "load_model" CATEGORY = "JoyAI-Echo" def load_model(self, checkpoint_path: str = "", gemma_path: str = "", lora_path: str = "", lora_strength: float = 1.0, lora_stack=None, low_vram: bool = False, fp8_transformer: bool = False, model_file: str = _MODEL_FILE_MANUAL, lora_file: str = _LORA_FILE_MANUAL, fp8_scaled_mm: bool = False, encoder_fp8: bool = False, gemma_file: str = _GEMMA_FILE_MANUAL): from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps from ltx_core.quantization import QuantizationPolicy from ltx_distillation.models.ltx_wrapper import create_ltx2_wrapper from ltx_distillation.models.text_encoder_wrapper import create_text_encoder_wrapper from ltx_distillation.models.vae_wrapper import create_vae_wrappers dropdown_lora = None if lora_file and lora_file != _LORA_FILE_MANUAL: dropdown_lora = _resolve_lora_file(lora_file) gguf_dit_path = None if model_file and model_file != _MODEL_FILE_MANUAL: _resolved = _resolve_model_file(model_file) if _resolved.lower().endswith(".gguf"): gguf_dit_path = _resolved print(f"[JoyEcho] model_file: DiT from GGUF {_resolved}; VAEs/vocoder/" f"connectors from checkpoint_path.", flush=True) else: checkpoint_path = _resolved print(f"[JoyEcho] model_file: full checkpoint {_resolved}.", flush=True) # A real pick is always "category: relative/path"; the sentinel (any # "(use ...)" placeholder) has no ": ", so this guard is robust to the # sentinel wording and to stale saved values from an older node version. if gemma_file and ": " in gemma_file: gemma_path = _resolve_gemma_file(gemma_file) print(f"[JoyEcho] gemma_file: {gemma_path}", flush=True) if not str(gemma_path).strip(): raise ValueError( "No text encoder selected. Pick a Gemma in gemma_file, or type its " "path/directory in gemma_path.") if not str(checkpoint_path).strip(): raise ValueError( "checkpoint_path is empty. It must point at a FULL safetensors checkpoint" + (" - with a GGUF picked in model_file it still supplies the VAEs, " "vocoder and text connectors (e.g. echo-longvideo-release.safetensors)." if gguf_dit_path else " (or pick a .safetensors in model_file).")) checkpoint_path = str(Path(checkpoint_path).expanduser().resolve()) gemma_path = str(Path(gemma_path).expanduser().resolve()) # ComfyUI-quantized checkpoints ("fp8mixed learned" builds, marked by # .comfy_quant tensors) are packaged for the standard ComfyUI loader. # This ledger path never applies their weight_scale at runtime (the # scaled-mm consumer needs tensorrt_llm) and LoRA fusion assumes the # LTX transposed fp8 convention - the model would load MIS-SCALED and # LoRA fusion crashes with shape errors. Refuse early and clearly. if checkpoint_path.lower().endswith(".safetensors"): try: import json as _json import struct as _struct with open(checkpoint_path, "rb") as _f: _n = _struct.unpack(" About > Advanced system" "\n settings > Performance Settings > Advanced > Virtual" "\n memory > Change. Set a custom size of 64-128 GB on an" "\n SSD. The peak is brief; it spills instead of dying." "\n 2. Use a GGUF in model_file instead of a bf16/fp8" "\n checkpoint - GGUF weights stay packed and memory-" "\n mapped rather than copied into RAM." "\n 3. Close other applications, and load no LoRAs (fusion" "\n allocates additional transients at the peak)." "\n" + "=" * 72 + "\n", flush=True) except Exception as _e_pf: print(f"[JoyEcho] RAM preflight skipped ({type(_e_pf).__name__}).", flush=True) print("[JoyEcho] Loading DiT generator...", flush=True) _lora_entries = [] if dropdown_lora: _lora_entries.append((dropdown_lora, float(lora_strength))) _lora_entries.extend(_parse_lora_entries(lora_path, lora_strength)) if lora_stack: _lora_entries.extend((str(p), float(st)) for p, st in lora_stack) # dedupe identical paths (dropdown pick repeated in lora_path) _seen = set() _lora_entries = [e for e in _lora_entries if not (str(Path(e[0]).expanduser()) in _seen or _seen.add(str(Path(e[0]).expanduser())))] loras = tuple( LoraPathStrengthAndSDOps(str(Path(p).expanduser()), float(st), LTXV_LORA_COMFY_RENAMING_MAP) for p, st in _lora_entries) if _lora_entries: print("[JoyEcho] LoRAs: " + ", ".join( f"{Path(p).name}@{st}" for p, st in _lora_entries), flush=True) if gguf_dit_path is not None: # DiT from GGUF via the Rebels loader machinery; the wrapper class # is identical to create_ltx2_wrapper's, so Generate can't tell. from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as _Builder from ltx_core.model.transformer import LTXModelConfigurator, X0Model from ltx_distillation.models.ltx_wrapper import LTX2DiffusionWrapper from .rebels_loaders import ( _LOADER_CFG, _SWAP_MAP, _dit_module_ops, _full_config, _GGUFDiTLoader, _gguf_entries, _materialize_meta, _rebind_swapped, ) if fp8_transformer: print("[JoyEcho] fp8_transformer ignored: GGUF DiT is already quantized.", flush=True) if loras: print("[JoyEcho] WARNING: lora_path is ignored on the GGUF DiT path.", flush=True) try: _cfg = _full_config(checkpoint_path) except Exception: _cfg = _full_config(_LOADER_CFG) _SWAP_MAP.clear() _entries = _gguf_entries(gguf_dit_path) _consumed = set() _builder = _Builder( model_class_configurator=LTXModelConfigurator, model_path=gguf_dit_path, model_sd_ops=None, module_ops=_dit_module_ops(_entries, _consumed, dtype), model_loader=_GGUFDiTLoader(_cfg, _entries, _consumed, dtype), ) _transformer = _builder.build(device=torch.device("cpu"), dtype=dtype) generator = LTX2DiffusionWrapper( model=X0Model(_transformer), video_height=736, video_width=1280) generator.eval() _materialize_meta(generator, _entries, _consumed, dtype) _rebind_swapped(generator) _SWAP_MAP.clear() else: quantization = None if fp8_scaled_mm: quantization = QuantizationPolicy.fp8_scaled_mm_torch() print("[JoyEcho] fp8_scaled_mm ON: DiT linears stored float8_e4m3fn and " "COMPUTED in fp8 via torch._scaled_mm (no upcast tax; ~22GB resident).", flush=True) elif fp8_transformer: quantization = QuantizationPolicy.fp8_cast() print("[JoyEcho] fp8_transformer ON: quantizing DiT linear weights to " "float8_e4m3fn (upcast per-layer at inference).", flush=True) generator = create_ltx2_wrapper( checkpoint_path=checkpoint_path, # A single-file gemma has no model*.safetensors folder for the # ledger's eager text-encoder builder; that builder is never # used here (the encoder is built separately below), so skip it. gemma_path=None if gemma_single_file else gemma_path, device=torch.device("cpu"), dtype=dtype, video_height=736, video_width=1280, loras=loras, quantization=quantization, ) generator.eval() # Free quantize/fuse transients before the next mmap-heavy stage. import gc gc.collect() # Load VAEs to CPU print("[JoyEcho] Loading VAEs...", flush=True) video_vae, audio_vae = create_vae_wrappers( checkpoint_path=checkpoint_path, device=torch.device("cpu"), dtype=dtype, with_video_encoder=True, with_audio_encoder=True, decoder_device=torch.device("cpu"), ) video_vae.eval() audio_vae.eval() gc.collect() # Text encoder: built LAZILY via this closure. On conditioning-cache-HIT # runs the encoder is never used, and eager loading cost 21-33GB of host # RAM plus up to a minute of load for nothing - on a 64GB box that alone # pushed the whole run into pagefile thrash. TextEncode resolves the # builder only after a cache MISS. text_encoder_device = torch.device("cpu") if low_vram else device def _build_text_encoder(): if gemma_single_file: # Single-file gemma routes through the Rebels TextEncoder node: # .gemma_virtual_folder + sidecars, scale-aware fp8/GGUF weight # swap, session model cache. Returns the same # GemmaTextEncoderWrapper class as the folder path, so # TextEncode's GPU hot-swap and release work unchanged. # (Rebels local patch: gemma_single_file routing) from .rebels_loaders import RebelsJE_TextEncoder, _full_config, _LOADER_CFG if encoder_fp8: print("[JoyEcho] encoder_fp8 ignored for a single-file gemma: " "fp8 files already carry their own quantization; bf16 " "files stay bf16.", flush=True) print(f"[JoyEcho] Loading text encoder (single file) on " f"{text_encoder_device} via Rebels routing...", flush=True) try: _cfg = _full_config(checkpoint_path) except Exception: _cfg = _full_config(_LOADER_CFG) wrapper = RebelsJE_TextEncoder().run( _cfg, gemma_path, "our_fp8", checkpoint_path, low_vram)[0] wrapper.eval() return wrapper print(f"[JoyEcho] Loading text encoder (bf16) on {text_encoder_device}...", flush=True) text_encoder = create_text_encoder_wrapper( checkpoint_path=checkpoint_path, gemma_path=gemma_path, device=text_encoder_device, dtype=dtype, ) text_encoder.eval() if encoder_fp8: # The vision tower + projector are NEVER touched by text encoding # (the conditioning is a trained mix over the LANGUAGE model's # hidden states only - feature_extractor stacks hidden_states from # the text stack). Drop them entirely: ~1GB weights + buffers, # which is exactly the margin that decides whether a 24GB card can # host the encode pass on GPU. _stripped = 0 for _mname, _mod in text_encoder.named_modules(): for _attr in ("vision_tower", "multi_modal_projector"): _sub = getattr(_mod, _attr, None) if isinstance(_sub, torch.nn.Module): _stripped += sum(p.nbytes for p in _sub.parameters()) _stripped += sum(b.nbytes for b in _sub.buffers()) setattr(_mod, _attr, None) if _stripped: gc.collect() print(f"[JoyEcho] encoder_fp8: dropped the unused vision tower/projector " f"({_stripped/1e9:.1f}GB).", flush=True) # Halve the Gemma footprint: store the remaining linear weights as # fp8, upcasting per layer at encode. Encode runs ONCE per queue # item, so the upcast tax that makes fp8_transformer slow on the # DiT is irrelevant here. JD's embeddings processor / connector # projections stay bf16. from ltx_core.quantization.fp8_cast import _replace_fwd_with_upcast _n = 0 for _name, _m in text_encoder.named_modules(): if (isinstance(_m, torch.nn.Linear) and ("language_model" in _name or "vision_tower" in _name) and _m.weight.dtype in (torch.bfloat16, torch.float16)): _m.weight.data = _m.weight.data.to(torch.float8_e4m3fn) if _m.bias is not None: _m.bias.data = _m.bias.data.to(torch.float8_e4m3fn) _replace_fwd_with_upcast(_m) _n += 1 _gb = (sum(p.nbytes for p in text_encoder.parameters()) + sum(b.nbytes for b in text_encoder.buffers())) / 1e9 print(f"[JoyEcho] encoder_fp8 ON: {_n} Gemma linears stored float8_e4m3fn; " f"wrapper now {_gb:.1f}GB (upcast per-layer at encode).", flush=True) return text_encoder audio_sample_rate = audio_vae.get_output_sample_rate() or 24000 model = { "text_encoder": None, # resolved lazily from text_encoder_builder "text_encoder_builder": _build_text_encoder, "generator": generator, "video_vae": video_vae, "audio_vae": audio_vae, "audio_sample_rate": audio_sample_rate, "device": device, "dtype": dtype, "checkpoint_path": checkpoint_path, "gemma_path": gemma_path, "encoder_fp8": bool(encoder_fp8), } print(f"[JoyEcho] Model loaded. Audio sample rate: {audio_sample_rate}", flush=True) return (model,) # Default negative for the DMD (no-CFG) pipeline: steers each shot's conditioning # away from these in embedding space. Covers BOTH failure modes seen on the # multishot path: burned-in captions/subtitles (video context) and invented # sung/musical audio from the Hat Man etc. (audio context). Kept as the FUNCTION # default too, so it still fires when a stale graph node lacks the new widget. _DEFAULT_JOYECHO_NEGATIVE = ( "subtitles, captions, closed captions, on-screen text, text characters, glyphs, " "letters, words, writing, garbled text, lyrics, signatures, watermark, timestamp, " "logo, music, singing, song, humming, melody, chanting, vocalizing, score, " "soundtrack, musical, instrumental" ) # Split per-domain defaults. The encoder emits SEPARATE video_context / # audio_context tensors, so each domain gets its own negative text + scale: # - video: burned-in captions live here -> can be pushed hard # - audio: music lives here, but SPEECH does too ("subtitles" also correlates # with speech in training data) -> push gently, music tokens ONLY, no # voice-adjacent words (humming/chanting/vocalizing strangle whispers). _DEFAULT_JOYECHO_NEGATIVE_VIDEO = ( "subtitles, captions, closed captions, on-screen text, text characters, glyphs, " "letters, words, writing, garbled text, lyrics, signatures, watermark, timestamp, logo" ) _DEFAULT_JOYECHO_NEGATIVE_AUDIO = ( "music, singing, song, melody, score, soundtrack, musical, instrumental, " "background music" ) class JoyEcho_TextEncode: """Encode text prompts using Gemma-3-12b. Supports: - One prompt per line (multi-line text, each line = one shot) - JSON format: {"prompts": ["shot1", "shot2", ...]} (official format) - JSON file path (*.json) After encoding, the text encoder is released from GPU to free ~24GB VRAM. """ @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("JOYECHO_MODEL",), "prompts": ("STRING", { "multiline": True, "default": "", "tooltip": "One prompt per line, JSON object, or path to .json file", }), }, "optional": { "negative_prompt_video": ("STRING", { "multiline": True, "default": _DEFAULT_JOYECHO_NEGATIVE_VIDEO, "tooltip": "Steered away from in VIDEO context only (burned-in captions/subtitles/text). Safe to push hard - does not touch the audio lane. Empty or scale 0 disables.", }), "negative_scale_video": ("FLOAT", { "default": 0.8, "min": 0.0, "max": 3.0, "step": 0.05, "tooltip": "Video-context steering strength. Renormalized, so higher values no longer degrade the image the way the old shared lever did.", }), "negative_prompt_audio": ("STRING", { "multiline": True, "default": _DEFAULT_JOYECHO_NEGATIVE_AUDIO, "tooltip": "Steered away from in AUDIO context only. Music tokens ONLY - do NOT add caption words (captions correlate with speech; steering audio away from them kills dialogue). Empty or scale 0 disables.", }), "negative_scale_audio": ("FLOAT", { "default": 0.3, "min": 0.0, "max": 3.0, "step": 0.05, "tooltip": "Audio-context steering strength. Keep LOW (~0.2-0.4) or dialogue suffers.", }), "release_text_encoder": ("BOOLEAN", {"default": True}), }, } RETURN_TYPES = ("JOYECHO_MODEL", "JOYECHO_COND",) RETURN_NAMES = ("model", "conditioning",) FUNCTION = "encode" CATEGORY = "JoyAI-Echo" @staticmethod def _parse_prompts(prompts: str) -> list[str]: """Parse prompts from text, JSON string, or JSON file path.""" text = prompts.strip() # Check if it's a file path to a .json if text.endswith(".json") and not text.startswith("{"): p = Path(text).expanduser() if not p.is_absolute(): p = Path(__file__).resolve().parent / p p = p.resolve() if p.exists(): with open(p, "r", encoding="utf-8") as f: data = json.load(f) return JoyEcho_TextEncode._extract_from_json(data) # Check if it's a JSON object if text.startswith("{"): try: data = json.loads(text) return JoyEcho_TextEncode._extract_from_json(data) except json.JSONDecodeError: pass # Fall back to one-prompt-per-line return [line.strip() for line in text.split("\n") if line.strip()] @staticmethod def _extract_from_json(data: dict) -> list[str]: """Extract prompt list from JSON (supports 'prompts' or 'shots' key).""" if isinstance(data.get("prompts"), list): return [str(p).strip() for p in data["prompts"] if str(p).strip()] if isinstance(data.get("shots"), list): return [str(p).strip() for p in data["shots"] if str(p).strip()] raise ValueError("JSON must contain a 'prompts' or 'shots' array.") def encode(self, model: dict, prompts: str, negative_prompt: str = _DEFAULT_JOYECHO_NEGATIVE, negative_scale: float = 0.5, release_text_encoder: bool = True, negative_prompt_video: str = None, negative_scale_video: float = None, negative_prompt_audio: str = None, negative_scale_audio: float = None): text_encoder = model.get("text_encoder") if text_encoder is None and not callable(model.get("text_encoder_builder")): raise RuntimeError( "Text encoder not available. It may have been released already. " "Reload the model to encode new prompts." ) prompt_list = self._parse_prompts(prompts) if not prompt_list: raise ValueError("No prompts provided. Enter text, JSON, or a .json file path.") # Speakers + voice anchors ride INSIDE the conditioning from here on. # They are properties of THIS text, so they belong on the data path, not # in module globals: the old stash pattern leaked a previous script's # speakers/anchors into any graph without a picker node (audit finding # 2026-07-29). Derived here, attached to conds[0] as "joyecho_meta", # cached WITH the conditioning (correct: same text = same speakers), # consumed by Generate, and stripped before the transformer sees it. _je_meta = None try: import json as _json_meta _md = _json_meta.loads(prompts) if str(prompts).lstrip().startswith("{") else None if isinstance(_md, dict): try: from .joyecho_script_picker import derive_speakers as _derive_spk except ImportError: from joyecho_script_picker import derive_speakers as _derive_spk _arr = _md.get("prompts") or _md.get("shots") or [] _je_meta = { "speakers": _derive_spk(_md, _arr), "voice_refs": dict(_md.get("voice_refs") or {}), } except Exception: _je_meta = None device = model["device"] # --- Conditioning disk cache: the same script + negatives + encoder # always produces the same conditioning, so encode it ONCE per machine. # Re-renders (the A/B loop) skip the encoder entirely - on a box whose # encoder can only run on CPU (24GB cards), this turns a many-minutes # encode into a sub-second load. import hashlib as _hashlib import json as _json_cc _cc_key = _hashlib.sha256(_json_cc.dumps([ prompt_list, str(negative_prompt), float(negative_scale or 0), str(negative_prompt_video), str(negative_scale_video), str(negative_prompt_audio), str(negative_scale_audio), str(model.get("gemma_path")), bool(model.get("encoder_fp8")), # checkpoint identity: the text-embedding connectors come from the # checkpoint and co-shape conditioning, so a model swap with the # same text must MISS, not serve the previous model's tensors # (audit finding 2026-07-29; invalidates all earlier cache files). str(model.get("checkpoint_path")), ]).encode()).hexdigest()[:16] _cc_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "cond_cache") _cc_path = os.path.join(_cc_dir, f"conds_{_cc_key}.pt") if os.path.isfile(_cc_path): try: cached_conds = torch.load(_cc_path, map_location="cpu", weights_only=True) print(f"[JoyEcho] Conditioning cache HIT ({os.path.basename(_cc_path)}) - " f"{len(cached_conds)} shot(s), encode skipped entirely.", flush=True) if release_text_encoder and text_encoder is not None: print("[JoyEcho] Releasing text encoder to free VRAM...", flush=True) del text_encoder model["text_encoder"] = None gc.collect() _empty_cache() return (model, cached_conds,) except Exception as _e: print(f"[JoyEcho] Conditioning cache unreadable ({_e}); re-encoding.", flush=True) print(f"[JoyEcho] Encoding {len(prompt_list)} prompt(s)...", flush=True) # Cache MISS: now (and only now) the encoder is actually needed. if text_encoder is None: text_encoder = model["text_encoder_builder"]() model["text_encoder"] = text_encoder # Hot-swap: a low_vram (CPU-resident) encoder costs ~10s+ per shot to run # the 12B gemma on CPU, while the GPU sits idle during this phase. Borrow # the GPU for the encode pass when the encoder fits in free VRAM, and hand # it back before the denoise phase - denoise-time VRAM is untouched. moved_to_gpu = False if getattr(device, "type", "") == "cuda" and torch.cuda.is_available(): try: enc_dev = next(text_encoder.parameters()).device except (StopIteration, AttributeError): enc_dev = None if enc_dev is not None and enc_dev.type == "cpu": need = sum(p.nbytes for p in text_encoder.parameters()) need += sum(b.nbytes for b in text_encoder.buffers()) free = torch.cuda.mem_get_info(device)[0] # Absolute headroom, not a multiplier: encode activations for a # <=1k-token prompt are well under 1GB, and a 1.15x margin on a # ~21GB encoder demanded ~3GB of phantom headroom - exactly what # kept 24GB cards on CPU encode. The OOM fallback in _enc() # backstops a miss. if free > need + 1.2e9: print(f"[JoyEcho] Text encoder -> GPU for the encode pass " f"({need/1e9:.1f}GB weights, {free/1e9:.1f}GB free)...", flush=True) try: _move(text_encoder, device) moved_to_gpu = True except Exception as _e: _move(text_encoder, torch.device("cpu")) _empty_cache() print(f"[JoyEcho] GPU encode swap failed ({_e}); encoding on CPU.", flush=True) else: print(f"[JoyEcho] Encoder stays on CPU for encode " f"({need/1e9:.1f}GB weights vs {free/1e9:.1f}GB free VRAM).", flush=True) if any(p.dtype == torch.float8_e4m3fn for p in text_encoder.parameters()): print("[JoyEcho] WARNING: encoder is fp8-stored but must encode on CPU - " "the per-layer CPU upcast makes this MUCH slower than plain bf16. " "On GPUs where the encoder cannot fit (24GB cards), set " "encoder_fp8=False.", flush=True) def _enc(texts): nonlocal moved_to_gpu try: return text_encoder(texts) except RuntimeError as _e: # A hard allocator failure surfaces as a GENERIC RuntimeError # ("CUDA error: out of memory"), not torch.cuda.OutOfMemoryError # - on 2026-07-19 that escaped this handler and killed the whole # queue item instead of degrading to CPU encode. Treat any # out-of-memory RuntimeError as the fallback trigger. # (Rebels local patch: broad OOM fallback) _is_oom = (isinstance(_e, torch.cuda.OutOfMemoryError) or "out of memory" in str(_e).lower()) if not _is_oom or not moved_to_gpu: raise print(f"[JoyEcho] GPU encode OOM ({type(_e).__name__}) - " "falling back to CPU for the rest.", flush=True) _move(text_encoder, torch.device("cpu")) _empty_cache() moved_to_gpu = False return text_encoder(texts) # PER-DOMAIN NEGATIVE STEERING: the DMD-distilled pipeline has no CFG, so # we extrapolate conditioning away from a negative in embedding space: # cond' = cond + scale * (cond - neg) (then renormalized, below) # The encoder emits SEPARATE video_context / audio_context tensors, so # each domain gets its own negative text and scale: # - video_context: burned-in captions live here -> push hard # - audio_context: music lives here, but speech does too -> push gently # The old SHARED lever coupled the two: raising it past ~0.5 to kill # subtitles also strangled dialogue (captions correlate with speech in # training data) and drove the audio context off-manifold (the hum). # RENORM: raw extrapolation grows the context norm by ~(1+scale); the DiT # never saw conditioning at that magnitude -> audio hum / image drift. # Restoring the original per-token norm keeps only the DIRECTION change. if negative_prompt_video is None and negative_prompt_audio is None: # Stale graph still carrying the old single-lever widgets: preserve # the old behavior (same negative, same scale, both domains). negative_prompt_video = negative_prompt negative_prompt_audio = negative_prompt negative_scale_video = negative_scale negative_scale_audio = negative_scale domains = [] # (context_key, negative_text, scale) for key, txt, sc in (("video_context", negative_prompt_video, negative_scale_video), ("audio_context", negative_prompt_audio, negative_scale_audio)): try: sc = float(sc) except (TypeError, ValueError): sc = 0.0 txt = str(txt).strip() if txt is not None else "" if sc > 0.0 and txt: domains.append((key, txt, sc)) neg_ctx = {} # context_key -> (neg_tensor, scale) if domains: encoded = {} # one encoder pass per DISTINCT negative text for key, txt, sc in domains: if txt not in encoded: _nc = _enc([txt]) encoded[txt] = {k: (t.detach() if isinstance(t, torch.Tensor) else t) for k, t in _nc.items()} del _nc nv = encoded[txt].get(key) if isinstance(nv, torch.Tensor) and nv.is_floating_point(): neg_ctx[key] = (nv, sc) print(f"[JoyEcho] Negative for {key}: scale={sc}.", flush=True) def _steer(v, nv, scale): out = v + scale * (v - nv.to(v.device)) # Norm-preserving rescale (same idea as RescaleCFG): keep the # direction change, restore the original per-token magnitude. norm_in = v.norm(dim=-1, keepdim=True) norm_out = out.norm(dim=-1, keepdim=True).clamp_min(1e-6) return out * (norm_in / norm_out) cached_conds = [] for i, prompt in enumerate(prompt_list): cond = _enc([prompt]) if neg_ctx: cond = dict(cond) for key, (nv, sc) in neg_ctx.items(): v = cond.get(key) if (isinstance(v, torch.Tensor) and v.is_floating_point() and v.shape == nv.shape): cond[key] = _steer(v, nv, sc) elif i == 0: print(f"[JoyEcho] WARNING: negative SKIPPED for {key} " f"(shape {getattr(v, 'shape', None)} vs {tuple(nv.shape)}).", flush=True) if i == 0: applied = ", ".join(f"{k}@{sc}" for k, (_, sc) in neg_ctx.items()) print(f"[JoyEcho] negative applied per-domain: {applied}.", flush=True) cached_conds.append( {k: (v.detach().cpu() if isinstance(v, torch.Tensor) else v) for k, v in cond.items()} ) del cond print(f"[JoyEcho] Encoded shot {i+1}/{len(prompt_list)}", flush=True) if neg_ctx: neg_ctx.clear() if moved_to_gpu: _move(text_encoder, torch.device("cpu")) _empty_cache() print("[JoyEcho] Text encoder -> CPU (encode pass done).", flush=True) # Attach the script-derived speakers/anchors so they persist in the # cache and reach Generate on the data path (see comment at derivation). if _je_meta and cached_conds: cached_conds[0]["joyecho_meta"] = _je_meta # Persist the conditioning for instant re-runs of the same script. try: _cc_bytes = sum(v.nbytes for c in cached_conds for v in c.values() if isinstance(v, torch.Tensor)) if _cc_bytes < 4_000_000_000: # sanity cap os.makedirs(_cc_dir, exist_ok=True) torch.save(cached_conds, _cc_path) print(f"[JoyEcho] Conditioning cached ({os.path.basename(_cc_path)}, " f"{_cc_bytes/1e6:.0f}MB) - future runs of this script skip the encode.", flush=True) except Exception as _e: print(f"[JoyEcho] Conditioning cache save failed ({_e}); continuing.", flush=True) if release_text_encoder: print("[JoyEcho] Releasing text encoder to free VRAM...", flush=True) del text_encoder model["text_encoder"] = None gc.collect() _empty_cache() return (model, cached_conds,) class JoyEcho_Generate: """Generate multi-shot video + audio using DMD few-step denoising with memory bank. Implements the same hot-swap memory management as official inference.py: - Denoise phase: generator on GPU, VAE on CPU - Decode phase: generator on CPU, VAE on GPU """ @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("JOYECHO_MODEL",), "conditioning": ("JOYECHO_COND",), "seed": ("INT", {"default": 12345, "min": 0, "max": 2**31 - 1}), "num_frames": ("INT", {"default": 241, "min": 9, "max": 1441, "step": 8, "tooltip": "Must be 1 + 8*k (e.g. 121, 241, 361)"}), # Rebels local patch: portrait resolutions. Height was capped at # 1088 while width allowed 1920, which silently forbade portrait # (e.g. 1088x1920 to match a portrait Z-Image first frame). Both # axes now cap at 1920; step 32 keeps the latent packing valid. "video_height": ("INT", {"default": 736, "min": 256, "max": 1920, "step": 32}), "video_width": ("INT", {"default": 1280, "min": 256, "max": 1920, "step": 32}), }, "optional": { "video_fps": ("INT", {"default": 24, "min": 1, "max": 60, "tooltip": "KEEP AT 24. The LTX joint audio-video prior is 24fps-native: any other value (25 included) systematically drifts spoken voices toward Commonwealth accents (British/Australian) and overrides accent wording in the prompt. Verified A/B 2026-07-29."}), "v2a_grad_scale": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.1}), "memory_max_size": ("INT", {"default": 7, "min": 0, "max": 20}), "num_fix_frames": ("INT", {"default": 3, "min": 0, "max": 10}), "enable_audio_memory": ("BOOLEAN", { "default": True, "tooltip": "Feed previous shots' audio latents as cross-shot conditioning. " "ON keeps the same voice across every shot - the point of " "multishot - at a small lip-sync cost on long dialogue. OFF " "gives the tightest sync but the voice can drift between " "shots; use only for sync-critical single-voice pieces with " "a strong voice description repeated in every shot.", }), "audio_memory_window_size": ("INT", {"default": 96, "min": 16, "max": 256}), "memory_video_anchor": (["last", "loudest"], { "default": "last", "tooltip": "Which part of a shot becomes the VIDEO half of its memory " "slot. last = the final frames, so the next shot continues " "from where this one ended. loudest = the original behaviour: " "the frames around the shot's most articulated speech, which " "is usually mid-shot and makes the next shot look like it " "rewound a second or two. The AUDIO half is chosen by " "max_response either way - that is the right voice exemplar.", }), "speaker_order": ("STRING", { "default": "", "tooltip": "Who SPEAKS in each shot, comma or space separated " "(e.g. 'A,B'). Tags each memory slot with its speaker so a " "shot only inherits audio memory from the character talking " "in it - the other character's slots are silenced for that " "shot while their face memory is kept. This is the fix for " "two same-gender voices merging across shots. The list CYCLES, " "so 'A,B' covers any number of alternating shots. Leave EMPTY " "for the previous behaviour (every voice in every shot).", }), "sequential_offload": ("BOOLEAN", { "default": False, "tooltip": "Enable layer-by-layer GPU offloading for DiT. " "Reduces VRAM from ~30GB to ~3GB at the cost of slower inference.", }), "output_prefix": ("STRING", { "default": "joyecho/shot", "tooltip": "Prefix for per-shot video files saved immediately after each shot completes.", }), "reference_image": ("IMAGE", { "tooltip": "Optional identity reference (e.g. a Z-Image render). Pre-seeds the " "cross-shot memory bank as a permanent anchor slot, so every shot is " "conditioned on this face/look - reference-driven I2V. Uses one of the " "num_fix_frames anchor slots.", }), "transition": (["cut", "dissolve", "vhs_glitch"], { "default": "cut", "tooltip": "Shot-boundary treatment. cut = hard cuts (original). dissolve = " "overlap cross-dissolve + equal-power audio crossfade (shortens total " "by transition_frames per boundary). vhs_glitch = analog static burst " "at each cut: snow, tearing bands, dropout lines + a tape-noise audio " "hit (length unchanged).", }), "transition_frames": ("INT", { "default": 8, "min": 1, "max": 48, "tooltip": "Length of the transition in frames. Dissolve: 8-12 is natural. " "VHS glitch: 3-6 reads as a head-switch stutter, 8-12 as a violent burst.", }), "glitch_intensity": ("FLOAT", { "default": 0.7, "min": 0.1, "max": 1.0, "step": 0.05, "tooltip": "vhs_glitch only: how hard the burst hits (snow mix, tear count, " "audio static level).", }), "head_trim_frames": ("INT", { "default": 14, "min": 0, "max": 24, "tooltip": "Trim this many frames (plus matching audio) from the START of every " "shot. The model's first frames morph out of the reference/memory " "content, and shot-start lip sync is at its worst there. 14 is the " "production-proven value (THE WITNESS ran it). 0 = auto: trims 8 " "when a reference_image is wired, none otherwise.", }), "decode_tiling": (["auto", "on", "off"], { "default": "auto", "tooltip": "Temporal-chunked VAE decode (64-frame chunks, 24-frame blended " "overlap, no spatial tiles = no spatial seams). Caps decode peak " "memory at ~one chunk instead of the whole shot - fixes the hard " "crash decoding 241f at 1280x736. auto = only when " "height*width*frames exceeds the known-safe budget; small renders " "keep the original single-pass decode.", }), "head_trim_first_shot": ("INT", { "default": 0, "min": 0, "max": 64, "tooltip": "Extra-long head trim for SHOT 1 ONLY. Shot 1 has an empty memory " "bank, so its conditioning is 100% the reference clips - it opens " "practically ON the reference and takes far longer than later " "shots to morph out. 0 = auto: double head_trim_frames (min 24) " "when a reference_image is wired. Raise to 32-40 if the first " "frame still shows the reference.", }), "reference_zoom": ("FLOAT", { "default": 1.2, "min": 1.0, "max": 2.0, "step": 0.05, "tooltip": "Over-zoom applied to reference images before conditioning. A " "reference at EXACTLY the render size is pixel-continuable - the " "model opens the shot ON it (a long literal ghost that no head " "trim covers). Cropping in ~20% keeps identity and wide " "composition but breaks the wholesale-continuation shortcut. " "1.0 = old passthrough behavior.", }), "reference_shots": ("STRING", { "default": "", "forceInput": True, "tooltip": "Comma list from RefPicker's ref_shots output: 0-based shot " "index where each reference injects (per-character entry " "shots). Unwired/empty = all references at shot 1.", }), "resident_blocks": ("INT", { "default": 0, "min": 0, "max": 48, "tooltip": "Sequential offload only: keep the first N of 48 transformer " "blocks permanently on the GPU and stream the rest. Each " "streamed block is a PCIe round-trip per denoise step, so " "N=24 halves the streaming time for N x per-block VRAM " "(bf16 ~0.9GB/block, fp8 ~0.45GB). Raise until VRAM is " "nearly full; 0 = stream everything (old behavior).", }), "hires_factor": ("FLOAT", { "default": 1.0, "min": 1.0, "max": 2.0, "step": 0.05, "tooltip": "ROUTING switch, not a quality slider - it decides who builds " "your master. 1.0 (RECOMMENDED): shots stay base-res and the " "AutoFinish worker upscales them deterministically " "(bicubic+CAS) - zero per-frame detail shimmer. Above 1.0: the " "selected hires_denoise pass runs in-render and the master is " "built from ITS output instead (AutoFinish detects the hires " "shots and skips its own upscale). Only raise this when you " "deliberately want the refine modes' synthesized detail " "(slight texture shimmer) or the spatial mode (even /32 dims " "only, e.g. 768-height). Judge results from the *_MASTER.mp4 " "file, never the in-canvas preview.", }), "hires_denoise": (["subtle (1 step)", "medium (2 steps)", "strong (tenstrip 4-step)", "spatial (LTX latent upsampler, no churn)"], { "default": "subtle (1 step)", "tooltip": "How the hires pass works. The three refine modes re-noise and " "re-denoise at target res - they SYNTHESIZE real detail but " "reshuffle fine texture slightly per frame (subtle = sigma 0.42 " "one step, most faithful; medium = 0.725 two steps; strong = " "tenstrip's 0.92/0.725/0.42 ladder, deepest). spatial = the LTX " "latent upsampler on the shot's own latents: deterministic, " "zero temporal churn, fixed 1.5x (hires_factor just enables " "it) - use heights whose /32 is EVEN (768, not 736) or it " "smears one edge.", }), "temporal_upscale": (["off", "2x (48 fps master)"], { "default": "off", "tooltip": "LTX temporal latent upsampler (x2) on each shot's OWN latents " "after sampling: per-shot files and the master come out at " "double fps (24 -> ~48) with the audio untouched - same " "duration, twice the motion samples, no optical-flow ghosting. " "The render itself stays at video_fps 24, so the accent/sync " "law is unaffected; the in-canvas preview also stays at base " "fps (only the saved files double). v1 limitation: requires " "hires_factor 1.0 (skipped with a warning otherwise). Keep OFF " "for found-footage looks - interpolated camcorder reads as " "soap opera.", }), }, } RETURN_TYPES = ("IMAGE", "AUDIO",) RETURN_NAMES = ("images", "audio",) FUNCTION = "generate" CATEGORY = "JoyAI-Echo" OUTPUT_NODE = True @classmethod def VALIDATE_INPUTS(cls, hires_factor=None, reference_zoom=None, **kwargs): # Turn ComfyUI's cryptic pre-execution validation error into an # actionable one for the widget-scramble case. Loading a workflow saved # under an OLDER node layout shifts values into the wrong slots (values # serialize by position): the tell is hires_factor holding a # hires_denoise STRING ("could not convert string to float") and/or # reference_zoom holding a big resident_blocks number ("bigger than max # of 2.0"). Listing these two here makes ComfyUI defer their validation # to us so we can name the real fix instead of a type error. generate() # still clamps at runtime as a backstop. _scramble = ("JoyEcho Generate: this node's widget values are scrambled - " "you loaded a workflow saved under an OLDER node layout, so " "values shifted into the wrong slots. FIX: right-click this " "node and remove it, add a fresh 'JoyEcho Generate' node, " "reconnect its inputs and re-enter your settings (or load the " "workflow from the current release zip). ") if hires_factor is not None: try: float(hires_factor) except (TypeError, ValueError): return _scramble + f"(hires_factor got {hires_factor!r})" if reference_zoom is not None: try: if not (1.0 <= float(reference_zoom) <= 2.0): return _scramble + f"(reference_zoom got {reference_zoom})" except (TypeError, ValueError): return _scramble + f"(reference_zoom got {reference_zoom!r})" return True def generate( self, model: dict, conditioning: list, seed: int = 12345, num_frames: int = 241, video_height: int = 736, video_width: int = 1280, video_fps: int = 25, v2a_grad_scale: float = 2.0, memory_max_size: int = 7, num_fix_frames: int = 3, enable_audio_memory: bool = True, audio_memory_window_size: int = 96, memory_video_anchor: str = "last", speaker_order: str = "", sequential_offload: bool = False, output_prefix: str = "joyecho/shot", reference_image=None, transition: str = "cut", transition_frames: int = 8, glitch_intensity: float = 0.7, head_trim_frames: int = 0, decode_tiling: str = "auto", head_trim_first_shot: int = 0, reference_zoom: float = 1.2, reference_shots: str = "", resident_blocks: int = 0, hires_factor: float = 1.0, hires_denoise: str = "subtle (1 step)", temporal_upscale: str = "off", ): from ltx_distillation.inference.bidirectional_pipeline import BidirectionalAVInferencePipeline from ltx_distillation.inference.memory_bidirectional_pipeline import BidirectionalMemoryAVInferencePipeline from ltx_distillation.inference.memory_multishot import ( PairedAudioVideoMemoryBank, build_paired_audio_memory_kwargs, video_uint8_to_pil_frames, ) from ltx_distillation.utils import ( add_noise, compute_latent_shapes, decode_benchmark_sample, encode_memory_frames_batch, ) # WIDGET-SHIFT SELF-HEAL. ComfyUI serializes widget values by POSITION, # so a graph saved under an older node layout maps its tail values into # the wrong slots after an upgrade (classic signature: reference_zoom # gets an old resident_blocks value like 24, hires_factor gets a STRING # from hires_denoise and arrives as NaN). Clamp the two unambiguous # cases loudly instead of rendering garbage; the real fix is deleting # and re-adding the node (fresh, correctly-ordered widgets). import math as _math try: _hf = float(hires_factor) except (TypeError, ValueError): _hf = float("nan") if _math.isnan(_hf) or _math.isinf(_hf) or not (1.0 <= _hf <= 2.0): print(f"[JoyEcho] WARNING: hires_factor={hires_factor!r} is not a " "sane value - your saved graph was made under an OLDER node " "layout and the widget values have shifted slots. Using 1.0 " "(hires off). Fix: delete + re-add the Generate node, then " "re-enter your settings.", flush=True) hires_factor = 1.0 if not (0.3 <= float(reference_zoom) <= 4.0): print(f"[JoyEcho] WARNING: reference_zoom={reference_zoom} is not a " "sane value (widget-shift from an older graph layout). Using " "1.2. Fix: delete + re-add the Generate node.", flush=True) reference_zoom = 1.2 generator = model["generator"] video_vae = model["video_vae"] audio_vae = model["audio_vae"] audio_sample_rate = model["audio_sample_rate"] device = model["device"] dtype = model["dtype"] # Validate num_frames if (num_frames - 1) % 8 != 0: num_frames = 1 + ((num_frames - 1) // 8) * 8 print(f"[JoyEcho] Adjusted num_frames to {num_frames} (must be 1 + 8*k)", flush=True) # Update generator resolution if changed generator.video_height = video_height generator.video_width = video_width generator.latent_height = video_height // 32 generator.latent_width = video_width // 32 generator.video_frame_seqlen = generator.latent_height * generator.latent_width # RENDER fps must drive the video RoPE clock (wrapper class default is # 24.0). At 25fps the hardcoded 24 skewed video rope-time 4% fast vs # audio's true-seconds rope -> ~40ms/s growing mouth-ahead-of-audio # drift, crossing visibility at ~9.6s into every shot (the "10s lip # sync cliff", 2026-07-23). generator.VIDEO_FPS = float(video_fps) if int(video_fps) != 24: print(f"[JoyEcho] WARNING: video_fps={video_fps}. The joint AV prior is " f"24fps-native - non-24 fps drifts voices toward Commonwealth " f"accents and overrides accent wording (verified 2026-07-29). " f"Use 24 unless you specifically want that.", flush=True) # Compute latent shapes video_shape, audio_shape = compute_latent_shapes( num_frames=num_frames, video_height=video_height, video_width=video_width, batch_size=1, video_fps=float(video_fps), ) # Build pipelines denoising_sigmas = torch.tensor(DENOISING_SIGMAS, device=device, dtype=torch.float32) base_pipeline = BidirectionalAVInferencePipeline( generator=generator, add_noise_fn=add_noise, denoising_sigmas=denoising_sigmas, ) memory_pipeline = BidirectionalMemoryAVInferencePipeline( generator=generator, add_noise_fn=add_noise, denoising_sigmas=denoising_sigmas, memory_downscale_factor=1, ) # Memory bank memory_bank = PairedAudioVideoMemoryBank( max_size=memory_max_size, save_mode="random_every_shot_frame", num_fix_frames=num_fix_frames, ) # REFERENCE CONDITIONING (I2V-as-reference): references are VIDEO-ONLY # conditioning clips kept OUT of the paired audio-video memory bank. # (The earlier design seeded them into the bank with zero-filled audio # latents; with 2+ refs those silent paired slots audibly polluted the # audio lane. The bank now holds only real generated shots; refs are # prepended at the video-memory encode step, invisible to all audio # machinery, and persist for the whole run.) _ref_clips = [] _ref_sched = [] # must exist even with NO reference image (a no-character # script legitimately produces none - first hit 2026-07-18) if reference_image is not None and memory_max_size > 0: import numpy as np from PIL import Image as _PILImage try: _sched_ints = [int(x) for x in str(reference_shots).split(",") if x.strip() != ""] except ValueError: _sched_ints = [] def _sched_of(_j): return _sched_ints[_j] if _j < len(_sched_ints) else 0 # Dedupe on (pixels, scheduled shot): the same picked image delivered # twice by a wiring quirk collapses, but the same image scheduled at # TWO different shots is a deliberate re-entry injection (character # returning after an absence) and every copy must survive. _uniq_idx = [] _seen = [] for _i in range(int(reference_image.shape[0])): _t = reference_image[_i] if not any(_sched_of(_i) == _s and _t.shape == _u.shape and torch.equal(_t, _u) for _u, _s in _seen): _seen.append((_t, _sched_of(_i))) _uniq_idx.append(_i) if len(_uniq_idx) < int(reference_image.shape[0]): print(f"[JoyEcho] Reference batch: {int(reference_image.shape[0])} images, " f"{len(_uniq_idx)} unique after dedupe.", flush=True) _tw, _th = int(video_width), int(video_height) _zoom = max(1.0, float(reference_zoom)) _ref_sched = [] for _ri in _uniq_idx[:6]: _arr = reference_image[_ri].detach().cpu().numpy() _arr = (np.clip(_arr, 0.0, 1.0) * 255.0).astype(np.uint8) _ref_pil = _PILImage.fromarray(_arr) # The memory encoder requires frames at EXACTLY the render size # (frames_to_video_tensor raises on mismatch). ALWAYS cover-fit # with a deliberate over-zoom: a reference at exactly the render # size is pixel-continuable, and the model opens the shot ON it # (a long literal full-frame ghost that no head trim covers). # Cropping in keeps identity and wide composition but breaks the # wholesale-continuation shortcut. Crop is mildly top-biased so # faces survive. _scale = max(_tw / _ref_pil.width, _th / _ref_pil.height) * _zoom _rw = max(_tw, int(round(_ref_pil.width * _scale))) _rh = max(_th, int(round(_ref_pil.height * _scale))) if (_rw, _rh) != _ref_pil.size: _ref_pil = _ref_pil.resize((_rw, _rh), _PILImage.LANCZOS) if (_rw, _rh) != (_tw, _th): _left = (_rw - _tw) // 2 _top = int((_rh - _th) * 0.25) # bias crop toward the top (faces) _ref_pil = _ref_pil.crop((_left, _top, _left + _tw, _top + _th)) _ref_clips.append([_ref_pil] * 9) _ref_sched.append(_sched_ints[_ri] if _ri < len(_sched_ints) else 0) print(f"[JoyEcho] {len(_ref_clips)} reference image(s) prepared as VIDEO-ONLY " f"conditioning clips ({_tw}x{_th}); audio lane untouched.", flush=True) all_video_frames = [] all_audio_waveforms = [] _hires_audio_lats = [] _hires_video_lats = [] # spatial hires mode only (~10MB/shot on CPU) _hires_spatial = str(hires_denoise).lower().startswith("spatial") _temporal_on = str(temporal_upscale).lower().startswith("2x") if _temporal_on and hires_factor > 1.0: print("[JoyEcho] temporal_upscale requires hires_factor 1.0 in this " "version - the hires passes re-save shots at base fps and would " "undo it. Temporal upscale SKIPPED.", flush=True) _temporal_on = False num_shots = len(conditioning) offloader = None if sequential_offload: offloader = SequentialOffloader(generator, device, resident_blocks=resident_blocks) elif device.type == "cuda": # PRE-FLIGHT for no-offload runs. A config that cannot fit does not # raise a normal OOM here - cudaMallocAsync hard-aborts the entire # ComfyUI process mid-denoise (observed dying on the tiny timestep # embedding once weights + the AdaLN working set pinned VRAM full). # Refuse certain death with a catchable error; warn on tight fits. try: _free, _total = torch.cuda.mem_get_info() _pbytes = sum(p.numel() * p.element_size() for p in generator.parameters()) _tokens = (video_height // 32) * (video_width // 32) * (num_frames // 8 + 1) _act_est = _tokens * 4096 * 2 * 32 # empirical AdaLN/attention working-set floor except Exception: _pbytes = 0 if _pbytes: if _pbytes + 4 * 2**30 > _total: raise ValueError( f"sequential_offload=False but the DiT weighs {_pbytes/2**30:.1f}GiB and this " f"GPU has {_total/2**30:.1f}GiB total. This cannot fit and would hard-abort " f"the whole ComfyUI process (fatal CUDA error, not a catchable OOM). Enable " f"sequential_offload, or shrink the weights (fp8_transformer / a GGUF DiT).") if _pbytes + _act_est + 2 * 2**30 > _free: print(f"[JoyEcho] WARNING: no-offload budget is tight: DiT {_pbytes/2**30:.1f}GiB " f"+ ~{_act_est/2**30:.1f}GiB est. activations ({_tokens} tokens) vs " f"{_free/2**30:.1f}GiB free. A fatal process abort mid-denoise is likely - " f"consider sequential_offload=True or lower resolution/frames.", flush=True) # Temporal-chunked VAE decode: 241f at 1280x736 decoded in ONE pass # hard-crashes a 32GB card (cuDNN abort mid-conv); 361f at 544x960 # (~189M pixels*frames) is render-proven safe, so auto kicks in just # above that. Temporal-only tiling = no spatial seams; 24-frame # blended overlap. _decode_tiling_config = None if decode_tiling == "on" or ( decode_tiling == "auto" and video_height * video_width * num_frames > 195_000_000 ): from ltx_core.model.video_vae import TemporalTilingConfig, TilingConfig _decode_tiling_config = TilingConfig( spatial_config=None, temporal_config=TemporalTilingConfig( tile_size_in_frames=64, tile_overlap_in_frames=24), ) print("[JoyEcho] Tiled VAE decode ON (temporal 64f chunks, 24f overlap).", flush=True) print(f"[JoyEcho] Generating {num_shots} shot(s) at {video_width}x{video_height}, " f"{num_frames} frames{' [sequential offload]' if sequential_offload else ''}...", flush=True) # PER-CHARACTER AUDIO MEMORY. The bank stores one slot per SHOT and is # otherwise character-blind, so with two characters alternating shots the # audio lane hands BOTH voices to every shot and the model may continue # either one - which is what makes two same-gender voices converge. # Tagging each slot with its speaker lets a shot inherit audio memory # only from the character actually talking in it. Empty = old behaviour. # The speaker_order WIDGET is an override. Left empty (the normal case in # a queue-driven workflow, where nobody retypes it per script), the order # comes from the script itself - either an explicit "speakers" array or # the " is talking" attribution already present in every shot. The # widget only exists for the rare case of overriding a script by hand. import re as _re_spk # Conditioning-carried metadata is the AUTHORITATIVE script-derived # source: it was computed from the exact text that produced these # embeddings and travels with them through the cache. The module-global # stash is retained only as a last-resort legacy fallback (it can be # stale when a loader node is ComfyUI-cached - audit finding 2026-07-29). _je_meta = {} try: _m0 = conditioning[0].get("joyecho_meta") if conditioning else None if isinstance(_m0, dict): _je_meta = _m0 except (AttributeError, IndexError, TypeError): _je_meta = {} _speakers = [t for t in _re_spk.split(r"[,\s]+", str(speaker_order).strip()) if t] _spk_src = "speaker_order widget" if not _speakers and _je_meta.get("speakers"): _speakers = [str(x) for x in _je_meta["speakers"]] _spk_src = "script (via conditioning)" if not _speakers: try: from .joyecho_script_picker import LAST_SPEAKERS as _auto except Exception: try: from joyecho_script_picker import LAST_SPEAKERS as _auto except Exception: _auto = [] if _auto: _speakers = list(_auto) _spk_src = "script (legacy stash - may be stale)" if _speakers: _speakers = [_speakers[i % len(_speakers)] for i in range(num_shots)] print(f"[JoyEcho] per-character audio memory ON (from {_spk_src}): " f"{' '.join(_speakers)}", flush=True) else: print("[JoyEcho] per-character audio memory OFF - no speakers declared " "and no ' is talking' attribution found in the script.", flush=True) # VOICE ANCHORS (2026-07-28). The memory bank guarantees CONSISTENCY, not # CORRECTNESS: shot 1 rolls its voice from text conditioning alone (empty # bank), and whatever region/timbre that roll lands on, the bank then # carries faithfully - measured twice today: two renders on different # seeds AND different DiT files produced the same wrong voices at the # same shots, because both merges share audio branches and the roll is # text-deterministic. No prompt wording pins it. # # So pin it with DATA: the script may carry # "voice_refs": {"Zara": "path/to/clip.mp4", ...} # and each clip's audio is encoded (the pack's audio-VAE encoder, loaded # but never previously called) and stored as a bank slot TAGGED with the # character BEFORE the shot loop. Shot 1 then CONTINUES an approved voice # instead of auditioning a new one, and the same file re-casts the same # voice in every future render. Keys must exactly match the script's # speaker tags or per-character filtering will zero the anchor once the # speaker owns a generated slot. _voice_refs = {} if isinstance(_je_meta.get("voice_refs"), dict) and _je_meta["voice_refs"]: _voice_refs = dict(_je_meta["voice_refs"]) # authoritative: rode else: # with the conditioning try: from .joyecho_script_picker import LAST_VOICE_REFS as _lvr except ImportError: try: from joyecho_script_picker import LAST_VOICE_REFS as _lvr except ImportError: _lvr = {} if isinstance(_lvr, dict): _voice_refs = dict(_lvr) # TIER-1 AUTO-CAST (2026-07-28): a speaker whose tag matches a folder # under input/joyecho_voices/ is cast automatically - no voice_refs key # needed. Explicit voice_refs entries WIN over the folder scan. The pick # is DETERMINISTIC (alphabetically first file): same character, same # file, same voice, every render. Dropping a file into the folder IS # casting the character; a random per-render pick would reintroduce the # per-video voice roulette this whole system exists to kill. try: import folder_paths as _fp _vroot = os.path.join(_fp.get_input_directory(), "joyecho_voices") except Exception: _fp, _vroot = None, "" if _vroot and os.path.isdir(_vroot) and _speakers: _vexts = (".mp4", ".wav", ".mp3", ".flac", ".mov", ".m4a", ".webm") for _sp in dict.fromkeys(_speakers): if _sp in _voice_refs: continue _vd = os.path.join(_vroot, str(_sp).lower()) if not os.path.isdir(_vd): continue _vfiles = sorted(f for f in os.listdir(_vd) if f.lower().endswith(_vexts)) if _vfiles: _voice_refs[_sp] = os.path.join(_vd, _vfiles[0]) print(f"[JoyEcho] voice anchor {_sp!r}: auto-cast from folder " f"({_vfiles[0]}).", flush=True) if _voice_refs and not enable_audio_memory: print("[JoyEcho] voice anchors present but enable_audio_memory is off - " "anchors SKIPPED (the bank is the injection path).", flush=True) elif _voice_refs and memory_max_size > 0: from ltx_pipelines.utils.media_io import decode_audio_from_file from PIL import Image as _PILImage import av as _av for _vchar, _vpath in _voice_refs.items(): try: _aud = decode_audio_from_file(str(_vpath), torch.device("cpu"), 0.0, 6.0) if _aud is None: print(f"[JoyEcho] voice anchor {_vchar!r}: no audio stream " f"in {_vpath}; skipped.", flush=True) continue # conv_in expects 2 channels; the bank's own normalizer is the # canonical mono->stereo/downmix path. _wav = type(memory_bank)._normalize_waveform_channels( _aud.waveform, target_channels=2) _lat = audio_vae.encode(_wav, int(_aud.sampling_rate)) # Generated slots store the pipeline's bf16 latents; a # float32 seed would crash torch.cat at the next shot's # kwargs build (belt: get_memory_audio also harmonizes). _lat = _lat.to(torch.bfloat16).detach().cpu().contiguous() # Video half: 9 frames from the SAME clip, cover-fit to the # render size - the slot pairs the voice with the face that # was producing it, same contract as generated slots. _frames = [] _cont = _av.open(str(_vpath)) try: _vs = next((s for s in _cont.streams if s.type == "video"), None) if _vs is not None: _all = [f.to_image() for _i, f in zip(range(150), _cont.decode(_vs))] if _all: _step = max(1, len(_all) // 9) _frames = _all[::_step][:9] finally: _cont.close() if not _frames and _fp is not None: # Audio-only file (library wav/flac): pair the voice with # the character's REF IMAGE - the same face the video # lane is conditioned with - so the slot keeps its # face+voice contract instead of being skipped. try: _rr = os.path.join(_fp.get_input_directory(), "joyecho_refs", str(_vchar).lower()) _imgs = sorted( f for f in os.listdir(_rr) if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))) if _imgs: _frames = [_PILImage.open( os.path.join(_rr, _imgs[0])).convert("RGB")] print(f"[JoyEcho] voice anchor {_vchar!r}: audio-" f"only file; video half from ref image " f"{_imgs[0]}.", flush=True) except OSError: pass if not _frames: print(f"[JoyEcho] voice anchor {_vchar!r}: no video frames " f"and no ref image under joyecho_refs/" f"{str(_vchar).lower()}/; skipped.", flush=True) continue _tw, _th = int(video_width), int(video_height) _fitted = [] # loop var must NOT be _fp - that's the folder_paths alias, # and shadowing it silently killed the SECOND character's # audio-only ref-image fallback (audit finding, 2026-07-29) for _frm in _frames: _scale = max(_tw / _frm.width, _th / _frm.height) _rw = max(_tw, int(round(_frm.width * _scale))) _rh = max(_th, int(round(_frm.height * _scale))) if (_rw, _rh) != _frm.size: _frm = _frm.resize((_rw, _rh), _PILImage.LANCZOS) if (_rw, _rh) != (_tw, _th): _left = (_rw - _tw) // 2 _top = int((_rh - _th) * 0.25) _frm = _frm.crop((_left, _top, _left + _tw, _top + _th)) _fitted.append(_frm) while len(_fitted) < 9: _fitted.append(_fitted[-1]) # Seeded TWICE: memory is cross-attention context, and with # one slot the anchor only CONTESTS the text prior - measured # 2026-07-28 evening: same config, different seeds, shot 1 # flipped American then Australian. Doubling the anchor's # share of the context tips the coin toward the cast voice. # (The real fix is AV-extend continuation of the anchor at # each character's first shot - context can be ignored, a # waveform being extended cannot. Until that lands, weight.) for _rep in range(2): memory_bank.save_memory_slot( _fitted, _lat, audio_window_size=audio_memory_window_size, video_clip_num_frames=9, audio_waveform=_aud.waveform, audio_sample_rate=int(_aud.sampling_rate), video_fps=float(video_fps), audio_window_selection_mode="max_response", video_frame_selection_mode="center", character=str(_vchar), ) # Tag as anchor: the anchor+latest policy keeps these # slots live and protected from drift-outvoting. memory_bank.memory[-1].metadata["voice_anchor"] = True print(f"[JoyEcho] voice anchor {_vchar!r}: seeded from " f"{os.path.basename(str(_vpath))} " f"({_lat.shape[1]} audio tokens).", flush=True) except Exception as _vexc: print(f"[JoyEcho] voice anchor {_vchar!r} FAILED " f"({type(_vexc).__name__}: {_vexc}); continuing without.", flush=True) for shot_idx in range(num_shots): _shot_speaker = _speakers[shot_idx] if _speakers else None prompt_seed = seed + shot_idx conditional_dict = { k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in conditioning[shot_idx].items() if k != "joyecho_meta" # node-level metadata, not model input } print(f"[JoyEcho] Shot {shot_idx+1}/{num_shots}, seed={prompt_seed}, " f"memory_size={len(memory_bank)}", flush=True) # --- Phase A: Denoise (generator on GPU, VAE on CPU) --- _move(video_vae.encoder, "cpu") _move(video_vae.decoder, "cpu") _move(audio_vae.encoder, "cpu") _move(audio_vae.decoder, "cpu") _move(audio_vae.vocoder, "cpu") if sequential_offload: offloader.install() else: _move(generator, device) _empty_cache() with torch.random.fork_rng(devices=[device] if device.type == "cuda" else []): torch.manual_seed(prompt_seed) if device.type == "cuda": torch.cuda.manual_seed(prompt_seed) # References seed IDENTITY on shot 1 only. From shot 2 the bank # holds real rendered frames carrying identity AND staging; # re-injecting a wide full-scene reference every shot keeps # offering the whole room as alternative staging and makes the # subject relocate between shots (observed: the character in a different # spot per shot). Bank-only continuity held identity fine in # every bank-era run. _shot_refs = [c for c, s in zip(_ref_clips, _ref_sched) if s == shot_idx] if _shot_refs: print(f"[JoyEcho] injecting {len(_shot_refs)} reference(s) at shot {shot_idx+1}.", flush=True) if _shot_refs or len(memory_bank) > 0: # Encode memory frames (briefly bring video encoder to GPU). # References prepend as pure video conditioning; the bank # contributes only real generated shots. _mem_frames = list(_shot_refs) + (memory_bank.get_memory_frames() if len(memory_bank) > 0 else []) _move(video_vae.encoder, device) memory_video = encode_memory_frames_batch( video_vae=video_vae, batch_memory_frames=[_mem_frames], target_h=video_height, target_w=video_width, device=device, dtype=dtype, ) _move(video_vae.encoder, "cpu") _empty_cache() # Audio memory kwargs come from the BANK ONLY (never refs); # an empty bank means no audio-memory kwargs at all. memory_audio_kwargs = {} if len(memory_bank) > 0: memory_audio_kwargs = build_paired_audio_memory_kwargs( memory_bank, enable_audio_memory=enable_audio_memory, v2a_grad_scale=v2a_grad_scale, memory_position_mode="reference", speaker=_shot_speaker, ) # Reference clips are prepended as VIDEO-only slots, but the # audio side comes from the bank alone. The paired cross-mask # zips the two slot lists POSITIONALLY, and # _memory_slot_ranges_from_lengths falls back to an EVEN split # whenever len(segment_lengths) != num_slots - so a single # injected ref re-slices the whole audio memory at the wrong # boundaries and every stored voice is smeared across the # wrong slots. That is the "completely different voice on the # re-entry shot" bug, not drift. # # Fix: give each ref its own silent audio slot at the FRONT so # the counts match and the bank's real segments line up with # their own video slots again. Zero-LENGTH padding cannot work # here - _memory_slot_ranges_from_lengths drops empty ranges # (`if end > start`) and the alignment breaks a second time. if _shot_refs and memory_audio_kwargs: _npad = len(_shot_refs) _ma = memory_audio_kwargs.get("memory_audio") _segs = memory_audio_kwargs.get("memory_audio_segment_lengths") if _ma is not None and _segs: _PAD_T = 8 # tokens per ref slot; small, contributes silence _pad = torch.zeros( (_ma.shape[0], _PAD_T * _npad, _ma.shape[2]), dtype=_ma.dtype, device=_ma.device) memory_audio_kwargs["memory_audio"] = torch.cat([_pad, _ma], dim=1) memory_audio_kwargs["memory_audio_segment_lengths"] = tuple( tuple([_PAD_T] * _npad) + tuple(row) for row in _segs) memory_audio_kwargs["memory_audio_timestep"] = torch.zeros( memory_audio_kwargs["memory_audio"].shape[:2], dtype=torch.float32) print(f"[JoyEcho] paired-memory realign: {_npad} reference slot(s) " f"padded with silent audio so {len(_mem_frames)} video slots " f"match {len(_mem_frames)} audio slots.", flush=True) video_latent, audio_latent = memory_pipeline.generate( video_shape=tuple(video_shape), audio_shape=tuple(audio_shape), conditional_dict=conditional_dict, memory_video=memory_video, seed=prompt_seed, **memory_audio_kwargs, ) del memory_video else: video_latent, audio_latent = base_pipeline.generate( video_shape=tuple(video_shape), audio_shape=tuple(audio_shape), conditional_dict=conditional_dict, seed=prompt_seed, ) if device.type == "cuda": torch.cuda.synchronize() del conditional_dict _empty_cache() # Save audio latent for memory before decode moves things around. # NOTE: storage is deliberately NOT gated on enable_audio_memory — # the paired bank needs an audio slot to save the VIDEO slot, and # skipping the save silently disabled ALL cross-shot identity memory # whenever audio memory was off (diagnosed 2026-07-14: memory_size=0 # every shot). enable_audio_memory still gates the INJECTION path # (build_paired_audio_memory_kwargs), which is where the wrong-rate # drone bug lived, so storing here reintroduces no audio artifacts. audio_memory_latent = ( audio_latent.detach().cpu().contiguous() if audio_latent is not None else None ) # --- Phase B: Decode (generator off GPU, VAE on GPU) --- if sequential_offload: offloader.remove() _move(generator, "cpu") _empty_cache() _move(video_vae.decoder, device) _move(audio_vae.decoder, device) _move(audio_vae.vocoder, device) video_uint8, audio_waveform = decode_benchmark_sample( video_vae, audio_vae, video_latent, audio_latent, video_tiling_config=_decode_tiling_config, ) if device.type == "cuda": torch.cuda.synchronize() # Temporal x2: second decode from the shot's OWN latents through the # LTX temporal upsampler while the decoder is still on GPU. Only the # SAVED shot files (and thus the worker-built master) get the doubled # fps; the bank, refs and in-graph frames stay base-fps native. _tu_uint8 = None if _temporal_on: try: _tu_uint8 = self._temporal_upsample_decode( video_vae, video_latent, _decode_tiling_config, device) print(f"[JoyEcho] temporal x2: shot {shot_idx+1} " f"{video_uint8.shape[0]}f -> {_tu_uint8.shape[0]}f " f"(saved at {video_fps * 2} fps).", flush=True) except Exception as _tu_e: print(f"[JoyEcho] temporal x2 FAILED on shot {shot_idx+1}: " f"{_tu_e} - saving base-fps shot instead.", flush=True) _tu_uint8 = None # Move VAE back to CPU _move(video_vae.decoder, "cpu") _move(audio_vae.decoder, "cpu") _move(audio_vae.vocoder, "cpu") _empty_cache() # Update memory bank memory_frames_pil = video_uint8_to_pil_frames(video_uint8) if audio_memory_latent is not None: memory_bank.save_memory_slot( memory_frames_pil, audio_memory_latent, audio_window_size=audio_memory_window_size, video_clip_num_frames=9, audio_waveform=audio_waveform, audio_sample_rate=16000, video_fps=float(video_fps), audio_window_selection_mode="max_response", video_frame_selection_mode=( "last" if str(memory_video_anchor) == "last" else "center"), audio_memory_mel_bins=128, audio_memory_mel_hop_length=160, audio_memory_n_fft=1024, audio_memory_downsample_factor=4, audio_memory_is_causal=True, character=_shot_speaker, ) # HEAD TRIM: each shot's first frames morph out of the memory / # reference content (a split-second flash of the reference image). # Trim the DECODED shot ONCE, up front, with matching audio, so the # trim reaches EVERYTHING downstream: the final concatenated output, # the per-shot preview files, and any external concat of them. (The # memory-bank save above uses center-frame selection, so it never # saw the head flash and is left on the full shot.) _trim = max(0, int(head_trim_frames)) if _trim == 0 and _ref_clips: _trim = 8 # auto when references are wired if _shot_refs: # A shot that RECEIVES references (its character's entry shot) # opens morphing out of the reference and needs the longer trim. _t1 = max(0, int(head_trim_first_shot)) if _t1 == 0: _t1 = max(2 * _trim, 24) # auto _trim = max(_trim, _t1) if _trim > 0 and video_uint8.shape[0] > _trim + 16: video_uint8 = video_uint8[_trim:] if _tu_uint8 is not None: _tu_uint8 = _tu_uint8[2 * _trim:] # same cut in doubled frames if audio_waveform is not None: _cut = int(round(_trim / float(video_fps) * audio_sample_rate)) if audio_waveform.shape[-1] > _cut: audio_waveform = audio_waveform[..., _cut:] print(f"[JoyEcho] head-trim: dropped {_trim} frames from shot {shot_idx+1} " f"(all outputs + per-shot preview).", flush=True) else: _trim = 0 # Collect outputs (trimmed) as uint8 [F,H,W,3]. Float conversion is # deferred to final assembly (progressive fill there): storing shots # as float32 cost 4x the RAM, and at 10-shot hires the accumulated # float frames (~75GB at 1920x1088) plus the offloader's pinned host # pages hard-crashed the box at end-of-refine (2026-07-19). all_video_frames.append(video_uint8) if hires_factor > 1.0: # keep the shot's audio latent for the joint hires refine pass _hires_audio_lats.append( audio_memory_latent.clone() if audio_memory_latent is not None else None) if _hires_spatial: # spatial mode upscales the NATIVE latents (no VAE # roundtrip, no re-generation) - stash them pre-decode, # plus this shot's head-trim so the upscaled decode can be # cut identically (latents are untrimmed) _hires_video_lats.append( (video_latent.detach().to("cpu", torch.float32).contiguous(), int(_trim))) if audio_waveform is not None: from ltx_distillation.inference.memory_multishot import normalize_audio_waveform_for_media all_audio_waveforms.append(normalize_audio_waveform_for_media(audio_waveform)) # Save per-shot video immediately for real-time preview (now trimmed, # so it matches the final output frame-for-frame). if _tu_uint8 is not None: # doubled-fps shot file: the AutoFinish worker probes r_frame_rate # per shot, so the master assembles at the doubled rate untouched self._save_shot_video( _tu_uint8, audio_waveform, shot_idx, video_fps * 2, audio_sample_rate, output_prefix ) del _tu_uint8 else: self._save_shot_video( video_uint8, audio_waveform, shot_idx, video_fps, audio_sample_rate, output_prefix ) del video_latent, audio_latent, audio_memory_latent, video_uint8, audio_waveform _empty_cache() print(f"[JoyEcho] Shot {shot_idx+1}/{num_shots} done.", flush=True) # --- Optional hires-fix second pass: refine every shot at higher res. # Runs AFTER the loop so the memory bank and per-shot previews stay at # base res, and a failure here falls back to the base frames instead of # killing the render. if hires_factor > 1.0: try: if _hires_spatial: self._hires_spatial_pass( video_vae=video_vae, audio_vae=audio_vae, video_lats=_hires_video_lats, frames_list=all_video_frames, device=device, waveforms=all_audio_waveforms, fps=video_fps, audio_sr=audio_sample_rate, out_prefix=str(output_prefix) + "_hires", ) else: self._hires_refine_pass( generator=generator, video_vae=video_vae, audio_vae=audio_vae, conditioning=conditioning, frames_list=all_video_frames, audio_lats=_hires_audio_lats, device=device, factor=float(hires_factor), mode=str(hires_denoise), sequential_offload=bool(sequential_offload), resident_blocks=int(resident_blocks), waveforms=all_audio_waveforms, fps=video_fps, audio_sr=audio_sample_rate, out_prefix=str(output_prefix) + "_hires", base_seed=int(seed), ) except Exception: import traceback traceback.print_exc() print("[JoyEcho] WARNING: hires refine pass FAILED - delivering the " "base-resolution frames instead.", flush=True) finally: _hires_audio_lats.clear() _hires_video_lats.clear() _empty_cache() # Concatenate all shots with the selected boundary treatment. xf = max(1, int(transition_frames)) if transition == "dissolve" else 0 paired_audio = bool(all_audio_waveforms) and len(all_audio_waveforms) == len(all_video_frames) if xf > 0 and len(all_video_frames) > 1: vids = all_video_frames auds = all_audio_waveforms if paired_audio else None # Shots are stored uint8; the blend math needs float 0..1 - # convert lazily as each shot is consumed. out_v = vids[0].float().div_(255.0) out_a = auds[0] if auds else None for i in range(1, len(vids)): b_v = vids[i].float().div_(255.0) n = min(xf, out_v.shape[0], b_v.shape[0]) if n <= 0: out_v = torch.cat([out_v, b_v], dim=0) if auds is not None: out_a = torch.cat([out_a, auds[i]], dim=-1) continue w = torch.linspace(0.0, 1.0, n, dtype=out_v.dtype).view(n, 1, 1, 1) blend = out_v[-n:] * (1.0 - w) + b_v[:n] * w out_v = torch.cat([out_v[:-n], blend, b_v[n:]], dim=0) if auds is not None: b_a = auds[i] n_s = min(int(round(n / float(video_fps) * audio_sample_rate)), out_a.shape[-1], b_a.shape[-1]) if n_s > 0: t = torch.linspace(0.0, 1.0, n_s, dtype=out_a.dtype) fade_out = torch.cos(t * torch.pi / 2.0) fade_in = torch.sin(t * torch.pi / 2.0) a_blend = out_a[..., -n_s:] * fade_out + b_a[..., :n_s] * fade_in out_a = torch.cat([out_a[..., :-n_s], a_blend, b_a[..., n_s:]], dim=-1) else: out_a = torch.cat([out_a, b_a], dim=-1) images = out_v print(f"[JoyEcho] Crossfaded {len(vids)-1} shot boundaries ({xf} frames each).", flush=True) audio_out = None if paired_audio: audio_out = {"waveform": out_a.unsqueeze(0), "sample_rate": audio_sample_rate} elif all_audio_waveforms: combined_waveform = torch.cat(all_audio_waveforms, dim=-1) audio_out = {"waveform": combined_waveform.unsqueeze(0), "sample_rate": audio_sample_rate} else: # Progressive uint8 -> float fill. A plain torch.cat(...).float() # holds the uint8 total AND the float total simultaneously; this # allocates the final float tensor once and frees each uint8 shot # as it lands, so peak = float total + one shot. _shot_lengths = [int(v.shape[0]) for v in all_video_frames] _total_f = sum(_shot_lengths) _H0 = int(all_video_frames[0].shape[1]) _W0 = int(all_video_frames[0].shape[2]) images = torch.empty((_total_f, _H0, _W0, 3), dtype=torch.float32) _pos = 0 for _i in range(len(all_video_frames)): _v = all_video_frames[_i] images[_pos:_pos + _v.shape[0]] = _v.float().div_(255.0) _pos += _v.shape[0] all_video_frames[_i] = _v[:0] # drop the uint8 payload audio_out = None if all_audio_waveforms: combined_waveform = torch.cat(all_audio_waveforms, dim=-1) # [2, total_samples] audio_out = { "waveform": combined_waveform.unsqueeze(0), # [1, 2, samples] "sample_rate": audio_sample_rate, } # VHS GLITCH transition: corrupt the frames AROUND each boundary in # place (snow, tearing bands, dropout lines) + a tape-static audio # hit. Total length unchanged; deterministic per seed+boundary. if transition == "vhs_glitch" and len(all_video_frames) > 1: n = max(1, int(transition_frames)) amt_base = float(max(0.1, min(1.0, glitch_intensity))) boundaries = [] acc = 0 # _shot_lengths, not the list shapes - the progressive fill # above empties each shot tensor after copying it out. for _len in _shot_lengths[:-1]: acc += _len boundaries.append(acc) # first frame index of the NEXT shot total_f = images.shape[0] H, W = images.shape[1], images.shape[2] for bi, b in enumerate(boundaries): g = torch.Generator().manual_seed(int(seed) * 1009 + bi) start = max(0, b - n // 2) end = min(total_f, start + n) span = max(1, end - start - 1) for k, fidx in enumerate(range(start, end)): env = 1.0 - abs((k - span / 2.0) / (span / 2.0 or 1.0)) amt = amt_base * (0.35 + 0.65 * max(0.0, env)) f = images[fidx] # snow (monochrome noise mix) snow = torch.rand((H, W, 1), generator=g).expand(H, W, 3) f = f * (1.0 - amt * 0.8) + snow * (amt * 0.8) # horizontal tearing bands for _ in range(int(1 + amt * 6)): y0 = int(torch.randint(0, max(1, H - 8), (1,), generator=g)) bh = int(torch.randint(2, max(3, H // 20), (1,), generator=g)) dx = int(torch.randint(-W // 6, W // 6 + 1, (1,), generator=g)) f[y0:y0 + bh] = torch.roll(f[y0:y0 + bh], shifts=dx, dims=1) # dropout scanlines for _ in range(int(amt * 4)): y = int(torch.randint(0, H, (1,), generator=g)) f[y:y + 1] = float(torch.rand((1,), generator=g)) images[fidx] = f.clamp(0.0, 1.0) # audio: tape-static bed over a WIDER window than the video # burst - JoyAI's per-shot room tone fades out at shot edges, # so the static must SPAN that dead seam (>=1.2s), not just # the few glitched frames, or the cut reads as burst->dead # air->tone. Quieter hit with a smooth raised-cosine envelope # (the old short triangular hit at 0.22 was sharp and loud). if audio_out is not None: wav = audio_out["waveform"][0] # [2, samples] c = int(round(b / float(video_fps) * audio_sample_rate)) n_s = max(int(round(n / float(video_fps) * audio_sample_rate)), int(round(1.2 * audio_sample_rate))) s0 = max(0, c - n_s // 2) s1 = min(wav.shape[-1], s0 + n_s) if s1 > s0: ln = s1 - s0 t = torch.linspace(0.0, 1.0, ln) env_a = 0.5 - 0.5 * torch.cos(t * 2.0 * torch.pi) # raised cosine noise = (torch.rand((wav.shape[0], ln), generator=g) * 2.0 - 1.0) wav[..., s0:s1] = (wav[..., s0:s1] * (1.0 - 0.35 * amt_base * env_a) + noise * (0.10 * amt_base) * env_a).clamp(-1.0, 1.0) print(f"[JoyEcho] VHS glitch applied at {len(boundaries)} boundaries " f"({n} frames, intensity {amt_base}).", flush=True) # Sidecar for the AutoFinish worker: the worker rebuilds the master from # the PRE-glitch per-shot files and cannot see these widgets, so record # them next to the shots. Written for every transition type so the # worker honors "cut"/"dissolve" too (its own defaults only apply when # no same-run sidecar exists). try: import json as _json import folder_paths as _fpaths _parts = str(output_prefix).rsplit("/", 1) _sdir = os.path.join(_fpaths.get_output_directory(), _parts[0] if len(_parts) == 2 else "") os.makedirs(_sdir, exist_ok=True) with open(os.path.join(_sdir, "_transition.json"), "w", encoding="utf-8") as _fh: _json.dump({"transition": str(transition), "frames": int(transition_frames), "intensity": float(glitch_intensity), "seed": int(seed)}, _fh) except Exception as _e: # noqa: BLE001 print(f"[JoyEcho] WARNING: transition sidecar not written ({_e}).", flush=True) print(f"[JoyEcho] Generation complete. {images.shape[0]} frames, " f"{num_shots} shot(s).", flush=True) return (images, audio_out,) def _hires_spatial_pass(self, *, video_vae, audio_vae, video_lats, frames_list, device, waveforms=None, fps=25, audio_sr=48000, out_prefix=""): """Deterministic hires: run each shot's NATIVE latents through the LTX x1.5 spatial upscaler (one forward pass, no noise, no re-generation), then decode at the higher resolution. The standard-workflow upscale path, minus the VAE roundtrip the post-pass version pays. Temporal reflect-pad mirrors the RIFT padded upsampler (the x1.5-1.0 model corrupts sequence ends; TPAD=8 latents puts the damage in disposable mirrors). Factor is FIXED at 1.5 by the model - hires_factor is ignored beyond enabling the pass. x1.5 needs EVEN spatial latent dims (heights like 704/768, not 736) or it fabricates an edge band. """ import time as _time import json as _json import folder_paths import comfy.utils as _cutils from comfy.ldm.lightricks.latent_upsampler import LatentUpsampler from ltx_core.model.video_vae import TemporalTilingConfig, TilingConfig from ltx_distillation.utils import decode_benchmark_sample _t0 = _time.time() _dt = torch.bfloat16 UPSCALER_FILE = "ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors" TPAD = 8 path = folder_paths.get_full_path_or_raise("latent_upscale_models", UPSCALER_FILE) sd, metadata = _cutils.load_torch_file(path, safe_load=True, return_metadata=True) up_model = LatentUpsampler.from_config(_json.loads(metadata["config"])).to(dtype=_dt) up_model.load_state_dict(sd) up_model.eval().to(device) del sd stats = video_vae.decoder.per_channel_statistics _move(video_vae.decoder, device) tiling = TilingConfig( spatial_config=None, temporal_config=TemporalTilingConfig(tile_size_in_frames=64, tile_overlap_in_frames=24)) try: for si, (lat, trim) in enumerate(video_lats): # [1,F,C,h,w] normalized (pipeline layout) -> [1,C,F,h,w] raw lat = lat.permute(0, 2, 1, 3, 4) h_lat, w_lat = int(lat.shape[3]), int(lat.shape[4]) if h_lat % 2 or w_lat % 2: print(f"[JoyEcho] Hires-spatial WARNING: odd latent dim " f"(h={h_lat}, w={w_lat}) - the x1.5 model will smear one " f"edge. Use heights/widths whose /32 is EVEN (768, not 736).", flush=True) # stats is a decoder submodule - it moved to device with it raw = stats.un_normalize(lat.to(device)) # temporal reflect pad (no edge repeat), crop after f = raw.shape[2] p = min(TPAD, f - 1) if p > 0: raw = torch.cat([raw[:, :, 1:p + 1].flip(2), raw, raw[:, :, -(p + 1):-1].flip(2)], dim=2) with torch.no_grad(): up = up_model(raw.to(device=device, dtype=_dt)) del raw if p > 0: up = up[:, :, p:-p] up = stats.normalize(up.to(torch.float32)).to(_dt) up = up.permute(0, 2, 1, 3, 4).contiguous() # back to [1,F,C,H,W] _empty_cache() u8, _ = decode_benchmark_sample(video_vae, audio_vae, up.to(device), None, video_tiling_config=tiling) del up _empty_cache() if trim > 0 and u8.shape[0] > trim: u8 = u8[trim:] frames_list[si] = u8 if out_prefix: try: _wav = (waveforms[si] if waveforms is not None and si < len(waveforms) else None) self._save_shot_video(u8, _wav, si, fps, audio_sr, out_prefix) except Exception as _e: print(f"[JoyEcho] Hires-spatial: per-shot save failed for " f"shot {si+1}: {_e}", flush=True) print(f"[JoyEcho] Hires-spatial: shot {si+1}/{len(video_lats)} " f"upscaled ({u8.shape[2]}x{u8.shape[1]}).", flush=True) finally: up_model.cpu() _move(video_vae.decoder, "cpu") _empty_cache() print(f"[JoyEcho] Hires-spatial pass done in {_time.time()-_t0:.0f}s.", flush=True) def _hires_refine_pass(self, *, generator, video_vae, audio_vae, conditioning, frames_list, audio_lats, device, factor, mode, sequential_offload, resident_blocks, waveforms=None, fps=25, audio_sr=48000, out_prefix="", base_seed=0): """Second-pass hires fix: per shot, upscale (bicubic) -> VAE re-encode -> re-noise at a tail sigma -> re-denoise through the DMD ladder tail at the TARGET resolution -> decode. The model synthesizes genuine detail the base render never had (RTX-class upscalers only sharpen what exists). Video-only output: audio latents ride along for the joint AV model's cross-attention, then the refined audio is discarded (the original track is untouched). Runs in 65-frame windows with an 8-frame linear cross-fade so a 24GB card survives 1920x1088 refines. """ import time as _time from PIL import Image as _PILImage import numpy as _np from ltx_distillation.utils import add_noise, decode_benchmark_sample, frames_to_video_tensor from ltx_core.model.video_vae import TemporalTilingConfig, TilingConfig _t0 = _time.time() _dt = torch.bfloat16 # fixed: never derive dtype from parameters (fp8 trap) SIGMA_TAILS = {"subtle": [0.421875, 0.0], "medium": [0.725, 0.421875, 0.0], # tenstrip's published UPSCALE ladder (HF, ~Jul 1) - # purpose-built for a refine pass, unlike the generation # tails above. (Shallower/frozen/reused-noise variants # were tested 2026-07-23 and removed: none reduced the # per-frame detail churn - see project notes.) "strong": [0.92, 0.725, 0.421875, 0.0]} sigmas = SIGMA_TAILS.get(str(mode).split()[0].lower(), SIGMA_TAILS["subtle"]) base_h, base_w = int(frames_list[0].shape[1]), int(frames_list[0].shape[2]) th = max(32, int(round(base_h * factor / 32.0)) * 32) tw = max(32, int(round(base_w * factor / 32.0)) * 32) if (th, tw) == (base_h, base_w): print("[JoyEcho] Hires: factor rounds to the base resolution; skipping.", flush=True) return # 33-frame windows (8n+1): a 65f window at 1920x1088 hard-aborts the VAE # ENCODER (uncatchable cuDNN abort - the encode sibling of the decode # crash that tiled decode fixed). 33f keeps encode peaks ~half the # proven-safe decode chunk size. WIN, OVL = 33, 9 print(f"[JoyEcho] Hires refine: {base_w}x{base_h} -> {tw}x{th}, " f"start sigma {sigmas[0]} ({len(sigmas)-1} step(s)), " f"{WIN}f windows / {OVL}f blend.", flush=True) tiling = TilingConfig( spatial_config=None, temporal_config=TemporalTilingConfig(tile_size_in_frames=64, tile_overlap_in_frames=24)) # Device setup once for the whole pass. resident_blocks is deliberately # 0 here: pinned DiT blocks (~12GB at 30 resident) compete with the VAE # encoder's conv activations at target res - exactly what pushed the # 65f encode into a fatal abort. The refine is 1-2 steps per window, so # streaming all blocks costs seconds while freeing the VRAM the VAE needs. # The wrapper caches the RENDER resolution (video_frame_seqlen drives the # per-token timestep expansion; latent_height/width drive unpatchify). # generate() sets them once from the base res, so a refine at target res # feeds 2040-token frames against a 1008-token modulation vector -> # "size of tensor a (10200) must match tensor b (5040)". Retarget for the # pass, restore in the finally so the memory bank / next queue item are # unaffected. _res_saved = (generator.video_height, generator.video_width, generator.latent_height, generator.latent_width, generator.video_frame_seqlen) generator.video_height = th generator.video_width = tw generator.latent_height = th // 32 generator.latent_width = tw // 32 generator.video_frame_seqlen = generator.latent_height * generator.latent_width generator.VIDEO_FPS = float(fps) # refine RoPE clock = render fps too offl = None if sequential_offload: offl = SequentialOffloader(generator, device, resident_blocks=0) offl.install() else: _move(generator, device) _move(video_vae.encoder, device) _move(video_vae.decoder, device) try: for si in range(len(frames_list)): if audio_lats[si] is None: print(f"[JoyEcho] Hires: shot {si+1} has no audio latent; skipped.", flush=True) continue vf = frames_list[si] # [T,H,W,3] uint8 T = int(vf.shape[0]) if T < WIN: print(f"[JoyEcho] Hires: shot {si+1} shorter than a window ({T}f); skipped.", flush=True) continue cond = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in conditioning[si].items() if k != "joyecho_meta"} alat_full = audio_lats[si] # [1,Fa,C] cpu Fa = int(alat_full.shape[1]) starts = list(range(0, T - WIN + 1, WIN - OVL)) if starts[-1] + WIN < T: starts.append(T - WIN) out = torch.zeros((T, th, tw, 3), dtype=torch.float32) wsum = torch.zeros((T, 1, 1, 1), dtype=torch.float32) # HIRES V2 (2026-07-23): every window of a shot slices from ONE # per-shot noise field per noising event, indexed by global # latent position, instead of drawing independent per-window # noise. Independent noise was the structural failure of v1: # each window synthesized a different texture identity and the # cross-fade blended two textures instead of reconciling one # (measured detail-correlation collapse at exactly the window # stride). Fields are created lazily at the first window (latent # shapes are known only after the first VAE encode), seeded # per-shot, kept on CPU (~45MB each) and sliced to GPU per window. shot_fields = None for wi, s in enumerate(starts): win = vf[s:s + WIN].float().div_(255.0) # uint8 -> float 0..1 x = win.permute(0, 3, 1, 2) # [33,3,H,W] x = torch.nn.functional.interpolate( x, size=(th, tw), mode="bicubic", antialias=True).clamp(0, 1) pil = [_PILImage.fromarray( (f.permute(1, 2, 0) * 255).round().to(torch.uint8).numpy()) for f in x] del x fv = frames_to_video_tensor(pil, th, tw).unsqueeze(0).to( device=device, dtype=_dt) # [1,3,65,th,tw] in [-1,1] del pil lat = video_vae.encode(fv).permute(0, 2, 1, 3, 4).to(_dt) # [1,Fl,C,h,w] del fv _empty_cache() Fl = int(lat.shape[1]) a0 = int(round(Fa * s / float(T))) a1 = max(a0 + 2, int(round(Fa * (s + WIN) / float(T)))) alat = alat_full[:, a0:min(a1, Fa)].to(device=device, dtype=_dt) Fa_w = int(alat.shape[1]) if shot_fields is None: n_events = len(sigmas) - 1 # initial noising + each re-corruption NG = T // 8 + Fl + 2 # global latent-frame budget with margin # The refine noise was seeded from a CONSTANT (1234) for # months - every render's detail layer shared identical # noise regardless of the seed widget (audit, 2026-07-29). # base_seed routes the user's seed in; the +1234 offset # keeps refine noise decorrelated from the base pass. gs = torch.Generator(device="cpu").manual_seed( int(base_seed) + 1234 + si * 100) shot_fields = { "v": [torch.randn((1, NG) + tuple(lat.shape[2:]), generator=gs) for _ in range(n_events)], "a": [torch.randn((1, Fa) + tuple(alat.shape[2:]), generator=gs) for _ in range(n_events)], } # window latent j -> global latent ls+j; window starts are # multiples of (WIN-OVL)=24 pixel frames = 3 latent frames, # so overlapping windows land on the same global indices. ls = min(int(round(s / 8.0)), shot_fields["v"][0].shape[1] - Fl) nv = shot_fields["v"][0][:, ls:ls + Fl].to(device=device, dtype=_dt) na = shot_fields["a"][0][:, a0:a0 + Fa_w].to(device=device, dtype=_dt) s0 = torch.full((1, Fl), sigmas[0]) sa0 = torch.full((1, Fa_w), sigmas[0]) v = add_noise(lat.flatten(0, 1), nv.flatten(0, 1), s0.flatten(0, 1)).unflatten(0, (1, Fl)) a = add_noise(alat, na, sa0) del lat, nv, na for i_s, sig in enumerate(sigmas[:-1]): vs = sig * torch.ones((1, Fl), device=device) as_ = sig * torch.ones((1, Fa_w), device=device) pred_v, pred_a = generator( noisy_image_or_video=v, conditional_dict=cond, timestep=vs, noisy_audio=a, audio_timestep=as_) nxt = sigmas[i_s + 1] if nxt > 0: # re-corruption noise from the step's shared field, # same global slice - never randn_like (that was # fresh unseeded noise per window per step in v1) fnv = shot_fields["v"][i_s + 1][:, ls:ls + Fl].to( device=device, dtype=_dt) fna = shot_fields["a"][i_s + 1][:, a0:a0 + Fa_w].to( device=device, dtype=_dt) nvs = nxt * torch.ones((1, Fl), device=device) nas = nxt * torch.ones((1, Fa_w), device=device) v = add_noise(pred_v.flatten(0, 1), fnv.flatten(0, 1), nvs.flatten(0, 1)).unflatten(0, (1, Fl)) a = add_noise(pred_a, fna, nas) else: v = pred_v a = pred_a del a u8w, _ = decode_benchmark_sample(video_vae, audio_vae, v, None, video_tiling_config=tiling) del v _empty_cache() w = torch.ones((WIN, 1, 1, 1), dtype=torch.float32) if s > 0: w[:OVL, 0, 0, 0] = torch.linspace(0.0, 1.0, OVL) out[s:s + WIN] += (u8w.float() / 255.0) * w wsum[s:s + WIN] += w del u8w refined = (out / wsum.clamp_min(1e-6)).mul_(255.0).round_().clamp_(0, 255).to(torch.uint8) frames_list[si] = refined del out, wsum, cond _empty_cache() # Persist the refined shot IMMEDIATELY, same crash-safety # contract as the base per-shot saves: a crash later in the # refine (or in final assembly) costs at most one shot's # refine, never the pass. 2026-07-19: the box hard-crashed at # end-of-refine and all refined frames (RAM-only) were lost. if out_prefix: try: _wav = (waveforms[si] if waveforms is not None and si < len(waveforms) else None) self._save_shot_video(refined, _wav, si, fps, audio_sr, out_prefix) except Exception as _e: print(f"[JoyEcho] Hires: per-shot save failed for shot " f"{si+1}: {_e}", flush=True) print(f"[JoyEcho] Hires: shot {si+1}/{len(frames_list)} refined " f"({len(starts)} windows).", flush=True) finally: (generator.video_height, generator.video_width, generator.latent_height, generator.latent_width, generator.video_frame_seqlen) = _res_saved if offl is not None: offl.remove() _move(generator, "cpu") _move(video_vae.encoder, "cpu") _move(video_vae.decoder, "cpu") _empty_cache() print(f"[JoyEcho] Hires refine pass done in {_time.time()-_t0:.0f}s.", flush=True) def _temporal_upsample_decode(self, video_vae, video_latent, tiling_config, device): """Decode a shot's video latents through the LTX temporal x2 upsampler. Mirrors the hires-spatial pass: pipeline latents [1,F,C,h,w] -> un-normalize -> reflect-pad T (the upsampler family corrupts sequence ends; pads absorb it) -> upsample -> crop pad -> re-normalize -> tiled decode. Deterministic, video-only - the audio lane is never touched. Returns uint8 frames with F_out ~= 2*F-1. Decoder must already be on `device` (call inside Phase B). """ import json as _json import folder_paths import comfy.utils as _cutils from comfy.ldm.lightricks.latent_upsampler import LatentUpsampler from ltx_core.model.video_vae import TemporalTilingConfig, TilingConfig from ltx_distillation.utils import decode_benchmark_sample # The doubled sequence (~2x frames) MUST decode temporally tiled no # matter what the shot's own tiling resolved to - an untiled 449-frame # decode is exactly the uncatchable native cuDNN abort the tiled-decode # work fixed (killed the whole process on first try, 2026-07-30). tiling_config = TilingConfig( spatial_config=None, temporal_config=TemporalTilingConfig(tile_size_in_frames=64, tile_overlap_in_frames=24)) _dt = torch.bfloat16 cls = type(self) if getattr(cls, "_tu_model", None) is None: path = folder_paths.get_full_path_or_raise( "latent_upscale_models", "ltx-2.3-temporal-upscaler-x2-1.0.safetensors") sd, metadata = _cutils.load_torch_file( path, safe_load=True, return_metadata=True) m = LatentUpsampler.from_config( _json.loads(metadata["config"])).to(dtype=_dt) m.load_state_dict(sd) m.eval() cls._tu_model = m del sd up_model = cls._tu_model.to(device) stats = video_vae.decoder.per_channel_statistics # on device with decoder raw = video_latent.permute(0, 2, 1, 3, 4) # [1,C,F,h,w] raw = stats.un_normalize(raw.to(device=device, dtype=torch.float32)) f_in = raw.shape[2] p = min(8, f_in - 1) if p > 0: raw = torch.cat([raw[:, :, 1:p + 1].flip(2), raw, raw[:, :, -(p + 1):-1].flip(2)], dim=2) with torch.no_grad(): up = up_model(raw.to(dtype=_dt)) _scale = up.shape[2] / raw.shape[2] del raw if p > 0: _po = int(round(p * _scale)) up = up[:, :, _po:up.shape[2] - _po] up = stats.normalize(up.to(torch.float32)).to(_dt) up = up.permute(0, 2, 1, 3, 4).contiguous() # [1,F2,C,h,w] u8, _ = decode_benchmark_sample(video_vae, None, up.to(device), None, video_tiling_config=tiling_config) del up up_model.to("cpu") return u8 @staticmethod def _save_shot_video(video_uint8, audio_waveform, shot_idx, fps, audio_sr, prefix): """Save a single shot as mp4 immediately after generation.""" import av import numpy as np try: import folder_paths output_dir = folder_paths.get_output_directory() except Exception: output_dir = Path("/root/ComfyUI/output") # Build output path parts = prefix.rsplit("/", 1) if len(parts) == 2: sub_dir = Path(output_dir) / parts[0] name_prefix = parts[1] else: sub_dir = Path(output_dir) name_prefix = prefix sub_dir.mkdir(parents=True, exist_ok=True) out_path = sub_dir / f"{name_prefix}_{shot_idx:03d}.mp4" frames_np = video_uint8.cpu().numpy() if isinstance(video_uint8, torch.Tensor) else video_uint8 container = av.open(str(out_path), mode="w") stream = container.add_stream("h264", rate=fps) stream.height = frames_np.shape[1] stream.width = frames_np.shape[2] stream.pix_fmt = "yuv420p" # crf18/fast produced visible per-frame quality pumping on the # grain-heavy analog-horror content (2026-07-19, output_00153): # preset fast enables b-pyramid, so the GOP alternates # well-fed P / reference-B frames with starved outer B-frames - # sharp/soft flicker with a period of 2, measured at 14/14 # sign-flips in per-frame Laplacian variance. bf=0 removes # B-frames entirely (the only full cure), tune=grain keeps the # noise field stable across frames, crf 16 + medium feeds it. # These per-shot files are the MASTERS (finals should be # stream-copy concats of them - SaveVideo exposes no quality # knobs); the ~40% size increase is the cost of no pumping. stream.options = {"crf": "16", "preset": "medium", "tune": "grain", "bf": "0"} for frame_data in frames_np: frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24") for packet in stream.encode(frame): container.mux(packet) for packet in stream.encode(): container.mux(packet) container.close() # Save audio sidecar. DELIBERATE ORDER: soundfile FIRST, torchaudio never. # torchaudio.save on torch 2.8+ dispatches to torchcodec, which dlopens # the system FFmpeg shared libs (libavutil.so.60 = FFmpeg 7/8). On a box # whose FFmpeg is older (Ubuntu 24.04 ships 6.x = libavutil.so.58) that # dlopen does not always raise a catchable Python error - it can trip the # glibc loader assertion "ld.so: Assertion `listp != NULL' failed!", which # ABORTS the process. A try/except cannot save you from that: ComfyUI just # dies, mid-run, after shots have already rendered (reported 2026-07-24). # soundfile bundles libsndfile and never touches FFmpeg, so it cannot trip # it; the builtin wave module has no deps at all. Neither needs torchaudio, # which buys nothing for a plain PCM WAV write. Do not "restore" torchaudio # to the front of this chain. if audio_waveform is not None: wav_path = sub_dir / f"{name_prefix}_{shot_idx:03d}.wav" waveform = audio_waveform.cpu() if waveform.dim() == 1: waveform = waveform.unsqueeze(0) _saved = False try: import soundfile as _sf _sf.write(str(wav_path), waveform.transpose(0, 1).numpy(), int(audio_sr)) _saved = True except Exception as _e_sf: import wave as _wave pcm = (waveform.clamp(-1.0, 1.0) * 32767.0).round().to( torch.int16).transpose(0, 1).contiguous().numpy() with _wave.open(str(wav_path), "wb") as _w: _w.setnchannels(int(waveform.shape[0])) _w.setsampwidth(2) _w.setframerate(int(audio_sr)) _w.writeframes(pcm.tobytes()) _saved = True print(f"[JoyEcho] wav via builtin wave, PCM16 (soundfile " f"unavailable: {type(_e_sf).__name__}). " f"pip install soundfile for float-wav output.", flush=True) print(f"[JoyEcho] Shot {shot_idx} saved → {out_path}", flush=True) class JoyEcho_SingleShotGenerate: """Generate a single shot with memory bank input/output for chaining. Each instance has its own editable prompt text box and outputs video frames that can be previewed immediately via CreateVideo → SaveVideo. Chain multiple instances via the memory output → next shot's memory input. """ @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("JOYECHO_MODEL",), "prompt": ("STRING", { "multiline": True, "default": "", "tooltip": "Single shot prompt text", }), "seed": ("INT", {"default": 12345, "min": 0, "max": 2**31 - 1}), "num_frames": ("INT", {"default": 241, "min": 9, "max": 1441, "step": 8, "tooltip": "Must be 1 + 8*k (e.g. 121, 241, 361)"}), # Rebels local patch: portrait resolutions. Height was capped at # 1088 while width allowed 1920, which silently forbade portrait # (e.g. 1088x1920 to match a portrait Z-Image first frame). Both # axes now cap at 1920; step 32 keeps the latent packing valid. "video_height": ("INT", {"default": 736, "min": 256, "max": 1920, "step": 32}), "video_width": ("INT", {"default": 1280, "min": 256, "max": 1920, "step": 32}), }, "optional": { "memory": ("JOYECHO_MEMORY",), "video_fps": ("INT", {"default": 24, "min": 1, "max": 60, "tooltip": "KEEP AT 24. The LTX joint audio-video prior is 24fps-native: any other value (25 included) systematically drifts spoken voices toward Commonwealth accents (British/Australian) and overrides accent wording in the prompt. Verified A/B 2026-07-29."}), "v2a_grad_scale": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.1}), "memory_max_size": ("INT", {"default": 7, "min": 0, "max": 20}), "num_fix_frames": ("INT", {"default": 3, "min": 0, "max": 10}), "enable_audio_memory": ("BOOLEAN", { "default": True, "tooltip": "Feed previous shots' audio latents as cross-shot conditioning. " "ON keeps the same voice across every shot - the point of " "multishot - at a small lip-sync cost on long dialogue. OFF " "gives the tightest sync but the voice can drift between " "shots; use only for sync-critical single-voice pieces with " "a strong voice description repeated in every shot.", }), "audio_memory_window_size": ("INT", {"default": 96, "min": 16, "max": 256}), "sequential_offload": ("BOOLEAN", { "default": False, "tooltip": "Enable layer-by-layer GPU offloading for DiT.", }), }, } RETURN_TYPES = ("IMAGE", "AUDIO", "JOYECHO_MEMORY", "JOYECHO_MODEL",) RETURN_NAMES = ("images", "audio", "memory", "model",) FUNCTION = "generate_shot" CATEGORY = "JoyAI-Echo" def generate_shot( self, model: dict, prompt: str, seed: int = 12345, num_frames: int = 241, video_height: int = 736, video_width: int = 1280, memory: dict | None = None, video_fps: int = 25, v2a_grad_scale: float = 2.0, memory_max_size: int = 7, num_fix_frames: int = 3, enable_audio_memory: bool = True, audio_memory_window_size: int = 96, sequential_offload: bool = False, ): from ltx_distillation.inference.bidirectional_pipeline import BidirectionalAVInferencePipeline from ltx_distillation.inference.memory_bidirectional_pipeline import BidirectionalMemoryAVInferencePipeline from ltx_distillation.inference.memory_multishot import ( PairedAudioVideoMemoryBank, build_paired_audio_memory_kwargs, video_uint8_to_pil_frames, ) from ltx_distillation.utils import ( add_noise, compute_latent_shapes, decode_benchmark_sample, encode_memory_frames_batch, ) if not prompt.strip(): raise ValueError("Prompt is empty. Enter a shot description.") text_encoder = model.get("text_encoder") if text_encoder is None and callable(model.get("text_encoder_builder")): text_encoder = model["text_encoder_builder"]() model["text_encoder"] = text_encoder if text_encoder is None: raise RuntimeError( "Text encoder not available. It may have been released by a previous shot. " "Set release_text_encoder=False on earlier shots." ) generator = model["generator"] video_vae = model["video_vae"] audio_vae = model["audio_vae"] audio_sample_rate = model["audio_sample_rate"] device = model["device"] dtype = model["dtype"] # Validate num_frames if (num_frames - 1) % 8 != 0: num_frames = 1 + ((num_frames - 1) // 8) * 8 # Update generator resolution generator.video_height = video_height generator.video_width = video_width generator.latent_height = video_height // 32 generator.latent_width = video_width // 32 generator.video_frame_seqlen = generator.latent_height * generator.latent_width # render fps drives the video RoPE clock (see multishot Generate note) generator.VIDEO_FPS = float(video_fps) if int(video_fps) != 24: print(f"[JoyEcho] WARNING: video_fps={video_fps}. The joint AV prior is " f"24fps-native - non-24 fps drifts voices toward Commonwealth " f"accents and overrides accent wording (verified 2026-07-29). " f"Use 24 unless you specifically want that.", flush=True) # Compute latent shapes video_shape, audio_shape = compute_latent_shapes( num_frames=num_frames, video_height=video_height, video_width=video_width, batch_size=1, video_fps=float(video_fps), ) # Get or create memory bank. # CLONE the incoming bank, never adopt it (audit CRITICAL, 2026-07-29): # ComfyUI caches node outputs, so the upstream node hands this node the # SAME live bank object on every queue run. Mutating it in place made # each run condition on the PREVIOUS run's rendered frames and audio - # a literal identity/voice carryover between unrelated queue presses. # A shallow instance copy with a fresh slot list breaks the aliasing; # the entries themselves are immutable after save and safe to share. if memory is not None: import copy as _copy _src_bank = memory["bank"] memory_bank = _copy.copy(_src_bank) memory_bank.memory = list(_src_bank.memory) else: memory_bank = PairedAudioVideoMemoryBank( max_size=memory_max_size, save_mode="random_every_shot_frame", num_fix_frames=num_fix_frames, ) print(f"[JoyEcho] SingleShot: encoding prompt, seed={seed}, " f"memory_size={len(memory_bank)}", flush=True) # --- Phase 0: Encode (text encoder on GPU, everything else off) --- _move(generator, "cpu") _move(video_vae.encoder, "cpu") _move(video_vae.decoder, "cpu") _move(audio_vae.encoder, "cpu") _move(audio_vae.decoder, "cpu") _move(audio_vae.vocoder, "cpu") _move(text_encoder, device) _empty_cache() cond = text_encoder([prompt.strip()]) conditional_dict = { k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in cond.items() } del cond # Offload text encoder immediately after encoding _move(text_encoder, "cpu") _empty_cache() # Build pipelines denoising_sigmas = torch.tensor(DENOISING_SIGMAS, device=device, dtype=torch.float32) base_pipeline = BidirectionalAVInferencePipeline( generator=generator, add_noise_fn=add_noise, denoising_sigmas=denoising_sigmas, ) memory_pipeline = BidirectionalMemoryAVInferencePipeline( generator=generator, add_noise_fn=add_noise, denoising_sigmas=denoising_sigmas, memory_downscale_factor=1, ) offloader = None if sequential_offload: offloader = SequentialOffloader(generator, device) # --- Phase A: Denoise (generator on GPU, everything else off) --- if sequential_offload: offloader.install() else: _move(generator, device) _empty_cache() with torch.random.fork_rng(devices=[device] if device.type == "cuda" else []): torch.manual_seed(seed) if device.type == "cuda": torch.cuda.manual_seed(seed) if len(memory_bank) > 0: _move(video_vae.encoder, device) memory_video = encode_memory_frames_batch( video_vae=video_vae, batch_memory_frames=[memory_bank.get_memory_frames()], target_h=video_height, target_w=video_width, device=device, dtype=dtype, ) _move(video_vae.encoder, "cpu") _empty_cache() memory_audio_kwargs = build_paired_audio_memory_kwargs( memory_bank, enable_audio_memory=enable_audio_memory, v2a_grad_scale=v2a_grad_scale, memory_position_mode="reference", ) video_latent, audio_latent = memory_pipeline.generate( video_shape=tuple(video_shape), audio_shape=tuple(audio_shape), conditional_dict=conditional_dict, memory_video=memory_video, seed=seed, **memory_audio_kwargs, ) del memory_video else: video_latent, audio_latent = base_pipeline.generate( video_shape=tuple(video_shape), audio_shape=tuple(audio_shape), conditional_dict=conditional_dict, seed=seed, ) if device.type == "cuda": torch.cuda.synchronize() del conditional_dict _empty_cache() # Storage not gated on enable_audio_memory (same memory-bank fix as the # multishot node above): the flag gates INJECTION only. audio_memory_latent = ( audio_latent.detach().cpu().contiguous() if audio_latent is not None else None ) # --- Phase B: Decode --- if sequential_offload: offloader.remove() _move(generator, "cpu") _empty_cache() _move(video_vae.decoder, device) _move(audio_vae.decoder, device) _move(audio_vae.vocoder, device) video_uint8, audio_waveform = decode_benchmark_sample( video_vae, audio_vae, video_latent, audio_latent ) if device.type == "cuda": torch.cuda.synchronize() _move(video_vae.decoder, "cpu") _move(audio_vae.decoder, "cpu") _move(audio_vae.vocoder, "cpu") _empty_cache() # Update memory bank memory_frames_pil = video_uint8_to_pil_frames(video_uint8) if audio_memory_latent is not None: memory_bank.save_memory_slot( memory_frames_pil, audio_memory_latent, audio_window_size=audio_memory_window_size, video_clip_num_frames=9, audio_waveform=audio_waveform, audio_sample_rate=16000, video_fps=float(video_fps), audio_window_selection_mode="max_response", video_frame_selection_mode="center", audio_memory_mel_bins=128, audio_memory_mel_hop_length=160, audio_memory_n_fft=1024, audio_memory_downsample_factor=4, audio_memory_is_causal=True, ) # Build outputs images = video_uint8.float() / 255.0 # [F, H, W, 3] audio_out = None if audio_waveform is not None: from ltx_distillation.inference.memory_multishot import normalize_audio_waveform_for_media audio_norm = normalize_audio_waveform_for_media(audio_waveform) audio_out = { "waveform": audio_norm.unsqueeze(0), # [1, C, samples] "sample_rate": audio_sample_rate, } memory_out = {"bank": memory_bank} del video_latent, audio_latent, audio_memory_latent, video_uint8, audio_waveform _empty_cache() print(f"[JoyEcho] SingleShot done. {images.shape[0]} frames.", flush=True) return (images, audio_out, memory_out, model,) _PROMPTS_DIR = Path(__file__).resolve().parent / "prompts" _DEFAULT_LONG_STORY_SYSTEM_PROMPT = "" _long_sp_path = _PROMPTS_DIR / "long_story_writer_system_prompt.md" if _long_sp_path.exists(): _DEFAULT_LONG_STORY_SYSTEM_PROMPT = _long_sp_path.read_text(encoding="utf-8").strip() def _load_system_prompt(mode: str) -> str: """Load the full system prompt from the bundled markdown file.""" if "long" in mode: fp = _PROMPTS_DIR / "long_story_writer_system_prompt.md" else: fp = _PROMPTS_DIR / "short_story_writer_system_prompt.md" if fp.exists(): return fp.read_text(encoding="utf-8").strip() raise FileNotFoundError(f"System prompt not found: {fp}") class JoyEcho_PromptFormat: """Helper node providing the official prompt writing system prompts. Use this with any LLM node in ComfyUI to generate properly formatted shot prompts from a short story description. The output can be fed directly into JoyEcho_TextEncode. """ @classmethod def INPUT_TYPES(cls): return { "required": { "mode": (["long_story (multi-shot)", "short_story (single-shot)"],), }, } RETURN_TYPES = ("STRING",) RETURN_NAMES = ("system_prompt",) FUNCTION = "get_prompt" CATEGORY = "JoyAI-Echo" def get_prompt(self, mode: str): return (_load_system_prompt(mode),) def _extract_json_object(text: str): """Return the first top-level {...} JSON object substring in `text`, matching braces while ignoring any that sit inside a JSON string (the shot prompts themselves contain no raw braces, but escaped quotes / stray prose might). Returns None if no balanced object is found. Lets the enhancer survive a model that wraps its answer in preamble or trailing commentary.""" start = text.find("{") if start < 0: return None depth, in_str, esc = 0, False, False for i in range(start, len(text)): c = text[i] if in_str: if esc: esc = False elif c == "\\": esc = True elif c == '"': in_str = False continue if c == '"': in_str = True elif c == "{": depth += 1 elif c == "}": depth -= 1 if depth == 0: return text[start:i + 1] return None class JoyEcho_LLMEnhance: """Call a cloud LLM API to expand a short story idea into JoyAI-Echo shot prompts. Supports OpenAI-compatible APIs (OpenAI, DeepSeek, etc.). The output JSON can be fed directly into JoyEcho_TextEncode or split via JoyEcho_PromptAtIndex. Uses only cloud API calls — zero local GPU memory. """ @classmethod def INPUT_TYPES(cls): return { "required": { "story_idea": ("STRING", { "multiline": True, "default": "A young woman records a quiet evening vlog in her cozy room, reflecting on life and finding warmth in small things.", "tooltip": "Describe your story or scene idea in a few sentences.", }), "mode": (["long_story (multi-shot)", "short_story (single-shot)", "passthrough (raw JSON, skip LLM)", "revise (keep structure, rewrite prose)"],), "api_key": ("STRING", { "default": "", "tooltip": "Your provider's API key (OpenAI, DeepSeek, GLM, Gemini...). " "LEAVE BLANK for a local endpoint (localhost / 192.168.x / .local) " "- those ignore it and a placeholder is sent automatically. Also " "not needed in passthrough mode, which skips the LLM entirely.", }), "system_prompt": ("STRING", { "multiline": True, "default": _DEFAULT_LONG_STORY_SYSTEM_PROMPT, "tooltip": "System prompt for the LLM. Edit to customize prompt generation style.", }), }, "optional": { "base_url": ("STRING", { "default": "https://api.openai.com/v1", "tooltip": "API base URL. Use https://api.deepseek.com/v1 for DeepSeek, etc.", }), "model_name": ("STRING", { "default": "gpt-4o", "tooltip": "Model name (gpt-4o, deepseek-chat, claude-3-5-sonnet, etc.)", }), "num_shots": ("INT", { "default": 0, "min": 0, "max": 30, "tooltip": "Number of shots to generate (0 = let LLM decide, default 15 for long story).", }), "temperature": ("FLOAT", { "default": 0.7, "min": 0.0, "max": 2.0, "step": 0.05, }), "num_frames": ("INT", { "default": 0, "min": 0, "max": 100000, "tooltip": "Frames per shot (match JoyEcho_Generate). 0 = " "disabled (assume ~10s clips). When set, the LLM is " "told each shot's real duration and scales action + " "dialogue length to fit.", }), "fps": ("FLOAT", { "default": 25.0, "min": 1.0, "max": 120.0, "step": 1.0, "tooltip": "Playback fps used to turn num_frames into seconds " "per shot (match your CreateVideo / Generate fps).", }), }, } RETURN_TYPES = ("STRING",) RETURN_NAMES = ("prompts_json",) FUNCTION = "enhance" CATEGORY = "JoyAI-Echo" def enhance( self, story_idea: str, mode: str, api_key: str, system_prompt: str, base_url: str = "https://api.openai.com/v1", model_name: str = "gpt-4o", num_shots: int = 0, temperature: float = 0.7, num_frames: int = 0, fps: float = 25.0, ): import urllib.request import urllib.error # AUTO-DETECT: if story_idea is already a valid {"prompts":[...]} payload, # pass it through regardless of the mode widget - so one wiring serves both # briefs (LLM-enhanced) and finished scripts without flipping the mode. _looks_json = story_idea.strip().startswith("{") # 'revise' deliberately EXPECTS a finished JSON script, so it must not # be hijacked by the auto-passthrough below (2026-07-20). if _looks_json and "revise" not in mode.lower() and "passthrough" not in mode.lower(): try: _probe = json.loads(story_idea.strip()) if isinstance(_probe.get("prompts") or _probe.get("shots"), list): print("[JoyEcho] LLMEnhance: story_idea is a finished prompts JSON - " "auto-passthrough (mode widget ignored).", flush=True) mode = "passthrough (auto)" except (json.JSONDecodeError, AttributeError): pass # PASSTHROUGH: feed straight {"prompts":[...]} JSON in story_idea and skip the # LLM entirely. Lets the same node/wiring accept either an enhanced brief or a # finished script (e.g. from the Script Picker) via the mode toggle. if "passthrough" in mode.lower(): text = story_idea.strip() data = None if text.startswith("{"): try: data = json.loads(text) except json.JSONDecodeError: data = None if isinstance(data, dict): arr = data.get("prompts") or data.get("shots") if not isinstance(arr, list) or not arr: raise ValueError( 'Passthrough mode: JSON input must contain a non-empty ' '"prompts" (or "shots") array.' ) print(f"[JoyEcho] LLMEnhance PASSTHROUGH: {len(arr)} shots, " f"no LLM call.", flush=True) return (text,) # PLAIN TEXT (the LPFF/txt PromptSource path): a raw prompt IS a # finished one-shot script. Refusing it forced every txt batch # through the LLM rewrite, which silently destroyed any experiment # depending on exact prompt text (found by the accent census, # 2026-07-29). Wrap it; downstream sees the same JSON shape as ever. if not text: raise ValueError("Passthrough mode: story_idea is empty.") print("[JoyEcho] LLMEnhance PASSTHROUGH: plain-text prompt wrapped " "as 1 shot, no LLM call.", flush=True) return (json.dumps({"prompts": [text]}, ensure_ascii=True),) # ── REVISE: rewrite the PROSE of an existing script, keep its SHAPE ── # Rebels local patch 2026-07-20. Built for vision-native reasoning # models (minimax-m3) that are strong at physical/spatial/camera # language but, like every LLM, will happily reword an identity # sentence or drop a shot. So structure is NOT trusted to the model: # the shot count, the byte-identical identity sentences and the ASCII # rule are all re-imposed in code after the call, and a structurally # bad reply falls back to the ORIGINAL script rather than failing the # render. Pairs with frame-aware pacing (num_frames/fps), which the # passthrough path can never use because it skips the LLM entirely. if "revise" in mode.lower(): _src = story_idea.strip() _oshots = None if _src.startswith("{"): try: _orig = json.loads(_src) _oshots = _orig.get("prompts") or _orig.get("shots") except (json.JSONDecodeError, AttributeError): _oshots = None if not isinstance(_oshots, list) or not _oshots: # PLAIN TEXT is valid input too (2026-07-20): a single LPFF # block off JoyEcho_PromptSource is one finished LTX prompt # paragraph, not a JSON script. Treat it as a one-shot script # so the same revise pass works on RIFT corpus prompts. Output # is still {"prompts":[...]} because that is what TextEncode # downstream consumes. if not _src: raise ValueError( "Revise mode needs something to revise: either a " '{"prompts": [...]} JSON script or a single prompt ' "paragraph in story_idea.") _oshots = [_src] print("[JoyEcho] REVISE: input is plain text - treating as a " "single-shot script.", flush=True) _oshots = [str(s) for s in _oshots] if num_shots and num_shots > 0 and num_shots != len(_oshots): print(f"[JoyEcho] REVISE: num_shots={num_shots} IGNORED - revise " f"preserves the input's {len(_oshots)} shot(s) by design " f"(a reply with a different count is rejected). To EXPAND " f"into {num_shots} shots, use mode 'long_story' instead - " f"it treats the input as a premise and writes a new " f"script.", flush=True) # Fallback value is ALWAYS the JSON form, even when the input # was a plain LPFF paragraph - downstream TextEncode expects # {"prompts":[...]}, so a rejected revision must not hand it # back raw prose (2026-07-20). _fallback = json.dumps({"prompts": _oshots}, ensure_ascii=True) print(f"[JoyEcho] LLMEnhance REVISE: {len(_oshots)} shots in, " f"structure will be re-imposed after the call.", flush=True) _rev_sys = ( "You revise shot prompts for a multi-shot AI video harness. You are " "given a JSON object {\"prompts\":[...]} where each array item is ONE " "shot's complete prompt text.\n\n" "HARD RULES - violating any of these makes your output unusable:\n" "1. Return the SAME NUMBER of shots, in the SAME ORDER. Never add, " "drop, merge or reorder shots.\n" "2. If a shot contains a character-identity sentence (an 'ID_A is " "...' sentence describing appearance, wardrobe and voice), copy it " "BYTE-FOR-BYTE into your revision of that shot - never reword it " "even slightly, because cross-shot identity depends on it being " "character-identical. Not every script uses them; if there is none, " "keep the subject's described appearance and wardrobe unchanged " "instead. Likewise keep any trained-LoRA trigger token (a lowercase " "word ending in '_rift') exactly as written and in place.\n" "3. Keep each shot's quoted dialogue MEANING and its position in the " "story. You may re-word a line for rhythm, but never change what it " "says or move it to another shot.\n" "4. Keep the capture-medium sentence and the continuous-ambient-sound " "sentence in every shot.\n" "5. Pure ASCII only. No em dashes, smart quotes or unicode.\n\n" "WHAT YOU SHOULD IMPROVE: the physical, spatial and camera writing. " "Make the scene description more concrete and more renderable - named " "materials with a condition ('cracked wet asphalt', 'rust-streaked " "steel'), explicit light sources and their direction and falloff, one " "clear camera move stated in real cinematography terms, and beats that " "are physically possible in the shot's duration. Prefer literal " "physical description over poetic or emotional abstraction; emotion " "belongs only in the performance/voice cue and the dialogue itself. " "Never describe anything the camera cannot see.\n\n" "Output ONLY the JSON object. No commentary, no markdown fence." ) _rev_user = "SCRIPT TO REVISE:\n" + json.dumps( {"prompts": _oshots}, ensure_ascii=True) _fps_r = fps if (fps and fps > 0) else 25.0 if num_frames and num_frames > 0: _secs = num_frames / _fps_r _lo = max(6, int(round(_secs * 1.2))) _hi = int(round(_secs * 2.0)) _rev_user += ( f"\n\nEach shot renders as one continuous clip about " f"{_secs:.1f} seconds long at {int(round(_fps_r))} fps. Pace " f"every shot's action to fill that duration, and size each " f"spoken line to roughly {_lo}-{_hi} words so the speech fits " f"the clip with room for breath - lengthen or tighten the " f"existing lines as needed WITHOUT changing what they say.") _payload = json.dumps({ "model": model_name, "messages": [{"role": "system", "content": _rev_sys}, {"role": "user", "content": _rev_user}], "temperature": temperature, "max_tokens": 16384, }).encode("utf-8") _url = base_url.rstrip("/") + "/chat/completions" _hdrs = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key.strip()}"} print(f"[JoyEcho] REVISE: calling {model_name}...", flush=True) try: _rq = urllib.request.Request(_url, data=_payload, headers=_hdrs, method="POST") with urllib.request.urlopen(_rq, timeout=240) as _rs: _rr = json.loads(_rs.read().decode("utf-8")) _rtext = _rr["choices"][0]["message"]["content"].strip() except Exception as _e: # noqa: BLE001 print(f"[JoyEcho] REVISE: call failed ({_e}); keeping ORIGINAL " f"script.", flush=True) return (_fallback,) import re as _re2 _rtext = _re2.sub(r".*?", "", _rtext, flags=_re2.DOTALL | _re2.IGNORECASE).strip() if _rtext.startswith("```"): _rtext = "\n".join(l for l in _rtext.split("\n") if not l.strip().startswith("```")).strip() if not (_rtext.startswith("{") and _rtext.endswith("}")): _carved = _extract_json_object(_rtext) if _carved is not None: _rtext = _carved try: _new = json.loads(_rtext) _nshots = _new.get("prompts") or _new.get("shots") assert isinstance(_nshots, list) _nshots = [str(s) for s in _nshots] except Exception as _e: # noqa: BLE001 print(f"[JoyEcho] REVISE: unparsable reply ({_e}); keeping " f"ORIGINAL script.", flush=True) return (_fallback,) if len(_nshots) != len(_oshots): print(f"[JoyEcho] REVISE: shot count changed " f"({len(_oshots)} -> {len(_nshots)}); keeping ORIGINAL " f"script.", flush=True) return (_fallback,) # Re-impose structure: ASCII-fold, then force every identity # sentence back to the ORIGINAL wording per shot. _idre = _re2.compile(r"(ID_[A-Z] is .*?with every word\.)", _re2.S) _fixed, _n_id = [], 0 for _o, _n in zip(_oshots, _nshots): _n = (_n.replace("—", "-").replace("–", "-") .replace("‘", "'").replace("’", "'") .replace("“", '"').replace("”", '"') .replace("…", "...")) _n = "".join(c for c in _n if ord(c) < 128) _om = _idre.search(_o) _nm = _idre.search(_n) if _om and _nm and _om.group(1) != _nm.group(1): _n = _n.replace(_nm.group(1), _om.group(1), 1) _n_id += 1 elif _om and not _nm: print("[JoyEcho] REVISE: a shot lost its identity sentence; " "keeping ORIGINAL script.", flush=True) return (_fallback,) # LoRA trigger guard: RIFT corpus prompts carry trained-face # triggers (e.g. alice_rift, bob_rift). Dropping one renders a # generic face, so a revision that loses any trigger the # original had is rejected outright (2026-07-20). _otrig = set(_re2.findall(r"\b[a-z][a-z0-9_]*_rift\b", _o)) if _otrig and not _otrig.issubset( set(_re2.findall(r"\b[a-z][a-z0-9_]*_rift\b", _n))): print(f"[JoyEcho] REVISE: revision dropped LoRA trigger(s) " f"{sorted(_otrig)}; keeping ORIGINAL script.", flush=True) return (_fallback,) _fixed.append(_n) print(f"[JoyEcho] REVISE: {len(_fixed)} shots revised " f"({_n_id} identity sentence(s) restored).", flush=True) return (json.dumps({"prompts": _fixed}, ensure_ascii=True),) if not api_key.strip(): # Local OpenAI-compatible servers (Ollama, LM Studio, llama.cpp, # vLLM) ignore the Authorization header, but one still has to be # sent. Synthesize a placeholder instead of making every local user # discover they must type a fake key into a field their endpoint # never reads - that empty-field error was the single most common # first-run stumble on this node. from urllib.parse import urlparse as _urlparse _host = (_urlparse(base_url).hostname or "").lower() _is_local = (_host in ("localhost", "127.0.0.1", "0.0.0.0", "::1", "host.docker.internal") or _host.startswith("192.168.") or _host.startswith("10.") or _host.endswith(".local")) if _is_local: api_key = "local" print(f"[JoyEcho] LLMEnhance: no api_key set and {base_url} is a local " f"endpoint - sending a placeholder key.", flush=True) else: raise ValueError( f"API key is required for {base_url}. Enter your provider's key " f"(OpenAI, DeepSeek, GLM, Gemini, etc.). Local endpoints such as " f"http://localhost:11434/v1 do not need one - leave this blank.") if system_prompt.strip(): sys_prompt = system_prompt.strip() else: sys_prompt = _load_system_prompt(mode) user_msg = story_idea.strip() if num_shots > 0: user_msg += f"\n\nGenerate exactly {num_shots} shots." # Rebels local patch: frame-aware pacing. The enhancer otherwise assumes # a fixed ~10s clip and sizes dialogue for it; when the workflow renders # longer shots (e.g. 361 frames), tell the LLM the REAL per-shot duration # so it fills the time with beats + proportionally longer speech instead # of a 10s line stranded in a 14s clip. num_frames=0 -> unchanged. _fps = fps if (fps and fps > 0) else 25.0 if num_frames and num_frames > 0: secs = num_frames / _fps lo = max(6, int(round(secs * 1.2))) hi = int(round(secs * 2.0)) user_msg += ( f"\n\nEach shot is a single continuous clip about {secs:.1f} " f"seconds long (at {int(round(_fps))} fps). Pace every shot to " f"fill roughly {secs:.0f} seconds: give the action enough small " f"beats to occupy the full duration without any fast or complex " f"motion, and do not leave long dead air. For a speaking shot, " f"scale the dialogue to this length — about {lo}-{hi} words, " f"delivered as one natural line or a short two-line exchange, with " f"room for pauses, breath, and reaction. These per-shot timing " f"numbers override any default clip length mentioned above." ) url = base_url.rstrip("/") + "/chat/completions" payload = json.dumps({ "model": model_name, "messages": [ {"role": "system", "content": sys_prompt}, {"role": "user", "content": user_msg}, ], "temperature": temperature, "max_tokens": 16384, }).encode("utf-8") headers = { "Content-Type": "application/json", "Authorization": f"Bearer {api_key.strip()}", } print(f"[JoyEcho] Calling LLM ({model_name}) to enhance prompt...", flush=True) req = urllib.request.Request(url, data=payload, headers=headers, method="POST") try: with urllib.request.urlopen(req, timeout=240) as resp: result = json.loads(resp.read().decode("utf-8")) except urllib.error.HTTPError as e: body = e.read().decode("utf-8", errors="replace") raise RuntimeError(f"LLM API error {e.code}: {body}") _raw = result["choices"][0]["message"]["content"].strip() content = _raw # Rebels local patch: robust enhancer JSON extraction. Reasoning / cloud # models (e.g. minimax-m3) may wrap the answer in ... # traces or add preamble/trailing commentary; the pipeline needs ONLY the # {"prompts":[...]} object. A reasoning trace cannot be reliably prompted # away, so peel it here instead of failing the whole enhance. import re as _re content = _re.sub(r".*?", "", content, flags=_re.DOTALL | _re.IGNORECASE).strip() # Strip markdown code fences if present if content.startswith("```"): lines = content.split("\n") lines = [l for l in lines if not l.strip().startswith("```")] content = "\n".join(lines).strip() # Still-surrounding prose? Carve out the outermost balanced {...} object. if not (content.startswith("{") and content.endswith("}")): _obj = _extract_json_object(content) if _obj is not None: content = _obj # Validate JSON try: data = json.loads(content) if "prompts" not in data or not isinstance(data["prompts"], list): raise ValueError("LLM output missing 'prompts' array") num = len(data["prompts"]) except (json.JSONDecodeError, ValueError) as e: raise RuntimeError( f"LLM returned invalid JSON: {e}\n\nRaw output:\n{_raw[:800]}" ) print(f"[JoyEcho] LLM generated {num} shot prompt(s).", flush=True) # Persist + echo the generated prompts so you can inspect exactly what # the enhancer produced (this JSON is what feeds JoyEcho_TextEncode). try: import os import folder_paths _outdir = os.path.join(folder_paths.get_output_directory(), "joyecho") os.makedirs(_outdir, exist_ok=True) _dump = os.path.join(_outdir, "enhanced_prompts_latest.json") with open(_dump, "w", encoding="utf-8") as _f: _f.write(content) print(f"[JoyEcho] enhancer output written to: {_dump}", flush=True) except Exception as _e: print(f"[JoyEcho] could not write enhancer output file: {_e}", flush=True) print("[JoyEcho] ---------- enhancer output (prompts) ----------", flush=True) print(content, flush=True) print("[JoyEcho] ---------- end enhancer output ----------", flush=True) return (content,) class JoyEcho_PromptAtIndex: """Extract a single prompt from a JSON prompts array by index. Connect the output to a SingleShotGenerate node's prompt input to override the text box with LLM-generated content. This is optional — if not connected, the SingleShot node uses its own text box. """ @classmethod def INPUT_TYPES(cls): return { "required": { "prompts_json": ("STRING", { "multiline": True, "default": "", "tooltip": "JSON string with 'prompts' array (from LLM Enhance or file)", }), "index": ("INT", { "default": 0, "min": 0, "max": 29, "tooltip": "0-based shot index to extract", }), }, } RETURN_TYPES = ("STRING",) RETURN_NAMES = ("prompt",) FUNCTION = "extract" CATEGORY = "JoyAI-Echo" def extract(self, prompts_json: str, index: int): text = prompts_json.strip() if not text: raise ValueError("No prompts JSON provided.") try: data = json.loads(text) except json.JSONDecodeError as e: raise ValueError(f"Invalid JSON: {e}") prompt_list = data.get("prompts") or data.get("shots") or [] if not prompt_list: raise ValueError("JSON must contain a 'prompts' or 'shots' array.") if index >= len(prompt_list): raise ValueError( f"Index {index} out of range (only {len(prompt_list)} prompts available)." ) return (str(prompt_list[index]).strip(),)