flow_upscaler_pipeline.py

"""Latent-space 2x upscaling with FlowUpscaler (single-step rectified flow).

Diffusers-format pipeline. The components are the FlowUpscaler UNet
(`unet/upscaler_unet.py`, weights in `unet/diffusion_pytorch_model.safetensors`),
a shared Flux.2 VAE, and a `FlowMatchEulerDiscreteScheduler`. The VAE is not
bundled with this repo (it is shared with any Flux.2 checkpoint), so pass it
explicitly:

    vae = AutoencoderKLFlux2.from_pretrained(
        "black-forest-labs/FLUX.2-dev", subfolder="vae", torch_dtype=torch.bfloat16
    )
    pipeline = DiffusionPipeline.from_pretrained(
        "MinhNH232331M/FlowUpscaler-diffusers",
        vae=vae, torch_dtype=torch.bfloat16, trust_remote_code=True,
    ).to("cuda")
    image = pipeline(image, num_passes=2).images[0]

Inference convention replicated from the reference ComfyUI node
(github.com/TensorForger/comfyui-flow-upscaler) and the training notebooks in
github.com/tensorforger/CTGMWorkshop (notebooks/flow_upscaler): the UNet
consumes *unpatchified* 32-channel Flux.2 latents normalized with the VAE's
BatchNorm running stats, and denoises pure noise into the 2x latent in a
single FlowMatchEuler step. Passes can be chained for 4x/8x — the output of a
pass lives in the same normalized latent space as its conditioning input.

On top of the reference convention, two color corrections are applied: each
pass ends with a low-frequency transplant from the conditioning latent
(`_lowfreq_transplant`, fixes per-channel mean drift) and the decoded image
gets a low-band L/a/b transplant from the input (`_lab_stat_match`, fixes
the ~8%/pass chroma amplification and the tone-curve stretch the latent
transplant cannot see).

This file is executed standalone by diffusers' remote-code loader, so it must
stay at the repo root and must not import sibling modules at the top level
(the UNet class arrives as an instantiated component; `from_single_file` loads
its module by explicit path).
"""

import importlib.util
import os

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from safetensors.torch import load_file

from diffusers import (
    AutoencoderKLFlux2,
    DiffusionPipeline,
    FlowMatchEulerDiscreteScheduler,
    ImagePipelineOutput,
    ModelMixin,
)
from diffusers.utils import logging
from diffusers.utils.torch_utils import randn_tensor

logger = logging.get_logger(__name__)


def patchify_latents(latents: torch.Tensor) -> torch.Tensor:
    """(B, C, H, W) -> (B, 4C, H/2, W/2) via 2x2 space-to-channel."""
    batch_size, num_channels, height, width = latents.shape
    latents = latents.view(batch_size, num_channels, height // 2, 2, width // 2, 2)
    latents = latents.permute(0, 1, 3, 5, 2, 4)
    return latents.reshape(batch_size, num_channels * 4, height // 2, width // 2)


def unpatchify_latents(latents: torch.Tensor) -> torch.Tensor:
    """(B, 4C, H, W) -> (B, C, 2H, 2W), inverse of :func:`patchify_latents`."""
    batch_size, num_channels, height, width = latents.shape
    latents = latents.reshape(batch_size, num_channels // 4, 2, 2, height, width)
    latents = latents.permute(0, 1, 4, 2, 5, 3)
    return latents.reshape(batch_size, num_channels // 4, height * 2, width * 2)


# sRGB <-> Lab (D65), cv2 float conventions: L in [0, 100], a/b around 0.
_RGB2XYZ = torch.tensor(
    [[0.412453, 0.357580, 0.180423],
     [0.212671, 0.715160, 0.072169],
     [0.019334, 0.119193, 0.950227]]
)
_XYZ2RGB = torch.linalg.inv(_RGB2XYZ)
_LAB_WHITE = torch.tensor([0.950456, 1.0, 1.088754])


def _rgb_to_lab(rgb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """(B, 3, H, W) sRGB in [0, 1] -> (L, a, b) planes."""
    lin = torch.where(rgb <= 0.04045, rgb / 12.92, ((rgb + 0.055) / 1.055) ** 2.4)
    xyz = torch.einsum("ij,bjhw->bihw", _RGB2XYZ.to(rgb.device, rgb.dtype), lin)
    xyz = xyz / _LAB_WHITE.to(rgb.device, rgb.dtype).view(1, 3, 1, 1)
    f = torch.where(
        xyz > 0.008856, xyz.clamp(min=0.0) ** (1.0 / 3.0), 7.787 * xyz + 16.0 / 116.0
    )
    fx, fy, fz = f[:, 0], f[:, 1], f[:, 2]
    return 116.0 * fy - 16.0, 500.0 * (fx - fy), 200.0 * (fy - fz)


def _lab_to_rgb(lightness: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """Inverse of :func:`_rgb_to_lab`, output clamped to [0, 1]."""
    fy = (lightness + 16.0) / 116.0
    f = torch.stack([fy + a / 500.0, fy, fy - b / 200.0], dim=1)
    xyz = torch.where(f > 0.206893, f**3, (f - 16.0 / 116.0) / 7.787)
    xyz = xyz * _LAB_WHITE.to(f.device, f.dtype).view(1, 3, 1, 1)
    lin = torch.einsum("ij,bjhw->bihw", _XYZ2RGB.to(f.device, f.dtype), xyz)
    return torch.where(
        lin <= 0.0031308, lin * 12.92, 1.055 * lin.clamp(min=0.0) ** (1.0 / 2.4) - 0.055
    ).clamp(0.0, 1.0)


class FlowUpscalerPipeline(DiffusionPipeline):
    r"""Image-in / image-out 2x latent upscaling with the FlowUpscaler UNet.

    Each pass denoises pure noise into the 2x latent in a single
    FlowMatchEuler step, conditioned on the input latent; passes chain for
    4x/8x. The UNet defaults to float32 (the training dtype); bf16 matches it
    to ~54 dB PSNR while halving the forward peak and runtime (pass
    `torch_dtype=torch.bfloat16` to `from_pretrained`). The shared VAE is
    used in its own dtype.

    Args:
        unet (`UpscalerUNet`):
            Attention-free flow-matching UNet predicting velocity in the
            bn-normalized, unpatchified Flux.2 latent space
            (`unet/upscaler_unet.py` in this repo).
        vae (`AutoencoderKLFlux2`):
            A Flux.2 VAE (or a drop-in replacement operating in the same
            latent space, e.g. MageFlow). Its BatchNorm running stats define
            the latent normalization.
        scheduler (`FlowMatchEulerDiscreteScheduler`):
            Scheduler for the single Euler step of each pass.
    """

    model_cpu_offload_seq = "unet->vae"

    # Encode/decode images larger than this (pixels per side) with VAE tiling
    # to keep peak activation memory bounded on 24GB-class GPUs.
    TILED_VAE_THRESHOLD = 2048

    # Hard ceiling on the output side. Peak memory scales with latent *area*;
    # with the chunked inference executor (unet/upscaler_unet.py) the bf16
    # UNet pass measured 2.5GiB at an 8K-UHD output and 7.6GiB at 16K-UHD
    # (decode 5.5GiB), so even 16K fits a 23GB L4 alongside resident
    # pipelines. fp32 needs roughly double.
    MAX_OUTPUT_SIDE = 16384

    def __init__(
        self,
        unet: ModelMixin,  # an UpscalerUNet, loaded dynamically from unet/upscaler_unet.py
        vae: AutoencoderKLFlux2,
        scheduler: FlowMatchEulerDiscreteScheduler,
    ):
        super().__init__()
        self.register_modules(unet=unet, vae=vae, scheduler=scheduler)

    @classmethod
    def from_single_file(
        cls,
        model_path: str,
        vae,
        device: str = "cuda",
        dtype: torch.dtype = torch.float32,
    ) -> "FlowUpscalerPipeline":
        """Build the pipeline from the flat `flow_upscaler.safetensors` checkpoint.

        Mirrors the pre-diffusers constructor
        ``FlowUpscalerPipeline(model_path, vae, device, dtype)``. Needs
        `unet/upscaler_unet.py` (or `upscaler_unet.py`) next to this file —
        i.e. a local checkout of the repo; from the Hub use `from_pretrained`.
        """
        here = os.path.dirname(os.path.abspath(__file__))
        candidates = [
            os.path.join(here, "unet", "upscaler_unet.py"),
            os.path.join(here, "upscaler_unet.py"),
        ]
        module_path = next((p for p in candidates if os.path.isfile(p)), None)
        if module_path is None:
            raise FileNotFoundError(
                "from_single_file requires unet/upscaler_unet.py next to "
                f"{__file__}; use FlowUpscalerPipeline.from_pretrained(...) instead."
            )
        spec = importlib.util.spec_from_file_location("upscaler_unet", module_path)
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)

        unet = module.UpscalerUNet()
        unet.load_state_dict(load_file(model_path))
        unet.to(device, dtype).eval()
        return cls(unet=unet, vae=vae, scheduler=FlowMatchEulerDiscreteScheduler())

    def _bn_stats(self) -> tuple[torch.Tensor, torch.Tensor]:
        mean = self.vae.bn.running_mean.view(1, -1, 1, 1)
        std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps)
        device, dtype = self._execution_device, self.unet.dtype
        return mean.to(device, dtype), std.to(device, dtype)

    def normalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
        """Raw VAE latents (B, 32, H, W) -> bn-normalized, still unpatchified."""
        mean, std = self._bn_stats()
        latents = patchify_latents(latents.to(self._execution_device, self.unet.dtype))
        return unpatchify_latents((latents - mean) / std)

    def denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
        mean, std = self._bn_stats()
        latents = patchify_latents(latents)
        return unpatchify_latents(latents * std + mean)

    # Low-pass cutoff for the color fix: the low band is what survives an
    # antialiased downsample of the output latent by this factor. 8 pins the
    # decoded channel means to <1/255 (4 left ~1.5/255 residue in tests).
    LOWFREQ_FACTOR = 8

    def _lowfreq_transplant(
        self, latents: torch.Tensor, cond: torch.Tensor, strength: float = 1.0
    ) -> torch.Tensor:
        """Replace the low band of `latents` with the conditioning's, in place.

        `strength` interpolates between the latents' own low band (0.0 — no
        transplant: softer output, truest to the input's structure) and the
        conditioning's (1.0, default — on-manifold conditioning for the next
        pass, visibly crisper). Color stays correct at any strength; the
        pixel-space stage handles that.

        The single-step flow drifts per-channel latent means by ~0.1 sigma per
        pass, which compounds over chained passes into a visible color shift.
        The conditioning latent is ground truth for everything below its
        Nyquist, so swapping the low band in pins global and regional color
        while keeping the UNet's detail — and hands the next pass an
        on-manifold conditioning (measurably sharper output). Low bands come
        from antialiased downsamples, so the full-res upsampled conditioning
        is never materialized (+33 MiB peak at an 8K-UHD output).
        """
        if strength <= 0.0:
            return latents
        height, width = latents.shape[-2:]
        small = (max(1, height // self.LOWFREQ_FACTOR), max(1, width // self.LOWFREQ_FACTOR))

        def low_band(x: torch.Tensor) -> torch.Tensor:
            x = F.interpolate(x, size=small, mode="bilinear", antialias=True)
            return F.interpolate(x, size=(height, width), mode="bilinear")

        diff = low_band(cond).sub_(low_band(latents))
        return latents.add_(diff, alpha=float(strength))

    @torch.no_grad()
    def upscale_latents(
        self,
        latents_small: torch.Tensor,
        generator: torch.Generator | None = None,
        color_fix: bool = True,
        transplant_strength: float = 1.0,
    ) -> torch.Tensor:
        """One 2x pass in normalized latent space: (B, 32, H, W) -> (B, 32, 2H, 2W)."""
        batch_size, _, height, width = latents_small.shape
        device, dtype = self._execution_device, self.unet.dtype
        latents_small = latents_small.to(device, dtype)

        # set_timesteps resets the scheduler's step state, so chained passes
        # each run a fresh single-step schedule.
        self.scheduler.set_timesteps(1, device=device, mu=1.0)

        latents = randn_tensor(
            (batch_size, self.unet.config.sample_channels, height * 2, width * 2),
            generator=generator,
            device=device,
            dtype=dtype,
        )
        for t in self.scheduler.timesteps:
            t = t.view(1)
            velocity = self.unet(sample=latents, timestep=t, latents_small=latents_small)
            latents = self.scheduler.step(velocity, t, latents).prev_sample
        if color_fix:
            latents = self._lowfreq_transplant(latents, latents_small, strength=transplant_strength)
        return latents

    @torch.no_grad()
    def encode_image(self, image: Image.Image) -> torch.Tensor:
        """PIL image -> normalized unpatchified latents (1, 32, H/8, W/8).

        The image is resized to a multiple of 16 so the /8 latent stays
        patchifiable (this mirrors how training conditioning latents were made:
        decode -> downscale in pixel space -> encode).
        """
        image = image.convert("RGB")
        width = max(16, round(image.width / 16) * 16)
        height = max(16, round(image.height / 16) * 16)
        if (width, height) != image.size:
            image = image.resize((width, height), Image.LANCZOS)

        pixels = torch.from_numpy(np.array(image)).float().div(127.5).sub(1.0)
        pixels = pixels.permute(2, 0, 1).unsqueeze(0).to(self._execution_device, self.vae.dtype)
        previous_tiling = self.vae.use_tiling
        if max(image.size) > self.TILED_VAE_THRESHOLD:
            self.vae.use_tiling = True
        try:
            latents = self.vae.encode(pixels).latent_dist.mode()
        finally:
            self.vae.use_tiling = previous_tiling
        return self.normalize_latents(latents)

    @torch.no_grad()
    def _decode_to_tensor(self, latents: torch.Tensor) -> torch.Tensor:
        """Normalized latents -> (1, 3, H, W) pixels in [0, 1], still on GPU."""
        latents = self.denormalize_latents(latents).to(self.vae.dtype)
        output_side = max(latents.shape[-2:]) * 8
        previous_tiling = self.vae.use_tiling
        if output_side > self.TILED_VAE_THRESHOLD:
            self.vae.use_tiling = True
        try:
            decoded = self.vae.decode(latents, return_dict=False)[0]
        finally:
            self.vae.use_tiling = previous_tiling
        return decoded.add_(1.0).div_(2.0).clamp_(0.0, 1.0)

    @staticmethod
    def _tensor_to_pil(pixels: torch.Tensor) -> Image.Image:
        # Quantize on-device: the download is uint8, a third of the float
        # traffic the old CPU-side conversion paid.
        array = pixels[0].permute(1, 2, 0).float().mul(255.0).round().to(torch.uint8)
        return Image.fromarray(array.cpu().numpy())

    @torch.no_grad()
    def decode_latents(self, latents: torch.Tensor) -> Image.Image:
        return self._tensor_to_pil(self._decode_to_tensor(latents))

    # Row-strip height for the color-fix Lab math. The ops are elementwise
    # (no halo), so strips exist purely to bound the fp16 working set.
    LAB_STRIP_ROWS = 256

    @torch.no_grad()
    def _lab_stat_match(
        self, image: torch.Tensor, reference: torch.Tensor, strength: float = 1.0
    ) -> torch.Tensor:
        """Pin decoded pixels' regional color and tone to the reference's.

        `strength` linearly interpolates the low band between the image's own
        (0.0) and the reference's (1.0, default): drift metrics scale with
        1 - strength, so intermediate values trade source fidelity for the
        model's punchier rendition.

        Two spread errors survive the latent transplant (a mean fix): the
        flow UNet amplifies chroma ~8% per pass concentrated on already
        colorful regions (a global gain pinned mean chroma but left dC_p99
        +14 on lips/skin after 3 passes), and the latent transplant itself
        stretches the tone curve (raw UNet output is slightly *flat*; the
        transplant's sharpening side effect crushed shadows dp5 -2.75 and
        pushed highlights dp95 +2 after 3 passes, which a global L match
        only half-fixed). Below the reference's Nyquist both are ground
        truth, so the full L/a/b low band is corrected exactly: diff maps
        computed at the reference's resolution (area-downsampled proxy vs
        reference, in fp32 — full-res stats vs an upsampled reference would
        be blur-biased), bilinearly upsampled and added to the full-res
        planes. Measured: chroma p95 +4.9 -> +0.0, p99 +13.8 -> +0.6; tone
        dp5/dp95 -> 0.00 exactly; no halos at the strongest edges, detail
        above the reference Nyquist untouched. Full-res strips run in fp16 —
        within 1/255 of fp32 output; bf16's 8-bit mantissa steps exceed
        uint8 resolution and band on gradients.

        `image`: (1, 3, H, W) in [0, 1]; `reference`: same layout at its own
        smaller size. Returns (H, W, 3) uint8 on the same device.
        """
        proxy = F.interpolate(image.float(), size=reference.shape[-2:], mode="area")
        diff = torch.stack(
            [r - p for r, p in zip(_rgb_to_lab(reference.float()), _rgb_to_lab(proxy))],
            dim=1,
        )
        del proxy
        if strength != 1.0:
            diff = diff * strength

        height, width = image.shape[-2:]
        diff = F.interpolate(diff.to(torch.float16), size=(height, width), mode="bilinear")
        out = torch.empty(height, width, 3, dtype=torch.uint8, device=image.device)
        for r0 in range(0, height, self.LAB_STRIP_ROWS):
            r1 = min(r0 + self.LAB_STRIP_ROWS, height)
            l, a, b = _rgb_to_lab(image[:, :, r0:r1].to(torch.float16))
            rgb = _lab_to_rgb(
                l + diff[:, 0, r0:r1], a + diff[:, 1, r0:r1], b + diff[:, 2, r0:r1]
            )
            out[r0:r1] = rgb[0].permute(1, 2, 0).float().mul_(255.0).round_().to(torch.uint8)
        return out

    @torch.no_grad()
    def __call__(
        self,
        image: Image.Image,
        num_passes: int = 1,
        generator: torch.Generator | None = None,
        color_fix: bool = True,
        transplant_strength: float = 1.0,
        color_fix_strength: float = 1.0,
        output_type: str = "pil",
        return_dict: bool = True,
    ) -> ImagePipelineOutput | tuple:
        r"""Upscale `image` by 2x per pass.

        Args:
            image (`PIL.Image.Image`):
                Input image. Resized to a multiple of 16 before encoding.
            num_passes (`int`, defaults to 1):
                Number of chained 2x passes (2 -> 4x, 3 -> 8x). Reduced
                automatically (with a warning) if the output would exceed
                `MAX_OUTPUT_SIDE` pixels per side.
            generator (`torch.Generator`, *optional*):
                RNG for the per-pass noise. Use a generator on the pipeline's
                device for reproducible results.
            color_fix (`bool`, defaults to `True`):
                Master switch for both drift corrections (latent low-band
                transplant each pass + pixel-space L/a/b low-band transplant
                from the input after decoding).
            transplant_strength (`float`, defaults to 1.0):
                Latent-transplant stage; also a fidelity/realism control
                (lower = truer to input, higher = crisper).
            color_fix_strength (`float`, defaults to 1.0):
                Pixel-space stage; residual color drift scales with
                `1 - strength`.
            output_type (`str`, defaults to `"pil"`):
                `"pil"`, `"np"` (float32 in [0, 1], NHWC) or `"pt"` (float32
                in [0, 1], NCHW, on the pipeline's device).
            return_dict (`bool`, defaults to `True`):
                Return an `ImagePipelineOutput` instead of a plain tuple.

        Returns:
            [`~pipelines.ImagePipelineOutput`] or `tuple`: the upscaled image.
        """
        if output_type not in ("pil", "np", "pt"):
            raise ValueError(f"`output_type` must be 'pil', 'np' or 'pt', got {output_type!r}.")

        # Reduce passes if the result would exceed MAX_OUTPUT_SIDE (e.g. a
        # 2048px generation with 2 passes requested runs only one). Checked
        # before encoding so oversized inputs are rejected without GPU work.
        requested = max(1, int(num_passes))
        num_passes = requested
        side = max(round(image.width / 16), round(image.height / 16)) * 16
        while num_passes > 0 and side * (2 ** num_passes) > self.MAX_OUTPUT_SIDE:
            num_passes -= 1
        if num_passes == 0:
            raise ValueError(
                f"Input of {image.size} cannot be upscaled within the "
                f"{self.MAX_OUTPUT_SIDE}px output limit — downscale it first."
            )
        if num_passes < requested:
            logger.warning(
                "Reduced upscale passes %d -> %d to keep output within %dpx",
                requested, num_passes, self.MAX_OUTPUT_SIDE,
            )

        latents = self.encode_image(image)
        for _ in range(num_passes):
            latents = self.upscale_latents(
                latents, generator=generator, color_fix=color_fix,
                transplant_strength=transplant_strength,
            )
        decoded = self._decode_to_tensor(latents)

        if color_fix:
            reference = torch.from_numpy(np.asarray(image.convert("RGB")).copy())
            reference = reference.to(self._execution_device).permute(2, 0, 1)[None].float().div_(255.0)
            matched = self._lab_stat_match(decoded, reference, strength=float(color_fix_strength))
            if output_type == "pil":
                images = [Image.fromarray(matched.cpu().numpy())]
            elif output_type == "np":
                images = matched.float().div_(255.0)[None].cpu().numpy()
            else:
                images = matched.permute(2, 0, 1)[None].float().div_(255.0)
        else:
            if output_type == "pil":
                images = [self._tensor_to_pil(decoded)]
            elif output_type == "np":
                images = decoded.permute(0, 2, 3, 1).float().cpu().numpy()
            else:
                images = decoded.float()

        if not return_dict:
            return (images,)
        return ImagePipelineOutput(images=images)

    @torch.no_grad()
    def upscale_image(
        self,
        image: Image.Image,
        num_passes: int = 1,
        seed: int = 42,
        color_fix: bool = True,
        transplant_strength: float = 1.0,
        color_fix_strength: float = 1.0,
    ) -> Image.Image:
        """Seed-based convenience wrapper around `__call__` returning a PIL image.

        Kept API-compatible with the pre-diffusers pipeline: the generator is
        created on the execution device, so a given seed reproduces the same
        output as before.
        """
        generator = torch.Generator(device=self._execution_device).manual_seed(seed)
        return self(
            image,
            num_passes=num_passes,
            generator=generator,
            color_fix=color_fix,
            transplant_strength=transplant_strength,
            color_fix_strength=color_fix_strength,
        ).images[0]