Upload 11 files

README.md

@@ -1,3 +1,173 @@
 ---
-license: apache-2.0
+license: other
+license_name: ltx-2-community-license
+license_link: https://huggingface.co/Lightricks/LTX-2/blob/main/LICENSE.txt
+tags:
+  - comfyui
+  - joyai-echo
+  - ltx-video
+  - multishot
+  - patch
 ---
+
+# JoyAI-Echo GGUF nodes — multishot fixes + automation patch
+
+A set of bug fixes and features layered on top of the community
+`ComfyUI_JoyAI_Echo_GGUF_Nodes` pack (the Rebels GGUF loader stack around
+JoyAI-Echo). Everything here targets the **multi-shot** path (`JoyEcho_Generate`
++ the discrete Rebels loaders / `JoyEcho_ModelLoader`).
+
+This is a **patch drop**, not a standalone pack: copy these files over a working
+install of the same pack (back up first). The files are interdependent — in
+particular `nodes.py` calls new signatures added to the two `libs/` files, so
+apply them together.
+
+Tested on an RTX 5090 (32 GB) and a 3090 (24 GB), ComfyUI 0.26–0.27,
+torch 2.8–2.11, with the JoyAI-Echo bf16 release and self-built Q8 GGUFs.
+
+---
+
+## Files in this package
+
+```
+nodes.py                                     # JoyEcho_TextEncode / _Generate / _ModelLoader / _LLMEnhance
+__init__.py                                  # registrations for the new nodes
+rebels_loaders.py                            # discrete GGUF loaders (text-encoder fixes)
+joyecho_prompt_source.py   (new node)        # one dropdown: .txt briefs + .json scripts
+joyecho_ref_picker.py      (new node)        # auto reference-image picker by character name
+joyecho_ref_batch.py       (new node)        # None-tolerant image batcher
+joyecho_script_picker.py   (new node)        # JSON dropdown (superseded by PromptSource)
+libs/ltx_distillation/models/ltx_wrapper.py  # fp8 quantization passthrough
+libs/ltx_distillation/utils.py               # tiled VAE decode
+prompts/long_story_writer_system_prompt.md   # (optional) de-musicked + character-age edits
+```
+
+---
+
+## Bug fixes
+
+### 1. `enable_audio_memory=False` silently disabled ALL cross-shot memory
+The pack computed `audio_memory_latent=None` when audio memory was off, and the
+video **memory-bank save was gated on that latent being non-None** — so with
+audio memory off (the standard anti-drone setting) the bank never filled and
+cross-shot **identity** silently died (symptom: `memory_size=0` every shot even
+with `memory_max_size=7`; a new face each shot).
+Fix: memory storage is now unconditional; `enable_audio_memory` gates only the
+audio-memory **injection** path. Verify: console `memory_size=` should climb
+0,1,2,… capped at your `memory_max_size`. (`nodes.py`)
+
+### 2. GGUF text-encoder loader (`RebelsJE_TextEncoder`)
+Two fixes so a text-only Gemma-3 GGUF loads cleanly:
+- **meta-strip**: drop `vision_tower` / `multi_modal_projector` / `lm_head`
+  (the text-only GGUF has no weights for them → "Cannot copy out of meta tensor").
+- **device-unify**: pin the embeddings-processor to the encoder's actual device
+  (GGUF Gemma runs on CPU while the connector was on cuda → addmm device mismatch).
+(`rebels_loaders.py`)
+
+---
+
+## Features
+
+### 3. Split per-domain negative lever (`JoyEcho_TextEncode`)
+The DMD pipeline has no CFG, so the only steering lever is embedding-space.
+Instead of one `negative_prompt`/`negative_scale` that steers both branches,
+this splits it:
+- `negative_prompt_video` / `negative_scale_video` — kills burned-in
+  captions/subtitles. Working value ~0.5. **Above ~0.8 it over-rotates the
+  video context and locks every shot to shot 1's composition** (scene-lock).
+- `negative_prompt_audio` / `negative_scale_audio` — kills invented
+  music/score. Keep ≤ ~0.4 or dialogue suffers.
+Steering is norm-preserving (RescaleCFG-style): `cond' = renorm(cond + s*(cond − neg))`.
+Old single-widget names still work as a fallback. (`nodes.py`)
+
+### 4. Passthrough mode (`JoyEcho_LLMEnhance`)
+`mode = "passthrough (raw JSON, skip LLM)"` — feed a finished
+`{"prompts":[...]}` script straight through with no LLM call / no API key.
+Auto-detects when `story_idea` already parses as that JSON. (`nodes.py`)
+
+### 5. Reference-image conditioning — I2V-as-reference (`JoyEcho_Generate`)
+New `reference_image` (IMAGE batch, up to 4). Identity references are prepended
+as **video-only conditioning clips** at the memory-encode step — they are
+**never** written into the paired audio/video bank. (An earlier attempt that
+seeded refs into the bank with zero-filled audio latents injected loud
+background noise with 2+ refs; video-only conditioning avoids it entirely.)
+Also new: `head_trim_frames` (auto 8 with refs) drops the first N frames of each
+shot, where the model morphs out of the reference/memory content. (`nodes.py`)
+
+### 6. Shot transitions (`JoyEcho_Generate`)
+`transition`: `cut` (original) / `dissolve` (overlap cross-dissolve + equal-power
+audio crossfade) / `vhs_glitch` (analog static burst at each boundary: snow,
+tear bands, dropout lines + a raised-cosine tape-noise audio bed).
+`transition_frames`, `glitch_intensity` tune it. (`nodes.py`)
+
+### 7. fp8 transformer quantization (`JoyEcho_ModelLoader`)
+New `fp8_transformer` toggle. Quantizes the DiT's attention/FF linear weights to
+`float8_e4m3fn` **at load, from the normal bf16 checkpoint** (uses the vendored
+`ltx_core.quantization.QuantizationPolicy.fp8_cast()` — upcasts per-layer at
+inference). Roughly halves DiT weight memory and halves sequential-offload PCIe
+traffic; keeps memory training + all tensors; VAEs/text-encoder/non-linears stay
+bf16. Ignored when a GGUF DiT is selected (already quantized).
+(`nodes.py` + `libs/ltx_distillation/models/ltx_wrapper.py` — new `quantization`
+param; the quantized build path skips the post-load dtype cast that would
+otherwise silently upcast fp8 back to bf16.)
+
+### 8. Tiled VAE decode (`JoyEcho_Generate`)
+Decoding a long high-res shot (e.g. 241f @ 1280×736) in one pass hard-aborts the
+VAE decode on a 24–32 GB card (fatal cuDNN abort mid-conv, not a catchable OOM).
+New `decode_tiling` (`auto`/`on`/`off`) routes decode through the vendored
+`VideoDecoder.tiled_decode` — **temporal-only** 64-frame chunks with 24-frame
+blended overlap (no spatial tiles → no spatial seams), streaming each chunk to
+CPU. `auto` engages only above a size threshold, so small renders keep the
+original single-pass decode bit-for-bit.
+(`nodes.py` + `libs/ltx_distillation/utils.py` — `decode_benchmark_sample` gains
+a `video_tiling_config` kwarg + `_decode_video_tiled_uint8`.)
+
+### 9. Model dropdown (`JoyEcho_ModelLoader`)
+New `model_file` combo lists every `.safetensors` / `.gguf` under the ComfyUI
+`checkpoints` / `diffusion_models` / `unet` dirs. Pick a `.safetensors` → full
+checkpoint (replaces `checkpoint_path`); pick a `.gguf` → DiT loaded from GGUF
+while `checkpoint_path` still supplies the VAEs / vocoder / text connectors.
+`"(use checkpoint_path)"` keeps the old typed-path behavior. Plus a clear
+early error if `gemma_path` is a `.gguf`/file/sidecar-less dir (this loader
+needs the HF `gemma-3-12b-it` folder; GGUF Gemma only works via
+`RebelsJE_TextEncoder`). (`nodes.py`)
+
+### 10. Automation / batching nodes (new)
+- **`JoyEcho_PromptSource`** — one dropdown listing LPFF-style `.txt` briefs
+  (from the inspire-pack prompts tree) **and** passthrough `.json` scripts
+  (`input/joyecho_prompts/`). Multi-block briefs fan out like
+  LoadPromptsFromFile. Emits `story_idea` (→ LLMEnhance) + `character`
+  (→ RefPicker) + `count`. Replaces the LPFF→UnzipPrompt chain and lets you
+  switch prompt sources with one dropdown instead of rewiring.
+- **`JoyEcho_RefPicker`** — auto-selects a character reference image from a
+  folder tree keyed by character name (typed, or wired from PromptSource, or
+  scanned from the prompt — dialogue mentions are stripped so only the on-screen
+  subject wins). `on_no_match=no_reference` returns nothing so a batch keeps
+  running.
+- **`JoyEcho_RefBatch`** — None-tolerant image batcher: combines up to 4
+  optional IMAGE inputs (e.g. two RefPickers for a two-character shot), skips
+  missing refs, resizes mismatched sizes to the first image, outputs `None` if
+  all are missing (Generate then just skips identity seeding). The stock KJNodes
+  `ImageBatchMulti` crashes with `'NoneType' has no attribute 'shape'` on a
+  missing ref; this replaces it.
+- **`JoyEcho_ScriptPicker`** — JSON dropdown (superseded by PromptSource; kept
+  for compatibility).
+
+---
+
+## Applying
+
+1. Back up your existing pack folder.
+2. Copy each file over the same relative path in
+   `ComfyUI/custom_nodes/ComfyUI_JoyAI_Echo_GGUF_Nodes/`.
+3. Restart ComfyUI. New widgets append at the **end** of existing nodes, so
+   saved graphs keep their values; the four new nodes appear under the
+   `JoyAI-Echo` category. Press `R` after adding model files to refresh the
+   `model_file` dropdown.
+
+The `libs/` files must match the vendored `ltx_core` / `ltx_distillation` in
+your pack (same JoyAI-Echo release). If your `libs/` differ substantially,
+cherry-pick the changes described above rather than overwriting.
+
+Not included (intentionally): model weights, the `gemma_assets/` tokenizer
+binaries, `.bak` snapshots, and `__pycache__`.

__init__.py

@@ -0,0 +1,121 @@
+"""ComfyUI nodes for JoyAI-Echo: minute-level multi-shot audio-video generation.
+
+Registers:
+  - the 7 upstream JoyEcho nodes (nodes.py)
+  - the discrete Rebels loader nodes (rebels_loaders.py)
+  - the Rebels staged single-node pipeline for 16GB RAM (rebels_staged.py)
+
+The libs/ folder MUST be on sys.path before anything imports ltx_core /
+ltx_distillation, so that setup happens first.
+"""
+
+import sys
+from pathlib import Path
+
+_NODE_ROOT = Path(__file__).resolve().parent
+_LIBS = str(_NODE_ROOT / "libs")
+
+if _LIBS not in sys.path:
+    sys.path.insert(0, _LIBS)
+
+# ---------------------------------------------------------------- base maps
+NODE_CLASS_MAPPINGS = {}
+NODE_DISPLAY_NAME_MAPPINGS = {}
+
+# ---------------------------------------------------------------- upstream JoyEcho nodes
+from .nodes import (
+    JoyEcho_ModelLoader,
+    JoyEcho_TextEncode,
+    JoyEcho_Generate,
+    JoyEcho_SingleShotGenerate,
+    JoyEcho_PromptFormat,
+    JoyEcho_LLMEnhance,
+    JoyEcho_PromptAtIndex,
+)
+
+NODE_CLASS_MAPPINGS.update({
+    "JoyEcho_ModelLoader": JoyEcho_ModelLoader,
+    "JoyEcho_TextEncode": JoyEcho_TextEncode,
+    "JoyEcho_Generate": JoyEcho_Generate,
+    "JoyEcho_SingleShotGenerate": JoyEcho_SingleShotGenerate,
+    "JoyEcho_PromptFormat": JoyEcho_PromptFormat,
+    "JoyEcho_LLMEnhance": JoyEcho_LLMEnhance,
+    "JoyEcho_PromptAtIndex": JoyEcho_PromptAtIndex,
+})
+NODE_DISPLAY_NAME_MAPPINGS.update({
+    "JoyEcho_ModelLoader": "JoyEcho Model Loader",
+    "JoyEcho_TextEncode": "JoyEcho Text Encode",
+    "JoyEcho_Generate": "JoyEcho Generate (Multi-Shot)",
+    "JoyEcho_SingleShotGenerate": "JoyEcho Single Shot Generate",
+    "JoyEcho_PromptFormat": "JoyEcho Prompt Format (Helper)",
+    "JoyEcho_LLMEnhance": "JoyEcho LLM Enhance",
+    "JoyEcho_PromptAtIndex": "JoyEcho Prompt At Index",
+})
+
+# ---------------------------------------------------------------- Rebels discrete loaders
+try:
+    from .rebels_loaders import (
+        NODE_CLASS_MAPPINGS as _RL_CM,
+        NODE_DISPLAY_NAME_MAPPINGS as _RL_DM,
+    )
+    NODE_CLASS_MAPPINGS.update(_RL_CM)
+    NODE_DISPLAY_NAME_MAPPINGS.update(_RL_DM)
+except Exception as e:
+    print(f"[Rebels JE] rebels_loaders failed to load: {e!r}", flush=True)
+
+# ---------------------------------------------------------------- Rebels staged pipeline (16GB)
+try:
+    from .rebels_staged import (
+        NODE_CLASS_MAPPINGS as _ST_CM,
+        NODE_DISPLAY_NAME_MAPPINGS as _ST_DM,
+    )
+    NODE_CLASS_MAPPINGS.update(_ST_CM)
+    NODE_DISPLAY_NAME_MAPPINGS.update(_ST_DM)
+except Exception as e:
+    print(f"[Rebels JE] rebels_staged failed to load: {e!r}", flush=True)
+
+# ---------------------------------------------------------------- Script picker (JSON dropdown)
+try:
+    from .joyecho_script_picker import (
+        NODE_CLASS_MAPPINGS as _SP_CM,
+        NODE_DISPLAY_NAME_MAPPINGS as _SP_DM,
+    )
+    NODE_CLASS_MAPPINGS.update(_SP_CM)
+    NODE_DISPLAY_NAME_MAPPINGS.update(_SP_DM)
+except Exception as e:
+    print(f"[Rebels JE] joyecho_script_picker failed to load: {e!r}", flush=True)
+
+# ---------------------------------------------------------------- Prompt source (unified txt+json dropdown)
+try:
+    from .joyecho_prompt_source import (
+        NODE_CLASS_MAPPINGS as _PS_CM,
+        NODE_DISPLAY_NAME_MAPPINGS as _PS_DM,
+    )
+    NODE_CLASS_MAPPINGS.update(_PS_CM)
+    NODE_DISPLAY_NAME_MAPPINGS.update(_PS_DM)
+except Exception as e:
+    print(f"[Rebels JE] joyecho_prompt_source failed to load: {e!r}", flush=True)
+
+# ---------------------------------------------------------------- Reference batch (None-tolerant)
+try:
+    from .joyecho_ref_batch import (
+        NODE_CLASS_MAPPINGS as _RB_CM,
+        NODE_DISPLAY_NAME_MAPPINGS as _RB_DM,
+    )
+    NODE_CLASS_MAPPINGS.update(_RB_CM)
+    NODE_DISPLAY_NAME_MAPPINGS.update(_RB_DM)
+except Exception as e:
+    print(f"[Rebels JE] joyecho_ref_batch failed to load: {e!r}", flush=True)
+
+# ---------------------------------------------------------------- Reference picker (auto by character)
+try:
+    from .joyecho_ref_picker import (
+        NODE_CLASS_MAPPINGS as _RP_CM,
+        NODE_DISPLAY_NAME_MAPPINGS as _RP_DM,
+    )
+    NODE_CLASS_MAPPINGS.update(_RP_CM)
+    NODE_DISPLAY_NAME_MAPPINGS.update(_RP_DM)
+except Exception as e:
+    print(f"[Rebels JE] joyecho_ref_picker failed to load: {e!r}", flush=True)
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

joyecho_prompt_source.py

@@ -0,0 +1,151 @@
+"""JoyEcho Prompt Source - ONE dropdown for both prompt pipelines.
+
+Lists LPFF-style brief/JSON-block .txt files (from the inspire-pack prompts
+tree) AND passthrough .json scripts (from <ComfyUI>/input/joyecho_prompts/) in
+a single combo, and always emits the same outputs:
+
+  story_idea (STRING, list) -> JoyEcho_LLMEnhance.story_idea
+  character  (STRING, list) -> JoyEcho_RefPicker.character
+  count      (INT)
+
+.txt files are parsed LPFF-style (blocks split on ---, positive:/negative:/
+name: fields); each block becomes one queue item, so multi-block files fan out
+exactly like LoadPromptsFromFile. .json files load as a single passthrough item
+(the raw JSON text). Blocks without a name: line emit "" (RefPicker falls
+through to prompt-scan / fallback), never the filename.
+
+Replaces the LPFF -> UnzipPrompt chain and the Script Picker: wire once,
+switch sources by picking a different file.
+"""
+
+import json
+import re
+from pathlib import Path
+
+import folder_paths
+
+_JSON_SUBDIR = "joyecho_prompts"
+_TXT_PREFIX = "TXT: "
+_JSON_PREFIX = "JSON: "
+_EMPTY = "(no prompt files found)"
+
+_BLOCK_SPLIT = re.compile(r"\n\s*-+\s*\n")
+_BLOCK_PATTERN = re.compile(
+    r"^(?:(?:name:(?P<name>.*?)|positive:(?P<positive>.*?)|negative:(?P<negative>.*?))\n*)+$",
+    re.DOTALL,
+)
+
+
+def _txt_root() -> Path | None:
+    try:
+        roots = folder_paths.get_folder_paths("inspire_prompts")
+        return Path(roots[0]) if roots else None
+    except Exception:
+        return None
+
+
+def _json_root() -> Path:
+    d = Path(folder_paths.get_input_directory()) / _JSON_SUBDIR
+    try:
+        d.mkdir(parents=True, exist_ok=True)
+    except OSError:
+        pass
+    return d
+
+
+def _list_files() -> list[str]:
+    out = []
+    root = _txt_root()
+    if root and root.is_dir():
+        for p in sorted(root.rglob("*.txt")):
+            out.append(_TXT_PREFIX + str(p.relative_to(root)))
+    jroot = _json_root()
+    for p in sorted(jroot.glob("*.json")):
+        out.append(_JSON_PREFIX + p.name)
+    return out or [_EMPTY]
+
+
+def _resolve(choice: str) -> Path | None:
+    if choice.startswith(_TXT_PREFIX):
+        root = _txt_root()
+        return (root / choice[len(_TXT_PREFIX):]) if root else None
+    if choice.startswith(_JSON_PREFIX):
+        return _json_root() / choice[len(_JSON_PREFIX):]
+    return None
+
+
+class JoyEcho_PromptSource:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "source_file": (_list_files(),),
+            },
+            "optional": {
+                "load_cap": ("INT", {"default": 0, "min": 0,
+                                     "tooltip": "TXT files only: 0 = all blocks, N = first N from start_index."}),
+                "start_index": ("INT", {"default": 0, "min": 0,
+                                        "tooltip": "TXT files only: first block index to load."}),
+                "character_override": ("STRING", {"default": "",
+                                                  "tooltip": "Force this character for EVERY emitted item "
+                                                             "(useful for .json scripts, which carry no name field)."}),
+            },
+        }
+
+    RETURN_TYPES = ("STRING", "STRING", "INT")
+    RETURN_NAMES = ("story_idea", "character", "count")
+    OUTPUT_IS_LIST = (True, True, False)
+    FUNCTION = "load"
+    CATEGORY = "JoyAI-Echo"
+
+    @classmethod
+    def IS_CHANGED(cls, source_file, load_cap=0, start_index=0, character_override=""):
+        p = _resolve(source_file)
+        try:
+            return f"{source_file}:{p.stat().st_mtime}:{load_cap}:{start_index}:{character_override}"
+        except Exception:
+            return source_file
+
+    def load(self, source_file, load_cap=0, start_index=0, character_override=""):
+        if source_file == _EMPTY:
+            raise ValueError("PromptSource: no prompt files found in the inspire prompts tree "
+                             "or input/joyecho_prompts/.")
+        p = _resolve(source_file)
+        if p is None or not p.exists():
+            raise FileNotFoundError(f"PromptSource: {source_file} -> {p} not found. "
+                                    "Refresh the node list (R) after adding files.")
+        text = p.read_text(encoding="utf-8")
+        override = character_override.strip().lower()
+
+        if source_file.startswith(_JSON_PREFIX):
+            data = json.loads(text)  # fail early with a clear error
+            arr = data.get("prompts") or data.get("shots")
+            if not isinstance(arr, list) or not arr:
+                raise ValueError(f"{p.name} must contain a non-empty 'prompts' (or 'shots') array.")
+            print(f"[JoyEcho] PromptSource: {p.name} (json, {len(arr)} shots, 1 item).", flush=True)
+            return ([text], [override], 1)
+
+        # TXT: LPFF-style blocks
+        items, names = [], []
+        for blk in _BLOCK_SPLIT.split(text):
+            m = _BLOCK_PATTERN.search(blk)
+            if not m or m.group("positive") is None:
+                continue
+            items.append(m.group("positive").strip())
+            nm = (m.group("name") or "").strip().lower()
+            names.append(override or nm)
+        total = len(items)
+        items = items[start_index:]
+        names = names[start_index:]
+        if load_cap > 0:
+            items = items[:load_cap]
+            names = names[:load_cap]
+        if not items:
+            raise ValueError(f"PromptSource: {p.name} yielded no blocks "
+                             f"(total {total}, start_index {start_index}, load_cap {load_cap}).")
+        print(f"[JoyEcho] PromptSource: {p.name} (txt, {len(items)}/{total} blocks).", flush=True)
+        return (items, names, len(items))
+
+
+NODE_CLASS_MAPPINGS = {"JoyEcho_PromptSource": JoyEcho_PromptSource}
+NODE_DISPLAY_NAME_MAPPINGS = {"JoyEcho_PromptSource": "JoyEcho Prompt Source (txt briefs + json scripts)"}

joyecho_ref_batch.py

@@ -0,0 +1,57 @@
+"""JoyEcho Reference Batch - None-tolerant image batcher for reference images.
+
+Combines up to 4 optional IMAGE inputs into one batch for
+JoyEcho_Generate.reference_image. Unlike generic batch nodes (KJNodes
+ImageBatchMulti crashes on None), inputs that are missing - e.g. a RefPicker
+that found no character match and emitted its "no reference" None - are simply
+skipped. If ALL inputs are missing, outputs None, which Generate treats as
+"no reference wired": the item renders without identity seeding instead of
+killing the queue.
+
+Mixed sizes are resized (lanczos, center-crop semantics via common_upscale)
+to the first present image's dimensions; Generate cover-fits every ref frame
+to the video dimensions afterwards anyway.
+"""
+
+import torch
+
+import comfy.utils
+
+
+class JoyEcho_RefBatch:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {},
+            "optional": {
+                "image_1": ("IMAGE",),
+                "image_2": ("IMAGE",),
+                "image_3": ("IMAGE",),
+                "image_4": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("reference_image",)
+    FUNCTION = "batch"
+    CATEGORY = "JoyAI-Echo"
+
+    def batch(self, image_1=None, image_2=None, image_3=None, image_4=None):
+        imgs = [i for i in (image_1, image_2, image_3, image_4)
+                if i is not None and i.shape[0] > 0]
+        if not imgs:
+            print("[JoyEcho] RefBatch: no reference images present; "
+                  "passing None (Generate skips identity seeding).", flush=True)
+            return (None,)
+        h, w = imgs[0].shape[1], imgs[0].shape[2]
+        out = []
+        for img in imgs:
+            if img.shape[1] != h or img.shape[2] != w:
+                img = comfy.utils.common_upscale(
+                    img.movedim(-1, 1), w, h, "lanczos", "center").movedim(1, -1)
+            out.append(img)
+        return (torch.cat(out, dim=0),)
+
+
+NODE_CLASS_MAPPINGS = {"JoyEcho_RefBatch": JoyEcho_RefBatch}
+NODE_DISPLAY_NAME_MAPPINGS = {"JoyEcho_RefBatch": "JoyEcho Reference Batch (None-tolerant)"}

joyecho_ref_picker.py

@@ -0,0 +1,186 @@
+"""JoyEcho Reference Picker - auto-select a character reference image.
+
+Feeds JoyEcho_Generate's reference_image input in LPFF batch queues:
+  LPFF block carries `name: zara` -> UnzipPrompt name output -> this node ->
+  picks an image from <ComfyUI>/input/joyecho_refs/zara/ -> IMAGE out.
+
+Resolution order per run:
+  1. `character` input (usually UnzipPrompt's name output), lowercased.
+     LPFF quirk: blocks WITHOUT a name: line emit the prompt FILENAME here -
+     that never matches a folder, so it falls through cleanly.
+  2. scan `prompt_text` for any refs-folder name as a whole word (longest first).
+  3. `fallback_image` input if wired.
+  4. clear error.
+
+Pick strategies: by_seed (sorted files, seed % count - reproducible, vary the
+seed to vary the ref), first, newest.
+"""
+
+import os
+import re
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image, ImageOps
+
+import folder_paths
+
+_REFS_SUBDIR = "joyecho_refs"
+_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
+
+
+def _refs_root(custom_root: str = "") -> Path:
+    if custom_root and custom_root.strip():
+        return Path(custom_root.strip())
+    d = Path(folder_paths.get_input_directory()) / _REFS_SUBDIR
+    try:
+        d.mkdir(parents=True, exist_ok=True)
+    except OSError:
+        pass
+    return d
+
+
+def _character_dirs(root: Path) -> list[str]:
+    try:
+        return sorted(p.name for p in root.iterdir() if p.is_dir())
+    except OSError:
+        return []
+
+
+def _images_in(folder: Path) -> list[Path]:
+    try:
+        return sorted(p for p in folder.iterdir()
+                      if p.is_file() and p.suffix.lower() in _EXTS)
+    except OSError:
+        return []
+
+
+def _load_image(path: Path) -> torch.Tensor:
+    img = Image.open(path)
+    img = ImageOps.exif_transpose(img).convert("RGB")
+    arr = np.asarray(img).astype(np.float32) / 255.0
+    return torch.from_numpy(arr)[None, ...]  # [1, H, W, C]
+
+
+class JoyEcho_RefPicker:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "pick": (["by_seed", "first", "newest"],),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 2**31 - 1,
+                                 "tooltip": "Used by by_seed: index = seed % image count."}),
+            },
+            "optional": {
+                "refs_root": ("STRING", {
+                    "default": "G:\\RIFT Assets\\Rift Character Reference Images",
+                    "tooltip": "Root folder holding one subfolder per character. Empty = ComfyUI/input/joyecho_refs/. Folder-name matching is case-insensitive on Windows.",
+                }),
+                "on_no_match": (["no_reference", "error"], {
+                    "default": "no_reference",
+                    "tooltip": "When no character matches and no fallback_image is wired: "
+                               "no_reference = output nothing (Generate simply skips identity "
+                               "seeding for this item; the batch keeps running). error = stop the run.",
+                }),
+                "character": ("STRING", {"default": "",
+                                         "tooltip": "Character folder name (e.g. marcus). TYPE it here for a "
+                                                    "manual pick, or right-click the node > 'Convert character "
+                                                    "to input' and wire PromptSource's character output for "
+                                                    "automatic per-item picks."}),
+                "prompt_text": ("STRING", {"default": "", "forceInput": True,
+                                           "tooltip": "Fallback: scanned for any refs folder name as a whole word."}),
+                "fallback_image": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE", "STRING",)
+    RETURN_NAMES = ("reference_image", "picked_path",)
+    FUNCTION = "pick_ref"
+    CATEGORY = "JoyAI-Echo"
+
+    @classmethod
+    def IS_CHANGED(cls, pick, seed, on_no_match="no_reference", refs_root="", character="", prompt_text="", fallback_image=None):
+        # Re-run when the resolved folder's contents change.
+        root = _refs_root(refs_root)
+        sig = [pick, str(seed), str(root), character.strip().lower()]
+        for d in _character_dirs(root):
+            folder = root / d
+            imgs = _images_in(folder)
+            sig.append(f"{d}:{len(imgs)}:{max((p.stat().st_mtime for p in imgs), default=0)}")
+        return "|".join(sig)
+
+    def pick_ref(self, pick, seed, on_no_match="no_reference", refs_root="", character="", prompt_text="", fallback_image=None):
+        root = _refs_root(refs_root)
+        dirs = _character_dirs(root)
+
+        chosen_dir = None
+        want = (character or "").strip().lower()
+        # LPFF quirk: blocks without a `name:` line emit the prompt FILENAME as
+        # the name - anything path/file-shaped is not a character.
+        if any(s in want for s in ("\\", "/", ".txt", ".json")):
+            want = ""
+        if want and (root / want).is_dir():
+            chosen_dir = root / want
+
+        if chosen_dir is None and prompt_text:
+            # Ignore names spoken INSIDE dialogue: absent characters get
+            # mentioned in quotes ("Alana thinks I am imagining it"), while
+            # on-screen characters are named in the descriptive prose. Strip
+            # JSON-escaped quotes, plain double quotes, and says,-introduced
+            # single-quoted lines before scanning.
+            scrub = prompt_text
+            scrub = re.sub(r'\\"(?:[^"\\]|\\.)*?\\"', " ", scrub)      # \"...\" (JSON-escaped)
+            scrub = re.sub(r'"(?:[^"\\]|\\.)*?"', " ", scrub)          # "..."
+            scrub = re.sub(r"says,\s*'(?:[^'])*?'", " ", scrub)        # says, '...'
+            low = scrub.lower()
+            # The subject of a brief dominates its text: most-mentioned folder
+            # name wins; earliest first-mention breaks ties. (Longest-match-first
+            # wrongly picked a side character once - MARCUS over ZARA.)
+            best = None  # (count, -first_pos, dirname)
+            for d in dirs:
+                hits = [m.start() for m in re.finditer(r"\b" + re.escape(d.lower()) + r"\b", low)]
+                if hits:
+                    key = (len(hits), -hits[0])
+                    if best is None or key > best[0]:
+                        best = (key, d)
+            if best is not None:
+                chosen_dir = root / best[1]
+
+        if chosen_dir is None:
+            if fallback_image is not None:
+                print("[JoyEcho] RefPicker: no character match; using fallback_image.", flush=True)
+                return (fallback_image, "(fallback_image)")
+            if on_no_match == "no_reference":
+                print(f"[JoyEcho] RefPicker: no character match (character={character!r}); "
+                      f"continuing WITHOUT a reference.", flush=True)
+                return (None, "(no reference)")
+            raise ValueError(
+                f"RefPicker: no reference folder matched. character={character!r}, "
+                f"available folders in {root}: {dirs or '(none - create input/joyecho_refs/<name>/)'}"
+            )
+
+        imgs = _images_in(chosen_dir)
+        if not imgs:
+            if fallback_image is not None:
+                print(f"[JoyEcho] RefPicker: {chosen_dir.name}/ is empty; using fallback_image.", flush=True)
+                return (fallback_image, "(fallback_image)")
+            if on_no_match == "no_reference":
+                print(f"[JoyEcho] RefPicker: {chosen_dir.name}/ is empty; continuing WITHOUT a reference.", flush=True)
+                return (None, "(no reference)")
+            raise ValueError(f"RefPicker: no images in {chosen_dir} (put .png/.jpg refs there).")
+
+        if pick == "first":
+            path = imgs[0]
+        elif pick == "newest":
+            path = max(imgs, key=lambda p: p.stat().st_mtime)
+        else:  # by_seed
+            path = imgs[seed % len(imgs)]
+
+        print(f"[JoyEcho] RefPicker: {chosen_dir.name} -> {path.name} "
+              f"({pick}, {len(imgs)} candidates).", flush=True)
+        return (_load_image(path), str(path))
+
+
+NODE_CLASS_MAPPINGS = {"JoyEcho_RefPicker": JoyEcho_RefPicker}
+NODE_DISPLAY_NAME_MAPPINGS = {"JoyEcho_RefPicker": "JoyEcho Reference Picker (auto by character)"}

joyecho_script_picker.py

@@ -0,0 +1,89 @@
+"""JoyEcho Script Picker — a dropdown node for selecting a saved prompt JSON.
+
+The stock JoyEcho_TextEncode 'prompts' field only accepts pasted JSON / a typed
+path, which is painful in the canvas. This node lists every *.json in
+  <ComfyUI>/input/joyecho_prompts/
+as a COMBO dropdown and outputs the file's contents, wired straight into
+JoyEcho_TextEncode's 'prompts' input (which accepts inline {"prompts":[...]} JSON).
+
+Add a .json to that folder, hit the ComfyUI refresh button (or press R) to
+repopulate the dropdown, pick it, run. Editing the file re-triggers execution
+automatically (IS_CHANGED tracks mtime) — no need to reselect.
+"""
+
+import json
+from pathlib import Path
+
+import folder_paths
+
+_PROMPTS_SUBDIR = "joyecho_prompts"
+_EMPTY = "(no .json in input/joyecho_prompts)"
+
+
+def _scripts_dir() -> Path:
+    d = Path(folder_paths.get_input_directory()) / _PROMPTS_SUBDIR
+    try:
+        d.mkdir(parents=True, exist_ok=True)
+    except OSError:
+        pass
+    return d
+
+
+def _list_scripts() -> list[str]:
+    d = _scripts_dir()
+    try:
+        files = sorted(p.name for p in d.glob("*.json"))
+    except OSError:
+        files = []
+    return files if files else [_EMPTY]
+
+
+class JoyEcho_ScriptPicker:
+    """Pick a prompt-script .json from input/joyecho_prompts via a dropdown."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {"script": (_list_scripts(),)}}
+
+    RETURN_TYPES = ("STRING", "STRING",)
+    RETURN_NAMES = ("prompts_json", "path",)
+    FUNCTION = "load"
+    CATEGORY = "JoyAI-Echo"
+
+    @classmethod
+    def IS_CHANGED(cls, script):
+        # Re-run when the selected file changes on disk, so edits are picked up.
+        p = _scripts_dir() / script
+        try:
+            return f"{script}:{p.stat().st_mtime}"
+        except OSError:
+            return script
+
+    def load(self, script):
+        if script == _EMPTY:
+            raise ValueError(
+                f"No .json scripts found. Put your prompt JSON in {_scripts_dir()} "
+                "and press the refresh button (or R) to repopulate the dropdown."
+            )
+        p = _scripts_dir() / script
+        if not p.exists():
+            raise FileNotFoundError(
+                f"Script not found: {p}. Refresh the node list (R) after adding files."
+            )
+        text = p.read_text(encoding="utf-8")
+        # Fail early with a clear message rather than deep in the text encoder.
+        try:
+            data = json.loads(text)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"{script} is not valid JSON: {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(f"{script} must contain a non-empty 'prompts' (or 'shots') array.")
+        print(f"[JoyEcho] ScriptPicker: {script} ({len(arr)} shots).", flush=True)
+        return (text, str(p),)
+
+
+NODE_CLASS_MAPPINGS = {"JoyEcho_ScriptPicker": JoyEcho_ScriptPicker}
+NODE_DISPLAY_NAME_MAPPINGS = {"JoyEcho_ScriptPicker": "JoyEcho Script Picker (JSON dropdown)"}

libs/ltx_distillation/models/ltx_wrapper.py

@@ -0,0 +1,835 @@
+"""
+LTX-2 Diffusion Model Wrapper for DMD distillation.
+
+This wrapper adapts LTX-2's audio-video joint generation model for use in
+DMD (Distribution Matching Distillation) training.
+
+Model Architecture:
+- patch_size = (1, 1, 1): No spatial/temporal grouping
+- Patchification: Simple reshape [B, C, F, H, W] → [B, F*H*W, C]
+- Each token: 128-dimensional latent vector (one per spatial-temporal position)
+- Model input projection: Linear(128, 4096)
+"""
+
+from dataclasses import replace
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from ltx_core.components.patchifiers import (
+    AudioPatchifier,
+    VideoLatentPatchifier,
+    get_pixel_coords,
+)
+from ltx_core.guidance.perturbations import (
+    BatchedPerturbationConfig,
+    Perturbation,
+    PerturbationConfig,
+    PerturbationType,
+)
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer import LTXModel, X0Model
+from ltx_core.model.transformer.modality import Modality
+from ltx_core.types import (
+    AudioLatentShape,
+    SpatioTemporalScaleFactors,
+    VideoLatentShape,
+)
+
+
+class LTX2DiffusionWrapper(nn.Module):
+    """
+    Wrapper for LTX-2 model to provide DMD-compatible interface.
+
+    Handles:
+    - Input format conversion: [B, F, C, H, W] -> Modality
+    - Timestep handling: sigma values for all tokens
+    - Position computation for video (3D) and audio (1D)
+    - Output format: x0 predictions for both video and audio
+
+    Uses official LTX-2 patchifiers (patch_size=1) to ensure consistency
+    with the pretrained model weights.
+    """
+
+    # Time alignment constants
+    VIDEO_LATENT_FPS = 3.0  # 24fps / 8 (VAE compression)
+    AUDIO_LATENT_FPS = 25.0  # 16kHz / 160 / 4 (mel hop / VAE compression)
+    ALIGNMENT_RATIO = AUDIO_LATENT_FPS / VIDEO_LATENT_FPS  # ~8.33
+
+    # Video FPS for position computation
+    VIDEO_FPS = 24.0
+
+    # VAE scale factors (temporal=8, height=32, width=32)
+    DEFAULT_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
+
+    def __init__(
+        self,
+        model: LTXModel,
+        video_height: int = 512,
+        video_width: int = 768,
+        vae_spatial_compression: int = 32,
+    ):
+        """
+        Args:
+            model: X0Model instance (wraps velocity model, returns x0 predictions)
+            video_height: Video height in pixels
+            video_width: Video width in pixels
+            vae_spatial_compression: VAE spatial compression factor
+        """
+        super().__init__()
+        self.model = model
+        self.video_height = video_height
+        self.video_width = video_width
+        self.vae_spatial_compression = vae_spatial_compression
+
+        # Compute latent dimensions
+        self.latent_height = video_height // vae_spatial_compression  # 16
+        self.latent_width = video_width // vae_spatial_compression    # 24
+
+        # Official patchifiers with patch_size=1 (no spatial grouping)
+        self.video_patchifier = VideoLatentPatchifier(patch_size=1)
+        self.audio_patchifier = AudioPatchifier(patch_size=1)
+
+        # Frame sequence length: with patch_size=1, each spatial position is one token
+        # For 512x768: H'*W' = 16*24 = 384 tokens per frame
+        self.video_frame_seqlen = self.latent_height * self.latent_width  # 384
+
+    def set_module_grad(self, module_grad: Dict[str, bool]) -> None:
+        """
+        Set gradient requirements for model components.
+
+        Args:
+            module_grad: Dict mapping component names to requires_grad flags
+        """
+        if module_grad.get("model", True):
+            self.model.requires_grad_(True)
+        else:
+            self.model.requires_grad_(False)
+            self.model.eval()
+
+    def enable_gradient_checkpointing(self) -> None:
+        """Enable gradient checkpointing for memory efficiency."""
+        if hasattr(self.model, "velocity_model"):
+            self.model.velocity_model.set_gradient_checkpointing(True)
+        elif hasattr(self.model, "set_gradient_checkpointing"):
+            self.model.set_gradient_checkpointing(True)
+
+    def _flatten_video_latent(
+        self,
+        video_latent: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Flatten video latent from [B, F, C, H, W] to [B, T, C] using patch_size=1.
+
+        With patch_size=1, this is a simple reshape — no spatial grouping.
+        The official VideoLatentPatchifier(patch_size=1) does:
+            "b c (f 1) (h 1) (w 1) -> b (f h w) (c 1 1 1)" = "b c f h w -> b (f h w) c"
+
+        Args:
+            video_latent: Shape [B, F, C, H, W] where
+                - F: number of latent frames
+                - C: latent channels (128)
+                - H, W: latent spatial dimensions (16, 24)
+
+        Returns:
+            Flattened tensor [B, T, C] where:
+            - T = F * H * W (e.g., 16 * 16 * 24 = 6144)
+            - C = 128 (unchanged, since patch_size=1)
+        """
+        B, F, C, H, W = video_latent.shape
+        assert C == 128, (
+            f"Expected video latent C=128 at dim 2, got shape {video_latent.shape}. "
+            f"Input should be [B, F, C, H, W] with C=128."
+        )
+
+        # Convert from [B, F, C, H, W] to [B, C, F, H, W] (official format)
+        video_latent = video_latent.permute(0, 2, 1, 3, 4)
+
+        # Use official patchifier: [B, C, F, H, W] -> [B, F*H*W, C]
+        # With patch_size=1 this is equivalent to:
+        #   einops.rearrange(x, "b c f h w -> b (f h w) c")
+        video_latent = self.video_patchifier.patchify(video_latent)
+
+        return video_latent
+
+    def _unflatten_video_latent(
+        self,
+        flat_latent: torch.Tensor,
+        num_frames: int,
+    ) -> torch.Tensor:
+        """
+        Unflatten video latent from [B, T, C] back to [B, F, C, H, W].
+
+        Args:
+            flat_latent: Shape [B, T, C] where C = 128 (patch_size=1)
+            num_frames: Number of latent frames F
+
+        Returns:
+            Video latent [B, F, C, H, W]
+        """
+        B, T, C = flat_latent.shape
+        H = self.latent_height
+        W = self.latent_width
+        F = num_frames
+
+        # Use official unpatchifier: [B, T, C] -> [B, C, F, H, W]
+        output_shape = VideoLatentShape(
+            batch=B, channels=C, frames=F, height=H, width=W
+        )
+        video_latent = self.video_patchifier.unpatchify(flat_latent, output_shape)
+
+        # Convert from [B, C, F, H, W] to [B, F, C, H, W] (DMD format)
+        video_latent = video_latent.permute(0, 2, 1, 3, 4)
+
+        return video_latent
+
+    def _compute_video_positions(
+        self,
+        video_latent: torch.Tensor,
+        downscale_factor: int = 1,
+        start_frame: int = 0,
+    ) -> torch.Tensor:
+        """
+        Compute 3D position indices for video tokens with [start, end) bounds.
+
+        Uses the official VideoLatentPatchifier.get_patch_grid_bounds() and
+        get_pixel_coords() to ensure consistency with the pretrained model.
+
+        The RoPE computation expects positions in the format [B, 3, T, 2] where:
+        - dim 1 (size 3): temporal, height, width dimensions
+        - dim 3 (size 2): [start, end) bounds for each patch
+
+        Returns:
+            Position tensor of shape [B, 3, T, 2] with patch bounds in pixel space
+        """
+        B, F, C, H, W = video_latent.shape
+        device = video_latent.device
+
+        # Build VideoLatentShape for the patchifier
+        video_shape = VideoLatentShape(
+            batch=B, channels=C, frames=F, height=H, width=W
+        )
+
+        # Get patch grid bounds in latent coordinates: [B, 3, T, 2]
+        # With patch_size=1, each token covers [i, i+1) in each dimension
+        latent_coords = self.video_patchifier.get_patch_grid_bounds(
+            output_shape=video_shape,
+            device=device,
+        )
+        if start_frame != 0:
+            latent_coords = latent_coords.clone()
+            latent_coords[:, 0, :, :] += int(start_frame)
+
+        # Convert to pixel coordinates using official helper
+        # Applies scale_factors (temporal=8, height=32, width=32)
+        # and causal_fix (first frame temporal offset)
+        pixel_coords = get_pixel_coords(
+            latent_coords=latent_coords,
+            scale_factors=self.DEFAULT_SCALE_FACTORS,
+            causal_fix=True,
+        ).float()
+
+        # Convert temporal dimension from frames to seconds (divide by fps=24)
+        # This matches VideoLatentTools.create_initial_state
+        pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / self.VIDEO_FPS
+
+        if downscale_factor != 1:
+            pixel_coords = pixel_coords.clone()
+            pixel_coords[:, 1, ...] *= downscale_factor
+            pixel_coords[:, 2, ...] *= downscale_factor
+
+        return pixel_coords
+
+    # Audio timing constants (from AudioPatchifier defaults)
+    AUDIO_SAMPLE_RATE = 16000
+    AUDIO_HOP_LENGTH = 160
+    AUDIO_LATENT_DOWNSAMPLE_FACTOR = 4
+    AUDIO_IS_CAUSAL = True
+
+    def _get_audio_latent_time_in_sec(
+        self,
+        start_latent: int,
+        end_latent: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> torch.Tensor:
+        """
+        Converts latent indices into real-time seconds while honoring causal
+        offsets and the configured hop length.
+
+        Matches AudioPatchifier._get_audio_latent_time_in_sec exactly.
+        """
+        audio_latent_frame = torch.arange(start_latent, end_latent, dtype=dtype, device=device)
+        audio_mel_frame = audio_latent_frame * self.AUDIO_LATENT_DOWNSAMPLE_FACTOR
+
+        if self.AUDIO_IS_CAUSAL:
+            # Frame offset for causal alignment.
+            causal_offset = 1
+            audio_mel_frame = (audio_mel_frame + causal_offset - self.AUDIO_LATENT_DOWNSAMPLE_FACTOR).clip(min=0)
+
+        return audio_mel_frame * self.AUDIO_HOP_LENGTH / self.AUDIO_SAMPLE_RATE
+
+    def _compute_audio_positions(
+        self,
+        audio_latent: torch.Tensor,
+        start_frame: int = 0,
+    ) -> torch.Tensor:
+        """
+        Compute 1D temporal positions for audio tokens with [start, end) bounds.
+
+        The RoPE computation expects positions in the format [B, 1, T, 2] where:
+        - dim 1 (size 1): temporal dimension only (audio is 1D)
+        - dim 3 (size 2): [start, end) bounds in seconds
+
+        Returns:
+            Position tensor of shape [B, 1, T, 2] with temporal bounds in seconds
+        """
+        B, T, C = audio_latent.shape
+        device = audio_latent.device
+
+        # Compute start timings for each audio frame
+        start_timings = self._get_audio_latent_time_in_sec(
+            int(start_frame), int(start_frame) + T, torch.float32, device
+        )
+        start_timings = start_timings.unsqueeze(0).expand(B, -1).unsqueeze(1)  # [B, 1, T]
+
+        # Compute end timings for each audio frame (shifted by 1)
+        end_timings = self._get_audio_latent_time_in_sec(
+            int(start_frame) + 1, int(start_frame) + T + 1, torch.float32, device
+        )
+        end_timings = end_timings.unsqueeze(0).expand(B, -1).unsqueeze(1)  # [B, 1, T]
+
+        # Stack to create [B, 1, T, 2] with [start, end) bounds
+        positions = torch.stack([start_timings, end_timings], dim=-1)
+
+        return positions
+
+    def _compute_timesteps_for_tokens(
+        self,
+        sigma: torch.Tensor,
+        num_tokens: int,
+        tokens_per_frame: int,
+    ) -> torch.Tensor:
+        """
+        Expand sigma to per-token timesteps.
+
+        In the official pipeline, timesteps = denoise_mask * sigma, producing
+        shape [B, T, 1]. Here we replicate sigma to each token belonging to
+        the same frame and add a trailing dimension for broadcasting with
+        the latent channels.
+
+        Args:
+            sigma: Shape [B] or [B, F] - sigma values per frame
+            num_tokens: Total number of tokens
+            tokens_per_frame: Number of tokens per frame
+
+        Returns:
+            Timesteps tensor [B, T, 1] for correct broadcasting with [B, T, C]
+        """
+        B = sigma.shape[0]
+
+        if sigma.dim() == 1:
+            # Single sigma per sample -> expand to all tokens
+            return sigma.view(B, 1, 1).expand(B, num_tokens, 1)
+        else:
+            # Per-frame sigma [B, F] -> expand to per-token [B, T, 1]
+            F = sigma.shape[1]
+            expanded = sigma.unsqueeze(2).expand(B, F, tokens_per_frame).reshape(B, -1)
+            return expanded.unsqueeze(-1)  # [B, T, 1]
+
+    @staticmethod
+    def _memory_slot_ranges(total_seq_len: int, num_slots: int) -> list[tuple[int, int]]:
+        if total_seq_len <= 0 or num_slots <= 0:
+            return []
+
+        ranges: list[tuple[int, int]] = []
+        start = 0
+        for slot_idx in range(num_slots):
+            end = round((slot_idx + 1) * total_seq_len / num_slots)
+            if end > start:
+                ranges.append((start, end))
+            start = end
+        return ranges
+
+    @staticmethod
+    def _memory_slot_ranges_from_lengths(
+        lengths: tuple[int, ...] | None,
+        *,
+        total_seq_len: int,
+        num_slots: int,
+    ) -> list[tuple[int, int]]:
+        if not lengths or len(lengths) != num_slots:
+            return LTX2DiffusionWrapper._memory_slot_ranges(total_seq_len, num_slots)
+
+        ranges: list[tuple[int, int]] = []
+        start = 0
+        for raw_length in lengths:
+            length = max(0, int(raw_length))
+            end = min(start + length, total_seq_len)
+            if end > start:
+                ranges.append((start, end))
+            start = end
+        if start != total_seq_len:
+            return LTX2DiffusionWrapper._memory_slot_ranges(total_seq_len, num_slots)
+        return ranges
+
+    @classmethod
+    def _build_paired_memory_cross_mask(
+        cls,
+        *,
+        batch_size: int,
+        query_memory_seq_len: int,
+        query_target_seq_len: int,
+        kv_memory_seq_len: int,
+        kv_target_seq_len: int,
+        num_memory_slots: int,
+        device: torch.device,
+        query_segment_lengths: tuple[tuple[int, ...], ...] | None = None,
+        kv_segment_lengths: tuple[tuple[int, ...], ...] | None = None,
+    ) -> torch.Tensor:
+        query_total_seq_len = query_memory_seq_len + query_target_seq_len
+        kv_total_seq_len = kv_memory_seq_len + kv_target_seq_len
+        mask = torch.zeros(
+            batch_size,
+            query_total_seq_len,
+            kv_total_seq_len,
+            dtype=torch.bool,
+            device=device,
+        )
+
+        for batch_idx in range(batch_size):
+            query_lengths = (
+                query_segment_lengths[batch_idx]
+                if query_segment_lengths is not None and batch_idx < len(query_segment_lengths)
+                else None
+            )
+            kv_lengths = (
+                kv_segment_lengths[batch_idx]
+                if kv_segment_lengths is not None and batch_idx < len(kv_segment_lengths)
+                else None
+            )
+            query_ranges = cls._memory_slot_ranges_from_lengths(
+                query_lengths,
+                total_seq_len=query_memory_seq_len,
+                num_slots=num_memory_slots,
+            )
+            kv_ranges = cls._memory_slot_ranges_from_lengths(
+                kv_lengths,
+                total_seq_len=kv_memory_seq_len,
+                num_slots=num_memory_slots,
+            )
+            for (q_start, q_end), (k_start, k_end) in zip(query_ranges, kv_ranges, strict=False):
+                mask[batch_idx, q_start:q_end, k_start:k_end] = True
+
+        if query_target_seq_len > 0 and kv_target_seq_len > 0:
+            mask[:, query_memory_seq_len:, kv_memory_seq_len:] = True
+        return mask
+
+    @staticmethod
+    def _build_memory_self_attention_block_mask(
+        *,
+        batch_size: int,
+        memory_seq_len: int,
+        target_seq_len: int,
+        device: torch.device,
+    ) -> torch.Tensor | None:
+        if memory_seq_len <= 0:
+            return None
+
+        total_seq_len = memory_seq_len + target_seq_len
+        attention_mask = torch.ones(
+            batch_size,
+            total_seq_len,
+            total_seq_len,
+            dtype=torch.bool,
+            device=device,
+        )
+        attention_mask[:, :, :memory_seq_len] = False
+        attention_mask[:, :memory_seq_len, :] = False
+        attention_mask[:, :memory_seq_len, :memory_seq_len] = True
+        return attention_mask
+
+    def forward(
+        self,
+        noisy_image_or_video: torch.Tensor,
+        conditional_dict: Dict[str, Any],
+        timestep: torch.Tensor,
+        noisy_audio: Optional[torch.Tensor] = None,
+        audio_timestep: Optional[torch.Tensor] = None,
+        memory_video: Optional[torch.Tensor] = None,
+        memory_audio: Optional[torch.Tensor] = None,
+        memory_audio_timestep: Optional[torch.Tensor] = None,
+        memory_audio_segment_lengths: tuple[tuple[int, ...], ...] | None = None,
+        paired_audio_memory: bool = False,
+        v2a_grad_scale: float = 1.0,
+        memory_position_mode: str = "reference",
+        memory_downscale_factor: int = 1,
+        skip_a2v_cross_attn: bool = False,
+        skip_v2a_cross_attn: bool = False,
+        skip_video_self_attn: bool = False,
+        skip_audio_self_attn: bool = False,
+        use_causal_timestep: bool = False,  # ignored, for API compatibility
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """
+        Forward pass for DMD distillation.
+
+        Args:
+            noisy_image_or_video: Noisy video latent [B, F, C, H, W]
+            conditional_dict: Dictionary containing:
+                - video_context: [B, seq_len, dim]
+                - audio_context: [B, seq_len, dim]
+                - attention_mask: [B, seq_len]
+            timestep: Sigma values [B] or [B, F]
+            noisy_audio: Noisy audio latent [B, F_a, C_audio] (optional)
+                where C_audio = 128 (= 8 channels * 16 mel_bins, post-patchify)
+            audio_timestep: Audio sigma values [B] or [B, F_a] (optional)
+            memory_audio: Optional clean memory-audio prefix [B, F_mem_a, C_audio]
+            memory_audio_timestep: Optional memory-audio sigma values [B] or [B, F_mem_a]
+
+        Returns:
+            Tuple of (video_x0_pred, audio_x0_pred)
+            - video_x0_pred: [B, F, C, H, W]
+            - audio_x0_pred: [B, F_a, C_audio] or None
+        """
+        B = noisy_image_or_video.shape[0]
+        num_video_frames = noisy_image_or_video.shape[1]
+        device = noisy_image_or_video.device
+        memory_position_mode = str(memory_position_mode).lower()
+        if memory_position_mode == "reference":
+            memory_position_mode = "legacy"
+        if memory_position_mode not in {"legacy", "prefix_continuous"}:
+            raise ValueError(
+                "memory_position_mode must be one of {'reference', 'legacy', 'prefix_continuous'}, "
+                f"got {memory_position_mode}"
+            )
+        if memory_video is not None and int(memory_video.shape[1]) == 0:
+            memory_video = None
+
+        # Flatten target video latent: [B, F, C, H, W] -> [B, T, C]
+        # With patch_size=1: T = F*H*W, C = 128
+        target_video_flat = self._flatten_video_latent(noisy_image_or_video)
+        num_target_video_tokens = target_video_flat.shape[1]
+
+        # Compute target video positions / timesteps
+        target_video_position_start = (
+            int(memory_video.shape[1])
+            if memory_position_mode == "prefix_continuous" and memory_video is not None
+            else 0
+        )
+        target_video_positions = self._compute_video_positions(
+            noisy_image_or_video,
+            start_frame=target_video_position_start,
+        )
+        target_video_timesteps = self._compute_timesteps_for_tokens(
+            timestep, num_target_video_tokens, self.video_frame_seqlen
+        )
+
+        memory_seq_len = 0
+        if memory_video is not None:
+            memory_video_flat = self._flatten_video_latent(memory_video)
+            memory_video_positions = self._compute_video_positions(
+                memory_video, downscale_factor=memory_downscale_factor
+            )
+            memory_video_timesteps = torch.zeros(
+                B,
+                memory_video_flat.shape[1],
+                1,
+                device=device,
+                dtype=target_video_timesteps.dtype,
+            )
+            video_flat = torch.cat([memory_video_flat, target_video_flat], dim=1)
+            video_positions = torch.cat([memory_video_positions, target_video_positions], dim=2)
+            video_timesteps = torch.cat([memory_video_timesteps, target_video_timesteps], dim=1)
+            memory_seq_len = memory_video_flat.shape[1]
+        else:
+            video_flat = target_video_flat
+            video_positions = target_video_positions
+            video_timesteps = target_video_timesteps
+
+        # Build video modality
+        video_sigma = timestep if timestep.dim() == 1 else timestep[:, 0]
+        video_modality = Modality(
+            latent=video_flat,
+            sigma=video_sigma,
+            timesteps=video_timesteps,
+            positions=video_positions,
+            context=conditional_dict["video_context"],
+            context_mask=conditional_dict.get("attention_mask"),
+            enabled=True,
+        )
+
+        # Build audio modality if provided
+        audio_modality = None
+        memory_audio_seq_len = 0
+        if noisy_audio is None and (memory_audio is not None or memory_audio_timestep is not None):
+            raise ValueError("memory_audio requires noisy_audio")
+        if noisy_audio is not None:
+            target_audio = noisy_audio
+            target_audio_frames = target_audio.shape[1]
+
+            # Use provided audio timestep or derive from video timestep
+            if audio_timestep is None:
+                # In bidirectional mode, audio uses same sigma as video.
+                # video timestep could be [B] or [B, F_v]. For audio we need [B]
+                # or [B, F_a]. If timestep is [B, F_v] (per-frame video), take the
+                # first frame's sigma since bidirectional uses uniform sigma anyway.
+                if timestep.dim() == 1:
+                    audio_timestep = timestep  # [B]
+                else:
+                    # All video frames have same sigma in bidirectional mode,
+                    # take the first frame's value and broadcast to audio frames
+                    audio_timestep = timestep[:, 0]  # [B]
+
+            if audio_timestep.dim() == 1:
+                target_audio_timestep = audio_timestep[:, None].expand(B, target_audio_frames)
+            elif audio_timestep.shape == (B, target_audio_frames):
+                target_audio_timestep = audio_timestep
+            else:
+                raise ValueError(
+                    "audio_timestep must have shape [B] or [B, F_a], "
+                    f"got {tuple(audio_timestep.shape)} vs {(B, target_audio_frames)}"
+                )
+
+            if memory_audio_timestep is not None and memory_audio is None:
+                raise ValueError("memory_audio_timestep requires memory_audio")
+
+            if memory_audio is not None:
+                memory_audio = memory_audio.to(device=device, dtype=target_audio.dtype)
+                memory_audio_seq_len = memory_audio.shape[1]
+                if memory_audio_timestep is None:
+                    prefix_audio_timestep = torch.zeros(
+                        B,
+                        memory_audio_seq_len,
+                        device=device,
+                        dtype=target_audio_timestep.dtype,
+                    )
+                elif memory_audio_timestep.dim() == 1:
+                    prefix_audio_timestep = memory_audio_timestep[:, None].expand(B, memory_audio_seq_len)
+                elif memory_audio_timestep.shape == (B, memory_audio_seq_len):
+                    prefix_audio_timestep = memory_audio_timestep
+                else:
+                    raise ValueError(
+                        "memory_audio_timestep must have shape [B] or [B, F_mem_a], "
+                        f"got {tuple(memory_audio_timestep.shape)} vs {(B, memory_audio_seq_len)}"
+                    )
+
+                noisy_audio = torch.cat([memory_audio, target_audio], dim=1)
+                combined_audio_timestep = torch.cat([prefix_audio_timestep, target_audio_timestep], dim=1)
+            else:
+                noisy_audio = target_audio
+                combined_audio_timestep = target_audio_timestep
+
+            num_audio_tokens = noisy_audio.shape[1]
+            audio_timesteps = self._compute_timesteps_for_tokens(combined_audio_timestep, num_audio_tokens, 1)
+            if memory_audio_seq_len > 0:
+                memory_audio_positions = self._compute_audio_positions(memory_audio)
+                target_audio_position_start = memory_audio_seq_len if memory_position_mode == "prefix_continuous" else 0
+                target_audio_positions = self._compute_audio_positions(
+                    target_audio,
+                    start_frame=target_audio_position_start,
+                )
+                audio_positions = torch.cat([memory_audio_positions, target_audio_positions], dim=2)
+            else:
+                audio_positions = self._compute_audio_positions(noisy_audio)
+            audio_sigma = target_audio_timestep[:, 0]
+            audio_modality = Modality(
+                latent=noisy_audio,
+                sigma=audio_sigma,
+                timesteps=audio_timesteps,
+                positions=audio_positions,
+                context=conditional_dict.get("audio_context", conditional_dict["video_context"]),
+                context_mask=conditional_dict.get("attention_mask"),
+                enabled=True,
+                v2a_grad_scale=float(v2a_grad_scale),
+            )
+
+        if bool(paired_audio_memory) and memory_seq_len > 0 and audio_modality is not None and memory_audio_seq_len > 0:
+            num_memory_slots = int(memory_video.shape[1]) if memory_video is not None else 0
+            if num_memory_slots > 0:
+                target_audio_seq_len = int(audio_modality.latent.shape[1] - memory_audio_seq_len)
+                a2v_pairwise_mask = self._build_paired_memory_cross_mask(
+                    batch_size=B,
+                    query_memory_seq_len=memory_seq_len,
+                    query_target_seq_len=num_target_video_tokens,
+                    kv_memory_seq_len=memory_audio_seq_len,
+                    kv_target_seq_len=target_audio_seq_len,
+                    num_memory_slots=num_memory_slots,
+                    device=device,
+                    kv_segment_lengths=memory_audio_segment_lengths,
+                )
+                v2a_pairwise_mask = self._build_paired_memory_cross_mask(
+                    batch_size=B,
+                    query_memory_seq_len=memory_audio_seq_len,
+                    query_target_seq_len=target_audio_seq_len,
+                    kv_memory_seq_len=memory_seq_len,
+                    kv_target_seq_len=num_target_video_tokens,
+                    num_memory_slots=num_memory_slots,
+                    device=device,
+                    query_segment_lengths=memory_audio_segment_lengths,
+                )
+                video_cross_query_mask = torch.ones(
+                    B,
+                    video_modality.latent.shape[1],
+                    device=device,
+                    dtype=torch.bool,
+                )
+                audio_cross_query_mask = torch.ones(
+                    B,
+                    audio_modality.latent.shape[1],
+                    device=device,
+                    dtype=torch.bool,
+                )
+                audio_attention_mask = self._build_memory_self_attention_block_mask(
+                    batch_size=B,
+                    memory_seq_len=memory_audio_seq_len,
+                    target_seq_len=target_audio_seq_len,
+                    device=device,
+                )
+                video_modality = replace(
+                    video_modality,
+                    cross_kv_mask=v2a_pairwise_mask,
+                    cross_query_mask=video_cross_query_mask,
+                    late_cross_kv_mask=v2a_pairwise_mask,
+                    late_cross_query_mask=video_cross_query_mask,
+                )
+                audio_modality = replace(
+                    audio_modality,
+                    attention_mask=audio_attention_mask,
+                    cross_kv_mask=a2v_pairwise_mask,
+                    cross_query_mask=audio_cross_query_mask,
+                    late_cross_kv_mask=a2v_pairwise_mask,
+                    late_cross_query_mask=audio_cross_query_mask,
+                )
+
+        # Forward through model. The optional perturbation flags let inference
+        # freeze one direction of cross-modal interaction without modifying the
+        # shared core transformer implementation.
+        perturbation_items: list[Perturbation] = []
+        if skip_a2v_cross_attn:
+            perturbation_items.append(
+                Perturbation(
+                    type=PerturbationType.SKIP_A2V_CROSS_ATTN,
+                    blocks=None,
+                )
+            )
+        if skip_v2a_cross_attn:
+            perturbation_items.append(
+                Perturbation(
+                    type=PerturbationType.SKIP_V2A_CROSS_ATTN,
+                    blocks=None,
+                )
+            )
+        if skip_video_self_attn:
+            perturbation_items.append(
+                Perturbation(
+                    type=PerturbationType.SKIP_VIDEO_SELF_ATTN,
+                    blocks=None,
+                )
+            )
+        if skip_audio_self_attn:
+            perturbation_items.append(
+                Perturbation(
+                    type=PerturbationType.SKIP_AUDIO_SELF_ATTN,
+                    blocks=None,
+                )
+            )
+
+        if perturbation_items:
+            perturbation_config = PerturbationConfig(perturbations=perturbation_items)
+            perturbations = BatchedPerturbationConfig(
+                [perturbation_config for _ in range(B)]
+            )
+        else:
+            perturbations = BatchedPerturbationConfig.empty(batch_size=B)
+
+        # The model returns x0 predictions (X0Model wraps velocity model)
+        video_x0, audio_x0 = self.model(
+            video=video_modality,
+            audio=audio_modality,
+            perturbations=perturbations,
+        )
+
+        # Unflatten video output: [B, T, C] -> [B, F, C, H, W]
+        if video_x0 is not None:
+            if memory_seq_len > 0:
+                video_x0 = video_x0[:, memory_seq_len:, :]
+            video_x0 = self._unflatten_video_latent(video_x0, num_video_frames)
+        if audio_x0 is not None and memory_audio_seq_len > 0:
+            audio_x0 = audio_x0[:, memory_audio_seq_len:, :]
+
+        return video_x0, audio_x0
+
+    def load_state_dict(self, state_dict: Dict[str, Any], strict: bool = True) -> None:
+        """Load state dict, handling potential key mismatches."""
+        # Remove 'model.' prefix if present
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            if k.startswith("model."):
+                new_state_dict[k] = v
+            else:
+                new_state_dict[f"model.{k}"] = v
+
+        super().load_state_dict(new_state_dict, strict=strict)
+
+
+def create_ltx2_wrapper(
+    checkpoint_path: str,
+    gemma_path: str,
+    device: torch.device,
+    dtype: torch.dtype = torch.bfloat16,
+    video_height: int = 512,
+    video_width: int = 768,
+    loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
+    registry: Registry | None = None,
+    quantization=None,
+) -> LTX2DiffusionWrapper:
+    """
+    Factory function to create LTX2DiffusionWrapper from checkpoint.
+
+    Args:
+        checkpoint_path: Path to LTX-2 checkpoint
+        gemma_path: Path to Gemma text encoder
+        device: Target device
+        dtype: Model dtype
+        video_height: Video height
+        video_width: Video width
+
+    Returns:
+        Configured LTX2DiffusionWrapper
+    """
+    from ltx_pipelines.utils.model_ledger import ModelLedger
+
+    # IMPORTANT: Load to CPU first, then move to target device
+    # safetensors doesn't support device indices like "cuda:4"
+    # It only accepts "cuda" or "cpu"
+    ledger = ModelLedger(
+        dtype=dtype,
+        device=torch.device("cpu"),  # Load to CPU first
+        checkpoint_path=checkpoint_path,
+        gemma_root_path=gemma_path,
+        loras=loras,
+        registry=registry,
+        quantization=quantization,
+    )
+
+    # Get X0Model (wraps velocity model)
+    x0_model = ledger.transformer()
+
+    # Move to target device. A dtype cast here would silently upcast
+    # quantized fp8 weights back to bf16, so only move the device when a
+    # quantization policy is active.
+    if quantization is None:
+        x0_model = x0_model.to(device=device, dtype=dtype)
+    else:
+        x0_model = x0_model.to(device=device)
+
+    wrapper = LTX2DiffusionWrapper(
+        model=x0_model,
+        video_height=video_height,
+        video_width=video_width,
+    )
+
+    return wrapper

libs/ltx_distillation/utils.py

@@ -0,0 +1,261 @@
+"""Shared utilities for inference: latent computation, noise, media I/O, video concat."""
+
+from __future__ import annotations
+
+import subprocess
+import tempfile
+from pathlib import Path
+from typing import Any, Optional
+
+import torch
+import torchaudio
+try:
+    from torchvision.io import write_video
+except ImportError:  # torchvision >=0.23 removed video I/O (moved to torchcodec).
+    # Only needed when saving a benchmark mp4 to disk; the ComfyUI staged node
+    # consumes the returned tensors directly. (Rebels local patch)
+    write_video = None
+from torchvision.transforms import functional as TVF
+
+from ltx_distillation.inference.memory_multishot import (
+    audio_waveform_stats,
+    normalize_audio_waveform_for_media,
+)
+
+
+def compute_latent_shapes(
+    *,
+    num_frames: int,
+    video_height: int,
+    video_width: int,
+    batch_size: int = 1,
+    latent_channels: int = 128,
+    vae_temporal_compression: int = 8,
+    vae_spatial_compression: int = 32,
+    video_fps: float = 24.0,
+    audio_sample_rate: int = 16000,
+    audio_hop_length: int = 160,
+    audio_latent_downsample: int = 4,
+) -> tuple[list[int], list[int]]:
+    if (num_frames - 1) % vae_temporal_compression != 0:
+        raise ValueError(f"num_frames must be 1 + 8*k, got {num_frames}")
+
+    latent_frames = 1 + (num_frames - 1) // vae_temporal_compression
+    latent_h = video_height // vae_spatial_compression
+    latent_w = video_width // vae_spatial_compression
+
+    video_duration = float(num_frames) / float(video_fps)
+    audio_latent_fps = float(audio_sample_rate) / float(audio_hop_length) / float(audio_latent_downsample)
+    audio_frames = round(video_duration * audio_latent_fps)
+
+    return (
+        [batch_size, latent_frames, latent_channels, latent_h, latent_w],
+        [batch_size, audio_frames, latent_channels],
+    )
+
+
+def add_noise(original: torch.Tensor, noise: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
+    sigma = sigma.to(device=original.device, dtype=original.dtype)
+    if sigma.dim() == 1:
+        sigma = sigma.reshape(-1, *[1] * (original.dim() - 1))
+    elif sigma.dim() == 2:
+        sigma = sigma.reshape(*sigma.shape, *[1] * (original.dim() - 2))
+    return (1 - sigma) * original + sigma * noise
+
+
+def frames_to_video_tensor(frames, target_h: int, target_w: int) -> torch.Tensor:
+    tensors = []
+    for idx, image in enumerate(frames):
+        if image.size != (target_w, target_h):
+            raise ValueError(
+                f"Frame size mismatch at index {idx}: got={image.size}, expected={(target_w, target_h)}"
+            )
+        tensor = TVF.to_tensor(image)
+        tensors.append(tensor * 2.0 - 1.0)
+    return torch.stack(tensors, dim=1).contiguous()
+
+
+@torch.no_grad()
+def encode_memory_frames_batch(
+    *,
+    video_vae,
+    batch_memory_frames,
+    target_h: int,
+    target_w: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    if getattr(video_vae, "encoder", None) is None:
+        raise RuntimeError("video VAE encoder is not initialized for memory encoding")
+
+    latents = []
+    for memory_frames in batch_memory_frames:
+        if not memory_frames:
+            raise ValueError("memory_frames cannot be empty when encoding memory video")
+        per_frame_latents = []
+        for memory_item in memory_frames:
+            is_clip_memory = isinstance(memory_item, list)
+            frame_video = frames_to_video_tensor(
+                memory_item if is_clip_memory else [memory_item],
+                target_h,
+                target_w,
+            ).unsqueeze(0).to(device=device, dtype=dtype)
+            latent = video_vae.encode(frame_video)
+            del frame_video
+            latent = latent.permute(0, 2, 1, 3, 4).to(dtype=dtype)
+            if is_clip_memory:
+                latent = latent[:, -1:, :, :, :].contiguous()
+            per_frame_latents.append(latent)
+        latents.append(torch.cat(per_frame_latents, dim=1))
+        del per_frame_latents
+    return torch.cat(latents, dim=0)
+
+
+@torch.no_grad()
+def _decode_video_tiled_uint8(video_vae, video_latent, tiling_config) -> torch.Tensor:
+    """Decode a video latent in temporal/spatial tiles via the official
+    tiled_decode path, streaming each chunk to CPU so peak GPU memory stays
+    bounded by one tile instead of the full clip."""
+    from ltx_core.model.video_vae.video_vae import decode_video
+
+    decoder = video_vae.decoder
+    latent = video_latent
+    # Same layout fix as VideoVAEWrapper.decode: DMD stores [B, F, C, H, W].
+    if latent.dim() == 5 and latent.shape[2] == 128:
+        latent = latent.permute(0, 2, 1, 3, 4)
+    dec_param = next(decoder.parameters())
+    latent = latent.to(device=dec_param.device, dtype=dec_param.dtype)
+
+    chunks = []
+    for chunk in decode_video(latent, decoder, tiling_config):  # yields [f, h, w, c] uint8
+        chunks.append(chunk.cpu())
+        del chunk
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    return torch.cat(chunks, dim=0).contiguous()
+
+
+@torch.no_grad()
+def decode_benchmark_sample(video_vae, audio_vae, video_latent, audio_latent,
+                            video_tiling_config=None):
+    if video_tiling_config is not None:
+        video_uint8 = _decode_video_tiled_uint8(video_vae, video_latent, video_tiling_config)
+    else:
+        video_pixel = video_vae.decode_to_pixel(video_latent)
+        video_uint8 = video_pixel[0]
+        if video_uint8.shape[0] == 3:
+            video_uint8 = video_uint8.permute(1, 0, 2, 3)
+        video_uint8 = video_uint8.permute(0, 2, 3, 1)
+        video_uint8 = (video_uint8.clamp(0, 1) * 255).cpu().to(torch.uint8).contiguous()
+
+    audio_waveform = audio_vae.decode_to_waveform(audio_latent) if audio_latent is not None else None
+    audio_float = normalize_audio_waveform_for_media(audio_waveform)
+    return video_uint8, audio_float
+
+
+def write_benchmark_media(
+    *,
+    output_path: Path,
+    video_uint8: torch.Tensor,
+    audio_waveform: Optional[torch.Tensor],
+    fps: int,
+    audio_sr: int,
+) -> dict[str, Any]:
+    if write_video is None:
+        raise RuntimeError(
+            "torchvision.io.write_video is unavailable (removed in torchvision >=0.23); "
+            "write_benchmark_media cannot save an mp4 on this install. The ComfyUI staged "
+            "pipeline does not need this function - only standalone benchmark saving does."
+        )
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    audio_waveform = normalize_audio_waveform_for_media(audio_waveform)
+    stats = audio_waveform_stats(audio_waveform)
+
+    wrote_with_audio = False
+    wrote_sidecar_wav = False
+    if audio_waveform is not None:
+        try:
+            write_video(
+                str(output_path),
+                video_uint8,
+                fps=fps,
+                audio_array=audio_waveform,
+                audio_fps=audio_sr,
+                audio_codec="aac",
+            )
+            wrote_with_audio = True
+        except Exception as exc:
+            print(f"[warn] write_video with audio failed for {output_path}: {exc}; audio_stats={stats}", flush=True)
+
+    if not wrote_with_audio:
+        write_video(str(output_path), video_uint8, fps=fps)
+        if audio_waveform is not None:
+            try:
+                torchaudio.save(str(output_path.with_suffix(".wav")), audio_waveform, audio_sr)
+                wrote_sidecar_wav = True
+            except Exception as exc:
+                print(f"[warn] torchaudio.save failed for {output_path}: {exc}; audio_stats={stats}", flush=True)
+
+    return {
+        "wrote_audio_in_mp4": wrote_with_audio,
+        "wrote_sidecar_wav": wrote_sidecar_wav,
+        "audio_stats": stats,
+    }
+
+
+def save_memory_bank_frames(memory_frames: list[Any], save_dir: Path) -> None:
+    save_dir.mkdir(parents=True, exist_ok=True)
+    for old_file in save_dir.glob("*.jpg"):
+        old_file.unlink()
+    for idx, frame in enumerate(memory_frames):
+        if isinstance(frame, list):
+            frame = frame[len(frame) // 2]
+        frame.convert("RGB").save(save_dir / f"memory_{idx:03d}.jpg")
+
+
+def concat_shot_videos(shot_paths: list[Path], output_path: Path) -> None:
+    if not shot_paths:
+        raise ValueError("No shot videos provided for concatenation")
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False, encoding="utf-8") as fp:
+        concat_file = Path(fp.name)
+        for shot_path in shot_paths:
+            fp.write(f"file '{shot_path.resolve().as_posix()}'\n")
+
+    try:
+        cmd = [
+            "ffmpeg", "-y", "-f", "concat", "-safe", "0",
+            "-i", str(concat_file), "-c", "copy", str(output_path),
+        ]
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        if result.returncode != 0:
+            fallback_cmd = [
+                "ffmpeg", "-y", "-f", "concat", "-safe", "0",
+                "-i", str(concat_file),
+                "-c:v", "libx264", "-preset", "medium", "-crf", "18",
+                "-c:a", "aac", "-b:a", "192k",
+                str(output_path),
+            ]
+            fallback_result = subprocess.run(fallback_cmd, capture_output=True, text=True)
+            if fallback_result.returncode != 0:
+                raise RuntimeError(
+                    "Failed to concatenate shot videos with ffmpeg.\n"
+                    f"copy stderr:\n{result.stderr}\n"
+                    f"reencode stderr:\n{fallback_result.stderr}"
+                )
+    finally:
+        concat_file.unlink(missing_ok=True)
+
+
+def concat_shot_audios(audios: list[torch.Tensor]) -> Optional[torch.Tensor]:
+    if not audios:
+        return None
+    audio = audios[0]
+    if audio.ndim == 1:
+        sample_dim = 0
+    elif audio.ndim == 2:
+        sample_dim = 1 if audio.shape[0] <= audio.shape[1] else 0
+    else:
+        raise ValueError(f"Expected audio tensor with 1 or 2 dims, got shape={tuple(audio.shape)}")
+    return torch.cat([a.contiguous() for a in audios], dim=sample_dim).contiguous()

nodes.py

@@ -0,0 +1,1757 @@
+"""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(),)

prompts/long_story_writer_system_prompt.md

@@ -0,0 +1,63 @@
+You are a professional shot-prompt writer for a joint audio-video generation model. Given a user's story (a premise, theme, or outline), expand it into an ordered sequence of shot prompts — ALL shots in one response. Each shot is one ~10-second clip the model renders with synchronized video and audio.
+
+## STRICT OUTPUT FORMAT (MUST FOLLOW EXACTLY)
+- Output MUST be a single valid JSON object and NOTHING else:
+  {"prompts": ["<shot 1 prompt>", "<shot 2 prompt>", ...]}
+- No text before or after the JSON. No explanations, no comments, no markdown code fences (```), no trailing commas.
+- "prompts" is a JSON array of STRINGS. Each string = exactly one shot.
+- Each string is ONE single continuous English paragraph. Inside a string there must be NO field names, NO keys, NO labels, NO bullet points, and NO line breaks (no "\n") — merge everything into one flowing paragraph.
+- The spoken line, when present, is embedded inside the paragraph with escaped double quotes: ... ID_A says, \"...\" ...
+- The number of array elements equals the number of shots. Everything is written in English.
+
+## SPEECH IS OPTIONAL PER SHOT (IMPORTANT)
+- Not every shot has spoken dialogue. Decide per shot whether characters speak, and how many.
+- Use NON-SPEAKING shots for establishing, mood, reaction, object-detail, or transition beats — this varies the rhythm and strengthens the dramatic arc.
+- A shot may have one speaker, two speakers exchanging lines, or no speaker at all.
+- Only when a character speaks do you add that character's voice sentence, a lip-sync note, and the spoken line. For a non-speaking shot, omit all three and let the action and environmental sound carry it.
+
+## CHARACTER CONSISTENCY (CRITICAL — DO NOT LET IDENTITY DRIFT)
+- Give each recurring character a stable ID (ID_A, ID_B, ...). Reserve IDs for PEOPLE only; never label an object with an ID.
+- For each recurring visible character, repeat the EXACT SAME base identity sentence and clothing sentence in every shot where the character appears (and, when the character speaks, the EXACT SAME voice sentence). Copy these sentences verbatim — do not paraphrase, reorder, or change a single word between shots.
+- The base identity sentence describes only stable appearance (age, build, hair, face). It must NOT contain expression or mood.
+- Expression, gaze, posture, and emotional state may vary ONLY AFTER the base identity sentence, written as a separate sentence. The fixed identity/clothing/voice wording itself never changes, so each generated person stays identical across the whole story.
+
+## CHARACTER AGE (STATE IT — THE MODEL DEFAULTS TO ADULTS)
+- Put each character's age band in the base identity sentence; the model renders adults by default, so youth must be asserted or a teen comes out looking like a grown adult. For a teenager, say it plainly (e.g., "a fifteen-year-old teenage girl with a soft youthful round face, clearly a teenager, not an adult") and dress them the way the brief describes — contemporary teen clothing reads as teen and helps age fidelity, while blazers, suit jackets, and professional outerwear read as adult. For a grown character, label them "adult". Never make anyone look younger than a teenager.
+
+## KEEP SIMILAR CHARACTERS DISTINCT (PREVENT IDENTITY MERGE)
+- When two or more visible characters could look alike (same age range, same hair color, both wearing glasses, etc.), give each a BOLD, unmistakable distinguishing feature — a strong hair-color or hairstyle contrast, a distinctive accessory, or a clear facial mark — and restate it in every shot. State plainly that they are different people who must never blend, merge, average, or swap faces, hair, glasses, or clothing. Without a bold differentiator the model averages similar-looking people into a single hybrid face.
+
+## WHAT EACH SHOT PARAGRAPH CONTAINS (woven as natural prose, in this order)
+ALWAYS, for every visible character:
+1. The character's fixed base identity sentence (verbatim) + fixed clothing sentence (verbatim), then optionally one separate sentence for the current expression/gaze/posture/emotion.
+Then:
+2. Action: begin with "At normal speed, " then the action in temporal order.
+3. Style: visual aesthetic, palette, mood, realistic film look.
+4. Camera: framing and motion (keep speaking faces readable).
+5. Background: setting/location and lighting.
+6. Sound effects: the diegetic environmental sounds that are audible.
+7. Sound only, never music: Every shot's audio is strictly quiet, realistic diegetic sound (room tone, ambience, HVAC hum, footsteps, fabric, breathing). There is NEVER any background music, score, soundtrack, melody, humming, or singing in any shot. Do NOT add music of any kind, and do NOT write the words music, score, soundtrack, melody, BGM, humming, or singing anywhere in your output — simply describe the quiet diegetic sounds in a few plain words and stop.
+FOR EACH CHARACTER WHO SPEAKS IN THE SHOT, also add:
+- that character's fixed voice sentence (verbatim): "ID_X's voice is a ... [register, tone, pacing].";
+- a lip-sync note: the mouth movement is clearly visible in frame and stays naturally synchronized with the spoken line (reads well on slower, emotional phrases); for two speakers, state that both mouths stay synced to their own lines;
+- inside the action, reaffirm that the lip movement aligns closely with the audio;
+- the line itself: In a [voice description], ID_X says, \"<the spoken line>\". For two speakers, order the lines naturally (ID_A speaks, then ID_B answers).
+
+## DRAMATIC ARC
+- Build a clear emotional arc: an opening that sets mood, a rising line of realization or tension, a turning point or discovery, a vulnerable low beat, and a resolution. Let the meaning escalate shot by shot. Alternate speaking and non-speaking shots to control pacing and tension.
+
+## DIALOGUE (FOR SPEAKING SHOTS ONLY)
+- Each spoken line is short, roughly 10–20 words, natural and reflective, in the character's own voice; each line pushes the emotional arc forward. In a two-speaker shot keep it to one short line each. English only.
+
+## MODEL-FRIENDLY (AVOID GENERATION FAILURE)
+- Favor gentle, simple, physically plausible actions (standing, sitting, slow turning, walking slowly, reaching, holding, small gestures, speaking to camera). Avoid fast/complex motion (running, fighting, collisions, acrobatics, flying) — the model distorts or collapses.
+- Limit how many characters appear together (two is usually the safe maximum in one shot); keep each shot one clear scene with no mid-shot location jumps. Keep the world realistic; avoid on-screen text, UI, or subtitles.
+
+## NUMBER OF SHOTS
+- Produce exactly the number of shots the user requests. If the user does NOT specify a number, default to exactly 15 shots.
+
+## EXAMPLE OF THE EXACT OUTPUT (two speaking shots and one non-speaking shot; note ID_A's base identity, clothing, and voice sentences are byte-identical across all shots — only the expression sentence and the action change)
+{"prompts": ["ID_A is a young woman in her twenties with shoulder-length dark brown hair and a slim build. ID_A wears a loose light beige knit top and relaxed dark trousers. ID_A's voice is a clear young female voice with a soft mid-high register, gentle breathiness, and intimate vlog-style pacing. Her expression is calm and thoughtful. The mouth movement is clearly visible in the frame and stays naturally synchronized with the spoken line, especially on slower reflective phrases. At normal speed, ID_A steps into the center of the frame, settles her posture, and begins speaking, the lip movement aligning closely with the audio throughout the sentence. In a soft young female voice with reflective warmth, ID_A says, \"I did not plan to record tonight, but this room feels different now.\" The shot uses realistic indoor imagery with soft practical light, neutral warm tones, and a calm introspective domestic mood. A stable medium shot frames ID_A from the waist up, keeping the face and visible lip movement clearly readable while preserving some of the room behind. The background includes a white curtain, soft string lights, part of a small table, and the warm interior of a well-kept room. Very soft indoor room tone, light fabric movement, and subtle foot placement are audible; the sound stays on speech and subtle room ambience only.", "ID_A is a young woman in her twenties with shoulder-length dark brown hair and a slim build. ID_A wears a loose light beige knit top and relaxed dark trousers. ID_A's voice is a clear young female voice with a soft mid-high register, gentle breathiness, and intimate vlog-style pacing. Her expression is quiet and sincere. The mouth movement stays clearly visible and naturally synchronized with the spoken line through the slower delivery. At normal speed, ID_A rests one hand on the notebook on the desk and lets it stay there for a beat before speaking, the lip movement aligning closely with the audio. In a soft young female voice with quiet sincerity, ID_A says, \"This notebook has waited here for months, like a version of me waiting to be answered.\" The shot stays realistic and tactile, with warm desk light and close domestic detail that makes ordinary objects feel emotionally loaded. A close-medium shot keeps one hand on the notebook and part of the face in frame so speech and object interaction stay connected. The background includes the desk surface, the notebook, a soft lamp glow, and the blurred warm curtain lights. Soft contact with the notebook cover, a slight paper shift, and low room ambience are audible; the soundscape stays minimal and speech-centered.", "ID_A is a young woman in her twenties with shoulder-length dark brown hair and a slim build. ID_A wears a loose light beige knit top and relaxed dark trousers. Her gaze is quiet and introspective. At normal speed, ID_A reaches up, grips the curtain gently, and draws it partway, her head turning slightly as the evening settles outside the window. The shot stays naturalistic and warm, emphasizing calm everyday movement inside a softly lit room with a transitional, reflective mood. A medium shot with a slight pan keeps the curtain action and ID_A's face readable while revealing more of the side wall. The background includes the curtain, a faint edge of the window, small warm decorative lights, and a partial glimpse of the table nearby. Curtain fabric sliding, light hand contact with cloth, and low indoor ambience are audible; only faint room tone and soft fabric movement fill the quiet, transitional mood."]}
+
+## PROCESS
+- Read the user's story, decide the shot count per the rule above, break it into a coherent, well-paced emotional sequence. Keep each character's base identity, clothing, and voice sentences byte-identical across all their shots; vary only the separate expression sentence. Mix one-speaker, two-speaker, and non-speaking shots. Output ONLY the {"prompts": [...]} JSON in one response.

rebels_loaders.py

@@ -0,0 +1,840 @@
+"""
+rebels_loaders.py — discrete ComfyUI loader nodes for JoyAI-Echo on low VRAM.
+Patched for Single-File Gemma intake and Key Remapping.
+"""
+from __future__ import annotations
+import os, json, gc
+_LOADER_DIR = os.path.dirname(os.path.abspath(__file__))
+_LOADER_CFG = os.path.join(_LOADER_DIR, "configs", "joyai_echo_config.json")
+import dataclasses
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import gguf
+from gguf import GGUFReader, GGMLQuantizationType as QT
+
+from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
+from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
+from ltx_core.loader.primitives import StateDict
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.model.transformer import LTXV_MODEL_COMFY_RENAMING_MAP, LTXModelConfigurator, X0Model
+from ltx_core.model.video_vae import (VAE_DECODER_COMFY_KEYS_FILTER, VAE_ENCODER_COMFY_KEYS_FILTER,
+                                      VideoDecoderConfigurator, VideoEncoderConfigurator)
+from ltx_core.model.audio_vae import (AUDIO_VAE_DECODER_COMFY_KEYS_FILTER, AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
+                                      VOCODER_COMFY_KEYS_FILTER, AudioDecoderConfigurator,
+                                      AudioEncoderConfigurator, VocoderConfigurator)
+from ltx_core.text_encoders.gemma import (EMBEDDINGS_PROCESSOR_KEY_OPS, EmbeddingsProcessorConfigurator)
+from ltx_distillation.models.ltx_wrapper import LTX2DiffusionWrapper
+from ltx_distillation.models.vae_wrapper import VideoVAEWrapper, AudioVAEWrapper
+from ltx_distillation.models.text_encoder_wrapper import GemmaTextEncoderWrapper
+
+CAT = "Rebels/JoyAI-Echo"
+
+# ---------------------------------------------------------------- config
+def _full_config(src: str) -> dict:
+    src = src.strip().strip('"')
+    if src.lower().endswith(".json"):
+        with open(src, "r", encoding="utf-8") as f: return json.load(f)
+    from safetensors import safe_open
+    with safe_open(src, framework="pt") as f:
+        meta = f.metadata() or {}
+    if "config" not in meta:
+        raise ValueError(f"No 'config' in metadata of {src}. Point at the checkpoint or a config.json.")
+    return json.loads(meta["config"])
+
+_GGUF_MARKER = "__rebels_gguf_no_safetensors__"
+
+def _gemma_remap_key(k):
+    """Map raw checkpoint key names onto the wrapper's module tree. Shared by
+    the safetensors load path and the GGUF swap so both match identically."""
+    nk = k.replace("cond_stage_model.", "").replace("text_model.", "").replace("text_encoder.", "")
+    if "embed_tokens" in nk:
+        nk = "model.model.language_model.embed_tokens.weight"
+    elif "layers" in nk:
+        nk = nk.replace("model.layers", "model.model.language_model.layers")
+    elif "norm" in nk and "language_model" not in nk:
+        nk = nk.replace("model.norm", "model.model.language_model.norm")
+    return nk
+
+class _CfgLoader(SafetensorsModelStateDictLoader):
+    """Stock safetensors weight load, but metadata() returns the shared config."""
+    def __init__(self, config: dict, map_gemma=False, *a, **k):
+        super().__init__(*a, **k)
+        self._cfg = config
+        self._map_gemma = map_gemma
+        
+    def metadata(self, path): return self._cfg
+    
+    def load(self, paths, sd_ops=None, device=None):
+        plist = [str(x) for x in (paths if isinstance(paths, (list, tuple)) else [paths])]
+        if any(_GGUF_MARKER in x or x.lower().endswith(".gguf") for x in plist):
+            import types
+            return types.SimpleNamespace(sd={})  # build() only touches .sd
+        sd_obj = super().load(paths, sd_ops, device)
+        
+        if self._map_gemma:
+            new_sd = {}
+            for k, v in sd_obj.sd.items():
+                new_sd[_gemma_remap_key(k)] = v
+            return dataclasses.replace(sd_obj, sd=new_sd)
+        return sd_obj
+
+# ---------------------------------------------------------------- gguf dit
+# Keep GGUFReaders alive for the process lifetime so their memory-mapped data
+# (which the GGUFLinear weights below reference WITHOUT copying) stays valid.
+_OPEN_GGUF_READERS = []
+
+
+def _gguf_entries(path):
+    r = GGUFReader(path)
+    _OPEN_GGUF_READERS.append(r)
+    out = {}
+    for t in r.tensors:
+        out[t.name] = {"data": np.asarray(t.data), "qtype": t.tensor_type,
+                       "shape": tuple(int(d) for d in reversed(t.shape))}
+    return out
+
+def _dequant(entry, dtype):
+    data = np.asarray(entry["data"])
+    q = int(entry["qtype"])
+    # Unquantized tensors (norms etc. are stored F32/F16 in the GGUF): skip the
+    # old `.astype(np.float32)` which COPIED every one of them into a fresh f32
+    # array before making a second bf16 copy. from_numpy on the memmap view is
+    # zero-copy; the single .to(dtype) below is the only allocation.
+    if q in (int(QT.F32), int(QT.F16)):
+        return torch.from_numpy(data).to(dtype)
+    # Quantized tensors: prefer city96's pure-torch kernels (they run fine on
+    # CPU) dequanting STRAIGHT to the target dtype -- no numpy grouped-rows
+    # machinery and no f32 staging copy. This halves peak RAM per tensor vs the
+    # old path, which is what was tipping the Windows commit limit during load.
+    if _CITY_DEQUANT is not None and _GPU_DEQUANT_OK:
+        try:
+            raw = torch.from_numpy(data)
+            t = _CITY_DEQUANT(raw, QT(q), tuple(entry["shape"]), dtype=dtype)
+            return t.to(dtype)
+        except Exception:
+            pass  # fall through to the numpy reference path
+    deq = gguf.quants.dequantize(data, QT(q))
+    t = torch.from_numpy(deq).to(dtype)
+    del deq
+    return t
+    # NOTE: we deliberately do NOT delete entry["data"] anymore. It is a view
+    # into the memory-mapped file (costs no resident RAM), and the post-build
+    # meta-materialization sweep needs entries to stay readable.
+
+# --- optional GPU dequant kernels, borrowed from city96's ComfyUI-GGUF -------
+# city96 dequantizes packed GGUF weights with pure-torch kernels that run ON THE
+# GPU. That is the single biggest reason his loader is fast and ours was slow:
+# our old path dequantized every weight on the CPU through numpy on EVERY
+# forward. If the user has ComfyUI-GGUF installed (Noah does), import its
+# dequant module and use it; otherwise fall back to the numpy path.
+_CITY_DEQUANT = None
+_GPU_DEQUANT_OK = True  # legacy flag (kept for _dequant)
+_GPU_DEQUANT_BAD = set()  # qtypes whose GPU kernels failed; per-type, never global
+
+
+def _patch_gemma3_rope_compat():
+    """transformers >=~4.56 moved per-layer RoPE attrs (rope_local_base_freq)
+    into a rope_parameters dict; the LTX/JoyAI libs read the old attribute
+    directly. Install a __getattr__ fallback on Gemma3TextConfig that derives
+    the value from rope_parameters or returns the Gemma-3 default (10000.0),
+    chaining to any pre-existing __getattr__. No-op on transformers versions
+    that still have the attribute."""
+    try:
+        from transformers.models.gemma3.configuration_gemma3 import Gemma3TextConfig
+    except Exception:
+        return
+    if getattr(Gemma3TextConfig, "_rebels_rope_compat", False):
+        return
+    _orig_ga = Gemma3TextConfig.__getattr__ if "__getattr__" in vars(Gemma3TextConfig) else None
+    _base_ga = getattr(super(Gemma3TextConfig, Gemma3TextConfig), "__getattr__", None)
+
+    def _ga(self, name):
+        if name == "rope_local_base_freq":
+            rp = self.__dict__.get("rope_parameters")
+            if isinstance(rp, dict):
+                for key in ("sliding_attention", "local_attention", "local"):
+                    sub = rp.get(key)
+                    if isinstance(sub, dict) and sub.get("rope_theta"):
+                        return sub["rope_theta"]
+            return 10000.0
+        if _orig_ga is not None:
+            return _orig_ga(self, name)
+        if _base_ga is not None:
+            return _base_ga(self, name)
+        raise AttributeError(name)
+
+    Gemma3TextConfig.__getattr__ = _ga
+    Gemma3TextConfig._rebels_rope_compat = True
+
+
+_patch_gemma3_rope_compat()
+
+
+def gpu_dequant_supported(qtype_value):
+    """Probe whether city96's GPU kernels can handle this qtype, WITHOUT
+    running a real layer. Probes a tiny zeros tensor on CUDA; failures are
+    remembered in _GPU_DEQUANT_BAD. Used as a pre-flight so GPU encode/denoise
+    never silently grinds on CPU fallback for an unsupported quant type."""
+    q = int(qtype_value)
+    if q in (int(QT.F32), int(QT.F16)):
+        return True
+    if _CITY_DEQUANT is None or not torch.cuda.is_available():
+        return False
+    if q in _GPU_DEQUANT_BAD:
+        return False
+    try:
+        block_size, type_size = gguf.GGML_QUANT_SIZES[QT(q)]
+        data = torch.zeros((1, type_size), dtype=torch.uint8, device="cuda")
+        _CITY_DEQUANT(data, QT(q), (1, block_size), dtype=torch.bfloat16)
+        return True
+    except Exception:
+        _GPU_DEQUANT_BAD.add(q)
+        return False
+try:
+    import importlib.util as _ilu
+    _cn_dir = os.path.dirname(_LOADER_DIR)  # .../custom_nodes
+    for _cand in ("ComfyUI-GGUF", "ComfyUI-GGUF-main", "comfyui-gguf"):
+        _dq = os.path.join(_cn_dir, _cand, "dequant.py")
+        if os.path.isfile(_dq):
+            _spec = _ilu.spec_from_file_location("rebels_city96_dequant", _dq)
+            _mod = _ilu.module_from_spec(_spec); _spec.loader.exec_module(_mod)
+            _CITY_DEQUANT = getattr(_mod, "dequantize", None)
+            if _CITY_DEQUANT:
+                print(f"[Rebels JE] GPU dequant kernels loaded from {_cand}/dequant.py", flush=True)
+            break
+except Exception as _e:
+    print(f"[Rebels JE] city96 dequant unavailable ({_e}); using CPU numpy dequant.", flush=True)
+
+
+class GGUFLinear(nn.Module):
+    def __init__(self, entry, bias=None, compute_dtype=torch.bfloat16):
+        super().__init__()
+        self.qtype_value = int(entry["qtype"]); self.weight_shape = tuple(entry["shape"])
+        # The packed weight is kept as a PLAIN attribute (not a registered buffer):
+        #  - it stays a numpy view into the memory-mapped GGUF file, so it costs
+        #    ~zero resident RAM (the .copy() that duplicated the whole DiT in RAM
+        #    is gone for good);
+        #  - module.to(device) / state_dict / pin_memory all ignore it, so the
+        #    sequential offloader can shuttle blocks to the GPU without dragging
+        #    9GB of packed weights along or pinning them.
+        # Each forward streams just this layer's packed bytes to the GPU and
+        # dequantizes there (city96 kernels) -- the same per-layer streaming that
+        # makes Noah's other big LTX GGUF models run fine on 8GB.
+        self._qweight = entry["data"]
+        self.bias = nn.Parameter(bias.to(compute_dtype), requires_grad=False) if bias is not None else None
+
+    def forward(self, x):
+        q = self.qtype_value
+        # F16 / F32 tensors (llama-quantize leaves some in K-quant files) are
+        # not in city96's kernel table -- they don't need kernels at all.
+        # Previously ONE of these raised a KeyError that tripped a GLOBAL kill
+        # switch, silently dropping every linear in the whole run to CPU numpy
+        # dequant (the 40-minute generations). Handle them directly:
+        if q in (int(QT.F32), int(QT.F16)):
+            w = torch.from_numpy(np.asarray(self._qweight))
+            w = w.reshape(self.weight_shape).to(device=x.device, dtype=x.dtype)
+            return F.linear(x, w, self.bias.to(x.dtype) if self.bias is not None else None)
+        if _CITY_DEQUANT is not None and q not in _GPU_DEQUANT_BAD:
+            try:
+                data = torch.from_numpy(np.asarray(self._qweight)).to(x.device, non_blocking=True)
+                w = _CITY_DEQUANT(data, QT(q), tuple(self.weight_shape), dtype=x.dtype)
+                if tuple(w.shape) != tuple(self.weight_shape):
+                    w = w.reshape(self.weight_shape)
+                w = w.to(dtype=x.dtype)
+                return F.linear(x, w, self.bias.to(x.dtype) if self.bias is not None else None)
+            except Exception as e:
+                # disable the GPU path for THIS qtype only -- never globally
+                _GPU_DEQUANT_BAD.add(q)
+                print(f"[Rebels JE] GPU dequant unavailable for {QT(q).name} "
+                      f"({type(e).__name__}: {e}); that qtype uses CPU dequant.", flush=True)
+        raw = np.asarray(self._qweight)
+        w = torch.from_numpy(gguf.quants.dequantize(raw, QT(self.qtype_value)).astype(np.float32))
+        w = w.reshape(self.weight_shape).to(device=x.device, dtype=x.dtype)
+        return F.linear(x, w, self.bias.to(x.dtype) if self.bias is not None else None)
+
+def _set_sub(root, dotted, new):
+    *parents, leaf = dotted.split("."); p = root
+    for a in parents: p = getattr(p, a)
+    setattr(p, leaf, new)
+
+class _GGUFDiTLoader(SafetensorsModelStateDictLoader):
+    def __init__(self, config, entries, consumed, dtype):
+        super().__init__(); self._cfg = config; self._e = entries; self._consumed = consumed; self._dt = dtype
+    def metadata(self, path): return self._cfg
+    def load(self, paths, sd_ops=None, device=None):
+        sd, size, n = {}, 0, 0
+        for k, e in self._e.items():
+            if k in self._consumed: continue
+            t = _dequant(e, self._dt)
+            # THE META-TENSOR FIX: leftover GGUF keys still carry JD's checkpoint
+            # prefix ("model.diffusion_model."), but the bare LTXModel's module
+            # names do not. With model_sd_ops=None nothing strips it, so
+            # load_state_dict(strict=False) silently skipped every one of these
+            # tensors and they stayed empty 'meta' placeholders -- which is the
+            # "Tensor on device meta" crash in patchify_proj. Emit the stripped
+            # key alongside the raw one (same tensor object, costs nothing).
+            sd[k] = t
+            for p in ("model.diffusion_model.", "diffusion_model."):
+                if k.startswith(p):
+                    sd[k[len(p):]] = t
+                    break
+            size += t.numel() * t.element_size()
+            n += 1
+            if n % 32 == 0:
+                # Periodic (not per-tensor) gc: keeps Windows commit pressure
+                # down on a 16GB machine without 1700 collector passes.
+                gc.collect()
+        gc.collect()
+        return StateDict(sd=sd, device=device or torch.device("cpu"), size=size, dtype=self._dt)
+
+# The bare LTXModel modules are named e.g. "transformer_blocks.0.attn1.to_q",
+# but the GGUF keys keep JD's checkpoint prefix "model.diffusion_model.". The
+# configurator strips that prefix at load time, so we must match across it here
+# or the swap fires on nothing and the whole DiT dequantizes into RAM -> OOM.
+_DIT_PREFIXES = ("", "model.diffusion_model.", "diffusion_model.")
+
+
+def _find_entry(entries, base):
+    for p in _DIT_PREFIXES:
+        k = p + base
+        if k in entries:
+            return k
+    return None
+
+
+# Some LTX components (the transformer-args preprocessors) capture DIRECT
+# OBJECT REFERENCES to modules like patchify_proj at construction time, outside
+# the registered module tree. When the mutator swaps those modules for
+# GGUFLinear, the preprocessor keeps pointing at the ORIGINAL meta nn.Linear --
+# invisible to load_state_dict AND to the meta sweep (which is why the sweep
+# reports 0 while the forward still hits a meta tensor). We record every
+# old->new swap and then re-bind stale references across the whole object graph.
+_SWAP_MAP = {}
+
+def _rebind_swapped(root, max_objs=50000):
+    import types as _types
+    if not _SWAP_MAP:
+        return 0
+    _SKIP = (type, _types.FunctionType, _types.MethodType, _types.BuiltinFunctionType,
+             _types.ModuleType, str, bytes, int, float, bool, complex, torch.Tensor,
+             np.ndarray)
+    seen, queue, fixed = set(), [root], 0
+    while queue and len(seen) < max_objs:
+        obj = queue.pop()
+        oid = id(obj)
+        if oid in seen:
+            continue
+        seen.add(oid)
+        containers = []
+        d = getattr(obj, "__dict__", None)
+        if isinstance(d, dict):
+            containers.append(d)
+        m = getattr(obj, "_modules", None)
+        if isinstance(m, dict) and m is not d:
+            containers.append(m)
+        for cont in containers:
+            for k, v in list(cont.items()):
+                nv = _SWAP_MAP.get(id(v))
+                if nv is not None and v is not nv:
+                    cont[k] = nv
+                    fixed += 1
+                    continue
+                if isinstance(v, (list,)):
+                    for i, item in enumerate(v):
+                        nvi = _SWAP_MAP.get(id(item))
+                        if nvi is not None and item is not nvi:
+                            v[i] = nvi
+                            fixed += 1
+                        elif isinstance(item, nn.Module):
+                            queue.append(item)
+                    continue
+                if isinstance(v, dict):
+                    for dk, item in list(v.items()):
+                        nvi = _SWAP_MAP.get(id(item))
+                        if nvi is not None and item is not nvi:
+                            v[dk] = nvi
+                            fixed += 1
+                        elif isinstance(item, nn.Module):
+                            queue.append(item)
+                    continue
+                if isinstance(v, _SKIP) or v is None:
+                    continue
+                if isinstance(v, nn.Module) or hasattr(v, "__dict__"):
+                    queue.append(v)
+    if fixed:
+        print(f"[Rebels JE] re-bound {fixed} stale reference(s) to swapped GGUF layers.", flush=True)
+    return fixed
+
+
+def _dit_module_ops(entries, consumed, compute_dtype):
+    def mutator(model):
+        n_lin = n_hit = 0
+        miss = []
+        for name, mod in list(model.named_modules()):
+            if not isinstance(mod, nn.Linear):
+                continue
+            n_lin += 1
+            wk = _find_entry(entries, name + ".weight")
+            if wk is None:
+                if len(miss) < 5:
+                    miss.append(name)
+                continue
+            bk = _find_entry(entries, name + ".bias")
+            bias = _dequant(entries[bk], compute_dtype) if bk else None
+            new_mod = GGUFLinear(entries[wk], bias, compute_dtype)
+            _SWAP_MAP[id(mod)] = new_mod
+            _set_sub(model, name, new_mod)
+            consumed.add(wk)
+            if bk:
+                consumed.add(bk)
+            n_hit += 1
+        print(f"[Rebels JE] DiT GGUF swap: matched {n_hit}/{n_lin} Linear layers "
+              f"({len(consumed)} tensors kept packed).", flush=True)
+        if n_hit == 0 and n_lin:
+            print(f"[Rebels JE]   NO matches -> whole DiT would dequantize. "
+                  f"sample model Linears={miss}", flush=True)
+            print(f"[Rebels JE]   sample GGUF keys={list(entries.keys())[:5]}", flush=True)
+        return model
+    return (ModuleOps("gguf_linear_swap", matcher=lambda m: True, mutator=mutator),)
+
+# ---------------------------------------------------------------- gemma fp8
+class Fp8Linear(nn.Module):
+    def __init__(self, qweight_u8, shape, scale, bias=None, compute_dtype=torch.bfloat16):
+        super().__init__(); self.weight_shape = tuple(shape)
+        self.register_buffer("qweight", qweight_u8)
+        self.register_buffer("scale_weight", torch.tensor(float(scale), dtype=torch.float32))
+        self.bias = nn.Parameter(bias.to(compute_dtype), requires_grad=False) if bias is not None else None
+    def forward(self, x):
+        w = self.qweight.view(torch.float8_e4m3fn).reshape(self.weight_shape).to(torch.float32) * self.scale_weight
+        w = w.to(device=x.device, dtype=x.dtype)
+        return F.linear(x, w, self.bias.to(x.dtype) if self.bias is not None else None)
+
+def _swap_gemma_gguf(model, gguf_path, compute_dtype):
+    """Load Gemma from OUR quantized GGUF (made by make_gemma_gguf.py, which
+    preserves HF key names -- so matching is direct, no llama.cpp renaming).
+    Every nn.Linear becomes a packed GGUFLinear backed by the memory-mapped
+    file (~zero resident RAM); embeddings/norms are dequanted to compute dtype.
+    Resident footprint ~2.5GB vs ~12GB for the fp8 path."""
+    raw = _gguf_entries(gguf_path)
+    entries = {}
+    for k, v in raw.items():
+        entries[_gemma_remap_key(k)] = v
+        entries.setdefault(k, v)  # keep originals too; harmless duplicates
+    names = list(entries.keys())
+
+    def find(base):
+        if base in entries:
+            return base
+        cands = [n for n in names if n.endswith("." + base) or base.endswith("." + n)]
+        return cands[0] if len(cands) == 1 else None
+
+    used = set()
+    n_hit = n_lin = 0
+    for name, mod in list(model.named_modules()):
+        if not isinstance(mod, nn.Linear):
+            continue
+        n_lin += 1
+        wk = find(name + ".weight")
+        if wk is None:
+            continue
+        bk = find(name + ".bias")
+        if bk is not None:
+            bias = _dequant(entries[bk], compute_dtype)
+            used.add(bk)
+        elif mod.bias is not None and mod.bias.device.type != "meta":
+            bias = mod.bias.detach()
+        else:
+            bias = None
+        _set_sub(model, name, GGUFLinear(entries[wk], bias, compute_dtype))
+        used.add(wk)
+        n_hit += 1
+    print(f"[Rebels JE] Gemma GGUF swap: matched {n_hit}/{n_lin} Linear layers "
+          f"(packed, memory-mapped).", flush=True)
+    if n_hit == 0 and n_lin:
+        print(f"[Rebels JE]   sample GGUF keys={names[:5]}", flush=True)
+    # Fill everything else (embeddings, norms) from the GGUF. vision_tower /
+    # multi_modal_projector / lm_head have no weights anywhere and get stripped
+    # by the staged node right after this -- skip them instead of aborting.
+    _materialize_meta(model, entries, used, compute_dtype, strict=False,
+                      skip_substrings=("vision_tower", "multi_modal_projector", "lm_head"))
+    return model
+
+
+def _swap_gemma_fp8(model, fp8_dir, compute_dtype):
+    from safetensors import safe_open
+    from pathlib import Path
+    shards = sorted(Path(fp8_dir).glob("model*.safetensors"))
+    scales = {}
+    for sh in shards:
+        with safe_open(str(sh), framework="pt") as f:
+            for k in f.keys():
+                if k.endswith(".scale_weight"): scales[k[:-len(".scale_weight")]] = float(f.get_tensor(k))
+    n = 0
+    for name, mod in list(model.named_modules()):
+        if isinstance(mod, nn.Linear) and name in scales:
+            for sh in shards:
+                with safe_open(str(sh), framework="pt") as f:
+                    if name in f.keys():
+                        qw = f.get_tensor(name); bias = mod.bias.detach() if mod.bias is not None else None
+                        _set_sub(model, name, Fp8Linear(qw.view(torch.uint8), qw.shape, scales[name], bias, compute_dtype))
+                        n += 1; break
+    return model
+
+# ================================================================ NODES
+def _dev(): return torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+class RebelsJE_Config:
+    CATEGORY = CAT; RETURN_TYPES = ("JOYECHO_CONFIG",); RETURN_NAMES = ("config",); FUNCTION = "run"
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {"config_source": ("STRING", {"default": _LOADER_CFG})}}
+    def run(self, config_source): return (_full_config(config_source),)
+
+def _materialize_meta(root, entries, consumed, dtype, strict=True, skip_substrings=()):
+    """Safety net: find every parameter/buffer still on the 'meta' device after
+    build and fill it with real data from the GGUF. Resolves names by trying the
+    known checkpoint prefixes first, then by unique longest-suffix match (covers
+    modules the configurator registers under different paths, e.g. preprocessor
+    wrappers around patchify_proj). Raises a CLEAR error naming any weight it
+    cannot find, instead of letting a cryptic 'Tensor on device meta' crash
+    happen 40 minutes into a run."""
+    leftover = [k for k in entries if k not in consumed]
+    fixed, missing = 0, []
+
+    def resolve(pname):
+        k = _find_entry(entries, pname)
+        if k is not None and k not in consumed:
+            return k
+        parts = pname.split(".")
+        for i in range(len(parts)):
+            suf = ".".join(parts[i:])
+            cands = [n for n in leftover if n == suf or n.endswith("." + suf)]
+            if len(cands) == 1:
+                return cands[0]
+        return None
+
+    items = [(n, p, True) for n, p in root.named_parameters()] \
+          + [(n, b, False) for n, b in root.named_buffers()]
+    skipped = 0
+    for name, t, is_param in items:
+        if t is None or t.device.type != "meta":
+            continue
+        if any(sub in name for sub in skip_substrings):
+            skipped += 1
+            continue
+        k = resolve(name)
+        if k is None:
+            missing.append(name)
+            continue
+        new = _dequant(entries[k], dtype)
+        if new.numel() == t.numel() and tuple(new.shape) != tuple(t.shape):
+            new = new.reshape(t.shape)
+        *path, leaf = name.split(".")
+        mod = root
+        for a in path:
+            mod = getattr(mod, a)
+        if is_param:
+            mod._parameters[leaf] = nn.Parameter(new, requires_grad=False)
+        else:
+            mod._buffers[leaf] = new
+        fixed += 1
+
+    print(f"[Rebels JE] materialized {fixed} meta tensors from GGUF"
+          + (f" ({skipped} skipped by filter)" if skipped else "") + ".", flush=True)
+    if missing and not strict:
+        print(f"[Rebels JE] WARNING: {len(missing)} tensors left unresolved "
+              f"(non-strict): {missing[:6]}", flush=True)
+    if missing and strict:
+        raise RuntimeError(
+            f"[Rebels JE] {len(missing)} model weights are still empty (meta) and "
+            f"could not be located in the GGUF: {missing[:8]}. The GGUF may be "
+            f"missing these tensors -- re-check the quantization export.")
+
+
+class RebelsJE_DiTLoader:
+    CATEGORY = CAT; RETURN_TYPES = ("JOYECHO_GENERATOR",); RETURN_NAMES = ("generator",); FUNCTION = "run"
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {
+            "config": ("JOYECHO_CONFIG",),
+            "dit_gguf": ("STRING", {"default": ""}),
+            "video_height": ("INT", {"default": 736}), "video_width": ("INT", {"default": 1280})}}
+    def run(self, config, dit_gguf, video_height, video_width):
+        dtype = torch.bfloat16
+        _SWAP_MAP.clear()  # ids are only valid for THIS build
+        entries = _gguf_entries(dit_gguf); consumed = set()
+        builder = Builder(
+            model_class_configurator=LTXModelConfigurator,
+            model_path=dit_gguf,
+            model_sd_ops=None,
+            module_ops=_dit_module_ops(entries, consumed, dtype),
+            model_loader=_GGUFDiTLoader(config, entries, consumed, dtype),
+        )
+        transformer = builder.build(device=torch.device("cpu"), dtype=dtype)
+        gen = LTX2DiffusionWrapper(model=X0Model(transformer), video_height=video_height, video_width=video_width)
+        gen.eval()
+        # Sweep the FULL wrapper (not just the transformer) so anything the
+        # configurator or wrapper registered late gets real weights too.
+        _materialize_meta(gen, entries, consumed, dtype)
+        # Fix stale direct references (e.g. args-preprocessor patchify_proj)
+        # that still point at pre-swap meta modules.
+        _rebind_swapped(gen)
+        _SWAP_MAP.clear()
+        return (gen,)
+
+class RebelsJE_TextEncoder:
+    CATEGORY = CAT; RETURN_TYPES = ("JOYECHO_TEXTENC",); RETURN_NAMES = ("text_encoder",); FUNCTION = "run"
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {
+            "config": ("JOYECHO_CONFIG",),
+            "gemma_path": ("STRING", {"default": ""}),
+            "gemma_format": (["our_fp8", "bf16"], {"default": "our_fp8"}),
+            "connector_path": ("STRING", {"default": ""}),
+            "low_vram": ("BOOLEAN", {"default": True})}}
+    def run(self, config, gemma_path, gemma_format, connector_path, low_vram):
+        from ltx_core.text_encoders.gemma import GemmaTextEncoderConfigurator, GEMMA_MODEL_OPS, module_ops_from_gemma_root
+        from ltx_core.utils import find_matching_file
+        from pathlib import Path
+        
+        dtype = torch.bfloat16; dev = torch.device("cpu") if low_vram else _dev()
+        
+        # --- SINGLE FILE INTAKE PATCH START ---
+        gemma_path_str = str(gemma_path)
+        is_gguf = gemma_path_str.lower().endswith(".gguf")
+        if os.path.isfile(gemma_path_str):
+            parent_dir = os.path.dirname(gemma_path_str)
+            temp_folder = os.path.join(parent_dir, ".gemma_virtual_folder")
+            os.makedirs(temp_folder, exist_ok=True)
+            
+            temp_model = os.path.join(temp_folder, "model.safetensors")
+            if os.path.exists(temp_model):
+                try: os.remove(temp_model)
+                except OSError: pass
+            
+            # GGUF gemma: do NOT link the weights as model.safetensors (the
+            # builder would try to parse a GGUF as safetensors). The virtual
+            # folder only carries the HF sidecars; weights come from the GGUF.
+            if not is_gguf:
+                try: os.link(gemma_path_str, temp_model)
+                except OSError:
+                    import shutil
+                    shutil.copyfile(gemma_path_str, temp_model)
+                
+            # Gemma needs its HF sidecar files (tokenizer + config jsons) next to
+            # the weights. A single-file fp8 download has none of them, so we search
+            # several places, in order: the weights' own folder, a 'gemma_assets' or
+            # 'gemma' subfolder beside them, and a 'gemma_assets' folder bundled in
+            # this node pack (so they can ship with the pack).
+            _node_dir = os.path.dirname(os.path.abspath(__file__))
+            sidecar_sources = [
+                parent_dir,
+                os.path.join(parent_dir, "gemma_assets"),
+                os.path.join(parent_dir, "gemma"),
+                os.path.join(_node_dir, "gemma_assets"),
+            ]
+            sidecar_files = ["tokenizer.model", "tokenizer_config.json", "config.json",
+                             "special_tokens_map.json", "preprocessor_config.json"]
+            def _normalize_gemma_config(path):
+                """Load ANY Gemma-3 config.json (our bundled text-only one OR the
+                multimodal one from google/gemma-3-12b-it that users often
+                download) and return the text-only dict this pipeline needs,
+                with every field JD's encoder configurator reads guaranteed
+                present. Returns None if the file isn't a Gemma-3 config at all."""
+                import json as _json
+                try:
+                    with open(path, "r", encoding="utf-8") as fh:
+                        cj = _json.load(fh)
+                except Exception:
+                    return None
+                if "gemma3" not in str(cj.get("model_type", "")) and "text_config" not in cj:
+                    return None
+                # multimodal google layout -> pull the nested text config
+                if "text_config" in cj and isinstance(cj["text_config"], dict):
+                    cj = dict(cj["text_config"])
+                cj["model_type"] = "gemma3_text"
+                cj.setdefault("architectures", ["Gemma3ForCausalLM"])
+                # fields the encoder configurator / rotary init require:
+                rs = cj.get("rope_scaling")
+                if not isinstance(rs, dict):
+                    rs = {"rope_type": "linear", "factor": 8.0}
+                if "rope_type" not in rs:
+                    rs["rope_type"] = rs.get("type", "linear")
+                cj["rope_scaling"] = rs
+                cj.setdefault("rope_local_base_freq", 10000.0)
+                cj.setdefault("rope_theta", 1000000.0)
+                return cj
+
+            for t_file in sidecar_files:
+                dst = os.path.join(temp_folder, t_file)
+                if os.path.exists(dst):
+                    continue
+                for srcdir in sidecar_sources:
+                    src = os.path.join(srcdir, t_file)
+                    if os.path.exists(src):
+                        if t_file == "config.json":
+                            cj = _normalize_gemma_config(src)
+                            if cj is None:
+                                print(f"[Rebels JE] skipping {src}: not a Gemma-3 "
+                                      f"config.", flush=True)
+                                continue
+                            import json as _json
+                            with open(dst, "w", encoding="utf-8") as fh:
+                                _json.dump(cj, fh, indent=2)
+                            if "text_config" not in cj:
+                                pass  # already text-only
+                            print(f"[Rebels JE] config.json normalized from {src} "
+                                  f"(text-only layout, rotary fields ensured).", flush=True)
+                            break
+                        try: os.link(src, dst)
+                        except OSError:
+                            import shutil
+                            shutil.copyfile(src, dst)
+                        break
+
+            # Fail with a clear, actionable message instead of a cryptic one later.
+            missing = [f for f in sidecar_files
+                       if not os.path.exists(os.path.join(temp_folder, f))]
+            if missing:
+                raise FileNotFoundError(
+                    "Gemma sidecar files missing: " + ", ".join(missing) + ".\n"
+                    "Put them in one of:\n"
+                    f"  - the same folder as your Gemma file ({parent_dir})\n"
+                    f"  - {os.path.join(_node_dir, 'gemma_assets')}  (ships with the node pack)\n"
+                    "These ship inside the node pack's gemma_assets folder -- if they're "
+                    "missing, re-download the pack from GitHub rather than pulling configs "
+                    "from the google repo (google's config.json is the multimodal variant "
+                    "and is NOT compatible with this pipeline)."
+                )
+                        
+            model_folder = Path(temp_folder)
+            gemma_op_path = str(temp_folder)
+        else:
+            model_folder = find_matching_file(gemma_path, "model*.safetensors").parent
+            gemma_op_path = gemma_path
+        # --- SINGLE FILE INTAKE PATCH END ---
+
+        weight_paths = (_GGUF_MARKER,) if is_gguf else tuple(str(p) for p in model_folder.rglob("*.safetensors"))
+        te_builder = Builder(
+            model_class_configurator=GemmaTextEncoderConfigurator,
+            model_path=weight_paths,
+            module_ops=(GEMMA_MODEL_OPS, *module_ops_from_gemma_root(gemma_op_path)),
+            model_loader=_CfgLoader(config, map_gemma=True),
+        )
+        text_encoder = te_builder.build(device=dev, dtype=dtype)
+        if is_gguf:
+            _swap_gemma_gguf(text_encoder, gemma_path_str, dtype)
+        elif gemma_format == "our_fp8":
+            _swap_gemma_fp8(text_encoder, str(model_folder), dtype)
+
+        # META-TENSOR FIX: the text-only Gemma file has no vision_tower /
+        # multi_modal_projector / lm_head, so those modules stay on the meta
+        # device and poison model.device (vision_tower is the first registered
+        # param). base_encoder.encode() then builds input_ids/attention_mask on
+        # meta and dies with "Cannot copy out of meta tensor". None are used for
+        # text encoding -- drop them so the first real parameter (language_model)
+        # defines the device. (StagedJE does this after calling .run(); doing it
+        # here fixes the discrete multishot path too. Idempotent: StagedJE's own
+        # strip becomes a no-op.)
+        try:
+            gm = getattr(text_encoder, "model", None)
+            if gm is not None:
+                inner = getattr(gm, "model", None)
+                for parent, attr in ((inner, "vision_tower"),
+                                     (inner, "multi_modal_projector"),
+                                     (gm, "lm_head")):
+                    if parent is not None and getattr(parent, attr, None) is not None:
+                        try: setattr(parent, attr, None)
+                        except Exception: pass
+                try:
+                    print(f"[Rebels JE] text-encoder device after meta-strip = "
+                          f"{next(gm.parameters()).device}", flush=True)
+                except StopIteration:
+                    pass
+        except Exception as e:
+            print(f"[Rebels JE] meta-strip skipped: {e}", flush=True)
+
+
+        ep_builder = Builder(
+            model_class_configurator=EmbeddingsProcessorConfigurator,
+            model_path=connector_path, model_sd_ops=EMBEDDINGS_PROCESSOR_KEY_OPS,
+            model_loader=_CfgLoader(config, map_gemma=False),
+        )
+        embeddings_processor = ep_builder.build(device=dev, dtype=dtype)
+
+        # DEVICE UNIFY: the wrapper's forward feeds Gemma hidden_states straight
+        # into embeddings_processor with no device move (it never uses
+        # self.device). A GGUF-swapped Gemma effectively runs on CPU (packed
+        # layers stream per-layer), so hidden_states land on CPU while the
+        # connector was built on `dev` (cuda when low_vram=False) -> a
+        # cpu-vs-cuda addmm crash in feature_extractor. Pin embeddings_processor
+        # to the text encoder's ACTUAL device so both halves match. (StagedJE
+        # unifies both halves the same way after calling .run().)
+        try:
+            _te_dev = next(text_encoder.model.parameters()).device
+            embeddings_processor.to(_te_dev)
+            print(f"[Rebels JE] embeddings_processor pinned to {_te_dev} "
+                  f"(matches text encoder).", flush=True)
+        except StopIteration:
+            _te_dev = _dev()
+
+        wrapper = GemmaTextEncoderWrapper(text_encoder=text_encoder, embeddings_processor=embeddings_processor,
+                                          device=_te_dev, dtype=dtype)
+        return (wrapper,)
+
+class RebelsJE_VAELoader:
+    CATEGORY = CAT; RETURN_TYPES = ("JOYECHO_VVAE", "JOYECHO_AVAE", "INT"); RETURN_NAMES = ("video_vae", "audio_vae", "audio_sample_rate"); FUNCTION = "run"
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {
+            "config": ("JOYECHO_CONFIG",),
+            "video_vae_path": ("STRING", {"default": ""}),
+            "audio_vae_path": ("STRING", {"default": ""}),
+            "vocoder_path": ("STRING", {"default": ""}),
+            "with_encoders": ("BOOLEAN", {"default": True})}}
+    def _build(self, cfg, configurator, sd_ops, path):
+        return Builder(model_class_configurator=configurator, model_path=path,
+                       model_sd_ops=sd_ops, model_loader=_CfgLoader(cfg)).build(
+                       device=torch.device("cpu"), dtype=torch.bfloat16)
+    def run(self, config, video_vae_path, audio_vae_path, vocoder_path, with_encoders):
+        dtype = torch.bfloat16
+        v_dec = self._build(config, VideoDecoderConfigurator, VAE_DECODER_COMFY_KEYS_FILTER, video_vae_path)
+        a_dec = self._build(config, AudioDecoderConfigurator, AUDIO_VAE_DECODER_COMFY_KEYS_FILTER, audio_vae_path)
+        voc   = self._build(config, VocoderConfigurator, VOCODER_COMFY_KEYS_FILTER, vocoder_path)
+        v_enc = self._build(config, VideoEncoderConfigurator, VAE_ENCODER_COMFY_KEYS_FILTER, video_vae_path) if with_encoders else None
+        a_enc = self._build(config, AudioEncoderConfigurator, AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER, audio_vae_path) if with_encoders else None
+        video_vae = VideoVAEWrapper(encoder=v_enc, decoder=v_dec, device=_dev(), dtype=dtype)
+        audio_vae = AudioVAEWrapper(encoder=a_enc, decoder=a_dec, vocoder=voc, device=_dev(), dtype=dtype)
+        video_vae.eval(); audio_vae.eval()
+        sr = audio_vae.get_output_sample_rate() or 24000
+        return (video_vae, audio_vae, sr)
+
+class RebelsJE_Assemble:
+    CATEGORY = CAT; RETURN_TYPES = ("JOYECHO_MODEL",); RETURN_NAMES = ("model",); FUNCTION = "run"
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {"generator": ("JOYECHO_GENERATOR",), "text_encoder": ("JOYECHO_TEXTENC",),
+                             "video_vae": ("JOYECHO_VVAE",), "audio_vae": ("JOYECHO_AVAE",),
+                             "audio_sample_rate": ("INT", {"default": 24000})}}
+    def run(self, generator, text_encoder, video_vae, audio_vae, audio_sample_rate):
+        model = {"text_encoder": text_encoder, "generator": generator, "video_vae": video_vae,
+                 "audio_vae": audio_vae, "audio_sample_rate": audio_sample_rate,
+                 "device": _dev(), "dtype": torch.bfloat16}
+        return (model,)
+
+NODE_CLASS_MAPPINGS = {
+    "RebelsJE_Config": RebelsJE_Config,
+    "RebelsJE_DiTLoader": RebelsJE_DiTLoader,
+    "RebelsJE_TextEncoder": RebelsJE_TextEncoder,
+    "RebelsJE_VAELoader": RebelsJE_VAELoader,
+    "RebelsJE_Assemble": RebelsJE_Assemble,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "RebelsJE_Config": "Rebels JE • Config",
+    "RebelsJE_DiTLoader": "Rebels JE • DiT GGUF Loader (UNet)",
+    "RebelsJE_TextEncoder": "Rebels JE • Text Encoder (Gemma fp8 + Connector)",
+    "RebelsJE_VAELoader": "Rebels JE • VAE Loader (video+audio)",
+    "RebelsJE_Assemble": "Rebels JE • Assemble Model",
+          }