from __future__ import annotations import json import logging import os import folder_paths import comfy.sd import comfy.utils from nodes import CLIPTextEncode, EmptyLatentImage, ControlNetApplyAdvanced, ControlNetLoader, LoraLoader from . import flake_io def _resolve_lora_name(stem_or_name: str) -> str: available = folder_paths.get_filename_list("loras") # Map normalized (forward-slash) paths back to original paths available_norm = {p.replace("\\", "/"): p for p in available} norm = stem_or_name.replace("\\", "/") if norm in available_norm: return available_norm[norm] norm_stem, _ = os.path.splitext(norm) for cand_norm, candidate in available_norm.items(): stem, _ = os.path.splitext(cand_norm) if stem == norm_stem: return candidate raise FileNotFoundError(f"LoRA '{stem_or_name}' not found in models/loras/") def _load_preset_bundle(preset_name: str): """Load a model preset and return (model_bundle, generation_data, sampling_preset).""" preset_data = flake_io.load_preset(preset_name) # --- Load checkpoint ---------------------------------------------------- ckpt_path = folder_paths.get_full_path("checkpoints", preset_data.checkpoint) if not ckpt_path or not os.path.isfile(ckpt_path): raise FileNotFoundError( f"Checkpoint '{preset_data.checkpoint}' not found in models/checkpoints/" ) embedding_dir = folder_paths.get_folder_paths("embeddings") model, clip, vae, _ = comfy.sd.load_checkpoint_guess_config( ckpt_path, output_vae=True, output_clip=True, embedding_directory=embedding_dir, ) # --- Clip skip ---------------------------------------------------------- if preset_data.clip_skip: clip = clip.clone() clip.clip_layer(preset_data.clip_skip) # --- Optional VAE override ---------------------------------------------- if preset_data.vae: vae_path = folder_paths.get_full_path("vae", preset_data.vae) if vae_path and os.path.isfile(vae_path): vae_sd = comfy.utils.load_torch_file(vae_path) vae = comfy.sd.VAE(sd=vae_sd) # --- Encode prompts ----------------------------------------------------- encoder = CLIPTextEncode() pos_text = preset_data.positive.strip() neg_text = preset_data.negative.strip() positive = encoder.encode(clip, pos_text)[0] if pos_text else encoder.encode(clip, "")[0] negative = encoder.encode(clip, neg_text)[0] if neg_text else encoder.encode(clip, "")[0] # --- Latent ------------------------------------------------------------- width, height = preset_data.width, preset_data.height latent = EmptyLatentImage().generate(width, height, 1)[0] model_bundle = (model, clip, vae) generation_data = (positive, negative, latent, width, height, pos_text, neg_text) sampling_preset = (preset_data.steps, preset_data.cfg, preset_data.sampler, preset_data.scheduler) return model_bundle, generation_data, sampling_preset class FlakeModelPreset: @classmethod def INPUT_TYPES(cls): try: preset_names = flake_io.list_presets() except Exception: preset_names = [] if not preset_names: preset_names = ["No model preset is selected"] return { "required": { "preset": (["Select a preset..."] + preset_names,), }, } RETURN_TYPES = ( "FLAKES_MODEL", "FLAKES_COND", "FLAKES_SAMPLER", ) RETURN_NAMES = ( "model_bundle", "generation_data", "sampling_preset", ) FUNCTION = "execute" CATEGORY = "flakes" DESCRIPTION = ( "Load a model preset (checkpoint, VAE, prompts, resolution, sampler settings). " "Outputs are bundled for wiring into FlakeStack / FlakeCombo nodes." ) def execute(self, preset: str): preset_name = preset.strip() if preset else "" if not preset_name or preset_name in ("Select a preset...", "No model preset is selected"): raise ValueError("No model preset is selected — pick one from the dropdown.") return _load_preset_bundle(preset_name) class FlakeStack: @classmethod def INPUT_TYPES(cls): return { "required": { "model_bundle": ("FLAKES_MODEL",), "generation_data": ("FLAKES_COND",), "sampling_preset": ("FLAKES_SAMPLER",), "flakes_json": ("STRING", { "multiline": True, "default": "[]", "tooltip": "JSON-encoded list of flake entries. Managed by the Full Flakes widget.", }), }, } RETURN_TYPES = ( "FLAKES_MODEL", "FLAKES_COND", "FLAKES_SAMPLER", ) RETURN_NAMES = ( "model_bundle", "generation_data", "sampling_preset", ) FUNCTION = "execute" CATEGORY = "flakes" DESCRIPTION = ( "Compose flakes on top of an incoming model/generation/sampler bundle. " "Applies LoRAs, merges prompts, resolution overrides and ControlNets. " "Outputs updated bundles for chaining into downstream nodes." ) def execute(self, model_bundle, generation_data, sampling_preset, flakes_json: str): model, clip, vae = model_bundle positive_cond, negative_cond, latent, width, height, pos_text, neg_text = generation_data steps, cfg, sampler, scheduler = sampling_preset # --- Parse flakes_json -------------------------------------------------- try: entries = json.loads(flakes_json) if flakes_json else [] except json.JSONDecodeError as exc: raise ValueError(f"flakes_json is not valid JSON: {exc}") from exc if not isinstance(entries, list): raise ValueError("flakes_json must be a JSON list") normalized: list[dict] = [] for i, entry in enumerate(entries): if not isinstance(entry, dict): logging.warning("[FlakeStack] skipping non-object entry %d: %r", i, entry) continue if entry.get("inline") or entry.get("name"): normalized.append(entry) else: logging.warning("[FlakeStack] skipping entry %d (needs 'name' or 'inline')", i) # --- Resolve flakes ----------------------------------------------------- flakes = [flake_io.resolve(e) for e in normalized] # --- Apply LoRAs -------------------------------------------------------- lora_loader = LoraLoader() for f in flakes: if not f.lora_path: continue lora_name = _resolve_lora_name(f.lora_path) model, clip = lora_loader.load_lora(model, clip, lora_name, f.strength, f.strength) # --- Build prompt text (existing → flakes) ------------------------------ pos_parts: list[str] = [] neg_parts: list[str] = [] if pos_text and pos_text.strip(): pos_parts.append(pos_text.strip()) if neg_text and neg_text.strip(): neg_parts.append(neg_text.strip()) for f in flakes: if f.positive and f.positive.strip(): pos_parts.append(f.positive.strip()) if f.negative and f.negative.strip(): neg_parts.append(f.negative.strip()) new_pos_text = " BREAK ".join(pos_parts) if pos_parts else "" new_neg_text = ", ".join(neg_parts) if neg_parts else "" # --- Re-encode prompts -------------------------------------------------- encoder = CLIPTextEncode() positive_cond = encoder.encode(clip, new_pos_text)[0] if new_pos_text else encoder.encode(clip, "")[0] negative_cond = encoder.encode(clip, new_neg_text)[0] if new_neg_text else encoder.encode(clip, "")[0] # --- Apply ControlNets (from flakes) ------------------------------------ from .flake_compose import _load_cn_image cn_model_cache = {} cn_loader = ControlNetLoader() cn_apply = ControlNetApplyAdvanced() for f in flakes: for cn in f.controlnets: if cn.strength == 0: continue if cn.model_name not in cn_model_cache: cn_model_cache[cn.model_name] = cn_loader.load_controlnet(cn.model_name)[0] cn_model = cn_model_cache[cn.model_name] image = _load_cn_image(cn.image_name) positive_cond, negative_cond = cn_apply.apply_controlnet( positive_cond, negative_cond, cn_model, image, cn.strength, cn.start_percent, cn.end_percent, ) # --- Resolution --------------------------------------------------------- new_width, new_height = width, height for f in flakes: if f.resolution is not None: new_width, new_height = f.resolution break if (new_width, new_height) != (width, height): latent = EmptyLatentImage().generate(new_width, new_height, 1)[0] logging.info( "[FlakeStack] %dx%s steps=%s cfg=%s flakes=%d", new_width, new_height, steps, cfg, len(flakes), ) out_model_bundle = (model, clip, vae) out_generation_data = (positive_cond, negative_cond, latent, new_width, new_height, new_pos_text, new_neg_text) out_sampling_preset = (steps, cfg, sampler, scheduler) return ( out_model_bundle, out_generation_data, out_sampling_preset, ) class FlakeCombo: @classmethod def INPUT_TYPES(cls): return { "required": { "model_bundle": ("FLAKES_MODEL",), "generation_data": ("FLAKES_COND",), "sampling_preset": ("FLAKES_SAMPLER",), "flakes_json": ("STRING", { "multiline": True, "default": "[]", "tooltip": "JSON-encoded list of mutually exclusive flake entries.", }), }, } RETURN_TYPES = ( "FLAKES_MODEL", "FLAKES_COND", "FLAKES_SAMPLER", ) RETURN_NAMES = ( "model_bundle", "generation_data", "sampling_preset", ) FUNCTION = "execute" CATEGORY = "flakes" DESCRIPTION = ( "Mutually-exclusive flake stack. Each flake is applied in isolation, " "generating a separate output bundle per flake. (Sequential execution TBD.)" ) def execute(self, model_bundle, generation_data, sampling_preset, flakes_json: str): # TODO: implement combinatorial / sequential generation. # For now, behave like FlakeStack so the node is usable while the # iteration strategy is being clarified. return FlakeStack().execute(model_bundle, generation_data, sampling_preset, flakes_json) class FlakeModelCombo: @classmethod def INPUT_TYPES(cls): try: preset_names = flake_io.list_presets() except Exception: preset_names = [] if not preset_names: preset_names = ["No model preset is selected"] return { "required": { "presets_json": ("STRING", { "multiline": True, "default": "[]", "tooltip": "JSON-encoded list of preset names.", }), }, "optional": { "model_bundle": ("FLAKES_MODEL",), "generation_data": ("FLAKES_COND",), "sampling_preset": ("FLAKES_SAMPLER",), }, } RETURN_TYPES = ( "FLAKES_MODEL", "FLAKES_COND", "FLAKES_SAMPLER", ) RETURN_NAMES = ( "model_bundle", "generation_data", "sampling_preset", ) FUNCTION = "execute" CATEGORY = "flakes" DESCRIPTION = ( "Load multiple model presets and generate the same flake setup for each. " "(Sequential execution TBD.)" ) def execute(self, presets_json: str, model_bundle=None, generation_data=None, sampling_preset=None): # TODO: implement multi-preset iteration. # For now, load the first preset and pass through any upstream bundles. try: names = json.loads(presets_json) if presets_json else [] except json.JSONDecodeError as exc: raise ValueError(f"presets_json is not valid JSON: {exc}") from exc if not isinstance(names, list) or not names: raise ValueError("presets_json must be a non-empty JSON list of preset names") first_name = names[0].strip() if isinstance(names[0], str) else "" if not first_name: raise ValueError("first preset name is empty") return _load_preset_bundle(first_name)