flow_upscaler_pipeline.py

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

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).
"""

import logging

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 FlowMatchEulerDiscreteScheduler

from upscaler_unet import UpscalerUNet

logger = logging.getLogger(__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:
    """Wraps UpscalerUNet + the Flux.2 VAE for image-in / image-out 2x upscaling.

    The UNet defaults to float32 (the training dtype); bf16 matches it to
    ~54 dB PSNR while halving the forward peak and runtime (app_local passes
    bf16). The shared VAE is used as-is (bf16 in this project).
    """

    # 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 (upscaler_unet) 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 the app's resident
    # pipelines. fp32 needs roughly double.
    MAX_OUTPUT_SIDE = 16384

    def __init__(
        self,
        model_path: str,
        vae,
        device: str = "cuda",
        dtype: torch.dtype = torch.float32,
    ):
        self.unet = UpscalerUNet()
        self.unet.load_state_dict(load_file(model_path))
        self.unet.to(device, dtype).eval()
        self.vae = vae
        self.device = device
        self.dtype = dtype

    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)
        return mean.to(self.device, self.dtype), std.to(self.device, self.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.device, self.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
        latents_small = latents_small.to(self.device, self.dtype)

        scheduler = FlowMatchEulerDiscreteScheduler()
        scheduler.set_timesteps(1, mu=1.0)

        latents = torch.normal(
            mean=0.0,
            std=1.0,
            size=(batch_size, 32, height * 2, width * 2),
            dtype=self.dtype,
            device=self.device,
            generator=generator,
        )
        for t in scheduler.timesteps:
            t = t.to(self.device).view(1)
            velocity = self.unet(sample=latents, timestep=t, latents_small=latents_small)
            latents = 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.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 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:
        # 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,
            )

        generator = torch.Generator(device=self.device).manual_seed(seed)
        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 not color_fix:
            return self._tensor_to_pil(decoded)
        reference = torch.from_numpy(np.asarray(image.convert("RGB")).copy())
        reference = reference.to(self.device).permute(2, 0, 1)[None].float().div_(255.0)
        matched = self._lab_stat_match(decoded, reference, strength=float(color_fix_strength))
        return Image.fromarray(matched.cpu().numpy())