Add files via upload
This commit is contained in:
@@ -0,0 +1,153 @@
|
||||
import os
|
||||
import re
|
||||
import json
|
||||
from datetime import datetime
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image, PngImagePlugin
|
||||
import folder_paths
|
||||
|
||||
class ImageSaverMango:
|
||||
|
||||
CATEGORY = "Mango Node Pack/Metadata"
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "save_images"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE", {"tooltip": "Images to save"}),
|
||||
"metadata": ("METADATA", {"tooltip": "Metadata dictionary from upstream (e.g. from KSampler (Mango))"}),
|
||||
"filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "Prefix for the saved file. Supports placeholders like %date:hhmmss%."}),
|
||||
"subdirectory_name": ("STRING", {"default": "", "tooltip": "Custom subdirectory (inside the output folder) to save the image."}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
DESCRIPTION = (
|
||||
"Saves images with embedded metadata. Filenames are numbered as prefix_00001, prefix_00002, etc. "
|
||||
"Also creates a single 'parameters' text chunk containing the typical KSampler info (prompt, seed, etc.)."
|
||||
)
|
||||
|
||||
pattern_format = re.compile(r"(%[^%]+%)")
|
||||
|
||||
def __init__(self):
|
||||
self.output_dir = folder_paths.get_output_directory()
|
||||
self.compress_level = 4
|
||||
self.prefix_append = ""
|
||||
|
||||
def save_images(self, images, metadata, filename_prefix="ComfyUI", subdirectory_name=""):
|
||||
param_string = self.build_param_string(metadata)
|
||||
|
||||
base_name = self.format_filename(filename_prefix, metadata) + self.prefix_append
|
||||
|
||||
if subdirectory_name:
|
||||
subdirectory_name = self.format_filename(subdirectory_name, metadata)
|
||||
full_output_folder = os.path.join(self.output_dir, subdirectory_name)
|
||||
else:
|
||||
full_output_folder = self.output_dir
|
||||
os.makedirs(full_output_folder, exist_ok=True)
|
||||
|
||||
base_format = "png"
|
||||
start_suffix = self.next_available_suffix(base_name, full_output_folder, base_format)
|
||||
|
||||
results = []
|
||||
for i, image in enumerate(images):
|
||||
arr = self.to_uint8(image)
|
||||
pil_img = Image.fromarray(arr)
|
||||
# Build a PngInfo object.
|
||||
pnginfo = self.prepare_pnginfo(metadata, i, len(images))
|
||||
pnginfo.add_text("parameters", param_string)
|
||||
suffix_number = start_suffix + i
|
||||
filename_with_suffix = f"{base_name}_{suffix_number:05d}"
|
||||
file = f"{filename_with_suffix}.{base_format}"
|
||||
full_path = os.path.join(full_output_folder, file)
|
||||
pil_img.save(full_path, pnginfo=pnginfo, compress_level=self.compress_level)
|
||||
results.append({
|
||||
"filename": file,
|
||||
"subfolder": full_output_folder,
|
||||
"type": "output",
|
||||
})
|
||||
|
||||
return {"ui": {"images": results}}
|
||||
|
||||
def build_param_string(self, meta):
|
||||
pos_prompt = meta.get("Positive prompt", "")
|
||||
neg_prompt = meta.get("Negative prompt", "")
|
||||
steps = meta.get("Steps", 20)
|
||||
sampler = meta.get("Sampler", "euler")
|
||||
cfg_scale = meta.get("CFG scale", 7.0)
|
||||
seed = meta.get("Seed", 0)
|
||||
scheduler = meta.get("Scheduler", "normal")
|
||||
denoise = meta.get("Denoise", 1.0)
|
||||
model_file = meta.get("Model", "UnknownModel")
|
||||
model_hash = meta.get("Model hash", "")
|
||||
lora_hashes = meta.get("Lora hashes", {})
|
||||
hashes = meta.get("Hashes", "{}")
|
||||
|
||||
param_string = (
|
||||
f"{pos_prompt}\n"
|
||||
f"Negative prompt: {neg_prompt}\n"
|
||||
f"Steps: {steps}, Sampler: {sampler}, CFG scale: {cfg_scale}, Seed: {seed}, "
|
||||
f"Scheduler: {scheduler}, Denoise: {denoise}, "
|
||||
f"Model: {model_file}, Model hash: {model_hash}"
|
||||
)
|
||||
if lora_hashes:
|
||||
lora_str = ", ".join(lora_hashes.values())
|
||||
param_string += f", Lora hashes: \"{lora_str}\""
|
||||
if hashes:
|
||||
param_string += f", Hashes: {hashes}"
|
||||
return param_string
|
||||
|
||||
|
||||
def to_uint8(self, image):
|
||||
if hasattr(image, "cpu"):
|
||||
arr = image.cpu().numpy() * 255.0
|
||||
else:
|
||||
arr = np.array(image) * 255.0
|
||||
return np.clip(arr, 0, 255).astype(np.uint8)
|
||||
|
||||
def prepare_pnginfo(self, meta_dict, index, total):
|
||||
pnginfo = PngImagePlugin.PngInfo()
|
||||
if total > 1:
|
||||
pnginfo.add_text("Batch index", str(index))
|
||||
pnginfo.add_text("Batch size", str(total))
|
||||
for key, value in meta_dict.items():
|
||||
pnginfo.add_text(str(key), str(value))
|
||||
return pnginfo
|
||||
|
||||
def format_filename(self, filename, meta_dict):
|
||||
matches = self.pattern_format.findall(filename)
|
||||
now = datetime.now()
|
||||
date_table = {
|
||||
"yyyy": str(now.year),
|
||||
"yy": str(now.year)[-2:],
|
||||
"MM": str(now.month).zfill(2),
|
||||
"dd": str(now.day).zfill(2),
|
||||
"hh": str(now.hour).zfill(2),
|
||||
"mm": str(now.minute).zfill(2),
|
||||
"ss": str(now.second).zfill(2),
|
||||
}
|
||||
for match in matches:
|
||||
if match.startswith("%date"):
|
||||
parts = match.strip("%").split(":")
|
||||
fmt = parts[1] if len(parts) > 1 else "yyyyMMddhhmmss"
|
||||
for k, v in date_table.items():
|
||||
fmt = fmt.replace(k, v)
|
||||
filename = filename.replace(match, fmt)
|
||||
return filename
|
||||
|
||||
def next_available_suffix(self, base_name, folder, ext):
|
||||
pattern = re.compile(rf"^{re.escape(base_name)}_(\d+)\.{ext}$")
|
||||
max_num = 0
|
||||
for fname in os.listdir(folder):
|
||||
m = pattern.match(fname)
|
||||
if m:
|
||||
try:
|
||||
num = int(m.group(1))
|
||||
if num > max_num:
|
||||
max_num = num
|
||||
except Exception:
|
||||
pass
|
||||
return max_num + 1
|
||||
@@ -0,0 +1,189 @@
|
||||
import os
|
||||
import re
|
||||
import json
|
||||
import torch
|
||||
import hashlib
|
||||
from comfy import samplers
|
||||
import folder_paths
|
||||
from .MangoTriggerExporter import get_lora_metadata
|
||||
|
||||
LAST_USED_SEED = None
|
||||
|
||||
def short_sha256(file_path):
|
||||
if not os.path.exists(file_path):
|
||||
return "no_file"
|
||||
h = hashlib.sha256()
|
||||
with open(file_path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(4096), b""):
|
||||
h.update(chunk)
|
||||
return h.hexdigest()[:10]
|
||||
|
||||
def find_model_file_and_hash(model_name):
|
||||
candidate = folder_paths.get_full_path("checkpoints", model_name)
|
||||
if candidate and os.path.exists(candidate):
|
||||
return (os.path.basename(candidate), short_sha256(candidate))
|
||||
return (model_name, "no_file")
|
||||
|
||||
def parse_loras_from_stack(lora_stack):
|
||||
|
||||
results = {}
|
||||
if not lora_stack:
|
||||
return results
|
||||
for item in lora_stack:
|
||||
if isinstance(item, (tuple, list)) and len(item) > 0:
|
||||
lora_name = item[0]
|
||||
if lora_name and lora_name != "None":
|
||||
base = os.path.basename(lora_name)
|
||||
lora_path = folder_paths.get_full_path("loras", lora_name)
|
||||
if lora_path and os.path.exists(lora_path):
|
||||
lora_hash = short_sha256(lora_path)
|
||||
results[base] = f"{base}: {lora_hash}"
|
||||
else:
|
||||
results[base] = base
|
||||
return results
|
||||
|
||||
class BaseNode:
|
||||
pass
|
||||
|
||||
class KSamplerMango(BaseNode):
|
||||
CATEGORY = "Mango Node Pack/Sampling"
|
||||
FUNCTION = "sample"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"latent_image": ("LATENT", {"tooltip": "Latent to sample from"}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"reuse_seed": ("BOOLEAN", {"default": False}),
|
||||
"steps": ("INT", {"default": 20, "min": 1, "max": 200}),
|
||||
"cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0}),
|
||||
"sampler_name": (samplers.KSampler.SAMPLERS,),
|
||||
"scheduler": (samplers.SCHEDULER_NAMES,),
|
||||
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"positive": ("CONDITIONING", {"tooltip": "Positive conditioning tensor"}),
|
||||
"negative": ("CONDITIONING", {"tooltip": "Negative conditioning tensor"}),
|
||||
"positive_prompt_text": ("STRING", {"default": "", "multiline": True}),
|
||||
"negative_prompt_text": ("STRING", {"default": "", "multiline": True}),
|
||||
"ckpt_name": ("STRING", {"default": "", "multiline": False, "tooltip": "Checkpoint name (from MangoLoader)"}),
|
||||
"lora_stack": ("LORA_STACK", {"tooltip": "Optional LoRA stack (from MangoLoader)"}),
|
||||
},
|
||||
"hidden": {
|
||||
"prompt": "PROMPT",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO",
|
||||
"my_unique_id": "UNIQUE_ID",
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT", "METADATA", "STRING")
|
||||
RETURN_NAMES = ("LATENT", "METADATA", "METADATA_TEXT")
|
||||
DESCRIPTION = (
|
||||
"KSampler (Mango) that assembles metadata including prompts, checkpoint info, and LoRA info. "
|
||||
"It uses the 'ckpt_name' input to look up the checkpoint file and compute its hash, and processes the LoRA stack to produce LoRA info."
|
||||
)
|
||||
|
||||
def sample(
|
||||
self,
|
||||
model,
|
||||
latent_image,
|
||||
seed,
|
||||
reuse_seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
denoise,
|
||||
positive,
|
||||
negative,
|
||||
positive_prompt_text,
|
||||
negative_prompt_text,
|
||||
ckpt_name,
|
||||
lora_stack,
|
||||
prompt=None,
|
||||
extra_pnginfo=None,
|
||||
my_unique_id=None
|
||||
):
|
||||
global LAST_USED_SEED
|
||||
used_seed = LAST_USED_SEED if (reuse_seed and LAST_USED_SEED is not None) else seed
|
||||
|
||||
if hasattr(model, "device"):
|
||||
device = model.device
|
||||
else:
|
||||
try:
|
||||
device = model.get_model_object("model").device
|
||||
except Exception:
|
||||
device = torch.device("cpu")
|
||||
|
||||
if not isinstance(latent_image, dict) or "samples" not in latent_image:
|
||||
raise ValueError("Invalid latent_image input. Provide a valid LATENT node output.")
|
||||
latent_tensor = latent_image["samples"]
|
||||
|
||||
generator = torch.Generator(device=device).manual_seed(used_seed)
|
||||
noise = torch.randn(latent_tensor.shape, generator=generator, device=device, dtype=latent_tensor.dtype)
|
||||
|
||||
ksampler = samplers.KSampler(
|
||||
model,
|
||||
steps,
|
||||
device,
|
||||
sampler=sampler_name,
|
||||
scheduler=scheduler,
|
||||
denoise=denoise,
|
||||
model_options={}
|
||||
)
|
||||
out_latent = ksampler.sample(
|
||||
noise,
|
||||
positive,
|
||||
negative,
|
||||
cfg,
|
||||
latent_image=latent_tensor,
|
||||
seed=used_seed
|
||||
)
|
||||
LAST_USED_SEED = used_seed
|
||||
|
||||
metadata = {}
|
||||
metadata["Seed"] = used_seed
|
||||
metadata["Steps"] = steps
|
||||
metadata["CFG scale"] = cfg
|
||||
metadata["Sampler"] = sampler_name
|
||||
metadata["Scheduler"] = scheduler
|
||||
metadata["Denoise"] = denoise
|
||||
|
||||
used_ckpt = str(ckpt_name).strip() if ckpt_name and str(ckpt_name).strip() else "UnknownModel"
|
||||
model_filename, model_short_hash = find_model_file_and_hash(used_ckpt)
|
||||
metadata["Model"] = model_filename
|
||||
metadata["Model hash"] = model_short_hash
|
||||
|
||||
metadata["Positive prompt"] = positive_prompt_text
|
||||
metadata["Negative prompt"] = negative_prompt_text
|
||||
|
||||
if isinstance(lora_stack, tuple):
|
||||
lora_list = lora_stack[0]
|
||||
else:
|
||||
lora_list = lora_stack
|
||||
lora_info = parse_loras_from_stack(lora_list)
|
||||
if lora_info:
|
||||
metadata["Lora hashes"] = lora_info
|
||||
|
||||
hash_dict = {"model": model_short_hash}
|
||||
if lora_list:
|
||||
for item in lora_list:
|
||||
if isinstance(item, (tuple, list)) and len(item) > 0:
|
||||
ln = item[0]
|
||||
if ln and ln != "None":
|
||||
base = os.path.basename(ln)
|
||||
lora_path = folder_paths.get_full_path("loras", ln)
|
||||
if lora_path and os.path.exists(lora_path):
|
||||
hash_dict["lora:" + base] = short_sha256(lora_path)
|
||||
metadata["Hashes"] = json.dumps(hash_dict)
|
||||
|
||||
if extra_pnginfo is not None and isinstance(extra_pnginfo, dict):
|
||||
for k, v in extra_pnginfo.items():
|
||||
s = str(v)
|
||||
metadata[k] = s if len(s) < 1000 else s[:1000] + " [truncated]"
|
||||
if prompt is not None:
|
||||
s = str(prompt)
|
||||
metadata["prompt"] = s if len(s) < 1000 else s[:1000] + " [truncated]"
|
||||
|
||||
metadata_str = json.dumps(metadata, indent=2)
|
||||
return ({"samples": out_latent}, metadata, metadata_str)
|
||||
@@ -0,0 +1,54 @@
|
||||
import torch
|
||||
|
||||
class LatentImageMango:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
dimensions = [
|
||||
"1:1 square 1024x1024",
|
||||
"3:4 portrait 896x1152",
|
||||
"5:8 portrait 832x1216",
|
||||
"9:16 portrait 768x1344",
|
||||
"9:21 portrait 640x1536",
|
||||
"4:3 landscape 1152x896",
|
||||
"3:2 landscape 1216x832",
|
||||
"16:9 landscape 1344x768",
|
||||
"21:9 landscape 1536x640",
|
||||
"32:9 landscape 1600x450"
|
||||
]
|
||||
return {
|
||||
"required": {
|
||||
"dimensions": (dimensions,)
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT", "INT", "INT")
|
||||
RETURN_NAMES = ("latent_image", "width", "height")
|
||||
FUNCTION = "compute_dimensions"
|
||||
CATEGORY = "Mango Node Pack/Loaders"
|
||||
|
||||
def compute_dimensions(self, dimensions):
|
||||
if dimensions == "1:1 square 1024x1024":
|
||||
width, height = 1024, 1024
|
||||
elif dimensions == "3:4 portrait 896x1152":
|
||||
width, height = 896, 1152
|
||||
elif dimensions == "5:8 portrait 832x1216":
|
||||
width, height = 832, 1216
|
||||
elif dimensions == "9:16 portrait 768x1344":
|
||||
width, height = 768, 1344
|
||||
elif dimensions == "9:21 portrait 640x1536":
|
||||
width, height = 640, 1536
|
||||
elif dimensions == "4:3 landscape 1152x896":
|
||||
width, height = 1152, 896
|
||||
elif dimensions == "3:2 landscape 1216x832":
|
||||
width, height = 1216, 832
|
||||
elif dimensions == "16:9 landscape 1344x768":
|
||||
width, height = 1344, 768
|
||||
elif dimensions == "21:9 landscape 1536x640":
|
||||
width, height = 1536, 640
|
||||
elif dimensions == "32:9 landscape 1600x450":
|
||||
width, height = 1600, 450
|
||||
else:
|
||||
width, height = 1024, 1024
|
||||
latent_image = torch.zeros([1, 4, height // 8, width // 8])
|
||||
return {"samples": latent_image}, width, height
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
import os
|
||||
import hashlib
|
||||
import folder_paths
|
||||
import comfy.sd
|
||||
|
||||
def short_hash(path):
|
||||
if not os.path.exists(path):
|
||||
return "no_file"
|
||||
h = hashlib.sha256()
|
||||
with open(path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(4096), b""):
|
||||
h.update(chunk)
|
||||
return h.hexdigest()[:10]
|
||||
|
||||
def load_ckpt(ckpt_name):
|
||||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||||
if not ckpt_path or not os.path.exists(ckpt_path):
|
||||
raise ValueError(f"Checkpoint not found: {ckpt_name}")
|
||||
loaded = comfy.sd.load_checkpoint_guess_config(
|
||||
ckpt_path,
|
||||
output_vae=True,
|
||||
output_clip=True,
|
||||
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
||||
)
|
||||
model, clip, vae = loaded[:3]
|
||||
return model, clip, vae, short_hash(ckpt_path)
|
||||
|
||||
def load_single_lora(model, clip, lora_name, weight):
|
||||
lora_path = folder_paths.get_full_path("loras", lora_name)
|
||||
if not lora_path or not os.path.exists(lora_path):
|
||||
return model, clip
|
||||
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
||||
model, clip = comfy.sd.load_lora_for_models(model, clip, lora, weight, weight)
|
||||
return model, clip
|
||||
|
||||
class MangoLoader:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
ckpts = folder_paths.get_filename_list("checkpoints")
|
||||
loras = ["None"] + folder_paths.get_filename_list("loras")
|
||||
return {
|
||||
"required": {
|
||||
"ckpt_name": (ckpts, {"default": ckpts[0] if ckpts else None, "tooltip": "Select a checkpoint"}),
|
||||
"LoraName1": (loras, {"default": "None", "tooltip": "LoRA 1 filename"}),
|
||||
"LoraWeight1": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05, "tooltip": "LoRA 1 weight"}),
|
||||
"LoraName2": (loras, {"default": "None", "tooltip": "LoRA 2 filename"}),
|
||||
"LoraWeight2": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05, "tooltip": "LoRA 2 weight"}),
|
||||
"LoraName3": (loras, {"default": "None", "tooltip": "LoRA 3 filename"}),
|
||||
"LoraWeight3": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05, "tooltip": "LoRA 3 weight"}),
|
||||
"LoraName4": (loras, {"default": "None", "tooltip": "LoRA 4 filename"}),
|
||||
"LoraWeight4": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05, "tooltip": "LoRA 4 weight"}),
|
||||
"LoraName5": (loras, {"default": "None", "tooltip": "LoRA 5 filename"}),
|
||||
"LoraWeight5": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05, "tooltip": "LoRA 5 weight"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE", "LORA_STACK", "STRING", "STRING")
|
||||
RETURN_NAMES = ("model", "clip", "vae", "lora_stack", "ckpt_name", "ckpt_hash")
|
||||
FUNCTION = "load_checkpoint_and_loras"
|
||||
CATEGORY = "Mango Node Pack/Loaders"
|
||||
|
||||
def load_checkpoint_and_loras(self, ckpt_name, **kwargs):
|
||||
model, clip, vae, ckpt_hash = load_ckpt(ckpt_name)
|
||||
lora_stack = []
|
||||
for i in range(1, 6):
|
||||
name = kwargs.get(f"LoraName{i}")
|
||||
weight = kwargs.get(f"LoraWeight{i}", 1.0)
|
||||
if name and name != "None":
|
||||
model, clip = load_single_lora(model, clip, name, weight)
|
||||
lora_stack.append((name, weight, weight))
|
||||
return (model, clip, vae, lora_stack, ckpt_name, ckpt_hash)
|
||||
@@ -0,0 +1,126 @@
|
||||
"""
|
||||
Dependencies:
|
||||
- pip install safetensors
|
||||
- ComfyUI's 'folder_paths' for locating LoRA files.
|
||||
"""
|
||||
|
||||
import os
|
||||
import hashlib
|
||||
import json
|
||||
import requests
|
||||
|
||||
try:
|
||||
from safetensors import safe_open
|
||||
except ImportError:
|
||||
safe_open = None
|
||||
|
||||
import folder_paths
|
||||
|
||||
#METADATA FETCHING
|
||||
|
||||
def calculate_sha256(file_path):
|
||||
sha256_hash = hashlib.sha256()
|
||||
with open(file_path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(4096), b""):
|
||||
sha256_hash.update(chunk)
|
||||
return sha256_hash.hexdigest()
|
||||
|
||||
def get_model_version_info(hash_value):
|
||||
api_url = f"https://civitai.com/api/v1/model-versions/by-hash/{hash_value}"
|
||||
response = requests.get(api_url)
|
||||
if response.status_code == 200:
|
||||
return response.json()
|
||||
else:
|
||||
return {}
|
||||
|
||||
def parse_local_safetensors_metadata(lora_path):
|
||||
if not safe_open:
|
||||
return {}
|
||||
try:
|
||||
with safe_open(lora_path, framework="pt", device="cpu") as f:
|
||||
meta = f.metadata() or {}
|
||||
print(f"DEBUG: Safetensors metadata for {lora_path} -> {meta}")
|
||||
return meta
|
||||
except Exception as e:
|
||||
print(f"Failed reading local safetensors metadata for {lora_path}: {e}")
|
||||
return {}
|
||||
|
||||
def extract_trigger_words_from_metadata(meta):
|
||||
triggers_found = set()
|
||||
|
||||
if "trainedWords" in meta and meta["trainedWords"]:
|
||||
val = meta["trainedWords"]
|
||||
if isinstance(val, str):
|
||||
triggers_found.update(x.strip() for x in val.split(",") if x.strip())
|
||||
elif isinstance(val, list):
|
||||
triggers_found.update(val)
|
||||
elif isinstance(val, dict):
|
||||
triggers_found.update(val.keys())
|
||||
|
||||
if "modelspec.trigger_phrase" in meta:
|
||||
trigger_phrase = meta["modelspec.trigger_phrase"]
|
||||
if isinstance(trigger_phrase, str) and trigger_phrase.strip():
|
||||
triggers_found.add(trigger_phrase.strip())
|
||||
|
||||
return list(triggers_found)
|
||||
|
||||
def get_lora_metadata(lora_name):
|
||||
db_path = os.path.join(os.path.dirname(__file__), 'lora_metadata_db.json')
|
||||
try:
|
||||
with open(db_path, 'r') as f:
|
||||
db = json.load(f)
|
||||
except Exception:
|
||||
db = {}
|
||||
if lora_name in db:
|
||||
return db[lora_name]
|
||||
lora_path = folder_paths.get_full_path("loras", lora_name)
|
||||
if not os.path.exists(lora_path):
|
||||
db[lora_name] = {"triggerWords": ""}
|
||||
with open(db_path, 'w') as f:
|
||||
json.dump(db, f, indent=4)
|
||||
return db[lora_name]
|
||||
meta = parse_local_safetensors_metadata(lora_path)
|
||||
local_triggers = extract_trigger_words_from_metadata(meta)
|
||||
if not local_triggers:
|
||||
LORAsha256 = calculate_sha256(lora_path)
|
||||
model_info = get_model_version_info(LORAsha256)
|
||||
if model_info.get("trainedWords"):
|
||||
local_triggers = model_info["trainedWords"]
|
||||
triggers_str = ", ".join(local_triggers)
|
||||
db[lora_name] = {"triggerWords": triggers_str}
|
||||
try:
|
||||
with open(db_path, 'w') as f:
|
||||
json.dump(db, f, indent=4)
|
||||
except Exception as e:
|
||||
print(f"Error saving metadata cache: {e}")
|
||||
return db[lora_name]
|
||||
|
||||
#NODE
|
||||
|
||||
class MangoTriggerExporter:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"lora_stack": ("LORA_STACK",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
FUNCTION = "export_triggerwords"
|
||||
CATEGORY = "Mango Node Pack/Metadata"
|
||||
|
||||
def export_triggerwords(self, lora_stack):
|
||||
triggerwords_list = []
|
||||
for item in lora_stack:
|
||||
lora_name = item[0] if isinstance(item, (tuple, list)) and item else None
|
||||
if lora_name and lora_name != "None":
|
||||
try:
|
||||
metadata = get_lora_metadata(lora_name)
|
||||
tw = metadata.get("triggerWords", "")
|
||||
if tw:
|
||||
triggerwords_list.append(tw)
|
||||
except Exception as e:
|
||||
print(f"Error processing '{lora_name}': {e}")
|
||||
combined_triggerwords = ", ".join(triggerwords_list)
|
||||
return (combined_triggerwords,)
|
||||
@@ -0,0 +1,48 @@
|
||||
class PromptMango:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"QualityTags": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"rows": 3
|
||||
}
|
||||
),
|
||||
"SceneDescription": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"rows": 3
|
||||
}
|
||||
),
|
||||
"Tags": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"rows": 3
|
||||
}
|
||||
),
|
||||
"string": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"forceInput": True,
|
||||
}
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "Mango Node Pack/Metadata"
|
||||
|
||||
def execute(self, QualityTags, SceneDescription, Tags, string):
|
||||
parts = [s.strip() for s in (QualityTags, SceneDescription, Tags, string) if s.strip()]
|
||||
concatenated = ", ".join(parts)
|
||||
return (concatenated,)
|
||||
+24
@@ -0,0 +1,24 @@
|
||||
from .MangoTriggerExporter import MangoTriggerExporter
|
||||
from .KSamplerMango import KSamplerMango
|
||||
from .MangoLoader import MangoLoader
|
||||
from .ImageSaverMango import ImageSaverMango
|
||||
from .LatentImageMango import LatentImageMango
|
||||
from .PromptMango import PromptMango
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"MangoTriggerExporter": MangoTriggerExporter,
|
||||
"KSamplerMango": KSamplerMango,
|
||||
"MangoLoader": MangoLoader,
|
||||
"ImageSaverMango": ImageSaverMango,
|
||||
"LatentImageMango": LatentImageMango,
|
||||
"PromptMango": PromptMango
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MangoTriggerExporter": "Trigger Exporter (Mango)",
|
||||
"KSamplerMango": "KSampler (Mango)",
|
||||
"MangoLoader": "Loader (Mango)",
|
||||
"ImageSaverMango": "Image Saver (Mango)",
|
||||
"LatentImageMango": "Latent Image (Mango)",
|
||||
"PromptMango": "Prompt (Mango)"
|
||||
}
|
||||
Reference in New Issue
Block a user