Eric Rollei
d6dd6470cf
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
444 lines
16 KiB
Python
444 lines
16 KiB
Python
"""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",
|
|
]
|