Upload folder using huggingface_hub

LICENSE

@@ -0,0 +1,24 @@
+This is free and unencumbered software released into the public domain.
+
+Anyone is free to copy, modify, publish, use, compile, sell, or
+distribute this software, either in source code form or as a compiled
+binary, for any purpose, commercial or non-commercial, and by any
+means.
+
+In jurisdictions that recognize copyright laws, the author or authors
+of this software dedicate any and all copyright interest in the
+software to the public domain. We make this dedication for the benefit
+of the public at large and to the detriment of our heirs and
+successors. We intend this dedication to be an overt act of
+relinquishment in perpetuity of all present and future rights to this
+software under copyright law.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
+OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
+ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+OTHER DEALINGS IN THE SOFTWARE.
+
+For more information, please refer to <https://unlicense.org/>

README.md

@@ -0,0 +1,246 @@
+---
+license: unlicense
+tags:
+  - flux2
+  - flux
+  - image-to-image
+  - image-upscaling
+  - latent-upscaling
+  - super-resolution
+  - diffusion
+  - flow-matching
+  - rectified-flow
+  - generative-models
+  - image-generation
+  - pytorch
+library_name: pytorch
+base_model:
+  - TensorForger/FlowUpscaler
+pipeline_tag: image-to-image
+---
+
+
+# Flow Upscaler
+
+
+
+**Flow Upscaler** is a fast latent upscaler model that works in the [Flux.2](https://bfl.ai/models/flux-2) latent space.
+
+Under the hood, it is a lightweight **Rectified Flow** model with **59M** parameters that generates upscaled latents in a single denoising step.
+
+
+## Inference pipeline (`flow_upscaler_pipeline.py`)
+
+`flow_upscaler_pipeline.py` is a standalone, diffusers-based image-in /
+image-out wrapper around the model — the plain-Python equivalent of the
+ComfyUI node. `FlowUpscalerPipeline` combines the UNet defined in
+`upscaler_unet.py` (weights: `flow_upscaler.safetensors`) with a shared
+Flux.2 VAE and runs the full loop:
+
+1. **Encode** (`encode_image`): the input image is resized to a multiple of
+   16, VAE-encoded, and the 32-channel latents are normalized with the VAE's
+   BatchNorm running stats — the same convention the model was trained with.
+2. **Upscale** (`upscale_latents`): pure noise is denoised into the 2x latent
+   in a **single FlowMatchEuler step**, conditioned on the small latent.
+   Passes chain for 4x/8x because a pass's output lives in the same
+   normalized latent space as its conditioning input.
+3. **Decode + color fix** (`upscale_image`): latents are denormalized and
+   VAE-decoded. With the default `color_fix=True`, each pass ends with a
+   latent low-frequency transplant and the decoded image gets a low-band
+   L/a/b transplant from the input — the two-stage drift correction detailed
+   in the next section.
+
+Memory is bounded for large outputs: the VAE switches to tiling above 2048px,
+the UNet runs through a chunked inference executor (`upscaler_unet.py`), and
+outputs are capped at `MAX_OUTPUT_SIDE = 16384` px per side (excess passes
+are dropped with a warning). This keeps even 16K outputs within a 24GB-class
+GPU in bf16.
+
+```python
+import torch
+from diffusers import AutoencoderKLFlux2  # any Flux.2 checkpoint's VAE
+from PIL import Image
+
+from flow_upscaler_pipeline import FlowUpscalerPipeline
+
+vae = AutoencoderKLFlux2.from_pretrained(
+    "black-forest-labs/FLUX.2-dev", subfolder="vae", torch_dtype=torch.bfloat16
+).to("cuda")
+pipeline = FlowUpscalerPipeline("flow_upscaler.safetensors", vae, dtype=torch.bfloat16)
+
+image = Image.open("input.png")
+result = pipeline.upscale_image(image, num_passes=2)  # 2 passes = 4x
+result.save("output.png")
+```
+
+For the VAE, consider
+[MageFlow-VAE-diffusers](https://huggingface.co/MinhNH232331M/MageFlow-VAE-diffusers)
+— an efficient drop-in replacement for the Flux.2 VAE that operates in the
+same latent space (~10x faster encode, ~6x faster decode, roughly half the
+peak VRAM, with a chunked decoder that keeps memory near-constant at large
+resolutions). Since upscaling cost here is dominated by the encode/decode
+ends, it speeds up the whole loop.
+
+`upscale_image` knobs: `num_passes` (2x per pass), `seed`, `color_fix`
+(master switch for both correction stages), and the per-stage
+`transplant_strength` / `color_fix_strength` sliders described below.
+`transplant_strength` also acts as a fidelity/realism control — lower values
+track the input more faithfully, higher values give the model's crisper
+rendition.
+
+
+## Color drift and corrections (this repo's pipeline)
+
+Chained single-step passes drift measurably. The local wrapper
+(`flow_upscaler_pipeline.py`, used by `app_local.py`) measures and corrects
+three drift modes, all enabled by the default `color_fix=True`:
+
+| Drift mode | Magnitude (uncorrected) | Cause |
+|---|---|---|
+| Per-channel latent mean drift | ~0.1 sigma/pass -> 2-5/255 RGB cast after 2 passes | single-step Euler DC error (VAE round-trip is clean) |
+| Chroma amplification | ~8%/pass mean chroma; p99 +14/255 after 3 passes, concentrated on already-colorful regions (lips, skin) | model bias, rides on high-band latent content through the nonlinear decoder |
+| Tone-curve stretch | dL_std +1.1, shadows -2.75, highlights +2 after 3 passes | side effect of the latent transplant below (raw model output is slightly *flat*) |
+
+Two-stage correction:
+
+1. **Latent low-frequency transplant** (`_lowfreq_transplant`), after every
+   pass: the output latent's low band (antialiased downsample by
+   `LOWFREQ_FACTOR = 8`) is replaced with the conditioning's. Fixes the mean
+   drift and keeps the next pass's conditioning on-manifold, which measurably
+   sharpens results. Cost: ~1.5 ms / +33 MiB at an 8K-UHD output.
+   `transplant_strength` interpolates this stage from off (softer, truest to input structure) to
+   full (default, crisper); color is unaffected at any setting. Beyond drift
+   correction, this knob doubles as a **fidelity/realism control**: lower
+   values stay measurably closer to the input (PSNR to input 43.9 dB at 0.0
+   vs 38.4 dB at 1.0 in the 2-pass test below), higher values favor the
+   model's sharper, more confident rendition.
+2. **Pixel-space low-band L/a/b transplant** (`_lab_stat_match`), after the
+   final decode: Lab diff maps computed at the input's resolution
+   (area-downsampled output vs input, fp32), bilinearly upsampled and added
+   to the decoded image in fp16 row strips on GPU. Below the input's Nyquist
+   the input is ground truth, so regional color *and* tone are pinned exactly
+   while the model keeps all detail above it. Cost incl. uint8 download:
+   13 / 54 / 235 ms and +51 / +128 / +404 MiB transient VRAM at 2K / 4K / 8K.
+   `color_fix_strength` interpolates this stage linearly; residual drift scales with `1 - strength`.
+
+### Measured results
+
+Protocol: output area-downsampled (`INTER_AREA`) back to input size and
+compared against the input in Lab; worst cases from a 512px-base photo set,
+2-3 pass chains, seed 42, NVIDIA L4, bf16. Reproduce with
+`test_flow_upscaler.ipynb` (one section per design choice, with plots).
+
+| Metric (at input scale) | Uncorrected | Both fixes on |
+|---|---|---|
+| RGB mean cast, 2 passes | up to 5/255 per channel | <= 0.05/255 |
+| Mean chroma (Lab C_mean) | +7...+15% | within +/-0.5% |
+| Chroma p99 (worst case, 3 passes) | +13.8/255 | +0.6/255 |
+| Tone contrast (L_std), 3 passes | +1.08 | +0.03 |
+| Shadows p5 / highlights p95 | -2.75 / +1.96 | 0.00 / 0.00 |
+
+Knob response (2-pass restore-photo test, other knob at its 1.0 default —
+the two stages are orthogonal by construction and by measurement):
+
+| `transplant_strength` | 1.0 | 0.5 | 0.0 |
+|---|---|---|---|
+| chroma error dC | -0.02 | -0.02 | -0.02 |
+| sharpness (mean abs Laplacian) | 7.94 | 4.86 | 3.28 |
+| PSNR to input (dB) | 38.4 | 41.7 | 43.9 |
+
+| `color_fix_strength` | 1.0 | 0.5 | 0.0 |
+|---|---|---|---|
+| chroma error dC | -0.02 | +0.53 | +1.16 |
+| sharpness | 7.94 | 7.98 | 8.07 |
+
+Compounding is per pass: before the pixel-space fix, mean chroma ran +8.6%
+after 1 pass and +14.8% after 2; the latent mean drift is ~0.1 sigma per pass on average with
+0.35-0.48 sigma outlier channels. The VAE round-trip itself is clean (chroma
+-0.02, RGB means <= 0.9/255) — every drift mode above is created by the flow
+UNet or by the correction stack itself.
+
+### Nyquist / frequency-band analysis
+
+Band bookkeeping, with frequencies written as fractions of each grid's
+Nyquist: one latent pixel decodes to 8 image pixels, so the latent grid's
+Nyquist sits at 1/16 cycles/px in pixel space; a conditioning at half the
+output resolution owns everything below 0.5x the output-latent Nyquist. Each
+correction stage operates in the widest band where its reference is ground
+truth *and* its domain assumption holds:
+
+* **Stage 1** corrects below output-latent Nyquist / `LOWFREQ_FACTOR` (= /8).
+  The limiting factor is not the conditioning (ground truth up to 0.5) but
+  how faithfully latent-space bands map to pixel-space bands — see below.
+* **Stage 2** corrects below the *input image's* Nyquist, in pixel space,
+  where the low band is the pixel-space low band by definition — no mapping
+  assumption needed. Everything above the input Nyquist is model-invented
+  detail with no reference, and is deliberately left untouched.
+
+The latent->pixel mapping was measured directly on the Mage decoder: inject a
+band-limited 0.1-sigma perturbation into one latent spectral band, decode,
+and integrate the luma response spectrum per pixel band (bands as fractions
+of the latent Nyquist; pixel content above 1.0 is finer than the latent grid
+can represent spatially):
+
+| Latent band | -> same pixel band | -> finer than latent Nyquist |
+|---|---|---|
+| 0 - 1/16 | 87.7% | 7.8% |
+| 1/16 - 1/8 | 90.6% | 4.2% |
+| 1/8 - 1/4 | 77.1% | 15.0% |
+| 1/4 - 1/2 | 64.7% | 26.2% |
+| 1/2 - 1 | 38.0% | 58.1% |
+
+Three design rules fall out of this table:
+
+1. Latent band surgery is sound only in the top two rows (~88-91%
+   correspondence) — exactly where `LOWFREQ_FACTOR` 8-16 operates. Below 77%
+   the spliced content scatters across unrelated pixel frequencies and reads
+   as ghosting (matches the failed wide-band experiments below).
+2. The ~8% upward leakage from the lowest band is the latent transplant's
+   sharpening / tone-stretch side effect, quantified — the energy it injects
+   into finer-than-latent-Nyquist pixel detail.
+3. Energy fractions understate *statistical* bias: high-band latent chroma
+   texture rectifies through the nonlinear decoder into a mean-saturation
+   shift carrying almost no low-band energy (the mechanism behind drift
+   mode 2). Exact guarantees therefore only come from the pixel-space stage.
+
+
+
+
+
+**[ComfyUI Node](https://github.com/TensorForger/comfyui-flow-upscaler)**
+
+Features:
+
+* Upscaling from **512x512** to **1024x1024** takes **8ms*** 
+* The model is trained for **2X** upscaling, but multiple passes can be chained to reach up to **8K** resolution
+* A full pipeline with Flux generation, upscaling to **8K**, and decoding runs in just **25 seconds** (on RTX 5090)
+* The training process uses **Flow Distillation** with Flux.2 as a teacher, forcing the model to learn strong image semantics
+
+*On RTX 5090, in latent space, without decoding, see benchmark [here](https://github.com/tensorforger/CTGMWorkshop).
+
+Here is one **4X** upscaled image (two passes):
+
+![comparison](https://raw.githubusercontent.com/tensorforger/tensorforger/main/assets/upscaler_comparison.png)
+
+## How it works
+
+Architecturally, Flow Upscaler is a U-Net with SDXL-style ResNet blocks. It takes a noisy sample as input and predicts velocity as output. The generation process happens directly in high-resolution latent space.
+
+The low-resolution latents are passed through a separate conditioning encoder that produces control signals, which are injected into the main U-Net encoder using FiLM conditioning.
+
+No attention layers are used, so compute scales linearly with image area. This makes generation at **8K** resolution possible.
+
+![example](https://raw.githubusercontent.com/tensorforger/tensorforger/main/assets/flow_upscaler_architecture.PNG)
+
+The model is trained using **Flow Distillation** with Flux.2-klein-4B as a teacher. We generated **20K** diverse images with Flux, storing the initial noise, generated latents, and downscaled latents used for conditioning.
+
+The downscaled latents are created by decoding high-resolution latents, downscaling them in pixel space, and encoding them back into latents. Direct latent downscaling introduces artifacts and breaks latent patterns, resulting in blurry decoded images.
+
+![example](https://raw.githubusercontent.com/tensorforger/tensorforger/main/assets/flow_upscaler_training_approach.PNG)
+
+
+## Training code
+
+If you want to explore the training code or use the model outside ComfyUI, see:
+
+`notebooks/flow_upscaler` in [https://github.com/tensorforger/CTGMWorkshop](https://github.com/tensorforger/CTGMWorkshop)

config.json

@@ -0,0 +1,3 @@
+{
+  "model_type": "custom"
+}

flow_upscaler.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:91eb93b40179441e569e01fe04f20fd9b951e434f88c84dd0ba523f830a81839
+size 237085968

flow_upscaler_pipeline.py

@@ -0,0 +1,350 @@
+"""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())

upscaler_unet.py

@@ -0,0 +1,621 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def make_group_norm(
+    channels: int, max_groups: int = 32, eps: float = 1e-6
+) -> nn.GroupNorm:
+    groups = min(max_groups, channels)
+    while channels % groups != 0 and groups > 1:
+        groups -= 1
+    return nn.GroupNorm(groups, channels, eps=eps)
+
+
+# ---- chunked inference execution -------------------------------------------
+# At large latents a block's transient temporaries (SiLU copies, full-res
+# FiLM scale/shift maps, conv workspaces) dwarf its input/output tensors.
+# The chunked path caps the live set at ~3 activation-sized tensors per
+# block: convs write row strips (1px halo for 3x3) into preallocated
+# outputs, FiLM applies in strips, and GN runs as the native full-tensor op
+# (its output just fills a live-set slot the convs already require, so
+# keeping it native costs no memory and keeps the chunked path bit-identical
+# to the standard one). Inference-only; callers gate on
+# torch.is_grad_enabled().
+
+CHUNK_MIN_PIXELS = 512 * 512     # engage the chunked path at/above this H*W
+CHUNK_STRIP_PIXELS = 128 * 1024  # target H*W of one conv row strip
+
+
+def _strip_rows(height: int, width: int) -> int:
+    return max(8, min(height, CHUNK_STRIP_PIXELS // max(1, width)))
+
+
+def _use_chunked(x: torch.Tensor) -> bool:
+    return not torch.is_grad_enabled() and x.shape[-2] * x.shape[-1] >= CHUNK_MIN_PIXELS
+
+
+def chunked_conv(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    out: torch.Tensor | None = None,
+    accumulate: bool = False,
+) -> torch.Tensor:
+    """Stride-1 odd-kernel conv2d in row strips.
+
+    With ``accumulate`` the strips are added into ``out``, which lets the last
+    conv of a block write straight into its residual buffer.
+    """
+    pad = weight.shape[-1] // 2
+    batch, _, height, width = x.shape
+    if out is None:
+        out = x.new_empty(batch, weight.shape[0], height, width)
+    rows = _strip_rows(height, width)
+    for r0 in range(0, height, rows):
+        r1 = min(r0 + rows, height)
+        a = max(r0 - pad, 0)
+        y = F.conv2d(x[:, :, a : min(r1 + pad, height)], weight, bias, padding=pad)
+        y = y[:, :, r0 - a : r0 - a + r1 - r0]
+        if accumulate:
+            out[:, :, r0:r1] += y
+        else:
+            out[:, :, r0:r1] = y
+    return out
+
+
+def chunked_film_(film: "FilmCond2D", x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+    """FilmCond2D applied in place in row strips.
+
+    Avoids materializing the full-res 2x-channel scale/shift map (0.74GiB at
+    an 8K-UHD latent for film_cond_1).
+    """
+    silu, proj = film.cond_proj[0], film.cond_proj[1]
+    height, width = x.shape[-2:]
+    rows = _strip_rows(height, width)
+    for r0 in range(0, height, rows):
+        r1 = min(r0 + rows, height)
+        scale_shift = proj(silu(cond[:, :, r0:r1]))
+        scale, shift = scale_shift.chunk(2, dim=1)
+        x[:, :, r0:r1].mul_(1 + scale).add_(shift)
+    return x
+
+
+class SinusoidalTimeEmbedding(nn.Module):
+    def __init__(self, dim: int = 128, max_period: int = 10000):
+        super().__init__()
+        self.dim = dim
+        self.max_period = max_period
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        half = self.dim // 2
+
+        freqs = torch.exp(
+            -torch.log(torch.tensor(float(self.max_period), device=timesteps.device))
+            * torch.arange(half, device=timesteps.device, dtype=timesteps.dtype)
+            / half
+        )
+        args = timesteps[:, None] * freqs[None]
+        emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
+
+        if self.dim % 2 == 1:
+            emb = F.pad(emb, (0, 1))
+
+        return emb
+
+
+class ConditioningEncoder(nn.Module):
+    def __init__(self, time_dim: int = 128, cond_dim: int = 256):
+        super().__init__()
+        self.time_embed = SinusoidalTimeEmbedding(time_dim)
+
+        self.time_proj = nn.Sequential(
+            nn.Linear(time_dim, cond_dim),
+            nn.SiLU(),
+            nn.Linear(cond_dim, cond_dim),
+        )
+
+    def forward(self, timestep: torch.Tensor) -> torch.Tensor:
+        # Sinusoidal embedding stays fp32; cast to the projection weight
+        # dtype so the model can run in bf16.
+        emb = self.time_embed(timestep).to(self.time_proj[0].weight.dtype)
+        return self.time_proj(emb)
+
+
+class ConditionedResidualBlock(nn.Module):
+    """
+    SDXL-style residual block:
+      GN -> SiLU -> Conv
+      + condition (scale/shift)
+      GN -> SiLU -> Dropout -> Conv
+      + skip connection
+    """
+
+    def __init__(
+        self,
+        input_channels: int,
+        output_channels: int,
+        cond_dim: int = 256,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+
+        self.norm1 = make_group_norm(input_channels)
+        self.conv1 = nn.Conv2d(
+            input_channels, output_channels, kernel_size=3, padding=1
+        )
+
+        self.cond_proj = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(cond_dim, 2 * output_channels),
+        )
+
+        self.norm2 = make_group_norm(output_channels)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(
+            output_channels, output_channels, kernel_size=3, padding=1
+        )
+
+        if input_channels != output_channels:
+            self.skip = nn.Conv2d(
+                input_channels, output_channels, kernel_size=1, bias=False
+            )
+        else:
+            self.skip = nn.Identity()
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        if _use_chunked(x):
+            return self._forward_chunked(x, cond)
+        # In-place ops need the grad guard: silu/mul_ backward reads the
+        # pre-op values. At inference they drop several activation-sized
+        # temporaries per block (~0.5GiB peak at an 8K-UHD latent).
+        inplace = not torch.is_grad_enabled()
+        residual = self.skip(x)
+
+        h = self.norm1(x)
+        h = F.silu(h, inplace=inplace)
+        h = self.conv1(h)
+
+        scale_shift = self.cond_proj(cond)
+        scale, shift = scale_shift.chunk(2, dim=1)
+        scale = scale[:, :, None, None]
+        shift = shift[:, :, None, None]
+
+        h = self.norm2(h)
+        if inplace:
+            h.mul_(1 + scale).add_(shift)
+        else:
+            h = h * (1 + scale) + shift
+        h = F.silu(h, inplace=inplace)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        return h.add_(residual) if inplace else h + residual
+
+    def _forward_chunked(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        # Large-latent inference: native out-of-place GN keeps x pristine so
+        # it doubles as the residual; convs run in row strips. Peak live set
+        # is 3 activation-sized tensors, and the output is bit-identical to
+        # the standard path. Dropout is identity at inference and is skipped.
+        residual = x if isinstance(self.skip, nn.Identity) else self.skip(x)
+        h = self.norm1(x)
+        F.silu(h, inplace=True)
+        h = chunked_conv(h, self.conv1.weight, self.conv1.bias)
+
+        scale_shift = self.cond_proj(cond)
+        scale, shift = scale_shift.chunk(2, dim=1)
+        h = self.norm2(h)
+        h.mul_(1 + scale[:, :, None, None]).add_(shift[:, :, None, None])
+        F.silu(h, inplace=True)
+        out = chunked_conv(h, self.conv2.weight, self.conv2.bias)
+        del h
+        return out.add_(residual)
+
+
+class DownStage(nn.Module):
+    def __init__(
+        self,
+        input_channels: int,
+        output_channels: int,
+        cond_dim: int = 256,
+        dropout: float = 0.0,
+        num_blocks: int = 1,
+        downsample_first: bool = False,
+    ):
+        super().__init__()
+        self.downsample_first = downsample_first
+
+        self.blocks = nn.ModuleList()
+        for i in range(num_blocks):
+            in_ch = input_channels if i == 0 else output_channels
+            self.blocks.append(
+                ConditionedResidualBlock(
+                    input_channels=in_ch,
+                    output_channels=output_channels,
+                    cond_dim=cond_dim,
+                    dropout=dropout,
+                )
+            )
+
+        self.downsample = nn.Conv2d(
+            output_channels, output_channels, kernel_size=3, stride=2, padding=1
+        )
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor):
+
+        if self.downsample_first:
+            x = self.downsample(x)
+
+        for block in self.blocks:
+            x = block(x, cond)
+        skip = x
+
+        if not self.downsample_first:
+            x = self.downsample(x)
+
+        return x, skip
+
+
+class UpStage(nn.Module):
+    def __init__(
+        self,
+        input_channels: int,
+        skip_channels: int,
+        output_channels: int,
+        cond_dim: int = 256,
+        dropout: float = 0.0,
+        num_blocks: int = 1,
+    ):
+        super().__init__()
+
+        self.upsample = nn.Upsample(
+            scale_factor=2, mode="bilinear", align_corners=False
+        )
+
+        self.blocks = nn.ModuleList()
+        for i in range(num_blocks):
+            in_ch = (input_channels + skip_channels) if i == 0 else output_channels
+            self.blocks.append(
+                ConditionedResidualBlock(
+                    input_channels=in_ch,
+                    output_channels=output_channels,
+                    cond_dim=cond_dim,
+                    dropout=dropout,
+                )
+            )
+
+    def forward(
+        self, x: torch.Tensor, skip: torch.Tensor, cond: torch.Tensor
+    ) -> torch.Tensor:
+        x = self.upsample(x)
+
+        if x.shape[-2:] != skip.shape[-2:]:
+            x = F.interpolate(
+                x, size=skip.shape[-2:], mode="bilinear", align_corners=False
+            )
+
+        if torch.is_grad_enabled() or not self._can_split(x, skip):
+            x = torch.cat([x, skip], dim=1)
+            for block in self.blocks:
+                x = block(x, cond)
+            return x
+
+        if _use_chunked(x):
+            # Hand the tensors over in a list so this frame stops pinning
+            # them — the method frees each one as soon as it is consumed.
+            tensors = [x, skip]
+            del x, skip
+            x = self._split_concat_block0_chunked(tensors, cond)
+        else:
+            x = self._split_concat_block0(x, skip, cond)
+        for block in self.blocks[1:]:
+            x = block(x, cond)
+        return x
+
+    def _can_split(self, x: torch.Tensor, skip: torch.Tensor) -> bool:
+        b0 = self.blocks[0]
+        ca, cb = x.shape[1], skip.shape[1]
+        groups = b0.norm1.num_groups
+        ga = groups * ca // (ca + cb)
+        return (
+            isinstance(b0.skip, nn.Conv2d)
+            and groups * ca % (ca + cb) == 0
+            and 0 < ga < groups
+            and ca % ga == 0
+            and cb % (groups - ga) == 0
+        )
+
+    def _split_concat_block0(
+        self, x: torch.Tensor, skip: torch.Tensor, cond: torch.Tensor
+    ) -> torch.Tensor:
+        """Run blocks[0] without materializing the (ca+cb)-channel concat.
+
+        Valid when the GroupNorm group boundaries align with the concat
+        boundary (checked by _can_split): GN/SiLU apply per half, and the
+        skip-conv/conv1 outputs are the sums of per-half convolutions. Only
+        the fp reduction order differs from the concat path — this is the
+        peak-memory point of the whole UNet at large latents (the concat is
+        full-res at double width), worth ~1.1GiB at an 8K-UHD output.
+        Inference-only (in-place ops), guarded by the grad check in forward.
+        """
+        b0 = self.blocks[0]
+        ca = x.shape[1]
+        cb = skip.shape[1]
+        n1 = b0.norm1
+        ga = n1.num_groups * ca // (ca + cb)
+
+        w = b0.skip.weight
+        residual = F.conv2d(x, w[:, :ca]).add_(F.conv2d(skip, w[:, ca:]))
+
+        h = F.silu(
+            F.group_norm(x, ga, n1.weight[:ca], n1.bias[:ca], n1.eps), inplace=True
+        )
+        del x
+        w1 = b0.conv1.weight
+        out = F.conv2d(h, w1[:, :ca], b0.conv1.bias, padding=b0.conv1.padding)
+        del h
+        h = F.silu(
+            F.group_norm(
+                skip, n1.num_groups - ga, n1.weight[ca:], n1.bias[ca:], n1.eps
+            ),
+            inplace=True,
+        )
+        del skip
+        out.add_(F.conv2d(h, w1[:, ca:], padding=b0.conv1.padding))
+        del h
+
+        scale_shift = b0.cond_proj(cond)
+        scale, shift = scale_shift.chunk(2, dim=1)
+        out = b0.norm2(out)
+        out.mul_(1 + scale[:, :, None, None]).add_(shift[:, :, None, None])
+        out = F.silu(out, inplace=True)
+        out = b0.conv2(b0.dropout(out))
+        return out.add_(residual)
+
+    def _split_concat_block0_chunked(
+        self, tensors: list, cond: torch.Tensor
+    ) -> torch.Tensor:
+        """_split_concat_block0 with chunked convs and native per-half GN;
+        conv2 accumulates into the residual buffer. Output is bit-identical
+        to _split_concat_block0.
+
+        ``tensors`` is [x, skip]; it is emptied so each half can be freed the
+        moment its contributions are computed.
+        """
+        b0 = self.blocks[0]
+        x, skip = tensors
+        tensors.clear()
+        ca = x.shape[1]
+        n1 = b0.norm1
+        ga = n1.num_groups * ca // (ca + skip.shape[1])
+
+        w = b0.skip.weight
+        residual = chunked_conv(x, w[:, :ca])
+        chunked_conv(skip, w[:, ca:], out=residual, accumulate=True)
+
+        h = F.group_norm(x, ga, n1.weight[:ca], n1.bias[:ca], n1.eps)
+        del x
+        F.silu(h, inplace=True)
+        w1 = b0.conv1.weight
+        out = chunked_conv(h, w1[:, :ca], b0.conv1.bias)
+        del h
+        h = F.group_norm(skip, n1.num_groups - ga, n1.weight[ca:], n1.bias[ca:], n1.eps)
+        del skip
+        F.silu(h, inplace=True)
+        chunked_conv(h, w1[:, ca:], out=out, accumulate=True)
+        del h
+
+        scale_shift = b0.cond_proj(cond)
+        scale, shift = scale_shift.chunk(2, dim=1)
+        h = F.group_norm(out, b0.norm2.num_groups, b0.norm2.weight, b0.norm2.bias, b0.norm2.eps)
+        del out
+        h.mul_(1 + scale[:, :, None, None]).add_(shift[:, :, None, None])
+        F.silu(h, inplace=True)
+        return chunked_conv(h, b0.conv2.weight, b0.conv2.bias, out=residual, accumulate=True)
+
+
+class LowResEncoder(nn.Module):
+    def __init__(
+        self,
+        sample_channels: int = 32,
+        base_channels: int = 128,
+        cond_dim: int = 1024,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+
+        self.in_conv = nn.Conv2d(
+            sample_channels, base_channels, kernel_size=1, padding=0
+        )
+
+        self.block_1 = ConditionedResidualBlock(
+            input_channels=base_channels,
+            output_channels=base_channels,
+            cond_dim=cond_dim,
+            dropout=dropout,
+        )
+
+        self.block_2 = DownStage(
+            input_channels=base_channels,
+            output_channels=base_channels,
+            cond_dim=cond_dim,
+            dropout=dropout,
+            num_blocks=1,
+            downsample_first=True,
+        )
+
+        self.block_3 = DownStage(
+            input_channels=base_channels,
+            output_channels=base_channels,
+            cond_dim=cond_dim,
+            dropout=dropout,
+            num_blocks=1,
+            downsample_first=True,
+        )
+
+    def forward(self, latents_small, cond):
+        x = self.in_conv(latents_small)
+        block_1_out = self.block_1(x, cond)
+        block_2_out, _ = self.block_2(block_1_out, cond)
+        block_3_out, _ = self.block_3(block_2_out, cond)
+
+        return block_1_out, block_2_out, block_3_out
+
+
+class FilmCond2D(nn.Module):
+    def __init__(self, base_channels: int = 256, cond_channels: int = 256):
+        super().__init__()
+
+        self.cond_proj = nn.Sequential(
+            nn.SiLU(),
+            nn.Conv2d(cond_channels, base_channels * 2, kernel_size=1),
+        )
+
+    def forward(self, x, cond):
+        if _use_chunked(x):
+            return chunked_film_(self, x, cond)
+
+        scale_shift = self.cond_proj(cond)
+        scale, shift = scale_shift.chunk(2, dim=1)
+
+        if torch.is_grad_enabled():
+            return x * (1 + scale) + shift
+        # Callers rebind x, so modifying it in place at inference is safe and
+        # avoids two full-res temporaries.
+        return x.mul_(1 + scale).add_(shift)
+
+
+class UpscalerUNet(nn.Module):
+    def __init__(
+        self,
+        sample_channels: int = 32,
+        base_channels: int = 384,
+        time_dim: int = 512,
+        cond_dim: int = 1024,
+        dropout: float = 0.01,
+    ):
+        super().__init__()
+
+        self.conditioning = ConditioningEncoder(
+            time_dim=time_dim,
+            cond_dim=cond_dim,
+        )
+
+        self.in_conv = nn.Conv2d(
+            sample_channels, base_channels, kernel_size=1, padding=0
+        )
+
+        self.low_res_encoder = LowResEncoder(base_channels=base_channels)
+
+        self.film_cond_1 = FilmCond2D(
+            base_channels=base_channels, cond_channels=base_channels
+        )
+        self.film_cond_2 = FilmCond2D(
+            base_channels=base_channels, cond_channels=base_channels
+        )
+        self.film_cond_3 = FilmCond2D(
+            base_channels=base_channels, cond_channels=base_channels
+        )
+
+        self.down_stages = nn.ModuleList(
+            [
+                DownStage(
+                    input_channels=base_channels,
+                    output_channels=base_channels,
+                    cond_dim=cond_dim,
+                    dropout=dropout,
+                    num_blocks=3,
+                ),
+                DownStage(
+                    input_channels=base_channels,
+                    output_channels=base_channels,
+                    cond_dim=cond_dim,
+                    dropout=dropout,
+                    num_blocks=2,
+                ),
+            ]
+        )
+
+        self.mid_stages = nn.ModuleList(
+            [
+                ConditionedResidualBlock(
+                    input_channels=base_channels,
+                    output_channels=base_channels,
+                    cond_dim=cond_dim,
+                    dropout=dropout,
+                )
+                for i in range(1)
+            ]
+        )
+
+        self.up_stages = nn.ModuleList(
+            [
+                UpStage(
+                    input_channels=base_channels,
+                    skip_channels=base_channels,
+                    output_channels=base_channels,
+                    cond_dim=cond_dim,
+                    dropout=dropout,
+                    num_blocks=2,
+                ),
+                UpStage(
+                    input_channels=base_channels,
+                    skip_channels=base_channels,
+                    output_channels=base_channels,
+                    cond_dim=cond_dim,
+                    dropout=dropout,
+                    num_blocks=3,
+                ),
+            ]
+        )
+
+        self.out_conv = nn.Conv2d(
+            base_channels, sample_channels, kernel_size=1, padding=0
+        )
+
+    def forward(
+        self, sample: torch.Tensor, timestep: torch.Tensor, latents_small: torch.Tensor
+    ) -> torch.Tensor:
+        cond = self.conditioning(timestep)
+
+        B, C, H, W = sample.shape
+
+        lr_cond_1, lr_cond_2, lr_cond_3 = self.low_res_encoder(latents_small, cond)
+
+        lr_cond_1 = torch.nn.functional.interpolate(lr_cond_1, (H, W), mode="bilinear")
+        lr_cond_2 = torch.nn.functional.interpolate(
+            lr_cond_2, (H // 2, W // 2), mode="bilinear"
+        )
+        lr_cond_3 = torch.nn.functional.interpolate(
+            lr_cond_3, (H // 4, W // 4), mode="bilinear"
+        )
+
+        x = self.in_conv(sample)
+        x = self.film_cond_1(x, lr_cond_1)
+        # Each lr_cond is dead after its film layer; without the dels the
+        # full-res lr_cond_1 (~0.4GiB at an 8K-UHD latent) stays pinned for
+        # the whole forward.
+        del lr_cond_1
+
+        skips = []
+
+        x, skip = self.down_stages[0](x, cond)
+        skips.append(skip)
+
+        x = self.film_cond_2(x, lr_cond_2)
+        del lr_cond_2
+
+        x, skip = self.down_stages[1](x, cond)
+        skips.append(skip)
+
+        x = self.film_cond_3(x, lr_cond_3)
+        del lr_cond_3
+
+        for mid in self.mid_stages:
+            x = mid(x, cond)
+
+        for up in self.up_stages:
+            x = up(x, skips.pop(), cond)
+
+        x = self.out_conv(x)
+        return x