Add files via upload

This commit is contained in:
mango1010
2025-02-28 23:47:19 +01:00
committed by GitHub
commit d80d88fb0a
7 changed files with 666 additions and 0 deletions
+153
View File
@@ -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
+189
View File
@@ -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)
+54
View File
@@ -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
+72
View File
@@ -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)
+126
View File
@@ -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,)
+48
View File
@@ -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
View File
@@ -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)"
}