diffbir_processing.py

"""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",
]