Upload folder using huggingface_hub

.vscode/settings.json

@@ -1,4 +0,0 @@
-{
-    "python-envs.defaultEnvManager": "ms-python.python:conda",
-    "python-envs.defaultPackageManager": "ms-python.python:conda"
-}

README.md

@@ -13,7 +13,7 @@ tags:
   - generative-models
   - image-generation
   - pytorch
-library_name: pytorch
+library_name: diffusers
 base_model:
   - TensorForger/FlowUpscaler
 pipeline_tag: image-to-image
@@ -31,11 +31,16 @@ Under the hood, it is a lightweight **Rectified Flow** model with **59M** parame
 
 ## 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:
+This repo is in **diffusers format**: `FlowUpscalerPipeline` is a
+`DiffusionPipeline` whose components are the UNet (`unet/upscaler_unet.py`,
+a diffusers `ModelMixin`, weights in
+`unet/diffusion_pytorch_model.safetensors`), a shared Flux.2 VAE, and a
+`FlowMatchEulerDiscreteScheduler` — the plain-Python equivalent of the
+ComfyUI node, loadable with `DiffusionPipeline.from_pretrained`. The VAE is
+intentionally not bundled (it is shared with any Flux.2 checkpoint, and
+FLUX.2-dev is gated), so it is passed in at load time;
+`trust_remote_code=True` is required because the pipeline and UNet classes
+are loaded from this repo. The pipeline 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
@@ -44,35 +49,44 @@ Flux.2 VAE and runs the full loop:
    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
+3. **Decode + color fix** (`__call__`): 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`,
+the UNet runs through a chunked inference executor (`unet/upscaler_unet.py`,
 detailed in the next section), 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 diffusers import AutoencoderKLFlux2, DiffusionPipeline
 from PIL import Image
 
-from flow_upscaler_pipeline import FlowUpscalerPipeline
-
-vae = AutoencoderKLFlux2.from_pretrained(
+vae = AutoencoderKLFlux2.from_pretrained(  # any Flux.2 checkpoint's VAE
     "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")
-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 = pipeline(image, num_passes=2).images[0]  # 2 passes = 4x
 result.save("output.png")
 ```
 
+Two compatibility paths are kept. `flow_upscaler.safetensors` at the repo
+root is the same flat checkpoint the ComfyUI node consumes, and
+`FlowUpscalerPipeline.from_single_file("flow_upscaler.safetensors", vae,
+dtype=...)` replicates the pre-diffusers constructor for local checkouts.
+The seed-based `pipeline.upscale_image(image, num_passes, seed=42, ...)`
+method also remains, and a given seed reproduces the pre-conversion outputs
+bit-for-bit (`__call__` takes a `generator` instead, diffusers-style).
+
 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
@@ -81,22 +95,24 @@ 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.
+Pipeline call knobs: `num_passes` (2x per pass), `generator` (RNG for the
+per-pass noise; `upscale_image` takes a `seed` instead), `color_fix` (master
+switch for both correction stages), the per-stage `transplant_strength` /
+`color_fix_strength` sliders described below, and `output_type` (`"pil"`
+default, `"np"`, `"pt"`). `transplant_strength` also acts as a
+fidelity/realism control — lower values track the input more faithfully,
+higher values give the model's crisper rendition.
 
 
-## UNet inference optimizations (`upscaler_unet.py`)
+## UNet inference optimizations (`unet/upscaler_unet.py`)
 
 The network is attention-free, so compute scales linearly with area — but the
 stock forward's *memory* did not scale gracefully: at large latents each
 residual block allocated 4-8 full-resolution temporaries (GroupNorm outputs,
 SiLU copies, FiLM scale/shift maps, conv workspaces), and the skip-concat
-up-block briefly held several 768-channel full-res tensors. `upscaler_unet.py`
-now layers four inference-only optimizations on the same weights. All of them
+up-block briefly held several 768-channel full-res tensors.
+`unet/upscaler_unet.py` now layers four inference-only optimizations on the
+same weights. All of them
 are gated on `torch.is_grad_enabled()`, so training behavior is bit-for-bit
 untouched.
 

config.json

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

flow_upscaler_pipeline.py

@@ -1,5 +1,20 @@
 """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
@@ -14,9 +29,15 @@ pass ends with a low-frequency transplant from the conditioning latent
 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 logging
+import importlib.util
+import os
 
 import numpy as np
 import torch
@@ -24,11 +45,17 @@ 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
+from diffusers import (
+    AutoencoderKLFlux2,
+    DiffusionPipeline,
+    FlowMatchEulerDiscreteScheduler,
+    ImagePipelineOutput,
+    ModelMixin,
+)
+from diffusers.utils import logging
+from diffusers.utils.torch_utils import randn_tensor
 
-logger = logging.getLogger(__name__)
+logger = logging.get_logger(__name__)
 
 
 def patchify_latents(latents: torch.Tensor) -> torch.Tensor:
@@ -81,48 +108,96 @@ def _lab_to_rgb(lightness: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> to
     ).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).
+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 (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
+    # 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,
-    ):
-        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
+    ) -> "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)
-        return mean.to(self.device, self.dtype), std.to(self.device, self.dtype)
+        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.device, self.dtype))
+        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:
@@ -177,23 +252,23 @@ class FlowUpscalerPipeline:
     ) -> 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)
+        device, dtype = self._execution_device, self.unet.dtype
+        latents_small = latents_small.to(device, dtype)
 
-        scheduler = FlowMatchEulerDiscreteScheduler()
-        scheduler.set_timesteps(1, mu=1.0)
+        # 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 = torch.normal(
-            mean=0.0,
-            std=1.0,
-            size=(batch_size, 32, height * 2, width * 2),
-            dtype=self.dtype,
-            device=self.device,
+        latents = randn_tensor(
+            (batch_size, self.unet.config.sample_channels, height * 2, width * 2),
             generator=generator,
+            device=device,
+            dtype=dtype,
         )
-        for t in scheduler.timesteps:
-            t = t.to(self.device).view(1)
+        for t in self.scheduler.timesteps:
+            t = t.view(1)
             velocity = self.unet(sample=latents, timestep=t, latents_small=latents_small)
-            latents = scheduler.step(velocity, t, latents).prev_sample
+            latents = self.scheduler.step(velocity, t, latents).prev_sample
         if color_fix:
             latents = self._lowfreq_transplant(latents, latents_small, strength=transplant_strength)
         return latents
@@ -213,7 +288,7 @@ class FlowUpscalerPipeline:
             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)
+        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
@@ -306,15 +381,51 @@ class FlowUpscalerPipeline:
         return out
 
     @torch.no_grad()
-    def upscale_image(
+    def __call__(
         self,
         image: Image.Image,
         num_passes: int = 1,
-        seed: int = 42,
+        generator: torch.Generator | None = None,
         color_fix: bool = True,
         transplant_strength: float = 1.0,
         color_fix_strength: float = 1.0,
-    ) -> Image.Image:
+        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.
@@ -334,7 +445,6 @@ class FlowUpscalerPipeline:
                 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(
@@ -342,9 +452,51 @@ class FlowUpscalerPipeline:
                 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())
+
+        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]

model_index.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": ["flow_upscaler_pipeline", "FlowUpscalerPipeline"],
+  "_diffusers_version": "0.37.1",
+  "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
+  "unet": ["upscaler_unet", "UpscalerUNet"],
+  "vae": ["diffusers", "AutoencoderKLFlux2"]
+}

scheduler/scheduler_config.json

@@ -0,0 +1,18 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.37.1",
+  "base_image_seq_len": 256,
+  "base_shift": 0.5,
+  "invert_sigmas": false,
+  "max_image_seq_len": 4096,
+  "max_shift": 1.15,
+  "num_train_timesteps": 1000,
+  "shift": 1.0,
+  "shift_terminal": null,
+  "stochastic_sampling": false,
+  "time_shift_type": "exponential",
+  "use_beta_sigmas": false,
+  "use_dynamic_shifting": false,
+  "use_exponential_sigmas": false,
+  "use_karras_sigmas": false
+}

unet/config.json

@@ -0,0 +1,9 @@
+{
+  "_class_name": "UpscalerUNet",
+  "_diffusers_version": "0.37.1",
+  "base_channels": 384,
+  "cond_dim": 1024,
+  "dropout": 0.01,
+  "sample_channels": 32,
+  "time_dim": 512
+}

unet/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7930adcbf8545c23fe60d5d1157de2fedfd497a0bed48018f315c110a3e3e330
+size 237086000

upscaler_unet.py → unet/upscaler_unet.py

@@ -2,6 +2,9 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+
 
 def make_group_norm(
     channels: int, max_groups: int = 32, eps: float = 1e-6
@@ -485,7 +488,17 @@ class FilmCond2D(nn.Module):
         return x.mul_(1 + scale).add_(shift)
 
 
-class UpscalerUNet(nn.Module):
+class UpscalerUNet(ModelMixin, ConfigMixin):
+    """FlowUpscaler velocity-prediction UNet (diffusers `ModelMixin`).
+
+    Attention-free U-Net with SDXL-style residual blocks conditioned on the
+    timestep (FiLM from a sinusoidal embedding) and on the low-resolution
+    latents (multi-scale FiLM from `LowResEncoder`). The module tree is
+    identical to the original standalone `nn.Module`, so the state dict is
+    interchangeable with the flat `flow_upscaler.safetensors` checkpoint.
+    """
+
+    @register_to_config
     def __init__(
         self,
         sample_channels: int = 32,