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