207 lines
11 KiB
Python
207 lines
11 KiB
Python
"""PromptForge Sampler: KSampler / KSampler (Advanced) / custom sigmas in one node.
|
|
|
|
It runs ComfyUI's own `comfy.sample.sample`, so every sampler, scheduler, model patch and preview works as in
|
|
the stock nodes. On top of that:
|
|
|
|
* `scheduler = auto` picks the family's schedule (LTX Video's token-count dependent one, otherwise simple),
|
|
* sigmas can come from a SIGMAS input or a typed list,
|
|
* `runs` samples several seeds in one go and returns them as one batch,
|
|
* settings can be saved as named presets, and recommended settings come from the model's sidecar and family
|
|
(see sampler_defaults.py).
|
|
"""
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
|
|
import comfy.model_management
|
|
import comfy.sample
|
|
import comfy.samplers
|
|
import comfy.utils
|
|
import folder_paths
|
|
import latent_preview
|
|
from comfy_extras import nodes_easycache
|
|
|
|
from . import families, sampler_defaults
|
|
|
|
AUTO = "auto"
|
|
PRESET_FIELDS = ("steps", "cfg", "sampler_name", "scheduler", "denoise", "runs", "step_cache", "cache_threshold")
|
|
# step caching skips whole model calls, which only pays off (and only stays clean) with enough steps
|
|
AUTO_CACHE_MIN_STEPS = 15
|
|
CACHES = {"easycache": nodes_easycache.EasyCacheNode, "lazycache": nodes_easycache.LazyCacheNode}
|
|
|
|
|
|
def _preset_file():
|
|
return Path(folder_paths.models_dir) / "PromptForge" / "sampler_presets.json"
|
|
|
|
|
|
def _legacy_preset_file():
|
|
return Path(folder_paths.get_user_directory()) / "promptforge" / "sampler_presets.json"
|
|
|
|
|
|
def load_presets():
|
|
for path in (_preset_file(), _legacy_preset_file()):
|
|
try:
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
except (OSError, ValueError):
|
|
continue
|
|
return {}
|
|
|
|
|
|
@PromptServer.instance.routes.get("/promptforge/sampler_presets")
|
|
async def _presets_route(request):
|
|
return web.json_response(load_presets())
|
|
|
|
|
|
@PromptServer.instance.routes.post("/promptforge/sampler_presets")
|
|
async def _save_preset_route(request):
|
|
body = await request.json()
|
|
name = str(body.get("name", "")).strip()
|
|
if not name:
|
|
return web.json_response({"error": "a preset needs a name"}, status=400)
|
|
presets = load_presets()
|
|
if body.get("delete"):
|
|
presets.pop(name, None)
|
|
else:
|
|
presets[name] = {field: body["settings"][field] for field in PRESET_FIELDS if field in body.get("settings", {})}
|
|
for key in ("family", "model"):
|
|
value = str(body.get(key, "")).strip()
|
|
if value:
|
|
presets[name][key] = value
|
|
_preset_file().parent.mkdir(parents=True, exist_ok=True)
|
|
_preset_file().write_text(json.dumps(presets, indent=1), encoding="utf-8")
|
|
return web.json_response(presets)
|
|
|
|
|
|
@PromptServer.instance.routes.get("/promptforge/recommend")
|
|
async def _recommend_route(request):
|
|
family, name, folder = (request.query.get(key, "") for key in ("family", "model", "folder"))
|
|
path = folder_paths.get_full_path(folder, name) if folder and name else None
|
|
settings, source = sampler_defaults.recommend(family, name, path)
|
|
return web.json_response({"settings": settings, "source": source})
|
|
|
|
|
|
def parse_sigmas(text):
|
|
"""'1, 0.8, 0.5' -> descending sigma tensor ending in 0."""
|
|
values = [float(part) for part in text.replace("\n", ",").split(",") if part.strip()]
|
|
if len(values) < 2 or any(a < b for a, b in zip(values, values[1:])):
|
|
raise ValueError("[PromptForge] custom_sigmas must be a descending list of at least two numbers, for example 1, 0.8, 0.5, 0.")
|
|
if values[-1] > 0:
|
|
values.append(0.0)
|
|
return torch.tensor(values, dtype=torch.float32)
|
|
|
|
|
|
class PROMPTFORGE_Sampler:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"latent_image": ("LATENT",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 2 ** 64 - 1}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler"}),
|
|
"scheduler": ([AUTO] + comfy.samplers.KSampler.SCHEDULERS, {"default": AUTO, "tooltip": "auto uses a simple schedule."}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"runs": ("INT", {"default": 1, "min": 1, "max": 64, "tooltip": "Samples this many seeds (seed, seed+1 ...) and returns them as one batch."}),
|
|
},
|
|
"optional": {
|
|
"pipe": ("PROMPTFORGE_PIPE", {"tooltip": "From a PromptForge Loader: sets the family and fills recommended settings when connected."}),
|
|
"sigmas": ("SIGMAS", {"tooltip": "Overrides the scheduler."}),
|
|
"custom_sigmas": ("STRING", {"default": "", "tooltip": "Comma separated descending sigmas, for example 1, 0.8, 0.5, 0. Overrides the scheduler."}),
|
|
"step_cache": (["off", "auto", "easycache", "lazycache"], {"default": "off", "tooltip": "Skips model calls on steps where the output barely changes. easycache is general and also handles audio+video models; lazycache is cheaper but rougher. auto = easycache when steps >= 15, else off. Useful for many-step runs; distilled 4-8 step models should stay off."}),
|
|
"cache_threshold": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 3.0, "step": 0.01, "tooltip": "Higher skips more steps (faster, softer). 0.1-0.2 is conservative."}),
|
|
"cache_start": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01, "advanced": True, "tooltip": "Fraction of the run before which nothing is skipped."}),
|
|
"cache_end": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01, "advanced": True}),
|
|
"start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000, "advanced": True}),
|
|
"end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000, "advanced": True}),
|
|
"add_noise": ("BOOLEAN", {"default": True, "advanced": True}),
|
|
"preview": (["default", "taesd", "latent2rgb", "none"], {"default": "default", "tooltip": "Live preview while sampling. default follows ComfyUI's --preview-method; taesd uses the model's tiny decoder when it is in models/vae_approx (taeh3 for MiniMax H3) and latent2rgb otherwise."}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "SIGMAS", "STRING")
|
|
RETURN_NAMES = ("LATENT", "SIGMAS", "INFO")
|
|
FUNCTION = "sample"
|
|
CATEGORY = "🔮 PromptForge/Sampling"
|
|
DESCRIPTION = "KSampler with family-aware schedules, custom sigmas, multiple seeds per run and saved presets."
|
|
|
|
@staticmethod
|
|
def _family(model, pipe):
|
|
if pipe and pipe.get("family"):
|
|
return pipe["family"]
|
|
config = getattr(getattr(model, "model", None), "model_config", None)
|
|
family = families.family_from_config(type(config).__name__) if config is not None else None
|
|
return family.key if family else ""
|
|
|
|
def _sigmas(self, model, sampler_name, scheduler, steps, denoise, sigmas, custom_sigmas):
|
|
"""The sigmas to sample with: an input, a typed list, or the named scheduler's."""
|
|
if sigmas is not None:
|
|
return sigmas
|
|
if custom_sigmas.strip():
|
|
return parse_sigmas(custom_sigmas)
|
|
return comfy.samplers.KSampler(model, steps=steps, device=model.load_device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options).sigmas
|
|
|
|
def sample(self, model, positive, negative, latent_image, seed, steps, cfg, sampler_name, scheduler, denoise, runs,
|
|
pipe=None, sigmas=None, custom_sigmas="", step_cache="off", cache_threshold=0.2, cache_start=0.15, cache_end=0.95,
|
|
start_at_step=0, end_at_step=10000, add_noise=True, preview="default"):
|
|
family = self._family(model, pipe)
|
|
if step_cache == "auto":
|
|
step_cache = "easycache" if steps >= AUTO_CACHE_MIN_STEPS else "off"
|
|
if step_cache != "off":
|
|
model = CACHES[step_cache].execute(model, cache_threshold, cache_start, cache_end, False).args[0]
|
|
if scheduler == AUTO:
|
|
scheduler = "simple"
|
|
|
|
samples = comfy.sample.fix_empty_latent_channels(model, latent_image["samples"], latent_image.get("downscale_ratio_spacial"), latent_image.get("downscale_ratio_temporal"))
|
|
used = self._sigmas(model, sampler_name, scheduler, steps, denoise, sigmas, custom_sigmas)
|
|
if runs > 1 and not isinstance(samples, torch.Tensor):
|
|
raise ValueError("[PromptForge] runs > 1 needs a plain latent; this model uses a combined audio+video latent.")
|
|
|
|
device = comfy.model_management.get_torch_device()
|
|
if device.type == "cuda":
|
|
torch.cuda.reset_peak_memory_stats(device)
|
|
if preview != "default":
|
|
latent_preview.set_preview_method(preview)
|
|
try:
|
|
try:
|
|
callback = latent_preview.prepare_callback(model, steps)
|
|
except AttributeError:
|
|
# the tiny decoder file is not a loadable VAE; the colour preview needs no file
|
|
print("[PromptForge] the tiny preview decoder could not be loaded; using latent2rgb previews.")
|
|
latent_preview.set_preview_method("latent2rgb")
|
|
callback = latent_preview.prepare_callback(model, steps)
|
|
progress = comfy.utils.PROGRESS_BAR_ENABLED
|
|
results = []
|
|
for index in range(runs):
|
|
current = seed + index
|
|
noise = comfy.sample.prepare_empty_noise(samples) if not add_noise else comfy.sample.prepare_noise(samples, current, latent_image.get("batch_index"))
|
|
results.append(comfy.sample.sample(
|
|
model, noise, steps, cfg, sampler_name, scheduler, positive, negative, samples,
|
|
denoise=denoise, disable_noise=not add_noise, start_step=start_at_step or None,
|
|
last_step=end_at_step if end_at_step < steps else None, force_full_denoise=end_at_step >= steps,
|
|
noise_mask=latent_image.get("noise_mask"), sigmas=used, callback=callback, disable_pbar=not progress, seed=current))
|
|
finally:
|
|
if preview != "default":
|
|
latent_preview.set_preview_method()
|
|
merged = results[0] if runs == 1 else torch.cat(results)
|
|
|
|
out = latent_image.copy()
|
|
out.pop("downscale_ratio_spacial", None)
|
|
out.pop("downscale_ratio_temporal", None)
|
|
out["samples"] = merged
|
|
schedule = "custom sigmas" if sigmas is not None or custom_sigmas.strip() else scheduler
|
|
info = f"{sampler_name} / {schedule}, {steps} steps, cfg {cfg:g}, denoise {denoise:g}, seed {seed}" + (f" x{runs}" if runs > 1 else "") + (f" ({family})" if family else "") + (f", {step_cache} {cache_threshold:g}" if step_cache != "off" else "")
|
|
if device.type == "cuda":
|
|
info += f", peak VRAM {torch.cuda.max_memory_allocated(device) / 1024 ** 3:.1f} GB"
|
|
return (out, used, info)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"PROMPTFORGE_Sampler": PROMPTFORGE_Sampler}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"PROMPTFORGE_Sampler": "PromptForge Sampler"}
|