1784 lines
77 KiB
Python
1784 lines
77 KiB
Python
import numpy as np
|
||
from PIL import Image, ImageOps, ImageDraw, ImageFilter, ImageEnhance
|
||
from PIL.PngImagePlugin import PngInfo
|
||
import torch
|
||
import torch.nn.functional as F
|
||
import json
|
||
import random
|
||
import os
|
||
import re
|
||
import urllib.request, pathlib, shutil, tempfile
|
||
from typing import Dict, Any, List
|
||
import importlib.util
|
||
import sys
|
||
|
||
|
||
# Use importlib to probe for the optional global `config` module without triggering static
|
||
# import-time linter errors when it does not exist in the current environment.
|
||
_cfg_spec = importlib.util.find_spec("config")
|
||
if _cfg_spec is not None:
|
||
_cfg = importlib.util.module_from_spec(_cfg_spec) # type: ignore[var-annotated]
|
||
_cfg_spec.loader.exec_module(_cfg) # type: ignore[arg-type]
|
||
CONTEXT_THUMBNAIL_SIZE = getattr(_cfg, "CONTEXT_THUMBNAIL_SIZE", (64, 64))
|
||
else:
|
||
CONTEXT_THUMBNAIL_SIZE = (64, 64)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Remote model download helper
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# Base URL for hosting – we append <family>/<version>/<filename>
|
||
_REMOTE_MODEL_BASE_URL = "https://raw.githubusercontent.com/SuperBeastsAI/SuperBeastsAI-Models/main/"
|
||
# License info for downloaded model weights
|
||
_SPCA_LICENSE_ID = "SPCA-Community-NoSaaS"
|
||
_SPCA_LICENSE_URL = (
|
||
"https://github.com/SuperBeastsAI/SuperBeastsAI-Models/"
|
||
"blob/main/SuperPopColorAdjustment/LICENSE.txt"
|
||
)
|
||
# File written into models dir after first download so we don't spam the console
|
||
_SPCA_LICENSE_STUB_FILENAME = "SPCA_LICENSE_STUB.txt"
|
||
# Env var: set to "1" (or any truthy string) to suppress the banner (CI / headless automation)
|
||
_SUPERBEASTS_SILENCE_LICENSE_ENV = "SUPERBEASTS_SILENCE_LICENSE"
|
||
|
||
|
||
# Registry of published model families → version → filename
|
||
_MODEL_REGISTRY: dict[str, dict[str, str]] = {
|
||
"SuperPopColorAdjustment": {
|
||
# Default to safetensors weights. Companion TorchScript (.pt) and
|
||
"v1.0": "SuperBeasts_ColorAdjustment_512px_V1.onnx",
|
||
"v2.0": "SuperBeasts_ColorAdjustment_512px_V2.safetensors",
|
||
},
|
||
# Future models / versions can be added here
|
||
}
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Known SHA-256 digests for published weight files. This allows us to verify
|
||
# that a downloaded (or pre-existing) file has not been tampered with. When
|
||
# a file's entry is missing we skip verification so development models can be
|
||
# iterated quickly – but **release** models **should always** have a digest.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_MODEL_SHA256: dict[str, str] = {
|
||
"SuperBeasts_ColorAdjustment_512px_V1.onnx": "a813fbb091dc7849f460f17e54c70cc6e6ef2984589a0ae94ab308d2957dc4ad",
|
||
"SuperBeasts_ColorAdjustment_512px_V2.safetensors": "1a7e0aaa54d565b827c619a452bdeb04fd12975178b39560b92fe34cb839a065",
|
||
}
|
||
|
||
# Helper – compute sha256 digest of a file (hex string)
|
||
def _sha256_of_file(path: str, chunk: int = 1 << 20) -> str:
|
||
import hashlib
|
||
h = hashlib.sha256()
|
||
with open(path, "rb") as f:
|
||
while True:
|
||
blk = f.read(chunk)
|
||
if not blk:
|
||
break
|
||
h.update(blk)
|
||
return h.hexdigest()
|
||
|
||
# Verify checksum if we have one for this filename.
|
||
def _verify_model_checksum(path: str):
|
||
fn = os.path.basename(path)
|
||
expected = _MODEL_SHA256.get(fn)
|
||
if not expected:
|
||
# No checksum known – skip (development weight)
|
||
return
|
||
actual = _sha256_of_file(path)
|
||
if actual != expected:
|
||
raise RuntimeError(
|
||
f"Checksum mismatch for '{fn}'.\nExpected: {expected}\nActual: {actual}\n"
|
||
"The file appears to be corrupted or has been tampered with. Delete it and rerun to re-download, "
|
||
"or obtain a clean copy from the official model repository."
|
||
)
|
||
|
||
|
||
def _latest_version(family: str) -> str:
|
||
versions = list(_MODEL_REGISTRY.get(family, {}).keys())
|
||
if not versions:
|
||
raise ValueError(f"Unknown model family: {family}")
|
||
# assume semantic ‘vX.Y' strings – sort lexically
|
||
return sorted(versions)[-1]
|
||
|
||
|
||
def _download_remote_model(rel_path: str, dest_path: str):
|
||
"""Download file at *rel_path* (relative to base URL) to *dest_path*.
|
||
Emits a one-time license banner and writes a stub file so we don't repeat warnings.
|
||
"""
|
||
url = _REMOTE_MODEL_BASE_URL + rel_path
|
||
print(f"[SuperBeasts] Downloading model weights from {url} …")
|
||
|
||
tmp_dir = tempfile.mkdtemp()
|
||
tmp_file = os.path.join(tmp_dir, rel_path.split('/')[-1] + ".part")
|
||
|
||
try:
|
||
# --- download file ---
|
||
with urllib.request.urlopen(url) as r, open(tmp_file, "wb") as f:
|
||
shutil.copyfileobj(r, f)
|
||
|
||
# --- move into place ---
|
||
os.makedirs(os.path.dirname(dest_path), exist_ok=True)
|
||
shutil.move(tmp_file, dest_path)
|
||
print(f"[SuperBeasts] Saved weights to {dest_path}")
|
||
|
||
# --- after successful download, print license notice (once per repo install) ---
|
||
_maybe_print_spca_license_notice(os.path.dirname(dest_path))
|
||
|
||
finally:
|
||
shutil.rmtree(tmp_dir, ignore_errors=True)
|
||
|
||
|
||
# Import helper to add extra channels same way training did
|
||
try:
|
||
from .extras.inference.inference_utils import add_extra_color_channels_np as _add_extra_color_channels_np # type: ignore
|
||
except Exception:
|
||
_add_extra_color_channels_np = None # will define below
|
||
|
||
# ------------------------------------------------------------------
|
||
# Ensure we always add the extra YCbCr (+3) and Saturation (+1) channels so
|
||
# the input tensor matches the 7-channel format expected by the SuperBeasts
|
||
# models. The helper below mirrors the logic from inference_utils but forces
|
||
# the extras on regardless of global config flags.
|
||
# ------------------------------------------------------------------
|
||
|
||
try:
|
||
import cv2 # Optional – used for fast RGB→HSV conversion
|
||
_HAS_CV2 = True
|
||
except ImportError: # linter environments might not have it
|
||
_HAS_CV2 = False
|
||
|
||
def _sb_add_extra_channels(arr: np.ndarray) -> np.ndarray:
|
||
"""Return RGB array with YCbCr and saturation channels appended (total 7)."""
|
||
if arr.ndim != 3 or arr.shape[2] != 3:
|
||
# Already has extra channels or malformed – return unchanged
|
||
return arr
|
||
|
||
r = arr[:, :, 0]
|
||
g = arr[:, :, 1]
|
||
b = arr[:, :, 2]
|
||
|
||
# YCbCr as in training
|
||
y = 0.299 * r + 0.587 * g + 0.114 * b
|
||
cb = -0.168736 * r - 0.331264 * g + 0.5 * b + 0.5
|
||
cr = 0.5 * r - 0.418688 * g - 0.081312 * b + 0.5
|
||
|
||
ycbcr = np.stack([y, cb, cr], axis=-1).astype(arr.dtype)
|
||
|
||
# Saturation channel – use OpenCV if available, otherwise approximate
|
||
if _HAS_CV2:
|
||
hsv = cv2.cvtColor((arr * 255.0).astype(np.uint8), cv2.COLOR_RGB2HSV)
|
||
sat = (hsv[:, :, 1] / 255.0).astype(arr.dtype)
|
||
else:
|
||
# Simple saturation approximation without HSV conversion
|
||
max_rgb = arr.max(axis=2)
|
||
min_rgb = arr.min(axis=2)
|
||
sat = np.where(max_rgb == 0, 0.0, (max_rgb - min_rgb) / max_rgb).astype(arr.dtype)
|
||
sat = sat[:, :, np.newaxis]
|
||
|
||
return np.concatenate([arr, ycbcr, sat], axis=-1).astype(arr.dtype)
|
||
|
||
# If we imported the original util and it already produces ≥7 channels we keep it.
|
||
if _add_extra_color_channels_np is None:
|
||
_add_extra_color_channels_np = _sb_add_extra_channels
|
||
else:
|
||
# Wrap the existing util so we guarantee at least 7 channels.
|
||
_orig_fn = _add_extra_color_channels_np
|
||
|
||
def _add_extra_color_channels_np(arr): # type: ignore[override]
|
||
out = _orig_fn(arr)
|
||
if out.shape[2] < 7:
|
||
out = _sb_add_extra_channels(out)
|
||
return out
|
||
|
||
# ---------------------------------------------
|
||
# SB MODEL LOADER NODE
|
||
# ---------------------------------------------
|
||
|
||
def _maybe_print_spca_license_notice(models_dir: str):
|
||
"""
|
||
Print the SPCA license banner once per install.
|
||
We create a stub file in the models_dir so we don't repeat on future runs.
|
||
Suppress by setting env var SUPERBEASTS_SILENCE_LICENSE=1.
|
||
"""
|
||
if os.environ.get(_SUPERBEASTS_SILENCE_LICENSE_ENV, "").strip():
|
||
return
|
||
|
||
stub_path = os.path.join(models_dir, _SPCA_LICENSE_STUB_FILENAME)
|
||
if os.path.exists(stub_path):
|
||
# Already shown
|
||
return
|
||
|
||
banner = (
|
||
"\n"
|
||
"============================================================\n"
|
||
" Super Pop Color Adjustment – Model License Notice\n"
|
||
"------------------------------------------------------------\n"
|
||
f" Weights licensed under {_SPCA_LICENSE_ID}.\n"
|
||
" Personal / client / local use OK.\n"
|
||
" Public SaaS/API redistribution requires a license.\n"
|
||
f" Full terms: {_SPCA_LICENSE_URL}\n"
|
||
"============================================================\n"
|
||
)
|
||
print(banner, file=sys.stderr)
|
||
|
||
# write stub file so we don't print again
|
||
try:
|
||
with open(stub_path, "w", encoding="utf-8") as f:
|
||
f.write(
|
||
"Super Pop Color Adjustment weights downloaded.\n"
|
||
"No public SaaS/API redistribution without license. Contact: Via DM on Instagram @SuperBeasts.AI"
|
||
f"License: {_SPCA_LICENSE_ID}\n"
|
||
f"See: {_SPCA_LICENSE_URL}\n"
|
||
)
|
||
except Exception as e:
|
||
# non-fatal; just warn
|
||
print(f"[SuperBeasts] WARNING: could not write license stub: {e}", file=sys.stderr)
|
||
|
||
|
||
def _discover_sb_models(models_dir: str):
|
||
"""Return a list of available model filenames (searches recursively)."""
|
||
if not os.path.isdir(models_dir):
|
||
return []
|
||
names = []
|
||
for root, _dirs, files in os.walk(models_dir):
|
||
for file in files:
|
||
if file.lower().endswith(('.keras', '.h5', '.safetensors', '.onnx', '.pt')):
|
||
names.append(file)
|
||
return sorted(set(names))
|
||
|
||
|
||
class SBLoadModel:
|
||
"""Load a SuperBeasts colour-adjustment model (downloads if missing)./n
|
||
|
||
Select a model from the dropdown. Use Family/Version entries (e.g., SuperPopColorAdjustment/latest)
|
||
to auto-download from the official model registry when not found locally./n
|
||
|
||
⚠ License: Downloaded weights are licensed under SPCA-Community-NoSaaS.
|
||
Local / personal / commercial use OK. SaaS/API redistribution using this model requires a license.
|
||
See: https://github.com/SuperBeastsAI/SuperBeastsAI-Models/tree/main/SuperPopColorAdjustment
|
||
"""
|
||
# Shown in ComfyUI when the user clicks the "?" icon on the node title-bar.
|
||
DESCRIPTION = __doc__
|
||
|
||
_model_cache: Dict[str, Dict[str, Any]] = {}
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
models_dir = os.path.join(os.path.dirname(__file__), 'models')
|
||
available = _discover_sb_models(models_dir)
|
||
|
||
dropdown_set = set(available)
|
||
# Add registry family/version keys so saved workflows remain valid
|
||
for fam, vers_dict in _MODEL_REGISTRY.items():
|
||
dropdown_set.add(f"{fam}/latest")
|
||
for ver in vers_dict.keys():
|
||
dropdown_set.add(f"{fam}/{ver}")
|
||
|
||
dropdown = sorted(dropdown_set)
|
||
return {
|
||
"required": {
|
||
"model_key": (dropdown, ),
|
||
"device": (["AUTO", "CPU", "GPU"], {"default": "AUTO"}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("SBMODEL", )
|
||
RETURN_NAMES = ("model", )
|
||
FUNCTION = "load"
|
||
CATEGORY = "SuperBeastsAI/Model"
|
||
|
||
def _parse_patch_size(self, filename: str, default: int = 512):
|
||
m = re.search(r"(\d+)px", filename)
|
||
if m:
|
||
try:
|
||
return int(m.group(1))
|
||
except ValueError:
|
||
pass
|
||
return default
|
||
|
||
def load(self, model_key: str, device: str = "AUTO"):
|
||
"""`model_key` can be either a raw filename (legacy) *or* the form
|
||
"Family/Version" (e.g. "SuperPopColorAdjustment/v1.0" or "SuperPopColorAdjustment/latest")."""
|
||
|
||
# ------------------------------------------------------------------
|
||
# Resolve model filename, family, version
|
||
# ------------------------------------------------------------------
|
||
if "/" in model_key:
|
||
family, version = model_key.split("/", 1)
|
||
if version.lower() == "latest":
|
||
version = _latest_version(family)
|
||
try:
|
||
filename = _MODEL_REGISTRY[family][version]
|
||
except KeyError:
|
||
raise ValueError(f"Unknown model/version combination: {model_key}")
|
||
else:
|
||
# Legacy direct filename path (assume SuperPop)…
|
||
family = "SuperPopColorAdjustment"
|
||
version = _latest_version(family)
|
||
filename = model_key
|
||
|
||
# Resolve model storage directories --------------------------------
|
||
models_root = os.path.join(os.path.dirname(__file__), 'models')
|
||
version_dir = os.path.join(models_root, version)
|
||
os.makedirs(version_dir, exist_ok=True)
|
||
|
||
# Look for existing file in version dir first, then root (for legacy installs)
|
||
candidate_dirs = [version_dir, models_root]
|
||
model_path = None # type: ignore[assignment]
|
||
for _dir in candidate_dirs:
|
||
cand_path = os.path.join(_dir, filename)
|
||
if os.path.isfile(cand_path):
|
||
model_path = cand_path
|
||
models_dir = _dir # use this dir for subsequent operations
|
||
break
|
||
|
||
if model_path is None:
|
||
# Default to version_dir for new downloads
|
||
models_dir = version_dir
|
||
model_path = os.path.join(models_dir, filename)
|
||
|
||
# --------------------------------------------------------------
|
||
# Prioritise *.safetensors → *.pt → *.onnx in that order. If the
|
||
# requested file is missing locally we look for any preferred
|
||
# alternative that already exists. Otherwise we attempt to
|
||
# download the best-quality format available from the official
|
||
# model registry, trying each extension until one succeeds.
|
||
# --------------------------------------------------------------
|
||
if not os.path.isfile(model_path):
|
||
base_name, _orig_ext = os.path.splitext(filename)
|
||
|
||
# Preferred search / download order
|
||
_EXT_PRIORITY = ['.safetensors', '.pt', '.onnx']
|
||
|
||
# Build ordered unique candidate list starting with preferred
|
||
cand_filenames = [base_name + ext for ext in _EXT_PRIORITY]
|
||
if filename not in cand_filenames:
|
||
# Include the originally requested filename at front if it
|
||
# does not conform to the priority list (for backwards compatibility)
|
||
cand_filenames.insert(0, filename)
|
||
|
||
# 1) Check for any existing local file following priority
|
||
model_found = False
|
||
for cand in cand_filenames:
|
||
cand_path = os.path.join(models_dir, cand)
|
||
if os.path.isfile(cand_path):
|
||
filename = cand
|
||
model_path = cand_path
|
||
model_found = True
|
||
break
|
||
|
||
# 2) Attempt remote download for each preferred extension
|
||
if not model_found:
|
||
last_exc = None
|
||
for cand in cand_filenames:
|
||
rel_path = f"{family}/{version}/{cand}"
|
||
cand_path = os.path.join(models_dir, cand)
|
||
try:
|
||
_download_remote_model(rel_path, cand_path)
|
||
filename = cand
|
||
model_path = cand_path
|
||
model_found = True
|
||
break
|
||
except Exception as e:
|
||
last_exc = e # keep last for context
|
||
continue
|
||
|
||
if not model_found:
|
||
raise FileNotFoundError(
|
||
f"Model file not found locally and automatic download failed for any supported extension (.safetensors/.pt/.onnx).\nLast error: {last_exc}"
|
||
)
|
||
else:
|
||
# File exists locally; nothing to do.
|
||
pass
|
||
|
||
# --------------------------------------------------------------
|
||
# Ensure both *.safetensors and companion *.pt files are present
|
||
# (useful for users who switch between Torch and TorchScript).
|
||
# --------------------------------------------------------------
|
||
base_name, _ = os.path.splitext(filename)
|
||
for _ext in ('.safetensors', '.pt'):
|
||
_comp_path = os.path.join(models_dir, base_name + _ext)
|
||
if not os.path.isfile(_comp_path):
|
||
try:
|
||
_download_remote_model(f"{family}/{version}/{base_name + _ext}", _comp_path)
|
||
except Exception:
|
||
# It's OK if one of the companion formats is missing from the
|
||
# remote registry – we just skip silently.
|
||
pass
|
||
|
||
# --------------------------------------------------------------
|
||
# Attempt to download optional metadata JSON (licence, provenance).
|
||
# The file is named <base>_MODEL_INFO.json in the remote repo.
|
||
# --------------------------------------------------------------
|
||
_meta_filename = f"{base_name}_MODEL_INFO.json"
|
||
_meta_path = os.path.join(models_dir, _meta_filename)
|
||
if not os.path.isfile(_meta_path):
|
||
try:
|
||
# Remote repo stores a generic MODEL_INFO.json per version.
|
||
_download_remote_model(f"{family}/{version}/MODEL_INFO.json", _meta_path)
|
||
except Exception:
|
||
# Non-critical: metadata missing – warn once.
|
||
print(f"[SuperBeasts] Warning: metadata file not found for {_meta_filename}")
|
||
|
||
# --- Security: check SHA-256 digest (raises if mismatch) ---
|
||
try:
|
||
_verify_model_checksum(model_path)
|
||
except Exception as _ck_err:
|
||
# Surface early so users see a clear error in ComfyUI
|
||
raise
|
||
|
||
# Use cache if possible
|
||
if model_path in self._model_cache:
|
||
return (self._model_cache[model_path], )
|
||
|
||
# We defer importing onnxruntime unless we end up loading an ONNX model
|
||
ort = None # type: ignore[assignment]
|
||
|
||
# Optionally set device context
|
||
dev_ctx = None
|
||
if device == "GPU":
|
||
dev_ctx = "/GPU:0"
|
||
elif device == "CPU":
|
||
dev_ctx = "/CPU:0"
|
||
|
||
# Determine patch size from filename
|
||
patch_size = self._parse_patch_size(filename)
|
||
|
||
# --------------------------------------------------------------
|
||
# 1. TorchScript checkpoint – no external architecture required
|
||
# --------------------------------------------------------------
|
||
if model_path.lower().endswith('.pt'):
|
||
import torch
|
||
|
||
device_str = 'cuda' if device == 'GPU' and torch.cuda.is_available() else 'cpu'
|
||
|
||
try:
|
||
model_pt = torch.jit.load(model_path, map_location=device_str)
|
||
except Exception as e:
|
||
raise RuntimeError(f"Failed to load TorchScript model: {e}")
|
||
|
||
model_pt.eval()
|
||
|
||
model_info = {
|
||
"model": model_pt,
|
||
"device": device_str,
|
||
"patch_size": patch_size,
|
||
"type": "torchscript",
|
||
}
|
||
|
||
# --------------------------------------------------------------
|
||
# 2. Safetensors + Python architecture (requires contextual_residual_unet_v4.py)
|
||
# --------------------------------------------------------------
|
||
elif model_path.lower().endswith('.safetensors'):
|
||
import torch
|
||
from safetensors.torch import load_file as safe_load
|
||
import importlib.util, sys, pathlib
|
||
model_dir = pathlib.Path(__file__).parent / 'torch_models'
|
||
model_dir.mkdir(exist_ok=True)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Ensure architecture implementation file is present – download on
|
||
# demand if missing so repo stays lean for users who only use ONNX.
|
||
# ------------------------------------------------------------------
|
||
arch_filename = 'contextual_residual_unet_v4_pt.py'
|
||
module_path = (model_dir / arch_filename).resolve()
|
||
if not module_path.exists():
|
||
try:
|
||
rel_path_arch = f"{family}/TorchModel/{arch_filename}"
|
||
_download_remote_model(rel_path_arch, str(module_path))
|
||
except Exception as e:
|
||
raise RuntimeError(
|
||
"Required architecture file was not found locally and automatic download failed.\n" + str(e)
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Safety guard – ensure we never import a file outside our plugin dir
|
||
# even if the relative path is somehow manipulated.
|
||
# ------------------------------------------------------------------
|
||
|
||
allowed_root = model_dir.resolve()
|
||
if not str(module_path).startswith(str(allowed_root)):
|
||
raise RuntimeError("Blocked unsafe module import: path escapes plugin directory")
|
||
|
||
sys.path.append(str(model_dir))
|
||
spec = importlib.util.spec_from_file_location('contextual_residual_unet_v4_pt', module_path)
|
||
module = importlib.util.module_from_spec(spec)
|
||
spec.loader.exec_module(module) # type: ignore
|
||
ContextualResidualUNetV4Torch = module.ContextualResidualUNetV4Torch
|
||
|
||
device_str = 'cuda' if device == 'GPU' and torch.cuda.is_available() else 'cpu'
|
||
|
||
# Prefer TorchScript if it exists next to the safetensors
|
||
base_no_ext = os.path.splitext(model_path)[0]
|
||
pt_path_candidates = [base_no_ext + '.pt', base_no_ext + '_scripted.pt']
|
||
pt_path = next((p for p in pt_path_candidates if os.path.isfile(p)), None)
|
||
if pt_path and os.path.isfile(pt_path):
|
||
try:
|
||
model_pt = torch.jit.load(pt_path, map_location=device_str)
|
||
model_pt.eval()
|
||
except Exception as e:
|
||
raise RuntimeError(f"Failed to load TorchScript companion file {pt_path}: {e}")
|
||
else:
|
||
# Fall back to python architecture + direct state_dict load (legacy path)
|
||
model_pt = ContextualResidualUNetV4Torch()
|
||
|
||
state_in = safe_load(model_path, device='cpu')
|
||
try:
|
||
model_pt.load_state_dict(state_in, strict=True)
|
||
except Exception as e:
|
||
raise RuntimeError("Checkpoint did not match architecture and no .pt file found. "
|
||
"Convert the model to TorchScript to avoid mismatches.\n" + str(e))
|
||
|
||
model_pt.eval()
|
||
|
||
model_info = {
|
||
"model": model_pt,
|
||
"device": device_str,
|
||
"patch_size": patch_size,
|
||
"type": "torch",
|
||
}
|
||
|
||
# Load ONNX with onnxruntime (import lazily so safetensors users don't need it)
|
||
else:
|
||
if ort is None:
|
||
try:
|
||
import onnxruntime as ort # type: ignore
|
||
except Exception as e:
|
||
raise RuntimeError(
|
||
"onnxruntime is required for SuperBeasts ONNX models. Install with:\n"
|
||
" pip install onnxruntime-gpu # (or onnxruntime for CPU)\n\n" + str(e)
|
||
)
|
||
|
||
providers: List[str]
|
||
avail_prov = ort.get_available_providers()
|
||
print(f"[SuperBeasts] ONNX Runtime providers available: {avail_prov}")
|
||
cuda_available = 'CUDAExecutionProvider' in avail_prov
|
||
dml_available = 'DmlExecutionProvider' in avail_prov
|
||
|
||
chosen_provider = 'CPUExecutionProvider'
|
||
|
||
if device == "GPU":
|
||
if cuda_available:
|
||
chosen_provider = 'CUDAExecutionProvider'
|
||
elif dml_available:
|
||
chosen_provider = 'DmlExecutionProvider'
|
||
elif device == "AUTO":
|
||
if cuda_available:
|
||
chosen_provider = 'CUDAExecutionProvider'
|
||
elif dml_available:
|
||
chosen_provider = 'DmlExecutionProvider'
|
||
|
||
if chosen_provider != 'CPUExecutionProvider':
|
||
providers = [chosen_provider, 'CPUExecutionProvider']
|
||
else:
|
||
providers = ['CPUExecutionProvider']
|
||
|
||
print(f"[SuperBeasts] Using ORT providers: {providers}")
|
||
|
||
# Create session
|
||
session_opts = ort.SessionOptions()
|
||
session_opts.log_severity_level = 3 # suppress info logs
|
||
session = ort.InferenceSession(model_path, session_options=session_opts, providers=providers)
|
||
|
||
input_names = [i.name for i in session.get_inputs()]
|
||
|
||
model_info = {
|
||
"session": session,
|
||
"input_names": input_names,
|
||
"patch_size": patch_size,
|
||
"path": model_path,
|
||
"type": "onnx",
|
||
}
|
||
|
||
# Attach metadata path if we fetched it earlier
|
||
if '_meta_path' in locals() and os.path.isfile(_meta_path):
|
||
model_info["metadata_path"] = _meta_path
|
||
|
||
# Cache and return
|
||
self._model_cache[model_path] = model_info
|
||
return (model_info, )
|
||
|
||
|
||
# Tensor to PIL
|
||
def tensor2pil(image):
|
||
arr = image.detach().cpu().numpy()
|
||
if arr.ndim == 4:
|
||
arr = arr[0]
|
||
if arr.ndim == 2:
|
||
return Image.fromarray(np.clip(arr * 255.0, 0, 255).astype(np.uint8), mode='L')
|
||
if arr.ndim == 3 and arr.shape[-1] == 1:
|
||
return Image.fromarray(np.clip(arr[..., 0] * 255.0, 0, 255).astype(np.uint8), mode='L')
|
||
if arr.ndim != 3:
|
||
raise ValueError(f"Expected image tensor with 2 or 3 dims, got shape {arr.shape}")
|
||
return Image.fromarray(np.clip(arr * 255.0, 0, 255).astype(np.uint8))
|
||
|
||
# PIL to Tensor
|
||
def pil2tensor(image):
|
||
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
||
|
||
|
||
_LAB_EPSILON = 216.0 / 24389.0
|
||
_LAB_KAPPA = 24389.0 / 27.0
|
||
_D65_WHITE = np.array([0.95047, 1.0, 1.08883], dtype=np.float32)
|
||
_SRGB_TO_XYZ = np.array([
|
||
[0.4124564, 0.3575761, 0.1804375],
|
||
[0.2126729, 0.7151522, 0.0721750],
|
||
[0.0193339, 0.1191920, 0.9503041],
|
||
], dtype=np.float32)
|
||
_XYZ_TO_SRGB = np.array([
|
||
[ 3.2404542, -1.5371385, -0.4985314],
|
||
[-0.9692660, 1.8760108, 0.0415560],
|
||
[ 0.0556434, -0.2040259, 1.0572252],
|
||
], dtype=np.float32)
|
||
|
||
|
||
def _srgb_to_linear(rgb):
|
||
rgb = np.asarray(rgb, dtype=np.float32)
|
||
return np.where(rgb <= 0.04045, rgb / 12.92, ((rgb + 0.055) / 1.055) ** 2.4)
|
||
|
||
|
||
def _linear_to_srgb(rgb):
|
||
rgb = np.asarray(rgb, dtype=np.float32)
|
||
return np.where(rgb <= 0.0031308, 12.92 * rgb, 1.055 * np.power(np.clip(rgb, 0.0, None), 1.0 / 2.4) - 0.055)
|
||
|
||
|
||
def _f_xyz_to_lab(t):
|
||
return np.where(t > _LAB_EPSILON, np.cbrt(t), (_LAB_KAPPA * t + 16.0) / 116.0)
|
||
|
||
|
||
def _f_lab_to_xyz(t):
|
||
t3 = t ** 3
|
||
return np.where(t3 > _LAB_EPSILON, t3, (116.0 * t - 16.0) / _LAB_KAPPA)
|
||
|
||
|
||
def _rgb_to_lab_components(img: Image.Image):
|
||
rgb_u8 = np.ascontiguousarray(np.asarray(img.convert('RGB'), dtype=np.uint8))
|
||
if rgb_u8.ndim != 3 or rgb_u8.shape[2] != 3:
|
||
raise ValueError(f"Expected RGB image array with shape HxWx3, got {rgb_u8.shape}")
|
||
|
||
height, width = rgb_u8.shape[:2]
|
||
l_uint8 = np.empty((height, width), dtype=np.uint8)
|
||
a_out = np.empty((height, width), dtype=np.float32)
|
||
b_out = np.empty((height, width), dtype=np.float32)
|
||
|
||
max_chunk_pixels = 1_048_576
|
||
rows_per_chunk = max(1, max_chunk_pixels // max(1, width))
|
||
rgb_to_xyz = _SRGB_TO_XYZ
|
||
white_x = _D65_WHITE[0]
|
||
white_y = _D65_WHITE[1]
|
||
white_z = _D65_WHITE[2]
|
||
|
||
for y0 in range(0, height, rows_per_chunk):
|
||
y1 = min(height, y0 + rows_per_chunk)
|
||
rgb = rgb_u8[y0:y1].astype(np.float32) / 255.0
|
||
linear_rgb = _srgb_to_linear(rgb)
|
||
|
||
r = linear_rgb[..., 0]
|
||
g = linear_rgb[..., 1]
|
||
bl = linear_rgb[..., 2]
|
||
|
||
x = (rgb_to_xyz[0, 0] * r + rgb_to_xyz[0, 1] * g + rgb_to_xyz[0, 2] * bl) / white_x
|
||
y = (rgb_to_xyz[1, 0] * r + rgb_to_xyz[1, 1] * g + rgb_to_xyz[1, 2] * bl) / white_y
|
||
z = (rgb_to_xyz[2, 0] * r + rgb_to_xyz[2, 1] * g + rgb_to_xyz[2, 2] * bl) / white_z
|
||
|
||
fx = _f_xyz_to_lab(x)
|
||
fy = _f_xyz_to_lab(y)
|
||
fz = _f_xyz_to_lab(z)
|
||
|
||
l = (np.float32(116.0) * fy) - np.float32(16.0)
|
||
a = np.float32(500.0) * (fx - fy)
|
||
b = np.float32(200.0) * (fy - fz)
|
||
|
||
l_uint8[y0:y1] = np.clip(l * (255.0 / 100.0), 0, 255).astype(np.uint8)
|
||
a_out[y0:y1] = a.astype(np.float32, copy=False)
|
||
b_out[y0:y1] = b.astype(np.float32, copy=False)
|
||
|
||
return l_uint8, a_out, b_out
|
||
|
||
|
||
def _linear_srgb_channel_to_u8(channel: np.ndarray) -> np.ndarray:
|
||
"""Convert one linear-light sRGB channel to uint8 sRGB.
|
||
|
||
Kept channel-wise so HDR Effects does not allocate several full-frame
|
||
HxWx3 float32 temporaries while converting LAB back to RGB. Some NumPy
|
||
builds have crashed in those native ufunc/matmul paths under ComfyUI's
|
||
large-image memory pressure instead of raising MemoryError.
|
||
"""
|
||
srgb = np.clip(_linear_to_srgb(channel), 0.0, 1.0)
|
||
return (srgb * 255.0 + 0.5).astype(np.uint8)
|
||
|
||
|
||
def _lab_components_to_rgb(l_uint8: np.ndarray, a: np.ndarray, b: np.ndarray) -> Image.Image:
|
||
l_uint8 = np.ascontiguousarray(l_uint8)
|
||
a = np.ascontiguousarray(a, dtype=np.float32)
|
||
b = np.ascontiguousarray(b, dtype=np.float32)
|
||
|
||
if l_uint8.ndim != 2:
|
||
raise ValueError(f"Expected 2D LAB luminance channel, got shape {l_uint8.shape}")
|
||
if a.shape != l_uint8.shape or b.shape != l_uint8.shape:
|
||
raise ValueError(
|
||
"LAB channel shape mismatch: "
|
||
f"L={l_uint8.shape}, a={a.shape}, b={b.shape}"
|
||
)
|
||
|
||
height, width = l_uint8.shape
|
||
rgb_u8 = np.empty((height, width, 3), dtype=np.uint8)
|
||
|
||
# Process rows in bounded chunks. This avoids the previous full-frame LAB,
|
||
# XYZ, linear-RGB and matrix-multiply temporaries while preserving the same
|
||
# D65/LAB math. Keep the chunk size large enough to avoid Python-loop cost
|
||
# on normal images but small enough for 4K/8K frames and batched Comfy runs.
|
||
max_chunk_pixels = 1_048_576
|
||
rows_per_chunk = max(1, max_chunk_pixels // max(1, width))
|
||
xyz_to_rgb = _XYZ_TO_SRGB
|
||
white_x = _D65_WHITE[0]
|
||
white_y = _D65_WHITE[1]
|
||
white_z = _D65_WHITE[2]
|
||
|
||
for y0 in range(0, height, rows_per_chunk):
|
||
y1 = min(height, y0 + rows_per_chunk)
|
||
|
||
l = l_uint8[y0:y1].astype(np.float32, copy=False) * (100.0 / 255.0)
|
||
fy = (l + 16.0) / 116.0
|
||
fx = a[y0:y1] / 500.0 + fy
|
||
fz = fy - b[y0:y1] / 200.0
|
||
|
||
x = _f_lab_to_xyz(fx) * white_x
|
||
y = _f_lab_to_xyz(fy) * white_y
|
||
z = _f_lab_to_xyz(fz) * white_z
|
||
|
||
# Inline the XYZ->linear sRGB matrix multiply channel-wise. This keeps
|
||
# execution inside simple 2D ufuncs instead of NumPy's 3D matmul path.
|
||
r = xyz_to_rgb[0, 0] * x + xyz_to_rgb[0, 1] * y + xyz_to_rgb[0, 2] * z
|
||
g = xyz_to_rgb[1, 0] * x + xyz_to_rgb[1, 1] * y + xyz_to_rgb[1, 2] * z
|
||
bl = xyz_to_rgb[2, 0] * x + xyz_to_rgb[2, 1] * y + xyz_to_rgb[2, 2] * z
|
||
|
||
rgb_u8[y0:y1, :, 0] = _linear_srgb_channel_to_u8(r)
|
||
rgb_u8[y0:y1, :, 1] = _linear_srgb_channel_to_u8(g)
|
||
rgb_u8[y0:y1, :, 2] = _linear_srgb_channel_to_u8(bl)
|
||
|
||
return Image.fromarray(rgb_u8, mode='RGB')
|
||
|
||
|
||
def adjust_shadows(luminance_array, shadow_intensity, hdr_intensity):
|
||
# Darken shadows more as shadow_intensity increases, scaled by hdr_intensity
|
||
return np.clip(luminance_array - luminance_array * shadow_intensity * hdr_intensity * 0.5, 0, 255)
|
||
|
||
def adjust_highlights(luminance_array, highlight_intensity, hdr_intensity):
|
||
# Brighten highlights more as highlight_intensity increases, scaled by hdr_intensity
|
||
return np.clip(luminance_array + (255 - luminance_array) * highlight_intensity * hdr_intensity * 0.5, 0, 255)
|
||
|
||
def apply_adjustment(base, factor, intensity_scale):
|
||
"""Apply positive adjustment scaled by intensity."""
|
||
# Ensure the adjustment increases values within [0, 1] range, scaling by intensity
|
||
adjustment = base + (base * factor * intensity_scale)
|
||
# Ensure adjustment stays within bounds
|
||
return np.clip(adjustment, 0, 1)
|
||
|
||
|
||
def multiply_blend(base, blend):
|
||
"""Multiply blend mode."""
|
||
return np.clip(base * blend, 0, 255)
|
||
|
||
def overlay_blend(base, blend):
|
||
"""Overlay blend mode."""
|
||
# Normalize base and blend to [0, 1] for blending calculation
|
||
base = base / 255.0
|
||
blend = blend / 255.0
|
||
return np.where(base < 0.5, 2 * base * blend, 1 - 2 * (1 - base) * (1 - blend)) * 255
|
||
|
||
def adjust_shadows_non_linear(luminance, shadow_intensity, max_shadow_adjustment=1.5):
|
||
lum_array = np.array(luminance, dtype=np.float32) / 255.0 # Normalize
|
||
# Apply a non-linear darkening effect based on shadow_intensity
|
||
shadows = lum_array ** (1 / (1 + shadow_intensity * max_shadow_adjustment))
|
||
return np.clip(shadows * 255, 0, 255).astype(np.uint8) # Re-scale to [0, 255]
|
||
|
||
def adjust_highlights_non_linear(luminance, highlight_intensity, max_highlight_adjustment=1.5):
|
||
lum_array = np.array(luminance, dtype=np.float32) / 255.0 # Normalize
|
||
# Brighten highlights more aggressively based on highlight_intensity
|
||
highlights = 1 - (1 - lum_array) ** (1 + highlight_intensity * max_highlight_adjustment)
|
||
return np.clip(highlights * 255, 0, 255).astype(np.uint8) # Re-scale to [0, 255]
|
||
|
||
def merge_adjustments_with_blend_modes(luminance, shadows, highlights, hdr_intensity, shadow_intensity, highlight_intensity):
|
||
# Ensure the data is in the correct format for processing
|
||
base = np.array(luminance, dtype=np.float32)
|
||
|
||
# Scale the adjustments based on hdr_intensity
|
||
scaled_shadow_intensity = shadow_intensity ** 2 * hdr_intensity
|
||
scaled_highlight_intensity = highlight_intensity ** 2 * hdr_intensity
|
||
|
||
# Create luminance-based masks for shadows and highlights
|
||
shadow_mask = np.clip((1 - (base / 255)) ** 2, 0, 1)
|
||
highlight_mask = np.clip((base / 255) ** 2, 0, 1)
|
||
|
||
# Apply the adjustments using the masks
|
||
adjusted_shadows = np.clip(base * (1 - shadow_mask * scaled_shadow_intensity), 0, 255)
|
||
adjusted_highlights = np.clip(base + (255 - base) * highlight_mask * scaled_highlight_intensity, 0, 255)
|
||
|
||
# Combine the adjusted shadows and highlights
|
||
adjusted_luminance = np.clip(adjusted_shadows + adjusted_highlights - base, 0, 255)
|
||
|
||
# Blend the adjusted luminance with the original luminance based on hdr_intensity
|
||
final_luminance = np.clip(base * (1 - hdr_intensity) + adjusted_luminance * hdr_intensity, 0, 255).astype(np.uint8)
|
||
|
||
return Image.fromarray(final_luminance)
|
||
|
||
def apply_gamma_correction(lum_array, gamma):
|
||
"""
|
||
Apply gamma correction to the luminance array.
|
||
:param lum_array: Luminance channel as a NumPy array.
|
||
:param gamma: Gamma value for correction.
|
||
"""
|
||
if gamma == 0:
|
||
return np.clip(lum_array, 0, 255).astype(np.uint8)
|
||
|
||
epsilon = 1e-7 # Small value to avoid dividing by zero
|
||
gamma_corrected = 1 / (1.1 - gamma)
|
||
adjusted = 255 * ((lum_array / 255) ** gamma_corrected)
|
||
return np.clip(adjusted, 0, 255).astype(np.uint8)
|
||
|
||
# create a wrapper function that can apply a function to multiple images in a batch while passing all other arguments to the function
|
||
def apply_to_batch(func):
|
||
def wrapper(self, image, *args, **kwargs):
|
||
images = []
|
||
for img in image:
|
||
images.append(func(self, img, *args, **kwargs))
|
||
batch_tensor = torch.cat(images, dim=0)
|
||
return (batch_tensor, )
|
||
return wrapper
|
||
|
||
class HDREffects:
|
||
DESCRIPTION = "Apply an HDR-style tone-mapping effect with separate control over shadows, highlights, gamma, contrast and colour boost. Accepts a batch of images and returns the processed batch."
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {'required': {'image': ('IMAGE', {'default': None}),
|
||
'hdr_intensity': ('FLOAT', {'default': 0.5, 'min': 0.0, 'max': 5.0, 'step': 0.01}),
|
||
'shadow_intensity': ('FLOAT', {'default': 0.25, 'min': 0.0, 'max': 1.0, 'step': 0.01}),
|
||
'highlight_intensity': ('FLOAT', {'default': 0.75, 'min': 0.0, 'max': 1.0, 'step': 0.01}),
|
||
'gamma_intensity': ('FLOAT', {'default': 0.0, 'min': 0.0, 'max': 1.0, 'step': 0.01}),
|
||
'contrast': ('FLOAT', {'default': 0.1, 'min': 0.0, 'max': 1.0, 'step': 0.01}),
|
||
'enhance_color': ('FLOAT', {'default': 0.25, 'min': 0.0, 'max': 1.0, 'step': 0.01})
|
||
}}
|
||
|
||
RETURN_TYPES = ('IMAGE',)
|
||
RETURN_NAMES = ('result_img',)
|
||
FUNCTION = 'apply_hdr2'
|
||
CATEGORY = 'SuperBeastsAI/Image'
|
||
|
||
@apply_to_batch
|
||
def apply_hdr2(self, image, hdr_intensity=0.5, shadow_intensity=0.25, highlight_intensity=0.75, gamma_intensity=0.25, contrast=0.1, enhance_color=0.25):
|
||
# Load the image and force a stable RGB input mode.
|
||
img = tensor2pil(image).convert('RGB')
|
||
|
||
# Convert to LAB without relying on Pillow/ImageCms. The CMS path can
|
||
# segfault inside the C extension on some Python/Pillow builds.
|
||
l_uint8, a_channel, b_channel = _rgb_to_lab_components(img)
|
||
luminance = Image.fromarray(l_uint8, mode='L')
|
||
|
||
# Convert luminance to a NumPy array for processing
|
||
lum_array = l_uint8.astype(np.float32)
|
||
|
||
# Preparing adjustment layers (shadows, midtones, highlights)
|
||
# This example assumes you have methods to extract or calculate these adjustments
|
||
shadows_adjusted = adjust_shadows_non_linear(luminance, shadow_intensity)
|
||
highlights_adjusted = adjust_highlights_non_linear(luminance, highlight_intensity)
|
||
|
||
|
||
merged_adjustments = merge_adjustments_with_blend_modes(lum_array, shadows_adjusted, highlights_adjusted, hdr_intensity, shadow_intensity, highlight_intensity)
|
||
|
||
# Apply gamma correction with a base_gamma value (define based on desired effect)
|
||
gamma_corrected = apply_gamma_correction(np.array(merged_adjustments), gamma_intensity)
|
||
if gamma_corrected.shape != l_uint8.shape:
|
||
gamma_corrected = np.array(
|
||
Image.fromarray(gamma_corrected, mode='L').resize((l_uint8.shape[1], l_uint8.shape[0]))
|
||
)
|
||
|
||
# Merge L channel back with original A/B components and convert back to RGB.
|
||
img_adjusted = _lab_components_to_rgb(gamma_corrected, a_channel, b_channel)
|
||
|
||
|
||
# Enhance contrast
|
||
enhancer = ImageEnhance.Contrast(img_adjusted)
|
||
contrast_adjusted = enhancer.enhance(1 + contrast)
|
||
|
||
|
||
# Enhance color saturation
|
||
enhancer = ImageEnhance.Color(contrast_adjusted)
|
||
color_adjusted = enhancer.enhance(1 + enhance_color * 0.2)
|
||
|
||
return pil2tensor(color_adjusted)
|
||
|
||
class MakeResizedMaskBatch:
|
||
DESCRIPTION = "Combine up to 12 individual masks/batches into one mask batch, automatically sizing/cropping each mask to the requested width/height. Useful for building consistent video masks."
|
||
def __init__(self):
|
||
pass
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"width": ("INT", {"default": 512, "min": 1, "step": 1}),
|
||
"height": ("INT", {"default": 768, "min": 1, "step": 1}),
|
||
},
|
||
"optional": {
|
||
"mask1": ("MASK",),
|
||
"mask2": ("MASK",),
|
||
"mask3": ("MASK",),
|
||
"mask4": ("MASK",),
|
||
"mask5": ("MASK",),
|
||
"mask6": ("MASK",),
|
||
"mask7": ("MASK",),
|
||
"mask8": ("MASK",),
|
||
"mask9": ("MASK",),
|
||
"mask10": ("MASK",),
|
||
"mask11": ("MASK",),
|
||
"mask12": ("MASK",),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("MASK",)
|
||
FUNCTION = "append"
|
||
CATEGORY = "SuperBeastsAI/Masks"
|
||
|
||
def append(self, width, height, mask1=None, mask2=None, mask3=None, mask4=None, mask5=None, mask6=None,
|
||
mask7=None, mask8=None, mask9=None, mask10=None, mask11=None, mask12=None):
|
||
masks = [mask1, mask2, mask3, mask4, mask5, mask6, mask7, mask8, mask9, mask10, mask11, mask12]
|
||
valid_masks = [mask for mask in masks if mask is not None]
|
||
|
||
if not valid_masks:
|
||
raise ValueError("At least one input mask must be provided.")
|
||
|
||
cropped_masks = []
|
||
for mask in valid_masks:
|
||
if mask.ndim == 2:
|
||
mask = mask.unsqueeze(0).unsqueeze(0) # Add batch and channel dimensions
|
||
elif mask.ndim == 3:
|
||
mask = mask.unsqueeze(0) # Add batch dimension
|
||
elif mask.ndim != 4 or (mask.ndim == 4 and mask.shape[1] != 1):
|
||
raise ValueError(f"Invalid mask shape: {mask.shape}. Expected (N, 1, H, W) or (1, H, W) or (H, W).")
|
||
|
||
# Scale the mask to match the desired width while maintaining the aspect ratio
|
||
scale_factor = width / mask.shape[-1]
|
||
scaled_height = int(mask.shape[-2] * scale_factor)
|
||
scaled_mask = F.interpolate(mask, size=(scaled_height, width), mode='bilinear', align_corners=False)
|
||
|
||
# Perform center cropping
|
||
if scaled_height < height:
|
||
# Pad the top and bottom of the mask
|
||
pad_top = (height - scaled_height) // 2
|
||
pad_bottom = height - scaled_height - pad_top
|
||
cropped_mask = F.pad(scaled_mask, (0, 0, pad_top, pad_bottom), mode='constant', value=0)
|
||
else:
|
||
# Crop the center of the mask
|
||
crop_top = (scaled_height - height) // 2
|
||
crop_bottom = crop_top + height
|
||
cropped_mask = scaled_mask[:, :, crop_top:crop_bottom, :]
|
||
|
||
cropped_masks.append(cropped_mask)
|
||
|
||
# Concatenate the cropped masks along the batch dimension
|
||
result = torch.cat(cropped_masks, dim=0)
|
||
|
||
return (result,)
|
||
|
||
|
||
|
||
def adjust_brightness(image, brightness_factor):
|
||
enhancer = ImageEnhance.Brightness(image)
|
||
adjusted_image = enhancer.enhance(brightness_factor)
|
||
return adjusted_image
|
||
|
||
def calculate_brightness_factor(target_brightness, current_brightness):
|
||
return target_brightness / current_brightness
|
||
|
||
def get_average_brightness(image):
|
||
grayscale_image = image.convert("L")
|
||
histogram = grayscale_image.histogram()
|
||
pixels = sum(histogram)
|
||
brightness = scale = len(histogram)
|
||
|
||
total_brightness = sum(i * w for i, w in enumerate(histogram))
|
||
return total_brightness / pixels
|
||
|
||
def apply_dithering(image):
|
||
return image.convert("P", palette=Image.ADAPTIVE, colors=256).convert("RGB")
|
||
|
||
def apply_noise_reduction(image, strength):
|
||
return image.filter(ImageFilter.GaussianBlur(radius=strength))
|
||
|
||
def apply_gradient_smoothing(image, strength):
|
||
return image.filter(ImageFilter.SMOOTH_MORE if strength > 1 else ImageFilter.SMOOTH)
|
||
|
||
def blend_images(image1, image2, alpha):
|
||
return Image.blend(image1, image2, alpha)
|
||
|
||
class Deflicker:
|
||
DESCRIPTION = "Experimental high-level deflicker pass for video/animation. Analyses brightness across neighbouring frames and blends, denoises and smooths gradients to reduce global flicker."
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"images": ("IMAGE",),
|
||
"context_length": ("INT", {"default": 5, "min": 1, "max": 20, "step": 1}),
|
||
"brightness_threshold": ("FLOAT", {"default": 0.05, "min": 0.01, "max": 0.5, "step": 0.01}),
|
||
"blending_strength": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"noise_reduction_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}),
|
||
"gradient_smoothing_strength": ("INT", {"default": 1, "min": 0, "max": 3, "step": 1}),
|
||
"batch_size": ("INT", {"default": 10, "min": 1, "max": 100, "step": 1})
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
FUNCTION = "deflicker"
|
||
CATEGORY = "SuperBeastsAI/Animation"
|
||
|
||
def deflicker(self, images, context_length=5, brightness_threshold=0.05, blending_strength=0.5,
|
||
noise_reduction_strength=1.0, gradient_smoothing_strength=1, batch_size=10):
|
||
num_frames = len(images)
|
||
adjusted_tensor = []
|
||
|
||
for i in range(0, num_frames, batch_size):
|
||
batch_images = images[i:i+batch_size]
|
||
|
||
# Convert batch tensor to a list of PIL images
|
||
pil_images = [tensor2pil(image) for image in batch_images]
|
||
|
||
adjusted_images = []
|
||
|
||
for j in range(len(pil_images)):
|
||
current_image = pil_images[j]
|
||
context_start = max(0, i + j - context_length // 2)
|
||
context_end = min(num_frames, i + j + context_length // 2 + 1)
|
||
context_images = images[context_start:context_end]
|
||
|
||
current_brightness = get_average_brightness(current_image)
|
||
context_brightnesses = [get_average_brightness(tensor2pil(img)) for img in context_images]
|
||
average_brightness = np.mean(context_brightnesses)
|
||
|
||
if abs(current_brightness - average_brightness) > brightness_threshold:
|
||
brightness_factor = calculate_brightness_factor(average_brightness, current_brightness)
|
||
adjusted_image = adjust_brightness(current_image, brightness_factor)
|
||
else:
|
||
adjusted_image = current_image
|
||
|
||
# Apply noise reduction to the adjusted image
|
||
denoised_image = apply_noise_reduction(adjusted_image, noise_reduction_strength)
|
||
|
||
# Apply gradient smoothing to the denoised image
|
||
smoothed_image = apply_gradient_smoothing(denoised_image, gradient_smoothing_strength)
|
||
|
||
# Apply dithering to the smoothed image
|
||
dithered_image = apply_dithering(smoothed_image)
|
||
|
||
# Blend the dithered image with the original image using adaptive blending
|
||
blending_alpha = min(1.0, blending_strength * (1.0 + abs(current_brightness - average_brightness)))
|
||
blended_image = blend_images(current_image, dithered_image, blending_alpha)
|
||
|
||
adjusted_images.append(blended_image)
|
||
|
||
# Convert the adjusted PIL images back to a tensor
|
||
adjusted_batch_tensor = torch.cat([pil2tensor(img) for img in adjusted_images], dim=0)
|
||
adjusted_tensor.append(adjusted_batch_tensor)
|
||
|
||
# Concatenate the adjusted batches along the first dimension
|
||
adjusted_tensor = torch.cat(adjusted_tensor, dim=0)
|
||
|
||
return (adjusted_tensor,)
|
||
|
||
|
||
def temporal_smoothing(frames, window_size):
|
||
num_frames = len(frames)
|
||
smoothed_frames = []
|
||
|
||
for i in range(num_frames):
|
||
start = max(0, i - window_size // 2)
|
||
end = min(num_frames, i + window_size // 2 + 1)
|
||
window_frames = frames[start:end]
|
||
|
||
smoothed_frame = np.mean(window_frames, axis=0)
|
||
smoothed_frames.append(smoothed_frame)
|
||
|
||
return smoothed_frames
|
||
|
||
class PixelDeflicker:
|
||
DESCRIPTION = "Experimental per-pixel temporal smoothing for animation. Operates in a sliding window to average noisy pixels while preserving detail, helping mitigate small-scale flicker."
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"images": ("IMAGE",),
|
||
"window_size": ("INT", {"default": 2, "min": 1, "max": 20, "step": 1}),
|
||
"blending_strength": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"batch_size": ("INT", {"default": 10, "min": 1, "max": 100, "step": 1})
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
FUNCTION = "pixelDeflicker"
|
||
|
||
CATEGORY = "SuperBeastsAI/Animation"
|
||
|
||
def pixelDeflicker(self, images, window_size=5, blending_strength=0.5, batch_size=10):
|
||
num_frames = len(images)
|
||
blended_tensor = []
|
||
|
||
for i in range(0, num_frames, batch_size):
|
||
batch_images = images[i:i+batch_size]
|
||
|
||
# Convert batch tensor to a list of PIL images
|
||
pil_images = [tensor2pil(image) for image in batch_images]
|
||
|
||
# Convert PIL images to numpy arrays
|
||
numpy_frames = [np.array(img) / 255.0 for img in pil_images]
|
||
|
||
# Apply temporal smoothing to the numpy frames
|
||
smoothed_frames = temporal_smoothing(numpy_frames, window_size)
|
||
|
||
# Blend the smoothed frames with the original frames
|
||
blended_frames = [
|
||
np.clip(original * (1 - blending_strength) + smoothed * blending_strength, 0, 1)
|
||
for original, smoothed in zip(numpy_frames, smoothed_frames)
|
||
]
|
||
|
||
# Convert the blended frames back to PIL images
|
||
blended_pil_images = [Image.fromarray((frame * 255).astype(np.uint8)) for frame in blended_frames]
|
||
|
||
# Convert the blended PIL images back to a tensor
|
||
blended_batch_tensor = torch.cat([pil2tensor(img) for img in blended_pil_images], dim=0)
|
||
|
||
blended_tensor.append(blended_batch_tensor)
|
||
|
||
# Concatenate the blended batches along the first dimension
|
||
blended_tensor = torch.cat(blended_tensor, dim=0)
|
||
|
||
return (blended_tensor,)
|
||
|
||
def resize_and_crop(pil_img, target_width, target_height):
|
||
"""Resize and crop an image to fit exactly the specified dimensions."""
|
||
original_width, original_height = pil_img.size
|
||
aspect_ratio = original_width / original_height
|
||
target_aspect_ratio = target_width / target_height
|
||
|
||
if target_aspect_ratio > aspect_ratio:
|
||
# Target is wider than the image
|
||
scale_factor = target_width / original_width
|
||
scaled_height = int(original_height * scale_factor)
|
||
scaled_width = target_width
|
||
else:
|
||
# Target is taller than the image
|
||
scale_factor = target_height / original_height
|
||
scaled_height = target_height
|
||
scaled_width = int(original_width * scale_factor)
|
||
|
||
# Resize the image
|
||
resized_img = pil_img.resize((scaled_width, scaled_height), Image.BILINEAR)
|
||
|
||
# Crop the image
|
||
if scaled_width != target_width or scaled_height != target_height:
|
||
left = (scaled_width - target_width) // 2
|
||
top = (scaled_height - target_height) // 2
|
||
right = left + target_width
|
||
bottom = top + target_height
|
||
cropped_img = resized_img.crop((left, top, right, bottom))
|
||
else:
|
||
cropped_img = resized_img
|
||
|
||
return cropped_img
|
||
|
||
class ImageBatchManagement:
|
||
DESCRIPTION = "Resize, crop, limit, optionally shuffle or manually reorder an image batch so all frames match the specified resolution and sequence order. Returns the new batch and a CSV string listing filenames."
|
||
def __init__(self):
|
||
pass
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"width": ("INT", {"default": 512, "order": 1}),
|
||
"height": ("INT", {"default": 768}),
|
||
"max_images": ("INT", {"default": 10}), # New INT input for maximum number of images
|
||
"random_order": ("BOOLEAN", {"default": False})
|
||
},
|
||
"optional": {
|
||
"new_manual_order": ("STRING", {"default": ""}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "STRING")
|
||
FUNCTION = "reorder"
|
||
CATEGORY = "SuperBeastsAI/Image"
|
||
|
||
def reorder(self, width, height, random_order, max_images, **kwargs):
|
||
images = [kwargs["image1"]] # Start with the required image1 input
|
||
|
||
i = 2
|
||
while f"image{i}" in kwargs:
|
||
images.append(kwargs[f"image{i}"])
|
||
i += 1
|
||
|
||
if max_images is not None:
|
||
images = images[:max_images]
|
||
|
||
# Default order_output if new_manual_order isn't provided or is empty
|
||
order_output = ",".join(str(idx + 1) for idx in range(len(images)))
|
||
|
||
# Retrieve and apply new_manual_order if it exists
|
||
if 'new_manual_order' in kwargs and kwargs['new_manual_order']:
|
||
order_indices = [int(idx) - 1 for idx in kwargs['new_manual_order'].split(',') if idx.strip()]
|
||
images = [images[idx] for idx in order_indices if idx < len(images)]
|
||
order_output = kwargs['new_manual_order']
|
||
|
||
processed_images = []
|
||
for img in images:
|
||
pil_img = tensor2pil(img)
|
||
resized_cropped_img = resize_and_crop(pil_img, width, height)
|
||
img_tensor = pil2tensor(resized_cropped_img)
|
||
processed_images.append(img_tensor)
|
||
|
||
result = torch.cat(processed_images, dim=0) if processed_images else torch.empty(0, 3, height, width)
|
||
return (result, order_output)
|
||
|
||
|
||
class MaskBatchManagement:
|
||
DESCRIPTION = "Resize, crop and re-order a batch of masks so they all share the same dimensions, or to match a new manual order list."
|
||
def __init__(self):
|
||
pass
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"width": ("INT", {"default": 512}),
|
||
"height": ("INT", {"default": 768})
|
||
},
|
||
"optional": {
|
||
"new_order": ("STRING", {"default": ""}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("MASK",)
|
||
FUNCTION = "append"
|
||
CATEGORY = "SuperBeastsAI/Masks"
|
||
|
||
def append(self, width, height, new_order, **kwargs):
|
||
masks = [kwargs["mask1"]] # Start with the required mask1 input
|
||
|
||
i = 2
|
||
while f"mask{i}" in kwargs:
|
||
masks.append(kwargs[f"mask{i}"])
|
||
i += 1
|
||
|
||
if new_order:
|
||
order_indices = [int(idx) - 1 for idx in new_order.split(',') if idx.strip()]
|
||
masks = [masks[idx] for idx in order_indices if idx < len(masks)]
|
||
|
||
processed_masks = []
|
||
for mask in masks:
|
||
pil_mask = tensor2pil(mask)
|
||
resized_cropped_mask = resize_and_crop(pil_mask, width, height)
|
||
mask_tensor = pil2tensor(resized_cropped_mask)
|
||
processed_masks.append(mask_tensor)
|
||
|
||
result = torch.cat(processed_masks, dim=0) if processed_masks else torch.empty(0, 1, height, width)
|
||
|
||
return (result,)
|
||
|
||
class StringListManager:
|
||
DESCRIPTION = "Re-order or time-expand a list of strings (e.g. prompts or captions) so each item can repeat for N frames. Outputs the adjusted list as a single newline-separated string."
|
||
def __init__(self):
|
||
pass
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"frames_per_image": ("INT", {"default": 1, "min": 1, "step": 1})
|
||
},
|
||
"optional": {
|
||
"new_order": ("STRING", {"default": ""}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
FUNCTION = "reorder_strings"
|
||
CATEGORY = "SuperBeastsAI/Utils"
|
||
|
||
def reorder_strings(self, frames_per_image, new_order, **kwargs):
|
||
strings = [kwargs["string1"]] # Start with the required string1 input
|
||
|
||
i = 2
|
||
while f"string{i}" in kwargs:
|
||
strings.append(kwargs[f"string{i}"])
|
||
i += 1
|
||
|
||
if new_order:
|
||
order_indices = [int(idx) - 1 for idx in new_order.split(',') if idx.strip()]
|
||
strings = [strings[idx] for idx in order_indices if idx < len(strings)]
|
||
|
||
result = []
|
||
for i, string in enumerate(strings):
|
||
result.append('"{frames}": "{string}"'.format(frames=frames_per_image * i, string=string))
|
||
|
||
return (",\n".join(result),)
|
||
|
||
class SuperPopColorAdjustment:
|
||
"""Generate a series of color-adjusted images by blending the residual produced by a model-based correction with the original image at different strengths.
|
||
|
||
The node calls the provided *SBModel* once to obtain a fully-corrected reference image (Requires use of "SB Load Model" node). It then blends that result back into the
|
||
original image `count` times using evenly spaced strength values between `max_strength/count` and `max_strength` (inclusive).
|
||
"""
|
||
DESCRIPTION = __doc__
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"model": ("SBMODEL", ),
|
||
"image": ("IMAGE", ),
|
||
"max_strength": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
|
||
"count": ("INT", {"default": 1, "min": 1, "max": 99999, "step": 1}),
|
||
"overlap": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 0.99, "step": 0.01}),
|
||
"initial_context_for_batch": ("BOOLEAN", {"default": False}),
|
||
},
|
||
"optional": {
|
||
"context": ("IMAGE", ), # Optional context image(s)
|
||
}
|
||
}
|
||
|
||
# primary output: processed image batch
|
||
# secondary output: per-frame filename prefix strings for Save Image
|
||
# third output: raw residual tensors (type SPCA_RESIDUAL) so that a
|
||
# dedicated blend node can re-apply any arbitrary strength later on.
|
||
RETURN_TYPES = ("IMAGE", "STRING", "SPCA_RESIDUAL")
|
||
RETURN_NAMES = ("images", "filename_prefix", "residuals")
|
||
FUNCTION = "apply_adjustment"
|
||
CATEGORY = "SuperBeastsAI/Image"
|
||
|
||
def _run_model(self, sb_model_info, pil_image, ctx_np=None):
|
||
"""Attempt to run the supplied SBModel on *pil_image* to obtain a corrected image.
|
||
The implementation gracefully falls back to a PIL colour enhancement if the model
|
||
does not expose a compatible callable interface yet. This allows the node to work
|
||
out-of-the-box while the dedicated SuperBeasts model loader is still under development."""
|
||
try:
|
||
# Build context thumbnail from global image once if not provided
|
||
if ctx_np is None:
|
||
thumb = pil_image.resize((CONTEXT_THUMBNAIL_SIZE[1], CONTEXT_THUMBNAIL_SIZE[0]), Image.Resampling.BILINEAR)
|
||
ctx_np_local = _add_extra_color_channels_np(np.asarray(thumb, dtype=np.float32) / 255.0)
|
||
else:
|
||
ctx_np_local = ctx_np
|
||
|
||
# --- Torch path ---
|
||
if isinstance(sb_model_info, dict) and sb_model_info.get("type") == "torch":
|
||
import torch
|
||
model_pt = sb_model_info["model"]
|
||
device_str = sb_model_info["device"]
|
||
|
||
# Prepare inputs NCHW float32
|
||
tile_np = _add_extra_color_channels_np(np.asarray(pil_image, dtype=np.float32) / 255.0)
|
||
|
||
tile_t = torch.from_numpy(tile_np.transpose(2,0,1)).unsqueeze(0).to(device_str)
|
||
ctx_t = torch.from_numpy(ctx_np_local.transpose(2,0,1)).unsqueeze(0).to(device_str)
|
||
|
||
with torch.no_grad():
|
||
residual = model_pt(tile_t, ctx_t)
|
||
residual_np = residual.squeeze(0).cpu().numpy().transpose(1,2,0)
|
||
|
||
tile_base = np.asarray(pil_image, dtype=np.float32) / 255.0
|
||
corrected_np = np.clip(tile_base + residual_np, 0.0, 1.0)
|
||
result = Image.fromarray((corrected_np * 255.0).astype(np.uint8))
|
||
|
||
# --- ONNX path ---
|
||
elif isinstance(sb_model_info, dict) and sb_model_info.get("type") == "onnx":
|
||
import onnxruntime as ort # type: ignore
|
||
session: ort.InferenceSession = sb_model_info["session"]
|
||
input_names: List[str] = sb_model_info["input_names"]
|
||
|
||
# Prepare inputs (NHWC float32)
|
||
ctx_np_batch = ctx_np_local[np.newaxis, ...] # add batch dim
|
||
|
||
tile_np = _add_extra_color_channels_np(np.asarray(pil_image, dtype=np.float32) / 255.0)
|
||
tile_np = tile_np[np.newaxis, ...]
|
||
|
||
# Some exported models keep NHWC; others are converted to NCHW. Detect by checking input shape channels position.
|
||
def maybe_transpose(inp: np.ndarray, shape) -> np.ndarray:
|
||
# shape may have None for batch, use length check
|
||
if len(shape) == 4 and shape[1] == 3: # likely NCHW
|
||
return np.transpose(inp, (0, 3, 1, 2))
|
||
return inp # assume NHWC
|
||
|
||
tile_ready = maybe_transpose(tile_np, session.get_inputs()[0].shape)
|
||
ctx_ready = maybe_transpose(ctx_np_batch, session.get_inputs()[1].shape)
|
||
|
||
ort_inputs = {
|
||
input_names[0]: tile_ready.astype(np.float32),
|
||
input_names[1]: ctx_ready.astype(np.float32),
|
||
}
|
||
|
||
pred_residual = session.run(None, ort_inputs)[0]
|
||
if pred_residual.shape[1] == 3 and pred_residual.ndim == 4: # NCHW output
|
||
pred_residual = np.transpose(pred_residual, (0, 2, 3, 1))
|
||
residual_np = pred_residual[0]
|
||
|
||
tile_base = np.asarray(pil_image, dtype=np.float32) / 255.0
|
||
corrected_np = np.clip(tile_base + residual_np, 0.0, 1.0)
|
||
result = Image.fromarray((corrected_np * 255.0).astype(np.uint8))
|
||
|
||
else:
|
||
# If the supplied model information structure is not recognised, raise an error so
|
||
# users are immediately aware of the mis-configuration instead of falling back.
|
||
raise TypeError("Unsupported SBMODEL object – expected dict with type 'torch' or 'onnx'.")
|
||
|
||
# Convert potential tensor outputs to PIL
|
||
try:
|
||
import torch as _torch
|
||
except ImportError:
|
||
_torch = None
|
||
|
||
if _torch is not None and isinstance(result, _torch.Tensor):
|
||
result = tensor2pil(result)
|
||
if not isinstance(result, Image.Image):
|
||
raise TypeError("Model output is not an image")
|
||
return result
|
||
except Exception:
|
||
# Re-raise any exception so the UI surfaces the real failure rather than silently
|
||
# producing an uncorrected image.
|
||
raise
|
||
|
||
def _get_patch_positions(self, img_len: int, patch_len: int, overlap_px: int):
|
||
"""Return list of starting indices so patches cover dimension with given overlap."""
|
||
if patch_len >= img_len:
|
||
return [0]
|
||
stride = max(1, patch_len - overlap_px)
|
||
positions = [0]
|
||
while True:
|
||
next_pos = positions[-1] + stride
|
||
if next_pos + patch_len >= img_len:
|
||
break
|
||
positions.append(next_pos)
|
||
last_start = img_len - patch_len
|
||
if positions[-1] != last_start:
|
||
positions.append(last_start)
|
||
return positions
|
||
|
||
# --- Patch weight map (mirrors extras/inference_utils._create_patch_weight_map) ---
|
||
def _create_patch_weight_map(self, ph: int, pw: int, overlap: int | tuple[int,int]):
|
||
if isinstance(overlap, (list, tuple)) and len(overlap) == 2:
|
||
oh, ow = int(overlap[0]), int(overlap[1])
|
||
else:
|
||
oh = ow = int(overlap)
|
||
|
||
oh = max(0, min(oh, ph))
|
||
ow = max(0, min(ow, pw))
|
||
|
||
if oh == 0 and ow == 0:
|
||
return np.ones((ph, pw), dtype=np.float32)
|
||
|
||
w = np.ones((ph, pw), dtype=np.float32)
|
||
if oh > 0:
|
||
for i in range(oh):
|
||
fade = (i + 1) / (oh + 1)
|
||
w[i, :] *= fade
|
||
w[ph - 1 - i, :] *= fade
|
||
if ow > 0:
|
||
for i in range(ow):
|
||
fade = (i + 1) / (ow + 1)
|
||
w[:, i] *= fade
|
||
w[:, pw - 1 - i] *= fade
|
||
return w
|
||
|
||
def apply_adjustment(self, model, image, max_strength=1.0, count=1, overlap=0.5,
|
||
initial_context_for_batch=False, context=None):
|
||
# Ensure *image* is a batch tensor – iterate over each frame
|
||
output_tensors: list[torch.Tensor] = []
|
||
filename_prefixes: list[str] = []
|
||
residual_tensors: list[torch.Tensor] = []
|
||
|
||
# Cache for the first computed context thumbnail when reusing across batch
|
||
ctx_np_cached = None
|
||
|
||
# Helper to obtain the source PIL for context, considering optional input
|
||
def _get_context_pil(idx, orig_pil):
|
||
if context is not None:
|
||
# Guard against length mismatch – fallback to first element if idx out of range
|
||
ctx_tensor = context[idx] if idx < context.shape[0] else context[0]
|
||
return tensor2pil(ctx_tensor)
|
||
# Default: use the original image as its own context
|
||
return orig_pil
|
||
|
||
pbar = None # progress bar placeholder
|
||
try:
|
||
import importlib
|
||
ProgressBar = getattr(importlib.import_module("comfy.utils"), "ProgressBar", None) # type: ignore[attr-defined,ignored-type]
|
||
if ProgressBar is not None:
|
||
pbar = ProgressBar(image.shape[0])
|
||
except Exception:
|
||
pbar = None
|
||
|
||
for idx, img_tensor in enumerate(image):
|
||
# Convert to PIL for easier colour operations
|
||
original_pil = tensor2pil(img_tensor)
|
||
|
||
# -------------------------------------------------
|
||
# Prepare / reuse context thumbnail
|
||
# -------------------------------------------------
|
||
if initial_context_for_batch and ctx_np_cached is not None:
|
||
ctx_np_global = ctx_np_cached
|
||
else:
|
||
ctx_source_pil = _get_context_pil(idx, original_pil)
|
||
# Resize exactly to CONTEXT_THUMBNAIL_SIZE using bicubic resampling
|
||
thumb_global = ctx_source_pil.resize((CONTEXT_THUMBNAIL_SIZE[1], CONTEXT_THUMBNAIL_SIZE[0]),
|
||
Image.Resampling.BICUBIC)
|
||
ctx_np_global = _add_extra_color_channels_np(np.asarray(thumb_global, dtype=np.float32) / 255.0)
|
||
|
||
if initial_context_for_batch:
|
||
ctx_np_cached = ctx_np_global
|
||
|
||
# ---------- Patch-based processing ----------
|
||
h, w = original_pil.size[1], original_pil.size[0]
|
||
patch_size = model.get("patch_size", 512) if isinstance(model, dict) else 512
|
||
patch_size = int(patch_size)
|
||
overlap_px = int(patch_size * overlap)
|
||
if overlap_px >= patch_size:
|
||
overlap_px = patch_size // 2
|
||
|
||
# Prepare accumulators for residual image
|
||
residual_acc = np.zeros((h, w, 3), dtype=np.float32)
|
||
counter = np.zeros((h, w, 1), dtype=np.float32)
|
||
|
||
x_positions = self._get_patch_positions(w, patch_size, overlap_px)
|
||
y_positions = self._get_patch_positions(h, patch_size, overlap_px)
|
||
|
||
orig_np_full = np.array(original_pil, dtype=np.float32) / 255.0
|
||
|
||
# ctx_np_global already prepared above
|
||
|
||
for y in y_positions:
|
||
for x in x_positions:
|
||
patch = original_pil.crop((x, y, x + patch_size, y + patch_size))
|
||
patch_corrected = self._run_model(model, patch, ctx_np_global)
|
||
|
||
# Resize patch_corrected to patch size in case model changes size
|
||
if patch_corrected.size != (patch_size, patch_size):
|
||
patch_corrected = patch_corrected.resize((patch_size, patch_size), Image.Resampling.BILINEAR)
|
||
|
||
orig_patch_np = np.array(patch, dtype=np.float32) / 255.0
|
||
corr_patch_np = np.array(patch_corrected, dtype=np.float32) / 255.0
|
||
|
||
# Determine overlap area within the original image boundaries
|
||
target_h = min(patch_size, h - y)
|
||
target_w = min(patch_size, w - x)
|
||
|
||
# If the patch is already full size use it directly -----------------------
|
||
if target_h == patch_size and target_w == patch_size:
|
||
residual_patch = corr_patch_np - orig_patch_np
|
||
else:
|
||
# --------------------------------------------------------------
|
||
# Pad the original patch to 512×512 instead of scaling to avoid
|
||
# introducing blur / ghosting. Use edge padding.
|
||
# --------------------------------------------------------------
|
||
# Crop to the valid region before padding
|
||
orig_cropped = orig_patch_np[:target_h, :target_w, :]
|
||
|
||
# --- Symmetric edge-padding approach (avoid stretched background artefacts) ---
|
||
pad_vert = patch_size - target_h
|
||
pad_horz = patch_size - target_w
|
||
pad_top = pad_vert // 2
|
||
pad_bottom = pad_vert - pad_top
|
||
pad_left = pad_horz // 2
|
||
pad_right = pad_horz - pad_left
|
||
|
||
# Edge-pad the crop up to (patch_size, patch_size)
|
||
orig_pad = np.pad(orig_cropped, (
|
||
(pad_top, pad_bottom),
|
||
(pad_left, pad_right),
|
||
(0, 0)), mode='edge')
|
||
|
||
# Convert to PIL for model inference
|
||
pad_pil = Image.fromarray((orig_pad * 255.0).astype(np.uint8))
|
||
|
||
# Run model on padded patch – override earlier corr_patch_np
|
||
patch_corrected = self._run_model(model, pad_pil, ctx_np_global)
|
||
|
||
if patch_corrected.size != (patch_size, patch_size):
|
||
patch_corrected = patch_corrected.resize((patch_size, patch_size), Image.Resampling.BILINEAR)
|
||
|
||
corrected_pad_np = np.array(patch_corrected, dtype=np.float32) / 255.0
|
||
|
||
# Crop back to original region
|
||
corrected_crop = corrected_pad_np[pad_top:pad_top + target_h, pad_left:pad_left + target_w, :]
|
||
|
||
# Residual relative to original crop
|
||
residual_patch = corrected_crop - orig_cropped
|
||
|
||
# Weight map for smooth blending (same crop)
|
||
weight_full = self._create_patch_weight_map(patch_size, patch_size, overlap_px)
|
||
if target_h != patch_size or target_w != patch_size:
|
||
weight = weight_full[pad_top:pad_top + target_h, pad_left:pad_left + target_w]
|
||
else:
|
||
weight = weight_full[:target_h, :target_w]
|
||
|
||
# Determine the actual location in the full image where this residual belongs
|
||
residual_acc[y:y + target_h, x:x + target_w, :] += residual_patch * weight[:, :, np.newaxis]
|
||
counter[y:y + target_h, x:x + target_w, :] += weight[:, :, np.newaxis]
|
||
|
||
# Use a tiny epsilon instead of 1.0 so that edge pixels retain their full residual
|
||
counter = np.maximum(counter, 1e-6)
|
||
residual_full = residual_acc / counter
|
||
|
||
# Generate blended outputs for this frame
|
||
strengths = [max_strength * (count - i) / count for i in range(count)]
|
||
for s in strengths:
|
||
blended_np = np.clip(orig_np_full + residual_full * s, 0.0, 1.0)
|
||
blended_pil = Image.fromarray((blended_np * 255.0).astype(np.uint8))
|
||
|
||
# Convert blended PIL → tensor and append to batch list
|
||
output_tensors.append(pil2tensor(blended_pil))
|
||
|
||
# Build filename prefix string per‐image (we will join later)
|
||
filename_prefixes.append(f"SPCA_{s:.2f}_")
|
||
|
||
# Store residual for this output so downstream nodes can re-blend
|
||
residual_tensor = torch.from_numpy(residual_full).unsqueeze(0).float()
|
||
residual_tensors.append(residual_tensor)
|
||
# Update progress bar
|
||
if pbar is not None:
|
||
pbar.update_absolute(idx + 1)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Finalise outputs – convert accumulated lists to batched tensors and
|
||
# join filename prefixes into a single string (Save Image expects str)
|
||
# ------------------------------------------------------------------
|
||
|
||
if not output_tensors:
|
||
# Fallback – no processing happened; return originals and zero residual.
|
||
images_batch = image
|
||
residual_batch = [torch.zeros_like(image)]
|
||
filename_prefix_str = ""
|
||
else:
|
||
images_batch = torch.cat(output_tensors, dim=0)
|
||
residual_batch = residual_tensors # keep list per-frame
|
||
filename_prefix_str = "".join(filename_prefixes)
|
||
|
||
return (images_batch, filename_prefix_str, residual_batch)
|
||
|
||
class SuperPopResidualBlend:
|
||
"""Interactively re-apply a residual predicted by **Super Pop Color Adjustment** at any arbitrary strength.
|
||
|
||
Inputs
|
||
------
|
||
image : original image batch to adjust
|
||
residual : batch of residual tensors output from SuperPopColorAdjustment
|
||
strength : scalar multiplier (e.g. 0.0-2.0) applied to the residual
|
||
"""
|
||
|
||
DESCRIPTION = __doc__
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE", ),
|
||
"residual": ("SPCA_RESIDUAL", ),
|
||
"strength": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "STRING")
|
||
RETURN_NAMES = ("images", "filename_prefix")
|
||
FUNCTION = "blend"
|
||
CATEGORY = "SuperBeastsAI/Image"
|
||
|
||
def blend(self, image, residual, strength=1.0):
|
||
"""Apply *residual × strength* to *image* per-frame."""
|
||
import torch
|
||
|
||
# ------------------------------------------------------------------
|
||
# Normalise inputs to per-frame lists --------------------------------
|
||
# ------------------------------------------------------------------
|
||
|
||
if isinstance(image, torch.Tensor):
|
||
# Split the batch tensor into per-frame tensors (keeps gradients off)
|
||
image_list = list(image)
|
||
elif isinstance(image, (list, tuple)):
|
||
image_list = list(image)
|
||
else:
|
||
raise TypeError("Unsupported image type for SuperPopResidualBlend – expected torch.Tensor or list of tensors.")
|
||
|
||
if isinstance(residual, torch.Tensor):
|
||
residual_list = list(residual)
|
||
elif isinstance(residual, (list, tuple)):
|
||
residual_list = list(residual)
|
||
else:
|
||
raise TypeError("Unsupported residual type for SuperPopResidualBlend – expected torch.Tensor or list of tensors.")
|
||
|
||
# Simple broadcast if we only got one residual for many images
|
||
if len(residual_list) == 1 and len(image_list) > 1:
|
||
residual_list = residual_list * len(image_list)
|
||
|
||
if len(image_list) != len(residual_list):
|
||
raise ValueError("Mismatch between number of images and residuals ({} vs {}).".format(len(image_list), len(residual_list)))
|
||
|
||
# ------------------------------------------------------------------
|
||
# Blend per frame ----------------------------------------------------
|
||
# ------------------------------------------------------------------
|
||
|
||
blended_frames: list[torch.Tensor] = []
|
||
prefixes: list[str] = []
|
||
|
||
for img_t, res_t in zip(image_list, residual_list):
|
||
# Ensure on same device (fallback to img_t's device)
|
||
res_t = res_t.to(img_t.device)
|
||
|
||
# Both tensors may include an extra leading batch/channel dimension; squeeze safely
|
||
img = img_t.squeeze().clamp(0,1)
|
||
res = res_t.squeeze()
|
||
|
||
# Quantise the source to 8-bit – SuperPopColorAdjustment converted the
|
||
# original tensor to PIL (uint8) before computing residuals. Without
|
||
# replicating that rounding we would blend residuals into a higher-
|
||
# precision version of the source, leading to small but noticeable
|
||
# mismatches (especially at strength = 1.0). By rounding to the
|
||
# same 0–255 integer grid first we guarantee that
|
||
# original + residual == corrected
|
||
img_q = (img * 255.0).round().div(255.0)
|
||
|
||
blended = (img_q + res * strength).clamp(0,1)
|
||
|
||
# Restore missing channel dimension if blend lost it (should stay (C,H,W))
|
||
if blended.ndim == 2:
|
||
blended = blended.unsqueeze(0)
|
||
|
||
blended_frames.append(blended.unsqueeze(0)) # keep batch dim so we can cat
|
||
prefixes.append(f"SPCA_{strength:.2f}_")
|
||
|
||
# Concatenate into a single batch tensor (N,C,H,W)
|
||
batch_tensor = torch.cat(blended_frames, dim=0)
|
||
|
||
joined_prefix = "".join(prefixes)
|
||
|
||
return (batch_tensor, joined_prefix)
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
'HDR Effects (SuperBeasts.AI)': HDREffects,
|
||
'Make Resized Mask Batch (SuperBeasts.AI)': MakeResizedMaskBatch,
|
||
'Mask Batch Manager (SuperBeasts.AI)': MaskBatchManagement,
|
||
'Image Batch Manager (SuperBeasts.AI)': ImageBatchManagement,
|
||
'String List Manager (SuperBeasts.AI)': StringListManager,
|
||
'Deflicker - Experimental (SuperBeasts.AI)': Deflicker,
|
||
'Pixel Deflicker - Experimental (SuperBeasts.AI)': PixelDeflicker,
|
||
'Super Pop Color Adjustment (SuperBeasts.AI)': SuperPopColorAdjustment,
|
||
'SB Load Model (SuperBeasts.AI)': SBLoadModel,
|
||
'Super Pop Residual Blend (SuperBeasts.AI)': SuperPopResidualBlend,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
'HDREffects': 'HDR Effects (SuperBeasts.AI)',
|
||
'MakeResizedMaskBatch': 'Make Resized Mask Batch (SuperBeasts.AI)',
|
||
'MaskBatchManagement':'Mask Batch Manager (SuperBeasts.AI)',
|
||
'ImageBatchManagement':'Image Batch Manager (SuperBeasts.AI)',
|
||
'StringListManager': 'String List Manager (SuperBeasts.AI)',
|
||
'Deflicker': 'Deflicker - Experimental (SuperBeasts.AI)',
|
||
'PixelDeflicker': 'Pixel Deflicker - Experimental (SuperBeasts.AI)',
|
||
'SuperPopColorAdjustment': 'Super Pop Color Adjustment (SuperBeasts.AI)',
|
||
'SBLoadModel': 'SB Load Model (SuperBeasts.AI)',
|
||
'SuperPopResidualBlend': 'Super Pop Residual Blend (SuperBeasts.AI)',
|
||
}
|