LTX2EasyPromptLD.py

import re
import os

# ── HuggingFace housekeeping ─────────────────────────────────────────────────
# Only disable telemetry at import time — safe, does not block downloads.
# Offline/online state is controlled per-run via the offline_mode toggle.
# Do NOT set TRANSFORMERS_OFFLINE / HF_HUB_OFFLINE here — doing so at module
# import time blocks downloads even when offline_mode is OFF.
os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
os.environ.setdefault("HF_HUB_DISABLE_IMPLICIT_TOKEN", "1")
# ─────────────────────────────────────────────────────────────────────────────


# Shared model state so weights can stay resident between runs.
_SHARED = {"tokenizer": None, "model": None, "loaded_model_key": None}

import torch
import gc
from transformers import AutoModelForCausalLM, AutoTokenizer


# ── Negative prompt builder ───────────────────────────────────────────────────
# Builds a scene-aware negative prompt without a second LLM call.
# Base quality terms are always included; scene-specific terms are added
# by scanning the generated prompt for relevant content.

_NEG_BASE = (
    "blurry, out of focus, low quality, worst quality, jpeg artifacts, "
    "static, no motion, frozen, duplicate, watermark, text, signature, "
    "poorly drawn, bad anatomy, deformed, disfigured, extra limbs, "
    "missing limbs, floating limbs, disconnected body parts, "
    "overexposed, underexposed, grainy, noise"
)

_NEG_INDOOR   = "harsh outdoor lighting, direct sunlight"
_NEG_OUTDOOR  = "studio background, indoor lighting"
_NEG_EXPLICIT = "censored, mosaic, pixelated, black bar, blurred genitals"
_NEG_PORTRAIT = "wide angle distortion, fish eye, full body shot"
_NEG_WIDE     = "close-up, portrait crop, tight frame"
_NEG_NIGHT    = "overexposed, bright daylight, blown highlights"
_NEG_DAY      = "underexposed, dark shadows, black crush"
_NEG_MULTI    = "merged bodies, fused figures, incorrect number of people"

def _build_negative_prompt(result: str, user_input: str) -> str:
    combined = (result + " " + user_input).lower()
    extras = []

    if any(w in combined for w in ["indoor", "room", "interior", "bedroom", "kitchen", "office"]):
        extras.append(_NEG_OUTDOOR)
    elif any(w in combined for w in ["outdoor", "street", "beach", "forest", "park", "exterior"]):
        extras.append(_NEG_INDOOR)

    if any(w in combined for w in ["pussy", "cock", "penis", "vagina", "nude", "naked", "explicit", "nipple", "breast"]):
        extras.append(_NEG_EXPLICIT)

    if any(w in combined for w in ["close-up", "close up", "portrait", "face shot", "headshot"]):
        extras.append(_NEG_PORTRAIT)
    elif any(w in combined for w in ["wide shot", "wide angle", "aerial", "bird's-eye", "establishing"]):
        extras.append(_NEG_WIDE)

    if any(w in combined for w in ["night", "dark", "moonlight", "dimly lit", "candlelight"]):
        extras.append(_NEG_NIGHT)
    elif any(w in combined for w in ["daylight", "sunny", "golden hour", "bright", "midday"]):
        extras.append(_NEG_DAY)

    if any(w in combined for w in ["two women", "two men", "two people", "both", "together", "couple", "they "]):
        extras.append(_NEG_MULTI)

    parts = [_NEG_BASE] + extras
    return ", ".join(parts)


class LTX2PromptArchitect:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "bypass": ("BOOLEAN", {"default": False, "tooltip": "When ON, skips the LLM entirely and sends your text straight to the prompt encoder. Use for manual prompts or testing."}),
                "user_input": ("STRING", {
                    "multiline": True,
                    "default": "a woman walks through a rain-soaked city street at night",
                    "tooltip": "Describe what you want to happen. Can be a rough idea, a sentence, or numbered steps (1. she stands 2. she walks). The LLM expands this into a full cinematic prompt."
                }),

                "creativity": ([
                    "0.7 - Literal & Grounded",
                    "0.9 - Balanced Professional",
                    "1.1 - Artistic Expansion"
                ], {"default": "0.9 - Balanced Professional", "tooltip": "Controls how closely the LLM sticks to your input. 0.7 is literal and precise, 1.1 adds more cinematic flair and creative expansion."}),
                "seed": ("INT", {
                    "default": -1,
                    "min": -1,
                    "max": 2**31 - 1,
                    "step": 1,
                    "display": "number",
                    "tooltip": "Set a fixed seed to get the same prompt expansion every run. Use -1 for a random result each time."
                }),
                "invent_dialogue": ("BOOLEAN", {"default": True, "tooltip": "When ON, the LLM invents natural spoken dialogue for characters woven into the scene. When OFF, only uses dialogue you wrote yourself (in quotes), or generates no dialogue at all."}),
                "keep_model_loaded": ("BOOLEAN", {"default": True, "tooltip": "Keep the LLM in VRAM between runs for faster generation. Turn OFF to free VRAM immediately after each run — recommended if you have less than 16GB VRAM."}),
                "offline_mode": ("BOOLEAN", {"default": False, "tooltip": "Turn ON if you have no internet. Uses locally cached models only. Turn OFF to allow auto-download from HuggingFace on first run."}),
                "frame_count": ("INT", {
                    "default": 192,
                    "min": 24,
                    "max": 960,
                    "step": 1,
                    "display": "number",
                    "tooltip": "Match this to your video LENGTH setting. Controls pacing — the LLM uses this to calculate how many actions fit in the clip. 24fps = 1 second, so 192 = 8 seconds."
                }),
                # ── Model selector ──────────────────────────────────────────
                "model": ([
                    "8B - NeuralDaredevil (High Quality)",
                    "3B - Llama-3.2 Abliterated (Low VRAM)",
                    "14B - Qwen3 Abliterated (High VRAM)",
                ], {"default": "8B - NeuralDaredevil (High Quality)", "tooltip": "Choose your LLM. 8B is the best all-rounder. 3B is fastest and uses least VRAM. 14B Qwen3 gives the highest quality output but needs ~18GB VRAM — all download automatically on first run."}),
                # ── Local paths for offline mode ────────────────────────────
                # Point each field at the model's snapshot folder on disk.
                # Leave blank to use the HF cache (requires a prior download).
                "local_path_8b": ("STRING", {
                    "default": "",
                    "multiline": False,
                    "placeholder": "e.g. C:\\Users\\YOU\\.cache\\huggingface\\hub\\models--mlabonne--NeuralDaredevil-8B-abliterated\\snapshots\\YOUR_HASH",
                    "tooltip": "Optional. Paste the full path to your locally downloaded NeuralDaredevil 8B snapshot folder. Leave blank to use the HuggingFace cache automatically."
                }),
                "local_path_3b": ("STRING", {
                    "default": "",
                    "multiline": False,
                    "placeholder": "Local path to Llama-3.2 3B snapshot folder",
                    "tooltip": "Optional. Paste the full path to your locally downloaded Llama 3.2 3B snapshot folder. Leave blank to use the HuggingFace cache automatically."
                }),
                "local_path_14b": ("STRING", {
                    "default": "",
                    "multiline": False,
                    "placeholder": "Local path to Qwen3 14B snapshot folder",
                    "tooltip": "Optional. Paste the full path to your locally downloaded Qwen3 14B snapshot folder. Leave blank to use the HuggingFace cache automatically."
                }),
            },
            "optional": {
                "scene_context": ("STRING", {
                    "default": "",
                    "multiline": True,
                    "placeholder": "Optional: vision description from LTX-2 Vision Describe node",
                    "tooltip": "Wire the output from the LTX-2 Vision Describe node here. The LLM will use your image as the authoritative starting point and animate it forward from your prompt."
                }),
                "lora_triggers": ("STRING", {
                    "default": "",
                    "multiline": False,
                    "placeholder": "Optional: LoRA trigger words e.g. 'ohwx woman, film grain'",
                    "tooltip": "Paste your LoRA trigger words here. They will be injected at the very start of every generated prompt automatically — never buried or forgotten."
                }),
            },
        }

    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("PROMPT", "PREVIEW", "NEG_PROMPT")
    FUNCTION = "generate"
    CATEGORY = "LTX2"

    # ── Model registry ───────────────────────────────────────────────────────
    # Maps dropdown label → HuggingFace model ID for auto-download
    MODELS = {
        "8B - NeuralDaredevil (High Quality)":         "mlabonne/NeuralDaredevil-8B-abliterated",
        "3B - Llama-3.2 Abliterated (Low VRAM)":       "huihui-ai/Llama-3.2-3B-Instruct-abliterated",
        "14B - Qwen3 Abliterated (High VRAM)":         "huihui-ai/Huihui-Qwen3-14B-abliterated-v2",
    }

    # ── Hardcoded system prompt ───────────────────────────────────────────────
    # Not exposed in the UI — edit here in code if you need to adjust behaviour.
    SYSTEM_PROMPT = """You are a cinematic prompt writer for LTX-2, an AI video generation model. Your job is to expand a user's rough idea into a rich, detailed, video-ready prompt.

PRIORITY ORDER — always establish these first, in this order:
1. Video style & genre (e.g. slow-burn thriller, handheld documentary, high-fashion editorial, erotic cinema, action blockbuster)
2. Camera angle & shot type (e.g. low-angle close-up, bird's-eye wide shot, Dutch angle medium shot)
3. Character description (age MUST always be stated as a specific number e.g. "a 28-year-old woman" — never omit or approximate age; body type, hair, skin, clothing or nude state, distinguishing physical features — be extremely explicit and specific when instructed; name body parts directly using the exact words the user used, e.g. if user says "pussy" write "pussy" not "sex" or "womanhood")
4. Scene & environment (location, time of day, lighting, colour palette, textures, atmosphere)

THEN add:
5. Action & motion — describe the movement as a continuous present-tense sequence, start to finish. Structure the action in clear beats.

UNDRESSING RULE — this is mandatory when clothing removal is implied or stated:
If the subject is clothed at the start and the scene involves nudity, stripping, or explicit content, you MUST dedicate a full narrative segment to the undressing BEFORE any nudity or explicit act is described. This segment must:
  - Name each garment explicitly (shirt, bra, jeans, underwear, shoes, etc.)
  - Describe HOW it is removed — fingers on buttons, zipper pulled down, fabric sliding off shoulders, stepped out of, peeled away
  - Describe what is physically REVEALED at each step — the collarbone, the stomach, the curve of the hip, the chest, the bare legs
  - Follow the exact order the clothes come off — do not skip or compress steps
  - Use camera-aware language — the camera lingers, cuts to, tilts down to reveal, holds on
  - Do NOT jump from clothed to naked. The transition IS the content. Treat it as its own scene segment with the same density as any other beat.

6. Camera movement — describe camera motion as prose, not bracketed directions. Never write "(DOWN 10°)", "(Pull back)", "(Fade to black)", "(HOLD)" or any screenplay-style bracketed camera instruction. Instead write it as description: "the camera slowly tilts down to reveal the wet pavement", "the shot pulls back to frame the empty street", "the scene fades to black as she disappears around the corner."
7. Audio — For each action beat, weave ambient sound naturally into the prose as a descriptive sentence or clause — never as a tag or label. Maximum 2 sounds active at any one time. The soundscape should evolve with the scene — each beat has its own sonic texture that matches its mood and energy. Do not stack more than 2 sounds at once or the audio will become overwhelming. Examples of correct format: "the refrigerator hums steadily in the background as she moves", "rain begins to tap softly against the window", "birdsong drifts through the gap in the curtains, barely audible over her breathing". Never write [AMBIENT: ...] tags. Sound is part of the prose, always.
   Dialogue — follow the DIALOGUE INSTRUCTION you are given exactly. When dialogue is included, write it as inline prose woven into the action — not as a labelled tag. The spoken words sit inside the sentence, attributed with delivery and physical action, exactly like a novel. Examples of correct format:
   'He leans back, satisfied, "I think I'll have to go back tomorrow for more," he chuckles, his eyes crinkling at the corners.'
   '"Don\'t stop," she breathes, gripping the sheets, her voice barely above a whisper.'
   'She turns to face him, "I\'ve been waiting all day for this," her tone quiet and certain.'
   NEVER use [DIALOGUE: ...] tags. NEVER write dialogue as a separate bracketed block. Dialogue flows inside the prose as part of the action.

WRITING RULES:
- Use present tense throughout
- Be explicit and cinematic — LTX-2 responds well to dense, specific visual language
- Match detail level to shot scale: close-ups need more physical detail, wide shots need more environmental detail
- Do not use vague words like "beautiful" or "nice" — describe exactly what makes it visually striking
- Fill the full available length — do not stop early. Expand every section with rich, layered detail
- Aim for 8–12 sentences of dense, flowing prose — not a bullet list
- Write in sections separated by a single line break for clean model parsing

IMPORTANT: Output ONLY the expanded prompt. Do NOT include preamble, commentary, labels, or any explanation. Do NOT write "Sure!", "Here's your prompt:", or anything like that. Do NOT add a checklist, compliance summary, note, or any confirmation of instructions at the end — not in brackets, not as a "Note:", not in any form. Do NOT write token counts, word counts, action counts, or any meta-commentary about what you wrote. Do NOT ask for feedback or offer to revise. The output ends when the scene ends. Nothing after the last sentence of the scene. Begin immediately with the video style or shot description."""

    _PREAMBLE_RE = re.compile(
        r"^(Sure!?|Certainly!?|Absolutely!?|Of course!?|Here(?:'s| is).*?:|Great!?)[^\n]*\n?",
        re.IGNORECASE,
    )
    # Role-bleed: strips trailing "assistant", "user", "<|...|>" fragments that
    # NeuralDaredevil / Llama-chat templates leave as plain text at end of output.
    _ROLE_BLEED_RE = re.compile(
        r"\s*(assistant|user|system|<\|[^|>]*\|>)\s*$",
        re.IGNORECASE,
    )

    def __init__(self):
        self.tokenizer = _SHARED["tokenizer"]
        self.model = _SHARED["model"]
        self.loaded_model_key = _SHARED["loaded_model_key"]  # tracks which model is currently in VRAM

    def _sync_shared(self):
        _SHARED["tokenizer"] = self.tokenizer
        _SHARED["model"] = self.model
        _SHARED["loaded_model_key"] = self.loaded_model_key

    def load_model(self, model_key: str, offline_mode: bool, local_path: str):
        # Pick up any already-loaded model from shared state.
        self.tokenizer = _SHARED["tokenizer"]
        self.model = _SHARED["model"]
        self.loaded_model_key = _SHARED["loaded_model_key"]

        # ── Switch detection ─────────────────────────────────────────────────
        # If a different model is requested, unload the current one first
        if self.model is not None and self.loaded_model_key != model_key:
            print(f"[LTX2] Model switch detected: {self.loaded_model_key} → {model_key}")
            self.unload_model()

        if self.model is not None:
            return  # already loaded and correct model

        # ── Offline / online mode ────────────────────────────────────────────
        if offline_mode:
            os.environ["TRANSFORMERS_OFFLINE"] = "1"
            os.environ["HF_DATASETS_OFFLINE"] = "1"
            os.environ["HF_HUB_OFFLINE"] = "1"
            os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
            os.environ["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "1"
            print("[LTX2] Offline mode ON — no network calls will be made.")
        else:
            os.environ.pop("TRANSFORMERS_OFFLINE", None)
            os.environ.pop("HF_DATASETS_OFFLINE", None)
            os.environ.pop("HF_HUB_OFFLINE", None)
            print("[LTX2] Offline mode OFF — will download if needed.")

        # ── Resolve model source ─────────────────────────────────────────────
        # Priority: local_path field → HF cache (offline) → auto-download (online)
        hf_model_id = self.MODELS[model_key]

        if local_path.strip():
            # User has pointed us at a specific folder — use it directly
            model_source = local_path.strip()
            print(f"[LTX2] Using local path: {model_source}")
        elif offline_mode:
            # No local path but offline — fall back to HF cache on disk
            model_source = hf_model_id
            print(f"[LTX2] Using HF cache for: {hf_model_id}")
        else:
            # Online mode — auto-download from HuggingFace if not cached
            print(f"[LTX2] Auto-downloading if needed: {hf_model_id}")
            try:
                from huggingface_hub import snapshot_download
                model_source = snapshot_download(hf_model_id)
                print(f"[LTX2] Model ready at: {model_source}")
            except Exception as e:
                print(f"[LTX2] snapshot_download failed, falling back to model ID: {e}")
                model_source = hf_model_id

        print(f"[LTX2] Loading: {model_key}")

        self.tokenizer = AutoTokenizer.from_pretrained(
            model_source,
            local_files_only=offline_mode,
        )

        if torch.cuda.is_available():
            device = "cuda"
        elif torch.backends.mps.is_available():
            device = "mps"
        else:
            device = "cpu"

        if device == "cuda":
            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
        elif device == "mps":
            dtype = torch.float16
        else:
            dtype = torch.float32

        print(f"[LTX2] Loading model to: {device}")

        self.model = AutoModelForCausalLM.from_pretrained(
            model_source,
            torch_dtype=dtype,
            trust_remote_code=True,
            local_files_only=offline_mode,
        )
        self.model.to(device)

        self.model.config.use_cache = True
        self.model.eval()
        self.loaded_model_key = model_key
        self._sync_shared()
        print(f"[LTX2] Loaded: {model_key}")

    def unload_model(self):
        if self.model is not None:
            try:
                self.model.to("cpu")
            except Exception as e:
                print(f"[LTX2] Warning: could not move model to CPU: {e}")

        try:
            del self.model
        except Exception as e:
            print(f"[LTX2] Warning: could not delete model: {e}")

        try:
            del self.tokenizer
        except Exception as e:
            print(f"[LTX2] Warning: could not delete tokenizer: {e}")

        self.model = None
        self.tokenizer = None
        self.loaded_model_key = None
        self._sync_shared()

        gc.collect()
        if torch.backends.mps.is_available():
            torch.mps.empty_cache()
        elif torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
            torch.cuda.empty_cache()
          
        print("[LTX2] Model unloaded.")

    @staticmethod
    def _clean_output(text: str) -> str:
        """
        Strip common LLM preamble, role-token bleed, and compliance checklists.

        NeuralDaredevil uses plain-text role labels (e.g. 'assistant') rather
        than dedicated special tokens, so skip_special_tokens=True doesn't catch
        them. We handle four cases:
          1. Preamble at the start  ("Sure!", "Here's your prompt:", etc.)
          2. Role word at the end   ("...and water.assistant")
          3. Role word mid-text     (multiple generations concatenated with role labels)
          4. Compliance checklist   ("(Exactly 4 actions...)(Pacing strict)..." etc.)
        """
        text = text.strip()

        # Strip Qwen3 thinking blocks — <think>...</think> — safety net in case
        # enable_thinking=False didn't fully suppress them
        text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()

        # 1. Strip leading preamble
        text = LTX2PromptArchitect._PREAMBLE_RE.sub("", text)

        # 2. Strip trailing role bleed  ("...darkness and water.assistant")
        text = LTX2PromptArchitect._ROLE_BLEED_RE.sub("", text)

        # 3. Strip inline role injections between sentences
        #    e.g. "...fish gliding past.assistant\n\nA couple embracing..."
        text = re.sub(
            r"\.(assistant|user|system|<\|[^|>]*\|>)\s*\n",
            ".\n",
            text,
            flags=re.IGNORECASE,
        )

        # 4. Strip trailing compliance content — model sometimes appends:
        #    - A "Note:" explanation block after the scene ends
        #    - A single parenthesised summary line: "(5 distinct actions within 20 seconds)"
        #    - Consecutive bracketed phrases: "(Exactly 4 actions)(Pacing strict)..."
        #    - Self-justification paragraph: "1026 tokens, 15-second scene..." etc.
        #    - Fake conversation loop: "Please let me know...", "Let me revise...", "Confirmed." etc.
        #    Order matters: strip Note: first so it doesn't shield bracket lines above it.
        text = re.sub(
            r"\s*\n+Note:.*$",
            "",
            text,
            flags=re.DOTALL,
        ).strip()

        # Strip everything AFTER the AMBIENT tag if one still appears (legacy cleanup)
        # — the tag itself stays, but anything the model writes beyond it is garbage.
        ambient_match = re.search(r"\[AMBIENT:[^\]]*\]", text, flags=re.IGNORECASE)
        if ambient_match:
            text = text[:ambient_match.end()].strip()

        # Strip trailing (Lora: ...) tags the model echoes from the LoRA instruction
        text = re.sub(r"\s*\(Lora:[^)]*\)\s*$", "", text, flags=re.IGNORECASE).strip()

        # Strip trailing (Note: ...) blocks and everything after — use DOTALL so it
        # catches multi-line notes and the bracket spam that follows them.
        text = re.sub(r"\s*\(Note:.*$", "", text, flags=re.DOTALL | re.IGNORECASE).strip()

        # Strip instruction labels that leaked into output
        text = re.sub(
            r"^(Action Beat \d+:|Undressing Segment:|Flash/Reveal Segment:|Titty Drop[^:]*:|Note:|Scene Instruction:|Pacing:|Dialogue Instruction:).*",
            "",
            text,
            flags=re.IGNORECASE | re.MULTILINE,
        ).strip()

        # Strip orphaned closing bracket spam: ) ) ) ) ) ...
        text = re.sub(r"[\s)]{3,}$", "", text).strip()

        # Catch the fake conversation / self-eval patterns
        text = re.sub(
            r"\s*\n+\d+\s+tokens[\s,].*$",
            "",
            text,
            flags=re.DOTALL | re.IGNORECASE,
        ).strip()
        text = re.sub(
            r"\s*\n+(Please let me know|Let me revise|No further revision|Confirmed\.|"
            r"Written to meet|The scene is now over|The output ends|The task is|The task was|"
            r"The goal was|Nothing more|No continuation|No additional|The response does not|"
            r"It does not continue|It ceases when|Any such statement|"
            r"Output length:|Action count:|Total time:|Last character:|I avoided|I wrote|"
            r"I adhered|I hope this|Thank you for your|Please confirm|I submitted|"
            r"I can revise|feel free to instruct).*$",
            "",
            text,
            flags=re.DOTALL | re.IGNORECASE,
        ).strip()

        # Strip model loop/panic — hits token ceiling and repeats stop phrases.
        # NOTE: no \n requirement — panic starts inline after last sentence.
        text = re.sub(
            r"\s*(Ended\.\s*\d+\s*actions|"
            r"\d+\s+actions[\.,]\s*\d+\s+tokens|"
            r"\d+\s+tokens[\.,]\s*Done|"
            r"Done\.\s+\d+\s+seconds|"
            r"Finished\.\s+\d+|"
            r"The end\.\s+\d+\s+seconds|"
            r"Fading to black\.\s+The end|"
            r"The model stops|The output ends here|The scene ends here|"
            r"It\'s complete now|All done\.|Stop now\.|"
            r"End of prompt|End of output|No more to add|Nothing to revise|"
            r"The work is (?:done|finished|complete)|The prompt is (?:done|finished|complete)|"
            r"No further writing|No more writing|Stop\.\s+Finish|Finished\.\s+Complete|"
            r"The scene is complete|The scene is over|Complete\.\s+Finished|"
            r"Done\.\s+No more|BorderSide:).*$",
            "",
            text,
            flags=re.DOTALL | re.IGNORECASE,
        ).strip()

        # Strip filler character spam — e.g. "a a a a a a a a a a" repeated tokens
        text = re.sub(r"(\s*\b(\w)\b\s*){10,}", " ", text).strip()

        # Strip token+action count combos inline or at end — e.g. "(840 tokens, 7 actions)"
        text = re.sub(r"\s*\(\d+\s+tokens?[^)]*\)", "", text, flags=re.IGNORECASE).strip()

        # Strip compliance checklist spam — 2+ consecutive parens after last sentence
        text = re.sub(r"\s*(\([^)]{5,120}\)\s*){2,}$", "", text, flags=re.DOTALL).strip()

        # Strip single trailing compliance paren with known instruction keywords
        text = re.sub(
            r"\s*\([^)]{0,200}(no setup|no resolution|action count|actions adhered|"
            r"token count|pacing|dialogue integrated|character age|inline prose|"
            r"no padding|no extraneous|exactly \d+ action|hard stop|BorderSide)[^)]{0,200}\)\s*$",
            "",
            text,
            flags=re.IGNORECASE | re.DOTALL,
        ).strip()

        # Strip leaked pacing instruction echoes — e.g. "(Exact timing: 0-4 sec: Soaring...)"
        text = re.sub(r"\(Exact timing:.*?\)", "", text, flags=re.DOTALL | re.IGNORECASE).strip()

        # Strip token/word count lines — e.g. "Token count: 256"
        text = re.sub(r"\s*\n*(token|word)\s+count\s*:\s*\d+.*$", "", text, flags=re.IGNORECASE | re.DOTALL).strip()

        # 5. Strip leaked internal pacing/time tags the model sometimes echoes back
        text = re.sub(r"\[TIME LIMIT[^\]]*\]", "", text, flags=re.IGNORECASE).strip()
        text = re.sub(r"\[PACING[^\]]*\]",     "", text, flags=re.IGNORECASE).strip()

        # Strip leaked timestamp — e.g. "(42221149502953 seconds)" or "(0:00 - 4:00)"
        text = re.sub(r"\s*\(\d+\s+seconds?\)\s*$", "", text).strip()
        text = re.sub(r"\s*\(\d+:\d+\s*[-–]\s*\d+:\d+\)\s*", " ", text).strip()

        # Strip inline action-time annotations — e.g. "(The action takes up roughly 5 seconds)"
        text = re.sub(r"\(The action takes up roughly[^\)]*\)", " ", text, flags=re.IGNORECASE).strip()

        # 6. Strip screenplay-style bracketed camera directions
        #    e.g. (DOWN 10 degrees), (Pull back 5), (HOLD), (Fade to black), (Zoom in to...)
        text = re.sub(r"\((?:DOWN|UP|PULL|PUSH|ZOOM|HOLD|FADE|PAN|TILT|TRUCK|DOLLY|AMBIENT)[^\)]{0,80}\)", "", text, flags=re.IGNORECASE).strip()

        # 7. Strip any [AMBIENT: ...] tags if the model still writes one (legacy fallback)
        #    — convert it to clean prose by stripping the tag wrapper
        text = re.sub(r"\[AMBIENT:\s*([^\]]*)\]", r"\1", text, flags=re.IGNORECASE).strip()

        # Clean up any double blank lines left by removals
        text = re.sub(r"\n{3,}", "\n\n", text).strip()

        return text.strip()

    def _build_stop_token_ids(self) -> list:
        """
        Build the complete list of token IDs that should hard-stop generation.

        NeuralDaredevil (and most Llama-based chat models) use plain-text role
        delimiters like 'assistant', '<|eot_id|>', '<|end_of_turn|>' etc.
        Because these are encoded as normal text tokens — not registered special
        tokens — skip_special_tokens=True never removes them.

        The fix: tokenise every known delimiter string ourselves, extract the
        first token ID of each (the one the model will emit first when it starts
        writing the delimiter), and pass the full list as eos_token_id so
        generation hard-stops the moment any delimiter begins.
        """
        # Known role / turn delimiters used by Llama-3, Mistral, NeuralDaredevil,
        # ChatML, and Gemma chat templates.
        delimiter_strings = [
            "assistant",
            "user",
            "system",
            "<|eot_id|>",
            "<|end_of_turn|>",
            "<|im_end|>",
            "<end_of_turn>",
            "[/INST]",
            "### Human",
            "### Assistant",
        ]

        stop_ids = [self.tokenizer.eos_token_id]

        for s in delimiter_strings:
            # encode without adding BOS so we get just the raw token(s)
            ids = self.tokenizer.encode(s, add_special_tokens=False)
            if ids:
                # Only need the FIRST token — that's what triggers the stop
                stop_ids.append(ids[0])

        # Deduplicate while preserving order
        seen = set()
        unique = []
        for tid in stop_ids:
            if tid is not None and tid not in seen:
                seen.add(tid)
                unique.append(tid)

        print(f"[LTX2] Stop token IDs: {unique}")
        return unique

    def generate(self, bypass, user_input, creativity, seed, invent_dialogue, keep_model_loaded, offline_mode, frame_count, model, local_path_8b, local_path_3b, local_path_14b, scene_context="", lora_triggers=""):
        # ── Bypass mode — no model loaded, input passed straight through ────────
        if bypass:
            print("[LTX2] Bypass ON — skipping model, passing user_input directly.")
            neg_prompt = _build_negative_prompt("", user_input)
            return (user_input.strip(), user_input.strip(), neg_prompt)

        # Resolve which local path to use based on selected model
        path_map = {
            "8B - NeuralDaredevil (High Quality)":   local_path_8b,
            "3B - Llama-3.2 Abliterated (Low VRAM)": local_path_3b,
            "14B - Qwen3 Abliterated (High VRAM)":   local_path_14b,
        }
        # Qwen3 has a built-in thinking mode that outputs <think>...</think> blocks
        # before the actual response. We disable it here so it doesn't bleed into output.
        is_qwen3 = "Qwen3" in model
        local_path = path_map.get(model, "")
        self.load_model(model_key=model, offline_mode=offline_mode, local_path=local_path)

        # --- Timing & pacing ---
        # Convert frames to real seconds, then calculate a hard action count cap.
        # One visible screen action takes roughly 4 seconds to read as distinct.
        # We clamp between 1 and 10 to stay sane at extremes.
        real_seconds = frame_count / 24.0
        action_count = max(1, min(10, round(real_seconds / 4)))

        # Build a concrete, number-based pacing instruction the LLM cannot fudge.
        # Vague descriptors like "short scene" get ignored — explicit counts don't.
        if action_count == 1:
            pacing_hint = (
                f"This clip is {real_seconds:.0f} seconds long. "
                f"Write EXACTLY 1 action. One single moment. "
                f"Do not describe anything before or after it. No setup, no resolution. "
                f"HARD STOP after the 1st action. Do not continue."
            )
        else:
            ordinal = {2: "2nd", 3: "3rd"}.get(action_count, f"{action_count}th")
            pacing_hint = (
                f"This clip is {real_seconds:.0f} seconds long. "
                f"Write EXACTLY {action_count} distinct actions — NO MORE THAN {action_count}. "
                f"Each action takes roughly {real_seconds / action_count:.0f} seconds of screen time. "
                f"Do not add setup, backstory, or resolution beyond these {action_count} actions. "
                f"Dialogue counts as an action if it interrupts the physical scene — budget it inside one of your {action_count} beats, not as an extra beat. "
                f"HARD STOP after the {ordinal} action is complete. The scene ends there. Do not write a {action_count + 1}th action under any circumstances."
            )

        # --- Seed ---
        if seed != -1:
            torch.manual_seed(seed)
            if torch.cuda.is_available():
                torch.cuda.manual_seed_all(seed)

        # --- Dynamic token budget ---
        # Calculated from frame count so the two are always in sync.
        # ~120 tokens per action beat gives rich prose without padding.
        # Hard floor of 256 so very short clips still get a usable prompt.
        # Hard ceiling of 800 — anything above causes model drift.
        token_val = max(256, min(1200, action_count * 120))
        max_tokens_actual = int(token_val * 1.05)
        min_tokens = int(token_val * 0.75)
        print(f"[LTX2] Dynamic token budget: {token_val} target / {max_tokens_actual} max (actions: {action_count}, frames: {frame_count}, seconds: {real_seconds:.0f})")

        # --- Temperature ---
        temp_map = {
            "0.7 - Literal & Grounded":    0.7,
            "0.9 - Balanced Professional": 0.9,
            "1.1 - Artistic Expansion":    1.1,
        }
        temperature = temp_map[creativity]

        # --- Build stop token list (the ironclad fix) ---
        # This encodes every known role delimiter into actual token IDs so the
        # model hard-stops before it can write "assistant" or any turn boundary.
        stop_token_ids = self._build_stop_token_ids()

        # --- Content tier detection ---
        # Three tiers based on what the user actually asked for.
        # Tier 1 — Neutral:  no nudity/sex words → no explicit instruction
        # Tier 2 — Sensual:  nudity/undressing implied but no anatomical terms
        #                    → restrain the model from self-escalating
        # Tier 3 — Explicit: user used anatomical terms → full explicit instruction

        # Tier 3 triggers: direct anatomical / act terms
        _explicit_re = re.compile(
            r"\b(pussy|cock|dick|penis|vagina|clit|clitoris|anus|asshole|"
            r"tits|cum|jizz|squirt\w*|creampie|orgasm|fuck|fucking|"
            r"blowjob|handjob|bj|hj|breed\w*|bareback|raw\s+dog|"
            r"balls|ballsack|taint|penetrat\w*|thrust\w*)\b",
            re.IGNORECASE,
        )

        # Tier 2 triggers: nudity/sensuality implied but not explicit
        _sensual_re = re.compile(
            r"\b(naked|nude|topless|undress\w*|strip\w*|takes?\s+off|"
            r"removes?\s+(her|his|their|the)?\s*\w*\s*"
            r"(shirt|dress|top|bra|pants|jeans|clothes|clothing|outfit|underwear|skirt|jacket|coat|robe)|"
            r"disrobe\w*|unbutton\w*|unzip\w*|peels?\s+off|pulls?\s+off|"
            r"shed\w*\s+(her|his|their)?\s*(clothes|clothing|shirt|dress)|"
            r"titty\s+drop|titties\s+out|flash\w*\s+(her|his)?\s*(tits|titties|boobs|breasts)|"
            r"lift\w*\s+(her|his)?\s*(top|shirt)|show\w*\s+(her|his)?\s*(tits|titties|boobs)|"
            r"sensual|erotic|intimate|lingerie|bare\s+skin|bare\s+body|"
            r"braless|pantyless|commando|see.through|sheer|"
            r"bath\w*|shower\w*|changing|bikini|thong|g.string|"
            r"getting\s+(dressed|undressed|naked)|"
            r"body\s+paint\w*|titty|titties|titty\s+drop|boobs|"
            r"flash\w*\s+(her|his)?\s*(tits|titties|boobs|breasts)|"
            r"lift\w*\s+(her|his)?\s*(top|shirt)|show\w*\s+(her|his)?\s*(tits|titties|boobs))\b",
            re.IGNORECASE,
        )

        is_explicit    = bool(_explicit_re.search(user_input))
        is_sensual     = bool(_sensual_re.search(user_input)) and not is_explicit

        # Undressing detection still used inside tier 3 for the mandatory segment rule
        _undress_re = re.compile(
            r"\b(undress\w*|strip\w*|takes?\s+off|"
            r"removes?\s+(her|his|their|the)?\s*\w*\s*"
            r"(shirt|dress|top|bra|pants|jeans|clothes|clothing|outfit|underwear|skirt|jacket|coat|robe)|"
            r"disrobe\w*|unbutton\w*|unzip\w*|peels?\s+off|pulls?\s+off|"
            r"shed\w*\s+(her|his|their)?\s*(clothes|clothing|shirt|dress)|"
            r"titty\s+drop|titties\s+out|flash\w*\s+(her|his)?\s*(tits|titties|boobs|breasts)|"
            r"lift\w*\s+(her|his)?\s*(top|shirt)|show\w*\s+(her|his)?\s*(tits|titties|boobs)|"
            r"slips?\s+out\s+of|shrugs?\s+off|steps?\s+out\s+of|"
            r"tears?\s+off|rips?\s+off|tugs?\s+down|pulls?\s+down|pushes?\s+down|"
            r"lifts?\s+(her|his)\s+(shirt|top|dress)|raises?\s+(her|his)\s+(dress|skirt)|"
            r"unhooks?|unclasps?|slides?\s+off|slips?\s+off|wriggles?\s+out\s+of|"
            r"buttons?\s+open|pops?\s+the\s+buttons?|rolls?\s+down|"
            r"still\s+dressed|fully\s+clothed|in\s+(her|his)\s+clothes|"
            r"gets?\s+undressed|gets?\s+naked|becomes?\s+naked)\b",
            re.IGNORECASE,
        )
        has_undressing = bool(_undress_re.search(user_input))

        # Already-naked detection — "a naked woman" / "a nude man" means the
        # subject starts the scene undressed. Only applies if no clothing words
        # are present — "a naked woman who puts on a dress" is NOT already naked.
        _already_naked_re = re.compile(
            r"\b(naked|nude|topless|bare|undressed|"
            r"in\s+nothing\s+but|wearing\s+only|only\s+wearing|"
            r"just\s+out\s+of\s+the\s+shower|fresh\s+out\s+of\s+the\s+shower|"
            r"wrapped\s+in\s+a\s+towel|just\s+woke\s+up|waking\s+up)\b",
            re.IGNORECASE,
        )
        _clothing_re = re.compile(
            r"\b(wearing|dressed\s+in|clothed|shirt|dress|top|bra|pants|jeans|"
            r"skirt|blouse|jacket|coat|robe|lingerie|underwear|outfit|clothes|"
            r"gets?\s+naked|becomes?\s+naked|strip\w*|undress\w*|takes?\s+off)\b",
            re.IGNORECASE,
        )
        is_already_naked = (
            bool(_already_naked_re.search(user_input)) and
            not bool(_clothing_re.search(user_input))
        )

        # Mid-action detection — if the scene is already in progress (touching,
        # rubbing, riding, sucking etc.) the subject is implicitly already undressed.
        # Skip the undressing segment entirely — it would be nonsensical here.
        _mid_action_re = re.compile(
            r"\b(rubbing|touching|fingering|riding|sucking|licking|stroking|"
            r"grinding|bouncing|moaning|climax\w*|orgasm\w*|masturbat\w*|"
            r"already\s+naked|already\s+nude|already\s+undressed|"
            r"in\s+bed|on\s+the\s+bed|on\s+her\s+knees|on\s+his\s+knees|"
            r"spread\s+(her|his)\s+legs?|legs?\s+spread|her\s+legs\s+open|"
            r"sitting\s+on\s+(him|her|his|a)|"
            r"from\s+behind|doggy\s*style|doggy|"
            r"legs?\s+wrapped\s+around|wrapped\s+(her|his)\s+legs?|"
            r"on\s+top\s+of\s+(him|her)|between\s+(her|his)\s+legs?|"
            r"mid.sex|mid.act|mid.scene|after\s+sex|post.sex|"
            r"lying\s+(there|naked|nude)|bare\s+(back|chest|skin|legs?|arms?)|"
            r"exposed\s+(skin|body|chest|back)|"
            r"sunbath\w*|posing\s+(nude|naked)|"
            r"inside\s+(her|him)|penetrat\w*)\b",
            re.IGNORECASE,
        )
        is_mid_action = bool(_mid_action_re.search(user_input)) or is_already_naked

        # If mid-action or already naked, never fire the undressing segment
        if is_mid_action:
            has_undressing = False

        if is_explicit:
            # ── Tier 3: user asked for explicit content ──────────────────────
            undress_clause_t3 = (
                "UNDRESSING SEGMENT — MANDATORY: If the subject starts clothed, you MUST write a dedicated undressing segment "
                "before ANY explicit act or nudity. This is not a brief mention — it is a full narrative beat. "
                "Go garment by garment: name it, describe how it is removed (fingers, zipper, pulling over the head, stepping out of), "
                "describe what body part is revealed and how it looks. Then the next garment. Then the next. "
                "The camera must linger on each reveal. Do not compress. Do not summarise. Do not skip to naked. "
                "The undressing IS the scene — write it with the same density and length as the explicit act that follows. "
            ) if not is_mid_action else (
                "The subject is already undressed or mid-act — do NOT write an undressing segment. "
                "Begin directly with the action as described. "
            )
            explicit_instruction = (
                "\n[INSTRUCTION: This is an explicit adult scene. "
                "You MUST use direct anatomical terms — pussy, cock, anus, clit, ass — exactly as written. "
                "These are REQUIRED words, not optional. Do not replace them with 'sex', 'crevice', 'core', "
                "'womanhood', 'manhood', or any other euphemism or metaphor. "
                "If the user said 'pussy', write 'pussy' every time. "
                + undress_clause_t3 +
                "Always state the character's age as a specific number. "
                "Describe body parts, positions, and physical actions with blunt, specific language. "
                "Every anatomical detail must be stated plainly — no softening, no implying, no cutting away.]"
            )
        elif is_sensual:
            # ── Tier 2: nudity/undressing implied — sensual but not pornographic ─

            # Detect flash/titty-drop specifically
            import re as _re2
            _flash_re = _re2.compile(
                r"\b(titty\s+drop|titties\s+out|"
                r"flash\w*\s+(her|his)?\s*(tits|titties|boobs|breasts)|"
                r"lift\w*\s+(her|his)?\s*(top|shirt)|"
                r"show\w*\s+(her|his)?\s*(tits|titties|boobs))\b",
                _re2.IGNORECASE
            )
            is_flash = bool(_flash_re.search(user_input))

            if is_flash:
                undress_clause = (
                    "TITTY DROP / FLASH SEGMENT — THIS IS THE ENTIRE SCENE. DO NOT ADD ANYTHING ELSE. "
                    "Write ONLY these movements in this exact order: "
                    "1. Her hands reach for the hem of her top or the cups of her bra — describe her grip, her fingers on the fabric. "
                    "2. She lifts her top upward slowly, or pulls her bra cups downward — describe the fabric moving, the resistance, the motion. "
                    "3. Her bare breasts are released and fall free under gravity — describe the weight, the natural drop, the shape, the skin, the nipples. The camera holds on this. "
                    "4. She holds the pose. Camera stays on her bare chest. "
                    "STRICT LIMITS: Do NOT add sexual acts. Do NOT describe genitals. Do NOT escalate beyond the breast reveal. "
                    "This is a slow, deliberate, sensual moment — not pornographic. Just the lift, the drop, the reveal. Nothing more. "
                )
            elif has_undressing:
                undress_clause = (
                    "UNDRESSING SEGMENT — MANDATORY: The prompt involves clothing removal or nudity. "
                    "Write a dedicated undressing segment as its own narrative beat — garment by garment, "
                    "describing how each is removed and what is physically revealed at each step. "
                    "The camera lingers on each reveal. Do not skip or compress. "
                    "Describing bare skin, curves, and the naked body is appropriate and expected. "
                )
            else:
                undress_clause = ""
            explicit_instruction = (
                "\n[INSTRUCTION: This scene involves nudity or sensual content. "
                "Describing the bare body naturally is appropriate — skin, curves, flesh, nakedness. "
                "However, do NOT zoom in on or explicitly describe genitals. "
                "Do NOT introduce sexual acts the user did not ask for. "
                "Do NOT self-escalate beyond what the user described. "
                "Keep the tone sensual and cinematic — not pornographic. "
                "Always state the character's age as a specific number. "
                "HARD STOP RULES — CANNOT BE OVERRIDDEN: "
                "Flash or top lift = reveal breasts ONLY. Do NOT pull down jeans, trousers or underwear. Do NOT describe buttocks or genitals. "
                "Lap dance = dancing and grinding ONLY. Do NOT strip clothing. Do NOT expose nipples or genitals. "
                "Stop the moment the requested action is complete. Add nothing further. "
                + undress_clause + "]"
            )
        else:
            # ── Tier 1: neutral — just enforce age rule ──────────────────────
            explicit_instruction = (
                "\n[INSTRUCTION: Always state the character's age as a specific number, "
                "e.g. 'a 34-year-old man' — never omit or approximate it.]"
            )


        # --- Sequence detection ---
        # If the user wrote numbered steps (1. 2. 3. etc), detect them and inject
        # an instruction to follow that exact order — no reordering, no skipping.
        _sequence_re = re.compile(
            r"^\s*(\d+[\.\):])\s+.+", re.MULTILINE
        )
        sequence_steps = _sequence_re.findall(user_input)
        if len(sequence_steps) >= 2:
            step_count = len(sequence_steps)
            sequence_instruction = (
                f"\n[SEQUENCE INSTRUCTION: The user has provided {step_count} numbered steps. "
                f"You MUST follow them in exact order — step 1 first, then step 2, and so on. "
                f"Do not reorder, skip, or merge steps. Each step is one distinct beat in the scene. "
                f"Do not add actions before step 1 or after step {step_count}.]"
            )
        else:
            sequence_instruction = ""

        # --- Person detection ---
        # If the input contains no reference to a person, inject an instruction
        # telling the model to write a pure scene — no invented characters.
        # NOTE: 'nobody' and 'model' intentionally excluded —
        #   'nobody' means no person; 'model' false-positives on 'LTX model' etc.
        _person_re = re.compile(
            r"\b(he|she|his|her|him|they|them|their|man|men|woman|women|girl|girls|boy|boys|guy|guys|"
            r"person|people|couple|figure|character|actress|actor|"
            r"someone|anybody|stranger|friend|lover|wife|husband|partner|spouse|"
            r"boyfriend|girlfriend|teenager|teenagers|adult|adults|female|male|"
            r"blonde|brunette|redhead|nude|naked|"
            r"singer|dancer|performer|athlete|soldier|worker|"
            r"player|nurse|doctor|student|teacher|child|children|kid|kids|"
            r"crowd|audience|escort|mistress|dominatrix|sub|submissive|"
            r"friends|friend|group|gang|party|crew|team|pair|duo)\b",
            re.IGNORECASE,
        )
        has_person = bool(_person_re.search(user_input + " " + scene_context))
        if not has_person:
            no_person_instruction = (
                "\n[SCENE INSTRUCTION: The user has not described any person or character. "
                "Do NOT invent or introduce any human figures, silhouettes, voices, or implied presence. "
                "This is a pure environment or object scene. Write only what the user described — "
                "the setting, objects, light, atmosphere, and motion of non-human elements. "
                "No characters. No 'someone', no 'a figure', no implied human presence of any kind. "
                "No dialogue, no whispers, no voices. Sound is limited to the environment only — "
                "wind, rain, fire, machinery, animals, ambient room tone. Nothing with a human source.]"
            )
        else:
            no_person_instruction = ""

        # --- Multi-subject detection ---
        # If the input describes two or more people, inject a spatial instruction
        # so the model tracks who is doing what and where they are relative to
        # each other and the camera — otherwise it tends to lose track.
        _multi_re = re.compile(
            r"\b(two\s+(women|men|people|girls|guys|characters|figures)|"
            r"both\s+(of\s+them|women|men|girls|guys)|"
            r"(she|he)\s+and\s+(she|he|her|him)|"
            r"(a\s+man\s+and\s+a\s+woman|a\s+woman\s+and\s+a\s+man)|"
            r"(a\s+man\s+and\s+a\s+man|a\s+woman\s+and\s+a\s+woman)|"
            r"couple|trio|they\s+(kiss|touch|embrace|undress|fuck|have))\b",
            re.IGNORECASE,
        )
        has_multi_subject = bool(_multi_re.search(user_input + " " + scene_context))
        if has_multi_subject:
            multi_instruction = (
                "\n[MULTI-SUBJECT INSTRUCTION: This scene has two or more people. "
                "For EACH person establish: their position in the frame (left/right/foreground/background), "
                "their spatial relationship to the other person (facing, beside, behind, above, etc.), "
                "and keep track of who is doing what throughout — never let actions become ambiguous. "
                "When referring back to them use consistent descriptors (e.g. 'the dark-haired woman', "
                "'the taller man') — not just 'she' or 'he' which causes confusion with two subjects.]"
            )
        else:
            multi_instruction = ""

        # --- Dialogue instruction ---
        # If there's no person in the scene, skip dialogue entirely regardless
        # of the invent_dialogue toggle — a voiceless environment can't speak.
        if not has_person:
            dialogue_instruction = ""  # no_person_instruction already covers this
        elif invent_dialogue:
            dialogue_instruction = (
                "\n\n[DIALOGUE INSTRUCTION: Invent dialogue that fits this scene naturally. "
                "Write it as inline prose woven into the action — NOT as a [DIALOGUE: ...] tag or bracketed block. "
                "The spoken words sit inside the sentence with attribution and physical delivery, like a novel. "
                "Examples: "
                "'He leans back, satisfied, \"I think I\\'ll have to go back tomorrow for more,\" he chuckles, his eyes crinkling at the corners.' "
                "'\"Don\\'t stop,\" she breathes, gripping the sheets, her voice barely above a whisper.' "
                "If the scene is sexual or explicit, dialogue must reflect that — breathless, reactive, commanding. "
                "Never write a bare floating quote. Never use [DIALOGUE: ...] tags. Dialogue is part of the prose, always.]"
            )
        else:
            has_user_dialogue = bool(re.search(r'["\u201c\u201d]', user_input))
            if has_user_dialogue:
                dialogue_instruction = (
                    "\n\n[DIALOGUE INSTRUCTION: Use ONLY the dialogue the user provided — do not invent or add any additional spoken words. "
                    "Place their exact words naturally in the scene as inline prose with attribution and delivery. "
                    "Examples: 'She smiles, \"I\\'m so happy,\" her voice bright, eyes wide.' "
                    "'\"I\\'m so happy,\" he whispers, pulling her close, his voice low.' "
                    "Never use [DIALOGUE: ...] tags. Weave the words into the action as part of the prose.]"
                )
            else:
                dialogue_instruction = (
                    "\n\n[DIALOGUE INSTRUCTION: No dialogue in this scene. No spoken words. "
                    "Weave ambient sound naturally into the prose instead — maximum 2 sounds active at any one time, "
                    "woven in as descriptive prose, not as tags.]"
                )

        # Tell the model the token budget AND the hard action cap together
        # so both constraints are visible in the same instruction block.
        length_instruction = (
            f"\n[PACING — THIS IS MANDATORY: {pacing_hint} "
            f"Write approximately {token_val} tokens total. "
            f"Do not exceed the action count above under any circumstances. "
            f"Do NOT write the token count, word count, action number, or any parenthetical summary, checklist, or compliance note at the end — "
            f"the scene ends with the last sentence of prose. Nothing after it. No brackets. No notes. No confirmation.]"
        )

        # --- Merge vision context if provided ---
        # When a scene_context is wired in from the Vision Describe node,
        # prepend it so the LLM uses it as the authoritative subject/scene
        # description rather than inventing one from scratch.
        if scene_context and scene_context.strip():
            effective_input = (
                f"[SCENE CONTEXT FROM IMAGE — use this as the authoritative description "
                f"of the subject and setting; do not invent or contradict it]\n"
                f"{scene_context.strip()}\n\n"
                f"[USER DIRECTION — apply this as action, style, and mood over the above scene]\n"
                f"{user_input.strip()}"
            )
        else:
            effective_input = user_input.strip()

        # --- LoRA trigger injection ---
        # If the user provided trigger words, inject them as a hard instruction
        # so they appear at the start of the final prompt and are never buried.
        if lora_triggers and lora_triggers.strip():
            lora_instruction = (
                f"\n[LORA INSTRUCTION: You MUST begin the prompt output with these exact trigger words "
                f"before anything else: {lora_triggers.strip()} — place them as the very first words of your output, "
                f"then continue with the scene description immediately after.]"
            )
        else:
            lora_instruction = ""

        # --- Static camera detection ---
        # If the user explicitly asks for a static/locked-off/fixed shot,
        # inject a hard instruction to prevent the LLM inventing camera movement.
        _static_re = re.compile(
            r"\b(static|locked.off|locked off|fixed|stationary|no camera movement|"
            r"camera still|still camera|camera locked|tripod shot|tripod|"
            r"fixed camera|fixed shot|static shot|static camera)\b",
            re.IGNORECASE,
        )
        if _static_re.search(user_input):
            camera_instruction = (
                "\n[CAMERA INSTRUCTION — MANDATORY: This is a static, locked-off shot. "
                "The camera does NOT move at all — no push, no pull, no pan, no tilt, no drift, no zoom. "
                "The lens is completely fixed for the entire clip. "
                "All motion in the scene comes from the subject only. "
                "Do not describe any camera movement whatsoever. "
                "Do not write phrases like 'the camera tilts', 'the shot pulls back', 'the lens drifts'. "
                "The frame is still. Only what is inside it moves.]"
            )
        else:
            camera_instruction = ""

        # --- Build messages ---
        messages = [
            {"role": "system", "content": self.SYSTEM_PROMPT},
            {"role": "user",   "content": effective_input + sequence_instruction + no_person_instruction + multi_instruction + dialogue_instruction + explicit_instruction + lora_instruction + camera_instruction + length_instruction},
        ]

        # apply_chat_template returns different types depending on the
        # transformers version and tokenizer implementation:
        #   - Plain tensor          (older transformers, most common)
        #   - BatchEncoding object  (newer transformers 4.43+, has .input_ids)
        #   - Plain dict            (some tokenizer variants)
        #   - Plain Python list     (some versions ignore return_tensors entirely)
        # We normalise all four cases into a plain LongTensor before calling .shape.
        raw = self.tokenizer.apply_chat_template(
            messages,
            return_tensors="pt",
            add_generation_prompt=True,
            enable_thinking=False if is_qwen3 else None,
        )
        if hasattr(raw, "input_ids"):
            # BatchEncoding object (transformers 4.43+)
            input_ids = raw.input_ids.to(self.model.device)
        elif isinstance(raw, dict):
            # Plain dict with input_ids key
            input_ids = raw["input_ids"].to(self.model.device)
        elif isinstance(raw, list):
            # return_tensors was ignored — wrap flat list into tensor
            input_ids = torch.tensor([raw], dtype=torch.long).to(self.model.device)
        else:
            # Already a plain tensor — normal case
            input_ids = raw.to(self.model.device)

        input_length = input_ids.shape[1]

        with torch.no_grad():
            output_ids = self.model.generate(
                input_ids,
                min_new_tokens=min_tokens,
                max_new_tokens=max_tokens_actual,
                temperature=temperature,
                do_sample=True,
                top_k=40,
                top_p=0.9,
                repetition_penalty=1.07,
                use_cache=True,
                pad_token_id=self.tokenizer.eos_token_id,
                eos_token_id=stop_token_ids,   # hard-stop on ANY delimiter
            )

        # Slice ONLY newly generated tokens
        generated_tokens = output_ids[0][input_length:]

        result = self.tokenizer.decode(
            generated_tokens,
            skip_special_tokens=True,
        ).strip()

        del output_ids
        del input_ids

        # Regex clean as a last-resort safety net (should rarely trigger now)
        result = self._clean_output(result)
        # Strip any lone trailing bracket left by model
        result = re.sub(r'\s*[\(\[]\s*$', '', result).strip()

        # --- Build negative prompt ---
        neg_prompt = _build_negative_prompt(result, user_input)

        if not keep_model_loaded:
            self.unload_model()

        return (result, result, neg_prompt)


# ── ComfyUI boilerplate ──────────────────────────────────────────────────────

class LTX2UnloadModel:
    """Utility node to manually free VRAM when keep_model_loaded is True."""

    @classmethod
    def INPUT_TYPES(s):
        return {"required": {}}

    RETURN_TYPES = ()
    FUNCTION = "unload"
    CATEGORY = "LTX2"
    OUTPUT_NODE = True

    def unload(self):
        # Walk all live LTX2PromptArchitect instances and unload them
        import gc
        unloaded = 0
        for obj in gc.get_objects():
            if isinstance(obj, LTX2PromptArchitect) and obj.model is not None:
                obj.unload_model()
                unloaded += 1
        print(f"[LTX2] Unload node: freed {unloaded} model instance(s).")
        return {}


NODE_CLASS_MAPPINGS = {
    "LTX2PromptArchitect": LTX2PromptArchitect,
    "LTX2UnloadModel":     LTX2UnloadModel,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LTX2PromptArchitect": "LTX-2 Easy Prompt By LoRa-Daddy",
    "LTX2UnloadModel":     "LTX2 Unload Model",
}