diff --git a/CHANGELOG.md b/CHANGELOG.md index 84bd8da0..bfdf12a0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,29 @@ # Changelog +## [2.3.0] - 2026-05-08 + +### Added + +Added `Wrong appearance` rating to `FunPack Video Refiner V2` for outputs contaminated by remembered clothing, character, subject, or background concepts. + +Added `FunPack Refinement Key Loader`, with a selectable key dropdown, create-on-load behavior, and browser-side JSON import/export buttons. + +Added a Discord-friendly Refiner V2 quick guide for new users. + +### Changed + +Updated Refiner V2 Prompt Repair so it only auto-adds safe repair concepts such as action, camera, details, quality, and style. Appearance, subject/character, and environment/background concepts are now blocked from Prompt Repair. + +Updated `I'm Feeling Lucky` in Refiner V2 so appearance, subject/character, and environment/background memory is not auto-injected unless the user explicitly includes that phrase in the current prompt. + +Updated legacy Void/Lucky token-bank selection to skip appearance, subject/character, and environment/background tokens. + +Updated Refiner V2 and `FunPack Apply LoRA Weights` so both can accept a linked refinement key from `FunPack Refinement Key Loader`. + +### Fixed + +Fixed appearance bleed-over where highly liked clothing or character tags could reappear in unrelated image-to-video prompts. + ## [2.2.1] - 2026-05-07 ### Fixed diff --git a/README.md b/README.md index bed2b069..fbc2d780 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,16 @@ A set of ComfyUI nodes for experimenting with video generation workflows based on WAN, HunyuanVideo, LTX, and similar models. +## Updates in 2.3.0 + +Refiner V2 now blocks learned appearance, character, subject, and background concepts from Prompt Repair, Lucky auto-injection, and legacy Void memory unless the current prompt explicitly asks for them. + +Added `Wrong appearance` for outputs polluted by remembered clothing or character traits. This suppresses the responsible appearance memory without penalizing unrelated action, camera, detail, or quality learning. + +Added `FunPack Refinement Key Loader` for selecting, creating, importing, and exporting Refiner V2 keys. The loader can feed both Refiner V2 and `FunPack Apply LoRA Weights`. + +Added a paste-friendly [`Refiner V2 quick guide`](docs/FunPackVideoRefinerV2QuickGuide.md) for new users. + ## Updates in 2.2.1 Fixed Refiner V2 prompt phrase categorization so background, appearance, quality, camera, and action phrases stay aligned before the node updates Lucky memory or LoRA suggestions. @@ -68,6 +78,7 @@ Per-node documentation is available in the [`docs`](docs) folder. Start with: - [`docs/FunPackVideoRefinerV2.md`](docs/FunPackVideoRefinerV2.md) for `FunPack Video Refiner V2` +- [`docs/FunPackVideoRefinerV2QuickGuide.md`](docs/FunPackVideoRefinerV2QuickGuide.md) for a short Discord-friendly Refiner V2 guide - [`docs/FunPackLoraWorkflow.md`](docs/FunPackLoraWorkflow.md) for the LoRA/refiner helper workflow Version history is available in [CHANGELOG.md](CHANGELOG.md). diff --git a/__init__.py b/__init__.py index a48eda75..9b7dc80a 100644 --- a/__init__.py +++ b/__init__.py @@ -18,7 +18,7 @@ if __package__: from .model_management import FunPackApplyLoraWeights, FunPackLoraLoader from .samplers import FunPackHybridEuler2SSampler from .context_transition import FunPackContextTransitionWindows - from .templates import FunPackTemplateManager + from .templates import FunPackRefinementKeyLoader, FunPackTemplateManager else: # Standalone tests may not have the full ComfyUI/CUDA runtime loaded. from conditioning import ( @@ -56,8 +56,9 @@ else: except Exception: FunPackContextTransitionWindows = None try: - from templates import FunPackTemplateManager + from templates import FunPackRefinementKeyLoader, FunPackTemplateManager except Exception: + FunPackRefinementKeyLoader = None FunPackTemplateManager = None WEB_DIRECTORY = "./web" @@ -77,6 +78,7 @@ NODE_CLASS_MAPPINGS = { "FunPackContextTransitionWindows": FunPackContextTransitionWindows, "FunPackApplyLoraWeights": FunPackApplyLoraWeights, "FunPackLoraLoader": FunPackLoraLoader, + "FunPackRefinementKeyLoader": FunPackRefinementKeyLoader, "FunPackTemplateManager": FunPackTemplateManager, } NODE_CLASS_MAPPINGS = {name: cls for name, cls in NODE_CLASS_MAPPINGS.items() if cls is not None} @@ -96,6 +98,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "FunPackContextTransitionWindows": "FunPack Context Transition Windows", "FunPackApplyLoraWeights": "FunPack Apply LoRA Weights", "FunPackLoraLoader": "FunPack LoRA Loader", + "FunPackRefinementKeyLoader": "FunPack Refinement Key Loader", "FunPackTemplateManager": "FunPack Template Manager", } diff --git a/conditioning.py b/conditioning.py index be04c176..1d4a3da0 100644 --- a/conditioning.py +++ b/conditioning.py @@ -1457,10 +1457,21 @@ class FunPackVideoRefiner: return None, "" return canvas.clone(), f"{best['label']} ({best['seq_len']} positions)" + def _legacy_auto_inject_token_allowed(self, token): + token = str(token or "").strip().lower() + if not token: + return False + category = self._infer_concept_category([token]) + if category in {"appearance", "subject", "environment", "character"}: + return False + return True + def _eligible_void_bank_items(self, global_adaptive, embedding_dim): bank = self._ensure_void_token_bank(global_adaptive) eligible = [] for token, item in bank.items(): + if not self._legacy_auto_inject_token_allowed(token): + continue if int(item.get("awful_count", 0)) >= 2 and int(item.get("liked_count", 0)) <= 0: continue score = self._void_bank_score(item) @@ -1479,6 +1490,8 @@ class FunPackVideoRefiner: bank = self._ensure_void_token_bank(global_adaptive) candidates = [] for token, item in bank.items(): + if not self._legacy_auto_inject_token_allowed(token): + continue score = self._void_bank_score(item) embedding_shape = self._serialized_tensor_shape(item.get("embedding")) if embedding_shape is None or len(embedding_shape) != 1 or int(embedding_shape[0]) != int(embedding_dim): @@ -1643,6 +1656,8 @@ class FunPackVideoRefiner: bank = self._ensure_void_token_bank(global_adaptive) lucky_pool = [] for token, item in bank.items(): + if not self._legacy_auto_inject_token_allowed(token): + continue embedding_shape = self._serialized_tensor_shape(item.get("embedding")) if embedding_shape is None or len(embedding_shape) != 1 or int(embedding_shape[0]) != int(embedding_dim): continue @@ -1871,6 +1886,8 @@ class FunPackVideoRefiner: token = full_word.strip().lower() if not self._is_valuable_token(token): continue + if not self._legacy_auto_inject_token_allowed(token): + continue if token_mask_list is not None and not all(token_mask_list[start:min(len(token_mask_list), end)]): continue score = float(word_importance.get(token, 1.0)) @@ -5477,6 +5494,7 @@ V2_RATING_LABELS = [ "Wrong details", "Wrong action", "Wrong details + action", + "Wrong appearance", "Missing details + quality", "Missing action + quality", "Awful", @@ -5493,6 +5511,13 @@ V2_RATING_PROFILES = { "Wrong details": {"key": "wrong_details", "reward": 0.20, "level": 5, "missing_axes": ["details"], "wrong_axes": ["details"]}, "Wrong action": {"key": "wrong_action", "reward": 0.10, "level": 4, "missing_axes": ["action"], "wrong_axes": ["action"]}, "Wrong details + action": {"key": "wrong_details_action", "reward": 0.00, "level": 3, "missing_axes": ["details", "action"], "wrong_axes": ["details", "action"]}, + "Wrong appearance": { + "key": "wrong_appearance", + "reward": 0.0, + "level": 4, + "missing_axes": [], + "wrong_categories": ["appearance", "subject", "environment"], + }, "Missing details + quality": {"key": "missing_details_quality", "reward": -0.40, "level": 2, "missing_axes": ["details", "quality"]}, "Missing action + quality": {"key": "missing_action_quality", "reward": -0.55, "level": 1, "missing_axes": ["action", "quality"]}, "Awful": {"key": "awful", "reward": -0.90, "level": 0, "missing_axes": ["details", "action", "quality"]}, @@ -5521,6 +5546,11 @@ V2_RATING_ALIASES = { "Wrong concepts": "Wrong action", "Wrong movement": "Wrong action", "Wrong micro-movements": "Wrong details", + "Wrong character": "Wrong appearance", + "Wrong characters": "Wrong appearance", + "Wrong outfit": "Wrong appearance", + "Wrong clothing": "Wrong appearance", + "Wrong clothes": "Wrong appearance", "Missing everything": "Awful", } @@ -5566,6 +5596,8 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): V2_STATE_PREFIX = "refine_v2" ACTION_LORA_TYPES = {"action", "concept"} + AUTO_INJECT_BLOCKED_CATEGORIES = {"appearance", "subject", "environment"} + AUTO_INJECT_ALLOWED_CATEGORIES = {"action", "camera", "details", "quality", "style"} CATEGORY_DESCRIPTIONS = { "action": "physical actions, body movement, animation, motion, gestures, moving subjects", "camera": "camera motion, framing, zoom, pan, dolly, close-up, wide shot, focus behavior", @@ -5603,6 +5635,7 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): "hands", "pose", "expression", "beard", "makeup", "body", "anatomy", "wearing", "wears", "dressed", "clothed", "costume", "shirt", "coat", "robe", "boots", "hat", "helmet", "gloves", "mask", "tattoo", "freckles", "scar", "silhouette", + "pants", "skirt", "tights", "stockings", "socks", "shoes", "heels", "uniform", "flowing", "curly", "straight", "long", "short", "blonde", "brunette", }, "environment": { @@ -5650,7 +5683,8 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): APPEARANCE_CONTEXT_WORDS = { "hair", "eyes", "face", "skin", "dress", "jacket", "armor", "outfit", "clothing", "shirt", "coat", "robe", "boots", "hat", "helmet", "gloves", "mask", "body", - "anatomy", "beard", "makeup", "costume", + "anatomy", "beard", "makeup", "costume", "pants", "skirt", "tights", "stockings", + "socks", "shoes", "heels", "uniform", } ENVIRONMENT_CONTEXT_WORDS = { "background", "setting", "environment", "landscape", "room", "street", "forest", @@ -5675,6 +5709,12 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "label": "Exploration Seed"}), "reset_session": ("BOOLEAN", {"default": False, "label": "Reset V2 Session"}), "lora_stack": ("FUNPACK_LORA_STACK", {"tooltip": "Optional stack from FunPack LoRA Loader. V2 writes prompt-specific suggested weights."}), + "refinement_key_input": ("STRING", { + "default": "", + "multiline": False, + "forceInput": True, + "tooltip": "Optional linked refinement key, for example from FunPack Refinement Key Loader. Overrides the refinement_key widget when connected.", + }), "user_intent_prompt": ("STRING", { "multiline": True, "default": "", @@ -5860,6 +5900,56 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): ) return primary, float(confidence) + def _v2_primary_category_for_text(self, text): + return self._v2_scores_primary(self._v2_heuristic_scores(text))[0] + + def _v2_item_primary_category(self, item): + if not isinstance(item, dict): + return self._v2_primary_category_for_text(item) + scores = ( + item.get("effective_category_scores") or + item.get("category_scores") or + item.get("clip_heuristic_scores") or + self._v2_heuristic_scores(item.get("text", "")) + ) + primary, _ = self._v2_scores_primary(scores) + return str(item.get("primary") or primary or "details").lower() + + def _v2_item_has_blocked_auto_category(self, item): + if isinstance(item, dict): + categories = { + str(item.get("primary", "")).lower(), + str(item.get("machine_primary", "")).lower(), + self._v2_item_primary_category(item), + } + scores = item.get("effective_category_scores", item.get("category_scores", {})) + if isinstance(scores, dict): + primary, confidence = self._v2_scores_primary(scores) + categories.add(primary) + for category, score in scores.items(): + if float(score or 0.0) >= max(0.42, float(confidence) * 0.82): + categories.add(str(category).lower()) + return bool(categories & self.AUTO_INJECT_BLOCKED_CATEGORIES) + return self._v2_primary_category_for_text(item) in self.AUTO_INJECT_BLOCKED_CATEGORIES + + def _v2_prompt_repair_text_allowed(self, text): + if not self._v2_clean_phrase_text(text): + return False + if self._v2_item_has_blocked_auto_category(text): + return False + return self._v2_primary_category_for_text(text) in self.AUTO_INJECT_ALLOWED_CATEGORIES + + def _v2_auto_inject_entry_allowed(self, entry, prompt=""): + text = str(entry.get("text", "") if isinstance(entry, dict) else entry).strip() + if not text: + return False + explicit = self._v2_prompt_contains_text(prompt, text) + if isinstance(entry, dict) and bool(entry.get("auto_inject_suppressed")) and not explicit: + return False + if self._v2_item_has_blocked_auto_category(entry) and not explicit: + return False + return True + def _v2_clip_similarity_scores(self, clip, phrase, category_vectors): phrase_cond, _, _ = self._v2_encode_prompt(clip, phrase) phrase_vector = self._v2_conditioning_vector(phrase_cond) @@ -6118,6 +6208,9 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): "resolved_count": 0, "regressed_count": 0, "wrong_count": 0, + "wrong_appearance_count": 0, + "auto_inject_blocked_count": 0, + "auto_inject_suppressed": False, "bad_count": 0, "wanted_axes": {}, "satisfied_axes": {}, @@ -6148,6 +6241,9 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): entry.setdefault("resolved_count", 0) entry.setdefault("regressed_count", 0) entry.setdefault("wrong_count", 0) + entry.setdefault("wrong_appearance_count", 0) + entry.setdefault("auto_inject_blocked_count", 0) + entry.setdefault("auto_inject_suppressed", False) entry.setdefault("bad_count", 0) entry.setdefault("wanted_axes", {}) entry.setdefault("satisfied_axes", {}) @@ -6260,13 +6356,31 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): add(" ".join(words[start:start + size]), "ngram", phrase, phrase.get("position", index) + start, phrase_context) prompt = str(last_run.get("prompt", "") or "") - prompt_segments = [ - {"text": segment.strip()} - for segment in re.split(r"[,.;\n]+", prompt) - if segment.strip() - ] - for position, segment in enumerate(prompt_segments): - add(segment.get("text", ""), "prompt_phrase", None, position, self._v2_context_for_phrase(prompt_segments, position, segment.get("text", ""))) + encoded_prompt = str(last_run.get("encoded_prompt", "") or "") + + def add_prompt_segments(source_prompt, kind, offset=0): + prompt_segments = [ + {"text": segment.strip()} + for segment in re.split(r"[,.;\n]+", source_prompt) + if segment.strip() + ] + for position, segment in enumerate(prompt_segments): + add( + segment.get("text", ""), + kind, + None, + offset + position, + self._v2_context_for_phrase(prompt_segments, position, segment.get("text", "")), + ) + + add_prompt_segments(prompt, "prompt_phrase") + if encoded_prompt and encoded_prompt != prompt: + add_prompt_segments(encoded_prompt, "auto_phrase", offset=128) + + for position, candidate in enumerate(last_run.get("repair_candidates", []) or []): + if not isinstance(candidate, dict): + continue + add(candidate.get("text", ""), "repair_candidate", candidate, 256 + position, {}) return units[:96] def _v2_classify_phrases(self, clip, phrases, global_state=None): @@ -6419,6 +6533,11 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): "Guidance: '-Just forget it-' skipped learning. Use it for workflow/model/seed failures; " "use a missing-axis rating when the prompt should teach the node." ) + if learning_profile.get("key") == "wrong_appearance": + return ( + "Guidance: wrong-appearance rating suppresses auto-injected character, clothing, " + "subject, and background memory without training unrelated action/detail/quality ratings." + ) if learning_profile.get("wrong_axes"): return ( "Guidance: wrong-intent rating preserved satisfied axes but marked the current " @@ -6504,6 +6623,10 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): def _v2_memory_kind_scale(self, kind): if kind in {"prompt_phrase", "phrase"}: return 1.0 + if kind == "auto_phrase": + return 0.72 + if kind == "repair_candidate": + return 0.64 if kind == "ngram": return 0.62 if kind == "token": @@ -6513,6 +6636,8 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): def _v2_update_preferred_context_memory(self, global_state, last_run, rating_profile, iter_num, axis_feedback=None): if not isinstance(global_state, dict) or not isinstance(last_run, dict) or rating_profile.get("skip_learning"): return "Preferred context: no update." + if rating_profile.get("key") == "wrong_appearance": + return "Preferred context: skipped for wrong appearance." phrases = [phrase for phrase in last_run.get("phrases", []) or [] if isinstance(phrase, dict)] if not phrases: return "Preferred context: no prompt phrases." @@ -6615,6 +6740,8 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): regressed_axes = set(axis_feedback.get("regressed_axes", [])) wrong_axes = set(axis_feedback.get("wrong_axes", [])) reward = float(rating_profile.get("reward", 0.0)) + wrong_appearance = rating_profile.get("key") == "wrong_appearance" + original_prompt = str(last_run.get("prompt", "") or "") preferred_context_status = self._v2_update_preferred_context_memory( global_state, last_run, @@ -6625,6 +6752,7 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): touched = [] trained = [] kind_counts = {} + suppressed = [] for phrase in concepts: text = str(phrase.get("text", "")).strip().lower() if not text: @@ -6664,6 +6792,7 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): matched_wrong = phrase_axes & wrong_axes rating_evidence = entry.setdefault("rating_evidence", {}) weights_before = dict(weights) + count_category_evidence = rating_profile.get("key") != "discover" def bump_evidence(name, axes): axis_counts = rating_evidence.setdefault(name, {}) @@ -6675,7 +6804,23 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): relation = self._v2_axis_category_weight(axis, category, effective_before if prefer_current else machine_scores) weights[category] = float(weights.get(category, 0.0)) + amount * relation - if rating_profile.get("key") == "like": + if wrong_appearance: + explicit = self._v2_prompt_contains_text(original_prompt, text) + auto_source = phrase.get("kind") in {"auto_phrase", "repair_candidate"} + should_suppress = self._v2_item_has_blocked_auto_category(phrase) and (auto_source or not explicit) + if should_suppress: + delta = -0.95 * kind_scale + entry["wrong_appearance_count"] = int(entry.get("wrong_appearance_count", 0)) + 1 + entry["auto_inject_blocked_count"] = int(entry.get("auto_inject_blocked_count", 0)) + 1 + entry["auto_inject_suppressed"] = True + rating_evidence["wrong_appearance"] = int(rating_evidence.get("wrong_appearance", 0)) + 1 + for category in self.AUTO_INJECT_BLOCKED_CATEGORIES: + weights[category] = float(weights.get(category, 0.0)) - 0.30 * kind_scale + suppressed.append(text) + else: + delta = 0.0 + count_category_evidence = False + elif rating_profile.get("key") == "like": delta = 0.55 * kind_scale entry["liked_count"] = int(entry.get("liked_count", 0)) + 1 rating_evidence["liked"] = int(rating_evidence.get("liked", 0)) + 1 @@ -6769,9 +6914,9 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): for category, value in weights.items() } entry["category_evidence_count"] = int(entry.get("category_evidence_count", 0)) + ( - 0 if rating_profile.get("key") == "discover" else 1 + 1 if count_category_evidence else 0 ) - if isinstance(context_sense, dict) and rating_profile.get("key") != "discover": + if isinstance(context_sense, dict) and count_category_evidence: sense_weights = self._v2_clean_category_scores(context_sense.get("category_weights", {})) for category, value in entry["category_weights"].items(): sense_weights[category] = float(sense_weights.get(category, 0.0)) + ( @@ -6816,10 +6961,15 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): if not touched: return "Lucky memory: no prompt phrases stored." kind_summary = ", ".join(f"{name}:{count}" for name, count in sorted(kind_counts.items())) or "none" + suppression_summary = ( + f" Appearance auto-injection suppressed: {', '.join(suppressed[:8])}{'...' if len(suppressed) > 8 else ''}." + if suppressed else + "" + ) return ( f"Category memory trained: {len(touched)} concept unit(s) ({kind_summary}). " f"Sample: {', '.join(trained) if trained else ', '.join(touched[:6])}. " - f"{preferred_context_status} " + f"{preferred_context_status}{suppression_summary} " f"{self._v2_axis_feedback_status(axis_feedback)}" ) @@ -6871,6 +7021,8 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): def _v2_update_conditioning_memory(self, global_state, last_run, rating_profile, axis_feedback=None): if not isinstance(last_run, dict) or rating_profile.get("skip_learning"): return + if rating_profile.get("key") == "wrong_appearance": + return payload = last_run.get("conditioning") if not isinstance(payload, dict): return @@ -7083,19 +7235,27 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): clean = self._v2_clean_phrase_text(text) if not clean: return + if not self._v2_prompt_repair_text_allowed(clean): + return words = self._v2_phrase_words(clean) if len(words) > 16: return axes = set(axes or set()) & missing_axes if not axes or self._v2_prompt_contains_text(prompt, clean): return + safe_cluster = [ + self._v2_clean_phrase_text(item) + for item in (cluster or []) + if self._v2_prompt_repair_text_allowed(item) + ] + safe_cluster = [item for item in safe_cluster if item] existing = candidates.get(clean) payload = { "text": clean, "axes": axes, "score": float(score), "source": source, - "cluster": list(cluster or []), + "cluster": safe_cluster, } if existing is None or payload["score"] > existing["score"]: candidates[clean] = payload @@ -7124,7 +7284,10 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): self._v2_clean_phrase_text(item) for item in entry.get("phrases", []) ] - cluster = [item for item in cluster if item and not self._v2_prompt_contains_text(prompt, item)] + cluster = [ + item for item in cluster + if item and self._v2_prompt_repair_text_allowed(item) and not self._v2_prompt_contains_text(prompt, item) + ] score = ( 2.0 + float(entry.get("score", 0.0)) * 0.34 + @@ -7179,7 +7342,11 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): for item in ranked: cluster = item.get("cluster", []) if item.get("source") == "preferred_context" else [] for cluster_text in cluster: - if cluster_text in selected_texts or self._v2_prompt_contains_text(prompt, cluster_text): + if ( + cluster_text in selected_texts or + self._v2_prompt_contains_text(prompt, cluster_text) or + not self._v2_prompt_repair_text_allowed(cluster_text) + ): continue cluster_item = dict(item) cluster_item["text"] = cluster_text @@ -7238,11 +7405,14 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): for text, entry in memory.items(): if not isinstance(entry, dict): continue + if not self._v2_auto_inject_entry_allowed(entry, prompt): + continue score = float(entry.get("score", 0.0)) score += min(1.0, int(entry.get("liked_count", 0)) * 0.14) score += min(1.2, sum(int(v) for v in entry.get("wanted_axes", {}).values()) * 0.10) score -= min(1.6, int(entry.get("bad_count", 0)) * 0.22) score -= min(1.1, int(entry.get("wrong_count", 0)) * 0.18) + score -= min(1.8, int(entry.get("auto_inject_blocked_count", 0)) * 0.65) if score <= 0.10: continue scored.append((score, entry.get("primary", "details"), str(entry.get("text", text)).strip())) @@ -7389,11 +7559,16 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): return "LoRA suggestions: " + " | ".join(parts) def refine_v2(self, positive_prompt, clip, rating, refinement_key, - seed=0, reset_session=False, lora_stack=None, im_feeling_lucky=False, user_intent_prompt=""): + seed=0, reset_session=False, lora_stack=None, im_feeling_lucky=False, user_intent_prompt="", + refinement_key_input=""): if seed != 0: torch.manual_seed(seed) random.seed(seed) + linked_refinement_key = str(refinement_key_input or "").strip() + if linked_refinement_key: + refinement_key = linked_refinement_key + rating_profile = normalize_refiner_v2_rating(rating) rating_label = rating_profile.get("label", str(rating)) state, state_status = self._v2_load_state(refinement_key, reset_session=reset_session) @@ -7558,6 +7733,12 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): "phrases": phrases, "intent_prompt": intent_prompt, "repair_candidates": state_repair_candidates, + "encoded_role": encoded_role, + "auto_injected_prompt_additions": [ + item.get("text", "") + for item in state_repair_candidates + if isinstance(item, dict) and str(item.get("text", "")).strip() + ], "rating_label": "Unrated", "iteration": int(global_state["total_iterations"]), } diff --git a/docs/FunPackVideoRefinerV2.md b/docs/FunPackVideoRefinerV2.md index 25ebdc4d..d7f8f7ab 100644 --- a/docs/FunPackVideoRefinerV2.md +++ b/docs/FunPackVideoRefinerV2.md @@ -17,6 +17,10 @@ V2 replaces the old `FunPack Video Refiner` public node in version `2.2.0`. It u - `Missing action` - `Missing quality` - `Missing details + action` +- `Wrong details` +- `Wrong action` +- `Wrong details + action` +- `Wrong appearance` - `Missing details + quality` - `Missing action + quality` - `Awful` @@ -24,8 +28,12 @@ V2 replaces the old `FunPack Video Refiner` public node in version `2.2.0`. It u Older `Missing concept` labels are accepted as aliases for `Missing action`. +Use `Wrong appearance` when the video is polluted by remembered clothing, character traits, subject identity, or background concepts that were not requested for the current image. This suppresses those concepts from future auto-injection without deleting them for explicit prompts. + **refinement_key**: V2 session key. V2 stores separate `refine_v2` state under this key. +**refinement_key_input**: Optional linked key input, usually from `FunPack Refinement Key Loader`. When connected, it overrides the text widget. + **seed**: Optional seed for deterministic exploration paths. **reset_session**: Clears the V2 session state for the selected key. @@ -34,6 +42,20 @@ Older `Missing concept` labels are accepted as aliases for `Missing action`. **im_feeling_lucky**: When enabled, V2 composes a learned prompt from rated phrase memory and encodes it through the connected CLIP. When disabled, V2 may still store Lucky memory from rated runs, but it does not compose or apply a Lucky prompt. +## Appearance Safety + +Refiner V2 treats image-to-video appearance as image-owned by default. Learned appearance, subject/character, and environment/background phrases are not auto-added by Prompt Repair or Lucky unless they are already present in the current prompt. + +Prompt Repair is only for motion, camera work, small non-identity details, quality, and style. It will not add clothing, body traits, character identity, or background phrases from memory. + +Legacy Void memory also skips appearance, subject/character, and environment/background tokens so old token-bank preferences cannot reintroduce those concepts. + +## Refinement Key Loader + +`FunPack Refinement Key Loader` lists existing V2 keys, creates a typed key when none is selected, and exposes `Import`, `Export`, and `Refresh` buttons in the browser UI. + +Connect its `refinement_key` output to Refiner V2's `refinement_key_input` and to `FunPack Apply LoRA Weights`' `refinement_key_input` so both nodes use the same memory key. + ## Outputs **modified_positive**: Refined positive conditioning. diff --git a/docs/FunPackVideoRefinerV2QuickGuide.md b/docs/FunPackVideoRefinerV2QuickGuide.md new file mode 100644 index 00000000..7486bf04 --- /dev/null +++ b/docs/FunPackVideoRefinerV2QuickGuide.md @@ -0,0 +1,37 @@ +# Refiner V2 Quick Guide + +`FunPack Video Refiner V2` is a prompt memory node. It takes your positive prompt, encodes it through the connected CLIP, remembers how you rated the previous output, and slowly learns which actions, camera moves, details, quality cues, and style cues work for your workflows. + +It does not look at the finished video pixels. It does not magically know what is inside your input image. It does not replace clear prompting. In image-to-video, it now treats character appearance, clothing, subject identity, and background as image-owned by default, so it should not inject old remembered outfits or characters into a new image unless you explicitly prompt them. + +## How To Use It + +1. Add `FunPack Video Refiner V2` before your sampler where positive conditioning is encoded. +2. Connect your `CLIP`. +3. Type a `refinement_key`, or use `FunPack Refinement Key Loader` and connect its output into `refinement_key_input`. +4. First run only seeds memory. Generate once, look at the result, then rate that result on the next run with the same key. +5. Use `Perfect` when the output is good and should be remembered. +6. Use `Missing action` when the requested motion did not happen. +7. Use `Missing details` when small non-identity details are weak, like smoke, reflections, particles, hand contact, object interaction, or secondary motion. +8. Use `Missing quality` when the idea is right but the result is messy or ugly. +9. Use `Wrong action`, `Wrong details`, or `Wrong details + action` when the video looks usable but did the wrong thing. +10. Use `Wrong appearance` when Refiner/Lucky brought in unwanted clothing, face, body, character, subject, or background memory. +11. Use `-Just forget it-` for broken runs, bad seeds, workflow mistakes, model failures, or anything you do not want learned. + +## Lucky Mode + +`I'm Feeling Lucky` composes a prompt from learned memory. It is useful when you want the node to suggest motion, camera, detail, quality, or style ideas from previous liked outputs. + +Lucky should not auto-add remembered outfits, character traits, subjects, or backgrounds anymore. If you want an appearance item, write it directly in the prompt. + +## Examples + +If your prompt says `person smoking and exhaling smoke` but the video only shows a puff with no exhale, rate `Missing details` or `Missing action`, depending on what matters most. This teaches the key that smoke/exhale matters in that prompt context, without globally banning smoking. + +If you liked `white tights` in older prompts but your new input image is only a portrait face, and Lucky tries to bring tights into the scene, rate `Wrong appearance`. That tells Refiner to suppress that appearance memory from auto-injection while still allowing `white tights` when you explicitly type it. + +If the camera should push in, orbit, pan, or change framing and it stays static, rate `Missing action` or `Wrong action`. Camera movement belongs to the action/camera side of memory, so it remains safe for repair and Lucky. + +## Key Loader + +Use `FunPack Refinement Key Loader` when you want a reusable key dropdown or want to move keys between workflows/machines. Its `Import` and `Export` buttons save/load the key as JSON. Connect the loader to both Refiner V2 and `FunPack Apply LoRA Weights` so LoRA suggestions and prompt memory use the same key. diff --git a/funpack.py b/funpack.py index c1afca82..4013e4aa 100644 --- a/funpack.py +++ b/funpack.py @@ -19,7 +19,7 @@ try: from .model_management import FunPackApplyLoraWeights, FunPackLoraLoader from .samplers import FunPackHybridEuler2SSampler, sample_funpack_hybrid_euler_2s from .context_transition import FunPackContextTransitionWindows - from .templates import FunPackTemplateManager + from .templates import FunPackRefinementKeyLoader, FunPackTemplateManager except ImportError: from conditioning import ( FunPackLorebookEnhancer, @@ -38,7 +38,7 @@ except ImportError: from model_management import FunPackApplyLoraWeights, FunPackLoraLoader from samplers import FunPackHybridEuler2SSampler, sample_funpack_hybrid_euler_2s from context_transition import FunPackContextTransitionWindows - from templates import FunPackTemplateManager + from templates import FunPackRefinementKeyLoader, FunPackTemplateManager NODE_CLASS_MAPPINGS = { "FunPackPromptCombiner": FunPackPromptCombiner, @@ -55,6 +55,7 @@ except ImportError: "FunPackContextTransitionWindows": FunPackContextTransitionWindows, "FunPackApplyLoraWeights": FunPackApplyLoraWeights, "FunPackLoraLoader": FunPackLoraLoader, + "FunPackRefinementKeyLoader": FunPackRefinementKeyLoader, "FunPackTemplateManager": FunPackTemplateManager, } @@ -73,6 +74,7 @@ except ImportError: "FunPackContextTransitionWindows": "FunPack Context Transition Windows", "FunPackApplyLoraWeights": "FunPack Apply LoRA Weights", "FunPackLoraLoader": "FunPack LoRA Loader", + "FunPackRefinementKeyLoader": "FunPack Refinement Key Loader", "FunPackTemplateManager": "FunPack Template Manager", } @@ -93,6 +95,7 @@ __all__ = [ "FunPackStoryMemLastFrameExtractor", "FunPackApplyLoraWeights", "FunPackLoraLoader", + "FunPackRefinementKeyLoader", "FunPackTemplateManager", "sample_funpack_hybrid_euler_2s", ] diff --git a/model_management.py b/model_management.py index 07038c92..5b63939b 100644 --- a/model_management.py +++ b/model_management.py @@ -189,6 +189,12 @@ class FunPackApplyLoraWeights: any_type, { "lora_list": ("STRING", {"default": "[]", "multiline": False}), + "refinement_key_input": ("STRING", { + "default": "", + "multiline": False, + "forceInput": True, + "tooltip": "Optional linked refinement key, for example from FunPack Refinement Key Loader. Overrides the refinement_key widget when connected.", + }), "lora_0": (loras, {"default": "None"}), "lora_0_type": (LORA_TYPES, {"default": "general"}), "lora_0_base_weight": ( @@ -358,6 +364,9 @@ class FunPackApplyLoraWeights: return abs(float(saved_base) - float(entry["base_model_weight"])) <= 1e-6 def apply_lora_weights(self, positive_prompt, refinement_key, mode, per_block=False, **kwargs): + linked_refinement_key = str(kwargs.pop("refinement_key_input", "") or "").strip() + if linked_refinement_key: + refinement_key = linked_refinement_key mode = (mode or "ltx2").lower() per_block = coerce_bool(per_block) prompt_key = prompt_key_for_mode(positive_prompt, mode) diff --git a/pyproject.toml b/pyproject.toml index 5b279c3a..043b7fce 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,7 +2,7 @@ [project] name = "ComfyUI-FunPack" # Unique identifier for your node. Immutable after creation. description = "A set of nodes for experiments with video generation models, including prompt refinement, adaptive LoRA weight workflows, story tools, video stitching, lorebook enhancement and more." -version = "2.2.1" # Custom Node version. Must be semantically versioned. +version = "2.3.0" # Custom Node version. Must be semantically versioned. license = { text = "GNU GPL v3" } dependencies = [] # Filled in from requirements.txt diff --git a/templates.py b/templates.py index 3d106b03..554b116b 100644 --- a/templates.py +++ b/templates.py @@ -27,6 +27,7 @@ except ImportError: TEMPLATE_NONE = "-None-" +REFINEMENT_KEY_NONE = "-None-" TEMPLATE_DB_VERSION = 1 WILDCARD_RE = re.compile(r"\{([^{}]*\|[^{}]*)\}") @@ -46,6 +47,120 @@ def template_store_path(): return os.path.join(template_store_dir(), "templates.json") +def refinement_store_dir(): + return os.path.join(os.path.dirname(os.path.abspath(__file__)), "refinements") + + +def normalize_refinement_key(value): + value = str(value or "").strip() + if not value or value == REFINEMENT_KEY_NONE: + return "" + return value + + +def empty_v2_refinement_state(refinement_key): + return { + "version": 2, + "refinement_key": refinement_key, + "state_namespace": "clip", + "global": { + "total_iterations": 0, + "avg_reward_ema": 0.0, + "good_streak": 0, + "bad_streak": 0, + "last_rating_label": "Initial discovery", + "last_missing_axes": [], + "phrase_memory": {}, + "axis_conditioning_memory": {}, + "lora_weight_memory": {}, + "preferred_context_memory": {}, + "loss_history": [], + }, + "prompt_histories": {}, + "last_run": None, + } + + +def refinement_key_path(refinement_key): + return refinement_state_path(refinement_key, "clip", prefix="refine_v2") + + +def normalize_refinement_state(data, fallback_key=""): + if not isinstance(data, dict): + return None, "" + key = normalize_refinement_key(data.get("refinement_key") or fallback_key) + if not key: + return None, "" + state = dict(data) + state["version"] = 2 + state["refinement_key"] = key + state["state_namespace"] = "clip" + state.setdefault("global", {}) + state["global"].setdefault("phrase_memory", {}) + state["global"].setdefault("axis_conditioning_memory", {}) + state["global"].setdefault("lora_weight_memory", {}) + state["global"].setdefault("preferred_context_memory", {}) + state["global"].setdefault("loss_history", []) + state.setdefault("prompt_histories", {}) + state.setdefault("last_run", None) + return state, key + + +def load_refinement_key_state(refinement_key, create=False): + key = normalize_refinement_key(refinement_key) + if not key: + return None, "", "No refinement key selected." + path = refinement_key_path(key) + if os.path.exists(path): + try: + with open(path, "r", encoding="utf-8") as file: + data = json.load(file) + state, loaded_key = normalize_refinement_state(data, key) + if state is None: + return None, key, f"Refinement key '{key}' is unreadable." + return state, loaded_key, f"Loaded refinement key '{loaded_key}'." + except (json.JSONDecodeError, OSError, ValueError): + return None, key, f"Refinement key '{key}' is unreadable." + if not create: + return None, key, f"Refinement key '{key}' does not exist." + state = empty_v2_refinement_state(key) + save_refinement_key_state(state, key) + return state, key, f"Created refinement key '{key}'." + + +def save_refinement_key_state(state, refinement_key): + key = normalize_refinement_key(refinement_key) + if not key: + return "" + state, key = normalize_refinement_state(state, key) + if state is None: + return "" + path = refinement_key_path(key) + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w", encoding="utf-8") as file: + json.dump(state, file, indent=2) + return path + + +def refinement_key_names(): + keys = set() + directory = refinement_store_dir() + if os.path.isdir(directory): + for filename in os.listdir(directory): + if not filename.startswith("refine_v2_") or not filename.endswith(".json"): + continue + path = os.path.join(directory, filename) + try: + with open(path, "r", encoding="utf-8") as file: + data = json.load(file) + key = normalize_refinement_key(data.get("refinement_key") if isinstance(data, dict) else "") + if key: + keys.add(key) + except (json.JSONDecodeError, OSError, ValueError): + continue + return [REFINEMENT_KEY_NONE] + sorted(keys) + + def empty_template_db(): return { "version": TEMPLATE_DB_VERSION, @@ -322,6 +437,79 @@ async def funpack_templates_import(request): return web.json_response({"imported": imported, "templates": template_names()}) +@PromptServer.instance.routes.get("/funpack/refinement_keys") +async def funpack_refinement_keys(_): + return web.json_response( + {"keys": refinement_key_names(), "path": refinement_store_dir()}, + headers={"Cache-Control": "no-store, max-age=0"}, + ) + + +@PromptServer.instance.routes.get("/funpack/refinement_keys/export") +async def funpack_refinement_keys_export(request): + key = normalize_refinement_key(request.query.get("key", "")) + state, loaded_key, status = load_refinement_key_state(key, create=False) + if state is None: + return web.json_response({"error": status}, status=404) + return web.json_response( + state, + headers={ + "Cache-Control": "no-store, max-age=0", + "Content-Disposition": f"attachment; filename=funpack_refinement_{loaded_key}.json", + }, + ) + + +@PromptServer.instance.routes.post("/funpack/refinement_keys/import") +async def funpack_refinement_keys_import(request): + incoming = await request.json() + if isinstance(incoming, dict) and isinstance(incoming.get("state"), dict): + incoming = incoming["state"] + state, key = normalize_refinement_state(incoming) + if state is None: + return web.json_response({"error": "Imported file is not a valid V2 refinement key JSON."}, status=400) + path = save_refinement_key_state(state, key) + if not path: + return web.json_response({"error": "Could not save imported refinement key."}, status=400) + return web.json_response({"imported": key, "keys": refinement_key_names()}) + + +class FunPackRefinementKeyLoader: + CATEGORY = "FunPack/Refinement" + RETURN_TYPES = ("STRING", "STRING") + RETURN_NAMES = ("refinement_key", "status") + FUNCTION = "load_refinement_key" + OUTPUT_NODE = True + DESCRIPTION = "Loads, creates, imports, and exports FunPack Video Refiner V2 refinement keys." + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "refinement_key": (refinement_key_names(), {"default": REFINEMENT_KEY_NONE}), + "key_name": ("STRING", {"default": "", "multiline": False}), + "create_if_missing": ("BOOLEAN", {"default": True}), + }, + } + + @classmethod + def IS_CHANGED(cls, **kwargs): + return float("nan") + + @classmethod + def VALIDATE_INPUTS(cls, refinement_key=None, **kwargs): + return True + + def load_refinement_key(self, refinement_key, key_name, create_if_missing=True): + selected = normalize_refinement_key(refinement_key) + typed = normalize_refinement_key(key_name) + target = selected or typed + if not target: + return ("", "No refinement key selected. Pick an existing key or type a new key name.") + _, loaded_key, status = load_refinement_key_state(target, create=bool(create_if_missing)) + return (loaded_key if loaded_key else target, status) + + class FunPackTemplateManager: CATEGORY = "FunPack/Templates" RETURN_TYPES = ("STRING", "STRING", "STRING", "SIGMAS", "CONDITIONING", "FUNPACK_LORA_STACK", "STRING", "STRING") diff --git a/tests/test_v2_category_alignment.py b/tests/test_v2_category_alignment.py index 75a38129..e8191bf4 100644 --- a/tests/test_v2_category_alignment.py +++ b/tests/test_v2_category_alignment.py @@ -201,7 +201,10 @@ def test_awful_lora_feedback_reduces_before_missing_axis_boosts(): def test_category_weights_are_recorded_for_every_learning_rating(): refiner = FunPackVideoRefinerV2() - learning_labels = [label for label in V2_RATING_LABELS if label != "-Just forget it-"] + learning_labels = [ + label for label in V2_RATING_LABELS + if label not in {"-Just forget it-", "Wrong appearance"} + ] for label in learning_labels: entry, _, _ = train_phrase(refiner, f"running test {label}", label) @@ -212,6 +215,158 @@ def test_category_weights_are_recorded_for_every_learning_rating(): assert set(entry["effective_category_scores"]) == set(refiner.CATEGORY_DESCRIPTIONS) +def test_wrong_appearance_rating_is_available(): + profile = normalize_refiner_v2_rating("Wrong appearance") + + assert "Wrong appearance" in V2_RATING_LABELS + assert profile["key"] == "wrong_appearance" + assert profile["wrong_categories"] == ["appearance", "subject", "environment"] + + +def test_prompt_repair_blocks_appearance_subject_and_environment(): + refiner = FunPackVideoRefinerV2() + global_state = { + "phrase_memory": { + "white tights": { + "text": "white tights", + "primary": "appearance", + "effective_category_scores": refiner._v2_heuristic_scores("white tights"), + "wanted_axes": {"details": 3}, + "score": 6.0, + "liked_count": 4, + }, + "detailed background": { + "text": "detailed background", + "primary": "environment", + "effective_category_scores": refiner._v2_heuristic_scores("detailed background"), + "wanted_axes": {"details": 3}, + "score": 6.0, + "liked_count": 4, + }, + "female character": { + "text": "female character", + "primary": "subject", + "effective_category_scores": refiner._v2_heuristic_scores("female character"), + "wanted_axes": {"details": 3}, + "score": 6.0, + "liked_count": 4, + }, + "tiny smoke curls": { + "text": "tiny smoke curls", + "primary": "details", + "effective_category_scores": refiner._v2_heuristic_scores("tiny smoke curls"), + "wanted_axes": {"details": 3}, + "score": 4.0, + "liked_count": 2, + }, + } + } + profile = normalize_refiner_v2_rating("Missing details") + feedback = refiner._v2_axis_feedback(profile, None) + + repaired, _, candidates = refiner._v2_repair_prompt_for_missing_axes( + "person smoking", + refiner._v2_classify_phrases(None, prompt_items(refiner, ["person smoking"]), global_state), + global_state, + None, + feedback, + ) + + assert "tiny smoke curls" in repaired + assert "white tights" not in repaired + assert "detailed background" not in repaired + assert "female character" not in repaired + assert all(candidate["text"] != "white tights" for candidate in candidates) + + +def test_lucky_skips_appearance_memory_unless_prompt_mentions_it(): + refiner = FunPackVideoRefinerV2() + global_state = {"phrase_memory": {}, "preferred_context_memory": {}} + phrases = refiner._v2_classify_phrases( + None, + prompt_items(refiner, ["white tights", "walking"]), + global_state, + ) + profile = normalize_refiner_v2_rating("Perfect") + feedback = refiner._v2_axis_feedback(profile, None) + refiner._v2_update_phrase_memory( + global_state, + {"prompt": "white tights walking", "phrases": phrases}, + profile, + 1, + feedback, + ) + + lucky_prompt, _ = refiner._v2_compose_lucky_prompt("", [], global_state) + explicit_prompt, _ = refiner._v2_compose_lucky_prompt( + "white tights", + refiner._v2_classify_phrases(None, prompt_items(refiner, ["white tights"]), global_state), + global_state, + ) + + assert "walking" in lucky_prompt + assert "white tights" not in lucky_prompt + assert "tights" not in lucky_prompt + assert "white tights" in explicit_prompt + + +def test_wrong_appearance_suppresses_only_auto_inserted_appearance_memory(): + refiner = FunPackVideoRefinerV2() + global_state = {"phrase_memory": {}, "preferred_context_memory": {}} + phrases = refiner._v2_classify_phrases(None, prompt_items(refiner, ["walking"]), global_state) + profile = normalize_refiner_v2_rating("Wrong appearance") + feedback = refiner._v2_axis_feedback(profile, None) + + refiner._v2_update_phrase_memory( + global_state, + { + "prompt": "walking", + "encoded_prompt": "walking, white tights", + "phrases": phrases, + "repair_candidates": [{"text": "white tights", "axes": ["details"], "score": 3.0, "source": "memory"}], + }, + profile, + 1, + feedback, + ) + + appearance = global_state["phrase_memory"]["white tights"] + motion = global_state["phrase_memory"]["walking"] + assert appearance["auto_inject_suppressed"] is True + assert appearance["wrong_appearance_count"] >= 1 + assert motion.get("wrong_appearance_count", 0) == 0 + assert motion["category_evidence_count"] == 0 + + +def test_repeated_wrong_appearance_keeps_entry_out_of_lucky(): + refiner = FunPackVideoRefinerV2() + global_state = {"phrase_memory": {}, "preferred_context_memory": {}} + phrases = refiner._v2_classify_phrases(None, prompt_items(refiner, ["walking"]), global_state) + profile = normalize_refiner_v2_rating("Wrong appearance") + feedback = refiner._v2_axis_feedback(profile, None) + + for iteration in range(2): + refiner._v2_update_phrase_memory( + global_state, + { + "prompt": "walking", + "encoded_prompt": "walking, white tights", + "phrases": phrases, + "repair_candidates": [{"text": "white tights", "axes": ["details"], "score": 3.0, "source": "memory"}], + }, + profile, + iteration + 1, + feedback, + ) + + global_state["phrase_memory"]["white tights"]["score"] = 6.0 + global_state["phrase_memory"]["white tights"]["liked_count"] = 8 + lucky_prompt, _ = refiner._v2_compose_lucky_prompt("", [], global_state) + + assert global_state["phrase_memory"]["white tights"]["auto_inject_blocked_count"] >= 2 + assert "white tights" not in lucky_prompt + + def test_forget_rating_skips_category_weight_learning(): refiner = FunPackVideoRefinerV2() global_state = {"phrase_memory": {}} diff --git a/web/funpack_refinement_key_loader.js b/web/funpack_refinement_key_loader.js new file mode 100644 index 00000000..dd05ffac --- /dev/null +++ b/web/funpack_refinement_key_loader.js @@ -0,0 +1,250 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +const NODE_NAME = "FunPackRefinementKeyLoader"; +const NONE_KEY = "-None-"; +const BUTTON_NAMES = new Set([ + "funpack_refinement_key_export", + "funpack_refinement_key_import", + "funpack_refinement_key_refresh", +]); + +let cachedKeys = [NONE_KEY]; +const trackedNodes = new Set(); +let pendingRefresh = null; + +function widgetByName(node, name) { + return (node.widgets || []).find((widget) => widget.name === name); +} + +function keyWidget(node) { + return widgetByName(node, "refinement_key"); +} + +function keyNameWidget(node) { + return widgetByName(node, "key_name"); +} + +function selectedKey(node) { + const selected = String(keyWidget(node)?.value || "").trim(); + if (selected && selected !== NONE_KEY) { + return selected; + } + return String(keyNameWidget(node)?.value || "").trim(); +} + +function fitString(ctx, text, maxWidth) { + text = String(text ?? ""); + if (ctx.measureText(text).width <= maxWidth) { + return text; + } + const ellipsis = "..."; + let low = 0; + let high = text.length; + while (low < high) { + const mid = Math.ceil((low + high) / 2); + if (ctx.measureText(text.slice(0, mid) + ellipsis).width <= maxWidth) { + low = mid; + } else { + high = mid - 1; + } + } + return text.slice(0, low) + ellipsis; +} + +function updateKeyWidget(node, values = cachedKeys) { + const widget = keyWidget(node); + if (!widget) { + return; + } + const current = widget.value; + const nextValues = current && !values.includes(current) ? [...values, current] : values; + if (widget.options) { + widget.options.values = nextValues; + } + if (!widget.value || !nextValues.includes(widget.value)) { + widget.value = nextValues[0] || NONE_KEY; + } +} + +async function fetchKeys() { + try { + const response = await api.fetchApi(`/funpack/refinement_keys?cache_bust=${Date.now()}`, { + cache: "no-store", + }); + if (!response.ok) { + throw new Error(`HTTP ${response.status}`); + } + const data = await response.json(); + const values = Array.isArray(data.keys) && data.keys.length ? data.keys : [NONE_KEY]; + cachedKeys = values; + return values; + } catch (error) { + console.warn("FunPack: failed to refresh refinement keys", error); + return cachedKeys; + } +} + +async function refreshNode(node) { + const values = await fetchKeys(); + updateKeyWidget(node, values); + node.setDirtyCanvas(true, true); +} + +async function refreshTrackedNodes() { + if (pendingRefresh) { + return pendingRefresh; + } + pendingRefresh = (async () => { + const values = await fetchKeys(); + for (const node of [...trackedNodes]) { + if (!node?.graph) { + trackedNodes.delete(node); + continue; + } + updateKeyWidget(node, values); + node.setDirtyCanvas(true, true); + } + })().finally(() => { + pendingRefresh = null; + }); + return pendingRefresh; +} + +async function exportKey(node) { + const key = selectedKey(node); + if (!key) { + app.canvas?.prompt?.("Refinement key export skipped", "Select or type a refinement key first.", () => {}); + return; + } + const response = await api.fetchApi(`/funpack/refinement_keys/export?key=${encodeURIComponent(key)}&cache_bust=${Date.now()}`, { + cache: "no-store", + }); + if (!response.ok) { + const error = await response.json().catch(() => ({})); + throw new Error(error.error || `Export failed with HTTP ${response.status}`); + } + const blob = await response.blob(); + const url = URL.createObjectURL(blob); + const link = document.createElement("a"); + link.href = url; + link.download = `funpack_refinement_${key}.json`; + document.body.appendChild(link); + link.click(); + link.remove(); + URL.revokeObjectURL(url); +} + +function importKey() { + const input = document.createElement("input"); + input.type = "file"; + input.accept = ".json,application/json"; + input.onchange = async () => { + const file = input.files?.[0]; + if (!file) { + return; + } + try { + const data = JSON.parse(await file.text()); + const response = await api.fetchApi("/funpack/refinement_keys/import", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(data), + }); + if (!response.ok) { + const error = await response.json().catch(() => ({})); + throw new Error(error.error || `Import failed with HTTP ${response.status}`); + } + await refreshTrackedNodes(); + } catch (error) { + console.warn("FunPack: refinement key import failed", error); + app.canvas?.prompt?.("Refinement key import failed", error.message || String(error), () => {}); + } + }; + input.click(); +} + +class FunPackRefinementKeyButton { + constructor(name, label, callback) { + this.name = name; + this.type = "custom"; + this.options = { serialize: false }; + this.value = ""; + this.label = label; + this.callback = callback; + } + + computeSize(width) { + return [width, LiteGraph.NODE_WIDGET_HEIGHT + 2]; + } + + draw(ctx, _node, width, y, height) { + const x = 20; + const w = width - 40; + const h = height - 2; + ctx.save(); + ctx.beginPath(); + ctx.roundRect(x, y + 1, w, h, [3]); + ctx.fillStyle = LiteGraph.WIDGET_BGCOLOR; + ctx.strokeStyle = LiteGraph.WIDGET_OUTLINE_COLOR; + ctx.fill(); + ctx.stroke(); + ctx.fillStyle = LiteGraph.WIDGET_TEXT_COLOR; + ctx.textAlign = "center"; + ctx.textBaseline = "middle"; + ctx.fillText(fitString(ctx, this.label, w - 12), width / 2, y + height / 2); + ctx.restore(); + } + + mouse(event, _pos, node) { + if (event.type === "pointerup") { + this.callback(node); + return true; + } + return event.type === "pointerdown"; + } +} + +function removeButtons(node) { + node.widgets = (node.widgets || []).filter((widget) => !BUTTON_NAMES.has(widget.name)); +} + +function addButtons(node) { + removeButtons(node); + const buttons = [ + new FunPackRefinementKeyButton("funpack_refinement_key_export", "Export", (currentNode) => void exportKey(currentNode)), + new FunPackRefinementKeyButton("funpack_refinement_key_import", "Import", () => importKey()), + new FunPackRefinementKeyButton("funpack_refinement_key_refresh", "Refresh", (currentNode) => void refreshNode(currentNode)), + ]; + for (const button of buttons) { + node.addCustomWidget(button); + } +} + +function setupNode(node) { + trackedNodes.add(node); + updateKeyWidget(node); + addButtons(node); + void refreshNode(node); +} + +app.registerExtension({ + name: "funpack.refinementKeyLoader", + beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name !== NODE_NAME) { + return; + } + + const originalOnNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + originalOnNodeCreated?.apply(this, arguments); + setupNode(this); + }; + + const originalConfigure = nodeType.prototype.configure; + nodeType.prototype.configure = function (info) { + originalConfigure?.apply(this, arguments); + setupNode(this, info); + }; + }, +});