nodes.py

"""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
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):
        self._generator = generator
        self._device = device
        self._hooks: list[torch.utils.hooks.RemovableHook] = []
        self._pin_memory = pin_memory
        self._installed = False

    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)

        # Move all blocks to CPU (optionally pinned)
        for block in blocks:
            block.to("cpu")
            if self._pin_memory and torch.cuda.is_available():
                for param in block.parameters():
                    param.data = param.data.pin_memory()
                for buf in block.buffers():
                    buf.data = buf.data.pin_memory()

        # 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 blocks:
            h1 = block.register_forward_pre_hook(make_pre_hook(block))
            h2 = block.register_forward_hook(make_post_hook(block))
            self._hooks.extend([h1, h2])

        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")


def _list_model_files() -> list:
    """Every *.safetensors / *.gguf under the ComfyUI model dirs, 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 [_MODEL_FILE_MANUAL]
    out, seen_dirs, seen = [], set(), set()
    for cat in _MODEL_FILE_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 ("*.safetensors", "*.gguf"):
                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 [_MODEL_FILE_MANUAL] + sorted(out)


def _resolve_model_file(choice: str) -> str:
    import folder_paths
    cat, _, rel = choice.partition(": ")
    if cat in _MODEL_FILE_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"model_file {choice!r} no longer exists on disk. Refresh the node "
        f"list (R) and re-pick, or use {_MODEL_FILE_MANUAL} + checkpoint_path.")


class JoyEcho_ModelLoader:
    """Load JoyAI-Echo model components: text encoder, DiT generator, and VAEs."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "checkpoint_path": ("STRING", {
                    "default": "",
                    "tooltip": "Path to echo-longvideo-release.safetensors",
                }),
                "gemma_path": ("STRING", {
                    "default": "",
                    "tooltip": "Path to gemma-3-12b-it directory (bf16 safetensors)",
                }),
            },
            "optional": {
                "lora_path": ("STRING", {"default": ""}),
                "lora_strength": ("FLOAT", {
                    "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05,
                }),
                "low_vram": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "Load text encoder on CPU for 24GB GPUs. "
                               "Encoding will be slower but uses no GPU memory.",
                }),
                "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).",
                }),
                "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.",
                }),
            },
        }

    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,
                   low_vram: bool = False, fp8_transformer: bool = False,
                   model_file: str = _MODEL_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

        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)

        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())

        # gemma_path must be the HF gemma-3-12b-it DIRECTORY (model*.safetensors +
        # tokenizer.model), not a GGUF or single file. This loader is the
        # safetensors-native path; a GGUF text encoder only works through the
        # Rebels discrete TextEncoder node (which has the GGUF weight-swap).
        # Fail early with a readable message instead of a deep tokenizer.model
        # traceback repeated once per queued item.
        _gp = Path(gemma_path)
        if _gp.suffix.lower() == ".gguf" or _gp.is_file():
            raise ValueError(
                f"gemma_path points at a file ({_gp.name}). It must be the "
                f"gemma-3-12b-it FOLDER (containing model-0000x-of-*.safetensors "
                f"and tokenizer.model), not a .gguf or single file. A GGUF text "
                f"encoder is only supported by the Rebels discrete TextEncoder node, "
                f"not this loader.")
        if not (_gp / "tokenizer.model").is_file():
            raise ValueError(
                f"gemma_path {_gp} is not a valid Gemma root: no tokenizer.model "
                f"inside. Point it at a full gemma-3-12b-it folder.")

        device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        dtype = torch.bfloat16

        # Load text encoder
        text_encoder_device = torch.device("cpu") if low_vram else device
        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()

        # Load generator
        print("[JoyEcho] Loading DiT generator...", flush=True)
        loras = ()
        if lora_path and lora_path.strip():
            loras = (
                LoraPathStrengthAndSDOps(
                    str(Path(lora_path).expanduser()),
                    float(lora_strength),
                    LTXV_LORA_COMFY_RENAMING_MAP,
                ),
            )

        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_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,
                gemma_path=gemma_path,
                device=torch.device("cpu"),
                dtype=dtype,
                video_height=736,
                video_width=1280,
                loras=loras,
                quantization=quantization,
            )
            generator.eval()

        # 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()

        audio_sample_rate = audio_vae.get_output_sample_rate() or 24000

        model = {
            "text_encoder": 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,
        }

        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:
            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.")

        device = model["device"]
        print(f"[JoyEcho] Encoding {len(prompt_list)} prompt(s)...", flush=True)

        # 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 = text_encoder([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 = text_encoder([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 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": 481, "step": 8,
                                       "tooltip": "Must be 1 + 8*k (e.g. 121, 241, 361)"}),
                "video_height": ("INT", {"default": 736, "min": 256, "max": 1088, "step": 32}),
                "video_width": ("INT", {"default": 1280, "min": 256, "max": 1920, "step": 32}),
            },
            "optional": {
                "video_fps": ("INT", {"default": 25, "min": 1, "max": 60}),
                "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}),
                "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. "
                               "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": 0, "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 - a split-second flash of the reference image. 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.",
                }),
            },
        }

    RETURN_TYPES = ("IMAGE", "AUDIO",)
    RETURN_NAMES = ("images", "audio",)
    FUNCTION = "generate"
    CATEGORY = "JoyAI-Echo"
    OUTPUT_NODE = 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,
        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",
    ):
        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,
        )

        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

        # 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 = []
        if reference_image is not None and memory_max_size > 0:
            import numpy as np
            from PIL import Image as _PILImage
            # Dedupe identical frames in the batch (wiring quirks can deliver
            # the same picked image 2-3x).
            _uniq_idx = []
            _seen = []
            for _i in range(int(reference_image.shape[0])):
                _t = reference_image[_i]
                if not any(_t.shape == _u.shape and torch.equal(_t, _u) for _u in _seen):
                    _seen.append(_t)
                    _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)
            for _ri in _uniq_idx[:4]:
                _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). Cover-fit: scale
                # to fill, then crop with a mild top bias so faces survive.
                if _ref_pil.size != (_tw, _th):
                    _scale = max(_tw / _ref_pil.width, _th / _ref_pil.height)
                    _rw, _rh = max(_tw, int(round(_ref_pil.width * _scale))), max(_th, int(round(_ref_pil.height * _scale)))
                    _ref_pil = _ref_pil.resize((_rw, _rh), _PILImage.LANCZOS)
                    _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)
            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 = []

        num_shots = len(conditioning)
        offloader = None
        if sequential_offload:
            offloader = SequentialOffloader(generator, device)

        # 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)

        for shot_idx in range(num_shots):
            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()
            }

            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)

                if _ref_clips 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(_ref_clips) + (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",
                        )
                    if _ref_clips and memory_audio_kwargs:
                        print("[JoyEcho] WARNING: reference clips + enable_audio_memory=True gives "
                              f"{len(_mem_frames)} video slots vs {len(memory_bank)} audio slots; "
                              "if slot pairing errors, set enable_audio_memory=False.", 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()

            # 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="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,
                )

            # Collect outputs
            # video_uint8: [F, H, W, 3] uint8 -> [F, H, W, 3] float32 [0, 1]
            video_float = video_uint8.float() / 255.0

            # HEAD TRIM: each shot's first frames morph out of the memory /
            # reference content (a split-second flash of the reference image).
            # Trim them, with matching audio samples so A/V stay aligned.
            _trim = max(0, int(head_trim_frames))
            if _trim == 0 and _ref_clips:
                _trim = 8  # auto when references are wired
            if _trim > 0 and video_float.shape[0] > _trim + 16:
                video_float = video_float[_trim:]
            else:
                _trim = 0
            all_video_frames.append(video_float)

            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)
                if _trim > 0:
                    _cut = int(round(_trim / float(video_fps) * audio_sample_rate))
                    if audio_norm.shape[-1] > _cut:
                        audio_norm = audio_norm[..., _cut:]
                all_audio_waveforms.append(audio_norm)

            # Save per-shot video immediately for real-time preview
            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)

        # 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
            out_v = vids[0]
            out_a = auds[0] if auds else None
            for i in range(1, len(vids)):
                b_v = vids[i]
                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:
            images = torch.cat(all_video_frames, dim=0)  # [total_frames, H, W, 3]
            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
                for v in all_video_frames[:-1]:
                    acc += v.shape[0]
                    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)

        print(f"[JoyEcho] Generation complete. {images.shape[0]} frames, "
              f"{num_shots} shot(s).", flush=True)

        return (images, audio_out,)

    @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"
        stream.options = {"crf": "18", "preset": "fast"}

        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
        if audio_waveform is not None:
            import torchaudio
            wav_path = sub_dir / f"{name_prefix}_{shot_idx:03d}.wav"
            waveform = audio_waveform.cpu()
            if waveform.dim() == 1:
                waveform = waveform.unsqueeze(0)
            torchaudio.save(str(wav_path), waveform, sample_rate=audio_sr)

        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": 481, "step": 8,
                                       "tooltip": "Must be 1 + 8*k (e.g. 121, 241, 361)"}),
                "video_height": ("INT", {"default": 736, "min": 256, "max": 1088, "step": 32}),
                "video_width": ("INT", {"default": 1280, "min": 256, "max": 1920, "step": 32}),
            },
            "optional": {
                "memory": ("JOYECHO_MEMORY",),
                "video_fps": ("INT", {"default": 25, "min": 1, "max": 60}),
                "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}),
                "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:
            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

        # 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
        if memory is not None:
            memory_bank = memory["bank"]
        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),)


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)"],),
                "api_key": ("STRING", {
                    "default": "",
                    "tooltip": "Your API key (OpenAI, DeepSeek, etc.). Not needed in passthrough mode.",
                }),
                "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,
                }),
            },
        }

    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,
    ):
        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("{")
        if _looks_json 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()
            try:
                data = json.loads(text)
            except json.JSONDecodeError as e:
                raise ValueError(
                    f"Passthrough mode expects raw JSON in story_idea, but it did not parse: {e}"
                )
            arr = data.get("prompts") if isinstance(data, dict) else None
            if arr is None and isinstance(data, dict):
                arr = data.get("shots")
            if not isinstance(arr, list) or not arr:
                raise ValueError(
                    'Passthrough mode expects {"prompts": [...]} JSON (non-empty array) in story_idea.'
                )
            print(f"[JoyEcho] LLMEnhance PASSTHROUGH: {len(arr)} shots, no LLM call.", flush=True)
            return (text,)

        if not api_key.strip():
            raise ValueError("API key is required. Enter your OpenAI/DeepSeek/etc. API key.")

        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."

        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=120) 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}")

        content = result["choices"][0]["message"]["content"].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()

        # 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{content[:500]}"
            )

        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(),)