Mehdi Chaid
72e784ae0c
The _shorten_filenames function was adding both the full filename and the stem (without extension) to dropdown lists, causing models like VAEs and text encoders to appear twice (e.g. 'ae' and 'ae.safetensors'). Now only the stem is returned. Backend resolution helpers (_resolve_model_name) are added to translate stem-only names back to full filenames when loading checkpoin/vae/ text_encoder/controlnet models, ensuring backwards compatibility with existing YAML files that store full filenames.
136 lines
4.3 KiB
Python
136 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageOps
|
|
|
|
import folder_paths
|
|
from nodes import (
|
|
CLIPTextEncode,
|
|
ControlNetApplyAdvanced,
|
|
ControlNetLoader,
|
|
EmptyLatentImage,
|
|
LoraLoader,
|
|
)
|
|
|
|
from . import flake_io
|
|
|
|
|
|
def _resolve_model_name(category: str, stem_or_name: str) -> str:
|
|
available = folder_paths.get_filename_list(category)
|
|
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
|
|
|
|
return stem_or_name
|
|
|
|
|
|
def _resolve_lora_name(stem_or_name: str) -> str:
|
|
result = _resolve_model_name("loras", stem_or_name)
|
|
available = folder_paths.get_filename_list("loras")
|
|
if result.replace("\\", "/") in {p.replace("\\", "/") for p in available}:
|
|
return result
|
|
raise FileNotFoundError(
|
|
f"LoRA '{stem_or_name}' not found in models/loras/. "
|
|
f"Provide the stem or full filename of an existing LoRA."
|
|
)
|
|
|
|
|
|
def _load_cn_image(image_name: str) -> torch.Tensor:
|
|
"""Load an image tensor shaped [1, H, W, 3] (ComfyUI's IMAGE format) from ComfyUI/input/.
|
|
|
|
`image_name` may be a stem ('standing_openpose') or a filename with extension.
|
|
Subdirectories under input/ are allowed ('cnet/standing_openpose').
|
|
"""
|
|
input_dir = folder_paths.get_input_directory()
|
|
norm = image_name.replace("\\", "/")
|
|
|
|
candidates: list[str] = []
|
|
direct = os.path.join(input_dir, norm)
|
|
if os.path.isfile(direct):
|
|
candidates.append(direct)
|
|
else:
|
|
for ext in (".png", ".jpg", ".jpeg", ".webp"):
|
|
alt = os.path.join(input_dir, f"{norm}{ext}")
|
|
if os.path.isfile(alt):
|
|
candidates.append(alt)
|
|
break
|
|
|
|
if not candidates:
|
|
raise FileNotFoundError(
|
|
f"ControlNet image '{image_name}' not found in {input_dir}. "
|
|
f"Place the image there (any of .png/.jpg/.jpeg/.webp) or use a full filename."
|
|
)
|
|
|
|
img = Image.open(candidates[0])
|
|
img = ImageOps.exif_transpose(img).convert("RGB")
|
|
arr = np.array(img).astype(np.float32) / 255.0
|
|
return torch.from_numpy(arr)[None,]
|
|
|
|
|
|
def compose(
|
|
model: Any,
|
|
clip: Any,
|
|
entries: list[dict[str, Any]],
|
|
) -> tuple:
|
|
flakes = [flake_io.resolve(e) for e in entries]
|
|
|
|
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)
|
|
|
|
pos_text = " BREAK ".join(f.positive.strip() for f in flakes if f.positive and f.positive.strip())
|
|
neg_text = ", ".join(f.negative.strip() for f in flakes if f.negative and f.negative.strip())
|
|
|
|
encoder = CLIPTextEncode()
|
|
positive = encoder.encode(clip, pos_text)[0]
|
|
negative = encoder.encode(clip, neg_text)[0]
|
|
|
|
cn_model_cache: dict[str, Any] = {}
|
|
cn_loader = ControlNetLoader()
|
|
cn_apply = ControlNetApplyAdvanced()
|
|
for f in flakes:
|
|
for cn in f.controlnets:
|
|
if cn.strength == 0:
|
|
continue
|
|
cn_resolved = _resolve_model_name("controlnet", cn.model_name)
|
|
if cn_resolved not in cn_model_cache:
|
|
cn_model_cache[cn_resolved] = cn_loader.load_controlnet(cn_resolved)[0]
|
|
cn_model = cn_model_cache[cn_resolved]
|
|
image = _load_cn_image(cn.image_name)
|
|
positive, negative = cn_apply.apply_controlnet(
|
|
positive, negative, cn_model, image,
|
|
cn.strength, cn.start_percent, cn.end_percent,
|
|
)
|
|
|
|
width, height = 1024, 1024
|
|
for f in flakes:
|
|
if f.resolution is not None:
|
|
width, height = f.resolution
|
|
break
|
|
|
|
latent = EmptyLatentImage().generate(width, height, 1)[0]
|
|
|
|
logging.info(
|
|
"[FlakeStack] composed %d flake(s), resolution %dx%d, %d controlnet(s)",
|
|
len(flakes), width, height,
|
|
sum(len(f.controlnets) for f in flakes),
|
|
)
|
|
|
|
return model, clip, positive, negative, latent, width, height
|