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