Add complete Eric Image Processing Nodes implementation
Features: - 40+ image processing algorithms - Traditional denoising: Wavelet, NLM, BM3D (GPU-accelerated) - AI models: DnCNN, SwinIR, SCUNet, NAFNet, Noise-DA - DeepInv diffusion denoising integration - Advanced film grain detection and processing - Frequency domain processing - Richardson-Lucy deconvolution - Wiener filter - External DeepInv service with FastAPI - Comprehensive documentation and guides
This commit is contained in:
Eric Rollei
@@ -0,0 +1,443 @@
|
||||
"""DiffBIR processing helpers for ComfyUI integration.
|
||||
|
||||
This module wraps the official DiffBIR reference implementation so that it can be
|
||||
invoked from a lightweight ComfyUI node without duplicating the training codebase.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import shutil
|
||||
import tempfile
|
||||
from argparse import Namespace
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, Optional, Tuple, Type
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
try: # Optional dependency check
|
||||
import torch
|
||||
except ImportError as exc: # pragma: no cover - runtime guard
|
||||
torch = None # type: ignore
|
||||
_TORCH_IMPORT_ERROR = exc
|
||||
else:
|
||||
_TORCH_IMPORT_ERROR = None
|
||||
|
||||
try: # DiffBIR relies on accelerate for deterministic seeding
|
||||
from accelerate.utils import set_seed # type: ignore[import-error]
|
||||
except ImportError as exc: # pragma: no cover - runtime guard
|
||||
set_seed = None # type: ignore
|
||||
_ACCELERATE_IMPORT_ERROR = exc
|
||||
else:
|
||||
_ACCELERATE_IMPORT_ERROR = None
|
||||
|
||||
|
||||
# Defaults copied from transformer_lab/third_party/diffbir/inference.py to avoid
|
||||
# importing face-specific dependencies at module import time.
|
||||
DEFAULT_POS_PROMPT = (
|
||||
"Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, "
|
||||
"hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, "
|
||||
"skin pore detailing, hyper sharpness, perfect without deformations."
|
||||
)
|
||||
|
||||
DEFAULT_NEG_PROMPT = (
|
||||
"painting, oil painting, illustration, drawing, art, sketch, oil painting, cartoon, "
|
||||
"CG Style, 3D render, unreal engine, blurring, dirty, messy, worst quality, low quality, frames, watermark, "
|
||||
"signature, jpeg artifacts, deformed, lowres, over-smooth."
|
||||
)
|
||||
|
||||
|
||||
def _safe_check_device(device: str) -> str:
|
||||
"""Mirror DiffBIR's device selection without pulling heavy dependencies."""
|
||||
|
||||
if torch is None:
|
||||
print("PyTorch not available; defaulting to CPU")
|
||||
return "cpu"
|
||||
|
||||
if device == "cuda":
|
||||
if not torch.cuda.is_available():
|
||||
print(
|
||||
"CUDA not available because the current PyTorch install was not built with CUDA enabled."
|
||||
)
|
||||
device = "cpu"
|
||||
elif device == "mps":
|
||||
mps_backend = getattr(torch.backends, "mps", None)
|
||||
if not mps_backend or not mps_backend.is_available():
|
||||
if mps_backend and not mps_backend.is_built():
|
||||
print(
|
||||
"MPS not available because the current PyTorch install was not built with MPS enabled."
|
||||
)
|
||||
else:
|
||||
print(
|
||||
"MPS not available because the current MacOS version is not 12.3+ and/or the device does not support MPS."
|
||||
)
|
||||
device = "cpu"
|
||||
else:
|
||||
# fall back to cpu for any unexpected value
|
||||
device = "cpu"
|
||||
|
||||
print(f"using device {device}")
|
||||
return device
|
||||
|
||||
|
||||
ROOT_DIR = Path(__file__).resolve().parent.parent
|
||||
TRANSFORMER_LAB_DIR = ROOT_DIR / "transformer_lab"
|
||||
DIFFBIR_REPO_DIR = TRANSFORMER_LAB_DIR / "third_party" / "diffbir"
|
||||
|
||||
# Manual weight cache locations (mirrors bootstrap.py search order)
|
||||
LOCAL_WEIGHT_DIRS: Tuple[Path, ...] = (
|
||||
TRANSFORMER_LAB_DIR / "models" / "pretrained_weights",
|
||||
TRANSFORMER_LAB_DIR / "models",
|
||||
ROOT_DIR / "models" / "pretrained_weights",
|
||||
ROOT_DIR / "models",
|
||||
)
|
||||
|
||||
# Files referenced by diffbir/inference/pretrained_models.py
|
||||
REQUIRED_WEIGHT_FILES: Dict[str, str] = {
|
||||
"DiffBIR_v2.1.pt": "DiffBIR v2.1 (IRControlNet)",
|
||||
"v2.pth": "DiffBIR v2 (IRControlNet)",
|
||||
"v1_general.pth": "DiffBIR v1 general (IRControlNet)",
|
||||
"v1_face.pth": "DiffBIR v1 face (IRControlNet)",
|
||||
"realesrgan_s4_swinir_100k.pth": "SwinIR RealESRGAN cleaner",
|
||||
"general_swinir_v1.ckpt": "SwinIR general cleaner",
|
||||
"face_swinir_v1.ckpt": "SwinIR face cleaner",
|
||||
"BSRNet.pth": "BSRNet cleaner",
|
||||
"scunet_color_real_psnr.pth": "SCUNet cleaner",
|
||||
"codeformer_swinir.ckpt": "CodeFormer SwinIR cleaner",
|
||||
"v2-1_512-ema-pruned.ckpt": "Stable Diffusion 2.1 base",
|
||||
"sd2.1-base-zsnr-laionaes5.ckpt": "Stable Diffusion 2.1 ZSNR variant",
|
||||
}
|
||||
|
||||
# Subset mapping for per-version requirements
|
||||
VERSION_SPECIFIC_FILES: Dict[str, Iterable[str]] = {
|
||||
"v1": (
|
||||
"v1_general.pth",
|
||||
"v1_face.pth",
|
||||
"general_swinir_v1.ckpt",
|
||||
"face_swinir_v1.ckpt",
|
||||
"v2-1_512-ema-pruned.ckpt",
|
||||
),
|
||||
"v2": (
|
||||
"v2.pth",
|
||||
"bsrnet.pth",
|
||||
"scunet_color_real_psnr.pth",
|
||||
"realesrgan_s4_swinir_100k.pth",
|
||||
"general_swinir_v1.ckpt",
|
||||
"codeformer_swinir.ckpt",
|
||||
"v2-1_512-ema-pruned.ckpt",
|
||||
),
|
||||
"v2.1": (
|
||||
"DiffBIR_v2.1.pt",
|
||||
"realesrgan_s4_swinir_100k.pth",
|
||||
"general_swinir_v1.ckpt",
|
||||
"face_swinir_v1.ckpt",
|
||||
"scunet_color_real_psnr.pth",
|
||||
"codeformer_swinir.ckpt",
|
||||
"sd2.1-base-zsnr-laionaes5.ckpt",
|
||||
),
|
||||
}
|
||||
|
||||
# Version mapping above contains a typo (bsrnet casing). Normalise keys.
|
||||
VERSION_SPECIFIC_FILES["v2"] = tuple(
|
||||
"BSRNet.pth" if name.lower() == "bsrnet.pth" else name
|
||||
for name in VERSION_SPECIFIC_FILES["v2"]
|
||||
)
|
||||
|
||||
# Additional cleaner dependencies shared across tasks
|
||||
SHARED_CLEANER_FILES = {
|
||||
"BSRNet.pth",
|
||||
"scunet_color_real_psnr.pth",
|
||||
"realesrgan_s4_swinir_100k.pth",
|
||||
"general_swinir_v1.ckpt",
|
||||
"face_swinir_v1.ckpt",
|
||||
"codeformer_swinir.ckpt",
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class DiffBIRConfig:
|
||||
"""Configuration payload controlling the DiffBIR inference call."""
|
||||
|
||||
task: str = "sr"
|
||||
version: str = "v2.1"
|
||||
sampler: str = "edm_dpm++_3m_sde"
|
||||
steps: int = 12
|
||||
upscale: float = 4.0
|
||||
captioner: str = "none"
|
||||
cfg_scale: float = 6.0
|
||||
noise_aug: int = 0
|
||||
start_point_type: str = "noise"
|
||||
rescale_cfg: bool = False
|
||||
strength: float = 1.0
|
||||
cleaner_tiled: bool = False
|
||||
cleaner_tile_size: int = 512
|
||||
cleaner_tile_stride: int = 256
|
||||
vae_encoder_tiled: bool = False
|
||||
vae_encoder_tile_size: int = 256
|
||||
vae_decoder_tiled: bool = False
|
||||
vae_decoder_tile_size: int = 256
|
||||
cldm_tiled: bool = False
|
||||
cldm_tile_size: int = 512
|
||||
cldm_tile_stride: int = 256
|
||||
guidance: bool = False
|
||||
guidance_loss: str = "w_mse"
|
||||
guidance_scale: float = 0.0
|
||||
guidance_start: float = 0.0
|
||||
guidance_stop: float = 1.0
|
||||
guidance_space: str = "rgb"
|
||||
guidance_repeat: int = 1
|
||||
batch_size: int = 1
|
||||
n_samples: int = 1
|
||||
seed: int = 231
|
||||
device_preference: str = "auto"
|
||||
precision: str = "fp16"
|
||||
llava_bit: str = "4"
|
||||
s_churn: float = 0.0
|
||||
s_tmin: float = 0.0
|
||||
s_tmax: float = 300.0
|
||||
s_noise: float = 1.0
|
||||
eta: float = 1.0
|
||||
order: int = 1
|
||||
pos_prompt: Optional[str] = None
|
||||
neg_prompt: Optional[str] = None
|
||||
|
||||
|
||||
class DiffBIRProcessor:
|
||||
"""Thin wrapper that executes the official DiffBIR inference loops."""
|
||||
|
||||
def __init__(self, config: DiffBIRConfig):
|
||||
self.config = config
|
||||
if _TORCH_IMPORT_ERROR is not None:
|
||||
raise RuntimeError(
|
||||
"PyTorch is required for DiffBIR inference"
|
||||
) from _TORCH_IMPORT_ERROR
|
||||
if _ACCELERATE_IMPORT_ERROR is not None:
|
||||
raise RuntimeError(
|
||||
"`accelerate` is required for DiffBIR inference (set_seed)"
|
||||
) from _ACCELERATE_IMPORT_ERROR
|
||||
if not DIFFBIR_REPO_DIR.exists():
|
||||
raise FileNotFoundError(
|
||||
"DiffBIR repository not found. Run `python transformer_lab\\bootstrap_models.py diffbir` first."
|
||||
)
|
||||
|
||||
# Ensure the repo root is importable
|
||||
if str(DIFFBIR_REPO_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(DIFFBIR_REPO_DIR))
|
||||
|
||||
# Copy weights into the expected `weights/` directory inside the repo
|
||||
self._prepare_weights()
|
||||
|
||||
# Late imports (after sys.path update)
|
||||
from diffbir.inference.bid_loop import BIDInferenceLoop # type: ignore
|
||||
from diffbir.inference.bsr_loop import BSRInferenceLoop # type: ignore
|
||||
|
||||
self._DEFAULT_POS_PROMPT = DEFAULT_POS_PROMPT
|
||||
self._DEFAULT_NEG_PROMPT = DEFAULT_NEG_PROMPT
|
||||
self._check_device = _safe_check_device
|
||||
|
||||
self._loops: Dict[str, Optional[Type[Any]]] = {
|
||||
"sr": BSRInferenceLoop,
|
||||
"denoise": BIDInferenceLoop,
|
||||
}
|
||||
|
||||
self._face_dependency_error: Optional[BaseException] = None
|
||||
|
||||
try:
|
||||
from diffbir.inference.bfr_loop import BFRInferenceLoop # type: ignore
|
||||
except ModuleNotFoundError as exc:
|
||||
if exc.name == "facexlib":
|
||||
self._face_dependency_error = exc
|
||||
self._loops["face"] = None
|
||||
else:
|
||||
raise
|
||||
else:
|
||||
self._loops["face"] = BFRInferenceLoop
|
||||
|
||||
try:
|
||||
from diffbir.inference.unaligned_bfr_loop import UnAlignedBFRInferenceLoop # type: ignore
|
||||
except ModuleNotFoundError as exc:
|
||||
if exc.name == "facexlib":
|
||||
self._face_dependency_error = exc
|
||||
self._loops["unaligned_face"] = None
|
||||
else:
|
||||
raise
|
||||
else:
|
||||
self._loops["unaligned_face"] = UnAlignedBFRInferenceLoop
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ------------------------------------------------------------------
|
||||
def _prepare_weights(self) -> None:
|
||||
weights_dir = DIFFBIR_REPO_DIR / "weights"
|
||||
weights_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
required = set(REQUIRED_WEIGHT_FILES.keys())
|
||||
# Include version-specific list if defined
|
||||
version = self.config.version
|
||||
if version in VERSION_SPECIFIC_FILES:
|
||||
required |= set(VERSION_SPECIFIC_FILES[version])
|
||||
# cleaners are shared across tasks; keep them available to avoid lazy downloads
|
||||
required |= SHARED_CLEANER_FILES
|
||||
|
||||
missing: list[str] = []
|
||||
for filename in sorted(required):
|
||||
src = self._find_local_weight(filename)
|
||||
if src is None:
|
||||
missing.append(filename)
|
||||
continue
|
||||
dst = weights_dir / filename
|
||||
try:
|
||||
if not dst.exists() or src.stat().st_mtime > dst.stat().st_mtime:
|
||||
shutil.copy2(src, dst)
|
||||
except PermissionError:
|
||||
# Fall back to copying via temporary file if the destination is locked
|
||||
tmp_path = dst.with_suffix(dst.suffix + ".tmp")
|
||||
shutil.copy2(src, tmp_path)
|
||||
tmp_path.replace(dst)
|
||||
if missing:
|
||||
pretty = ", ".join(missing)
|
||||
raise FileNotFoundError(
|
||||
"The following DiffBIR weights are missing. Place them under "
|
||||
"`transformer_lab/models/pretrained_weights` and retry: " + pretty
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _find_local_weight(filename: str) -> Optional[Path]:
|
||||
for directory in LOCAL_WEIGHT_DIRS:
|
||||
candidate = directory / filename
|
||||
if candidate.exists():
|
||||
return candidate
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
def process_image(self, image: np.ndarray) -> np.ndarray:
|
||||
if image.ndim != 3 or image.shape[2] not in (3, 4):
|
||||
raise ValueError("DiffBIR expects an HWC image with 3 or 4 channels")
|
||||
|
||||
# Drop alpha channel if present
|
||||
if image.shape[2] == 4:
|
||||
image = image[..., :3]
|
||||
|
||||
image_uint8 = np.clip(image, 0, 255).astype(np.uint8)
|
||||
|
||||
temp_root = TRANSFORMER_LAB_DIR / "temp"
|
||||
temp_root.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="diffbir_", dir=str(temp_root)) as workdir:
|
||||
workdir_path = Path(workdir)
|
||||
input_dir = workdir_path / "input"
|
||||
output_dir = workdir_path / "output"
|
||||
input_dir.mkdir(parents=True, exist_ok=True)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
input_path = input_dir / "input.png"
|
||||
Image.fromarray(image_uint8).save(input_path)
|
||||
|
||||
args = self._build_args(
|
||||
input_dir=input_dir,
|
||||
output_dir=output_dir,
|
||||
file_stem=input_path.stem,
|
||||
)
|
||||
|
||||
args.device = self._check_device(args.device)
|
||||
if set_seed is not None:
|
||||
set_seed(args.seed)
|
||||
|
||||
loop_cls = self._loops.get(args.task)
|
||||
if loop_cls is None:
|
||||
raise ValueError(f"Unsupported DiffBIR task: {args.task}")
|
||||
|
||||
loop = loop_cls(args)
|
||||
loop.run()
|
||||
|
||||
output_file = output_dir / f"{input_path.stem}.png"
|
||||
if not output_file.exists():
|
||||
raise FileNotFoundError(
|
||||
f"DiffBIR did not produce an output image at {output_file}"
|
||||
)
|
||||
|
||||
result = Image.open(output_file).convert("RGB")
|
||||
return np.array(result)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal
|
||||
# ------------------------------------------------------------------
|
||||
def _resolve_prompts(self) -> Tuple[str, str]:
|
||||
pos = self.config.pos_prompt if self.config.pos_prompt is not None else self._DEFAULT_POS_PROMPT
|
||||
neg = self.config.neg_prompt if self.config.neg_prompt is not None else self._DEFAULT_NEG_PROMPT
|
||||
return pos, neg
|
||||
|
||||
def _resolve_device(self) -> str:
|
||||
pref = (self.config.device_preference or "auto").lower()
|
||||
if pref == "auto":
|
||||
if torch is not None:
|
||||
if torch.cuda.is_available():
|
||||
return "cuda"
|
||||
mps_backend = getattr(torch.backends, "mps", None)
|
||||
if mps_backend and mps_backend.is_available():
|
||||
return "mps"
|
||||
return "cpu"
|
||||
return pref
|
||||
|
||||
def _build_args(self, *, input_dir: Path, output_dir: Path, file_stem: str) -> Namespace:
|
||||
pos_prompt, neg_prompt = self._resolve_prompts()
|
||||
return Namespace(
|
||||
task=self.config.task,
|
||||
upscale=float(self.config.upscale),
|
||||
version=self.config.version,
|
||||
train_cfg="",
|
||||
ckpt="",
|
||||
sampler=self.config.sampler,
|
||||
steps=int(self.config.steps),
|
||||
start_point_type=self.config.start_point_type,
|
||||
cleaner_tiled=bool(self.config.cleaner_tiled),
|
||||
cleaner_tile_size=int(self.config.cleaner_tile_size),
|
||||
cleaner_tile_stride=int(self.config.cleaner_tile_stride),
|
||||
vae_encoder_tiled=bool(self.config.vae_encoder_tiled),
|
||||
vae_encoder_tile_size=int(self.config.vae_encoder_tile_size),
|
||||
vae_decoder_tiled=bool(self.config.vae_decoder_tiled),
|
||||
vae_decoder_tile_size=int(self.config.vae_decoder_tile_size),
|
||||
cldm_tiled=bool(self.config.cldm_tiled),
|
||||
cldm_tile_size=int(self.config.cldm_tile_size),
|
||||
cldm_tile_stride=int(self.config.cldm_tile_stride),
|
||||
captioner=self.config.captioner,
|
||||
pos_prompt=pos_prompt,
|
||||
neg_prompt=neg_prompt,
|
||||
cfg_scale=float(self.config.cfg_scale),
|
||||
rescale_cfg=bool(self.config.rescale_cfg),
|
||||
noise_aug=int(self.config.noise_aug),
|
||||
s_churn=float(self.config.s_churn),
|
||||
s_tmin=float(self.config.s_tmin),
|
||||
s_tmax=float(self.config.s_tmax),
|
||||
s_noise=float(self.config.s_noise),
|
||||
eta=float(self.config.eta),
|
||||
order=int(self.config.order),
|
||||
strength=float(self.config.strength),
|
||||
batch_size=int(self.config.batch_size),
|
||||
guidance=bool(self.config.guidance),
|
||||
g_loss=self.config.guidance_loss,
|
||||
g_scale=float(self.config.guidance_scale),
|
||||
g_start=float(self.config.guidance_start),
|
||||
g_stop=float(self.config.guidance_stop),
|
||||
g_space=self.config.guidance_space,
|
||||
g_repeat=int(self.config.guidance_repeat),
|
||||
input=str(input_dir),
|
||||
n_samples=int(self.config.n_samples),
|
||||
output=str(output_dir),
|
||||
seed=int(self.config.seed),
|
||||
device=self._resolve_device(),
|
||||
precision=self.config.precision,
|
||||
llava_bit=str(self.config.llava_bit),
|
||||
file_stem=file_stem,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DiffBIRConfig",
|
||||
"DiffBIRProcessor",
|
||||
]
|
||||
Reference in New Issue
Block a user