251 lines
9.0 KiB
Python
251 lines
9.0 KiB
Python
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}),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
}
|
|
}
|
|
|
|
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,
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
):
|
|
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
|
|
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)
|
|
|
|
# Include extra PNG info and workflow metadata if provided
|
|
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["samples"]}, metadata, metadata_str)
|