From 9c388e6c5bb865ef682e97da86cbdb22c4952ed0 Mon Sep 17 00:00:00 2001 From: mang01010 <201291354+mang01010@users.noreply.github.com> Date: Mon, 10 Mar 2025 16:09:13 +0100 Subject: [PATCH] Updated __init__.py and added new nodes: CompositeMangoLoader, FluxGuidanceMango, FluxSampler --- CompositeMangoLoader.py | 145 +++++++++++++++++++++++++ FluxGuidanceMango.py | 18 ++++ FluxSamplerMango.py | 234 ++++++++++++++++++++++++++++++++++++++++ __init__.py | 9 ++ 4 files changed, 406 insertions(+) create mode 100644 CompositeMangoLoader.py create mode 100644 FluxGuidanceMango.py create mode 100644 FluxSamplerMango.py diff --git a/CompositeMangoLoader.py b/CompositeMangoLoader.py new file mode 100644 index 0000000..24f6d30 --- /dev/null +++ b/CompositeMangoLoader.py @@ -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) diff --git a/FluxGuidanceMango.py b/FluxGuidanceMango.py new file mode 100644 index 0000000..0149a4e --- /dev/null +++ b/FluxGuidanceMango.py @@ -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) diff --git a/FluxSamplerMango.py b/FluxSamplerMango.py new file mode 100644 index 0000000..7e9f0f2 --- /dev/null +++ b/FluxSamplerMango.py @@ -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) diff --git a/__init__.py b/__init__.py index 349a8a0..b89929f 100644 --- a/__init__.py +++ b/__init__.py @@ -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)", }