diff --git a/README.md b/README.md index 7c91732..da8e8b0 100644 --- a/README.md +++ b/README.md @@ -215,6 +215,7 @@ Outputs: Highlights: - Added `LoRA Stack (Model In→Out) WAN` node: directly applies up to 4 WAN / Flow / Standard LoRAs to an incoming MODEL and outputs a patched MODEL (ideal for WAN 2.2 workflows where a single node step is preferred). +- Added `LoRA Schedule (WAN, ranged)` node: activates extra LoRA stacks by generation index, with optional open-ended ranges for long loop workflows. - Extended internal WAN key remapping for seamless WAN 2.2 (Flow) + WAN 2.1 cross-compatibility. - Version bump across project files. @@ -241,6 +242,37 @@ Recommended Use (WAN 2.2 workflows): Why this node: Eliminates one extra node hop, reduces graph complexity and clarifies model lineage in large animation workflows. +### Additional Node: LoRA Schedule (WAN, ranged) + +Use this node when a WAN loop needs always-on LoRAs plus extra LoRAs that only apply on specific generations or generation ranges. + +Typical setup: +1. Keep your main always-on LoRAs in `IAMCCS_WanLoRAStack` or `IAMCCS_WanLoRAStackModelIO`. +2. Build extra LoRA stacks for alternate phases. +3. Feed the loop `index` into `generation_index`. +4. Set `slot_01_start/end`, `slot_02_start/end`, etc. to define which extra stack is active on which generations. +5. Send the scheduler `lora` output into the optional `lora` input of your main WAN LoRA stack node. + +Notes: +- `default_lora` stays active on every generation. +- `end = -1` means "from this generation onward". +- Multiple active slots stack together, so you can layer phase LoRAs if ranges overlap. + +Preset behavior: +- `manual_range`: uses `start/end`; if `end = -1`, the slot stays active from `start` onward. +- `all_generations`: always active. +- `only_first`: active only on generation `0`. +- `all_nonfirst`: active from generation `1` onward. +- `even_generations`: active on `0, 2, 4, ...`. +- `odd_generations`: active on `1, 3, 5, ...`. +- `every_2_from_start`: active on `start, start+2, start+4, ...`. +- `every_3_from_start`: active on `start, start+3, start+6, ...`. + +Logging: +- Each execution logs the current `generation_index`. +- Logs show the always-on `default_lora` entries. +- Logs show each active slot, the preset that matched, and the actual LoRA names/strengths injected on that generation. + ## Previous Versions ### Version 1.2.3 — New input lora - add another StackLoraModel (concatenate) + Extended Wan 2.1 Compatibility diff --git a/__init__.py b/__init__.py index faf23f8..c69fed7 100644 --- a/__init__.py +++ b/__init__.py @@ -19,6 +19,9 @@ from .iamccs_wan_lora_stack import ( from .iamccs_wan_lora_stack_simple import ( IAMCCS_WanLoRAStackModelIO, ) +from .iamccs_wan_lora_schedule import ( + IAMCCS_WanLoRASchedule, +) from .iamccs_ltx2_lora_stack import ( IAMCCS_LTX2_LoRAStack, @@ -204,6 +207,7 @@ NODE_CLASS_MAPPINGS = { "IAMCCS_WanLoRAStack": IAMCCS_WanLoRAStack, "IAMCCS_ModelWithLoRA": IAMCCS_ModelWithLoRA, "IAMCCS_WanLoRAStackModelIO": IAMCCS_WanLoRAStackModelIO, + "IAMCCS_WanLoRASchedule": IAMCCS_WanLoRASchedule, # Backward-compatible key (kept as-is for existing workflows) "iamccs_ltx2_lora_stack": IAMCCS_LTX2_LoRAStack, # Preferred explicit names @@ -336,6 +340,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "IAMCCS_WanLoRAStack": "LoRA Stack (WAN-style remap)", "IAMCCS_ModelWithLoRA": "Apply LoRA to MODEL (Native)", "IAMCCS_WanLoRAStackModelIO": "LoRA Stack (Model In→Out) WAN", + "IAMCCS_WanLoRASchedule": "LoRA Schedule (WAN, ranged)", "iamccs_ltx2_lora_stack": "iamccs_ltx2_lora_stack (3 slots)", "IAMCCS_LTX2_LoRAStack": "LoRA Stack (LTX-2, 3 slots)", "IAMCCS_LTX2_LoRAStackStaged": "LoRA Stack (LTX-2, staged: stage1+stage2) (BETA)", diff --git a/iamccs_flux_klein_multigen.py b/iamccs_flux_klein_multigen.py index 4b702c0..b4105ac 100644 --- a/iamccs_flux_klein_multigen.py +++ b/iamccs_flux_klein_multigen.py @@ -1,4 +1,30 @@ -import math +"""IAMCCS_FluxKleinMultiGen - Complete pipeline reconstruction. + +Pipeline graph (mirrors working reference workflow 1:1): + image + -> ImageScaleToTotalPixels(upscale_method, megapixels, resolution_steps) + -> scaled_image + + scaled_image + vae + -> VAEEncode + -> reference_latent <- shared between positive AND negative + + clip + positive_prompt -> CLIPTextEncode -> pos_cond + pos_cond + reference_latent -> ReferenceLatent -> positive_with_ref + + clip + negative_prompt -> CLIPTextEncode -> neg_cond + neg_cond + reference_latent -> ReferenceLatent -> negative_with_ref + + output_width, output_height -> Flux2Scheduler(steps, width, height) -> sigmas + output_width, output_height -> EmptyFlux2LatentImage(width, height, 1) -> latent_image + seed -> RandomNoise -> noise + sampler_name -> KSamplerSelect -> sampler + model + positive_with_ref + negative_with_ref + cfg -> CFGGuider -> guider + noise + guider + sampler + sigmas + latent_image -> SamplerCustomAdvanced -> output_latent + output_latent + vae -> VAEDecode -> image_out +""" + +import json import os import re @@ -6,187 +32,150 @@ import numpy as np import torch from PIL import Image as _PILImage -import comfy.utils import folder_paths import node_helpers -RESOLUTION_PRESETS = [ - "512x512", - "720x1280", - "768x768", - "768x1024", - "832x1216", - "864x1536", - "1024x768", - "1024x1024", - "1024x1536", - "1280x720", - "1360x768", - "1216x832", - "1536x1024", - "1536x864", - "1920x1080", -] +# --------------------------------------------------------------------------- +# Unpack helpers +# io.NodeOutput stores outputs in .args; old-style nodes return plain tuples. +# --------------------------------------------------------------------------- + +def _unpack(value): + """Unwrap io.NodeOutput or single-element tuple to the actual value.""" + if hasattr(value, "args"): # io.NodeOutput (new ComfyUI API) + if value.args: + return value.args[0] + return value + if isinstance(value, (tuple, list)) and len(value) == 1: + return value[0] + return value -def _safe_unpack(value): - current = value - while isinstance(current, tuple) and len(current) == 1: - current = current[0] - if hasattr(current, "args"): - args = getattr(current, "args", None) - if isinstance(args, (tuple, list)) and len(args) > 0: - return _safe_unpack(args[0]) - return current +def _unpack_sampler(value): + """SamplerCustomAdvanced returns (output, denoised_output); take output (index 0).""" + if hasattr(value, "args") and value.args: + return value.args[0] + if isinstance(value, (tuple, list)): + return value[0] + return value -def _encode_text(clip, prompt): - tokens = clip.tokenize(prompt or "") - return clip.encode_from_tokens_scheduled(tokens) +# --------------------------------------------------------------------------- +# One function per node in the working workflow (exact mirrors) +# --------------------------------------------------------------------------- - -def _apply_reference_latent(conditioning, latent, method): - conditioning = node_helpers.conditioning_set_values( - conditioning, - {"reference_latents": [latent]}, - append=True, - ) - if "uxo" in method or "uso" in method: - method = "uxo" - return node_helpers.conditioning_set_values( - conditioning, - {"reference_latents_method": method}, +def _node_ImageScaleToTotalPixels(image, upscale_method, megapixels, resolution_steps): + """Node: ImageScaleToTotalPixels""" + from comfy_extras.nodes_post_processing import ImageScaleToTotalPixels + return _unpack( + ImageScaleToTotalPixels.execute( + image, + str(upscale_method), + float(megapixels), + int(resolution_steps), + ) ) -def _prompt_to_slug(prompt, max_words=5): - text = re.sub(r"<[^>]+>", "", prompt or "").strip() - words = [re.sub(r"[^a-zA-Z0-9]", "", word) for word in text.split()[:max_words] if word] - filtered = [word for word in words if word] - return "_".join(filtered) if filtered else "prompt" +def _node_VAEEncode(vae, pixels): + """Node: VAEEncode""" + import nodes + return _unpack(nodes.VAEEncode().encode(vae, pixels)) -def _unique_png_path(output_dir, prefix): - counter = 1 +def _node_CLIPTextEncode(clip, text): + """Node: CLIPTextEncode""" + import nodes + return _unpack(nodes.CLIPTextEncode().encode(clip, text or "")) + + +def _node_ReferenceLatent(conditioning, latent): + """Node: ReferenceLatent - latent must be a {'samples': tensor} dict""" + from comfy_extras.nodes_edit_model import ReferenceLatent + return _unpack(ReferenceLatent.execute(conditioning, latent)) + + +def _node_Flux2Scheduler(steps, width, height): + """Node: Flux2Scheduler""" + from comfy_extras.nodes_flux import Flux2Scheduler + return _unpack(Flux2Scheduler.execute(int(steps), int(width), int(height))) + + +def _node_EmptyFlux2LatentImage(width, height): + """Node: EmptyFlux2LatentImage (batch_size always 1 per iteration)""" + from comfy_extras.nodes_flux import EmptyFlux2LatentImage + return _unpack(EmptyFlux2LatentImage.execute(int(width), int(height), 1)) + + +def _node_RandomNoise(seed): + """Node: RandomNoise""" + from comfy_extras.nodes_custom_sampler import RandomNoise + return _unpack(RandomNoise.execute(int(seed))) + + +def _node_KSamplerSelect(sampler_name): + """Node: KSamplerSelect""" + from comfy_extras.nodes_custom_sampler import KSamplerSelect + return _unpack(KSamplerSelect.execute(str(sampler_name))) + + +def _node_CFGGuider(model, positive, negative, cfg): + """Node: CFGGuider""" + from comfy_extras.nodes_custom_sampler import CFGGuider + return _unpack(CFGGuider.execute(model, positive, negative, float(cfg))) + + +def _node_SamplerCustomAdvanced(noise, guider, sampler, sigmas, latent_image): + """Node: SamplerCustomAdvanced - returns first output slot (output latent)""" + from comfy_extras.nodes_custom_sampler import SamplerCustomAdvanced + return _unpack_sampler( + SamplerCustomAdvanced.execute(noise, guider, sampler, sigmas, latent_image) + ) + + +def _node_VAEDecode(vae, latent): + """Node: VAEDecode""" + import nodes + return _unpack(nodes.VAEDecode().decode(vae, latent)) + + +# --------------------------------------------------------------------------- +# Debug helpers +# --------------------------------------------------------------------------- + +def _shape(t): + return list(t.shape) if hasattr(t, "shape") else None + + +def _write_debug(path, record): + with open(path, "a", encoding="utf-8") as f: + f.write(json.dumps(record, ensure_ascii=True) + "\n") + + +def _unique_path(directory, prefix, ext): + i = 1 while True: - path = os.path.join(output_dir, f"{prefix}_{counter:05d}.png") - if not os.path.exists(path): - return path - counter += 1 + p = os.path.join(directory, f"{prefix}_{i:05d}.{ext}") + if not os.path.exists(p): + return p + i += 1 -def _round_to_multiple(value, multiple=16): - return max(multiple, int(round(float(value) / float(multiple))) * multiple) +def _prompt_slug(text, max_words=5): + words = [re.sub(r"[^a-zA-Z0-9]", "", w) for w in re.sub(r"<[^>]+>", "", text or "").split()] + words = [w for w in words if w][:max_words] + return "_".join(words) or "prompt" -def _parse_resolution_preset(preset): - try: - width_text, height_text = str(preset).lower().split("x", 1) - return int(width_text.strip()), int(height_text.strip()) - except Exception as exc: - raise ValueError(f"Invalid resolution preset: {preset}") from exc - - -def _resize_image_to_dimensions(image, width, height, upscale_method): - target_width = _round_to_multiple(width) - target_height = _round_to_multiple(height) - samples = image.movedim(-1, 1) - scaled = comfy.utils.common_upscale(samples, target_width, target_height, upscale_method, "disabled") - return scaled.movedim(1, -1) - - -def _scale_image_to_megapixels(image, target_megapixels, upscale_method): - if target_megapixels <= 0: - return image - - height = int(image.shape[1]) - width = int(image.shape[2]) - total_pixels = max(1.0, float(height * width)) - target_pixels = max(256.0, float(target_megapixels) * 1000000.0) - scale = math.sqrt(target_pixels / total_pixels) - target_width = _round_to_multiple(width * scale) - target_height = _round_to_multiple(height * scale) - - samples = image.movedim(-1, 1) - scaled = comfy.utils.common_upscale(samples, target_width, target_height, upscale_method, "disabled") - return scaled.movedim(1, -1) - - -def _prepare_flux_image(image, resolution_mode, resolution_preset, custom_width, custom_height, target_megapixels, upscale_method): - mode = (resolution_mode or "auto").strip().lower() - - if mode == "auto": - return _scale_image_to_megapixels(image, float(target_megapixels), upscale_method) - - if mode == "match_input": - return _resize_image_to_dimensions(image, int(image.shape[2]), int(image.shape[1]), upscale_method) - - if mode == "preset": - width, height = _parse_resolution_preset(resolution_preset) - return _resize_image_to_dimensions(image, width, height, upscale_method) - - if mode == "custom": - return _resize_image_to_dimensions(image, int(custom_width), int(custom_height), upscale_method) - - return _scale_image_to_megapixels(image, float(target_megapixels), upscale_method) - - -def _build_flux2_latent(width, height, batch_size): - try: - import nodes - - node_cls = getattr(nodes, "EmptyFlux2LatentImage", None) - if node_cls is not None: - instance = node_cls() - if hasattr(instance, "generate"): - return _safe_unpack(instance.generate(width, height, batch_size)) - if hasattr(instance, "execute"): - return _safe_unpack(instance.execute(width, height, batch_size)) - if hasattr(node_cls, "execute"): - return _safe_unpack(node_cls.execute(width, height, batch_size)) - except Exception: - pass - - try: - from comfy_extras.nodes_flux import EmptyFlux2LatentImage - - return _safe_unpack(EmptyFlux2LatentImage.execute(width, height, batch_size)) - except Exception: - return None - - -def _sample_flux2(model, positive, negative, width, height, steps, cfg, sampler_name, seed): - latent_image = _build_flux2_latent(width, height, 1) - if latent_image is None: - raise RuntimeError("EmptyFlux2LatentImage is not available") - - try: - from comfy_extras.nodes_custom_sampler import CFGGuider, KSamplerSelect, RandomNoise, SamplerCustomAdvanced - from comfy_extras.nodes_flux import Flux2Scheduler - - noise = _safe_unpack(RandomNoise.execute(int(seed))) - guider = _safe_unpack(CFGGuider.execute(model, positive, negative, float(cfg))) - sampler = _safe_unpack(KSamplerSelect.execute(str(sampler_name))) - sigmas = _safe_unpack(Flux2Scheduler.execute(int(steps), int(width), int(height))) - return _safe_unpack(SamplerCustomAdvanced.execute(noise, guider, sampler, sigmas, latent_image)) - except Exception as exc: - raise RuntimeError(f"Flux2 advanced sampler path failed: {exc}") from exc - - -def _decode_latent(vae, latent): - latent_samples = latent.get("samples") if isinstance(latent, dict) else latent - decoded = vae.decode(latent_samples) - if decoded.ndim == 5: - decoded = decoded.reshape(-1, decoded.shape[-3], decoded.shape[-2], decoded.shape[-1]) - return decoded.cpu() - +# --------------------------------------------------------------------------- +# Node class +# --------------------------------------------------------------------------- class IAMCCS_FluxKleinMultiGen: DISPLAY_NAME = "IAMCCS Flux Klein Multi-Gen" CATEGORY = "IAMCCS/Flux" FUNCTION = "generate" - RETURN_TYPES = ("IMAGE", "INT") RETURN_NAMES = ("images", "count") @@ -194,38 +183,38 @@ class IAMCCS_FluxKleinMultiGen: def INPUT_TYPES(cls): try: import comfy.samplers - samplers = comfy.samplers.KSampler.SAMPLERS except Exception: samplers = ["euler"] return { "required": { - "model": ("MODEL",), - "clip": ("CLIP",), - "vae": ("VAE",), - "image": ("IMAGE",), + "model": ("MODEL",), + "clip": ("CLIP",), + "vae": ("VAE",), + "image": ("IMAGE",), "multi_prompt": ("STRING", {"forceInput": True}), - "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "control_after_generate": True}), - "steps": ("INT", {"default": 8, "min": 1, "max": 100}), - "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "control_after_generate": True}), + "steps": ("INT", {"default": 8, "min": 1, "max": 100}), + "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}), "sampler_name": (samplers,), }, "optional": { - "separator": ("STRING", {"default": "\\n", "multiline": False}), + "separator": ("STRING", {"default": "\\n", "multiline": False}), "reference_latents_method": ( - ["index_timestep_zero", "offset", "index", "uxo/uno"], - {"default": "index_timestep_zero"}, + ["workflow_default", "index_timestep_zero", "offset", "index", "uxo/uno"], + {"default": "workflow_default"}, ), - "target_megapixels": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 16.0, "step": 0.1}), - "upscale_method": (["lanczos", "nearest-exact", "bilinear", "area", "bicubic"], {"default": "lanczos"}), - "negative_prompt": ("STRING", {"default": "", "multiline": True}), - "output_prefix": ("STRING", {"default": "flux_klein_multi"}), - "save_images": ("BOOLEAN", {"default": False}), - "resolution_mode": (["auto", "match_input", "preset", "custom"], {"default": "auto"}), - "resolution_preset": (RESOLUTION_PRESETS, {"default": "1024x1024"}), - "custom_width": ("INT", {"default": 1024, "min": 16, "max": 8192, "step": 16}), - "custom_height": ("INT", {"default": 1024, "min": 16, "max": 8192, "step": 16}), + "target_megapixels": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 16.0, "step": 0.1}), + "upscale_method": (["nearest-exact", "bilinear", "area", "bicubic", "lanczos"], {"default": "nearest-exact"}), + "negative_prompt": ("STRING", {"default": "", "multiline": True}), + "output_prefix": ("STRING", {"default": "flux_klein_multi"}), + "save_images": ("BOOLEAN", {"default": False}), + "resolution_steps": ("INT", {"default": 1, "min": 1, "max": 256, "step": 1}), + "output_width": ("INT", {"default": 720, "min": 16, "max": 8192, "step": 16}), + "output_height": ("INT", {"default": 1024, "min": 16, "max": 8192, "step": 16}), + "debug_enabled": ("BOOLEAN", {"default": False}), + "debug_prefix": ("STRING", {"default": "flux_klein_debug"}), }, } @@ -241,64 +230,134 @@ class IAMCCS_FluxKleinMultiGen: cfg, sampler_name, separator="\\n", - reference_latents_method="index_timestep_zero", + reference_latents_method="workflow_default", target_megapixels=1.0, - upscale_method="lanczos", + upscale_method="nearest-exact", negative_prompt="", output_prefix="flux_klein_multi", save_images=False, - resolution_mode="auto", - resolution_preset="1024x1024", - custom_width=1024, - custom_height=1024, + resolution_steps=1, + output_width=720, + output_height=1024, + debug_enabled=False, + debug_prefix="flux_klein_debug", ): sep = separator.replace("\\n", "\n") - prompts = [prompt.strip() for prompt in (multi_prompt or "").split(sep) if prompt.strip()] + prompts = [p.strip() for p in (multi_prompt or "").split(sep) if p.strip()] if not prompts: blank = torch.zeros((1, image.shape[1], image.shape[2], 3), dtype=torch.float32) return (blank, 0) - prepared_image = _prepare_flux_image( - image, - resolution_mode, - resolution_preset, - custom_width, - custom_height, - target_megapixels, - upscale_method, + # -- debug setup -- + dbg = None + if debug_enabled: + out_dir = folder_paths.get_output_directory() + dbg = _unique_path(out_dir, debug_prefix, "jsonl") + _write_debug(dbg, { + "event": "run_start", + "prompt_count": len(prompts), + "seed": int(seed), "steps": int(steps), "cfg": float(cfg), + "sampler_name": str(sampler_name), + "reference_latents_method": str(reference_latents_method), + "target_megapixels": float(target_megapixels), + "upscale_method": str(upscale_method), + "resolution_steps": int(resolution_steps), + "output_width": int(output_width), + "output_height": int(output_height), + "input_image_shape": _shape(image), + }) + print(f"[IAMCCS_FluxKleinMultiGen] debug -> {dbg}") + + # ── STEP 1 ImageScaleToTotalPixels ────────────────────────────────── + # Mirrors node: image -> ImageScaleToTotalPixels -> scaled + scaled = _node_ImageScaleToTotalPixels( + image, upscale_method, target_megapixels, resolution_steps ) - reference_latent = vae.encode(prepared_image[:, :, :, :3]) - output_dir = folder_paths.get_output_directory() if save_images else None - results = [] + if dbg: + _write_debug(dbg, {"event": "scaled_image", "shape": _shape(scaled)}) - for index, prompt in enumerate(prompts): - positive = _encode_text(clip, prompt) - negative = _encode_text(clip, negative_prompt) - positive = _apply_reference_latent(positive, reference_latent, reference_latents_method) - negative = _apply_reference_latent(negative, reference_latent, reference_latents_method) + # ── STEP 2 VAEEncode ───────────────────────────────────────────────── + # Mirrors node: scaled + vae -> VAEEncode -> reference_latent + # ONE encode shared between BOTH positive and negative ReferenceLatent nodes + reference_latent = _node_VAEEncode(vae, scaled) + if dbg: + s = reference_latent.get("samples") if isinstance(reference_latent, dict) else reference_latent + _write_debug(dbg, {"event": "reference_latent", "samples_shape": _shape(s)}) - width = int(prepared_image.shape[2]) - height = int(prepared_image.shape[1]) - latent = _sample_flux2( - model, - positive, - negative, - width, - height, - int(steps), - float(cfg), - sampler_name, - int(seed) + index, + # ── STEP 3 Negative conditioning ───────────────────────────────────── + # Mirrors: clip + neg_text -> CLIPTextEncode -> neg_cond + # neg_cond + reference_latent -> ReferenceLatent -> neg_with_ref + neg_cond = _node_CLIPTextEncode(clip, negative_prompt) + neg_with_ref = _node_ReferenceLatent(neg_cond, reference_latent) + + # Optional method override (workflow_default = no override = matches reference pipeline) + _method = None + if reference_latents_method and reference_latents_method != "workflow_default": + _method = "uxo" if ("uxo" in reference_latents_method or "uso" in reference_latents_method) else reference_latents_method + neg_with_ref = node_helpers.conditioning_set_values( + neg_with_ref, {"reference_latents_method": _method} ) - decoded = _decode_latent(vae, latent) + + # ── STEP 4 Sampler components (constant across all prompts) ────────── + # Mirrors: sampler_name -> KSamplerSelect -> sampler + # steps + output_width + output_height -> Flux2Scheduler -> sigmas + # output_width + output_height -> EmptyFlux2LatentImage -> latent_image + sampler_obj = _node_KSamplerSelect(sampler_name) + sigmas = _node_Flux2Scheduler(steps, output_width, output_height) + latent_img = _node_EmptyFlux2LatentImage(output_width, output_height) + + # ── GENERATION LOOP ──────────────────────────────────────────────────── + results = [] + save_dir = folder_paths.get_output_directory() if save_images else None + + for idx, prompt in enumerate(prompts): + current_seed = int(seed) + idx + + # Mirrors: clip + prompt -> CLIPTextEncode -> pos_cond + # pos_cond + reference_latent -> ReferenceLatent -> pos_with_ref + pos_cond = _node_CLIPTextEncode(clip, prompt) + pos_with_ref = _node_ReferenceLatent(pos_cond, reference_latent) + + if _method is not None: + pos_with_ref = node_helpers.conditioning_set_values( + pos_with_ref, {"reference_latents_method": _method} + ) + + # Mirrors: seed -> RandomNoise -> noise + # model + pos_with_ref + neg_with_ref + cfg -> CFGGuider -> guider + # SamplerCustomAdvanced(noise, guider, sampler, sigmas, latent_image) + noise = _node_RandomNoise(current_seed) + guider = _node_CFGGuider(model, pos_with_ref, neg_with_ref, cfg) + output_latent = _node_SamplerCustomAdvanced( + noise, guider, sampler_obj, sigmas, latent_img + ) + + # Mirrors: output_latent + vae -> VAEDecode -> image_out + decoded = _node_VAEDecode(vae, output_latent) + if isinstance(decoded, (tuple, list)): + decoded = decoded[0] + if decoded.ndim == 3: + decoded = decoded.unsqueeze(0) + decoded = decoded.cpu() + + if dbg: + _write_debug(dbg, { + "event": "iteration_result", + "idx": idx, "seed": current_seed, + "prompt": prompt, + "decoded_shape": _shape(decoded), + }) + results.append(decoded) - if save_images and output_dir is not None: - slug = _prompt_to_slug(prompt) - prefix = f"{output_prefix}_{slug}" - path = _unique_png_path(output_dir, prefix) - arr = (decoded[0].numpy() * 255).clip(0, 255).astype(np.uint8) + if save_images and save_dir: + slug = _prompt_slug(prompt) + path = _unique_path(save_dir, f"{output_prefix}_{slug}", "png") + arr = (decoded[0].numpy() * 255).clip(0, 255).astype(np.uint8) _PILImage.fromarray(arr).save(path) - return (torch.cat(results, dim=0), len(prompts)) \ No newline at end of file + if dbg: + _write_debug(dbg, {"event": "run_end", "count": len(results)}) + + return (torch.cat(results, dim=0), len(prompts)) diff --git a/iamccs_wan_lora_schedule.py b/iamccs_wan_lora_schedule.py new file mode 100644 index 0000000..039a0e9 --- /dev/null +++ b/iamccs_wan_lora_schedule.py @@ -0,0 +1,396 @@ +import logging +from collections import OrderedDict + +import comfy.utils +import folder_paths + +from .iamccs_wan_lora_stack import standardize_wan_lora_keys + +_log = logging.getLogger("IAMCCS.LoRA.Schedule") + +_PRESET_OPTIONS = [ + "custom range", # applies from Start gen to End gen (-1 = forever) + "manual_range", # legacy alias kept for prompt validation compatibility + "all generations", # always active + "gen 0 only", # first generation only + "gen 1 onwards", # skip gen 0, apply to all subsequent + "even gens (0,2,4...)", # every even-numbered generation + "odd gens (1,3,5...)", # every odd-numbered generation + "every 2nd gen", # every 2 gens starting from Start gen + "every 3rd gen", # every 3 gens starting from Start gen +] +_DEFAULT_MODEL_TYPE = "flow" +_SLOT_COUNT = 64 +_LINX_TYPE = "IAMCCS_WAN_LORA_LINX" +_AUTO_INDEX_STATE: OrderedDict[tuple[str, str], int] = OrderedDict() +_MAX_AUTO_INDEX_STATE = 64 + + +def _annotate_lora_entry(entry: dict, *, origin: str, generation_index: int, log_prefix: str, prompt_id: str | None, unique_id, slot: int | None = None, rule: str | None = None) -> dict: + annotated = dict(entry) + annotated["_iamccs_lora_origin"] = str(origin or "unknown") + annotated["_iamccs_generation_index"] = int(generation_index) + annotated["_iamccs_schedule_log_prefix"] = str(log_prefix or "WAN LoRA schedule") + annotated["_iamccs_prompt_id"] = str(prompt_id or "") + annotated["_iamccs_schedule_node_id"] = "" if unique_id is None else str(unique_id) + if slot is not None: + annotated["_iamccs_schedule_slot"] = int(slot) + if rule: + annotated["_iamccs_schedule_rule"] = str(rule) + return annotated + + +def _cache_put(cache: OrderedDict, key, value, max_size: int): + cache[key] = value + cache.move_to_end(key) + while len(cache) > max_size: + cache.popitem(last=False) + + +def _get_current_prompt_id() -> str | None: + try: + import server # Imported lazily to avoid hard import issues during discovery. + + prompt_server = getattr(server.PromptServer, "instance", None) + prompt_queue = getattr(prompt_server, "prompt_queue", None) + currently_running = getattr(prompt_queue, "currently_running", None) + if not currently_running: + return None + current = next(iter(currently_running.values())) + if len(current) >= 2: + return str(current[1]) + except Exception: + return None + return None + + +def _next_auto_generation_index(unique_id) -> int: + prompt_id = _get_current_prompt_id() + if not prompt_id or unique_id is None: + return 0 + + key = (prompt_id, str(unique_id)) + current = int(_AUTO_INDEX_STATE.get(key, 0) or 0) + _cache_put(_AUTO_INDEX_STATE, key, current + 1, _MAX_AUTO_INDEX_STATE) + return current + + +def _coerce_lora_stack(value): + if value is None: + return [] + if isinstance(value, list): + return [entry for entry in value if isinstance(entry, dict)] + return [] + + +def _range_matches(index: int, start: int, end: int) -> bool: + start_i = max(0, int(start)) + end_i = int(end) + if end_i < 0: + return index >= start_i + if end_i < start_i: + end_i = start_i + return start_i <= index <= end_i + + +def _preset_matches(index: int, preset: str, start: int, end: int) -> tuple[bool, str]: + start_i = max(0, int(start)) + end_i = int(end) + preset = str(preset or "custom range") + + # Friendly names (new) + underscore aliases (backward compat with old saved workflows) + if preset in ("all generations", "all_generations"): + return True, "all generations" + if preset in ("gen 0 only", "only_first"): + return index == 0, "gen 0 only" + if preset in ("gen 1 onwards", "all_nonfirst"): + return index >= 1, "gen 1 onwards" + if preset in ("even gens (0,2,4...)", "even_generations"): + return (index % 2) == 0, "even gens" + if preset in ("odd gens (1,3,5...)", "odd_generations"): + return (index % 2) == 1, "odd gens" + if preset in ("every 2nd gen", "every_2_from_start"): + return index >= start_i and ((index - start_i) % 2) == 0, f"every 2nd gen (from {start_i})" + if preset in ("every 3rd gen", "every_3_from_start"): + return index >= start_i and ((index - start_i) % 3) == 0, f"every 3rd gen (from {start_i})" + + # custom range / manual_range fallback + if end_i < 0: + return index >= start_i, f"range {start_i}..forever" + return _range_matches(index, start_i, end_i), f"range {start_i}..{max(start_i, end_i)}" + + +def _stack_names(stack: list[dict]) -> str: + names = [] + for entry in stack: + name = str(entry.get("name") or "unnamed") + strength = entry.get("strength", 0.0) + names.append(f"{name}({strength})") + return ", ".join(names) if names else "empty" + + +def _candidate_low_names(name: str) -> list[str]: + source = str(name or "") + if not source or source == "no": + return [] + + replacements = [ + ("_HN_", "_LN_"), + ("-HN_", "-LN_"), + ("_HN-", "_LN-"), + ("_HIGH_", "_LOW_"), + ("-HIGH_", "-LOW_"), + ("_HIGH-", "_LOW-"), + ("HN", "LN"), + ("Hn", "Ln"), + ("hn", "ln"), + ("HIGH", "LOW"), + ("High", "Low"), + ("high", "low"), + ] + out: list[str] = [] + seen: set[str] = set() + for old, new in replacements: + if old not in source: + continue + candidate = source.replace(old, new) + if candidate not in seen: + out.append(candidate) + seen.add(candidate) + if source not in seen: + out.append(source) + return out + + +def _resolve_linx_name(local_name: str, slot: int, linx_payload, available_loras: set[str]) -> str: + local = str(local_name or "no") + if local != "no": + return local + if not isinstance(linx_payload, dict): + return local + + slot_map = linx_payload.get("slot_map") or {} + slot_data = slot_map.get(f"slot_{slot:02d}") or {} + suggested = str(slot_data.get("suggested_low_name") or "") + if suggested and suggested in available_loras: + return suggested + + source_name = str(slot_data.get("name") or "") + for candidate in _candidate_low_names(source_name): + if candidate in available_loras: + return candidate + return local + + +def _build_linx_payload(model_type: str, kwargs: dict, available_loras: set[str]) -> dict: + slot_map = {} + for slot in range(1, _SLOT_COUNT + 1): + name = str(kwargs.get(f"slot_{slot:02d}_lora_name") or "no") + suggested_low_name = "" + for candidate in _candidate_low_names(name): + if candidate in available_loras: + suggested_low_name = candidate + break + slot_map[f"slot_{slot:02d}"] = { + "name": name, + "strength": float(kwargs.get(f"slot_{slot:02d}_strength", 0.0) or 0.0), + "preset": str(kwargs.get(f"slot_{slot:02d}_preset", "custom range") or "custom range"), + "start": int(kwargs.get(f"slot_{slot:02d}_start", slot - 1) or 0), + "end": int(kwargs.get(f"slot_{slot:02d}_end", slot - 1) or 0), + "suggested_low_name": suggested_low_name, + } + return { + "type": _LINX_TYPE, + "model_type": str(model_type or _DEFAULT_MODEL_TYPE), + "slot_map": slot_map, + } + + +def _load_internal_lora_entry(name: str, strength: float, model_type: str, cache: dict) -> dict | None: + if not name or name == "no": + return None + + strength_f = float(strength) + if strength_f == 0.0: + return None + + cache_key = (str(name), str(model_type or _DEFAULT_MODEL_TYPE)) + if cache_key in cache: + state_dict = cache[cache_key] + else: + path = folder_paths.get_full_path_or_raise("loras", name) + state_dict = comfy.utils.load_torch_file(path, safe_load=True) + if str(model_type or _DEFAULT_MODEL_TYPE) != "standard": + state_dict = standardize_wan_lora_keys(state_dict) + cache[cache_key] = state_dict + + return { + "name": str(name), + "strength": strength_f, + "state_dict": state_dict, + } + + +class IAMCCS_WanLoRASchedule: + @classmethod + def INPUT_TYPES(cls): + lora_list = folder_paths.get_filename_list("loras") + ["no"] + required = { + # NOTE: Must NOT use lazy:True here. With lazy:True, ComfyUI caches the + # node output after the first call (gen_idx=None→0) and skips re-execution + # on subsequent loop iterations when only generation_index changes, causing + # the schedule to always behave as if generation_index==0. + "generation_index": ("INT", {"default": 0, "min": 0, "max": 1000000, "step": 1}), + "log_prefix": ("STRING", {"default": "WAN LoRA schedule"}), + "model_type": (["wan2x", "flow", "standard"], {"default": _DEFAULT_MODEL_TYPE}), + } + + optional = { + "default_lora": ("LORA",), + "linx": (_LINX_TYPE,), + } + + hidden = { + "unique_id": "UNIQUE_ID", + } + + for slot in range(1, _SLOT_COUNT + 1): + required[f"slot_{slot:02d}_lora_name"] = (lora_list, {"default": "no"}) + required[f"slot_{slot:02d}_strength"] = ( + "FLOAT", + {"default": 1.0, "min": -5.0, "max": 5.0, "step": 0.01}, + ) + required[f"slot_{slot:02d}_preset"] = (_PRESET_OPTIONS, {"default": "custom range"}) + required[f"slot_{slot:02d}_start"] = ( + "INT", + {"default": slot - 1, "min": 0, "max": 1000000, "step": 1}, + ) + required[f"slot_{slot:02d}_end"] = ( + "INT", + {"default": slot - 1, "min": -1, "max": 1000000, "step": 1}, + ) + + return {"required": required, "optional": optional, "hidden": hidden} + + RETURN_TYPES = ("LORA", "INT", "STRING", _LINX_TYPE) + RETURN_NAMES = ("lora", "active_slots", "report", "linx") + FUNCTION = "schedule" + CATEGORY = "IAMCCS/LoRA" + + @classmethod + def IS_CHANGED(cls, **kwargs): + # This node may use the internal loop-safe auto counter mode; cached reuse + # would freeze the LoRA selection on the first iteration. + return float("nan") + + def schedule(self, generation_index, log_prefix, model_type=_DEFAULT_MODEL_TYPE, default_lora=None, linx=None, unique_id=None, **kwargs): + # Loop-safe mode: when generation_index is left at its hidden default 0, + # use a per-prompt internal counter instead of wiring the easy-use loop + # index into the MODEL branch, which creates a prompt-validation cycle. + generation_source = "input" + if generation_index is None: + generation_index = _next_auto_generation_index(unique_id) + generation_source = "auto:none" + else: + generation_index = int(generation_index) + if generation_index == 0: + generation_index = _next_auto_generation_index(unique_id) + generation_source = "auto:zero" + generation_index = max(0, generation_index) + log_prefix = str(log_prefix or "WAN LoRA schedule") + model_type = str(model_type or _DEFAULT_MODEL_TYPE) + prompt_id = _get_current_prompt_id() + available_loras = set(folder_paths.get_filename_list("loras")) + out = [ + _annotate_lora_entry( + entry, + origin="default", + generation_index=generation_index, + log_prefix=log_prefix, + prompt_id=prompt_id, + unique_id=unique_id, + ) + for entry in _coerce_lora_stack(default_lora) + ] + active_slots = [] + cache: dict = {} + linx_payload = _build_linx_payload(model_type, kwargs, available_loras) + + if out: + _log.info( + "[%s] prompt=%s node=%s generation=%s source=%s | default_lora=%s", + log_prefix, + prompt_id or "unknown", + unique_id if unique_id is not None else "unknown", + generation_index, + generation_source, + _stack_names(out), + ) + + for slot in range(1, _SLOT_COUNT + 1): + preset = kwargs.get(f"slot_{slot:02d}_preset", "custom range") + start = kwargs.get(f"slot_{slot:02d}_start", slot - 1) + end = kwargs.get(f"slot_{slot:02d}_end", slot - 1) + matched, preset_report = _preset_matches(generation_index, preset, start, end) + if not matched: + continue + + slot_stack = [] + slot_name = _resolve_linx_name( + str(kwargs.get(f"slot_{slot:02d}_lora_name") or "no"), + slot, + linx, + available_loras, + ) + entry = _load_internal_lora_entry( + slot_name, + float(kwargs.get(f"slot_{slot:02d}_strength", 0.0) or 0.0), + model_type, + cache, + ) + if entry is not None: + slot_stack.append( + _annotate_lora_entry( + entry, + origin="scheduled", + generation_index=generation_index, + log_prefix=log_prefix, + prompt_id=prompt_id, + unique_id=unique_id, + slot=slot, + rule=preset_report, + ) + ) + + if not slot_stack: + continue + + out.extend(slot_stack) + active_slots.append(f"slot_{slot:02d}:{preset_report} => {_stack_names(slot_stack)}") + + report = ( + f"generation_index={generation_index} | active={' + '.join(active_slots) if active_slots else 'default_only'} | " + f"entries={len(out)}" + ) + + if out: + _log.info("[%s] %s", log_prefix, report) + _log.info( + "[%s] generation=%s resolved_stack=%s", + log_prefix, + generation_index, + _stack_names(out), + ) + else: + _log.warning("[%s] No LoRA active | generation_index=%s", log_prefix, generation_index) + + return (out, len(active_slots), report, linx_payload) + + +NODE_CLASS_MAPPINGS = { + "IAMCCS_WanLoRASchedule": IAMCCS_WanLoRASchedule, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "IAMCCS_WanLoRASchedule": "LoRA Schedule (WAN, ranged)", +} diff --git a/iamccs_wan_lora_stack.py b/iamccs_wan_lora_stack.py index 61f2722..c07c37e 100644 --- a/iamccs_wan_lora_stack.py +++ b/iamccs_wan_lora_stack.py @@ -5,11 +5,88 @@ # =============================================================== import logging +from collections import OrderedDict + import comfy.utils import comfy.sd import folder_paths +_STATE_DICT_CACHE: OrderedDict[tuple[str, str], dict] = OrderedDict() +_PATCHED_MODEL_CACHE: OrderedDict[tuple[int, tuple[tuple[str, float], ...]], object] = OrderedDict() +_MAX_STATE_DICT_CACHE = 16 +_MAX_PATCHED_MODEL_CACHE = 8 + + +def _lora_stack_debug_summary(lora) -> str: + if not lora: + return "empty" + + parts = [] + for entry in lora: + name = str(entry.get("name") or "unnamed") + strength = float(entry.get("strength", 0.0) or 0.0) + origin = str(entry.get("_iamccs_lora_origin") or "manual") + generation_index = entry.get("_iamccs_generation_index") + slot = entry.get("_iamccs_schedule_slot") + rule = str(entry.get("_iamccs_schedule_rule") or "") + + extras = [origin] + if generation_index is not None: + extras.append(f"gen={generation_index}") + if slot is not None: + extras.append(f"slot={int(slot):02d}") + if rule: + extras.append(rule) + + parts.append(f"{name}({strength}) [{' | '.join(extras)}]") + + return "; ".join(parts) + + +def _lora_stack_debug_context(lora) -> str: + if not lora: + return "" + + prompt_ids = sorted({str(entry.get("_iamccs_prompt_id") or "") for entry in lora if entry.get("_iamccs_prompt_id")}) + schedule_nodes = sorted({str(entry.get("_iamccs_schedule_node_id") or "") for entry in lora if entry.get("_iamccs_schedule_node_id")}) + schedule_names = sorted({str(entry.get("_iamccs_schedule_log_prefix") or "") for entry in lora if entry.get("_iamccs_schedule_log_prefix")}) + generation_indexes = sorted({int(entry.get("_iamccs_generation_index")) for entry in lora if entry.get("_iamccs_generation_index") is not None}) + + parts = [] + if prompt_ids: + parts.append(f"prompt={','.join(prompt_ids)}") + if schedule_nodes: + parts.append(f"schedule_node={','.join(schedule_nodes)}") + if schedule_names: + parts.append(f"schedule={','.join(schedule_names)}") + if generation_indexes: + parts.append(f"generation={','.join(str(value) for value in generation_indexes)}") + return " | ".join(parts) + + +def _cache_put(cache: OrderedDict, key, value, max_size: int): + cache[key] = value + cache.move_to_end(key) + while len(cache) > max_size: + cache.popitem(last=False) + + +def _state_dict_for_lora(name: str, model_type: str) -> dict: + cache_key = (str(name), str(model_type or "flow")) + cached = _STATE_DICT_CACHE.get(cache_key) + if cached is not None: + _STATE_DICT_CACHE.move_to_end(cache_key) + return cached + + path = folder_paths.get_full_path_or_raise("loras", name) + sd = comfy.utils.load_torch_file(path, safe_load=True) + if model_type != "standard": + sd = standardize_wan_lora_keys(sd) + _cache_put(_STATE_DICT_CACHE, cache_key, sd, _MAX_STATE_DICT_CACHE) + return sd + + # --- Log Filter per sopprimere spam img_* e diff_m keys --- class SuppressOptionalKeysFilter(logging.Filter): """Filtra i log 'lora key not loaded' per chiavi opzionali (img_*, diff_m, ecc.)""" @@ -139,10 +216,7 @@ class IAMCCS_WanLoRAStack: # Skip if "no" selected or strength is 0 if not name or name == "no" or strength == 0.0: continue - path = folder_paths.get_full_path_or_raise("loras", name) - sd = comfy.utils.load_torch_file(path, safe_load=True) - if model_type != "standard": - sd = standardize_wan_lora_keys(sd) + sd = _state_dict_for_lora(name, model_type) loras.append({"name": name, "strength": strength, "state_dict": sd}) # Concatena la stack LORA opzionale se fornita @@ -182,8 +256,26 @@ class IAMCCS_ModelWithLoRA: if not lora: return (model,) + stack_context = _lora_stack_debug_context(lora) + stack_summary = _lora_stack_debug_summary(lora) + signature = tuple((str(entry.get("name") or ""), float(entry.get("strength", 0.0) or 0.0)) for entry in lora) + cache_key = (id(model), signature) + cached_model = _PATCHED_MODEL_CACHE.get(cache_key) + if cached_model is not None: + _PATCHED_MODEL_CACHE.move_to_end(cache_key) + if stack_context: + logging.info(f"[IAMCCS_ModelWithLoRA] ♻ cache hit: {len(signature)} LoRA(s) reused | {stack_context}") + else: + logging.info(f"[IAMCCS_ModelWithLoRA] ♻ cache hit: {len(signature)} LoRA(s) reused") + logging.info(f"[IAMCCS_ModelWithLoRA] active_stack={stack_summary}") + return (cached_model,) + model_out = model + if stack_context: + logging.info(f"[IAMCCS_ModelWithLoRA] apply request | {stack_context}") + logging.info(f"[IAMCCS_ModelWithLoRA] active_stack={stack_summary}") + # Installa filtro per sopprimere spam di chiavi opzionali (img_*, diff_m, ecc.) logger = logging.getLogger() optional_filter = SuppressOptionalKeysFilter() @@ -201,6 +293,8 @@ class IAMCCS_ModelWithLoRA: keys_types = ", ".join(sorted(optional_filter.suppressed_keys)) logging.info(f"[IAMCCS_ModelWithLoRA] ℹ {optional_filter.suppressed_count} optional keys not present in LORA ({keys_types})") + _cache_put(_PATCHED_MODEL_CACHE, cache_key, model_out, _MAX_PATCHED_MODEL_CACHE) + finally: # Rimuovi filtro logger.removeFilter(optional_filter) diff --git a/web/iamccs_qwen_multigen_compat.js b/web/iamccs_qwen_multigen_compat.js index 59e95e2..370f17e 100644 --- a/web/iamccs_qwen_multigen_compat.js +++ b/web/iamccs_qwen_multigen_compat.js @@ -3,7 +3,7 @@ import { app } from "../../scripts/app.js"; const SEED_MODES = new Set(["fixed", "randomize", "increment", "decrement"]); const NODE_DEFAULTS = { IAMCCS_QwenMultiGen: [0, "randomize", 4, 1, "euler", "simple", 1, "\\n", "index_timestep_zero", "qwen_multi", true], - IAMCCS_FluxKleinMultiGen: [0, "randomize", 8, 1, "euler", "\\n", "index_timestep_zero", 1, "lanczos", "", "flux_klein_multi", false], + IAMCCS_FluxKleinMultiGen: [0, "randomize", 8, 1, "euler", "\\n", "workflow_default", 1, "nearest-exact", "", "flux_klein_multi", false, 1, 720, 1024, false, "flux_klein_debug"], }; function isFiniteNumber(value) { diff --git a/web/iamccs_wan_lora_schedule_ui.js b/web/iamccs_wan_lora_schedule_ui.js new file mode 100644 index 0000000..c201a23 --- /dev/null +++ b/web/iamccs_wan_lora_schedule_ui.js @@ -0,0 +1,1412 @@ +import { app } from "../../scripts/app.js"; + +const NODE_NAME = "IAMCCS_WanLoRASchedule"; +const MODE_WIDGET = "_iamccs_schedule_mode"; +const LOAD_WIDGET = "_iamccs_schedule_load_preset"; +const SAVE_WIDGET = "_iamccs_schedule_save_preset"; +const ADD_WIDGET = "_iamccs_schedule_add_slot"; +const DELETE_WIDGET = "_iamccs_schedule_delete_slot"; +const RESET_WIDGET = "_iamccs_schedule_reset_rules"; +const MODE_PROP = "iamccs_schedule_ui_mode"; +const SLOTS_PROP = "iamccs_schedule_visible_slots"; +const AUTO_PROP = "iamccs_schedule_auto_synced"; +const INIT_PROP = "iamccs_schedule_inited_slots"; +const DROP_HOVER_TIMEOUT_MS = 250; +const PRESET_KIND = "iamccs_wan_lora_schedule_preset"; +const PRESET_VERSION = 1; +const MAX_SLOTS = 64; +const LEGACY_MAX_SLOTS = 6; +const UI_WIDGET_COUNT = 4; +const EXPECTED_SERIALIZED_VALUES = 3 + (MAX_SLOTS * 5); +const LEGACY_EXPECTED_SERIALIZED_VALUES = 3 + (LEGACY_MAX_SLOTS * 5); +const SIMPLE_WIDTH = 480; +const ADVANCED_WIDTH = 620; +const BASE_HEIGHT = 230; +const SIMPLE_SLOT_HEIGHT = 54; +const ADVANCED_SLOT_HEIGHT = 132; +const EXTRA_HEIGHT = 30; +const SLOT_BLOCK = 5; +const SLOTS_START_IDX = 3; +const VISIBLE_PRESET_OPTIONS = ["custom range", "all generations", "gen 0 only", "gen 1 onwards"]; + +const SLOT_META = [ + { slot: 1, label: "Gen 0", preset: "gen 0 only", start: 0, end: 0 }, + { slot: 2, label: "Gen 1", preset: "custom range", start: 1, end: 1 }, + { slot: 3, label: "Gen 2+", preset: "custom range", start: 2, end: -1 }, +]; + +function getGraph() { + return app?.canvas?.getCurrentGraph?.() || app?.graph || null; +} + +function getGraphLink(graph, linkId) { + if (!graph || linkId == null) return null; + try { + if (typeof graph.links?.get === "function") return graph.links.get(linkId) || graph.links.get(String(linkId)) || null; + return graph.links?.[linkId] || graph.links?.[String(linkId)] || null; + } catch { + return null; + } +} + +function getGraphNode(graph, nodeId) { + if (!graph || nodeId == null) return null; + try { + return graph.getNodeById?.(nodeId) || null; + } catch { + return null; + } +} + +function getWidget(node, name) { + return node?.widgets?.find((widget) => widget?.name === name || widget?.label === name) || null; +} + +function markUiWidgetNonSerializable(widget) { + if (!widget) return widget; + widget.options = { ...(widget.options || {}), serialize: false }; + widget.serializeValue = () => undefined; + widget._iamccsUiOnly = true; + return widget; +} + +function isModeValue(value) { + return value === "simple" || value === "advanced"; +} + +function isModelTypeValue(value) { + return value === "wan2x" || value === "flow" || value === "standard"; +} + +// Correct serialized count: generation_index(1) + log_prefix(1) + model_type(1) + 6 slots × 5 = 33 +// (generation_index stays as a widget even when wired, in ComfyUI's serialization model) +const EXPECTED_SERIALIZED_VALUES_CONNECTED = EXPECTED_SERIALIZED_VALUES; // 33 – gen_idx present as wired widget + +function slotCountFromSerializedLength(length) { + if (typeof length !== "number" || length < 3) return 0; + return Math.max(0, Math.floor((length - 3) / SLOT_BLOCK)); +} + +function sanitizePresetValue(value) { + const raw = String(value || "custom range"); + const aliases = { + manual_range: "custom range", + even_generations: "custom range", + odd_generations: "custom range", + every_2_from_start: "custom range", + every_3_from_start: "custom range", + "even gens (0,2,4...)": "custom range", + "odd gens (1,3,5...)": "custom range", + "every 2nd gen": "custom range", + "every 3rd gen": "custom range", + }; + return aliases[raw] || raw; +} + +function normalizeLegacyWidgetValues(config) { + const values = config?.widgets_values; + if (!Array.isArray(values)) return; + + let modeValue = null; + let normalized = null; + + // Pattern A (37 values): pre-fix format where 4 UI widgets were serialized at front, + // gen_idx at [-2] and log_prefix at [-1]. + // "simple"/null/null/null model_type ...30 slots... gen_idx log_prefix + if ( + values.length === LEGACY_EXPECTED_SERIALIZED_VALUES + UI_WIDGET_COUNT + && isModeValue(values[0]) + && isModelTypeValue(values[UI_WIDGET_COUNT]) + ) { + modeValue = values[0]; + const generationIndex = values[values.length - 2] ?? 0; + const logPrefix = values[values.length - 1] ?? "WAN LoRA schedule"; + const modelType = values[UI_WIDGET_COUNT] ?? "flow"; + const slotValues = values.slice(UI_WIDGET_COUNT + 1, -2); + // Correct order: gen_idx log_prefix model_type ...30 slots... + normalized = [generationIndex, logPrefix, modelType, ...slotValues]; + + // Pattern B (33 values): intermediate format where model_type was widget[0], + // gen_idx at [-2] and log_prefix at [-1]. + // model_type ...30 slots... gen_idx log_prefix + } else if ( + values.length === LEGACY_EXPECTED_SERIALIZED_VALUES + && isModelTypeValue(values[0]) + && typeof values[values.length - 1] === "string" + ) { + const generationIndex = values[values.length - 2] ?? 0; + const logPrefix = values[values.length - 1] ?? "WAN LoRA schedule"; + const modelType = values[0] ?? "flow"; + const slotValues = values.slice(1, -2); + normalized = [generationIndex, logPrefix, modelType, ...slotValues]; + } + + // Determine what to work on: use normalized if we reordered, else work on original (in-place) + // for preset-migration-only case (already-correct 33-value arrays with legacy preset names). + const inferredSlots = slotCountFromSerializedLength(values.length); + const target = normalized ?? (inferredSlots > 0 ? values : null); + if (!target) return; + + // Migrate legacy internal preset names → current _PRESET_OPTIONS names. + // "manual_range" was an internal alias for "custom range" in early builds. + // ComfyUI server REJECTS it as an invalid combo value because it is not in _PRESET_OPTIONS. + let presetPatched = false; + for (let s = 0; s < slotCountFromSerializedLength(target.length); s++) { + const presetIdx = SLOTS_START_IDX + s * SLOT_BLOCK + 2; + const raw = target[presetIdx]; + const normalizedPreset = sanitizePresetValue(raw); + if (raw != null && normalizedPreset !== raw) { + target[presetIdx] = normalizedPreset; + presetPatched = true; + } + } + + if (!normalized && !presetPatched) return; // nothing changed + + if (normalized) { + config.widgets_values = normalized; + } else if (presetPatched) { + // values was modified in-place above (same array reference) + config.widgets_values = values; + } + config.properties = config.properties || {}; + if (isModeValue(modeValue) && !config.properties[MODE_PROP]) { + config.properties[MODE_PROP] = modeValue; + } +} + +function hideWidget(widget) { + if (!widget) return; + // Use `in` check so we save the REAL original only once, never overwrite with [0,0] + if (!("_iamccsOriginalComputeSize" in widget)) { + widget._iamccsOriginalComputeSize = widget.computeSize ?? null; + } + widget.hidden = true; + widget.disabled = true; + widget.computeSize = () => [0, 0]; +} + +function showWidget(widget) { + if (!widget) return; + widget.hidden = false; + widget.disabled = false; + if (widget._iamccsOriginalComputeSize) { + widget.computeSize = widget._iamccsOriginalComputeSize; + } else { + delete widget.computeSize; + } +} + +function setWidgetValue(node, name, value) { + const widget = getWidget(node, name); + if (!widget) return false; + widget.value = value; + try { + widget.callback?.(value, app.canvas, node); + } catch {} + return true; +} + +function ensureBoxShape(node) { + try { + if (typeof LiteGraph !== "undefined" && LiteGraph?.BOX_SHAPE != null) { + node.shape = LiteGraph.BOX_SHAPE; + } else { + node.shape = 0; + } + } catch {} +} + +function getVisibleSlots(node) { + node.properties = node.properties || {}; + const raw = Number(node.properties[SLOTS_PROP] ?? 3); + return Math.max(3, Math.min(MAX_SLOTS, Math.trunc(raw) || 3)); +} + +function setVisibleSlots(node, value) { + node.properties = node.properties || {}; + node.properties[SLOTS_PROP] = Math.max(3, Math.min(MAX_SLOTS, Math.trunc(Number(value) || 3))); +} + +function slotMeta(slot) { + const base = SLOT_META.find((item) => item.slot === slot); + if (base) return base; + return { slot, label: `Gen ${slot - 1}`, preset: "custom range", start: slot - 1, end: slot - 1 }; +} + +function ensureSlotDefaults(node, slot) { + node.properties = node.properties || {}; + node.properties[INIT_PROP] = node.properties[INIT_PROP] || {}; + if (node.properties[INIT_PROP][slot]) return; // already set; don't clobber user values + const meta = slotMeta(slot); + const prefix = `slot_${String(slot).padStart(2, "0")}`; + setWidgetValue(node, `${prefix}_preset`, meta.preset); + setWidgetValue(node, `${prefix}_start`, meta.start); + setWidgetValue(node, `${prefix}_end`, meta.end); + const strengthWidget = getWidget(node, `${prefix}_strength`); + if (strengthWidget && (strengthWidget.value == null || Number(strengthWidget.value) === 0)) { + strengthWidget.value = 1; + } + node.properties[INIT_PROP][slot] = true; +} + +function ensureVisibleDefaults(node) { + for (let slot = 1; slot <= getVisibleSlots(node); slot += 1) { + ensureSlotDefaults(node, slot); + } +} + +function resetScheduleRules(node) { + node.properties = node.properties || {}; + node.properties[INIT_PROP] = {}; + node.properties[AUTO_PROP] = {}; + setVisibleSlots(node, 3); + for (let slot = 1; slot <= MAX_SLOTS; slot += 1) { + const meta = slotMeta(slot); + const prefix = `slot_${String(slot).padStart(2, "0")}`; + setWidgetValue(node, `${prefix}_lora_name`, "no"); + setWidgetValue(node, `${prefix}_strength`, 1); + setWidgetValue(node, `${prefix}_preset`, meta.preset); + setWidgetValue(node, `${prefix}_start`, meta.start); + setWidgetValue(node, `${prefix}_end`, meta.end); + node.properties[INIT_PROP][slot] = true; + } + clearDropHoverSlot(node); +} + +function notifyScheduleUi(message, isError = false) { + const prefix = "[IAMCCS_WanLoRASchedule UI]"; + if (isError) { + console.error(prefix, message); + } else { + console.info(prefix, message); + } +} + +function isAbortError(error) { + return error?.name === "AbortError" || /aborted/i.test(String(error?.message || "")); +} + +function sanitizeFilename(value) { + return String(value || "") + .trim() + .replace(/[<>:"/\\|?*\x00-\x1F]+/g, "_") + .replace(/\s+/g, "_") + .replace(/^_+|_+$/g, "") + .slice(0, 120); +} + +function getScheduleSlotRecord(node, slot) { + const meta = slotMeta(slot); + const prefix = `slot_${String(slot).padStart(2, "0")}`; + return { + slot, + lora_name: String(getWidget(node, `${prefix}_lora_name`)?.value || "no"), + strength: Number(getWidget(node, `${prefix}_strength`)?.value ?? 1) || 1, + preset: sanitizePresetValue(getWidget(node, `${prefix}_preset`)?.value || meta.preset), + start: Math.trunc(Number(getWidget(node, `${prefix}_start`)?.value ?? meta.start) || meta.start), + end: Math.trunc(Number(getWidget(node, `${prefix}_end`)?.value ?? meta.end) || meta.end), + }; +} + +function buildSchedulePreset(node) { + const mode = isModeValue(node?.properties?.[MODE_PROP]) ? node.properties[MODE_PROP] : "simple"; + const modelType = String(getWidget(node, "model_type")?.value || "flow"); + const logPrefix = String(getWidget(node, "log_prefix")?.value || "WAN LoRA schedule"); + return { + kind: PRESET_KIND, + version: PRESET_VERSION, + node_type: NODE_NAME, + exported_at: new Date().toISOString(), + mode, + visible_slots: getVisibleSlots(node), + model_type: isModelTypeValue(modelType) ? modelType : "flow", + log_prefix: logPrefix, + slots: Array.from({ length: MAX_SLOTS }, (_, index) => getScheduleSlotRecord(node, index + 1)), + }; +} + +function normalizeLoadedSlotRecord(rawSlot, fallbackSlot) { + const slotNum = Math.max(1, Math.min(MAX_SLOTS, Math.trunc(Number(rawSlot?.slot ?? fallbackSlot) || fallbackSlot))); + const meta = slotMeta(slotNum); + return { + slot: slotNum, + lora_name: String(rawSlot?.lora_name || rawSlot?.loraName || "no"), + strength: Number(rawSlot?.strength ?? 1) || 1, + preset: sanitizePresetValue(rawSlot?.preset || meta.preset), + start: Math.trunc(Number(rawSlot?.start ?? meta.start) || meta.start), + end: Math.trunc(Number(rawSlot?.end ?? meta.end) || meta.end), + }; +} + +async function applySchedulePreset(node, payload) { + if (!payload || typeof payload !== "object") { + throw new Error("Preset file is empty or invalid."); + } + if (payload.kind && payload.kind !== PRESET_KIND) { + throw new Error(`Unsupported preset kind: ${payload.kind}`); + } + + const slotMap = new Map(); + if (Array.isArray(payload.slots)) { + for (let index = 0; index < payload.slots.length; index += 1) { + const normalized = normalizeLoadedSlotRecord(payload.slots[index], index + 1); + slotMap.set(normalized.slot, normalized); + } + } else if (payload.slots && typeof payload.slots === "object") { + for (const [key, value] of Object.entries(payload.slots)) { + const numericSlot = Math.trunc(Number(String(key).replace(/[^0-9]/g, "")) || 0); + if (!numericSlot) continue; + const normalized = normalizeLoadedSlotRecord(value, numericSlot); + slotMap.set(normalized.slot, normalized); + } + } + + const nextMode = isModeValue(payload.mode) ? payload.mode : "simple"; + const nextVisibleSlots = Math.max(3, Math.min(MAX_SLOTS, Math.trunc(Number(payload.visible_slots ?? payload.visibleSlots) || 3))); + const nextModelType = isModelTypeValue(payload.model_type || payload.modelType) ? String(payload.model_type || payload.modelType) : "flow"; + const nextLogPrefix = String(payload.log_prefix || payload.logPrefix || "WAN LoRA schedule"); + + resetScheduleRules(node); + node.properties = node.properties || {}; + node.properties[MODE_PROP] = nextMode; + setVisibleSlots(node, nextVisibleSlots); + + const modeWidget = getWidget(node, MODE_WIDGET); + if (modeWidget) modeWidget.value = nextMode; + const modelTypeWidget = getWidget(node, "model_type"); + if (modelTypeWidget) modelTypeWidget.value = nextModelType; + const logPrefixWidget = getWidget(node, "log_prefix"); + if (logPrefixWidget) logPrefixWidget.value = nextLogPrefix; + + node.properties[INIT_PROP] = node.properties[INIT_PROP] || {}; + for (let slot = 1; slot <= MAX_SLOTS; slot += 1) { + const record = slotMap.get(slot) || normalizeLoadedSlotRecord({}, slot); + const prefix = `slot_${String(slot).padStart(2, "0")}`; + const loraWidget = getWidget(node, `${prefix}_lora_name`); + const strengthWidget = getWidget(node, `${prefix}_strength`); + const presetWidget = getWidget(node, `${prefix}_preset`); + const startWidget = getWidget(node, `${prefix}_start`); + const endWidget = getWidget(node, `${prefix}_end`); + if (loraWidget) loraWidget.value = record.lora_name; + if (strengthWidget) strengthWidget.value = record.strength; + if (presetWidget) presetWidget.value = record.preset; + if (startWidget) startWidget.value = record.start; + if (endWidget) endWidget.value = record.end; + node.properties[INIT_PROP][slot] = true; + } + + clearDropHoverSlot(node); + applyLayout(node); + await syncAllLinkedLowNodes(node); + node.setDirtyCanvas?.(true, true); +} + +function getSuggestedPresetFilename(node) { + const logPrefix = String(getWidget(node, "log_prefix")?.value || node?.title || "wan-lora-schedule"); + const filename = sanitizeFilename(logPrefix) || "wan-lora-schedule"; + return `${filename}.json`; +} + +function downloadTextFile(filename, text) { + const blob = new Blob([text], { type: "application/json" }); + const url = URL.createObjectURL(blob); + const link = document.createElement("a"); + link.href = url; + link.download = filename; + link.style.display = "none"; + document.body.appendChild(link); + link.click(); + link.remove(); + URL.revokeObjectURL(url); +} + +async function saveSchedulePresetAs(node) { + const preset = buildSchedulePreset(node); + const text = JSON.stringify(preset, null, 2); + const suggestedName = getSuggestedPresetFilename(node); + + if (typeof window !== "undefined" && typeof window.showSaveFilePicker === "function") { + const handle = await window.showSaveFilePicker({ + suggestedName, + types: [{ + description: "IAMCCS Wan LoRA Schedule preset", + accept: { "application/json": [".json"] }, + }], + }); + const writable = await handle.createWritable(); + await writable.write(text); + await writable.close(); + } else { + downloadTextFile(suggestedName, text); + } + + notifyScheduleUi(`Preset saved: ${suggestedName}`); +} + +async function pickPresetTextFile() { + if (typeof window !== "undefined" && typeof window.showOpenFilePicker === "function") { + const [handle] = await window.showOpenFilePicker({ + multiple: false, + types: [{ + description: "IAMCCS Wan LoRA Schedule preset", + accept: { "application/json": [".json"] }, + }], + }); + const file = await handle.getFile(); + return await file.text(); + } + + return await new Promise((resolve, reject) => { + const input = document.createElement("input"); + input.type = "file"; + input.accept = ".json,application/json"; + input.style.display = "none"; + input.addEventListener("change", async () => { + try { + const file = input.files?.[0]; + if (!file) { + reject(new DOMException("User aborted file selection.", "AbortError")); + return; + } + resolve(await file.text()); + } catch (error) { + reject(error); + } finally { + input.remove(); + } + }, { once: true }); + document.body.appendChild(input); + input.click(); + }); +} + +async function loadSchedulePreset(node) { + const text = await pickPresetTextFile(); + const payload = JSON.parse(text); + await applySchedulePreset(node, payload); + notifyScheduleUi("Preset loaded."); +} + +function applyLabels(node) { + const modeWidget = getWidget(node, MODE_WIDGET); + if (modeWidget) modeWidget.label = "Mode"; + const modelTypeWidget = getWidget(node, "model_type"); + if (modelTypeWidget) modelTypeWidget.label = "Model type"; + + for (const input of node.inputs || []) { + if (input?.name === "default_lora") input.label = "Always-on LoRA"; + if (input?.name === "linx") input.label = "Linx in"; + if (input?.name === "generation_index") input.label = "Generation"; + } + for (const output of node.outputs || []) { + if (output?.name === "linx") output.label = "Linx out"; + } + + for (let slot = 1; slot <= MAX_SLOTS; slot += 1) { + const meta = slotMeta(slot); + const prefix = `slot_${String(slot).padStart(2, "0")}`; + const loraWidget = getWidget(node, `${prefix}_lora_name`); + const strengthWidget = getWidget(node, `${prefix}_strength`); + const presetWidget = getWidget(node, `${prefix}_preset`); + const startWidget = getWidget(node, `${prefix}_start`); + const endWidget = getWidget(node, `${prefix}_end`); + if (loraWidget) loraWidget.label = `${meta.label} LoRA`; + if (strengthWidget) strengthWidget.label = `${meta.label} Strength`; + if (presetWidget) presetWidget.label = `${meta.label} Apply when`; + if (startWidget) startWidget.label = `${meta.label} Start gen`; + if (endWidget) endWidget.label = `${meta.label} End gen (-1 = forever)`; + } +} + +function applyPresetOptions(node) { + for (let slot = 1; slot <= MAX_SLOTS; slot += 1) { + const presetWidget = getWidget(node, `slot_${String(slot).padStart(2, "0")}_preset`); + if (!presetWidget) continue; + presetWidget.options = { ...(presetWidget.options || {}), values: VISIBLE_PRESET_OPTIONS }; + const current = sanitizePresetValue(presetWidget.value); + if (!VISIBLE_PRESET_OPTIONS.includes(current)) { + presetWidget.value = "custom range"; + } else if (current !== presetWidget.value) { + presetWidget.value = current; + } + } +} + +function ensureModeWidget(node) { + if (getWidget(node, MODE_WIDGET)) return; + node.properties = node.properties || {}; + if (!node.properties[MODE_PROP]) { + node.properties[MODE_PROP] = "simple"; + } + const widget = node.addWidget( + "combo", + MODE_WIDGET, + node.properties[MODE_PROP], + (value) => { + node.properties[MODE_PROP] = String(value || "simple"); + applyLayout(node); + syncAllLinkedLowNodes(node); + node.setDirtyCanvas(true, true); + }, + { values: ["simple", "advanced"], serialize: false } + ); + widget.label = "Mode"; + markUiWidgetNonSerializable(widget); +} + +function ensureActionWidgets(node) { + if (!getWidget(node, LOAD_WIDGET)) { + const widget = node.addWidget("button", "Load Preset", null, () => { + void loadSchedulePreset(node).catch((error) => { + if (isAbortError(error)) return; + notifyScheduleUi(`Load preset failed: ${error?.message || error}`, true); + }); + }); + widget.name = LOAD_WIDGET; + markUiWidgetNonSerializable(widget); + } + if (!getWidget(node, SAVE_WIDGET)) { + const widget = node.addWidget("button", "Save As...", null, () => { + void saveSchedulePresetAs(node).catch((error) => { + if (isAbortError(error)) return; + notifyScheduleUi(`Save preset failed: ${error?.message || error}`, true); + }); + }); + widget.name = SAVE_WIDGET; + markUiWidgetNonSerializable(widget); + } + if (!getWidget(node, ADD_WIDGET)) { + const widget = node.addWidget("button", "+ Add Slot", null, () => { + const newSlot = getVisibleSlots(node) + 1; + node.properties[INIT_PROP] = node.properties[INIT_PROP] || {}; + delete node.properties[INIT_PROP][newSlot]; // allow fresh defaults for new slot + setVisibleSlots(node, newSlot); + ensureVisibleDefaults(node); + applyLayout(node); + syncAllLinkedLowNodes(node); + node.setDirtyCanvas(true, true); + }); + widget.name = ADD_WIDGET; + markUiWidgetNonSerializable(widget); + } + if (!getWidget(node, DELETE_WIDGET)) { + const widget = node.addWidget("button", "- Delete Slot", null, () => { + setVisibleSlots(node, getVisibleSlots(node) - 1); + applyLayout(node); + syncAllLinkedLowNodes(node); + node.setDirtyCanvas(true, true); + }); + widget.name = DELETE_WIDGET; + markUiWidgetNonSerializable(widget); + } + if (!getWidget(node, RESET_WIDGET)) { + const widget = node.addWidget("button", "Reset Rules", null, () => { + resetScheduleRules(node); + applyLayout(node); + syncAllLinkedLowNodes(node); + node.setDirtyCanvas(true, true); + }); + widget.name = RESET_WIDGET; + markUiWidgetNonSerializable(widget); + } +} + +function getOutputIndexByName(node, name) { + return (node?.outputs || []).findIndex((output) => output?.name === name); +} + +function getInputIndexByName(node, name) { + return (node?.inputs || []).findIndex((input) => input?.name === name); +} + +function getLinkedLinxTargets(node) { + const graph = getGraph(); + if (!graph) return []; + const outIndex = getOutputIndexByName(node, "linx"); + if (outIndex < 0) return []; + const output = node.outputs?.[outIndex]; + const links = Array.isArray(output?.links) ? output.links : []; + const out = []; + for (const linkId of links) { + const link = getGraphLink(graph, linkId); + const target = getGraphNode(graph, link?.target_id); + if (target) out.push(target); + } + return out; +} + +function candidateLowNames(name) { + const source = String(name || ""); + if (!source || source === "no") return []; + const pairs = [ + ["_HN_", "_LN_"], + ["-HN_", "-LN_"], + ["_HN-", "_LN-"], + ["_HIGH_", "_LOW_"], + ["-HIGH_", "-LOW_"], + ["_HIGH-", "_LOW-"], + ["HN", "LN"], + ["Hn", "Ln"], + ["hn", "ln"], + ["HIGH", "LOW"], + ["High", "Low"], + ["high", "low"], + ]; + const out = []; + const seen = new Set(); + for (const [oldValue, newValue] of pairs) { + if (!source.includes(oldValue)) continue; + const candidate = source.replace(oldValue, newValue); + if (!seen.has(candidate)) { + out.push(candidate); + seen.add(candidate); + } + } + if (!seen.has(source)) { + out.push(source); + } + return out; +} + +function syncLowNodePresentation(highNode, lowNode) { + lowNode.properties = lowNode.properties || {}; + lowNode.properties[MODE_PROP] = highNode?.properties?.[MODE_PROP] || "simple"; + setVisibleSlots(lowNode, getVisibleSlots(highNode)); + lowNode.flags = lowNode.flags || {}; + lowNode.flags.collapsed = !!highNode?.flags?.collapsed; +} + +function normalizePathSeparators(value) { + return String(value || "").replace(/\\/g, "/").replace(/^file:\/\//i, "").trim(); +} + +function basenameOfPath(value) { + const normalized = normalizePathSeparators(value); + const parts = normalized.split("/"); + return String(parts[parts.length - 1] || "").trim(); +} + +function collectDroppedLoraCandidates(event) { + const out = []; + const seen = new Set(); + const push = (value) => { + const normalized = normalizePathSeparators(value); + if (!normalized) return; + const base = basenameOfPath(normalized); + const variants = [normalized, base]; + for (const variant of variants) { + const key = String(variant || "").trim(); + if (!key || !key.toLowerCase().endsWith(".safetensors") || seen.has(key)) continue; + seen.add(key); + out.push(key); + } + }; + + const fileList = Array.from(event?.dataTransfer?.files || []); + for (const file of fileList) { + push(String(file?.name || "")); + push(String(file?.path || "")); + } + + const payloads = [ + String(event?.dataTransfer?.getData?.("text/uri-list") || ""), + String(event?.dataTransfer?.getData?.("DownloadURL") || ""), + String(event?.dataTransfer?.getData?.("text/x-moz-url") || ""), + String(event?.dataTransfer?.getData?.("text/plain") || ""), + ].filter(Boolean); + + for (const payload of payloads) { + const lines = payload.split(/\r?\n/).map((line) => line.trim()).filter(Boolean); + for (const line of lines) { + const chunks = line.split(/\|/).map((part) => part.trim()).filter(Boolean); + for (const chunk of chunks) { + push(chunk); + } + } + } + + return out; +} + +function getOptionSetForNode(node) { + const rawOpts = getWidget(node, "slot_01_lora_name")?.options?.values; + if (Array.isArray(rawOpts) && rawOpts.length > 1) { + return new Set(rawOpts); + } + if (_loraSetCache instanceof Set && _loraSetCache.size > 0) { + return _loraSetCache; + } + return null; +} + +function resolveDroppedLoraName(node, event) { + const candidates = collectDroppedLoraCandidates(event); + if (!candidates.length) return ""; + + const optionSet = getOptionSetForNode(node); + if (optionSet && optionSet.size > 0) { + for (const candidate of candidates) { + if (optionSet.has(candidate)) { + return candidate; + } + } + + const optionEntries = Array.from(optionSet); + for (const candidate of candidates) { + const candidateBase = basenameOfPath(candidate).toLowerCase(); + const matched = optionEntries.find((entry) => basenameOfPath(entry).toLowerCase() === candidateBase); + if (matched) { + return matched; + } + } + } + + return basenameOfPath(candidates[0]); +} + +function isExternalFileDrag(event) { + const dt = event?.dataTransfer; + if (!dt) return false; + + try { + const items = Array.from(dt.items || []); + if (items.some((item) => item?.kind === "file")) { + return true; + } + } catch {} + + try { + const types = Array.from(dt.types || []).map((value) => String(value)); + return types.includes("Files") || types.includes("application/x-moz-file"); + } catch {} + + return false; +} + +function hasDroppedLoraCandidate(event) { + return collectDroppedLoraCandidates(event).length > 0; +} + +// Returns {x, y} in the same coordinate space as widget.y / widget.last_y in LiteGraph. +// The canvas is translated to node.pos before drawNode/onDrawForeground, so the coordinate +// system for drawing starts at (0,0) = top-left of node (including title). +// widget.y is set by _arrangeWidgets and starts at NODE_TITLE_HEIGHT (approximately), +// so no extra correction is needed — just subtract node.pos to get node-local coords. +function resolveLocalDropPoint(node, event) { + try { + const canvas = app?.canvas?.canvas; + const ds = app?.canvas?.ds; + const rect = canvas?.getBoundingClientRect?.(); + if (rect && ds && Number.isFinite(event?.clientX) && Number.isFinite(event?.clientY)) { + const scale = ds.scale || 1; + const graphX = ((event.clientX - rect.left) / scale) - (ds.offset?.[0] || 0); + const graphY = ((event.clientY - rect.top) / scale) - (ds.offset?.[1] || 0); + return { + x: graphX - (node?.pos?.[0] || 0), + y: graphY - (node?.pos?.[1] || 0), + }; + } + } catch {} + + const explicitX = Number(event?.graphX ?? event?.canvasX ?? NaN); + const explicitY = Number(event?.graphY ?? event?.canvasY ?? NaN); + if (Number.isFinite(explicitX) && Number.isFinite(explicitY)) { + return { + x: explicitX - (node?.pos?.[0] || 0), + y: explicitY - (node?.pos?.[1] || 0), + }; + } + + const graphMouse = app?.canvas?.graph_mouse; + if (Array.isArray(graphMouse) && graphMouse.length >= 2 && Number.isFinite(graphMouse[0]) && Number.isFinite(graphMouse[1])) { + return { + x: graphMouse[0] - (node?.pos?.[0] || 0), + y: graphMouse[1] - (node?.pos?.[1] || 0), + }; + } + + return null; +} + +function resolveLocalDropY(node, event) { + return resolveLocalDropPoint(node, event)?.y ?? null; +} + +function getWidgetHeight(widget, nodeWidth) { + // Use LiteGraph's already-computed height if available + if (typeof widget?.computedHeight === "number" && Number.isFinite(widget.computedHeight) && widget.computedHeight > 0) { + return widget.computedHeight; + } + try { + const height = widget?.computeSize?.(nodeWidth)?.[1]; + if (typeof height === "number" && Number.isFinite(height)) return Math.max(0, height); + } catch {} + try { + if (typeof LiteGraph !== "undefined" && LiteGraph?.NODE_WIDGET_HEIGHT != null) { + return LiteGraph.NODE_WIDGET_HEIGHT; + } + } catch {} + return 20; +} + +// Get widget local-body Y using LiteGraph's own properties (set during arrange/draw) +// widget.y is set by _arrangeWidgets(), widget.last_y is set during drawWidgets(). +// Both are relative to the node body (not including the title bar). +function getWidgetBodyY(widget) { + // widget.y is set by _arrangeWidgets() — always > 0 for arranged widgets (title adds ~30px offset). + // A value of 0 means the widget hasn't been arranged yet (default from addWidget), treat as unknown. + const y = widget?.y; + if (typeof y === "number" && Number.isFinite(y) && y > 0) return y; + const ly = widget?.last_y; + if (typeof ly === "number" && Number.isFinite(ly) && ly > 0) return ly; + return null; +} + +function getWidgetAtBodyY(node, bodyY) { + for (const widget of node.widgets || []) { + if (!widget || widget.hidden) continue; + const wy = getWidgetBodyY(widget); + if (wy === null) continue; + const wh = getWidgetHeight(widget, node.size?.[0] || SIMPLE_WIDTH); + if (bodyY >= wy && bodyY <= wy + wh) return widget; + } + return null; +} + +function getSlotFromWidgetName(name) { + const match = /^slot_(\d+)_/.exec(String(name || "")); + return match ? Number(match[1]) : null; +} + +function getWidgetLayoutRows(node) { + const rows = []; + const nodeWidth = node.size?.[0] || SIMPLE_WIDTH; + for (const widget of node.widgets || []) { + if (!widget || widget.hidden) continue; + const y = getWidgetBodyY(widget); + if (y === null) continue; + const height = getWidgetHeight(widget, nodeWidth); + rows.push({ widget, y, height }); + } + // Sort by Y in case order is wrong + rows.sort((a, b) => a.y - b.y); + return rows; +} + +function getSlotRect(node, slot) { + if (!slot || slot < 1) return null; + const rows = getWidgetLayoutRows(node); + let minY = Infinity; + let maxY = -Infinity; + + for (const row of rows) { + const rowSlot = getSlotFromWidgetName(row.widget?.name); + if (rowSlot !== slot) continue; + minY = Math.min(minY, row.y); + maxY = Math.max(maxY, row.y + row.height); + } + + if (!Number.isFinite(minY) || !Number.isFinite(maxY)) return null; + return { + x: 8, + y: minY - 2, + width: Math.max(0, (node?.size?.[0] || SIMPLE_WIDTH) - 16), + height: Math.max(0, maxY - minY + 4), + }; +} + +function getFallbackEmptySlot(node) { + const visibleSlots = getVisibleSlots(node); + for (let slot = 1; slot <= visibleSlots; slot += 1) { + const widget = getWidget(node, `slot_${String(slot).padStart(2, "0")}_lora_name`); + if (widget && !widget.hidden && String(widget.value || "no") === "no") { + return slot; + } + } + return visibleSlots > 0 ? 1 : null; +} + +function getDropTargetSlot(node, event, options = {}) { + const { fallbackToEmpty = false } = options; + const visibleSlots = getVisibleSlots(node); + const localPoint = resolveLocalDropPoint(node, event); + const localX = Number(localPoint?.x); + const localY = Number(localPoint?.y); + if (Number.isFinite(localY)) { + const slotRects = []; + for (let slot = 1; slot <= visibleSlots; slot += 1) { + const rect = getSlotRect(node, slot); + if (rect) slotRects.push({ slot, rect }); + } + + for (const { slot, rect } of slotRects) { + const insideX = !Number.isFinite(localX) || (localX >= rect.x && localX <= rect.x + rect.width); + if (insideX && localY >= rect.y && localY <= rect.y + rect.height) { + return slot; + } + } + + const hit = getWidgetAtBodyY(node, localY); + const hitSlot = getSlotFromWidgetName(hit?.name); + if (hitSlot != null) return hitSlot; + + for (const row of getWidgetLayoutRows(node)) { + if (localY >= row.y && localY <= row.y + row.height + 4) { + const rowSlot = getSlotFromWidgetName(row.widget?.name); + if (rowSlot != null) return rowSlot; + } + } + + if (slotRects.length > 0) { + let bestSlot = slotRects[0].slot; + let bestDistance = Infinity; + for (const { slot, rect } of slotRects) { + const centerY = rect.y + (rect.height / 2); + const distance = Math.abs(localY - centerY); + if (distance < bestDistance) { + bestDistance = distance; + bestSlot = slot; + } + } + return bestSlot; + } + } + return fallbackToEmpty ? getFallbackEmptySlot(node) : null; +} + +function setDropHoverSlot(node, slot) { + node._iamccsDropHoverSlot = Number(slot) || 0; + node._iamccsDropHoverAt = Date.now(); +} + +function clearDropHoverSlot(node) { + if (!node) return; + if (!node._iamccsDropHoverSlot && !node._iamccsDropHoverAt) return; + node._iamccsDropHoverSlot = 0; + node._iamccsDropHoverAt = 0; + node.setDirtyCanvas?.(true, true); +} + +function drawDropHover(node, ctx) { + const slot = Number(node?._iamccsDropHoverSlot) || 0; + const at = Number(node?._iamccsDropHoverAt) || 0; + if (!slot || !at) return; + if ((Date.now() - at) > DROP_HOVER_TIMEOUT_MS) { + node._iamccsDropHoverSlot = 0; + node._iamccsDropHoverAt = 0; + return; + } + + const rect = getSlotRect(node, slot); + if (!rect) return; + + ctx.save(); + ctx.fillStyle = "rgba(94, 181, 120, 0.14)"; + ctx.strokeStyle = "rgba(94, 181, 120, 0.95)"; + ctx.lineWidth = 2; + ctx.fillRect(rect.x, rect.y, rect.width, rect.height); + ctx.strokeRect(rect.x, rect.y, rect.width, rect.height); + ctx.restore(); +} + +function assignDroppedLora(node, event) { + const loraName = resolveDroppedLoraName(node, event); + if (!loraName) return false; + + const targetSlot = Number(node?._iamccsDropHoverSlot) || getDropTargetSlot(node, event, { fallbackToEmpty: true }); + const targetWidget = targetSlot != null + ? getWidget(node, `slot_${String(targetSlot).padStart(2, "0")}_lora_name`) + : null; + if (!targetWidget) return false; + + const optionValues = Array.isArray(targetWidget?.options?.values) ? targetWidget.options.values : null; + if (optionValues && !optionValues.includes(loraName)) { + targetWidget.options.values = [loraName, ...optionValues]; + } + + targetWidget.value = loraName; + try { + targetWidget.callback?.(loraName, app.canvas, node); + } catch {} + syncAllLinkedLowNodes(node); + clearDropHoverSlot(node); + node.setDirtyCanvas?.(true, true); + return true; +} + +// Cache of LoRA filenames from server API (loaded once). +let _loraSetCache = null; + +async function getLoraSet() { + if (_loraSetCache !== null) return _loraSetCache; + try { + const resp = await app.api.fetchApi("/object_info/IAMCCS_WanLoRASchedule"); + if (resp.ok) { + const data = await resp.json(); + const req = data?.IAMCCS_WanLoRASchedule?.input?.required; + const slot1 = req?.slot_01_lora_name; + if (Array.isArray(slot1?.[0])) { + _loraSetCache = new Set(slot1[0]); + return _loraSetCache; + } + } + } catch {} + // Fallback: try widget.options.values from any lora_name combo on an existing node + try { + const nodeType = Object.keys(app.graph?._nodes_by_id || {}); + for (const id of nodeType) { + const n = app.graph.getNodeById(id); + if (n?.type !== NODE_NAME) continue; + const w = getWidget(n, "slot_01_lora_name"); + const v = w?.options?.values ?? (Array.isArray(w?.options) ? w.options : null); + if (Array.isArray(v) && v.length > 1) { + _loraSetCache = new Set(v); + return _loraSetCache; + } + } + } catch {} + _loraSetCache = new Set(); + return _loraSetCache; +} + +async function syncSlotToLowNode(highNode, lowNode, slot) { + const prefix = `slot_${String(slot).padStart(2, "0")}`; + const highWidget = getWidget(highNode, `${prefix}_lora_name`); + const lowWidget = getWidget(lowNode, `${prefix}_lora_name`); + if (!highWidget || !lowWidget) return; + + lowNode.properties = lowNode.properties || {}; + lowNode.properties[AUTO_PROP] = lowNode.properties[AUTO_PROP] || {}; + const autoKey = `${prefix}_lora_name`; + const lastAuto = String(lowNode.properties[AUTO_PROP][autoKey] || ""); + const lowValue = String(lowWidget.value || "no"); + const highValue = String(highWidget.value || "no"); + if (!highValue || highValue === "no") { + if (lowValue === lastAuto || lowValue === "no") { + lowWidget.value = "no"; + delete lowNode.properties[AUTO_PROP][autoKey]; + try { + lowWidget.callback?.("no", app.canvas, lowNode); + } catch {} + } + return; + } + + const candidates = candidateLowNames(highValue); + if (!candidates.length) return; // no _HIGH_/_LOW_ pattern in name + + // Build option set: prefer widget options, then server cache + let optionSet = null; + const rawOpts = lowWidget?.options?.values ?? (Array.isArray(lowWidget?.options) ? lowWidget.options : null); + if (Array.isArray(rawOpts) && rawOpts.length > 1) { + optionSet = new Set(rawOpts); + } else { + optionSet = await getLoraSet(); + } + + const suggested = (optionSet && optionSet.size > 0) + ? candidates.find((c) => optionSet.has(c)) + : candidates[0]; // no list available: try first candidate blindly + if (!suggested) return; + // If we have a list and the candidate is not in it, don't set an invalid value + if (optionSet && optionSet.size > 0 && !optionSet.has(suggested)) return; + + // Don't overwrite if user manually chose something other than "no" or our last auto-set value + if (lowValue !== "no" && lowValue !== lastAuto) return; + + lowWidget.value = suggested; + lowNode.properties[AUTO_PROP][autoKey] = suggested; + try { + lowWidget.callback?.(suggested, app.canvas, lowNode); + } catch {} +} + +async function syncAllLinkedLowNodes(highNode) { + const targets = getLinkedLinxTargets(highNode); + if (!targets.length) return; + for (const lowNode of targets) { + syncLowNodePresentation(highNode, lowNode); + ensureVisibleDefaults(lowNode); + applyLayout(lowNode); + for (let slot = 1; slot <= getVisibleSlots(highNode); slot += 1) { + await syncSlotToLowNode(highNode, lowNode, slot); + } + lowNode.setDirtyCanvas?.(true, true); + } +} + +function wrapSlotCallbacks(node) { + if (node._iamccsScheduleWrappedCallbacks) return; + for (let slot = 1; slot <= MAX_SLOTS; slot += 1) { + const prefix = `slot_${String(slot).padStart(2, "0")}`; + // Wrap lora_name → triggers linx sync + const loraWidget = getWidget(node, `${prefix}_lora_name`); + if (loraWidget && !loraWidget._iamccsWrapped) { + const origLora = loraWidget.callback; + loraWidget.callback = (value, canvas, targetNode) => { + try { origLora?.(value, canvas, targetNode); } finally { + syncAllLinkedLowNodes(node); + } + }; + loraWidget._iamccsWrapped = true; + } + // Wrap preset → refreshes start/end visibility + const presetWidget = getWidget(node, `${prefix}_preset`); + if (presetWidget && !presetWidget._iamccsWrapped) { + const origPreset = presetWidget.callback; + presetWidget.callback = (value, canvas, targetNode) => { + try { origPreset?.(value, canvas, targetNode); } finally { + setTimeout(() => { applyLayout(node); node.setDirtyCanvas(true, true); }, 0); + } + }; + presetWidget._iamccsWrapped = true; + } + } + node._iamccsScheduleWrappedCallbacks = true; +} + +function presetNeedsRange(preset) { + const s = String(preset || ""); + return !s || s === "custom range" || s === "manual_range"; +} + +function setSlotVisibility(node, slot, visible, advanced) { + const prefix = `slot_${String(slot).padStart(2, "0")}`; + const loraWidget = getWidget(node, `${prefix}_lora_name`); + const strengthWidget = getWidget(node, `${prefix}_strength`); + const presetWidget = getWidget(node, `${prefix}_preset`); + const startWidget = getWidget(node, `${prefix}_start`); + const endWidget = getWidget(node, `${prefix}_end`); + if (visible) { + showWidget(loraWidget); + showWidget(strengthWidget); + if (advanced) { + showWidget(presetWidget); + const currentPreset = String(presetWidget?.value ?? ""); + if (presetNeedsRange(currentPreset)) { + showWidget(startWidget); + showWidget(endWidget); + } else { + hideWidget(startWidget); + hideWidget(endWidget); + } + } else { + hideWidget(presetWidget); + hideWidget(startWidget); + hideWidget(endWidget); + } + } else { + hideWidget(loraWidget); + hideWidget(strengthWidget); + hideWidget(presetWidget); + hideWidget(startWidget); + hideWidget(endWidget); + } +} + +function reorderWidgets(node) { + const orderedNames = [ + "generation_index", + "log_prefix", + "model_type", + MODE_WIDGET, + LOAD_WIDGET, + SAVE_WIDGET, + ADD_WIDGET, + DELETE_WIDGET, + RESET_WIDGET, + ]; + for (let slot = 1; slot <= MAX_SLOTS; slot += 1) { + const prefix = `slot_${String(slot).padStart(2, "0")}`; + orderedNames.push(`${prefix}_lora_name`); + orderedNames.push(`${prefix}_strength`); + orderedNames.push(`${prefix}_preset`); + orderedNames.push(`${prefix}_start`); + orderedNames.push(`${prefix}_end`); + } + const byName = new Map((node.widgets || []).map((widget) => [widget?.name, widget])); + const ordered = []; + for (const name of orderedNames) { + const widget = byName.get(name); + if (widget) { + ordered.push(widget); + byName.delete(name); + } + } + node.widgets = [...ordered, ...Array.from(byName.values())]; +} + +function targetWidth(node) { + return (node?.properties?.[MODE_PROP] || "simple") === "advanced" ? ADVANCED_WIDTH : SIMPLE_WIDTH; +} + +function targetHeight(node) { + const mode = node?.properties?.[MODE_PROP] || "simple"; + const count = getVisibleSlots(node); + const slotHeight = mode === "advanced" ? ADVANCED_SLOT_HEIGHT : SIMPLE_SLOT_HEIGHT; + return BASE_HEIGHT + (count * slotHeight) + EXTRA_HEIGHT; +} + +function applyLayout(node) { + ensureBoxShape(node); + ensureVisibleDefaults(node); + applyLabels(node); + applyPresetOptions(node); + wrapSlotCallbacks(node); + + hideWidget(getWidget(node, "generation_index")); + hideWidget(getWidget(node, "log_prefix")); + + const advanced = (node?.properties?.[MODE_PROP] || "simple") === "advanced"; + if (advanced) { + showWidget(getWidget(node, "model_type")); + } else { + hideWidget(getWidget(node, "model_type")); + } + + const visibleSlots = getVisibleSlots(node); + for (let slot = 1; slot <= MAX_SLOTS; slot += 1) { + setSlotVisibility(node, slot, slot <= visibleSlots, advanced); + } + + reorderWidgets(node); + const width = targetWidth(node); + const height = targetHeight(node); + const computed = node.computeSize?.() || node.size || [width, height]; + node.size = [Math.max(width, computed[0] || width), Math.max(height, computed[1] || height)]; +} + +app.registerExtension({ + name: "iamccs.wan_lora_schedule_ui", + + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData?.name !== NODE_NAME) return; + + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + const result = onNodeCreated?.apply(this, arguments); + this.properties = this.properties || {}; + if (!this.properties[MODE_PROP]) this.properties[MODE_PROP] = "simple"; + if (!this.properties[SLOTS_PROP]) this.properties[SLOTS_PROP] = 3; + ensureModeWidget(this); + ensureActionWidgets(this); + setTimeout(() => { + applyLayout(this); + syncAllLinkedLowNodes(this); + this.setDirtyCanvas(true, true); + }, 0); + return result; + }; + + const onDragOver = nodeType.prototype.onDragOver; + nodeType.prototype.onDragOver = function (event) { + if (isExternalFileDrag(event)) { + try { + app?.canvas?.adjustMouseEvent?.(event); + } catch {} + const slot = getDropTargetSlot(this, event, { fallbackToEmpty: false }); + if (slot != null) { + setDropHoverSlot(this, slot); + } else { + clearDropHoverSlot(this); + } + try { + event?.preventDefault?.(); + event?.stopPropagation?.(); + if (event?.dataTransfer) event.dataTransfer.dropEffect = "copy"; + } catch {} + this.setDirtyCanvas?.(true, true); + return true; + } + clearDropHoverSlot(this); + return onDragOver?.apply(this, arguments); + }; + + const onDragDrop = nodeType.prototype.onDragDrop; + nodeType.prototype.onDragDrop = function (event) { + try { + app?.canvas?.adjustMouseEvent?.(event); + } catch {} + try { + event?.preventDefault?.(); + event?.stopPropagation?.(); + } catch {} + if (hasDroppedLoraCandidate(event)) { + assignDroppedLora(this, event); + return true; + } + clearDropHoverSlot(this); + return onDragDrop?.apply(this, arguments); + }; + + const onDrawForeground = nodeType.prototype.onDrawForeground; + nodeType.prototype.onDrawForeground = function (ctx) { + const result = onDrawForeground?.apply(this, arguments); + drawDropHover(this, ctx); + return result; + }; + + const onConfigure = nodeType.prototype.onConfigure; + nodeType.prototype.onConfigure = function (config) { + // 1. Normalize legacy corrupt widget-value arrays (modifies config in place) + normalizeLegacyWidgetValues(config); + + // 2. Re-apply normalized values to widgets explicitly. + // LiteGraph assigns widgets_values → widget.value BEFORE calling onConfigure, + // so if the values were corrupt, they are already set wrong. Re-assigning here + // (in the same iteration order, skipping non-serializable UI widgets) corrects them. + if (Array.isArray(config.widgets_values) && this.widgets) { + let j = 0; + for (const widget of this.widgets) { + if (!widget || widget.options?.serialize === false || widget._iamccsUiOnly) continue; + if (j < config.widgets_values.length) { + widget.value = config.widgets_values[j++]; + } + } + } + + const result = onConfigure?.apply(this, arguments); + this.properties = this.properties || {}; + if (!this.properties[MODE_PROP]) this.properties[MODE_PROP] = "simple"; + if (!this.properties[SLOTS_PROP]) this.properties[SLOTS_PROP] = 3; + ensureModeWidget(this); + ensureActionWidgets(this); + setTimeout(() => { + // Mark ALL slots as already-initialized so ensureSlotDefaults won't clobber + // values that were restored from the saved workflow. + this.properties[INIT_PROP] = this.properties[INIT_PROP] || {}; + for (let s = 1; s <= MAX_SLOTS; s++) { + this.properties[INIT_PROP][s] = true; + } + applyLayout(this); + syncAllLinkedLowNodes(this); + this.setDirtyCanvas(true, true); + }, 0); + return result; + }; + + const onConnectionsChange = nodeType.prototype.onConnectionsChange; + nodeType.prototype.onConnectionsChange = function () { + const result = onConnectionsChange?.apply(this, arguments); + setTimeout(() => { + applyLayout(this); + syncAllLinkedLowNodes(this); + this.setDirtyCanvas(true, true); + }, 0); + return result; + }; + + const collapse = nodeType.prototype.collapse; + if (typeof collapse === "function") { + nodeType.prototype.collapse = function () { + const result = collapse?.apply(this, arguments); + setTimeout(() => { + syncAllLinkedLowNodes(this); + this.setDirtyCanvas(true, true); + }, 0); + return result; + }; + } + }, +}); diff --git a/web/iamccs_workflow_persist_cleanup.js b/web/iamccs_workflow_persist_cleanup.js new file mode 100644 index 0000000..0d252ee --- /dev/null +++ b/web/iamccs_workflow_persist_cleanup.js @@ -0,0 +1,188 @@ +import { app } from "../../scripts/app.js"; + +const LEGACY_LOCAL_KEYS = [ + "workflow", + "Comfy.OpenWorkflowsPaths", + "Comfy.ActiveWorkflowIndex", +]; + +const LEGACY_PREFIXES = [ + "Comfy.Workflow.Drafts", + "Comfy.Workflow.DraftOrder", +]; + +const V2_PREFIX = { + draftIndex: "Comfy.Workflow.DraftIndex.v2:", + draftPayload: "Comfy.Workflow.Draft.v2:", + lastActivePath: "Comfy.Workflow.LastActivePath:", + lastOpenPaths: "Comfy.Workflow.LastOpenPaths:", +}; + +const MAX_DRAFTS_PER_WORKSPACE = 4; +const MAX_DRAFT_CHARS_PER_WORKSPACE = 2_500_000; + +function safeParseJson(json) { + if (!json) return null; + try { + return JSON.parse(json); + } catch { + return null; + } +} + +function listStorageKeys(storage) { + const keys = []; + try { + for (let index = 0; index < storage.length; index += 1) { + const key = storage.key(index); + if (key) keys.push(key); + } + } catch { + return []; + } + return keys; +} + +function removeKey(storage, key) { + try { + storage.removeItem(key); + return true; + } catch { + return false; + } +} + +function getPayloadKey(workspaceId, draftKey) { + return `${V2_PREFIX.draftPayload}${workspaceId}:${draftKey}`; +} + +function cleanupLegacyWorkflowStorage(localStorageRef) { + let removed = 0; + for (const key of LEGACY_LOCAL_KEYS) { + if (removeKey(localStorageRef, key)) removed += 1; + } + + for (const key of listStorageKeys(localStorageRef)) { + if (LEGACY_PREFIXES.some((prefix) => key.startsWith(prefix))) { + if (removeKey(localStorageRef, key)) removed += 1; + } + } + + return removed; +} + +function cleanupWorkspaceDrafts(localStorageRef, workspaceId) { + const indexKey = `${V2_PREFIX.draftIndex}${workspaceId}`; + const rawIndex = localStorageRef.getItem(indexKey); + if (!rawIndex) return { removed: 0, rewritten: false }; + + const index = safeParseJson(rawIndex); + if (!index || typeof index !== "object" || !Array.isArray(index.order) || typeof index.entries !== "object" || index.entries == null) { + let removed = removeKey(localStorageRef, indexKey) ? 1 : 0; + const payloadPrefix = `${V2_PREFIX.draftPayload}${workspaceId}:`; + for (const key of listStorageKeys(localStorageRef)) { + if (key.startsWith(payloadPrefix) && removeKey(localStorageRef, key)) { + removed += 1; + } + } + return { removed, rewritten: false }; + } + + const payloadPrefix = `${V2_PREFIX.draftPayload}${workspaceId}:`; + const seenPayloads = new Set(); + const keptEntries = {}; + const keptOrder = []; + let keptChars = 0; + let keptCount = 0; + let removed = 0; + + const newestFirst = [...index.order].reverse(); + for (const draftKey of newestFirst) { + const entry = index.entries[draftKey]; + if (!entry || typeof entry.path !== "string") continue; + + const payloadKey = getPayloadKey(workspaceId, draftKey); + const payloadJson = localStorageRef.getItem(payloadKey); + if (!payloadJson) continue; + + const payloadChars = payloadJson.length; + const canKeep = keptCount < MAX_DRAFTS_PER_WORKSPACE && ( + keptCount === 0 || keptChars + payloadChars <= MAX_DRAFT_CHARS_PER_WORKSPACE + ); + + if (canKeep) { + keptEntries[draftKey] = entry; + keptOrder.unshift(draftKey); + keptChars += payloadChars; + keptCount += 1; + seenPayloads.add(payloadKey); + } else if (removeKey(localStorageRef, payloadKey)) { + removed += 1; + } + } + + for (const key of listStorageKeys(localStorageRef)) { + if (key.startsWith(payloadPrefix) && !seenPayloads.has(key)) { + if (removeKey(localStorageRef, key)) removed += 1; + } + } + + const nextIndex = { + ...index, + updatedAt: Date.now(), + order: keptOrder, + entries: keptEntries, + }; + + const nextIndexJson = JSON.stringify(nextIndex); + const rewritten = nextIndexJson !== rawIndex; + if (rewritten) { + try { + localStorageRef.setItem(indexKey, nextIndexJson); + } catch { + // If index rewrite fails, leaving the trimmed payload set is still better + // than keeping localStorage saturated with stale drafts. + } + } + + return { removed, rewritten }; +} + +function cleanupWorkflowPersistenceStorage() { + try { + const localStorageRef = window.localStorage; + if (!localStorageRef) return; + + let removed = cleanupLegacyWorkflowStorage(localStorageRef); + let rewritten = 0; + + const workspaceIds = new Set(); + for (const key of listStorageKeys(localStorageRef)) { + if (key.startsWith(V2_PREFIX.draftIndex)) { + workspaceIds.add(key.slice(V2_PREFIX.draftIndex.length)); + } + } + + for (const workspaceId of workspaceIds) { + const result = cleanupWorkspaceDrafts(localStorageRef, workspaceId); + removed += result.removed; + if (result.rewritten) rewritten += 1; + } + + if (removed || rewritten) { + console.info("[IAMCCS] Workflow draft storage cleanup applied", { + removedKeys: removed, + rewrittenIndexes: rewritten, + }); + } + } catch (error) { + console.warn("[IAMCCS] Workflow draft storage cleanup skipped", error); + } +} + +app.registerExtension({ + name: "iamccs.workflow_persist_cleanup", + async setup() { + cleanupWorkflowPersistenceStorage(); + }, +});