Updated __init__.py and added new nodes: CompositeMangoLoader, FluxGuidanceMango, FluxSampler
This commit is contained in:
mang01010
@@ -0,0 +1,145 @@
|
||||
import os
|
||||
import hashlib
|
||||
import folder_paths
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
import torch
|
||||
|
||||
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_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 CompositeMangoLoader:
|
||||
@staticmethod
|
||||
def vae_list():
|
||||
vaes = folder_paths.get_filename_list("vae")
|
||||
approx_vaes = folder_paths.get_filename_list("vae_approx")
|
||||
taesd_combinations = {
|
||||
"taesd": ("taesd_encoder", "taesd_decoder"),
|
||||
"taesdxl": ("taesdxl_encoder", "taesdxl_decoder"),
|
||||
"taesd3": ("taesd3_encoder", "taesd3_decoder"),
|
||||
"taef1": ("taef1_encoder", "taef1_decoder")
|
||||
}
|
||||
for name, (enc_prefix, dec_prefix) in taesd_combinations.items():
|
||||
if any(v.startswith(enc_prefix) for v in approx_vaes) and any(v.startswith(dec_prefix) for v in approx_vaes):
|
||||
vaes.append(name)
|
||||
return vaes
|
||||
|
||||
@staticmethod
|
||||
def load_taesd(name):
|
||||
sd = {}
|
||||
approx_vaes = folder_paths.get_filename_list("vae_approx")
|
||||
enc_prefix, dec_prefix = {
|
||||
"taesd": ("taesd_encoder", "taesd_decoder"),
|
||||
"taesdxl": ("taesdxl_encoder", "taesdxl_decoder"),
|
||||
"taesd3": ("taesd3_encoder", "taesd3_decoder"),
|
||||
"taef1": ("taef1_encoder", "taef1_decoder")
|
||||
}[name]
|
||||
|
||||
encoder = next(v for v in approx_vaes if v.startswith(enc_prefix))
|
||||
decoder = next(v for v in approx_vaes if v.startswith(dec_prefix))
|
||||
|
||||
enc = comfy.utils.load_torch_file(folder_paths.get_full_path("vae_approx", encoder))
|
||||
dec = comfy.utils.load_torch_file(folder_paths.get_full_path("vae_approx", decoder))
|
||||
|
||||
for k in enc: sd[f"{enc_prefix}.{k}"] = enc[k]
|
||||
for k in dec: sd[f"{dec_prefix}.{k}"] = dec[k]
|
||||
|
||||
sd.update({
|
||||
"taesd": (0.18215, 0.0),
|
||||
"taesdxl": (0.13025, 0.0),
|
||||
"taesd3": (1.5305, 0.0609),
|
||||
"taef1": (0.3611, 0.1159)
|
||||
}[name])
|
||||
sd["vae_scale"], sd["vae_shift"] = torch.tensor(sd["vae_scale"]), torch.tensor(sd["vae_shift"])
|
||||
return sd
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"unet_name": (folder_paths.get_filename_list("diffusion_models"),),
|
||||
"weight_dtype": (["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"],),
|
||||
"clip_name1": (folder_paths.get_filename_list("text_encoders"),),
|
||||
"clip_name2": (folder_paths.get_filename_list("text_encoders"),),
|
||||
"type": (["sdxl", "sd3", "flux", "hunyuan_video"],),
|
||||
"vae_name": (cls.vae_list(),),
|
||||
"LoraName1": (["None"] + folder_paths.get_filename_list("loras"), {"default": "None"}),
|
||||
"LoraWeight1": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05}),
|
||||
"LoraName2": (["None"] + folder_paths.get_filename_list("loras"), {"default": "None"}),
|
||||
"LoraWeight2": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05}),
|
||||
"LoraName3": (["None"] + folder_paths.get_filename_list("loras"), {"default": "None"}),
|
||||
"LoraWeight3": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05}),
|
||||
"LoraName4": (["None"] + folder_paths.get_filename_list("loras"), {"default": "None"}),
|
||||
"LoraWeight4": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05}),
|
||||
"LoraName5": (["None"] + folder_paths.get_filename_list("loras"), {"default": "None"}),
|
||||
"LoraWeight5": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.05}),
|
||||
},
|
||||
"optional": {
|
||||
"device": (["default", "cpu"], {"advanced": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE", "LORA_STACK", "STRING", "STRING")
|
||||
RETURN_NAMES = ("model", "clip", "vae", "lora_stack", "unet_name", "unet_hash")
|
||||
FUNCTION = "load_all"
|
||||
CATEGORY = "Mango Node Pack/Loaders"
|
||||
|
||||
def load_all(self, unet_name, weight_dtype, clip_name1, clip_name2, type, vae_name, device="default", **kwargs):
|
||||
# Load UNET
|
||||
model_options = {}
|
||||
if weight_dtype.startswith("fp8"):
|
||||
model_options["dtype"] = getattr(torch, f"float{weight_dtype[2:]}")
|
||||
if "fast" in weight_dtype:
|
||||
model_options["fp8_optimizations"] = True
|
||||
|
||||
unet_path = folder_paths.get_full_path_or_raise("diffusion_models", unet_name)
|
||||
unet_hash = short_hash(unet_path)
|
||||
model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options)
|
||||
|
||||
# Load Dual CLIP
|
||||
clip_path1 = folder_paths.get_full_path_or_raise("text_encoders", clip_name1)
|
||||
clip_path2 = folder_paths.get_full_path_or_raise("text_encoders", clip_name2)
|
||||
clip_type = getattr(comfy.sd.CLIPType, type.upper())
|
||||
|
||||
model_options_clip = {}
|
||||
if device == "cpu":
|
||||
model_options_clip.update({"load_device": "cpu", "offload_device": "cpu"})
|
||||
|
||||
clip = comfy.sd.load_clip(
|
||||
ckpt_paths=[clip_path1, clip_path2],
|
||||
clip_type=clip_type,
|
||||
model_options=model_options_clip
|
||||
)
|
||||
|
||||
# Load VAE
|
||||
if vae_name in ["taesd", "taesdxl", "taesd3", "taef1"]:
|
||||
vae_sd = self.load_taesd(vae_name)
|
||||
else:
|
||||
vae_path = folder_paths.get_full_path_or_raise("vae", vae_name)
|
||||
vae_sd = comfy.utils.load_torch_file(vae_path)
|
||||
vae = comfy.sd.VAE(sd=vae_sd)
|
||||
|
||||
# Apply LORAs
|
||||
lora_stack = []
|
||||
for i in range(1, 6):
|
||||
lora_name = kwargs.get(f"LoraName{i}")
|
||||
lora_weight = kwargs.get(f"LoraWeight{i}", 1.0)
|
||||
if lora_name and lora_name != "None":
|
||||
model, clip = load_single_lora(model, clip, lora_name, lora_weight)
|
||||
lora_stack.append((lora_name, lora_weight, lora_weight))
|
||||
|
||||
return (model, clip, vae, lora_stack, unet_name, unet_hash)
|
||||
@@ -0,0 +1,18 @@
|
||||
import node_helpers
|
||||
|
||||
class FluxGuidanceMango:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"conditioning": ("CONDITIONING",),
|
||||
"guidance": ("FLOAT", {"default": 3.5, "min": 0.0, "max": 100.0, "step": 0.1}),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("CONDITIONING", "FLOAT")
|
||||
RETURN_NAMES = ("conditioned", "flux_guidance")
|
||||
FUNCTION = "append"
|
||||
CATEGORY = "Mango Node Pack/Metadata"
|
||||
|
||||
def append(self, conditioning, guidance):
|
||||
c = node_helpers.conditioning_set_values(conditioning, {"guidance": guidance})
|
||||
return (c, guidance)
|
||||
@@ -0,0 +1,234 @@
|
||||
import os
|
||||
import json
|
||||
import hashlib
|
||||
import torch
|
||||
import comfy.samplers
|
||||
import comfy.sample
|
||||
import comfy.utils
|
||||
import latent_preview
|
||||
import folder_paths
|
||||
|
||||
# Global variables for noise reuse
|
||||
LAST_USED_NOISE = None
|
||||
LAST_USED_NOISE_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 FluxNoise:
|
||||
"""
|
||||
A simple noise wrapper mimicking a random noise generator.
|
||||
"""
|
||||
def __init__(self, seed, noise_tensor):
|
||||
self.seed = seed
|
||||
self.noise_tensor = noise_tensor
|
||||
|
||||
def generate_noise(self, latent):
|
||||
return self.noise_tensor
|
||||
|
||||
def inline_basic_guider(model, conditioning):
|
||||
"""
|
||||
Build a simple guider using CFGGuider with a single 'positive' conditioning.
|
||||
(Note: guidance is not embedded in conditioning here; it comes via flux_guidance.)
|
||||
"""
|
||||
guider = comfy.samplers.CFGGuider(model)
|
||||
guider.inner_set_conds({"positive": conditioning})
|
||||
return guider
|
||||
|
||||
def inline_basic_scheduler(model, scheduler, steps, denoise):
|
||||
"""
|
||||
Compute sigmas using the same logic as BasicScheduler.
|
||||
"""
|
||||
total_steps = steps
|
||||
if denoise < 1.0:
|
||||
if denoise <= 0.0:
|
||||
return torch.FloatTensor([])
|
||||
total_steps = int(steps / denoise)
|
||||
sigmas = comfy.samplers.calculate_sigmas(
|
||||
model.get_model_object("model_sampling"), scheduler, total_steps
|
||||
).cpu()
|
||||
sigmas = sigmas[-(steps + 1):]
|
||||
return sigmas
|
||||
|
||||
class FluxSamplerMango:
|
||||
|
||||
CATEGORY = "Mango Node Pack/Sampling"
|
||||
FUNCTION = "sample"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"conditioning": ("CONDITIONING",),
|
||||
"flux_guidance": ("FLOAT", {"default": 3.5, "min": 0.0, "max": 100.0, "step": 0.1}),
|
||||
"lora_stack": ("LORA_STACK",),
|
||||
"latent_image": ("LATENT",),
|
||||
"prompt_text": ("STRING", {"default": "", "multiline": True}),
|
||||
"unet_name": ("STRING", {"default": "", "tooltip": "UNET/Checkpoint name"}),
|
||||
"unet_hash": ("STRING", {"default": "", "tooltip": "UNET hash from CompositeMangoLoader"}),
|
||||
"noise_seed": ("INT", {
|
||||
"default": 1234,
|
||||
"min": 0,
|
||||
"max": 0xffffffffffffffff,
|
||||
"control_after_generate": True
|
||||
}),
|
||||
"reuse_noise": ("BOOLEAN", {"default": False}),
|
||||
"sampler_name": (comfy.samplers.SAMPLER_NAMES,),
|
||||
"scheduler": (comfy.samplers.SCHEDULER_NAMES,),
|
||||
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
||||
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT", "METADATA", "STRING")
|
||||
RETURN_NAMES = ("output", "METADATA", "METADATA_TEXT")
|
||||
|
||||
def sample(
|
||||
self,
|
||||
model,
|
||||
conditioning,
|
||||
flux_guidance,
|
||||
lora_stack,
|
||||
latent_image,
|
||||
prompt_text,
|
||||
unet_name,
|
||||
unet_hash,
|
||||
noise_seed,
|
||||
reuse_noise,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
steps,
|
||||
denoise,
|
||||
):
|
||||
global LAST_USED_NOISE, LAST_USED_NOISE_SEED
|
||||
|
||||
# Determine device from the model
|
||||
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 = latent_image.copy()
|
||||
original_tensor = latent["samples"]
|
||||
|
||||
# Build the guider using the provided conditioning (without altering resolution)
|
||||
guider = inline_basic_guider(model, conditioning)
|
||||
# Fix latent channels using the model directly (as in SamplerCustomAdvanced)
|
||||
fixed_latent = comfy.sample.fix_empty_latent_channels(model, original_tensor)
|
||||
latent["samples"] = fixed_latent
|
||||
|
||||
if reuse_noise and (LAST_USED_NOISE is not None):
|
||||
used_noise_tensor = LAST_USED_NOISE
|
||||
used_seed = LAST_USED_NOISE_SEED
|
||||
else:
|
||||
used_seed = noise_seed
|
||||
generator = torch.Generator(device=device).manual_seed(used_seed)
|
||||
used_noise_tensor = torch.randn(fixed_latent.shape, generator=generator, device=device, dtype=fixed_latent.dtype)
|
||||
LAST_USED_NOISE = used_noise_tensor
|
||||
LAST_USED_NOISE_SEED = used_seed
|
||||
|
||||
noise = FluxNoise(used_seed, used_noise_tensor)
|
||||
sampler = comfy.samplers.sampler_object(sampler_name)
|
||||
sigmas = inline_basic_scheduler(model, scheduler, steps, denoise)
|
||||
x0_output = {}
|
||||
callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output)
|
||||
noise_mask = latent_image.get("noise_mask") if "noise_mask" in latent_image else None
|
||||
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
|
||||
|
||||
samples = guider.sample(
|
||||
noise.generate_noise(latent),
|
||||
fixed_latent,
|
||||
sampler,
|
||||
sigmas,
|
||||
denoise_mask=noise_mask,
|
||||
callback=callback,
|
||||
disable_pbar=disable_pbar,
|
||||
seed=noise.seed,
|
||||
)
|
||||
samples = samples.to(comfy.model_management.intermediate_device())
|
||||
|
||||
out = latent.copy()
|
||||
out["samples"] = samples
|
||||
if "x0" in x0_output:
|
||||
out_denoised = latent.copy()
|
||||
out_denoised["samples"] = model.model.process_latent_out(x0_output["x0"].cpu())
|
||||
else:
|
||||
out_denoised = out
|
||||
|
||||
metadata = {}
|
||||
metadata["Seed"] = used_seed
|
||||
metadata["Steps"] = steps
|
||||
metadata["Sampler"] = sampler_name
|
||||
metadata["Scheduler"] = scheduler
|
||||
metadata["Denoise"] = denoise
|
||||
|
||||
# Use the separate flux_guidance float for CFG scale
|
||||
metadata["CFG scale"] = flux_guidance
|
||||
|
||||
used_ckpt = str(unet_name).strip() if unet_name and str(unet_name).strip() else "UnknownModel"
|
||||
model_filename, model_short_hash = find_model_file_and_hash(used_ckpt)
|
||||
metadata["Model"] = unet_name.strip() if unet_name.strip() else "UnknownModel"
|
||||
metadata["Model hash"] = unet_hash if unet_hash else "no_hash"
|
||||
|
||||
metadata["Positive prompt"] = prompt_text
|
||||
metadata["Negative prompt"] = ""
|
||||
|
||||
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)
|
||||
else:
|
||||
hash_dict["lora:" + base] = base
|
||||
metadata["Hashes"] = json.dumps(hash_dict)
|
||||
metadata_str = json.dumps(metadata, indent=2)
|
||||
|
||||
return ({"samples": out["samples"]}, metadata, metadata_str)
|
||||
@@ -5,6 +5,9 @@ from .ImageSaverMango import ImageSaverMango
|
||||
from .LatentImageMango import LatentImageMango
|
||||
from .PromptMango import PromptMango
|
||||
from .PromptEmbedMango import PromptEmbedMango
|
||||
from .CompositeMangoLoader import CompositeMangoLoader
|
||||
from .FluxSamplerMango import FluxSamplerMango
|
||||
from .FluxGuidanceMango import FluxGuidanceMango
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"MangoTriggerExporter": MangoTriggerExporter,
|
||||
@@ -14,6 +17,9 @@ NODE_CLASS_MAPPINGS = {
|
||||
"LatentImageMango": LatentImageMango,
|
||||
"PromptMango": PromptMango,
|
||||
"PromptEmbedMango": PromptEmbedMango,
|
||||
"CompositeMangoLoader": CompositeMangoLoader,
|
||||
"FluxSamplerMango": FluxSamplerMango,
|
||||
"FluxGuidanceMango": FluxGuidanceMango,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
@@ -24,4 +30,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"LatentImageMango": "Latent Image (Mango)",
|
||||
"PromptMango": "Prompt (Mango)",
|
||||
"PromptEmbedMango": "Prompt /w Embedding (Mango)",
|
||||
"CompositeMangoLoader": "Diffusion Loader (Mango)",
|
||||
"FluxSamplerMango": "FluxSampler (Mango)",
|
||||
"FluxGuidanceMango": "FluxGuidance (Mango)",
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user