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