Upload 3 files

RebelsPromptEnhancer/__init__.py

@@ -0,0 +1,979 @@
+import torch
+import gc
+import os
+import re
+import base64
+from io import BytesIO
+
+import numpy as np
+from PIL import Image
+from llama_cpp import Llama
+
+# Vision chat handlers — best-effort import.
+_HANDLER_CLASSES = {}
+for _label, _modname in [
+    ("LLaVA 1.5", "Llava15ChatHandler"),
+    ("LLaVA 1.6", "Llava16ChatHandler"),
+    ("Moondream", "MoondreamChatHandler"),
+    ("MiniCPM-V 2.6", "MiniCPMv26ChatHandler"),
+    ("NanoLLaVA", "NanoLlavaChatHandler"),
+    ("Qwen2.5-VL", "Qwen25VLChatHandler"),
+]:
+    try:
+        _mod = __import__("llama_cpp.llama_chat_format", fromlist=[_modname])
+        _HANDLER_CLASSES[_label] = getattr(_mod, _modname)
+    except Exception:
+        pass
+
+
+# =========================================================================
+# Shared cleaning helpers
+# =========================================================================
+
+REASONING_MARKERS = (
+    "let me", "i'll ", "i will ", "i need", "i must", "i should",
+    "the prompt is", "the user", "key elements", "brainstorm",
+    "as per the rules", "according to the rules", "the rules say",
+    "let's", "okay so", "first,", "second,", "third,",
+    "i can say", "i might", "since this", "so i",
+)
+
+PREAMBLES = (
+    "here's", "here is", "sure,", "sure!", "certainly,", "of course,",
+    "okay,", "okay.", "alright,",
+    "enhanced prompt:", "expanded prompt:", "prompt:",
+    "output:", "answer:", "final prompt:", "final:", "example:",
+    "the image shows", "the image depicts", "this image shows",
+    "in the image", "i can see", "i see",
+)
+
+
+def _strip_thinking_tags(text):
+    for pat in (
+        r"<think(?:ing)?>.*?</think(?:ing)?>",
+        r"<\|thinking\|>.*?<\|/thinking\|>",
+        r"\[THINK(?:ING)?\].*?\[/THINK(?:ING)?\]",
+    ):
+        text = re.sub(pat, "", text, flags=re.DOTALL | re.IGNORECASE)
+    return text
+
+
+def _extract_final_paragraph(text):
+    paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
+    if len(paragraphs) < 2:
+        return text
+    full_lower = text.lower()
+    if sum(1 for m in REASONING_MARKERS if m in full_lower) < 2:
+        return text
+    for p in reversed(paragraphs):
+        lower = p.lower().lstrip()
+        if any(lower.startswith(m) for m in REASONING_MARKERS):
+            continue
+        lines = p.split("\n")
+        bullets = sum(
+            1 for l in lines
+            if l.strip().startswith(("-", "*", "•", "1.", "2.", "3."))
+        )
+        if bullets > 0 and bullets >= len(lines) / 2:
+            continue
+        if len(p) < 60:
+            continue
+        if sum(1 for m in REASONING_MARKERS if m in lower) >= 2:
+            continue
+        return p
+    return text
+
+
+def _strip_preambles(text):
+    for _ in range(3):
+        lowered = text.lower().lstrip()
+        stripped = False
+        for p in PREAMBLES:
+            if lowered.startswith(p):
+                colon = text.find(":")
+                newline = text.find("\n")
+                if 0 < colon < 60:
+                    text = text[colon + 1:].strip()
+                elif 0 < newline < 80:
+                    text = text[newline + 1:].strip()
+                else:
+                    text = text[len(p):].strip(" ,.:-")
+                stripped = True
+                break
+        if not stripped:
+            break
+    return text
+
+
+def _clean_output(text, original_input=""):
+    text = text.strip()
+    text = _strip_thinking_tags(text)
+    text = _extract_final_paragraph(text)
+    text = _strip_preambles(text)
+    if original_input:
+        raw = original_input.strip()
+        if raw and text.lower().startswith(raw.lower()):
+            text = text[len(raw):].lstrip(" ,.:-\"'\n")
+    return text.strip().strip('"\'')
+
+
+def _free_llm(llm):
+    try:
+        del llm
+    except Exception:
+        pass
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def _list_ggufs(exclude_mmproj=True):
+    try:
+        node_dir = os.path.dirname(os.path.abspath(__file__))
+        files = []
+        for f in os.listdir(node_dir):
+            if not f.lower().endswith(".gguf"):
+                continue
+            if exclude_mmproj and f.lower().startswith("mmproj"):
+                continue
+            files.append(f)
+        files.sort()
+        return files if files else ["NO_GGUF_FILES_IN_FOLDER"]
+    except Exception:
+        return ["NO_GGUF_FILES_IN_FOLDER"]
+
+
+def _list_mmproj():
+    try:
+        node_dir = os.path.dirname(os.path.abspath(__file__))
+        files = sorted(
+            f for f in os.listdir(node_dir)
+            if f.lower().endswith(".gguf") and f.lower().startswith("mmproj")
+        )
+        return files if files else ["NO_MMPROJ_FILE_FOUND"]
+    except Exception:
+        return ["NO_MMPROJ_FILE_FOUND"]
+
+
+# =========================================================================
+# Style system — Purpose + Model Format + Aesthetic
+# =========================================================================
+
+PURPOSE_OPTIONS = ["Image", "Video", "Edit (Inpainting/I2V)"]
+
+PURPOSE_FRAMING = {
+    "Image": (
+        "Rewrite the user's input as a detailed prompt for a static image. "
+        "Cover subject, environment, lighting, composition, and mood."
+    ),
+    "Video": (
+        "Rewrite the user's input as a detailed prompt for a video shot. "
+        "Cover subject motion, camera movement, pacing, lighting, and mood."
+    ),
+    "Edit (Inpainting/I2V)": (
+        "Rewrite the user's edit instruction as a description of the final transformed "
+        "scene as it appears after the edit. Do not describe the editing process."
+    ),
+}
+
+MODEL_FORMAT_OPTIONS = [
+    "Flux / Chroma (natural language)",
+    "Z-Image / Lumina-2 (LLM text encoder)",
+    "HiDream (hybrid prose + descriptors)",
+    "SDXL (tags + weights)",
+    "SD 1.5 (tags + weights)",
+    "Pony / Illustrious (booru tags + score)",
+    "LTX Video (motion-focused prose)",
+    "Hunyuan / Wan Video (cinematic motion prose)",
+    "Universal Natural Language",
+]
+
+MODEL_FORMAT_INSTRUCTIONS = {
+    "Flux / Chroma (natural language)": (
+        "Format the output as flowing natural language in long descriptive sentences "
+        "suitable for Flux or Chroma. No tag syntax. No parenthesis weights. "
+        "No 'masterpiece' or quality boosters. Describe the scene as prose."
+    ),
+    "Z-Image / Lumina-2 (LLM text encoder)": (
+        "Format the output as long, richly descriptive natural-language prose suitable "
+        "for LLM-based text encoders (Z-Image, Lumina-2). Use complex sentences and "
+        "vivid concrete detail. No tag syntax, no weights."
+    ),
+    "HiDream (hybrid prose + descriptors)": (
+        "Format the output as flowing natural language with concrete photographic and "
+        "material descriptors woven in (lighting type, lens feel, surface texture). "
+        "Suitable for HiDream's multi-encoder pipeline. No weight syntax."
+    ),
+    "SDXL (tags + weights)": (
+        "Format the output as a comma-separated list of descriptive tags for SDXL. "
+        "Start with quality tags (masterpiece, best quality, highly detailed). "
+        "Use parenthesis weight syntax for emphasis like (cinematic lighting:1.2). "
+        "Order: subject, action, environment, lighting, camera, style, quality."
+    ),
+    "SD 1.5 (tags + weights)": (
+        "Format the output as a comma-separated list of descriptive tags for SD 1.5. "
+        "Lead with quality boosters (masterpiece, best quality, ultra-detailed). "
+        "Use parenthesis weight syntax for emphasis. Keep tags compact and direct. "
+        "Order: subject, action, environment, lighting, style, quality."
+    ),
+    "Pony / Illustrious (booru tags + score)": (
+        "Format the output as comma-separated booru-style tags for Pony / Illustrious. "
+        "Lead with score tags: score_9, score_8_up, score_7_up. Include an appropriate "
+        "rating tag (rating_safe, rating_questionable, rating_explicit). Use underscores "
+        "for multi-word tags (long_hair, blue_eyes). "
+        "Order: score, rating, subject, character traits, action, setting, style."
+    ),
+    "LTX Video (motion-focused prose)": (
+        "Format the output as a video shot description for LTX Video. "
+        "Lead with a clear shot description, then describe subject motion explicitly "
+        "(what moves and how), camera movement (pan, dolly, zoom, tracking), pacing, "
+        "and atmosphere. Natural language sentences. No tag syntax."
+    ),
+    "Hunyuan / Wan Video (cinematic motion prose)": (
+        "Format the output as a cinematic video description with detailed motion. "
+        "Cover camera angle, camera movement, subject action, environmental motion, "
+        "atmosphere, and pacing. Natural language sentences. No tag syntax."
+    ),
+    "Universal Natural Language": (
+        "Format the output as a flowing natural language paragraph describing the scene "
+        "in concrete visual detail. No tags, no weights."
+    ),
+}
+
+AESTHETIC_OPTIONS = [
+    "None (no aesthetic injection)",
+    "Photorealistic",
+    "Cinematic Film",
+    "Anime / Manga",
+    "Studio Ghibli",
+    "Pixar / 3D Animation",
+    "Comic Book / Graphic Novel",
+    "Concept Art",
+    "Oil Painting",
+    "Watercolor",
+    "Pencil Sketch",
+    "Cyberpunk",
+    "Steampunk",
+    "Fantasy",
+    "Sci-Fi",
+    "Horror / Dark",
+    "Vintage / Retro Film",
+    "Film Noir",
+    "Glamour / Editorial",
+    "Minimalist",
+    "Surreal / Dreamy",
+    "3D Render / CGI",
+]
+
+AESTHETIC_DESCRIPTORS = {
+    "Photorealistic": (
+        "Visual style: photorealistic, sharp focus, accurate textures, lifelike skin and "
+        "materials, realistic lighting and shadows, shallow depth of field, 8k camera detail."
+    ),
+    "Cinematic Film": (
+        "Visual style: cinematic film aesthetic, dramatic lighting with strong key and rim, "
+        "filmic color grade, anamorphic framing, atmospheric haze, shallow depth of field."
+    ),
+    "Anime / Manga": (
+        "Visual style: anime/manga aesthetic, cel-shaded coloring, stylized features, "
+        "expressive large eyes, dynamic poses, vibrant saturated color palette, clean lineart."
+    ),
+    "Studio Ghibli": (
+        "Visual style: Studio Ghibli aesthetic, hand-painted watercolor backgrounds, "
+        "soft natural lighting, warm pastoral atmosphere, gentle character designs, "
+        "nostalgic and serene mood."
+    ),
+    "Pixar / 3D Animation": (
+        "Visual style: 3D animation aesthetic in the spirit of Pixar/Disney, expressive "
+        "stylized character proportions, polished CG surfaces, warm cinematic lighting, "
+        "vibrant family-friendly color palette."
+    ),
+    "Comic Book / Graphic Novel": (
+        "Visual style: comic book aesthetic, bold ink linework, halftone or hatching shading, "
+        "dynamic poses, exaggerated proportions, saturated panel colors."
+    ),
+    "Concept Art": (
+        "Visual style: digital concept art, painterly brushwork, atmospheric perspective, "
+        "value-driven composition, loose suggestion of detail over full rendering, "
+        "professional production-art feel."
+    ),
+    "Oil Painting": (
+        "Visual style: oil painting aesthetic, visible impasto brushwork, rich color depth, "
+        "painterly textures, classical composition, warm gallery lighting."
+    ),
+    "Watercolor": (
+        "Visual style: watercolor painting aesthetic, soft translucent washes, paper texture "
+        "visible, flowing pigment bleeds, gentle edges, limited palette."
+    ),
+    "Pencil Sketch": (
+        "Visual style: pencil sketch aesthetic, graphite linework, crosshatching shadows, "
+        "monochrome or restrained color accents, loose unfinished sketchbook feel."
+    ),
+    "Cyberpunk": (
+        "Visual style: cyberpunk aesthetic, neon signage, rain-slicked streets, holographic "
+        "interfaces, cybernetic implants, dystopian high-tech low-life atmosphere, "
+        "magenta-cyan color palette, deep shadows with glowing accents."
+    ),
+    "Steampunk": (
+        "Visual style: steampunk aesthetic, brass and copper machinery, Victorian-era styling, "
+        "exposed gears and clockwork, gas-lamp lighting, sepia and bronze color palette, "
+        "steam and soot atmosphere."
+    ),
+    "Fantasy": (
+        "Visual style: high fantasy aesthetic, medieval or magical setting, ethereal lighting, "
+        "ornate detail, mythological elements, lush detailed environments, painterly atmosphere."
+    ),
+    "Sci-Fi": (
+        "Visual style: science fiction aesthetic, advanced technology, sleek futuristic surfaces, "
+        "panel-screen lighting, industrial design, blue and white accent palette."
+    ),
+    "Horror / Dark": (
+        "Visual style: horror aesthetic, low-key dramatic lighting, deep shadows, "
+        "unsettling atmosphere, desaturated muted palette with occasional blood-red accents, "
+        "dread-filled mood."
+    ),
+    "Vintage / Retro Film": (
+        "Visual style: vintage film aesthetic, 35mm grain, faded warm color cast, soft contrast, "
+        "period-appropriate styling, light leaks and lens artifacts."
+    ),
+    "Film Noir": (
+        "Visual style: film noir aesthetic, high-contrast black-and-white or near-monochrome, "
+        "dramatic chiaroscuro lighting, venetian-blind shadows, urban night atmosphere, "
+        "cigarette smoke and rain."
+    ),
+    "Glamour / Editorial": (
+        "Visual style: high-fashion editorial aesthetic, polished beauty lighting, "
+        "magazine-shoot composition, soft skin rendering, dramatic backdrop, "
+        "professional studio styling."
+    ),
+    "Minimalist": (
+        "Visual style: minimalist aesthetic, clean composition, generous negative space, "
+        "limited color palette, simple geometric forms, calm uncluttered framing."
+    ),
+    "Surreal / Dreamy": (
+        "Visual style: surreal dreamlike aesthetic, soft hazy lighting, impossible compositions, "
+        "dream-logic juxtapositions, ethereal color shifts, painterly atmosphere."
+    ),
+    "3D Render / CGI": (
+        "Visual style: 3D render aesthetic, clean CGI surfaces, ray-traced lighting, "
+        "sharp reflections, subsurface scattering, Octane/Blender render quality."
+    ),
+}
+
+
+def build_layered_system_prompt(purpose, model_format, aesthetic,
+                                 extra_instructions="", append_no_think=False):
+    """Stitch purpose + aesthetic + model format into a single system prompt."""
+    parts = [PURPOSE_FRAMING[purpose]]
+    if aesthetic in AESTHETIC_DESCRIPTORS:
+        parts.append(AESTHETIC_DESCRIPTORS[aesthetic])
+    parts.append(MODEL_FORMAT_INSTRUCTIONS[model_format])
+    if extra_instructions.strip():
+        parts.append(extra_instructions.strip())
+    parts.append("Output only the prompt, nothing else.")
+    text = " ".join(parts)
+    if append_no_think:
+        text = text.rstrip() + " /no_think"
+    return text
+
+
+# =========================================================================
+# Node 1: Curated Qwen3.5-4B enhancer
+# =========================================================================
+
+class RebelsPromptEnhancer:
+    def __init__(self):
+        pass
+
+    _cache = {}
+
+    PRECISION_OPTIONS = ["Efficiency (UD-IQ2)", "Quality (UD-Q8)"]
+    PRECISION_SEARCH = {
+        "Efficiency (UD-IQ2)": ["qwen3.5-4b", "ud-iq2"],
+        "Quality (UD-Q8)":     ["qwen3.5-4b", "ud-q8"],
+    }
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "raw_prompt": ("STRING", {"multiline": True}),
+                "purpose": (PURPOSE_OPTIONS,),
+                "model_format": (MODEL_FORMAT_OPTIONS,),
+                "aesthetic": (AESTHETIC_OPTIONS,),
+                "precision": (s.PRECISION_OPTIONS,),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "lock_in": ("BOOLEAN", {
+                    "default": False,
+                    "label_on": "🔒 LOCKED (cached)",
+                    "label_off": "🔄 LIVE (generating)",
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("STRING", "STRING")
+    RETURN_NAMES = ("enhanced_prompt", "thought_process")
+    FUNCTION = "enhance"
+    CATEGORY = "Rebel AI"
+
+    @classmethod
+    def IS_CHANGED(cls, raw_prompt, purpose, model_format, aesthetic, precision, seed, lock_in):
+        if lock_in:
+            return f"LOCKED|{raw_prompt}|{purpose}|{model_format}|{aesthetic}|{precision}"
+        return float("nan")
+
+    def _find_files(self, node_dir, terms):
+        terms = [t.lower() for t in terms]
+        return [
+            f for f in os.listdir(node_dir)
+            if f.lower().endswith(".gguf") and all(t in f.lower() for t in terms)
+        ]
+
+    def enhance(self, raw_prompt, purpose, model_format, aesthetic, precision, seed, lock_in):
+        cache_key = (raw_prompt, purpose, model_format, aesthetic, precision)
+
+        if lock_in and cache_key in self._cache:
+            c = self._cache[cache_key]
+            thought = (
+                f"=== 🔒 LOCKED — Returning Cached Output ===\n"
+                f"No model load. No VRAM used. Seed ignored.\n\n"
+                f"=== Original Generation Info ===\n{c['meta']}\n"
+                f"=== Cached System Prompt ===\n{c['sys_prompt']}\n\n"
+                f"=== Cached Raw Output ===\n{c['raw_output']}\n\n"
+                f"=== Cached Final Prompt ===\n{c['final_prompt']}"
+            )
+            return (c['final_prompt'], thought)
+
+        node_dir = os.path.dirname(os.path.abspath(__file__))
+        terms = self.PRECISION_SEARCH[precision]
+        files = self._find_files(node_dir, terms)
+        if not files:
+            raise FileNotFoundError(
+                f"No .gguf in {node_dir} matching {terms}.\n"
+                f"Expected a filename containing: {' AND '.join(terms)}"
+            )
+        files.sort(key=lambda x: (len(x), x))
+        model_file = files[0]
+        model_path = os.path.join(node_dir, model_file)
+        n_ctx = 4096
+
+        sys_prompt = build_layered_system_prompt(
+            purpose, model_format, aesthetic, append_no_think=True
+        )
+
+        llm = Llama(model_path=model_path, n_gpu_layers=-1, verbose=False,
+                    n_ctx=n_ctx, seed=seed)
+        try:
+            output = llm.create_chat_completion(
+                messages=[
+                    {"role": "system", "content": sys_prompt},
+                    {"role": "user", "content": raw_prompt},
+                ],
+                max_tokens=500, temperature=0.7, top_p=0.9, repeat_penalty=1.15,
+                stop=["\n\nUser:", "\n\nAssistant:", "Human:", "</think>", "</thinking>"],
+            )
+            raw_output = output["choices"][0]["message"]["content"].strip()
+        finally:
+            _free_llm(llm)
+
+        final_prompt = _clean_output(raw_output, raw_prompt)
+        if not final_prompt.strip() or len(final_prompt) < 30:
+            final_prompt = raw_output
+
+        meta_block = (
+            f"File:         {model_file}\n"
+            f"Precision:    {precision}\n"
+            f"Purpose:      {purpose}\n"
+            f"Model Format: {model_format}\n"
+            f"Aesthetic:    {aesthetic}\n"
+            f"Context:      {n_ctx}\n"
+            f"Seed:         {seed}\n"
+            f"Raw chars:    {len(raw_output)}\n"
+            f"Clean chars:  {len(final_prompt)}\n"
+            f"Stripped:     {len(raw_output) - len(final_prompt)} chars\n"
+        )
+        thought = (
+            f"=== 🔄 LIVE — Fresh Generation ===\n\n"
+            f"=== Run Info ===\n{meta_block}\n"
+            f"=== Assembled System Prompt ===\n{sys_prompt}\n\n"
+            f"=== Raw Model Output ===\n{raw_output}\n\n"
+            f"=== Final Prompt ===\n{final_prompt}"
+        )
+        self._cache[cache_key] = {
+            "final_prompt": final_prompt, "raw_output": raw_output,
+            "meta": meta_block, "sys_prompt": sys_prompt,
+        }
+        return (final_prompt, thought)
+
+
+# =========================================================================
+# Node 2: Custom GGUF enhancer (any text model)
+# =========================================================================
+
+class RebelsPromptEnhancerCustom:
+    def __init__(self):
+        pass
+
+    _cache = {}
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "raw_prompt": ("STRING", {"multiline": True}),
+                "model_file": (_list_ggufs(),),
+                "purpose": (PURPOSE_OPTIONS,),
+                "model_format": (MODEL_FORMAT_OPTIONS,),
+                "aesthetic": (AESTHETIC_OPTIONS,),
+                "extra_instructions": ("STRING", {
+                    "multiline": True,
+                    "default": "",
+                    "placeholder": "Optional extra instructions appended to the system prompt (e.g. 'Avoid clichés.' or 'Emphasize hands.')",
+                }),
+                "system_prompt_override": ("STRING", {
+                    "multiline": True,
+                    "default": "",
+                    "placeholder": "If non-empty, this REPLACES the entire layered system prompt. Leave blank to use Purpose+Format+Aesthetic.",
+                }),
+                "append_no_think": ("BOOLEAN", {
+                    "default": False,
+                    "label_on": "Append /no_think",
+                    "label_off": "Don't append",
+                }),
+                "n_gpu_layers": ("INT", {"default": -1, "min": -1, "max": 999, "step": 1}),
+                "n_ctx": ("INT", {"default": 4096, "min": 512, "max": 32768, "step": 512}),
+                "max_tokens": ("INT", {"default": 500, "min": 50, "max": 4096, "step": 50}),
+                "temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.05}),
+                "top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.05}),
+                "repeat_penalty": ("FLOAT", {"default": 1.15, "min": 1.0, "max": 2.0, "step": 0.05}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "lock_in": ("BOOLEAN", {
+                    "default": False,
+                    "label_on": "🔒 LOCKED (cached)",
+                    "label_off": "🔄 LIVE (generating)",
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("STRING", "STRING")
+    RETURN_NAMES = ("enhanced_prompt", "thought_process")
+    FUNCTION = "enhance"
+    CATEGORY = "Rebel AI"
+
+    @classmethod
+    def IS_CHANGED(cls, raw_prompt, model_file, purpose, model_format, aesthetic,
+                   extra_instructions, system_prompt_override, append_no_think,
+                   n_gpu_layers, n_ctx, max_tokens, temperature, top_p,
+                   repeat_penalty, seed, lock_in):
+        if lock_in:
+            return (
+                f"LOCKED|{raw_prompt}|{model_file}|{purpose}|{model_format}|{aesthetic}|"
+                f"{extra_instructions}|{system_prompt_override}|{append_no_think}|"
+                f"{temperature}|{top_p}|{repeat_penalty}"
+            )
+        return float("nan")
+
+    def _build_sys_prompt(self, purpose, model_format, aesthetic,
+                          extra_instructions, override, append_no_think):
+        if override.strip():
+            base = override.strip()
+            if append_no_think:
+                base = base.rstrip() + " /no_think"
+            return base
+        return build_layered_system_prompt(
+            purpose, model_format, aesthetic,
+            extra_instructions=extra_instructions,
+            append_no_think=append_no_think,
+        )
+
+    def enhance(self, raw_prompt, model_file, purpose, model_format, aesthetic,
+                extra_instructions, system_prompt_override, append_no_think,
+                n_gpu_layers, n_ctx, max_tokens, temperature, top_p,
+                repeat_penalty, seed, lock_in):
+        cache_key = (
+            raw_prompt, model_file, purpose, model_format, aesthetic,
+            extra_instructions, system_prompt_override, append_no_think,
+            temperature, top_p, repeat_penalty,
+        )
+
+        if lock_in and cache_key in self._cache:
+            c = self._cache[cache_key]
+            thought = (
+                f"=== 🔒 LOCKED — Returning Cached Output ===\n"
+                f"No model load. No VRAM used. Seed ignored.\n\n"
+                f"=== Original Generation Info ===\n{c['meta']}\n"
+                f"=== Cached System Prompt ===\n{c['sys_prompt']}\n\n"
+                f"=== Cached Raw Output ===\n{c['raw_output']}\n\n"
+                f"=== Cached Final Prompt ===\n{c['final_prompt']}"
+            )
+            return (c['final_prompt'], thought)
+
+        if model_file == "NO_GGUF_FILES_IN_FOLDER":
+            raise FileNotFoundError("Drop a .gguf in the node folder and restart ComfyUI.")
+        node_dir = os.path.dirname(os.path.abspath(__file__))
+        model_path = os.path.join(node_dir, model_file)
+        if not os.path.isfile(model_path):
+            raise FileNotFoundError(f"Selected file not found: {model_path}. Restart ComfyUI to refresh.")
+
+        sys_prompt = self._build_sys_prompt(
+            purpose, model_format, aesthetic,
+            extra_instructions, system_prompt_override, append_no_think,
+        )
+
+        llm = Llama(model_path=model_path, n_gpu_layers=n_gpu_layers, verbose=False,
+                    n_ctx=n_ctx, seed=seed)
+        try:
+            output = llm.create_chat_completion(
+                messages=[
+                    {"role": "system", "content": sys_prompt},
+                    {"role": "user", "content": raw_prompt},
+                ],
+                max_tokens=max_tokens, temperature=temperature, top_p=top_p,
+                repeat_penalty=repeat_penalty,
+                stop=["\n\nUser:", "\n\nAssistant:", "Human:", "</think>", "</thinking>"],
+            )
+            raw_output = output["choices"][0]["message"]["content"].strip()
+        finally:
+            _free_llm(llm)
+
+        final_prompt = _clean_output(raw_output, raw_prompt)
+        if not final_prompt.strip() or len(final_prompt) < 30:
+            final_prompt = raw_output
+
+        gpu_label = "all" if n_gpu_layers < 0 else ("CPU only" if n_gpu_layers == 0 else f"{n_gpu_layers} layers")
+        override_status = "ACTIVE (custom override used)" if system_prompt_override.strip() else "inactive"
+        meta_block = (
+            f"File:           {model_file}\n"
+            f"Purpose:        {purpose}\n"
+            f"Model Format:   {model_format}\n"
+            f"Aesthetic:      {aesthetic}\n"
+            f"Override:       {override_status}\n"
+            f"GPU layers:     {gpu_label}\n"
+            f"Context:        {n_ctx}\n"
+            f"Max tokens:     {max_tokens}\n"
+            f"Temperature:    {temperature}\n"
+            f"top_p:          {top_p}\n"
+            f"repeat_penalty: {repeat_penalty}\n"
+            f"/no_think:      {'yes' if append_no_think else 'no'}\n"
+            f"Seed:           {seed}\n"
+            f"Raw chars:      {len(raw_output)}\n"
+            f"Clean chars:    {len(final_prompt)}\n"
+            f"Stripped:       {len(raw_output) - len(final_prompt)} chars\n"
+        )
+        thought = (
+            f"=== 🔄 LIVE — Fresh Generation (Custom) ===\n\n"
+            f"=== Run Settings ===\n{meta_block}\n"
+            f"=== Assembled System Prompt ===\n{sys_prompt}\n\n"
+            f"=== Raw Model Output ===\n{raw_output}\n\n"
+            f"=== Final Prompt ===\n{final_prompt}"
+        )
+        self._cache[cache_key] = {
+            "final_prompt": final_prompt, "raw_output": raw_output,
+            "meta": meta_block, "sys_prompt": sys_prompt,
+        }
+        return (final_prompt, thought)
+
+
+# =========================================================================
+# Node 3: Image-to-Prompt vision node
+# =========================================================================
+
+class RebelsImageToPrompt:
+    def __init__(self):
+        pass
+
+    _cache = {}
+
+    VISION_TASK_OPTIONS = [
+        "Caption (plain description)",
+        "Caption + Format (apply model_format below)",
+        "SD/Booru Tags",
+        "Pose & Anatomy Focus",
+        "Custom Instruction",
+    ]
+
+    VISION_TASK_INSTRUCTIONS = {
+        "Caption (plain description)": (
+            "Describe this image in one detailed paragraph covering subject, composition, "
+            "lighting, colors, mood, and notable details. Output only the description."
+        ),
+        "SD/Booru Tags": (
+            "Generate a comma-separated list of descriptive tags for this image. "
+            "Include subject, action, setting, lighting, mood, and style tags. "
+            "Output only the tags, comma-separated."
+        ),
+        "Pose & Anatomy Focus": (
+            "Describe the subject's pose, body position, expression, framing, and what's "
+            "visible in detail. Be precise about positioning. Output only the description."
+        ),
+    }
+
+    @classmethod
+    def INPUT_TYPES(s):
+        handler_options = ["Auto-detect"] + list(_HANDLER_CLASSES.keys())
+        if not _HANDLER_CLASSES:
+            handler_options = ["NO_VISION_HANDLERS_AVAILABLE"]
+
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "model_file": (_list_ggufs(),),
+                "mmproj_file": (_list_mmproj(),),
+                "chat_handler": (handler_options,),
+                "vision_task": (s.VISION_TASK_OPTIONS,),
+                "model_format": (MODEL_FORMAT_OPTIONS,),
+                "aesthetic": (AESTHETIC_OPTIONS,),
+                "custom_instruction": ("STRING", {
+                    "multiline": True,
+                    "default": "",
+                    "placeholder": "Used when vision_task is 'Custom Instruction'.",
+                }),
+                "n_gpu_layers": ("INT", {"default": -1, "min": -1, "max": 999, "step": 1}),
+                "n_ctx": ("INT", {"default": 4096, "min": 512, "max": 32768, "step": 512}),
+                "max_tokens": ("INT", {"default": 400, "min": 50, "max": 2048, "step": 50}),
+                "temperature": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 2.0, "step": 0.05}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "lock_in": ("BOOLEAN", {
+                    "default": False,
+                    "label_on": "🔒 LOCKED (cached)",
+                    "label_off": "🔄 LIVE (analyzing)",
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("STRING", "STRING")
+    RETURN_NAMES = ("image_prompt", "thought_process")
+    FUNCTION = "describe"
+    CATEGORY = "Rebel AI"
+
+    @classmethod
+    def IS_CHANGED(cls, image, model_file, mmproj_file, chat_handler, vision_task,
+                   model_format, aesthetic, custom_instruction, n_gpu_layers,
+                   n_ctx, max_tokens, temperature, seed, lock_in):
+        if lock_in:
+            return (f"LOCKED|{model_file}|{mmproj_file}|{chat_handler}|{vision_task}|"
+                    f"{model_format}|{aesthetic}|{custom_instruction}")
+        try:
+            img_hash = hash(image.detach().cpu().numpy().tobytes())
+        except Exception:
+            img_hash = "noimg"
+        return f"LIVE|{img_hash}|{seed}|{float('nan')}"
+
+    def _resolve_handler(self, chat_handler_choice, model_file):
+        if chat_handler_choice == "Auto-detect":
+            mf = model_file.lower()
+            if "moondream" in mf and "Moondream" in _HANDLER_CLASSES:
+                return _HANDLER_CLASSES["Moondream"], "Moondream"
+            if "qwen" in mf and "Qwen2.5-VL" in _HANDLER_CLASSES:
+                return _HANDLER_CLASSES["Qwen2.5-VL"], "Qwen2.5-VL"
+            if "minicpm" in mf and "MiniCPM-V 2.6" in _HANDLER_CLASSES:
+                return _HANDLER_CLASSES["MiniCPM-V 2.6"], "MiniCPM-V 2.6"
+            if "nano" in mf and "llava" in mf and "NanoLLaVA" in _HANDLER_CLASSES:
+                return _HANDLER_CLASSES["NanoLLaVA"], "NanoLLaVA"
+            if "llava" in mf:
+                if "LLaVA 1.6" in _HANDLER_CLASSES:
+                    return _HANDLER_CLASSES["LLaVA 1.6"], "LLaVA 1.6"
+                if "LLaVA 1.5" in _HANDLER_CLASSES:
+                    return _HANDLER_CLASSES["LLaVA 1.5"], "LLaVA 1.5"
+            if "LLaVA 1.5" in _HANDLER_CLASSES:
+                return _HANDLER_CLASSES["LLaVA 1.5"], "LLaVA 1.5 (fallback)"
+            raise RuntimeError("No vision chat handler available.")
+        if chat_handler_choice not in _HANDLER_CLASSES:
+            raise RuntimeError(
+                f"Chat handler '{chat_handler_choice}' isn't available. "
+                f"Available: {list(_HANDLER_CLASSES.keys())}"
+            )
+        return _HANDLER_CLASSES[chat_handler_choice], chat_handler_choice
+
+    def _tensor_to_data_uri(self, image_tensor):
+        if image_tensor.dim() == 4:
+            img = image_tensor[0]
+        else:
+            img = image_tensor
+        img_np = (img.detach().cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
+        pil = Image.fromarray(img_np)
+        buf = BytesIO()
+        pil.save(buf, format="PNG")
+        b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+        return f"data:image/png;base64,{b64}"
+
+    def _build_instruction(self, vision_task, model_format, aesthetic, custom_instruction):
+        if vision_task == "Custom Instruction":
+            return custom_instruction.strip() or "Describe this image."
+
+        if vision_task == "Caption + Format (apply model_format below)":
+            base = (
+                "Describe this image faithfully, then format the description according "
+                "to the rules below."
+            )
+            base += " " + MODEL_FORMAT_INSTRUCTIONS[model_format]
+            if aesthetic in AESTHETIC_DESCRIPTORS:
+                base += " " + AESTHETIC_DESCRIPTORS[aesthetic]
+            base += " Output only the formatted prompt."
+            return base
+
+        base = self.VISION_TASK_INSTRUCTIONS[vision_task]
+        if aesthetic in AESTHETIC_DESCRIPTORS:
+            base += " " + AESTHETIC_DESCRIPTORS[aesthetic]
+        return base
+
+    def describe(self, image, model_file, mmproj_file, chat_handler, vision_task,
+                 model_format, aesthetic, custom_instruction, n_gpu_layers,
+                 n_ctx, max_tokens, temperature, seed, lock_in):
+        cache_key = (model_file, mmproj_file, chat_handler, vision_task,
+                     model_format, aesthetic, custom_instruction, temperature)
+
+        if lock_in and cache_key in self._cache:
+            c = self._cache[cache_key]
+            thought = (
+                f"=== 🔒 LOCKED — Returning Cached Caption ===\n"
+                f"No model load. New image ignored.\n\n"
+                f"=== Original Run Info ===\n{c['meta']}\n"
+                f"=== Cached Instruction ===\n{c['instruction']}\n\n"
+                f"=== Cached Raw Output ===\n{c['raw_output']}\n\n"
+                f"=== Cached Image Prompt ===\n{c['final_prompt']}"
+            )
+            return (c['final_prompt'], thought)
+
+        if model_file == "NO_GGUF_FILES_IN_FOLDER":
+            raise FileNotFoundError("Drop a vision-capable .gguf in the node folder and restart ComfyUI.")
+        if mmproj_file == "NO_MMPROJ_FILE_FOUND":
+            raise FileNotFoundError("No mmproj file found. Vision models need a paired mmproj-*.gguf.")
+        if not _HANDLER_CLASSES:
+            raise RuntimeError("No vision chat handlers available. Update llama-cpp-python.")
+
+        node_dir = os.path.dirname(os.path.abspath(__file__))
+        model_path = os.path.join(node_dir, model_file)
+        mmproj_path = os.path.join(node_dir, mmproj_file)
+        if not os.path.isfile(model_path):
+            raise FileNotFoundError(f"Model not found: {model_path}")
+        if not os.path.isfile(mmproj_path):
+            raise FileNotFoundError(f"mmproj not found: {mmproj_path}")
+
+        handler_cls, handler_label = self._resolve_handler(chat_handler, model_file)
+        instruction = self._build_instruction(vision_task, model_format, aesthetic, custom_instruction)
+        img_uri = self._tensor_to_data_uri(image)
+
+        chat_handler_instance = handler_cls(clip_model_path=mmproj_path, verbose=False)
+        llm = Llama(
+            model_path=model_path,
+            chat_handler=chat_handler_instance,
+            n_gpu_layers=n_gpu_layers,
+            verbose=False,
+            n_ctx=n_ctx,
+            seed=seed,
+            logits_all=True,
+        )
+        try:
+            output = llm.create_chat_completion(
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "image_url", "image_url": {"url": img_uri}},
+                            {"type": "text", "text": instruction},
+                        ],
+                    },
+                ],
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=0.9,
+                repeat_penalty=1.1,
+            )
+            raw_output = output["choices"][0]["message"]["content"].strip()
+        finally:
+            _free_llm(llm)
+            try:
+                del chat_handler_instance
+            except Exception:
+                pass
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+        final_prompt = _clean_output(raw_output)
+        if not final_prompt.strip() or len(final_prompt) < 20:
+            final_prompt = raw_output
+
+        gpu_label = "all" if n_gpu_layers < 0 else ("CPU only" if n_gpu_layers == 0 else f"{n_gpu_layers} layers")
+        meta_block = (
+            f"Model file:    {model_file}\n"
+            f"mmproj file:   {mmproj_file}\n"
+            f"Chat handler:  {handler_label}\n"
+            f"Vision task:   {vision_task}\n"
+            f"Model Format:  {model_format}\n"
+            f"Aesthetic:     {aesthetic}\n"
+            f"GPU layers:    {gpu_label}\n"
+            f"Context:       {n_ctx}\n"
+            f"Max tokens:    {max_tokens}\n"
+            f"Temperature:   {temperature}\n"
+            f"Seed:          {seed}\n"
+            f"Raw chars:     {len(raw_output)}\n"
+            f"Clean chars:   {len(final_prompt)}\n"
+        )
+        thought = (
+            f"=== 🔄 LIVE — Image Analysis ===\n\n"
+            f"=== Run Settings ===\n{meta_block}\n"
+            f"=== Instruction Sent ===\n{instruction}\n\n"
+            f"=== Raw Model Output ===\n{raw_output}\n\n"
+            f"=== Image Prompt (downstream) ===\n{final_prompt}"
+        )
+        self._cache[cache_key] = {
+            "final_prompt": final_prompt, "raw_output": raw_output,
+            "meta": meta_block, "instruction": instruction,
+        }
+        return (final_prompt, thought)
+
+
+# =========================================================================
+# Node 4: Locker
+# =========================================================================
+
+class RebelsPromptLocker:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "text_input": ("STRING", {"forceInput": True}),
+                "lock_in_prompt": ("BOOLEAN", {
+                    "default": False,
+                    "label_on": "LOCKED IN",
+                    "label_off": "PAUSED",
+                }),
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("text_output",)
+    FUNCTION = "execute"
+    CATEGORY = "Rebel AI"
+    OUTPUT_NODE = True
+
+    def execute(self, text_input, lock_in_prompt):
+        if not lock_in_prompt:
+            raise ValueError(
+                "🛑 WORKFLOW PAUSED: Toggle 'lock_in_prompt' to LOCKED IN to pass the text through."
+            )
+        return {"ui": {"text": [text_input]}, "result": (text_input,)}
+
+
+NODE_CLASS_MAPPINGS = {
+    "RebelsPromptEnhancer": RebelsPromptEnhancer,
+    "RebelsPromptEnhancerCustom": RebelsPromptEnhancerCustom,
+    "RebelsImageToPrompt": RebelsImageToPrompt,
+    "RebelsPromptLocker": RebelsPromptLocker,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "RebelsPromptEnhancer": "🚀 Rebels Prompt Enhancer",
+    "RebelsPromptEnhancerCustom": "🧪 Rebels Prompt Enhancer (Custom GGUF)",
+    "RebelsImageToPrompt": "👁️ Rebels Image to Prompt",
+    "RebelsPromptLocker": "🔒 Rebels Prompt Locker",
+}
+
+WEB_DIRECTORY = "./web"
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]

RebelsPromptEnhancer/requirements.txt

@@ -0,0 +1,2 @@
+llama-cpp-python
+torch

RebelsPromptEnhancer/web/js/rebels_locker_display.js

@@ -0,0 +1,46 @@
+import { app } from "../../../scripts/app.js";
+import { ComfyWidgets } from "../../../scripts/widgets.js";
+
+app.registerExtension({
+    name: "RebelAI.PromptLocker.Display",
+    async beforeRegisterNodeDef(nodeType, nodeData, app) {
+        if (nodeData.name !== "RebelsPromptLocker") return;
+
+        const onExecuted = nodeType.prototype.onExecuted;
+        nodeType.prototype.onExecuted = function (message) {
+            onExecuted?.apply(this, arguments);
+
+            // Remove any existing display widget so we don't stack them on re-runs.
+            if (this.widgets) {
+                const idx = this.widgets.findIndex(w => w.name === "locked_text_display");
+                if (idx !== -1) {
+                    this.widgets[idx].onRemove?.();
+                    this.widgets.splice(idx, 1);
+                }
+            }
+
+            const text = (message?.text || []).join("");
+
+            const widget = ComfyWidgets["STRING"](
+                this,
+                "locked_text_display",
+                ["STRING", { multiline: true }],
+                app
+            ).widget;
+
+            // Read-only, but otherwise use ComfyUI's native widget styling so it
+            // matches the size/weight/font of a normal text widget.
+            widget.inputEl.readOnly = true;
+            widget.inputEl.style.fontSize = "13px";
+            widget.inputEl.style.lineHeight = "1.4";
+            widget.value = text;
+
+            requestAnimationFrame(() => {
+                const sz = this.computeSize();
+                if (sz[1] < this.size[1]) sz[1] = this.size[1];
+                this.onResize?.(sz);
+                app.graph.setDirtyCanvas(true, false);
+            });
+        };
+    },
+});