custom_nodes/plxr_deteriorate/__init__.py

"""PLXR Deteriorate — apply the training-time synthetic degradation to an image.

Wraps the recipe families from the qwen-edit-restore project's deteriorate.py
(bundled as deteriorate_core.py) so a workflow can force-degrade an input
before restoration, for testing / guaranteeing visible deterioration.
"""

import io
import random

import numpy as np
import torch
from PIL import Image

from . import deteriorate_core as core

_REMBG_SESSION = None


def _person_mask_from_array(img01):
    """Soft subject mask via rembg, from a float [H,W,3] array. None on failure."""
    global _REMBG_SESSION
    try:
        from rembg import remove, new_session
        if _REMBG_SESSION is None:
            _REMBG_SESSION = new_session("u2net")
        buf = io.BytesIO()
        Image.fromarray(core.to_uint8(img01)).save(buf, format="PNG")
        out = remove(buf.getvalue(), session=_REMBG_SESSION, only_mask=True)
        m = np.asarray(Image.open(io.BytesIO(out)).convert("L"),
                       dtype=np.float32) / 255.0
        if m.shape != img01.shape[:2]:
            import cv2
            m = cv2.resize(m, (img01.shape[1], img01.shape[0]))
        return m
    except Exception as e:
        print(f"[plxr_deteriorate] person mask failed ({e}); bg-blur ops degrade to global blur")
        return None


class PLXRDeteriorate:
    CATEGORY = "image/plxr"
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "run"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "family": (["random", "atmospheric", "digital"],),
                "severity_min": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.05}),
                "severity_max": ("FLOAT", {"default": 0.85, "min": 0.0, "max": 1.0, "step": 0.05}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 2**32 - 1}),
                "use_person_mask": ("BOOLEAN", {"default": True}),
            }
        }

    def run(self, image, family, severity_min, severity_max, seed, use_person_mask):
        out = []
        for b in range(image.shape[0]):
            x0 = image[b].cpu().numpy().astype(np.float32)  # [H,W,C] 0..1
            x0 = np.clip(x0[..., :3], 0.0, 1.0)

            item_seed = seed + b
            rng_np = np.random.default_rng(item_seed)
            pyrng = random.Random(item_seed ^ 0xABCD)

            class R:
                uniform = staticmethod(pyrng.uniform)
                choice = staticmethod(pyrng.choice)
                random = staticmethod(pyrng.random)
                normal = staticmethod(rng_np.normal)

            lo, hi = sorted((severity_min, severity_max))
            sev = pyrng.uniform(lo, hi)

            fam = family
            if fam == "random":
                fam = pyrng.choice(["atmospheric", "digital"])
            recipe = (core.recipe_atmospheric if fam == "atmospheric"
                      else core.recipe_digital)

            mask = _person_mask_from_array(x0) if use_person_mask else None
            y, label = recipe(x0.copy(), R, sev, mask)
            print(f"[plxr_deteriorate] applied {label} sev={sev:.2f} seed={item_seed}")
            y = np.clip(y, 0.0, 1.0)
            out.append(torch.from_numpy(y.astype(np.float32)))
        return (torch.stack(out).to(image.device),)


NODE_CLASS_MAPPINGS = {"PLXRDeteriorate": PLXRDeteriorate}
NODE_DISPLAY_NAME_MAPPINGS = {"PLXRDeteriorate": "PLXR Deteriorate (restore-lora test)"}