Upload folder using huggingface_hub

README.md

@@ -0,0 +1,171 @@
+---
+license: cc-by-nc-4.0
+license_name: cc-by-nc-4.0
+pipeline_tag: image-to-image
+library_name: diffusers
+tags:
+  - sar-to-eo
+  - remote-sensing
+  - flow-matching
+  - image-to-image
+  - synthetic-aperture-radar
+base_model: black-forest-labs/FLUX.2-klein-base-4B
+---
+
+# ReFlowSET
+
+SAR-to-EO image translation with a conditional flow-matching transformer trained
+from scratch inside a frozen high-fidelity autoencoder latent space.
+
+- **Paper:** https://arxiv.org/abs/{{ARXIV_ID}}
+- **Code:** https://github.com/KAIST-VICLab/ReFlowSET
+- **Project page:** https://kaist-viclab.github.io/ReFlowSET_site/
+- **Comparison-method weights:** [`JeonghyeokDo/ReFlowSET-baselines`](https://huggingface.co/JeonghyeokDo/ReFlowSET-baselines)
+
+## Checkpoints
+
+| Subfolder | Dataset | Resolution | Training | Deployed parameters |
+|---|---|---|---|---|
+| `qxs-saropt` | QXS-SAROPT | 256×256 | 40,000 steps × global batch 64 (2.56 M samples) | 509,324,417 + 84,046,115 frozen autoencoder |
+| `sar2opt` | SAR2Opt | 512×512 | 20,000 steps × global batch 32 (640 k samples) | same |
+
+Each subfolder is a complete `diffusers` pipeline: `transformer/`, `vae/`,
+`scheduler/` and `model_index.json`. The two arms share the same architecture and
+the same frozen autoencoder; they differ only in resolution, batch size and step
+count. SAR2Opt stops at 20,000 steps to hold a comparable sample budget on a
+1,450-image training set.
+
+The published weights are the **EMA** parameters. The training-only REPA projector
+is not included.
+
+## Usage
+
+```python
+import torch
+from PIL import Image
+from diffusers import DiffusionPipeline
+from huggingface_hub import snapshot_download
+
+# Both arms live in this one repository, one per subfolder. `DiffusionPipeline`
+# has no `subfolder` argument, so fetch the arm and load it as a local pipeline.
+ARM = "qxs-saropt"                                   # or "sar2opt"
+root = snapshot_download("JeonghyeokDo/ReFlowSET", allow_patterns=[f"{ARM}/*"])
+pipe = DiffusionPipeline.from_pretrained(
+    f"{root}/{ARM}", custom_pipeline=f"{root}/{ARM}", torch_dtype=torch.float32,
+).to("cuda")
+
+sar = Image.open("sar.png")                       # 1-channel SAR, 8-bit PNG
+eo = pipe(sar, num_inference_steps=50, guidance_scale=1.5,
+          generator=torch.Generator("cuda").manual_seed(2024)).images[0]
+eo.save("eo.png")
+```
+
+`custom_pipeline` points at the same directory because the pipeline, transformer,
+autoencoder and scheduler classes ship with the checkpoint rather than living in
+`diffusers`. The classes are also on GitHub under `src/reflowset/`.
+
+**Sampling settings are part of the reported result, not free knobs.** The paper's
+main table is NFE 50 with guidance scale 1.5. NFE 4 samples **11× faster at 256²**
+(163 ms vs 1824 ms) and **13× faster at 512²** (371 ms vs 4807 ms), batch 1 on one
+B200, and trades distribution metrics against pixel metrics; do not mix the two in
+one comparison.
+
+The SAR input is read without a colour conversion, collapsed to one channel,
+center-cropped (never resized), scaled by `x / 127.5 − 1`, replicated to three
+channels, and encoded by the same frozen autoencoder that defines the EO latent
+space. The pipeline does all of this; feed it the raw PNG.
+
+## Results
+
+Scored on the same test items as fifteen prior methods that we retrained under one
+protocol, by a single evaluator.
+
+| Dataset | n | FID↓ | DISTS↓ | LPIPS↓ | SSIM↑ | PSNR↑ |
+|---|---|---|---|---|---|---|
+| QXS-SAROPT @256 | 3,999 | 19.1 | **0.2310** | 0.5344 | 0.3554 | 16.09 |
+| SAR2Opt @512 | 627 | **66.3** | **0.1847** | **0.5217** | 0.2871 | 16.06 |
+
+Bold marks the best value among all sixteen methods in the paper's main table.
+The full table, with every comparison method's weights and licence, is in
+[`MODEL_ZOO.md`](https://github.com/KAIST-VICLab/ReFlowSET/blob/main/MODEL_ZOO.md).
+
+> **These numbers are not comparable with the ones printed in the source papers.**
+> Splits, resolutions and evaluator conventions differ. In particular **LPIPS has
+> two conventions in this literature that differ by ~0.05**: we feed `x*2−1` to
+> the LPIPS network, while several released evaluators feed `[0,1]` with
+> `normalize=False` and obtain a systematically lower number.
+
+## Architecture
+
+A DiT with hidden size 1024 and depth 24 — eight double-stream blocks that give
+the EO and SAR streams their own projections and joint attention, then sixteen
+single-stream blocks over the concatenated token sequence — with 16 heads of
+dimension 64 and 2-D RoPE over axes (32, 32).
+
+There is **no separate SAR encoder**: the SAR image goes through the same frozen
+autoencoder as the EO image. Training defines a linear bridge
+`z_t = (1−t)·ε + t·z_e` and regresses the velocity `u* = z_e − ε` conditioned on
+the SAR latent; sampling starts from `N(0, I)` and integrates `t: 0 → 1` with an
+explicit Euler step. Classifier-free guidance was trained by zeroing the SAR
+condition on 10 % of rows.
+
+## Training data
+
+QXS-SAROPT and SAR2Opt only, one dataset per arm, trained from scratch. No
+pretraining corpus. Neither dataset is redistributed. QXS-SAROPT requires citing
+arXiv:2103.08259 for research use.
+
+## Licence and provenance
+
+**Weights: CC BY-NC 4.0. Code: Apache-2.0.**
+
+The frozen autoencoder bundled in `vae/` is the **Apache-2.0** autoencoder from
+[`black-forest-labs/FLUX.2-klein-base-4B`](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-4B),
+re-serialised to the upstream layout and cast to bfloat16 — tensors paired by
+value, not by an assumed rename table, and bit-identical to that source through a
+full encode/decode.
+
+**It is a substitution, and here is exactly what was substituted.** Both arms were
+*trained and evaluated* with the `FLUX.2-dev` serialisation of the same network;
+that file is under the FLUX Non-Commercial License, whose §4(a)(iii) forbids
+"research and development related to surveillance" and whose §1(a) makes the
+restriction inherit permanently. The two serialisations are the same autoencoder
+— 250 of 251 tensors pair by value, worst absolute deviation 7.8e-03 (bfloat16
+rounding) — and swapping the Apache file into the released checkpoints changes
+QXS-SAROPT PSNR by **less than 0.004 dB in absolute value**. Four independent
+measurements at different guidance scales and sample sets land between −0.004 and
++0.002 dB, so the sign is not resolved and only the magnitude is meaningful.
+Changing only the evaluation seed moves the same number by +0.395 dB. Nothing
+reported here changes.
+
+Apache-2.0 §6 withholds trademark rights. This model is named ReFlowSET; it is not
+a FLUX product and is not endorsed by Black Forest Labs.
+
+**Why this autoencoder.** A latent generator cannot beat its codec's round trip,
+so the codec is a ceiling on every row of a latent-model comparison.
+[`vae_audit/`](https://github.com/KAIST-VICLab/ReFlowSET/tree/main/vae_audit)
+measures that ceiling for six autoencoders — SD2.1, SDXL, SD3.0, SD3.5, FLUX.1
+and FLUX.2 — on four SAR/EO benchmarks, EO and SAR scored separately. It ships as
+code with download links and licence terms; no imagery and no third-party
+autoencoder weights are redistributed.
+
+Training used a frozen **DINOv3** ViT-L/16 (LVD-1689M) as a representation-
+alignment teacher, acknowledged here as the DINOv3 License §1(b)(ii) requires. The
+teacher is not loaded at inference and **no DINOv3 weights are redistributed**;
+obtain them from [Meta's release](https://github.com/facebookresearch/dinov3)
+under its own terms if you intend to retrain.
+
+See [`LICENSE-WEIGHTS.md`](https://github.com/KAIST-VICLab/ReFlowSET/blob/main/LICENSE-WEIGHTS.md)
+for the full reasoning, including one open question about the datasets' optical
+imagery that we flag rather than resolve.
+
+## Citation
+
+```bibtex
+@article{do2026reflowset,
+  title   = {ReFlowSET: Representation-Aligned Latent Flow Matching for SAR-to-EO Image Translation},
+  author  = {Do, Jeonghyeok and Lee, Seungchul and Kim, Munchurl},
+  journal = {arXiv preprint arXiv:{{ARXIV_ID}}},
+  year    = {2026}
+}
+```

autoencoder_flux2.py

@@ -0,0 +1,426 @@
+"""Frozen FLUX.2 autoencoder — the latent endpoint of ReFlowSET.
+
+ReFlowSET never fine-tunes this module: it is loaded once, frozen, and used to
+encode the SAR condition and to decode the sampled EO latent.  The released
+weights are the **Apache-2.0** FLUX.2-klein-base-4B copy of the autoencoder,
+re-keyed to the layout below (see ``scripts/convert_flux2_ae.py``).
+
+Three details of the checkpoint are non-standard for `diffusers` and are
+preserved exactly, because the file must load with ``strict=True``:
+
+* ``quant_conv`` lives **inside** ``encoder.*`` and is the last op of the
+  encoder forward; ``post_quant_conv`` lives **inside** ``decoder.*`` and is the
+  first op of the decoder forward.  `diffusers`' ``AutoencoderKL`` makes both
+  siblings of the encoder/decoder.
+* The latent normaliser is a real ``BatchNorm2d(128, affine=False)`` whose
+  running statistics ship in the checkpoint under ``bn.*`` — a per-channel mean
+  **and** variance, not a scalar ``scaling_factor``/``shift_factor``.  Its
+  epsilon is ``1e-4``, not torch's ``1e-5``.
+* ``encode`` returns the posterior **mean**; the log-variance chunk of the
+  encoder's moments is discarded, so encoding is deterministic and there is no
+  ``DiagonalGaussianDistribution`` and no ``.sample()``.
+
+The public latent is ``[B, 128, H/16, W/16]``: an 8x convolutional stride
+followed by a 2x2 space-to-depth pack that is part of the *autoencoder*, not of
+the transformer.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+
+def swish(x: Tensor) -> Tensor:
+    """``x * sigmoid(x)`` — the activation used throughout the FLUX.2 AE."""
+    return x * torch.sigmoid(x)
+
+
+class AttnBlock(nn.Module):
+    """Single-head self-attention over the spatial grid (head dim == channels)."""
+
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q, k, v = self.q(h_), self.k(h_), self.v(h_)
+        b, c, h, w = q.shape
+        # "b c h w -> b 1 (h w) c": ONE head whose head-dim is the full channel
+        # count (flux2_ae.py:70-73).
+        q = q.reshape(b, c, h * w).transpose(1, 2).unsqueeze(1).contiguous()
+        k = k.reshape(b, c, h * w).transpose(1, 2).unsqueeze(1).contiguous()
+        v = v.reshape(b, c, h * w).transpose(1, 2).unsqueeze(1).contiguous()
+        h_ = F.scaled_dot_product_attention(q, k, v)
+        return h_.squeeze(1).transpose(1, 2).reshape(b, c, h, w)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if in_channels != out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x: Tensor) -> Tensor:
+        h = self.conv1(swish(self.norm1(x)))
+        h = self.conv2(swish(self.norm2(h)))
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+        return x + h
+
+
+class Downsample(nn.Module):
+    """Stride-2 conv with FLUX's asymmetric ``(0, 1, 0, 1)`` pad (flux2_ae.py:111-121)."""
+
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.conv(F.pad(x, (0, 1, 0, 1), mode="constant", value=0))
+
+
+class Upsample(nn.Module):
+    """Nearest-neighbour 2x followed by a 3x3 conv (flux2_ae.py:124-132)."""
+
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.conv(F.interpolate(x, scale_factor=2.0, mode="nearest"))
+
+
+class Encoder(nn.Module):
+    """FLUX.2 encoder.  Emits ``2 * z_channels`` moments; ``quant_conv`` is internal."""
+
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        z_channels: int,
+    ) -> None:
+        super().__init__()
+        # Declared first so the checkpoint key is `encoder.quant_conv.*`
+        # (flux2_ae.py:146) — diffusers keeps quant_conv outside the encoder.
+        self.quant_conv = nn.Conv2d(2 * z_channels, 2 * z_channels, 1)
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        self.conv_in = nn.Conv2d(in_channels, ch, kernel_size=3, stride=1, padding=1)
+
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.down = nn.ModuleList()
+        block_in = ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(num_res_blocks):
+                block.append(ResnetBlock(block_in, block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            # Empty at every level in this checkpoint: attention exists only in
+            # `mid` (flux2_ae.py:162).  Kept so the forward guard is meaningful.
+            down.attn = nn.ModuleList()
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in)
+            self.down.append(down)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(block_in, block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(block_in, block_in)
+
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        h = self.mid.block_2(self.mid.attn_1(self.mid.block_1(hs[-1])))
+        h = self.conv_out(swish(self.norm_out(h)))
+        return self.quant_conv(h)  # last op of the encoder (flux2_ae.py:207)
+
+
+class Decoder(nn.Module):
+    """FLUX.2 decoder.  ``post_quant_conv`` is internal and runs first."""
+
+    def __init__(
+        self,
+        ch: int,
+        out_ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+    ) -> None:
+        super().__init__()
+        # Checkpoint key `decoder.post_quant_conv.*` (flux2_ae.py:223).
+        self.post_quant_conv = nn.Conv2d(z_channels, z_channels, 1)
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(block_in, block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(block_in, block_in)
+
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(num_res_blocks + 1):
+                block.append(ResnetBlock(block_in, block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = nn.ModuleList()  # empty in this checkpoint (flux2_ae.py:249)
+            if i_level != 0:
+                up.upsample = Upsample(block_in)
+            self.up.insert(0, up)  # prepend so `up.<i>` indexes by resolution level
+
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: Tensor) -> Tensor:
+        z = self.post_quant_conv(z)  # first op of the decoder (flux2_ae.py:267)
+        upscale_dtype = next(self.up.parameters()).dtype
+
+        h = self.conv_in(z)
+        h = self.mid.block_2(self.mid.attn_1(self.mid.block_1(h)))
+        h = h.to(upscale_dtype)
+
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        return self.conv_out(swish(self.norm_out(h)))
+
+
+class AutoencoderFlux2(ModelMixin, ConfigMixin):
+    """Frozen FLUX.2 autoencoder with ReFlowSET's packed, BN-normalised latent.
+
+    ``encode`` maps ``[B, 3, H, W]`` in ``[-1, 1]`` to ``[B, 128, H/16, W/16]``
+    and ``decode`` inverts it.  The module is frozen: ``train()`` is a no-op that
+    always selects eval mode, and the latent BatchNorm is additionally forced to
+    eval on every call so no batch statistic can ever leak into the latent.
+
+    Args:
+        resolution: Nominal training resolution of the original autoencoder.
+            Only used to size bookkeeping attributes; any ``H``, ``W`` divisible
+            by 16 may be encoded.
+        in_channels: Input image channels (3).
+        ch: Base width.
+        out_ch: Output image channels (3).
+        ch_mult: Per-level width multipliers; ``len(ch_mult) - 1`` downsamples.
+        num_res_blocks: Residual blocks per level.
+        z_channels: Pre-pack latent channels (32).
+        patch_size: Space-to-depth factor applied after the encoder (2), which
+            takes the latent from 32 channels at ``H/8`` to 128 at ``H/16``.
+        bn_eps: Epsilon of the latent BatchNorm.  **1e-4**, not torch's 1e-5
+            (flux2_ae.py:331); using 1e-5 shifts the latent by up to 2.6e-5.
+    """
+
+    _supports_gradient_checkpointing = False
+
+    @register_to_config
+    def __init__(
+        self,
+        resolution: int = 256,
+        in_channels: int = 3,
+        ch: int = 128,
+        out_ch: int = 3,
+        ch_mult: tuple[int, ...] = (1, 2, 4, 4),
+        num_res_blocks: int = 2,
+        z_channels: int = 32,
+        patch_size: int = 2,
+        bn_eps: float = 1e-4,
+    ) -> None:
+        super().__init__()
+        ch_mult = list(ch_mult)
+        self.encoder = Encoder(
+            resolution=resolution,
+            in_channels=in_channels,
+            ch=ch,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            z_channels=z_channels,
+        )
+        self.decoder = Decoder(
+            ch=ch,
+            out_ch=out_ch,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            in_channels=in_channels,
+            resolution=resolution,
+            z_channels=z_channels,
+        )
+        # Per-channel latent normaliser with the checkpoint's running statistics.
+        # affine=False, so there is no weight/bias to load (flux2_ae.py:334-340).
+        self.bn = nn.BatchNorm2d(
+            patch_size * patch_size * z_channels,
+            eps=bn_eps,
+            momentum=0.1,
+            affine=False,
+            track_running_stats=True,
+        )
+
+    @property
+    def latent_channels(self) -> int:
+        """Channels of the public latent: ``patch_size**2 * z_channels`` = 128."""
+        return self.config.patch_size**2 * self.config.z_channels
+
+    @property
+    def spatial_factor(self) -> int:
+        """Total stride: 8x convolutional times ``patch_size`` packing = 16."""
+        return 2 ** (len(self.config.ch_mult) - 1) * self.config.patch_size
+
+    # ---- 2x2 space-to-depth pack / unpack -----------------------------------
+
+    def pack(self, z: Tensor) -> Tensor:
+        """``[B, C, H, W] -> [B, C*p*p, H/p, W/p]``, channel-major.
+
+        Bit-identical to the reference ``rearrange("... c (i pi) (j pj) -> ...
+        (c pi pj) i j")`` (flux2_ae.py:349-357).  Note this is **not** diffusers'
+        ``_pack_latents``, whose channel grouping is transposed.
+        """
+        return F.pixel_unshuffle(z, self.config.patch_size)
+
+    def unpack(self, z: Tensor) -> Tensor:
+        """Exact inverse of :meth:`pack` (flux2_ae.py:359-367)."""
+        return F.pixel_shuffle(z, self.config.patch_size)
+
+    # ---- latent normalisation ----------------------------------------------
+
+    def normalize(self, z: Tensor) -> Tensor:
+        """``(z - running_mean) / sqrt(running_var + bn_eps)``, per channel."""
+        self.bn.eval()  # forced every call (flux2_ae.py:372); train mode shifts z by ~1.67
+        return self.bn(z)
+
+    def inv_normalize(self, z: Tensor) -> Tensor:
+        """Exact inverse of :meth:`normalize` — same ``bn_eps`` (flux2_ae.py:375-379)."""
+        self.bn.eval()
+        s = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.config.bn_eps)
+        m = self.bn.running_mean.view(1, -1, 1, 1)
+        return z * s + m
+
+    # ---- public API ---------------------------------------------------------
+
+    @torch.no_grad()
+    def encode(self, x: Tensor) -> Tensor:
+        """Encode an image to the packed, normalised latent.
+
+        Args:
+            x: ``[B, 3, H, W]`` in ``[-1, 1]``; ``H`` and ``W`` divisible by 16.
+
+        Returns:
+            ``[B, 128, H/16, W/16]`` — the posterior **mean**, packed and
+            BN-normalised.  The encoder's log-variance chunk is discarded
+            (flux2_ae.py:396), so this is deterministic: there is no posterior
+            distribution object and nothing to sample.
+        """
+        if x.ndim != 4 or x.shape[1] != self.config.in_channels:
+            raise ValueError(
+                f"encode expects [B, {self.config.in_channels}, H, W], got {tuple(x.shape)}"
+            )
+        h, w = x.shape[-2:]
+        if h % self.spatial_factor or w % self.spatial_factor:
+            raise ValueError(
+                f"encode requires H and W divisible by {self.spatial_factor}, got {h}x{w}"
+            )
+        moments = self.encoder(x)
+        mean = torch.chunk(moments, 2, dim=1)[0]
+        return self.normalize(self.pack(mean))
+
+    @torch.no_grad()
+    def decode(self, z: Tensor) -> Tensor:
+        """Decode a packed, normalised latent ``[B, 128, h, w]`` to ``[B, 3, 16h, 16w]``.
+
+        The output is approximately ``[-1, 1]`` and is **not** clamped here; the
+        pipeline applies ``(x * 0.5 + 0.5).clamp(0, 1)``.
+        """
+        if z.ndim != 4 or z.shape[1] != self.latent_channels:
+            raise ValueError(
+                f"decode expects [B, {self.latent_channels}, h, w], got {tuple(z.shape)}"
+            )
+        return self.decoder(self.unpack(self.inv_normalize(z)))
+
+    # ---- construction / freezing -------------------------------------------
+
+    @classmethod
+    def from_single_file(
+        cls,
+        path: str | os.PathLike,
+        torch_dtype: torch.dtype = torch.float32,
+    ) -> "AutoencoderFlux2":
+        """Load the single-file ``ae.safetensors`` (BFL key names) with ``strict=True``.
+
+        The released file is the Apache-2.0 FLUX.2-klein-base-4B autoencoder
+        re-keyed to this layout; it is stored in bfloat16 and is upcast to
+        ``torch_dtype``.  ReFlowSET runs the autoencoder in float32.
+        """
+        from safetensors.torch import load_file
+
+        path = os.fspath(path)
+        if not os.path.isfile(path):
+            raise FileNotFoundError(
+                f"FLUX.2 autoencoder weights not found at: {path}. Expected the "
+                "single-file 'ae.safetensors' shipped with ReFlowSET."
+            )
+        model = cls()
+        model.load_state_dict(load_file(path, device="cpu"), strict=True)
+        model.to(dtype=torch_dtype)
+        model.eval()
+        model.requires_grad_(False)
+        return model
+
+    def train(self, mode: bool = True) -> "AutoencoderFlux2":
+        """The autoencoder is frozen: never leave eval mode (flux2_ae.py:437-439)."""
+        return super().train(False)

pipeline.py

@@ -0,0 +1,267 @@
+"""ReFlowSET SAR -> EO translation pipeline."""
+
+from __future__ import annotations
+
+import warnings
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from PIL import Image
+
+from .autoencoder_flux2 import AutoencoderFlux2
+from .scheduler_flow_bridge import FlowBridgeScheduler
+from .transformer_reflowset import ReFlowSETTransformer2DModel
+
+#: PIL modes the SAR loader accepts.  The reference loader calls ``np.array(im)``
+#: with no ``convert()`` (datasets.py:454, 574), so a 16-bit (``I;16``) or
+#: palette (``P``) raster would flow straight into ``x / 127.5 - 1`` and be
+#: badly out of range.  That is an unguarded trap upstream; it is guarded here.
+_ACCEPTED_SAR_MODES = ("L", "RGB", "RGBA")
+
+
+class ReFlowSETPipeline(DiffusionPipeline):
+    """Generate an EO image from a SAR image with ReFlowSET's flow bridge.
+
+    Args:
+        transformer: The velocity transformer.
+        vae: The frozen FLUX.2 autoencoder that defines the latent space.
+        scheduler: The Design-B flow-bridge Euler solver.
+
+    To reproduce the paper's numbers, sample at ``num_inference_steps=50``,
+    ``guidance_scale=1.5``, float32, one image per call, with a generator freshly
+    seeded to 2024 on the compute device before each call — every test image in
+    the reported evaluation starts from the same seeded noise draw, and CPU-drawn
+    noise does not reproduce a CUDA draw.
+    """
+
+    model_cpu_offload_seq = "transformer->vae"
+
+    def __init__(
+        self,
+        transformer: ReFlowSETTransformer2DModel,
+        vae: AutoencoderFlux2,
+        scheduler: FlowBridgeScheduler,
+    ) -> None:
+        super().__init__()
+        self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)
+
+    # ---- preprocessing (datasets.py:124-205, 452-474, 572-589) --------------
+
+    @staticmethod
+    def _sar_hwc(raster: Union[Image.Image, np.ndarray]) -> np.ndarray:
+        """SAR raster -> ``[H, W, C]`` float32 in ``[0, 255]``, collapsed to 1 channel."""
+        if isinstance(raster, Image.Image):
+            if raster.mode not in _ACCEPTED_SAR_MODES:
+                raise ValueError(
+                    f"SAR image mode {raster.mode!r} is not an 8-bit display raster; expected "
+                    f"one of {_ACCEPTED_SAR_MODES}. ReFlowSET was trained on 8-bit display "
+                    "quicklooks (sar_value_domain='display_png'); convert with .convert('L') "
+                    "and be aware that the contrast stretch you choose is part of the input."
+                )
+            # No .convert() on the SAR side, matching datasets.py:454, 574.
+            arr = np.array(raster)
+        else:
+            arr = np.asarray(raster)
+        if arr.ndim == 2:  # PIL mode "L" (datasets.py:171-172)
+            arr = arr[:, :, None]
+        if arr.shape[-1] == 4:  # drop a container alpha channel (datasets.py:173-174)
+            arr = arr[..., :3]
+        arr = arr.astype(np.float32)
+        if arr.shape[-1] > 1:
+            # Exact-equality test, tol=0.0 (datasets.py:100-111): a display RGB
+            # quicklook collapses to its single amplitude channel.
+            if np.abs(arr - arr[..., :1]).max() == 0.0:
+                arr = arr[..., :1]
+            else:
+                warnings.warn(
+                    "SAR raster has non-identical colour channels; feeding all 3 to the "
+                    "frozen encoder. The released arms were trained on single-channel "
+                    "amplitude quicklooks, so this is an undeclared input.",
+                    RuntimeWarning,
+                    stacklevel=3,
+                )
+        return arr
+
+    @staticmethod
+    def _center_crop(arr: np.ndarray, crop: int) -> np.ndarray:
+        """Center-crop ``[H, W, C]`` to ``crop`` — **never** resize (datasets.py:145-165).
+
+        The offsets are albumentations' ``CenterCrop`` arithmetic ``(n - c) // 2``:
+        the SAR2Opt protocol takes the central 512 of 600 at offset 44.
+        """
+        h, w = arr.shape[:2]
+        if h < crop or w < crop:
+            raise ValueError(
+                f"image {h}x{w} is smaller than the requested crop {crop}; ReFlowSET never "
+                "upscales an input"
+            )
+        top, left = (h - crop) // 2, (w - crop) // 2
+        return arr[top : top + crop, left : left + crop]
+
+    def preprocess(
+        self,
+        sar: Union[Image.Image, np.ndarray, torch.Tensor, list],
+        crop: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Build the model-boundary SAR tensor ``[B, 3, H, W]`` in ``[-1, 1]``.
+
+        Args:
+            sar: A PIL image, a list of PIL images, an ``[H, W]`` / ``[H, W, C]``
+                uint8 array, or a float tensor already in ``[-1, 1]`` shaped
+                ``[H, W]``, ``[C, H, W]`` or ``[B, C, H, W]``.
+            crop: Center-crop size applied before normalisation.  ``None``
+                center-crops to the arm's own training resolution when the
+                raster is larger and not already a multiple of the latent
+                stride -- which is exactly the SAR2Opt 600 -> 512 protocol the
+                reported numbers use.  Pass an explicit size to override, or
+                ``0`` to keep the native raster and fail loudly if it does not
+                fit.
+
+        Images are read as 8-bit display rasters and mapped to ``[-1, 1]`` by
+        ``x / 127.5 - 1`` (datasets.py:124-126) with no per-image statistics, no
+        percentile stretch and no resize.  The single SAR channel is then
+        replicated to 3 at the model boundary (evaluate.py:566-570), because the
+        frozen FLUX.2 encoder is the same one that encodes EO — ReFlowSET has no
+        separate SAR encoder.
+        """
+        if crop == 0:
+            crop = None
+        elif crop is None:
+            # Fall back to the resolution this arm was trained at. Cropping is
+            # the protocol (train.py random-crops, evaluate.py center-crops);
+            # ReFlowSET never resizes, so an un-croppable raster is an error
+            # rather than something to silently rescale.
+            crop = self.transformer.config.sample_size
+
+        if isinstance(sar, torch.Tensor):
+            x = sar.float()
+            if x.ndim == 2:
+                x = x[None, None]
+            elif x.ndim == 3:
+                x = x[None]
+            elif x.ndim != 4:
+                raise ValueError(f"SAR tensor must have 2, 3 or 4 dims, got {tuple(sar.shape)}")
+            h, w = x.shape[-2:]
+            if crop is not None and (h, w) != (crop, crop):
+                if h < crop or w < crop:
+                    raise ValueError(f"tensor {h}x{w} is smaller than the requested crop {crop}")
+                top, left = (h - crop) // 2, (w - crop) // 2
+                x = x[..., top : top + crop, left : left + crop]
+        else:
+            images = sar if isinstance(sar, list) else [sar]
+            arrays = []
+            for item in images:
+                if not isinstance(item, (Image.Image, np.ndarray)):
+                    raise TypeError(f"unsupported SAR input type {type(item)!r}")
+                arr = self._sar_hwc(item)
+                if crop is not None and arr.shape[:2] != (crop, crop):
+                    arr = self._center_crop(arr, crop)
+                arrays.append(np.ascontiguousarray(arr.transpose(2, 0, 1)))
+            x = torch.from_numpy(np.stack(arrays)) / 127.5 - 1.0
+
+        # Train-side clamp (train.py:770); a no-op on 8-bit input, which maps
+        # exactly onto [-1, 1].
+        x = x.clamp(-1.0, 1.0)
+        if x.shape[1] == 1:
+            x = x.repeat(1, 3, 1, 1)
+        elif x.shape[1] != 3:
+            raise ValueError(
+                f"the frozen FLUX.2 encoder takes 1 or 3 SAR channels, got {x.shape[1]}"
+            )
+        factor = self.vae.spatial_factor
+        if x.shape[-2] % factor or x.shape[-1] % factor:
+            raise ValueError(
+                f"SAR size {x.shape[-2]}x{x.shape[-1]} must be divisible by {factor}; pass "
+                "crop= to center-crop (ReFlowSET never resizes)"
+            )
+        return x
+
+    # ---- postprocessing -----------------------------------------------------
+
+    @staticmethod
+    def _to_pil(images: torch.Tensor) -> list[Image.Image]:
+        """``[B, 3, H, W]`` in ``[0, 1]`` -> PIL, quantised round-half-up.
+
+        ``255 * x + 0.5`` truncated is what ``torchvision.utils.save_image``
+        does and is therefore what the released PNGs contain; numpy's
+        ``round()`` is banker's rounding and would differ on exact halves.
+        """
+        arr = (images * 255 + 0.5).clamp(0, 255).to(torch.uint8)
+        arr = arr.permute(0, 2, 3, 1).cpu().numpy()
+        return [Image.fromarray(a) for a in arr]
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        sar: Union[Image.Image, np.ndarray, torch.Tensor, list],
+        num_inference_steps: int = 50,
+        guidance_scale: float = 1.5,
+        generator: Optional[Union[torch.Generator, list[torch.Generator]]] = None,
+        output_type: str = "pil",
+        crop: Optional[int] = None,
+        return_dict: bool = True,
+    ) -> Union[ImagePipelineOutput, tuple[list]]:
+        """Translate a SAR image into an EO image.
+
+        Args:
+            sar: SAR input; see :meth:`preprocess`.
+            num_inference_steps: NFE, the number of velocity evaluations.  The
+                paper's main results are NFE 50; NFE 4 is the efficiency
+                operating point and trades FID for PSNR/SSIM, so the two must
+                not be mixed in one comparison.
+            guidance_scale: Classifier-free guidance scale.  1.5 is the published
+                setting; 1.0 disables guidance and halves the cost.
+            generator: Generator for the initial noise.  Create it on the compute
+                device — CPU-drawn noise does not reproduce a CUDA draw.
+            output_type: ``"pil"``, ``"np"`` or ``"pt"``.
+            crop: Center-crop size applied to the SAR input before encoding.
+            return_dict: Return an ``ImagePipelineOutput`` instead of a tuple.
+
+        Returns:
+            The generated EO image(s) in ``[0, 1]`` (or as PIL).
+        """
+        if output_type not in ("pil", "np", "pt"):
+            raise ValueError(f"output_type must be 'pil', 'np' or 'pt', got {output_type!r}")
+
+        device = self._execution_device
+        dtype = self.transformer.dtype
+
+        sar_pm1 = self.preprocess(sar, crop=crop).to(device=device, dtype=self.vae.dtype)
+        # The SAR condition is encoded by the SAME frozen autoencoder that
+        # defines the EO latent space (evaluate.py:553-577).
+        z_s = self.vae.encode(sar_pm1).to(dtype)
+
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        # Design B: the bridge starts at t = 0 from pure Gaussian noise
+        # (bridge.py:409-433), NOT from the SAR latent.
+        latents = randn_tensor(z_s.shape, generator=generator, device=device, dtype=z_s.dtype)
+
+        for t in self.progress_bar(self.scheduler.timesteps):
+            timestep = t.expand(latents.shape[0])
+            velocity = self.transformer(latents, timestep, z_s, return_dict=False)[0]
+            if guidance_scale != 1.0:
+                # Two passes; the null branch is cond=None, which the transformer
+                # turns into an all-zero conditioning latent (bridge.py:530-535).
+                uncond = self.transformer(latents, timestep, None, return_dict=False)[0]
+                velocity = uncond + guidance_scale * (velocity - uncond)
+            latents = self.scheduler.step(velocity, t, latents, return_dict=False)[0]
+
+        image = self.vae.decode(latents.to(self.vae.dtype))
+        # `--denorm standard` (evaluate.py:292-295). The `legacy` C-DiffSET
+        # convention `(x + 0.5).clamp(0, 1)` is a 2x contrast stretch and must
+        # not be used with these numbers.
+        image = (image * 0.5 + 0.5).clamp(0.0, 1.0)
+
+        self.maybe_free_model_hooks()
+
+        if output_type == "pil":
+            image = self._to_pil(image)
+        elif output_type == "np":
+            image = image.permute(0, 2, 3, 1).float().cpu().numpy()
+
+        if not return_dict:
+            return (image,)
+        return ImagePipelineOutput(images=image)

pipeline_reflowset.py

@@ -0,0 +1,267 @@
+"""ReFlowSET SAR -> EO translation pipeline."""
+
+from __future__ import annotations
+
+import warnings
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from PIL import Image
+
+from .autoencoder_flux2 import AutoencoderFlux2
+from .scheduler_flow_bridge import FlowBridgeScheduler
+from .transformer_reflowset import ReFlowSETTransformer2DModel
+
+#: PIL modes the SAR loader accepts.  The reference loader calls ``np.array(im)``
+#: with no ``convert()`` (datasets.py:454, 574), so a 16-bit (``I;16``) or
+#: palette (``P``) raster would flow straight into ``x / 127.5 - 1`` and be
+#: badly out of range.  That is an unguarded trap upstream; it is guarded here.
+_ACCEPTED_SAR_MODES = ("L", "RGB", "RGBA")
+
+
+class ReFlowSETPipeline(DiffusionPipeline):
+    """Generate an EO image from a SAR image with ReFlowSET's flow bridge.
+
+    Args:
+        transformer: The velocity transformer.
+        vae: The frozen FLUX.2 autoencoder that defines the latent space.
+        scheduler: The Design-B flow-bridge Euler solver.
+
+    To reproduce the paper's numbers, sample at ``num_inference_steps=50``,
+    ``guidance_scale=1.5``, float32, one image per call, with a generator freshly
+    seeded to 2024 on the compute device before each call — every test image in
+    the reported evaluation starts from the same seeded noise draw, and CPU-drawn
+    noise does not reproduce a CUDA draw.
+    """
+
+    model_cpu_offload_seq = "transformer->vae"
+
+    def __init__(
+        self,
+        transformer: ReFlowSETTransformer2DModel,
+        vae: AutoencoderFlux2,
+        scheduler: FlowBridgeScheduler,
+    ) -> None:
+        super().__init__()
+        self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)
+
+    # ---- preprocessing (datasets.py:124-205, 452-474, 572-589) --------------
+
+    @staticmethod
+    def _sar_hwc(raster: Union[Image.Image, np.ndarray]) -> np.ndarray:
+        """SAR raster -> ``[H, W, C]`` float32 in ``[0, 255]``, collapsed to 1 channel."""
+        if isinstance(raster, Image.Image):
+            if raster.mode not in _ACCEPTED_SAR_MODES:
+                raise ValueError(
+                    f"SAR image mode {raster.mode!r} is not an 8-bit display raster; expected "
+                    f"one of {_ACCEPTED_SAR_MODES}. ReFlowSET was trained on 8-bit display "
+                    "quicklooks (sar_value_domain='display_png'); convert with .convert('L') "
+                    "and be aware that the contrast stretch you choose is part of the input."
+                )
+            # No .convert() on the SAR side, matching datasets.py:454, 574.
+            arr = np.array(raster)
+        else:
+            arr = np.asarray(raster)
+        if arr.ndim == 2:  # PIL mode "L" (datasets.py:171-172)
+            arr = arr[:, :, None]
+        if arr.shape[-1] == 4:  # drop a container alpha channel (datasets.py:173-174)
+            arr = arr[..., :3]
+        arr = arr.astype(np.float32)
+        if arr.shape[-1] > 1:
+            # Exact-equality test, tol=0.0 (datasets.py:100-111): a display RGB
+            # quicklook collapses to its single amplitude channel.
+            if np.abs(arr - arr[..., :1]).max() == 0.0:
+                arr = arr[..., :1]
+            else:
+                warnings.warn(
+                    "SAR raster has non-identical colour channels; feeding all 3 to the "
+                    "frozen encoder. The released arms were trained on single-channel "
+                    "amplitude quicklooks, so this is an undeclared input.",
+                    RuntimeWarning,
+                    stacklevel=3,
+                )
+        return arr
+
+    @staticmethod
+    def _center_crop(arr: np.ndarray, crop: int) -> np.ndarray:
+        """Center-crop ``[H, W, C]`` to ``crop`` — **never** resize (datasets.py:145-165).
+
+        The offsets are albumentations' ``CenterCrop`` arithmetic ``(n - c) // 2``:
+        the SAR2Opt protocol takes the central 512 of 600 at offset 44.
+        """
+        h, w = arr.shape[:2]
+        if h < crop or w < crop:
+            raise ValueError(
+                f"image {h}x{w} is smaller than the requested crop {crop}; ReFlowSET never "
+                "upscales an input"
+            )
+        top, left = (h - crop) // 2, (w - crop) // 2
+        return arr[top : top + crop, left : left + crop]
+
+    def preprocess(
+        self,
+        sar: Union[Image.Image, np.ndarray, torch.Tensor, list],
+        crop: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Build the model-boundary SAR tensor ``[B, 3, H, W]`` in ``[-1, 1]``.
+
+        Args:
+            sar: A PIL image, a list of PIL images, an ``[H, W]`` / ``[H, W, C]``
+                uint8 array, or a float tensor already in ``[-1, 1]`` shaped
+                ``[H, W]``, ``[C, H, W]`` or ``[B, C, H, W]``.
+            crop: Center-crop size applied before normalisation.  ``None``
+                center-crops to the arm's own training resolution when the
+                raster is larger and not already a multiple of the latent
+                stride -- which is exactly the SAR2Opt 600 -> 512 protocol the
+                reported numbers use.  Pass an explicit size to override, or
+                ``0`` to keep the native raster and fail loudly if it does not
+                fit.
+
+        Images are read as 8-bit display rasters and mapped to ``[-1, 1]`` by
+        ``x / 127.5 - 1`` (datasets.py:124-126) with no per-image statistics, no
+        percentile stretch and no resize.  The single SAR channel is then
+        replicated to 3 at the model boundary (evaluate.py:566-570), because the
+        frozen FLUX.2 encoder is the same one that encodes EO — ReFlowSET has no
+        separate SAR encoder.
+        """
+        if crop == 0:
+            crop = None
+        elif crop is None:
+            # Fall back to the resolution this arm was trained at. Cropping is
+            # the protocol (train.py random-crops, evaluate.py center-crops);
+            # ReFlowSET never resizes, so an un-croppable raster is an error
+            # rather than something to silently rescale.
+            crop = self.transformer.config.sample_size
+
+        if isinstance(sar, torch.Tensor):
+            x = sar.float()
+            if x.ndim == 2:
+                x = x[None, None]
+            elif x.ndim == 3:
+                x = x[None]
+            elif x.ndim != 4:
+                raise ValueError(f"SAR tensor must have 2, 3 or 4 dims, got {tuple(sar.shape)}")
+            h, w = x.shape[-2:]
+            if crop is not None and (h, w) != (crop, crop):
+                if h < crop or w < crop:
+                    raise ValueError(f"tensor {h}x{w} is smaller than the requested crop {crop}")
+                top, left = (h - crop) // 2, (w - crop) // 2
+                x = x[..., top : top + crop, left : left + crop]
+        else:
+            images = sar if isinstance(sar, list) else [sar]
+            arrays = []
+            for item in images:
+                if not isinstance(item, (Image.Image, np.ndarray)):
+                    raise TypeError(f"unsupported SAR input type {type(item)!r}")
+                arr = self._sar_hwc(item)
+                if crop is not None and arr.shape[:2] != (crop, crop):
+                    arr = self._center_crop(arr, crop)
+                arrays.append(np.ascontiguousarray(arr.transpose(2, 0, 1)))
+            x = torch.from_numpy(np.stack(arrays)) / 127.5 - 1.0
+
+        # Train-side clamp (train.py:770); a no-op on 8-bit input, which maps
+        # exactly onto [-1, 1].
+        x = x.clamp(-1.0, 1.0)
+        if x.shape[1] == 1:
+            x = x.repeat(1, 3, 1, 1)
+        elif x.shape[1] != 3:
+            raise ValueError(
+                f"the frozen FLUX.2 encoder takes 1 or 3 SAR channels, got {x.shape[1]}"
+            )
+        factor = self.vae.spatial_factor
+        if x.shape[-2] % factor or x.shape[-1] % factor:
+            raise ValueError(
+                f"SAR size {x.shape[-2]}x{x.shape[-1]} must be divisible by {factor}; pass "
+                "crop= to center-crop (ReFlowSET never resizes)"
+            )
+        return x
+
+    # ---- postprocessing -----------------------------------------------------
+
+    @staticmethod
+    def _to_pil(images: torch.Tensor) -> list[Image.Image]:
+        """``[B, 3, H, W]`` in ``[0, 1]`` -> PIL, quantised round-half-up.
+
+        ``255 * x + 0.5`` truncated is what ``torchvision.utils.save_image``
+        does and is therefore what the released PNGs contain; numpy's
+        ``round()`` is banker's rounding and would differ on exact halves.
+        """
+        arr = (images * 255 + 0.5).clamp(0, 255).to(torch.uint8)
+        arr = arr.permute(0, 2, 3, 1).cpu().numpy()
+        return [Image.fromarray(a) for a in arr]
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        sar: Union[Image.Image, np.ndarray, torch.Tensor, list],
+        num_inference_steps: int = 50,
+        guidance_scale: float = 1.5,
+        generator: Optional[Union[torch.Generator, list[torch.Generator]]] = None,
+        output_type: str = "pil",
+        crop: Optional[int] = None,
+        return_dict: bool = True,
+    ) -> Union[ImagePipelineOutput, tuple[list]]:
+        """Translate a SAR image into an EO image.
+
+        Args:
+            sar: SAR input; see :meth:`preprocess`.
+            num_inference_steps: NFE, the number of velocity evaluations.  The
+                paper's main results are NFE 50; NFE 4 is the efficiency
+                operating point and trades FID for PSNR/SSIM, so the two must
+                not be mixed in one comparison.
+            guidance_scale: Classifier-free guidance scale.  1.5 is the published
+                setting; 1.0 disables guidance and halves the cost.
+            generator: Generator for the initial noise.  Create it on the compute
+                device — CPU-drawn noise does not reproduce a CUDA draw.
+            output_type: ``"pil"``, ``"np"`` or ``"pt"``.
+            crop: Center-crop size applied to the SAR input before encoding.
+            return_dict: Return an ``ImagePipelineOutput`` instead of a tuple.
+
+        Returns:
+            The generated EO image(s) in ``[0, 1]`` (or as PIL).
+        """
+        if output_type not in ("pil", "np", "pt"):
+            raise ValueError(f"output_type must be 'pil', 'np' or 'pt', got {output_type!r}")
+
+        device = self._execution_device
+        dtype = self.transformer.dtype
+
+        sar_pm1 = self.preprocess(sar, crop=crop).to(device=device, dtype=self.vae.dtype)
+        # The SAR condition is encoded by the SAME frozen autoencoder that
+        # defines the EO latent space (evaluate.py:553-577).
+        z_s = self.vae.encode(sar_pm1).to(dtype)
+
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        # Design B: the bridge starts at t = 0 from pure Gaussian noise
+        # (bridge.py:409-433), NOT from the SAR latent.
+        latents = randn_tensor(z_s.shape, generator=generator, device=device, dtype=z_s.dtype)
+
+        for t in self.progress_bar(self.scheduler.timesteps):
+            timestep = t.expand(latents.shape[0])
+            velocity = self.transformer(latents, timestep, z_s, return_dict=False)[0]
+            if guidance_scale != 1.0:
+                # Two passes; the null branch is cond=None, which the transformer
+                # turns into an all-zero conditioning latent (bridge.py:530-535).
+                uncond = self.transformer(latents, timestep, None, return_dict=False)[0]
+                velocity = uncond + guidance_scale * (velocity - uncond)
+            latents = self.scheduler.step(velocity, t, latents, return_dict=False)[0]
+
+        image = self.vae.decode(latents.to(self.vae.dtype))
+        # `--denorm standard` (evaluate.py:292-295). The `legacy` C-DiffSET
+        # convention `(x + 0.5).clamp(0, 1)` is a 2x contrast stretch and must
+        # not be used with these numbers.
+        image = (image * 0.5 + 0.5).clamp(0.0, 1.0)
+
+        self.maybe_free_model_hooks()
+
+        if output_type == "pil":
+            image = self._to_pil(image)
+        elif output_type == "np":
+            image = image.permute(0, 2, 3, 1).float().cpu().numpy()
+
+        if not return_dict:
+            return (image,)
+        return ImagePipelineOutput(images=image)

qxs-saropt/autoencoder_flux2.py

@@ -0,0 +1,426 @@
+"""Frozen FLUX.2 autoencoder — the latent endpoint of ReFlowSET.
+
+ReFlowSET never fine-tunes this module: it is loaded once, frozen, and used to
+encode the SAR condition and to decode the sampled EO latent.  The released
+weights are the **Apache-2.0** FLUX.2-klein-base-4B copy of the autoencoder,
+re-keyed to the layout below (see ``scripts/convert_flux2_ae.py``).
+
+Three details of the checkpoint are non-standard for `diffusers` and are
+preserved exactly, because the file must load with ``strict=True``:
+
+* ``quant_conv`` lives **inside** ``encoder.*`` and is the last op of the
+  encoder forward; ``post_quant_conv`` lives **inside** ``decoder.*`` and is the
+  first op of the decoder forward.  `diffusers`' ``AutoencoderKL`` makes both
+  siblings of the encoder/decoder.
+* The latent normaliser is a real ``BatchNorm2d(128, affine=False)`` whose
+  running statistics ship in the checkpoint under ``bn.*`` — a per-channel mean
+  **and** variance, not a scalar ``scaling_factor``/``shift_factor``.  Its
+  epsilon is ``1e-4``, not torch's ``1e-5``.
+* ``encode`` returns the posterior **mean**; the log-variance chunk of the
+  encoder's moments is discarded, so encoding is deterministic and there is no
+  ``DiagonalGaussianDistribution`` and no ``.sample()``.
+
+The public latent is ``[B, 128, H/16, W/16]``: an 8x convolutional stride
+followed by a 2x2 space-to-depth pack that is part of the *autoencoder*, not of
+the transformer.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+
+def swish(x: Tensor) -> Tensor:
+    """``x * sigmoid(x)`` — the activation used throughout the FLUX.2 AE."""
+    return x * torch.sigmoid(x)
+
+
+class AttnBlock(nn.Module):
+    """Single-head self-attention over the spatial grid (head dim == channels)."""
+
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q, k, v = self.q(h_), self.k(h_), self.v(h_)
+        b, c, h, w = q.shape
+        # "b c h w -> b 1 (h w) c": ONE head whose head-dim is the full channel
+        # count (flux2_ae.py:70-73).
+        q = q.reshape(b, c, h * w).transpose(1, 2).unsqueeze(1).contiguous()
+        k = k.reshape(b, c, h * w).transpose(1, 2).unsqueeze(1).contiguous()
+        v = v.reshape(b, c, h * w).transpose(1, 2).unsqueeze(1).contiguous()
+        h_ = F.scaled_dot_product_attention(q, k, v)
+        return h_.squeeze(1).transpose(1, 2).reshape(b, c, h, w)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if in_channels != out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x: Tensor) -> Tensor:
+        h = self.conv1(swish(self.norm1(x)))
+        h = self.conv2(swish(self.norm2(h)))
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+        return x + h
+
+
+class Downsample(nn.Module):
+    """Stride-2 conv with FLUX's asymmetric ``(0, 1, 0, 1)`` pad (flux2_ae.py:111-121)."""
+
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.conv(F.pad(x, (0, 1, 0, 1), mode="constant", value=0))
+
+
+class Upsample(nn.Module):
+    """Nearest-neighbour 2x followed by a 3x3 conv (flux2_ae.py:124-132)."""
+
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.conv(F.interpolate(x, scale_factor=2.0, mode="nearest"))
+
+
+class Encoder(nn.Module):
+    """FLUX.2 encoder.  Emits ``2 * z_channels`` moments; ``quant_conv`` is internal."""
+
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        z_channels: int,
+    ) -> None:
+        super().__init__()
+        # Declared first so the checkpoint key is `encoder.quant_conv.*`
+        # (flux2_ae.py:146) — diffusers keeps quant_conv outside the encoder.
+        self.quant_conv = nn.Conv2d(2 * z_channels, 2 * z_channels, 1)
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        self.conv_in = nn.Conv2d(in_channels, ch, kernel_size=3, stride=1, padding=1)
+
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.down = nn.ModuleList()
+        block_in = ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(num_res_blocks):
+                block.append(ResnetBlock(block_in, block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            # Empty at every level in this checkpoint: attention exists only in
+            # `mid` (flux2_ae.py:162).  Kept so the forward guard is meaningful.
+            down.attn = nn.ModuleList()
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in)
+            self.down.append(down)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(block_in, block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(block_in, block_in)
+
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        h = self.mid.block_2(self.mid.attn_1(self.mid.block_1(hs[-1])))
+        h = self.conv_out(swish(self.norm_out(h)))
+        return self.quant_conv(h)  # last op of the encoder (flux2_ae.py:207)
+
+
+class Decoder(nn.Module):
+    """FLUX.2 decoder.  ``post_quant_conv`` is internal and runs first."""
+
+    def __init__(
+        self,
+        ch: int,
+        out_ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+    ) -> None:
+        super().__init__()
+        # Checkpoint key `decoder.post_quant_conv.*` (flux2_ae.py:223).
+        self.post_quant_conv = nn.Conv2d(z_channels, z_channels, 1)
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(block_in, block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(block_in, block_in)
+
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(num_res_blocks + 1):
+                block.append(ResnetBlock(block_in, block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = nn.ModuleList()  # empty in this checkpoint (flux2_ae.py:249)
+            if i_level != 0:
+                up.upsample = Upsample(block_in)
+            self.up.insert(0, up)  # prepend so `up.<i>` indexes by resolution level
+
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: Tensor) -> Tensor:
+        z = self.post_quant_conv(z)  # first op of the decoder (flux2_ae.py:267)
+        upscale_dtype = next(self.up.parameters()).dtype
+
+        h = self.conv_in(z)
+        h = self.mid.block_2(self.mid.attn_1(self.mid.block_1(h)))
+        h = h.to(upscale_dtype)
+
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        return self.conv_out(swish(self.norm_out(h)))
+
+
+class AutoencoderFlux2(ModelMixin, ConfigMixin):
+    """Frozen FLUX.2 autoencoder with ReFlowSET's packed, BN-normalised latent.
+
+    ``encode`` maps ``[B, 3, H, W]`` in ``[-1, 1]`` to ``[B, 128, H/16, W/16]``
+    and ``decode`` inverts it.  The module is frozen: ``train()`` is a no-op that
+    always selects eval mode, and the latent BatchNorm is additionally forced to
+    eval on every call so no batch statistic can ever leak into the latent.
+
+    Args:
+        resolution: Nominal training resolution of the original autoencoder.
+            Only used to size bookkeeping attributes; any ``H``, ``W`` divisible
+            by 16 may be encoded.
+        in_channels: Input image channels (3).
+        ch: Base width.
+        out_ch: Output image channels (3).
+        ch_mult: Per-level width multipliers; ``len(ch_mult) - 1`` downsamples.
+        num_res_blocks: Residual blocks per level.
+        z_channels: Pre-pack latent channels (32).
+        patch_size: Space-to-depth factor applied after the encoder (2), which
+            takes the latent from 32 channels at ``H/8`` to 128 at ``H/16``.
+        bn_eps: Epsilon of the latent BatchNorm.  **1e-4**, not torch's 1e-5
+            (flux2_ae.py:331); using 1e-5 shifts the latent by up to 2.6e-5.
+    """
+
+    _supports_gradient_checkpointing = False
+
+    @register_to_config
+    def __init__(
+        self,
+        resolution: int = 256,
+        in_channels: int = 3,
+        ch: int = 128,
+        out_ch: int = 3,
+        ch_mult: tuple[int, ...] = (1, 2, 4, 4),
+        num_res_blocks: int = 2,
+        z_channels: int = 32,
+        patch_size: int = 2,
+        bn_eps: float = 1e-4,
+    ) -> None:
+        super().__init__()
+        ch_mult = list(ch_mult)
+        self.encoder = Encoder(
+            resolution=resolution,
+            in_channels=in_channels,
+            ch=ch,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            z_channels=z_channels,
+        )
+        self.decoder = Decoder(
+            ch=ch,
+            out_ch=out_ch,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            in_channels=in_channels,
+            resolution=resolution,
+            z_channels=z_channels,
+        )
+        # Per-channel latent normaliser with the checkpoint's running statistics.
+        # affine=False, so there is no weight/bias to load (flux2_ae.py:334-340).
+        self.bn = nn.BatchNorm2d(
+            patch_size * patch_size * z_channels,
+            eps=bn_eps,
+            momentum=0.1,
+            affine=False,
+            track_running_stats=True,
+        )
+
+    @property
+    def latent_channels(self) -> int:
+        """Channels of the public latent: ``patch_size**2 * z_channels`` = 128."""
+        return self.config.patch_size**2 * self.config.z_channels
+
+    @property
+    def spatial_factor(self) -> int:
+        """Total stride: 8x convolutional times ``patch_size`` packing = 16."""
+        return 2 ** (len(self.config.ch_mult) - 1) * self.config.patch_size
+
+    # ---- 2x2 space-to-depth pack / unpack -----------------------------------
+
+    def pack(self, z: Tensor) -> Tensor:
+        """``[B, C, H, W] -> [B, C*p*p, H/p, W/p]``, channel-major.
+
+        Bit-identical to the reference ``rearrange("... c (i pi) (j pj) -> ...
+        (c pi pj) i j")`` (flux2_ae.py:349-357).  Note this is **not** diffusers'
+        ``_pack_latents``, whose channel grouping is transposed.
+        """
+        return F.pixel_unshuffle(z, self.config.patch_size)
+
+    def unpack(self, z: Tensor) -> Tensor:
+        """Exact inverse of :meth:`pack` (flux2_ae.py:359-367)."""
+        return F.pixel_shuffle(z, self.config.patch_size)
+
+    # ---- latent normalisation ----------------------------------------------
+
+    def normalize(self, z: Tensor) -> Tensor:
+        """``(z - running_mean) / sqrt(running_var + bn_eps)``, per channel."""
+        self.bn.eval()  # forced every call (flux2_ae.py:372); train mode shifts z by ~1.67
+        return self.bn(z)
+
+    def inv_normalize(self, z: Tensor) -> Tensor:
+        """Exact inverse of :meth:`normalize` — same ``bn_eps`` (flux2_ae.py:375-379)."""
+        self.bn.eval()
+        s = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.config.bn_eps)
+        m = self.bn.running_mean.view(1, -1, 1, 1)
+        return z * s + m
+
+    # ---- public API ---------------------------------------------------------
+
+    @torch.no_grad()
+    def encode(self, x: Tensor) -> Tensor:
+        """Encode an image to the packed, normalised latent.
+
+        Args:
+            x: ``[B, 3, H, W]`` in ``[-1, 1]``; ``H`` and ``W`` divisible by 16.
+
+        Returns:
+            ``[B, 128, H/16, W/16]`` — the posterior **mean**, packed and
+            BN-normalised.  The encoder's log-variance chunk is discarded
+            (flux2_ae.py:396), so this is deterministic: there is no posterior
+            distribution object and nothing to sample.
+        """
+        if x.ndim != 4 or x.shape[1] != self.config.in_channels:
+            raise ValueError(
+                f"encode expects [B, {self.config.in_channels}, H, W], got {tuple(x.shape)}"
+            )
+        h, w = x.shape[-2:]
+        if h % self.spatial_factor or w % self.spatial_factor:
+            raise ValueError(
+                f"encode requires H and W divisible by {self.spatial_factor}, got {h}x{w}"
+            )
+        moments = self.encoder(x)
+        mean = torch.chunk(moments, 2, dim=1)[0]
+        return self.normalize(self.pack(mean))
+
+    @torch.no_grad()
+    def decode(self, z: Tensor) -> Tensor:
+        """Decode a packed, normalised latent ``[B, 128, h, w]`` to ``[B, 3, 16h, 16w]``.
+
+        The output is approximately ``[-1, 1]`` and is **not** clamped here; the
+        pipeline applies ``(x * 0.5 + 0.5).clamp(0, 1)``.
+        """
+        if z.ndim != 4 or z.shape[1] != self.latent_channels:
+            raise ValueError(
+                f"decode expects [B, {self.latent_channels}, h, w], got {tuple(z.shape)}"
+            )
+        return self.decoder(self.unpack(self.inv_normalize(z)))
+
+    # ---- construction / freezing -------------------------------------------
+
+    @classmethod
+    def from_single_file(
+        cls,
+        path: str | os.PathLike,
+        torch_dtype: torch.dtype = torch.float32,
+    ) -> "AutoencoderFlux2":
+        """Load the single-file ``ae.safetensors`` (BFL key names) with ``strict=True``.
+
+        The released file is the Apache-2.0 FLUX.2-klein-base-4B autoencoder
+        re-keyed to this layout; it is stored in bfloat16 and is upcast to
+        ``torch_dtype``.  ReFlowSET runs the autoencoder in float32.
+        """
+        from safetensors.torch import load_file
+
+        path = os.fspath(path)
+        if not os.path.isfile(path):
+            raise FileNotFoundError(
+                f"FLUX.2 autoencoder weights not found at: {path}. Expected the "
+                "single-file 'ae.safetensors' shipped with ReFlowSET."
+            )
+        model = cls()
+        model.load_state_dict(load_file(path, device="cpu"), strict=True)
+        model.to(dtype=torch_dtype)
+        model.eval()
+        model.requires_grad_(False)
+        return model
+
+    def train(self, mode: bool = True) -> "AutoencoderFlux2":
+        """The autoencoder is frozen: never leave eval mode (flux2_ae.py:437-439)."""
+        return super().train(False)

qxs-saropt/model_index.json

@@ -0,0 +1,16 @@
+{
+  "_class_name": "ReFlowSETPipeline",
+  "_diffusers_version": "0.37.1",
+  "transformer": [
+    "transformer_reflowset",
+    "ReFlowSETTransformer2DModel"
+  ],
+  "vae": [
+    "autoencoder_flux2",
+    "AutoencoderFlux2"
+  ],
+  "scheduler": [
+    "scheduler_flow_bridge",
+    "FlowBridgeScheduler"
+  ]
+}

qxs-saropt/pipeline.py

@@ -0,0 +1,267 @@
+"""ReFlowSET SAR -> EO translation pipeline."""
+
+from __future__ import annotations
+
+import warnings
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from PIL import Image
+
+from .autoencoder_flux2 import AutoencoderFlux2
+from .scheduler_flow_bridge import FlowBridgeScheduler
+from .transformer_reflowset import ReFlowSETTransformer2DModel
+
+#: PIL modes the SAR loader accepts.  The reference loader calls ``np.array(im)``
+#: with no ``convert()`` (datasets.py:454, 574), so a 16-bit (``I;16``) or
+#: palette (``P``) raster would flow straight into ``x / 127.5 - 1`` and be
+#: badly out of range.  That is an unguarded trap upstream; it is guarded here.
+_ACCEPTED_SAR_MODES = ("L", "RGB", "RGBA")
+
+
+class ReFlowSETPipeline(DiffusionPipeline):
+    """Generate an EO image from a SAR image with ReFlowSET's flow bridge.
+
+    Args:
+        transformer: The velocity transformer.
+        vae: The frozen FLUX.2 autoencoder that defines the latent space.
+        scheduler: The Design-B flow-bridge Euler solver.
+
+    To reproduce the paper's numbers, sample at ``num_inference_steps=50``,
+    ``guidance_scale=1.5``, float32, one image per call, with a generator freshly
+    seeded to 2024 on the compute device before each call — every test image in
+    the reported evaluation starts from the same seeded noise draw, and CPU-drawn
+    noise does not reproduce a CUDA draw.
+    """
+
+    model_cpu_offload_seq = "transformer->vae"
+
+    def __init__(
+        self,
+        transformer: ReFlowSETTransformer2DModel,
+        vae: AutoencoderFlux2,
+        scheduler: FlowBridgeScheduler,
+    ) -> None:
+        super().__init__()
+        self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)
+
+    # ---- preprocessing (datasets.py:124-205, 452-474, 572-589) --------------
+
+    @staticmethod
+    def _sar_hwc(raster: Union[Image.Image, np.ndarray]) -> np.ndarray:
+        """SAR raster -> ``[H, W, C]`` float32 in ``[0, 255]``, collapsed to 1 channel."""
+        if isinstance(raster, Image.Image):
+            if raster.mode not in _ACCEPTED_SAR_MODES:
+                raise ValueError(
+                    f"SAR image mode {raster.mode!r} is not an 8-bit display raster; expected "
+                    f"one of {_ACCEPTED_SAR_MODES}. ReFlowSET was trained on 8-bit display "
+                    "quicklooks (sar_value_domain='display_png'); convert with .convert('L') "
+                    "and be aware that the contrast stretch you choose is part of the input."
+                )
+            # No .convert() on the SAR side, matching datasets.py:454, 574.
+            arr = np.array(raster)
+        else:
+            arr = np.asarray(raster)
+        if arr.ndim == 2:  # PIL mode "L" (datasets.py:171-172)
+            arr = arr[:, :, None]
+        if arr.shape[-1] == 4:  # drop a container alpha channel (datasets.py:173-174)
+            arr = arr[..., :3]
+        arr = arr.astype(np.float32)
+        if arr.shape[-1] > 1:
+            # Exact-equality test, tol=0.0 (datasets.py:100-111): a display RGB
+            # quicklook collapses to its single amplitude channel.
+            if np.abs(arr - arr[..., :1]).max() == 0.0:
+                arr = arr[..., :1]
+            else:
+                warnings.warn(
+                    "SAR raster has non-identical colour channels; feeding all 3 to the "
+                    "frozen encoder. The released arms were trained on single-channel "
+                    "amplitude quicklooks, so this is an undeclared input.",
+                    RuntimeWarning,
+                    stacklevel=3,
+                )
+        return arr
+
+    @staticmethod
+    def _center_crop(arr: np.ndarray, crop: int) -> np.ndarray:
+        """Center-crop ``[H, W, C]`` to ``crop`` — **never** resize (datasets.py:145-165).
+
+        The offsets are albumentations' ``CenterCrop`` arithmetic ``(n - c) // 2``:
+        the SAR2Opt protocol takes the central 512 of 600 at offset 44.
+        """
+        h, w = arr.shape[:2]
+        if h < crop or w < crop:
+            raise ValueError(
+                f"image {h}x{w} is smaller than the requested crop {crop}; ReFlowSET never "
+                "upscales an input"
+            )
+        top, left = (h - crop) // 2, (w - crop) // 2
+        return arr[top : top + crop, left : left + crop]
+
+    def preprocess(
+        self,
+        sar: Union[Image.Image, np.ndarray, torch.Tensor, list],
+        crop: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Build the model-boundary SAR tensor ``[B, 3, H, W]`` in ``[-1, 1]``.
+
+        Args:
+            sar: A PIL image, a list of PIL images, an ``[H, W]`` / ``[H, W, C]``
+                uint8 array, or a float tensor already in ``[-1, 1]`` shaped
+                ``[H, W]``, ``[C, H, W]`` or ``[B, C, H, W]``.
+            crop: Center-crop size applied before normalisation.  ``None``
+                center-crops to the arm's own training resolution when the
+                raster is larger and not already a multiple of the latent
+                stride -- which is exactly the SAR2Opt 600 -> 512 protocol the
+                reported numbers use.  Pass an explicit size to override, or
+                ``0`` to keep the native raster and fail loudly if it does not
+                fit.
+
+        Images are read as 8-bit display rasters and mapped to ``[-1, 1]`` by
+        ``x / 127.5 - 1`` (datasets.py:124-126) with no per-image statistics, no
+        percentile stretch and no resize.  The single SAR channel is then
+        replicated to 3 at the model boundary (evaluate.py:566-570), because the
+        frozen FLUX.2 encoder is the same one that encodes EO — ReFlowSET has no
+        separate SAR encoder.
+        """
+        if crop == 0:
+            crop = None
+        elif crop is None:
+            # Fall back to the resolution this arm was trained at. Cropping is
+            # the protocol (train.py random-crops, evaluate.py center-crops);
+            # ReFlowSET never resizes, so an un-croppable raster is an error
+            # rather than something to silently rescale.
+            crop = self.transformer.config.sample_size
+
+        if isinstance(sar, torch.Tensor):
+            x = sar.float()
+            if x.ndim == 2:
+                x = x[None, None]
+            elif x.ndim == 3:
+                x = x[None]
+            elif x.ndim != 4:
+                raise ValueError(f"SAR tensor must have 2, 3 or 4 dims, got {tuple(sar.shape)}")
+            h, w = x.shape[-2:]
+            if crop is not None and (h, w) != (crop, crop):
+                if h < crop or w < crop:
+                    raise ValueError(f"tensor {h}x{w} is smaller than the requested crop {crop}")
+                top, left = (h - crop) // 2, (w - crop) // 2
+                x = x[..., top : top + crop, left : left + crop]
+        else:
+            images = sar if isinstance(sar, list) else [sar]
+            arrays = []
+            for item in images:
+                if not isinstance(item, (Image.Image, np.ndarray)):
+                    raise TypeError(f"unsupported SAR input type {type(item)!r}")
+                arr = self._sar_hwc(item)
+                if crop is not None and arr.shape[:2] != (crop, crop):
+                    arr = self._center_crop(arr, crop)
+                arrays.append(np.ascontiguousarray(arr.transpose(2, 0, 1)))
+            x = torch.from_numpy(np.stack(arrays)) / 127.5 - 1.0
+
+        # Train-side clamp (train.py:770); a no-op on 8-bit input, which maps
+        # exactly onto [-1, 1].
+        x = x.clamp(-1.0, 1.0)
+        if x.shape[1] == 1:
+            x = x.repeat(1, 3, 1, 1)
+        elif x.shape[1] != 3:
+            raise ValueError(
+                f"the frozen FLUX.2 encoder takes 1 or 3 SAR channels, got {x.shape[1]}"
+            )
+        factor = self.vae.spatial_factor
+        if x.shape[-2] % factor or x.shape[-1] % factor:
+            raise ValueError(
+                f"SAR size {x.shape[-2]}x{x.shape[-1]} must be divisible by {factor}; pass "
+                "crop= to center-crop (ReFlowSET never resizes)"
+            )
+        return x
+
+    # ---- postprocessing -----------------------------------------------------
+
+    @staticmethod
+    def _to_pil(images: torch.Tensor) -> list[Image.Image]:
+        """``[B, 3, H, W]`` in ``[0, 1]`` -> PIL, quantised round-half-up.
+
+        ``255 * x + 0.5`` truncated is what ``torchvision.utils.save_image``
+        does and is therefore what the released PNGs contain; numpy's
+        ``round()`` is banker's rounding and would differ on exact halves.
+        """
+        arr = (images * 255 + 0.5).clamp(0, 255).to(torch.uint8)
+        arr = arr.permute(0, 2, 3, 1).cpu().numpy()
+        return [Image.fromarray(a) for a in arr]
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        sar: Union[Image.Image, np.ndarray, torch.Tensor, list],
+        num_inference_steps: int = 50,
+        guidance_scale: float = 1.5,
+        generator: Optional[Union[torch.Generator, list[torch.Generator]]] = None,
+        output_type: str = "pil",
+        crop: Optional[int] = None,
+        return_dict: bool = True,
+    ) -> Union[ImagePipelineOutput, tuple[list]]:
+        """Translate a SAR image into an EO image.
+
+        Args:
+            sar: SAR input; see :meth:`preprocess`.
+            num_inference_steps: NFE, the number of velocity evaluations.  The
+                paper's main results are NFE 50; NFE 4 is the efficiency
+                operating point and trades FID for PSNR/SSIM, so the two must
+                not be mixed in one comparison.
+            guidance_scale: Classifier-free guidance scale.  1.5 is the published
+                setting; 1.0 disables guidance and halves the cost.
+            generator: Generator for the initial noise.  Create it on the compute
+                device — CPU-drawn noise does not reproduce a CUDA draw.
+            output_type: ``"pil"``, ``"np"`` or ``"pt"``.
+            crop: Center-crop size applied to the SAR input before encoding.
+            return_dict: Return an ``ImagePipelineOutput`` instead of a tuple.
+
+        Returns:
+            The generated EO image(s) in ``[0, 1]`` (or as PIL).
+        """
+        if output_type not in ("pil", "np", "pt"):
+            raise ValueError(f"output_type must be 'pil', 'np' or 'pt', got {output_type!r}")
+
+        device = self._execution_device
+        dtype = self.transformer.dtype
+
+        sar_pm1 = self.preprocess(sar, crop=crop).to(device=device, dtype=self.vae.dtype)
+        # The SAR condition is encoded by the SAME frozen autoencoder that
+        # defines the EO latent space (evaluate.py:553-577).
+        z_s = self.vae.encode(sar_pm1).to(dtype)
+
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        # Design B: the bridge starts at t = 0 from pure Gaussian noise
+        # (bridge.py:409-433), NOT from the SAR latent.
+        latents = randn_tensor(z_s.shape, generator=generator, device=device, dtype=z_s.dtype)
+
+        for t in self.progress_bar(self.scheduler.timesteps):
+            timestep = t.expand(latents.shape[0])
+            velocity = self.transformer(latents, timestep, z_s, return_dict=False)[0]
+            if guidance_scale != 1.0:
+                # Two passes; the null branch is cond=None, which the transformer
+                # turns into an all-zero conditioning latent (bridge.py:530-535).
+                uncond = self.transformer(latents, timestep, None, return_dict=False)[0]
+                velocity = uncond + guidance_scale * (velocity - uncond)
+            latents = self.scheduler.step(velocity, t, latents, return_dict=False)[0]
+
+        image = self.vae.decode(latents.to(self.vae.dtype))
+        # `--denorm standard` (evaluate.py:292-295). The `legacy` C-DiffSET
+        # convention `(x + 0.5).clamp(0, 1)` is a 2x contrast stretch and must
+        # not be used with these numbers.
+        image = (image * 0.5 + 0.5).clamp(0.0, 1.0)
+
+        self.maybe_free_model_hooks()
+
+        if output_type == "pil":
+            image = self._to_pil(image)
+        elif output_type == "np":
+            image = image.permute(0, 2, 3, 1).float().cpu().numpy()
+
+        if not return_dict:
+            return (image,)
+        return ImagePipelineOutput(images=image)

qxs-saropt/pipeline_reflowset.py

@@ -0,0 +1,267 @@
+"""ReFlowSET SAR -> EO translation pipeline."""
+
+from __future__ import annotations
+
+import warnings
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from PIL import Image
+
+from .autoencoder_flux2 import AutoencoderFlux2
+from .scheduler_flow_bridge import FlowBridgeScheduler
+from .transformer_reflowset import ReFlowSETTransformer2DModel
+
+#: PIL modes the SAR loader accepts.  The reference loader calls ``np.array(im)``
+#: with no ``convert()`` (datasets.py:454, 574), so a 16-bit (``I;16``) or
+#: palette (``P``) raster would flow straight into ``x / 127.5 - 1`` and be
+#: badly out of range.  That is an unguarded trap upstream; it is guarded here.
+_ACCEPTED_SAR_MODES = ("L", "RGB", "RGBA")
+
+
+class ReFlowSETPipeline(DiffusionPipeline):
+    """Generate an EO image from a SAR image with ReFlowSET's flow bridge.
+
+    Args:
+        transformer: The velocity transformer.
+        vae: The frozen FLUX.2 autoencoder that defines the latent space.
+        scheduler: The Design-B flow-bridge Euler solver.
+
+    To reproduce the paper's numbers, sample at ``num_inference_steps=50``,
+    ``guidance_scale=1.5``, float32, one image per call, with a generator freshly
+    seeded to 2024 on the compute device before each call — every test image in
+    the reported evaluation starts from the same seeded noise draw, and CPU-drawn
+    noise does not reproduce a CUDA draw.
+    """
+
+    model_cpu_offload_seq = "transformer->vae"
+
+    def __init__(
+        self,
+        transformer: ReFlowSETTransformer2DModel,
+        vae: AutoencoderFlux2,
+        scheduler: FlowBridgeScheduler,
+    ) -> None:
+        super().__init__()
+        self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)
+
+    # ---- preprocessing (datasets.py:124-205, 452-474, 572-589) --------------
+
+    @staticmethod
+    def _sar_hwc(raster: Union[Image.Image, np.ndarray]) -> np.ndarray:
+        """SAR raster -> ``[H, W, C]`` float32 in ``[0, 255]``, collapsed to 1 channel."""
+        if isinstance(raster, Image.Image):
+            if raster.mode not in _ACCEPTED_SAR_MODES:
+                raise ValueError(
+                    f"SAR image mode {raster.mode!r} is not an 8-bit display raster; expected "
+                    f"one of {_ACCEPTED_SAR_MODES}. ReFlowSET was trained on 8-bit display "
+                    "quicklooks (sar_value_domain='display_png'); convert with .convert('L') "
+                    "and be aware that the contrast stretch you choose is part of the input."
+                )
+            # No .convert() on the SAR side, matching datasets.py:454, 574.
+            arr = np.array(raster)
+        else:
+            arr = np.asarray(raster)
+        if arr.ndim == 2:  # PIL mode "L" (datasets.py:171-172)
+            arr = arr[:, :, None]
+        if arr.shape[-1] == 4:  # drop a container alpha channel (datasets.py:173-174)
+            arr = arr[..., :3]
+        arr = arr.astype(np.float32)
+        if arr.shape[-1] > 1:
+            # Exact-equality test, tol=0.0 (datasets.py:100-111): a display RGB
+            # quicklook collapses to its single amplitude channel.
+            if np.abs(arr - arr[..., :1]).max() == 0.0:
+                arr = arr[..., :1]
+            else:
+                warnings.warn(
+                    "SAR raster has non-identical colour channels; feeding all 3 to the "
+                    "frozen encoder. The released arms were trained on single-channel "
+                    "amplitude quicklooks, so this is an undeclared input.",
+                    RuntimeWarning,
+                    stacklevel=3,
+                )
+        return arr
+
+    @staticmethod
+    def _center_crop(arr: np.ndarray, crop: int) -> np.ndarray:
+        """Center-crop ``[H, W, C]`` to ``crop`` — **never** resize (datasets.py:145-165).
+
+        The offsets are albumentations' ``CenterCrop`` arithmetic ``(n - c) // 2``:
+        the SAR2Opt protocol takes the central 512 of 600 at offset 44.
+        """
+        h, w = arr.shape[:2]
+        if h < crop or w < crop:
+            raise ValueError(
+                f"image {h}x{w} is smaller than the requested crop {crop}; ReFlowSET never "
+                "upscales an input"
+            )
+        top, left = (h - crop) // 2, (w - crop) // 2
+        return arr[top : top + crop, left : left + crop]
+
+    def preprocess(
+        self,
+        sar: Union[Image.Image, np.ndarray, torch.Tensor, list],
+        crop: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Build the model-boundary SAR tensor ``[B, 3, H, W]`` in ``[-1, 1]``.
+
+        Args:
+            sar: A PIL image, a list of PIL images, an ``[H, W]`` / ``[H, W, C]``
+                uint8 array, or a float tensor already in ``[-1, 1]`` shaped
+                ``[H, W]``, ``[C, H, W]`` or ``[B, C, H, W]``.
+            crop: Center-crop size applied before normalisation.  ``None``
+                center-crops to the arm's own training resolution when the
+                raster is larger and not already a multiple of the latent
+                stride -- which is exactly the SAR2Opt 600 -> 512 protocol the
+                reported numbers use.  Pass an explicit size to override, or
+                ``0`` to keep the native raster and fail loudly if it does not
+                fit.
+
+        Images are read as 8-bit display rasters and mapped to ``[-1, 1]`` by
+        ``x / 127.5 - 1`` (datasets.py:124-126) with no per-image statistics, no
+        percentile stretch and no resize.  The single SAR channel is then
+        replicated to 3 at the model boundary (evaluate.py:566-570), because the
+        frozen FLUX.2 encoder is the same one that encodes EO — ReFlowSET has no
+        separate SAR encoder.
+        """
+        if crop == 0:
+            crop = None
+        elif crop is None:
+            # Fall back to the resolution this arm was trained at. Cropping is
+            # the protocol (train.py random-crops, evaluate.py center-crops);
+            # ReFlowSET never resizes, so an un-croppable raster is an error
+            # rather than something to silently rescale.
+            crop = self.transformer.config.sample_size
+
+        if isinstance(sar, torch.Tensor):
+            x = sar.float()
+            if x.ndim == 2:
+                x = x[None, None]
+            elif x.ndim == 3:
+                x = x[None]
+            elif x.ndim != 4:
+                raise ValueError(f"SAR tensor must have 2, 3 or 4 dims, got {tuple(sar.shape)}")
+            h, w = x.shape[-2:]
+            if crop is not None and (h, w) != (crop, crop):
+                if h < crop or w < crop:
+                    raise ValueError(f"tensor {h}x{w} is smaller than the requested crop {crop}")
+                top, left = (h - crop) // 2, (w - crop) // 2
+                x = x[..., top : top + crop, left : left + crop]
+        else:
+            images = sar if isinstance(sar, list) else [sar]
+            arrays = []
+            for item in images:
+                if not isinstance(item, (Image.Image, np.ndarray)):
+                    raise TypeError(f"unsupported SAR input type {type(item)!r}")
+                arr = self._sar_hwc(item)
+                if crop is not None and arr.shape[:2] != (crop, crop):
+                    arr = self._center_crop(arr, crop)
+                arrays.append(np.ascontiguousarray(arr.transpose(2, 0, 1)))
+            x = torch.from_numpy(np.stack(arrays)) / 127.5 - 1.0
+
+        # Train-side clamp (train.py:770); a no-op on 8-bit input, which maps
+        # exactly onto [-1, 1].
+        x = x.clamp(-1.0, 1.0)
+        if x.shape[1] == 1:
+            x = x.repeat(1, 3, 1, 1)
+        elif x.shape[1] != 3:
+            raise ValueError(
+                f"the frozen FLUX.2 encoder takes 1 or 3 SAR channels, got {x.shape[1]}"
+            )
+        factor = self.vae.spatial_factor
+        if x.shape[-2] % factor or x.shape[-1] % factor:
+            raise ValueError(
+                f"SAR size {x.shape[-2]}x{x.shape[-1]} must be divisible by {factor}; pass "
+                "crop= to center-crop (ReFlowSET never resizes)"
+            )
+        return x
+
+    # ---- postprocessing -----------------------------------------------------
+
+    @staticmethod
+    def _to_pil(images: torch.Tensor) -> list[Image.Image]:
+        """``[B, 3, H, W]`` in ``[0, 1]`` -> PIL, quantised round-half-up.
+
+        ``255 * x + 0.5`` truncated is what ``torchvision.utils.save_image``
+        does and is therefore what the released PNGs contain; numpy's
+        ``round()`` is banker's rounding and would differ on exact halves.
+        """
+        arr = (images * 255 + 0.5).clamp(0, 255).to(torch.uint8)
+        arr = arr.permute(0, 2, 3, 1).cpu().numpy()
+        return [Image.fromarray(a) for a in arr]
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        sar: Union[Image.Image, np.ndarray, torch.Tensor, list],
+        num_inference_steps: int = 50,
+        guidance_scale: float = 1.5,
+        generator: Optional[Union[torch.Generator, list[torch.Generator]]] = None,
+        output_type: str = "pil",
+        crop: Optional[int] = None,
+        return_dict: bool = True,
+    ) -> Union[ImagePipelineOutput, tuple[list]]:
+        """Translate a SAR image into an EO image.
+
+        Args:
+            sar: SAR input; see :meth:`preprocess`.
+            num_inference_steps: NFE, the number of velocity evaluations.  The
+                paper's main results are NFE 50; NFE 4 is the efficiency
+                operating point and trades FID for PSNR/SSIM, so the two must
+                not be mixed in one comparison.
+            guidance_scale: Classifier-free guidance scale.  1.5 is the published
+                setting; 1.0 disables guidance and halves the cost.
+            generator: Generator for the initial noise.  Create it on the compute
+                device — CPU-drawn noise does not reproduce a CUDA draw.
+            output_type: ``"pil"``, ``"np"`` or ``"pt"``.
+            crop: Center-crop size applied to the SAR input before encoding.
+            return_dict: Return an ``ImagePipelineOutput`` instead of a tuple.
+
+        Returns:
+            The generated EO image(s) in ``[0, 1]`` (or as PIL).
+        """
+        if output_type not in ("pil", "np", "pt"):
+            raise ValueError(f"output_type must be 'pil', 'np' or 'pt', got {output_type!r}")
+
+        device = self._execution_device
+        dtype = self.transformer.dtype
+
+        sar_pm1 = self.preprocess(sar, crop=crop).to(device=device, dtype=self.vae.dtype)
+        # The SAR condition is encoded by the SAME frozen autoencoder that
+        # defines the EO latent space (evaluate.py:553-577).
+        z_s = self.vae.encode(sar_pm1).to(dtype)
+
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        # Design B: the bridge starts at t = 0 from pure Gaussian noise
+        # (bridge.py:409-433), NOT from the SAR latent.
+        latents = randn_tensor(z_s.shape, generator=generator, device=device, dtype=z_s.dtype)
+
+        for t in self.progress_bar(self.scheduler.timesteps):
+            timestep = t.expand(latents.shape[0])
+            velocity = self.transformer(latents, timestep, z_s, return_dict=False)[0]
+            if guidance_scale != 1.0:
+                # Two passes; the null branch is cond=None, which the transformer
+                # turns into an all-zero conditioning latent (bridge.py:530-535).
+                uncond = self.transformer(latents, timestep, None, return_dict=False)[0]
+                velocity = uncond + guidance_scale * (velocity - uncond)
+            latents = self.scheduler.step(velocity, t, latents, return_dict=False)[0]
+
+        image = self.vae.decode(latents.to(self.vae.dtype))
+        # `--denorm standard` (evaluate.py:292-295). The `legacy` C-DiffSET
+        # convention `(x + 0.5).clamp(0, 1)` is a 2x contrast stretch and must
+        # not be used with these numbers.
+        image = (image * 0.5 + 0.5).clamp(0.0, 1.0)
+
+        self.maybe_free_model_hooks()
+
+        if output_type == "pil":
+            image = self._to_pil(image)
+        elif output_type == "np":
+            image = image.permute(0, 2, 3, 1).float().cpu().numpy()
+
+        if not return_dict:
+            return (image,)
+        return ImagePipelineOutput(images=image)

qxs-saropt/scheduler/scheduler_config.json

@@ -0,0 +1,5 @@
+{
+  "_class_name": "FlowBridgeScheduler",
+  "_diffusers_version": "0.37.1",
+  "t_end": 1.0
+}

qxs-saropt/scheduler/scheduler_flow_bridge.py

@@ -0,0 +1,176 @@
+"""ReFlowSET's Design-B flow bridge and its explicit-Euler solver.
+
+Forward (training) process, with ``eps ~ N(0, I)`` and ``z_e`` the EO latent::
+
+    z_t = (1 - t) * eps + t * z_e          (bridge.py:311, sigma_b = 0)
+    u*  = z_e - eps                        (bridge.py:328 at sigma_b = 0)
+
+Sampling starts from ``z_0 ~ N(0, I)`` and integrates the predicted velocity
+with explicit Euler on a uniform grid ``linspace(0, t_end, nfe + 1)``
+(bridge.py:519, 536).  The bridge is deterministic: ``sigma_b = 0``, so no
+stochastic term ever executes, and the only randomness in a sample is the
+initial noise draw.
+
+**Time direction.**  ``t = 0`` is NOISE and ``t = 1`` is DATA, and the solver
+integrates ``t`` **ascending** (bridge.py:86-88).  That is the opposite of
+`diffusers`' ``sigma`` convention: setting ``sigma := 1 - t`` recovers
+``FlowMatchEulerDiscreteScheduler``'s interpolation, but then this bridge's
+velocity is the **negative** of the diffusers flow-matching target and the
+network must still be fed ``1 - sigma``.  This scheduler keeps ReFlowSET's own
+sign and direction so neither flip is needed; ``timesteps`` therefore *increase*
+from 0 towards 1, unlike every noise-schedule scheduler in `diffusers`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class FlowBridgeSchedulerOutput(BaseOutput):
+    """Output of :meth:`FlowBridgeScheduler.step`.
+
+    Args:
+        prev_sample: The bridge state at the next time on the grid.
+    """
+
+    prev_sample: torch.Tensor
+
+
+class FlowBridgeScheduler(SchedulerMixin, ConfigMixin):
+    """Explicit-Euler solver for ReFlowSET's Design-B flow bridge.
+
+    Args:
+        t_end: End time of the integration grid (1.0 — the EO endpoint).  The
+            model is evaluated at ``linspace(0, t_end, nfe + 1)[:-1]`` and the
+            final Euler step lands on ``t_end``; the network is never queried at
+            ``t = t_end``.
+    """
+
+    order = 1
+
+    @register_to_config
+    def __init__(self, t_end: float = 1.0) -> None:
+        if not 0.0 < t_end <= 1.0:
+            raise ValueError(f"t_end must lie in (0, 1], got {t_end}")
+        self._grid: Optional[torch.Tensor] = None
+        self._step_index: Optional[int] = None
+        self.num_inference_steps: Optional[int] = None
+
+    @property
+    def timesteps(self) -> torch.Tensor:
+        """The ``nfe`` bridge times at which the model is evaluated, ascending."""
+        if self._grid is None:
+            raise ValueError("call set_timesteps() before reading timesteps")
+        return self._grid[:-1]
+
+    @property
+    def step_index(self) -> Optional[int]:
+        """Index of the next grid interval; ``None`` until the first :meth:`step`."""
+        return self._step_index
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        device: Optional[Union[str, torch.device]] = None,
+    ) -> None:
+        """Build the uniform grid ``linspace(0, t_end, num_inference_steps + 1)``.
+
+        Args:
+            num_inference_steps: NFE — the number of velocity evaluations.
+                50 reproduces the paper's main results; 4 is the efficiency
+                operating point.
+            device: Device the grid is built on.
+
+        There is no shift, no dynamic shifting, no Karras or exponential
+        spacing, and no timestep-spacing option: the reference solver uses a
+        plain uniform grid (bridge.py:519).
+        """
+        if num_inference_steps < 1:
+            raise ValueError(f"num_inference_steps must be >= 1, got {num_inference_steps}")
+        self.num_inference_steps = num_inference_steps
+        self._grid = torch.linspace(
+            0.0, self.config.t_end, num_inference_steps + 1, device=device, dtype=torch.float32
+        )
+        self._step_index = 0
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[float, torch.Tensor],
+        sample: torch.Tensor,
+        return_dict: bool = True,
+    ) -> Union[FlowBridgeSchedulerOutput, tuple[torch.Tensor]]:
+        """One explicit-Euler step: ``z + (t_next - t_cur) * v`` (bridge.py:536).
+
+        Args:
+            model_output: The predicted velocity ``dz/dt`` at ``timestep``,
+                already classifier-free-guided by the caller.
+            timestep: The current bridge time.  Present for API compatibility and
+                checked against the grid; the step size comes from the grid.
+            sample: The current bridge state.
+            return_dict: Return a :class:`FlowBridgeSchedulerOutput` instead of a
+                tuple.
+
+        Steps must be taken in order, starting from the first entry of
+        :attr:`timesteps`.
+        """
+        if self._grid is None or self._step_index is None:
+            raise ValueError("call set_timesteps() before step()")
+        if self._step_index >= self.num_inference_steps:
+            raise ValueError(
+                f"already took {self.num_inference_steps} steps; call set_timesteps() again"
+            )
+        t_cur, t_next = self._grid[self._step_index], self._grid[self._step_index + 1]
+        if not torch.isclose(torch.as_tensor(timestep, dtype=torch.float32).to(t_cur.device), t_cur):
+            raise ValueError(
+                f"step {self._step_index} expects timestep {t_cur.item()}, got {float(timestep)}; "
+                "the flow bridge must be integrated in ascending grid order"
+            )
+
+        # The state is carried in float32 even if the model ran lower (bridge.py:515-517).
+        dtype = sample.dtype if sample.dtype in (torch.float32, torch.float64) else torch.float32
+        prev_sample = sample.to(dtype) + (t_next - t_cur) * model_output.to(dtype)
+        prev_sample = prev_sample.to(sample.dtype)
+
+        self._step_index += 1
+        if not return_dict:
+            return (prev_sample,)
+        return FlowBridgeSchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        """The training-side bridge state ``z_t = (1 - t) * eps + t * z_e`` (bridge.py:311).
+
+        Args:
+            original_samples: The EO latent ``z_e`` (the ``t = 1`` endpoint).
+            noise: ``eps ~ N(0, I)`` (the ``t = 0`` endpoint).
+            timesteps: Bridge times in ``[0, 1]``, broadcastable over the batch.
+        """
+        t = timesteps.to(original_samples.device, original_samples.dtype)
+        t = t.view(-1, *([1] * (original_samples.ndim - 1)))
+        return (1.0 - t) * noise + t * original_samples
+
+    def get_velocity(
+        self,
+        sample: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        """The training target ``u* = z_e - eps`` (bridge.py:328 at ``sigma_b = 0``).
+
+        Constant along the path, hence independent of ``timesteps``; the argument
+        is kept for `diffusers` API compatibility.
+        """
+        del timesteps
+        return sample - noise

qxs-saropt/scheduler_flow_bridge.py

@@ -0,0 +1,176 @@
+"""ReFlowSET's Design-B flow bridge and its explicit-Euler solver.
+
+Forward (training) process, with ``eps ~ N(0, I)`` and ``z_e`` the EO latent::
+
+    z_t = (1 - t) * eps + t * z_e          (bridge.py:311, sigma_b = 0)
+    u*  = z_e - eps                        (bridge.py:328 at sigma_b = 0)
+
+Sampling starts from ``z_0 ~ N(0, I)`` and integrates the predicted velocity
+with explicit Euler on a uniform grid ``linspace(0, t_end, nfe + 1)``
+(bridge.py:519, 536).  The bridge is deterministic: ``sigma_b = 0``, so no
+stochastic term ever executes, and the only randomness in a sample is the
+initial noise draw.
+
+**Time direction.**  ``t = 0`` is NOISE and ``t = 1`` is DATA, and the solver
+integrates ``t`` **ascending** (bridge.py:86-88).  That is the opposite of
+`diffusers`' ``sigma`` convention: setting ``sigma := 1 - t`` recovers
+``FlowMatchEulerDiscreteScheduler``'s interpolation, but then this bridge's
+velocity is the **negative** of the diffusers flow-matching target and the
+network must still be fed ``1 - sigma``.  This scheduler keeps ReFlowSET's own
+sign and direction so neither flip is needed; ``timesteps`` therefore *increase*
+from 0 towards 1, unlike every noise-schedule scheduler in `diffusers`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class FlowBridgeSchedulerOutput(BaseOutput):
+    """Output of :meth:`FlowBridgeScheduler.step`.
+
+    Args:
+        prev_sample: The bridge state at the next time on the grid.
+    """
+
+    prev_sample: torch.Tensor
+
+
+class FlowBridgeScheduler(SchedulerMixin, ConfigMixin):
+    """Explicit-Euler solver for ReFlowSET's Design-B flow bridge.
+
+    Args:
+        t_end: End time of the integration grid (1.0 — the EO endpoint).  The
+            model is evaluated at ``linspace(0, t_end, nfe + 1)[:-1]`` and the
+            final Euler step lands on ``t_end``; the network is never queried at
+            ``t = t_end``.
+    """
+
+    order = 1
+
+    @register_to_config
+    def __init__(self, t_end: float = 1.0) -> None:
+        if not 0.0 < t_end <= 1.0:
+            raise ValueError(f"t_end must lie in (0, 1], got {t_end}")
+        self._grid: Optional[torch.Tensor] = None
+        self._step_index: Optional[int] = None
+        self.num_inference_steps: Optional[int] = None
+
+    @property
+    def timesteps(self) -> torch.Tensor:
+        """The ``nfe`` bridge times at which the model is evaluated, ascending."""
+        if self._grid is None:
+            raise ValueError("call set_timesteps() before reading timesteps")
+        return self._grid[:-1]
+
+    @property
+    def step_index(self) -> Optional[int]:
+        """Index of the next grid interval; ``None`` until the first :meth:`step`."""
+        return self._step_index
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        device: Optional[Union[str, torch.device]] = None,
+    ) -> None:
+        """Build the uniform grid ``linspace(0, t_end, num_inference_steps + 1)``.
+
+        Args:
+            num_inference_steps: NFE — the number of velocity evaluations.
+                50 reproduces the paper's main results; 4 is the efficiency
+                operating point.
+            device: Device the grid is built on.
+
+        There is no shift, no dynamic shifting, no Karras or exponential
+        spacing, and no timestep-spacing option: the reference solver uses a
+        plain uniform grid (bridge.py:519).
+        """
+        if num_inference_steps < 1:
+            raise ValueError(f"num_inference_steps must be >= 1, got {num_inference_steps}")
+        self.num_inference_steps = num_inference_steps
+        self._grid = torch.linspace(
+            0.0, self.config.t_end, num_inference_steps + 1, device=device, dtype=torch.float32
+        )
+        self._step_index = 0
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[float, torch.Tensor],
+        sample: torch.Tensor,
+        return_dict: bool = True,
+    ) -> Union[FlowBridgeSchedulerOutput, tuple[torch.Tensor]]:
+        """One explicit-Euler step: ``z + (t_next - t_cur) * v`` (bridge.py:536).
+
+        Args:
+            model_output: The predicted velocity ``dz/dt`` at ``timestep``,
+                already classifier-free-guided by the caller.
+            timestep: The current bridge time.  Present for API compatibility and
+                checked against the grid; the step size comes from the grid.
+            sample: The current bridge state.
+            return_dict: Return a :class:`FlowBridgeSchedulerOutput` instead of a
+                tuple.
+
+        Steps must be taken in order, starting from the first entry of
+        :attr:`timesteps`.
+        """
+        if self._grid is None or self._step_index is None:
+            raise ValueError("call set_timesteps() before step()")
+        if self._step_index >= self.num_inference_steps:
+            raise ValueError(
+                f"already took {self.num_inference_steps} steps; call set_timesteps() again"
+            )
+        t_cur, t_next = self._grid[self._step_index], self._grid[self._step_index + 1]
+        if not torch.isclose(torch.as_tensor(timestep, dtype=torch.float32).to(t_cur.device), t_cur):
+            raise ValueError(
+                f"step {self._step_index} expects timestep {t_cur.item()}, got {float(timestep)}; "
+                "the flow bridge must be integrated in ascending grid order"
+            )
+
+        # The state is carried in float32 even if the model ran lower (bridge.py:515-517).
+        dtype = sample.dtype if sample.dtype in (torch.float32, torch.float64) else torch.float32
+        prev_sample = sample.to(dtype) + (t_next - t_cur) * model_output.to(dtype)
+        prev_sample = prev_sample.to(sample.dtype)
+
+        self._step_index += 1
+        if not return_dict:
+            return (prev_sample,)
+        return FlowBridgeSchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        """The training-side bridge state ``z_t = (1 - t) * eps + t * z_e`` (bridge.py:311).
+
+        Args:
+            original_samples: The EO latent ``z_e`` (the ``t = 1`` endpoint).
+            noise: ``eps ~ N(0, I)`` (the ``t = 0`` endpoint).
+            timesteps: Bridge times in ``[0, 1]``, broadcastable over the batch.
+        """
+        t = timesteps.to(original_samples.device, original_samples.dtype)
+        t = t.view(-1, *([1] * (original_samples.ndim - 1)))
+        return (1.0 - t) * noise + t * original_samples
+
+    def get_velocity(
+        self,
+        sample: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        """The training target ``u* = z_e - eps`` (bridge.py:328 at ``sigma_b = 0``).
+
+        Constant along the path, hence independent of ``timesteps``; the argument
+        is kept for `diffusers` API compatibility.
+        """
+        del timesteps
+        return sample - noise

qxs-saropt/transformer/config.json

@@ -0,0 +1,18 @@
+{
+  "_class_name": "ReFlowSETTransformer2DModel",
+  "_diffusers_version": "0.37.1",
+  "axes_dim": [
+    32,
+    32
+  ],
+  "depth": 24,
+  "double_blocks": 8,
+  "double_merge": "token",
+  "hidden_size": 1024,
+  "in_channels": 128,
+  "mlp_ratio": 4.0,
+  "num_heads": 16,
+  "out_channels": 128,
+  "sample_size": 256,
+  "theta": 10000
+}

qxs-saropt/transformer/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a549f999d19044f73cb62943dbfbc45fbd48cc0350b7670cb279fb10a3f4f1e
+size 2037319452

qxs-saropt/transformer/transformer_reflowset.py

@@ -0,0 +1,471 @@
+"""ReFlowSET velocity transformer — a latent DiT with an EO/SAR double stream.
+
+The network predicts the flow-bridge velocity ``dz/dt`` in the frozen FLUX.2
+latent space.  It takes the noisy EO latent ``[B, 128, h, w]``, a scalar bridge
+time ``t`` in ``[0, 1]``, and the SAR conditioning latent of the same shape; the
+first 8 of its 24 blocks are double-stream (one EO tower and one SAR tower over
+a single joint attention), the remaining 16 are single-stream over the
+concatenated ``[EO | SAR]`` sequence, and only the EO half is decoded.
+
+This is an inference-only port.  The training-only REPA projection head
+(``repa_proj``) is a separate module in the reference implementation and is
+deliberately absent here.
+
+Deviations from `diffusers`' FLUX blocks that this file has to keep — each one
+is silent if you get it wrong:
+
+* ``FinalLayer`` unpacks ``shift, scale`` (dit.py:360), the **opposite** order of
+  ``AdaLayerNormContinuous``.
+* The single-stream MLP is **SwiGLU** of width 2752, not a 4x GELU of width 4096.
+* ``linear1``/``linear2`` are **bias-free**, and the QK-norm parameter is called
+  ``scale``, not ``weight``.
+* The timestep is multiplied by 1000 *inside* the model and the sinusoid is
+  **cos first, then sin**.
+* RoPE runs on **two** axes of **centred half-integer** coordinates, not on
+  FLUX's three axes of integers starting at 0.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional, Union
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+#: Width of the sinusoidal timestep embedding fed to ``time_in`` (dit.py:44).
+#: A module constant, deliberately independent of ``hidden_size``.
+TIME_EMBED_DIM = 256
+
+
+def swiglu_hidden_dim(hidden_size: int, mlp_ratio: float) -> int:
+    """SwiGLU intermediate width (dit.py:120-127).
+
+    The canonical 2/3 rule rounded to a multiple of 64, so a gated MLP at
+    ``mlp_ratio=4.0`` costs the same parameters as a plain 4x GELU MLP.
+    ``hidden_size=1024, mlp_ratio=4.0 -> 2752``.
+    """
+    return int(round(hidden_size * mlp_ratio * 2 / 3 / 64)) * 64
+
+
+class SwiGLU(nn.Module):
+    """``silu(first half) * second half`` — gate first, value second (dit.py:130-133)."""
+
+    def forward(self, x: Tensor) -> Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        return F.silu(x1) * x2
+
+
+class RMSNorm(nn.Module):
+    """RMS norm computed in float32.  The parameter is named ``scale`` (dit.py:136-145)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        x_dtype = x.dtype
+        x = x.float()
+        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
+        return (x * rrms).to(dtype=x_dtype) * self.scale
+
+
+class QKNorm(nn.Module):
+    """Per-head query/key RMS norm, applied **before** RoPE (dit.py:148-155)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.query_norm = RMSNorm(dim)
+        self.key_norm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        return self.query_norm(q).to(v), self.key_norm(k).to(v)
+
+
+class MLPEmbedder(nn.Module):
+    """``Linear -> SiLU -> Linear`` time-embedding MLP (dit.py:158-166).
+
+    Checkpoint keys are ``time_in.in_layer.*`` / ``time_in.out_layer.*``, not
+    diffusers' ``time_text_embed.timestep_embedder.linear_{1,2}``.
+    """
+
+    def __init__(self, in_dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+        self.silu = nn.SiLU()
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class Modulation(nn.Module):
+    """AdaLN-Zero triple.  Order is ``shift, scale, gate`` (dit.py:169-181)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.lin = nn.Linear(dim, 3 * dim, bias=True)
+
+    def forward(self, vec: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        out = self.lin(F.silu(vec))
+        if out.ndim == 2:
+            out = out[:, None, :]
+        shift, scale, gate = out.chunk(3, dim=-1)
+        return shift, scale, gate
+
+
+def timestep_embedding(
+    t: Tensor, dim: int, max_period: int = 10000, time_factor: float = 1000.0
+) -> Tensor:
+    """Sinusoidal embedding of a fractional bridge time (dit.py:184-201).
+
+    Two things differ from `diffusers`' ``get_timestep_embedding`` defaults:
+    ``t`` is a fraction in ``[0, 1]`` that is scaled by ``time_factor = 1000``
+    **here**, and the concatenation order is ``[cos, sin]`` (FLUX's ordering,
+    i.e. ``flip_sin_to_cos=True``).
+    """
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period)
+        * torch.arange(start=0, end=half, device=t.device, dtype=torch.float32)
+        / half
+    )
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    """Per-axis rotation matrices ``[..., L, dim/2, 2, 2]`` (dit.py:204-211)."""
+    if dim % 2:
+        raise ValueError(f"RoPE axis dim must be even, got {dim}")
+    scale = torch.arange(0, dim, 2, dtype=pos.dtype, device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
+    return out.reshape(*out.shape[:-1], 2, 2).float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+    """Rotate consecutive dimension pairs — the interleaved (FLUX) convention.
+
+    ``(x0, x1) -> (cos*x0 - sin*x1, sin*x0 + cos*x1)`` on ``(x[2k], x[2k+1])``
+    (dit.py:214-219).  Equivalent to diffusers' ``apply_rotary_emb(...,
+    use_real_unbind_dim=-1)``; ``-2`` is the split-halves convention and is wrong
+    for these weights.
+    """
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+
+
+class EmbedND(nn.Module):
+    """Concatenates the per-axis RoPE ladders and inserts the head axis (dit.py:222-234).
+
+    Holds no parameters and no buffers: the grid is rebuilt on every forward,
+    which is what lets one checkpoint serve 256 and 512 inputs.
+    """
+
+    def __init__(self, theta: int, axes_dim: list[int]) -> None:
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        emb = torch.cat(
+            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(len(self.axes_dim))],
+            dim=-3,
+        )
+        return emb.unsqueeze(1)
+
+
+def latent_image_ids(h: int, w: int, device, dtype=torch.float32) -> Tensor:
+    """Centred ``(y, x)`` coordinates for an ``h x w`` latent grid, ``[h*w, 2]``.
+
+    ``arange(n) - (n - 1) / 2`` with unit spacing (dit.py:237-252), so for even
+    ``n`` the coordinates are half-integers and the central 16x16 region of a
+    32x32 grid carries exactly the coordinates a 256-trained model saw — RoPE
+    only extrapolates outwards, it never rescales.  Row-major, so token
+    ``p = y * w + x``.  This is **not** FLUX's 3-axis integer id grid.
+    """
+    y = torch.arange(h, device=device, dtype=dtype) - (h - 1) / 2
+    x = torch.arange(w, device=device, dtype=dtype) - (w - 1) / 2
+    ids = torch.zeros(h, w, 2, device=device, dtype=dtype)
+    ids[..., 0] = y[:, None]
+    ids[..., 1] = x[None, :]
+    return ids.reshape(h * w, 2)
+
+
+class SingleStreamBlock(nn.Module):
+    """Fused attention + SwiGLU MLP under one modulation and one residual.
+
+    ``linear1`` emits ``[q | k | v | mlp_gate | mlp_value]`` in that order; the
+    qkv slab is K-major (``(K H D)``).  Both linears are bias-free
+    (dit.py:255-300).
+    """
+
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.mlp_hidden_dim = swiglu_hidden_dim(hidden_size, mlp_ratio)
+
+        self.linear1 = nn.Linear(hidden_size, 3 * hidden_size + 2 * self.mlp_hidden_dim, bias=False)
+        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size, bias=False)
+        self.norm = QKNorm(head_dim)
+        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp_act = SwiGLU()
+        self.modulation = Modulation(hidden_size)
+
+    def pre_attention(self, x: Tensor, vec: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+        """Everything up to (not including) RoPE and attention (dit.py:272-288)."""
+        shift, scale, gate = self.modulation(vec)
+        x_mod = (1 + scale) * self.pre_norm(x) + shift
+
+        qkv, mlp = torch.split(
+            self.linear1(x_mod), [3 * self.hidden_size, 2 * self.mlp_hidden_dim], dim=-1
+        )
+        b, length, _ = qkv.shape
+        # "B L (K H D) -> K B H L D" with K=3, H=num_heads.
+        q, k, v = qkv.reshape(b, length, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        q, k = self.norm(q, k, v)
+        return q, k, v, mlp, gate
+
+    def post_attention(self, x: Tensor, attn: Tensor, mlp: Tensor, gate: Tensor) -> Tensor:
+        """Output projection and the single gated residual (dit.py:290-294)."""
+        b, heads, length, head_dim = attn.shape
+        attn = attn.transpose(1, 2).reshape(b, length, heads * head_dim)
+        out = self.linear2(torch.cat((attn, self.mlp_act(mlp)), dim=-1))
+        return x + gate * out
+
+    def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
+        q, k, v, mlp, gate = self.pre_attention(x, vec)
+        q, k = apply_rope(q, k, pe)
+        attn = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+        return self.post_attention(x, attn, mlp, gate)
+
+
+class DoubleStreamBlock(nn.Module):
+    """Two independent towers over **one** joint attention across ``[EO | SAR]``.
+
+    The towers have completely separate weights but share the modulation vector
+    ``vec`` and the RoPE grid, so an EO token and the SAR token at the same
+    ground position carry an identical phase (dit.py:303-343).  SAR plays the
+    structural role text plays in FLUX.
+    """
+
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0) -> None:
+        super().__init__()
+        self.eo = SingleStreamBlock(hidden_size, num_heads, mlp_ratio)
+        self.sar = SingleStreamBlock(hidden_size, num_heads, mlp_ratio)
+
+    def forward(self, eo: Tensor, sar: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]:
+        """``pe`` must already cover the joint 2P-token sequence."""
+        q_e, k_e, v_e, mlp_e, gate_e = self.eo.pre_attention(eo, vec)
+        q_s, k_s, v_s, mlp_s, gate_s = self.sar.pre_attention(sar, vec)
+
+        q = torch.cat((q_e, q_s), dim=2)
+        k = torch.cat((k_e, k_s), dim=2)
+        v = torch.cat((v_e, v_s), dim=2)
+        q, k = apply_rope(q, k, pe)
+
+        attn = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+        attn_e, attn_s = attn.split([q_e.shape[2], q_s.shape[2]], dim=2)
+        return (
+            self.eo.post_attention(eo, attn_e, mlp_e, gate_e),
+            self.sar.post_attention(sar, attn_s, mlp_s, gate_s),
+        )
+
+
+class FinalLayer(nn.Module):
+    """AdaLN output layer.
+
+    ``adaLN`` unpacks ``shift, scale`` — the **opposite** order of diffusers'
+    ``AdaLayerNormContinuous`` (dit.py:346-362).  ``logvar_proj`` belongs to a
+    beta-NLL loss that was never enabled (``loss.flow = mse``); its weights are
+    kept so the published checkpoint loads with ``strict=True``, but inference
+    never evaluates it — the sampler reads only the velocity (bridge.py:531).
+    """
+
+    def __init__(self, hidden_size: int, out_channels: int) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.adaLN = nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+        self.proj = nn.Linear(hidden_size, out_channels, bias=True)
+        self.logvar_proj = nn.Linear(hidden_size, 1, bias=True)
+
+    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+        mod = self.adaLN(F.silu(vec))
+        if mod.ndim == 2:
+            mod = mod[:, None, :]
+        shift, scale = mod.chunk(2, dim=-1)
+        return self.proj((1 + scale) * self.norm(x) + shift)
+
+
+class ReFlowSETTransformer2DModel(ModelMixin, ConfigMixin):
+    """ReFlowSET's flow-velocity transformer (509.32 M parameters as configured).
+
+    Args:
+        in_channels: Channels of the packed FLUX.2 latent (128).
+        out_channels: Channels of the predicted velocity (128).
+        hidden_size: Residual width (1024).
+        depth: **Total** blocks, double plus single (24).
+        num_heads: Attention heads (16), so ``head_dim = 64``.
+        mlp_ratio: Nominal MLP ratio; the SwiGLU width is derived from it.
+        axes_dim: RoPE dims for the ``(y, x)`` axes; must sum to ``head_dim``.
+        theta: RoPE base period (10000).
+        sample_size: Input image resolution the released arm was trained at
+            (256 for QXS-SAROPT, 512 for SAR2Opt).  Recorded for provenance
+            only: the forward pass derives every shape from its input and the
+            RoPE grid is rebuilt per call, so one checkpoint serves any size
+            divisible by 16.
+        double_blocks: Leading double-stream blocks (8); the remaining
+            ``depth - double_blocks`` are single-stream.
+        double_merge: How the two streams become one.  ``"token"`` (the released
+            setting) concatenates on the sequence axis, so the single stack runs
+            over 2P tokens and the SAR half is dropped only at the very end;
+            ``"channel"`` fuses per position and keeps P tokens.
+
+    Forward contract: ``forward(hidden_states, timestep, condition)`` where
+    ``hidden_states`` is the bridge state ``[B, 128, h, w]``, ``timestep`` is the
+    bridge time in ``[0, 1]`` (**not** an integer diffusion step), and
+    ``condition`` is the SAR latent of the same shape or ``None``.  ``None`` is
+    the classifier-free-guidance null branch and is turned into an all-zero
+    latent inside the model — there is no learned null token.
+    """
+
+    _supports_gradient_checkpointing = False
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 128,
+        out_channels: int = 128,
+        hidden_size: int = 1024,
+        depth: int = 24,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        axes_dim: tuple[int, ...] = (32, 32),
+        theta: int = 10000,
+        sample_size: Optional[int] = None,
+        double_blocks: int = 8,
+        double_merge: str = "token",
+    ) -> None:
+        super().__init__()
+        if hidden_size % num_heads != 0:
+            raise ValueError(f"hidden_size {hidden_size} must be divisible by num_heads {num_heads}")
+        pe_dim = hidden_size // num_heads
+        if sum(axes_dim) != pe_dim:
+            raise ValueError(f"axes_dim {list(axes_dim)} must sum to the per-head dim {pe_dim}")
+        if not 0 <= double_blocks < depth:
+            raise ValueError(f"double_blocks {double_blocks} must be in [0, depth={depth})")
+        if double_merge not in ("token", "channel"):
+            raise ValueError(f"double_merge must be 'token' or 'channel', got {double_merge!r}")
+
+        self.pe_embedder = EmbedND(theta=theta, axes_dim=list(axes_dim))
+        if double_blocks:
+            # Each stream gets its own 1x1 "patchify": they are two token
+            # sequences now, not two halves of one channel stack.
+            self.in_proj_eo = nn.Linear(in_channels, hidden_size, bias=True)
+            self.in_proj_sar = nn.Linear(in_channels, hidden_size, bias=True)
+            if double_merge == "channel":
+                self.merge = nn.Linear(2 * hidden_size, hidden_size, bias=True)
+        else:
+            self.in_proj = nn.Linear(2 * in_channels, hidden_size, bias=True)
+        self.time_in = MLPEmbedder(TIME_EMBED_DIM, hidden_size)
+        self.double_stream = nn.ModuleList(
+            [DoubleStreamBlock(hidden_size, num_heads, mlp_ratio) for _ in range(double_blocks)]
+        )
+        self.blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(hidden_size, num_heads, mlp_ratio)
+                for _ in range(depth - double_blocks)
+            ]
+        )
+        self.final_layer = FinalLayer(hidden_size, out_channels)
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        timestep: Tensor,
+        condition: Optional[Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[Transformer2DModelOutput, tuple[Tensor]]:
+        """Predict the flow velocity ``dz/dt``.
+
+        Args:
+            hidden_states: ``[B, in_channels, h, w]`` bridge state.
+            timestep: Bridge time in ``[0, 1]``; a scalar or ``[B]``.
+            condition: ``[B, in_channels, h, w]`` SAR latent, or ``None`` for the
+                null branch (an all-zero conditioning latent, dit.py:531-532).
+            return_dict: Return a ``Transformer2DModelOutput`` instead of a tuple.
+
+        Returns:
+            The velocity ``[B, out_channels, h, w]``.  This is a flow velocity,
+            not ``epsilon`` and not diffusers' ``v_prediction``.
+        """
+        if hidden_states.ndim != 4:
+            raise ValueError(f"hidden_states must be [B, C, h, w], got {tuple(hidden_states.shape)}")
+        batch, _, h, w = hidden_states.shape
+        if condition is None:
+            condition = torch.zeros_like(hidden_states)
+        elif condition.shape != hidden_states.shape:
+            raise ValueError(
+                f"condition shape {tuple(condition.shape)} must match "
+                f"hidden_states shape {tuple(hidden_states.shape)}"
+            )
+        if timestep.ndim == 0:
+            timestep = timestep.expand(batch)
+
+        n_double = self.config.double_blocks
+        if n_double:
+            eo = self.in_proj_eo(hidden_states.flatten(2).transpose(1, 2))  # [B, P, D]
+            sar = self.in_proj_sar(condition.flatten(2).transpose(1, 2))  # [B, P, D]
+            ref = eo
+        else:
+            x = torch.cat([hidden_states, condition], dim=1).flatten(2).transpose(1, 2)
+            x = self.in_proj(x)
+            ref = x
+
+        vec = self.time_in(timestep_embedding(timestep, TIME_EMBED_DIM).to(ref.dtype))
+
+        ids = latent_image_ids(h, w, device=hidden_states.device, dtype=torch.float32)
+        pe = self.pe_embedder(ids[None].expand(batch, -1, -1))
+
+        num_tokens = ref.shape[1]
+        pe_single = pe
+        if n_double:
+            # Token axis of pe is dim 2 ([B, 1, L, head_dim/2, 2, 2]); repeating
+            # the same P coordinates gives EO and SAR one shared grid.
+            pe_joint = torch.cat((pe, pe), dim=2)
+            for block in self.double_stream:
+                eo, sar = block(eo, sar, vec, pe_joint)
+            if self.config.double_merge == "token":
+                x = torch.cat((eo, sar), dim=1)  # [B, 2P, D]
+                pe_single = pe_joint
+            else:
+                x = self.merge(torch.cat((eo, sar), dim=-1))  # [B, P, D]
+
+        for block in self.blocks:
+            x = block(x, vec, pe_single)
+
+        if n_double and self.config.double_merge == "token":
+            x = x[:, :num_tokens]  # drop the SAR half: only EO is decoded
+
+        v = self.final_layer(x, vec)
+        v = v.transpose(1, 2).reshape(batch, self.config.out_channels, h, w)
+        if not return_dict:
+            return (v,)
+        return Transformer2DModelOutput(sample=v)

qxs-saropt/transformer_reflowset.py

@@ -0,0 +1,471 @@
+"""ReFlowSET velocity transformer — a latent DiT with an EO/SAR double stream.
+
+The network predicts the flow-bridge velocity ``dz/dt`` in the frozen FLUX.2
+latent space.  It takes the noisy EO latent ``[B, 128, h, w]``, a scalar bridge
+time ``t`` in ``[0, 1]``, and the SAR conditioning latent of the same shape; the
+first 8 of its 24 blocks are double-stream (one EO tower and one SAR tower over
+a single joint attention), the remaining 16 are single-stream over the
+concatenated ``[EO | SAR]`` sequence, and only the EO half is decoded.
+
+This is an inference-only port.  The training-only REPA projection head
+(``repa_proj``) is a separate module in the reference implementation and is
+deliberately absent here.
+
+Deviations from `diffusers`' FLUX blocks that this file has to keep — each one
+is silent if you get it wrong:
+
+* ``FinalLayer`` unpacks ``shift, scale`` (dit.py:360), the **opposite** order of
+  ``AdaLayerNormContinuous``.
+* The single-stream MLP is **SwiGLU** of width 2752, not a 4x GELU of width 4096.
+* ``linear1``/``linear2`` are **bias-free**, and the QK-norm parameter is called
+  ``scale``, not ``weight``.
+* The timestep is multiplied by 1000 *inside* the model and the sinusoid is
+  **cos first, then sin**.
+* RoPE runs on **two** axes of **centred half-integer** coordinates, not on
+  FLUX's three axes of integers starting at 0.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional, Union
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+#: Width of the sinusoidal timestep embedding fed to ``time_in`` (dit.py:44).
+#: A module constant, deliberately independent of ``hidden_size``.
+TIME_EMBED_DIM = 256
+
+
+def swiglu_hidden_dim(hidden_size: int, mlp_ratio: float) -> int:
+    """SwiGLU intermediate width (dit.py:120-127).
+
+    The canonical 2/3 rule rounded to a multiple of 64, so a gated MLP at
+    ``mlp_ratio=4.0`` costs the same parameters as a plain 4x GELU MLP.
+    ``hidden_size=1024, mlp_ratio=4.0 -> 2752``.
+    """
+    return int(round(hidden_size * mlp_ratio * 2 / 3 / 64)) * 64
+
+
+class SwiGLU(nn.Module):
+    """``silu(first half) * second half`` — gate first, value second (dit.py:130-133)."""
+
+    def forward(self, x: Tensor) -> Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        return F.silu(x1) * x2
+
+
+class RMSNorm(nn.Module):
+    """RMS norm computed in float32.  The parameter is named ``scale`` (dit.py:136-145)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        x_dtype = x.dtype
+        x = x.float()
+        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
+        return (x * rrms).to(dtype=x_dtype) * self.scale
+
+
+class QKNorm(nn.Module):
+    """Per-head query/key RMS norm, applied **before** RoPE (dit.py:148-155)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.query_norm = RMSNorm(dim)
+        self.key_norm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        return self.query_norm(q).to(v), self.key_norm(k).to(v)
+
+
+class MLPEmbedder(nn.Module):
+    """``Linear -> SiLU -> Linear`` time-embedding MLP (dit.py:158-166).
+
+    Checkpoint keys are ``time_in.in_layer.*`` / ``time_in.out_layer.*``, not
+    diffusers' ``time_text_embed.timestep_embedder.linear_{1,2}``.
+    """
+
+    def __init__(self, in_dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+        self.silu = nn.SiLU()
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class Modulation(nn.Module):
+    """AdaLN-Zero triple.  Order is ``shift, scale, gate`` (dit.py:169-181)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.lin = nn.Linear(dim, 3 * dim, bias=True)
+
+    def forward(self, vec: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        out = self.lin(F.silu(vec))
+        if out.ndim == 2:
+            out = out[:, None, :]
+        shift, scale, gate = out.chunk(3, dim=-1)
+        return shift, scale, gate
+
+
+def timestep_embedding(
+    t: Tensor, dim: int, max_period: int = 10000, time_factor: float = 1000.0
+) -> Tensor:
+    """Sinusoidal embedding of a fractional bridge time (dit.py:184-201).
+
+    Two things differ from `diffusers`' ``get_timestep_embedding`` defaults:
+    ``t`` is a fraction in ``[0, 1]`` that is scaled by ``time_factor = 1000``
+    **here**, and the concatenation order is ``[cos, sin]`` (FLUX's ordering,
+    i.e. ``flip_sin_to_cos=True``).
+    """
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period)
+        * torch.arange(start=0, end=half, device=t.device, dtype=torch.float32)
+        / half
+    )
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    """Per-axis rotation matrices ``[..., L, dim/2, 2, 2]`` (dit.py:204-211)."""
+    if dim % 2:
+        raise ValueError(f"RoPE axis dim must be even, got {dim}")
+    scale = torch.arange(0, dim, 2, dtype=pos.dtype, device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
+    return out.reshape(*out.shape[:-1], 2, 2).float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+    """Rotate consecutive dimension pairs — the interleaved (FLUX) convention.
+
+    ``(x0, x1) -> (cos*x0 - sin*x1, sin*x0 + cos*x1)`` on ``(x[2k], x[2k+1])``
+    (dit.py:214-219).  Equivalent to diffusers' ``apply_rotary_emb(...,
+    use_real_unbind_dim=-1)``; ``-2`` is the split-halves convention and is wrong
+    for these weights.
+    """
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+
+
+class EmbedND(nn.Module):
+    """Concatenates the per-axis RoPE ladders and inserts the head axis (dit.py:222-234).
+
+    Holds no parameters and no buffers: the grid is rebuilt on every forward,
+    which is what lets one checkpoint serve 256 and 512 inputs.
+    """
+
+    def __init__(self, theta: int, axes_dim: list[int]) -> None:
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        emb = torch.cat(
+            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(len(self.axes_dim))],
+            dim=-3,
+        )
+        return emb.unsqueeze(1)
+
+
+def latent_image_ids(h: int, w: int, device, dtype=torch.float32) -> Tensor:
+    """Centred ``(y, x)`` coordinates for an ``h x w`` latent grid, ``[h*w, 2]``.
+
+    ``arange(n) - (n - 1) / 2`` with unit spacing (dit.py:237-252), so for even
+    ``n`` the coordinates are half-integers and the central 16x16 region of a
+    32x32 grid carries exactly the coordinates a 256-trained model saw — RoPE
+    only extrapolates outwards, it never rescales.  Row-major, so token
+    ``p = y * w + x``.  This is **not** FLUX's 3-axis integer id grid.
+    """
+    y = torch.arange(h, device=device, dtype=dtype) - (h - 1) / 2
+    x = torch.arange(w, device=device, dtype=dtype) - (w - 1) / 2
+    ids = torch.zeros(h, w, 2, device=device, dtype=dtype)
+    ids[..., 0] = y[:, None]
+    ids[..., 1] = x[None, :]
+    return ids.reshape(h * w, 2)
+
+
+class SingleStreamBlock(nn.Module):
+    """Fused attention + SwiGLU MLP under one modulation and one residual.
+
+    ``linear1`` emits ``[q | k | v | mlp_gate | mlp_value]`` in that order; the
+    qkv slab is K-major (``(K H D)``).  Both linears are bias-free
+    (dit.py:255-300).
+    """
+
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.mlp_hidden_dim = swiglu_hidden_dim(hidden_size, mlp_ratio)
+
+        self.linear1 = nn.Linear(hidden_size, 3 * hidden_size + 2 * self.mlp_hidden_dim, bias=False)
+        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size, bias=False)
+        self.norm = QKNorm(head_dim)
+        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp_act = SwiGLU()
+        self.modulation = Modulation(hidden_size)
+
+    def pre_attention(self, x: Tensor, vec: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+        """Everything up to (not including) RoPE and attention (dit.py:272-288)."""
+        shift, scale, gate = self.modulation(vec)
+        x_mod = (1 + scale) * self.pre_norm(x) + shift
+
+        qkv, mlp = torch.split(
+            self.linear1(x_mod), [3 * self.hidden_size, 2 * self.mlp_hidden_dim], dim=-1
+        )
+        b, length, _ = qkv.shape
+        # "B L (K H D) -> K B H L D" with K=3, H=num_heads.
+        q, k, v = qkv.reshape(b, length, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        q, k = self.norm(q, k, v)
+        return q, k, v, mlp, gate
+
+    def post_attention(self, x: Tensor, attn: Tensor, mlp: Tensor, gate: Tensor) -> Tensor:
+        """Output projection and the single gated residual (dit.py:290-294)."""
+        b, heads, length, head_dim = attn.shape
+        attn = attn.transpose(1, 2).reshape(b, length, heads * head_dim)
+        out = self.linear2(torch.cat((attn, self.mlp_act(mlp)), dim=-1))
+        return x + gate * out
+
+    def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
+        q, k, v, mlp, gate = self.pre_attention(x, vec)
+        q, k = apply_rope(q, k, pe)
+        attn = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+        return self.post_attention(x, attn, mlp, gate)
+
+
+class DoubleStreamBlock(nn.Module):
+    """Two independent towers over **one** joint attention across ``[EO | SAR]``.
+
+    The towers have completely separate weights but share the modulation vector
+    ``vec`` and the RoPE grid, so an EO token and the SAR token at the same
+    ground position carry an identical phase (dit.py:303-343).  SAR plays the
+    structural role text plays in FLUX.
+    """
+
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0) -> None:
+        super().__init__()
+        self.eo = SingleStreamBlock(hidden_size, num_heads, mlp_ratio)
+        self.sar = SingleStreamBlock(hidden_size, num_heads, mlp_ratio)
+
+    def forward(self, eo: Tensor, sar: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]:
+        """``pe`` must already cover the joint 2P-token sequence."""
+        q_e, k_e, v_e, mlp_e, gate_e = self.eo.pre_attention(eo, vec)
+        q_s, k_s, v_s, mlp_s, gate_s = self.sar.pre_attention(sar, vec)
+
+        q = torch.cat((q_e, q_s), dim=2)
+        k = torch.cat((k_e, k_s), dim=2)
+        v = torch.cat((v_e, v_s), dim=2)
+        q, k = apply_rope(q, k, pe)
+
+        attn = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+        attn_e, attn_s = attn.split([q_e.shape[2], q_s.shape[2]], dim=2)
+        return (
+            self.eo.post_attention(eo, attn_e, mlp_e, gate_e),
+            self.sar.post_attention(sar, attn_s, mlp_s, gate_s),
+        )
+
+
+class FinalLayer(nn.Module):
+    """AdaLN output layer.
+
+    ``adaLN`` unpacks ``shift, scale`` — the **opposite** order of diffusers'
+    ``AdaLayerNormContinuous`` (dit.py:346-362).  ``logvar_proj`` belongs to a
+    beta-NLL loss that was never enabled (``loss.flow = mse``); its weights are
+    kept so the published checkpoint loads with ``strict=True``, but inference
+    never evaluates it — the sampler reads only the velocity (bridge.py:531).
+    """
+
+    def __init__(self, hidden_size: int, out_channels: int) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.adaLN = nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+        self.proj = nn.Linear(hidden_size, out_channels, bias=True)
+        self.logvar_proj = nn.Linear(hidden_size, 1, bias=True)
+
+    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+        mod = self.adaLN(F.silu(vec))
+        if mod.ndim == 2:
+            mod = mod[:, None, :]
+        shift, scale = mod.chunk(2, dim=-1)
+        return self.proj((1 + scale) * self.norm(x) + shift)
+
+
+class ReFlowSETTransformer2DModel(ModelMixin, ConfigMixin):
+    """ReFlowSET's flow-velocity transformer (509.32 M parameters as configured).
+
+    Args:
+        in_channels: Channels of the packed FLUX.2 latent (128).
+        out_channels: Channels of the predicted velocity (128).
+        hidden_size: Residual width (1024).
+        depth: **Total** blocks, double plus single (24).
+        num_heads: Attention heads (16), so ``head_dim = 64``.
+        mlp_ratio: Nominal MLP ratio; the SwiGLU width is derived from it.
+        axes_dim: RoPE dims for the ``(y, x)`` axes; must sum to ``head_dim``.
+        theta: RoPE base period (10000).
+        sample_size: Input image resolution the released arm was trained at
+            (256 for QXS-SAROPT, 512 for SAR2Opt).  Recorded for provenance
+            only: the forward pass derives every shape from its input and the
+            RoPE grid is rebuilt per call, so one checkpoint serves any size
+            divisible by 16.
+        double_blocks: Leading double-stream blocks (8); the remaining
+            ``depth - double_blocks`` are single-stream.
+        double_merge: How the two streams become one.  ``"token"`` (the released
+            setting) concatenates on the sequence axis, so the single stack runs
+            over 2P tokens and the SAR half is dropped only at the very end;
+            ``"channel"`` fuses per position and keeps P tokens.
+
+    Forward contract: ``forward(hidden_states, timestep, condition)`` where
+    ``hidden_states`` is the bridge state ``[B, 128, h, w]``, ``timestep`` is the
+    bridge time in ``[0, 1]`` (**not** an integer diffusion step), and
+    ``condition`` is the SAR latent of the same shape or ``None``.  ``None`` is
+    the classifier-free-guidance null branch and is turned into an all-zero
+    latent inside the model — there is no learned null token.
+    """
+
+    _supports_gradient_checkpointing = False
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 128,
+        out_channels: int = 128,
+        hidden_size: int = 1024,
+        depth: int = 24,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        axes_dim: tuple[int, ...] = (32, 32),
+        theta: int = 10000,
+        sample_size: Optional[int] = None,
+        double_blocks: int = 8,
+        double_merge: str = "token",
+    ) -> None:
+        super().__init__()
+        if hidden_size % num_heads != 0:
+            raise ValueError(f"hidden_size {hidden_size} must be divisible by num_heads {num_heads}")
+        pe_dim = hidden_size // num_heads
+        if sum(axes_dim) != pe_dim:
+            raise ValueError(f"axes_dim {list(axes_dim)} must sum to the per-head dim {pe_dim}")
+        if not 0 <= double_blocks < depth:
+            raise ValueError(f"double_blocks {double_blocks} must be in [0, depth={depth})")
+        if double_merge not in ("token", "channel"):
+            raise ValueError(f"double_merge must be 'token' or 'channel', got {double_merge!r}")
+
+        self.pe_embedder = EmbedND(theta=theta, axes_dim=list(axes_dim))
+        if double_blocks:
+            # Each stream gets its own 1x1 "patchify": they are two token
+            # sequences now, not two halves of one channel stack.
+            self.in_proj_eo = nn.Linear(in_channels, hidden_size, bias=True)
+            self.in_proj_sar = nn.Linear(in_channels, hidden_size, bias=True)
+            if double_merge == "channel":
+                self.merge = nn.Linear(2 * hidden_size, hidden_size, bias=True)
+        else:
+            self.in_proj = nn.Linear(2 * in_channels, hidden_size, bias=True)
+        self.time_in = MLPEmbedder(TIME_EMBED_DIM, hidden_size)
+        self.double_stream = nn.ModuleList(
+            [DoubleStreamBlock(hidden_size, num_heads, mlp_ratio) for _ in range(double_blocks)]
+        )
+        self.blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(hidden_size, num_heads, mlp_ratio)
+                for _ in range(depth - double_blocks)
+            ]
+        )
+        self.final_layer = FinalLayer(hidden_size, out_channels)
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        timestep: Tensor,
+        condition: Optional[Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[Transformer2DModelOutput, tuple[Tensor]]:
+        """Predict the flow velocity ``dz/dt``.
+
+        Args:
+            hidden_states: ``[B, in_channels, h, w]`` bridge state.
+            timestep: Bridge time in ``[0, 1]``; a scalar or ``[B]``.
+            condition: ``[B, in_channels, h, w]`` SAR latent, or ``None`` for the
+                null branch (an all-zero conditioning latent, dit.py:531-532).
+            return_dict: Return a ``Transformer2DModelOutput`` instead of a tuple.
+
+        Returns:
+            The velocity ``[B, out_channels, h, w]``.  This is a flow velocity,
+            not ``epsilon`` and not diffusers' ``v_prediction``.
+        """
+        if hidden_states.ndim != 4:
+            raise ValueError(f"hidden_states must be [B, C, h, w], got {tuple(hidden_states.shape)}")
+        batch, _, h, w = hidden_states.shape
+        if condition is None:
+            condition = torch.zeros_like(hidden_states)
+        elif condition.shape != hidden_states.shape:
+            raise ValueError(
+                f"condition shape {tuple(condition.shape)} must match "
+                f"hidden_states shape {tuple(hidden_states.shape)}"
+            )
+        if timestep.ndim == 0:
+            timestep = timestep.expand(batch)
+
+        n_double = self.config.double_blocks
+        if n_double:
+            eo = self.in_proj_eo(hidden_states.flatten(2).transpose(1, 2))  # [B, P, D]
+            sar = self.in_proj_sar(condition.flatten(2).transpose(1, 2))  # [B, P, D]
+            ref = eo
+        else:
+            x = torch.cat([hidden_states, condition], dim=1).flatten(2).transpose(1, 2)
+            x = self.in_proj(x)
+            ref = x
+
+        vec = self.time_in(timestep_embedding(timestep, TIME_EMBED_DIM).to(ref.dtype))
+
+        ids = latent_image_ids(h, w, device=hidden_states.device, dtype=torch.float32)
+        pe = self.pe_embedder(ids[None].expand(batch, -1, -1))
+
+        num_tokens = ref.shape[1]
+        pe_single = pe
+        if n_double:
+            # Token axis of pe is dim 2 ([B, 1, L, head_dim/2, 2, 2]); repeating
+            # the same P coordinates gives EO and SAR one shared grid.
+            pe_joint = torch.cat((pe, pe), dim=2)
+            for block in self.double_stream:
+                eo, sar = block(eo, sar, vec, pe_joint)
+            if self.config.double_merge == "token":
+                x = torch.cat((eo, sar), dim=1)  # [B, 2P, D]
+                pe_single = pe_joint
+            else:
+                x = self.merge(torch.cat((eo, sar), dim=-1))  # [B, P, D]
+
+        for block in self.blocks:
+            x = block(x, vec, pe_single)
+
+        if n_double and self.config.double_merge == "token":
+            x = x[:, :num_tokens]  # drop the SAR half: only EO is decoded
+
+        v = self.final_layer(x, vec)
+        v = v.transpose(1, 2).reshape(batch, self.config.out_channels, h, w)
+        if not return_dict:
+            return (v,)
+        return Transformer2DModelOutput(sample=v)

qxs-saropt/vae/autoencoder_flux2.py

@@ -0,0 +1,426 @@
+"""Frozen FLUX.2 autoencoder — the latent endpoint of ReFlowSET.
+
+ReFlowSET never fine-tunes this module: it is loaded once, frozen, and used to
+encode the SAR condition and to decode the sampled EO latent.  The released
+weights are the **Apache-2.0** FLUX.2-klein-base-4B copy of the autoencoder,
+re-keyed to the layout below (see ``scripts/convert_flux2_ae.py``).
+
+Three details of the checkpoint are non-standard for `diffusers` and are
+preserved exactly, because the file must load with ``strict=True``:
+
+* ``quant_conv`` lives **inside** ``encoder.*`` and is the last op of the
+  encoder forward; ``post_quant_conv`` lives **inside** ``decoder.*`` and is the
+  first op of the decoder forward.  `diffusers`' ``AutoencoderKL`` makes both
+  siblings of the encoder/decoder.
+* The latent normaliser is a real ``BatchNorm2d(128, affine=False)`` whose
+  running statistics ship in the checkpoint under ``bn.*`` — a per-channel mean
+  **and** variance, not a scalar ``scaling_factor``/``shift_factor``.  Its
+  epsilon is ``1e-4``, not torch's ``1e-5``.
+* ``encode`` returns the posterior **mean**; the log-variance chunk of the
+  encoder's moments is discarded, so encoding is deterministic and there is no
+  ``DiagonalGaussianDistribution`` and no ``.sample()``.
+
+The public latent is ``[B, 128, H/16, W/16]``: an 8x convolutional stride
+followed by a 2x2 space-to-depth pack that is part of the *autoencoder*, not of
+the transformer.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+
+def swish(x: Tensor) -> Tensor:
+    """``x * sigmoid(x)`` — the activation used throughout the FLUX.2 AE."""
+    return x * torch.sigmoid(x)
+
+
+class AttnBlock(nn.Module):
+    """Single-head self-attention over the spatial grid (head dim == channels)."""
+
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q, k, v = self.q(h_), self.k(h_), self.v(h_)
+        b, c, h, w = q.shape
+        # "b c h w -> b 1 (h w) c": ONE head whose head-dim is the full channel
+        # count (flux2_ae.py:70-73).
+        q = q.reshape(b, c, h * w).transpose(1, 2).unsqueeze(1).contiguous()
+        k = k.reshape(b, c, h * w).transpose(1, 2).unsqueeze(1).contiguous()
+        v = v.reshape(b, c, h * w).transpose(1, 2).unsqueeze(1).contiguous()
+        h_ = F.scaled_dot_product_attention(q, k, v)
+        return h_.squeeze(1).transpose(1, 2).reshape(b, c, h, w)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if in_channels != out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x: Tensor) -> Tensor:
+        h = self.conv1(swish(self.norm1(x)))
+        h = self.conv2(swish(self.norm2(h)))
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+        return x + h
+
+
+class Downsample(nn.Module):
+    """Stride-2 conv with FLUX's asymmetric ``(0, 1, 0, 1)`` pad (flux2_ae.py:111-121)."""
+
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.conv(F.pad(x, (0, 1, 0, 1), mode="constant", value=0))
+
+
+class Upsample(nn.Module):
+    """Nearest-neighbour 2x followed by a 3x3 conv (flux2_ae.py:124-132)."""
+
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.conv(F.interpolate(x, scale_factor=2.0, mode="nearest"))
+
+
+class Encoder(nn.Module):
+    """FLUX.2 encoder.  Emits ``2 * z_channels`` moments; ``quant_conv`` is internal."""
+
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        z_channels: int,
+    ) -> None:
+        super().__init__()
+        # Declared first so the checkpoint key is `encoder.quant_conv.*`
+        # (flux2_ae.py:146) — diffusers keeps quant_conv outside the encoder.
+        self.quant_conv = nn.Conv2d(2 * z_channels, 2 * z_channels, 1)
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        self.conv_in = nn.Conv2d(in_channels, ch, kernel_size=3, stride=1, padding=1)
+
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.down = nn.ModuleList()
+        block_in = ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(num_res_blocks):
+                block.append(ResnetBlock(block_in, block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            # Empty at every level in this checkpoint: attention exists only in
+            # `mid` (flux2_ae.py:162).  Kept so the forward guard is meaningful.
+            down.attn = nn.ModuleList()
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in)
+            self.down.append(down)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(block_in, block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(block_in, block_in)
+
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        h = self.mid.block_2(self.mid.attn_1(self.mid.block_1(hs[-1])))
+        h = self.conv_out(swish(self.norm_out(h)))
+        return self.quant_conv(h)  # last op of the encoder (flux2_ae.py:207)
+
+
+class Decoder(nn.Module):
+    """FLUX.2 decoder.  ``post_quant_conv`` is internal and runs first."""
+
+    def __init__(
+        self,
+        ch: int,
+        out_ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+    ) -> None:
+        super().__init__()
+        # Checkpoint key `decoder.post_quant_conv.*` (flux2_ae.py:223).
+        self.post_quant_conv = nn.Conv2d(z_channels, z_channels, 1)
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(block_in, block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(block_in, block_in)
+
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(num_res_blocks + 1):
+                block.append(ResnetBlock(block_in, block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = nn.ModuleList()  # empty in this checkpoint (flux2_ae.py:249)
+            if i_level != 0:
+                up.upsample = Upsample(block_in)
+            self.up.insert(0, up)  # prepend so `up.<i>` indexes by resolution level
+
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: Tensor) -> Tensor:
+        z = self.post_quant_conv(z)  # first op of the decoder (flux2_ae.py:267)
+        upscale_dtype = next(self.up.parameters()).dtype
+
+        h = self.conv_in(z)
+        h = self.mid.block_2(self.mid.attn_1(self.mid.block_1(h)))
+        h = h.to(upscale_dtype)
+
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        return self.conv_out(swish(self.norm_out(h)))
+
+
+class AutoencoderFlux2(ModelMixin, ConfigMixin):
+    """Frozen FLUX.2 autoencoder with ReFlowSET's packed, BN-normalised latent.
+
+    ``encode`` maps ``[B, 3, H, W]`` in ``[-1, 1]`` to ``[B, 128, H/16, W/16]``
+    and ``decode`` inverts it.  The module is frozen: ``train()`` is a no-op that
+    always selects eval mode, and the latent BatchNorm is additionally forced to
+    eval on every call so no batch statistic can ever leak into the latent.
+
+    Args:
+        resolution: Nominal training resolution of the original autoencoder.
+            Only used to size bookkeeping attributes; any ``H``, ``W`` divisible
+            by 16 may be encoded.
+        in_channels: Input image channels (3).
+        ch: Base width.
+        out_ch: Output image channels (3).
+        ch_mult: Per-level width multipliers; ``len(ch_mult) - 1`` downsamples.
+        num_res_blocks: Residual blocks per level.
+        z_channels: Pre-pack latent channels (32).
+        patch_size: Space-to-depth factor applied after the encoder (2), which
+            takes the latent from 32 channels at ``H/8`` to 128 at ``H/16``.
+        bn_eps: Epsilon of the latent BatchNorm.  **1e-4**, not torch's 1e-5
+            (flux2_ae.py:331); using 1e-5 shifts the latent by up to 2.6e-5.
+    """
+
+    _supports_gradient_checkpointing = False
+
+    @register_to_config
+    def __init__(
+        self,
+        resolution: int = 256,
+        in_channels: int = 3,
+        ch: int = 128,
+        out_ch: int = 3,
+        ch_mult: tuple[int, ...] = (1, 2, 4, 4),
+        num_res_blocks: int = 2,
+        z_channels: int = 32,
+        patch_size: int = 2,
+        bn_eps: float = 1e-4,
+    ) -> None:
+        super().__init__()
+        ch_mult = list(ch_mult)
+        self.encoder = Encoder(
+            resolution=resolution,
+            in_channels=in_channels,
+            ch=ch,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            z_channels=z_channels,
+        )
+        self.decoder = Decoder(
+            ch=ch,
+            out_ch=out_ch,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            in_channels=in_channels,
+            resolution=resolution,
+            z_channels=z_channels,
+        )
+        # Per-channel latent normaliser with the checkpoint's running statistics.
+        # affine=False, so there is no weight/bias to load (flux2_ae.py:334-340).
+        self.bn = nn.BatchNorm2d(
+            patch_size * patch_size * z_channels,
+            eps=bn_eps,
+            momentum=0.1,
+            affine=False,
+            track_running_stats=True,
+        )
+
+    @property
+    def latent_channels(self) -> int:
+        """Channels of the public latent: ``patch_size**2 * z_channels`` = 128."""
+        return self.config.patch_size**2 * self.config.z_channels
+
+    @property
+    def spatial_factor(self) -> int:
+        """Total stride: 8x convolutional times ``patch_size`` packing = 16."""
+        return 2 ** (len(self.config.ch_mult) - 1) * self.config.patch_size
+
+    # ---- 2x2 space-to-depth pack / unpack -----------------------------------
+
+    def pack(self, z: Tensor) -> Tensor:
+        """``[B, C, H, W] -> [B, C*p*p, H/p, W/p]``, channel-major.
+
+        Bit-identical to the reference ``rearrange("... c (i pi) (j pj) -> ...
+        (c pi pj) i j")`` (flux2_ae.py:349-357).  Note this is **not** diffusers'
+        ``_pack_latents``, whose channel grouping is transposed.
+        """
+        return F.pixel_unshuffle(z, self.config.patch_size)
+
+    def unpack(self, z: Tensor) -> Tensor:
+        """Exact inverse of :meth:`pack` (flux2_ae.py:359-367)."""
+        return F.pixel_shuffle(z, self.config.patch_size)
+
+    # ---- latent normalisation ----------------------------------------------
+
+    def normalize(self, z: Tensor) -> Tensor:
+        """``(z - running_mean) / sqrt(running_var + bn_eps)``, per channel."""
+        self.bn.eval()  # forced every call (flux2_ae.py:372); train mode shifts z by ~1.67
+        return self.bn(z)
+
+    def inv_normalize(self, z: Tensor) -> Tensor:
+        """Exact inverse of :meth:`normalize` — same ``bn_eps`` (flux2_ae.py:375-379)."""
+        self.bn.eval()
+        s = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.config.bn_eps)
+        m = self.bn.running_mean.view(1, -1, 1, 1)
+        return z * s + m
+
+    # ---- public API ---------------------------------------------------------
+
+    @torch.no_grad()
+    def encode(self, x: Tensor) -> Tensor:
+        """Encode an image to the packed, normalised latent.
+
+        Args:
+            x: ``[B, 3, H, W]`` in ``[-1, 1]``; ``H`` and ``W`` divisible by 16.
+
+        Returns:
+            ``[B, 128, H/16, W/16]`` — the posterior **mean**, packed and
+            BN-normalised.  The encoder's log-variance chunk is discarded
+            (flux2_ae.py:396), so this is deterministic: there is no posterior
+            distribution object and nothing to sample.
+        """
+        if x.ndim != 4 or x.shape[1] != self.config.in_channels:
+            raise ValueError(
+                f"encode expects [B, {self.config.in_channels}, H, W], got {tuple(x.shape)}"
+            )
+        h, w = x.shape[-2:]
+        if h % self.spatial_factor or w % self.spatial_factor:
+            raise ValueError(
+                f"encode requires H and W divisible by {self.spatial_factor}, got {h}x{w}"
+            )
+        moments = self.encoder(x)
+        mean = torch.chunk(moments, 2, dim=1)[0]
+        return self.normalize(self.pack(mean))
+
+    @torch.no_grad()
+    def decode(self, z: Tensor) -> Tensor:
+        """Decode a packed, normalised latent ``[B, 128, h, w]`` to ``[B, 3, 16h, 16w]``.
+
+        The output is approximately ``[-1, 1]`` and is **not** clamped here; the
+        pipeline applies ``(x * 0.5 + 0.5).clamp(0, 1)``.
+        """
+        if z.ndim != 4 or z.shape[1] != self.latent_channels:
+            raise ValueError(
+                f"decode expects [B, {self.latent_channels}, h, w], got {tuple(z.shape)}"
+            )
+        return self.decoder(self.unpack(self.inv_normalize(z)))
+
+    # ---- construction / freezing -------------------------------------------
+
+    @classmethod
+    def from_single_file(
+        cls,
+        path: str | os.PathLike,
+        torch_dtype: torch.dtype = torch.float32,
+    ) -> "AutoencoderFlux2":
+        """Load the single-file ``ae.safetensors`` (BFL key names) with ``strict=True``.
+
+        The released file is the Apache-2.0 FLUX.2-klein-base-4B autoencoder
+        re-keyed to this layout; it is stored in bfloat16 and is upcast to
+        ``torch_dtype``.  ReFlowSET runs the autoencoder in float32.
+        """
+        from safetensors.torch import load_file
+
+        path = os.fspath(path)
+        if not os.path.isfile(path):
+            raise FileNotFoundError(
+                f"FLUX.2 autoencoder weights not found at: {path}. Expected the "
+                "single-file 'ae.safetensors' shipped with ReFlowSET."
+            )
+        model = cls()
+        model.load_state_dict(load_file(path, device="cpu"), strict=True)
+        model.to(dtype=torch_dtype)
+        model.eval()
+        model.requires_grad_(False)
+        return model
+
+    def train(self, mode: bool = True) -> "AutoencoderFlux2":
+        """The autoencoder is frozen: never leave eval mode (flux2_ae.py:437-439)."""
+        return super().train(False)

qxs-saropt/vae/config.json

@@ -0,0 +1,18 @@
+{
+  "_class_name": "AutoencoderFlux2",
+  "_diffusers_version": "0.37.1",
+  "bn_eps": 0.0001,
+  "ch": 128,
+  "ch_mult": [
+    1,
+    2,
+    4,
+    4
+  ],
+  "in_channels": 3,
+  "num_res_blocks": 2,
+  "out_ch": 3,
+  "patch_size": 2,
+  "resolution": 256,
+  "z_channels": 32
+}

qxs-saropt/vae/diffusion_pytorch_model.safetensors

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

sar2opt/autoencoder_flux2.py

@@ -0,0 +1,426 @@
+"""Frozen FLUX.2 autoencoder — the latent endpoint of ReFlowSET.
+
+ReFlowSET never fine-tunes this module: it is loaded once, frozen, and used to
+encode the SAR condition and to decode the sampled EO latent.  The released
+weights are the **Apache-2.0** FLUX.2-klein-base-4B copy of the autoencoder,
+re-keyed to the layout below (see ``scripts/convert_flux2_ae.py``).
+
+Three details of the checkpoint are non-standard for `diffusers` and are
+preserved exactly, because the file must load with ``strict=True``:
+
+* ``quant_conv`` lives **inside** ``encoder.*`` and is the last op of the
+  encoder forward; ``post_quant_conv`` lives **inside** ``decoder.*`` and is the
+  first op of the decoder forward.  `diffusers`' ``AutoencoderKL`` makes both
+  siblings of the encoder/decoder.
+* The latent normaliser is a real ``BatchNorm2d(128, affine=False)`` whose
+  running statistics ship in the checkpoint under ``bn.*`` — a per-channel mean
+  **and** variance, not a scalar ``scaling_factor``/``shift_factor``.  Its
+  epsilon is ``1e-4``, not torch's ``1e-5``.
+* ``encode`` returns the posterior **mean**; the log-variance chunk of the
+  encoder's moments is discarded, so encoding is deterministic and there is no
+  ``DiagonalGaussianDistribution`` and no ``.sample()``.
+
+The public latent is ``[B, 128, H/16, W/16]``: an 8x convolutional stride
+followed by a 2x2 space-to-depth pack that is part of the *autoencoder*, not of
+the transformer.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+
+def swish(x: Tensor) -> Tensor:
+    """``x * sigmoid(x)`` — the activation used throughout the FLUX.2 AE."""
+    return x * torch.sigmoid(x)
+
+
+class AttnBlock(nn.Module):
+    """Single-head self-attention over the spatial grid (head dim == channels)."""
+
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q, k, v = self.q(h_), self.k(h_), self.v(h_)
+        b, c, h, w = q.shape
+        # "b c h w -> b 1 (h w) c": ONE head whose head-dim is the full channel
+        # count (flux2_ae.py:70-73).
+        q = q.reshape(b, c, h * w).transpose(1, 2).unsqueeze(1).contiguous()
+        k = k.reshape(b, c, h * w).transpose(1, 2).unsqueeze(1).contiguous()
+        v = v.reshape(b, c, h * w).transpose(1, 2).unsqueeze(1).contiguous()
+        h_ = F.scaled_dot_product_attention(q, k, v)
+        return h_.squeeze(1).transpose(1, 2).reshape(b, c, h, w)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if in_channels != out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x: Tensor) -> Tensor:
+        h = self.conv1(swish(self.norm1(x)))
+        h = self.conv2(swish(self.norm2(h)))
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+        return x + h
+
+
+class Downsample(nn.Module):
+    """Stride-2 conv with FLUX's asymmetric ``(0, 1, 0, 1)`` pad (flux2_ae.py:111-121)."""
+
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.conv(F.pad(x, (0, 1, 0, 1), mode="constant", value=0))
+
+
+class Upsample(nn.Module):
+    """Nearest-neighbour 2x followed by a 3x3 conv (flux2_ae.py:124-132)."""
+
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.conv(F.interpolate(x, scale_factor=2.0, mode="nearest"))
+
+
+class Encoder(nn.Module):
+    """FLUX.2 encoder.  Emits ``2 * z_channels`` moments; ``quant_conv`` is internal."""
+
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        z_channels: int,
+    ) -> None:
+        super().__init__()
+        # Declared first so the checkpoint key is `encoder.quant_conv.*`
+        # (flux2_ae.py:146) — diffusers keeps quant_conv outside the encoder.
+        self.quant_conv = nn.Conv2d(2 * z_channels, 2 * z_channels, 1)
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        self.conv_in = nn.Conv2d(in_channels, ch, kernel_size=3, stride=1, padding=1)
+
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.down = nn.ModuleList()
+        block_in = ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(num_res_blocks):
+                block.append(ResnetBlock(block_in, block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            # Empty at every level in this checkpoint: attention exists only in
+            # `mid` (flux2_ae.py:162).  Kept so the forward guard is meaningful.
+            down.attn = nn.ModuleList()
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in)
+            self.down.append(down)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(block_in, block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(block_in, block_in)
+
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        h = self.mid.block_2(self.mid.attn_1(self.mid.block_1(hs[-1])))
+        h = self.conv_out(swish(self.norm_out(h)))
+        return self.quant_conv(h)  # last op of the encoder (flux2_ae.py:207)
+
+
+class Decoder(nn.Module):
+    """FLUX.2 decoder.  ``post_quant_conv`` is internal and runs first."""
+
+    def __init__(
+        self,
+        ch: int,
+        out_ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+    ) -> None:
+        super().__init__()
+        # Checkpoint key `decoder.post_quant_conv.*` (flux2_ae.py:223).
+        self.post_quant_conv = nn.Conv2d(z_channels, z_channels, 1)
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(block_in, block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(block_in, block_in)
+
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(num_res_blocks + 1):
+                block.append(ResnetBlock(block_in, block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = nn.ModuleList()  # empty in this checkpoint (flux2_ae.py:249)
+            if i_level != 0:
+                up.upsample = Upsample(block_in)
+            self.up.insert(0, up)  # prepend so `up.<i>` indexes by resolution level
+
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: Tensor) -> Tensor:
+        z = self.post_quant_conv(z)  # first op of the decoder (flux2_ae.py:267)
+        upscale_dtype = next(self.up.parameters()).dtype
+
+        h = self.conv_in(z)
+        h = self.mid.block_2(self.mid.attn_1(self.mid.block_1(h)))
+        h = h.to(upscale_dtype)
+
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        return self.conv_out(swish(self.norm_out(h)))
+
+
+class AutoencoderFlux2(ModelMixin, ConfigMixin):
+    """Frozen FLUX.2 autoencoder with ReFlowSET's packed, BN-normalised latent.
+
+    ``encode`` maps ``[B, 3, H, W]`` in ``[-1, 1]`` to ``[B, 128, H/16, W/16]``
+    and ``decode`` inverts it.  The module is frozen: ``train()`` is a no-op that
+    always selects eval mode, and the latent BatchNorm is additionally forced to
+    eval on every call so no batch statistic can ever leak into the latent.
+
+    Args:
+        resolution: Nominal training resolution of the original autoencoder.
+            Only used to size bookkeeping attributes; any ``H``, ``W`` divisible
+            by 16 may be encoded.
+        in_channels: Input image channels (3).
+        ch: Base width.
+        out_ch: Output image channels (3).
+        ch_mult: Per-level width multipliers; ``len(ch_mult) - 1`` downsamples.
+        num_res_blocks: Residual blocks per level.
+        z_channels: Pre-pack latent channels (32).
+        patch_size: Space-to-depth factor applied after the encoder (2), which
+            takes the latent from 32 channels at ``H/8`` to 128 at ``H/16``.
+        bn_eps: Epsilon of the latent BatchNorm.  **1e-4**, not torch's 1e-5
+            (flux2_ae.py:331); using 1e-5 shifts the latent by up to 2.6e-5.
+    """
+
+    _supports_gradient_checkpointing = False
+
+    @register_to_config
+    def __init__(
+        self,
+        resolution: int = 256,
+        in_channels: int = 3,
+        ch: int = 128,
+        out_ch: int = 3,
+        ch_mult: tuple[int, ...] = (1, 2, 4, 4),
+        num_res_blocks: int = 2,
+        z_channels: int = 32,
+        patch_size: int = 2,
+        bn_eps: float = 1e-4,
+    ) -> None:
+        super().__init__()
+        ch_mult = list(ch_mult)
+        self.encoder = Encoder(
+            resolution=resolution,
+            in_channels=in_channels,
+            ch=ch,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            z_channels=z_channels,
+        )
+        self.decoder = Decoder(
+            ch=ch,
+            out_ch=out_ch,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            in_channels=in_channels,
+            resolution=resolution,
+            z_channels=z_channels,
+        )
+        # Per-channel latent normaliser with the checkpoint's running statistics.
+        # affine=False, so there is no weight/bias to load (flux2_ae.py:334-340).
+        self.bn = nn.BatchNorm2d(
+            patch_size * patch_size * z_channels,
+            eps=bn_eps,
+            momentum=0.1,
+            affine=False,
+            track_running_stats=True,
+        )
+
+    @property
+    def latent_channels(self) -> int:
+        """Channels of the public latent: ``patch_size**2 * z_channels`` = 128."""
+        return self.config.patch_size**2 * self.config.z_channels
+
+    @property
+    def spatial_factor(self) -> int:
+        """Total stride: 8x convolutional times ``patch_size`` packing = 16."""
+        return 2 ** (len(self.config.ch_mult) - 1) * self.config.patch_size
+
+    # ---- 2x2 space-to-depth pack / unpack -----------------------------------
+
+    def pack(self, z: Tensor) -> Tensor:
+        """``[B, C, H, W] -> [B, C*p*p, H/p, W/p]``, channel-major.
+
+        Bit-identical to the reference ``rearrange("... c (i pi) (j pj) -> ...
+        (c pi pj) i j")`` (flux2_ae.py:349-357).  Note this is **not** diffusers'
+        ``_pack_latents``, whose channel grouping is transposed.
+        """
+        return F.pixel_unshuffle(z, self.config.patch_size)
+
+    def unpack(self, z: Tensor) -> Tensor:
+        """Exact inverse of :meth:`pack` (flux2_ae.py:359-367)."""
+        return F.pixel_shuffle(z, self.config.patch_size)
+
+    # ---- latent normalisation ----------------------------------------------
+
+    def normalize(self, z: Tensor) -> Tensor:
+        """``(z - running_mean) / sqrt(running_var + bn_eps)``, per channel."""
+        self.bn.eval()  # forced every call (flux2_ae.py:372); train mode shifts z by ~1.67
+        return self.bn(z)
+
+    def inv_normalize(self, z: Tensor) -> Tensor:
+        """Exact inverse of :meth:`normalize` — same ``bn_eps`` (flux2_ae.py:375-379)."""
+        self.bn.eval()
+        s = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.config.bn_eps)
+        m = self.bn.running_mean.view(1, -1, 1, 1)
+        return z * s + m
+
+    # ---- public API ---------------------------------------------------------
+
+    @torch.no_grad()
+    def encode(self, x: Tensor) -> Tensor:
+        """Encode an image to the packed, normalised latent.
+
+        Args:
+            x: ``[B, 3, H, W]`` in ``[-1, 1]``; ``H`` and ``W`` divisible by 16.
+
+        Returns:
+            ``[B, 128, H/16, W/16]`` — the posterior **mean**, packed and
+            BN-normalised.  The encoder's log-variance chunk is discarded
+            (flux2_ae.py:396), so this is deterministic: there is no posterior
+            distribution object and nothing to sample.
+        """
+        if x.ndim != 4 or x.shape[1] != self.config.in_channels:
+            raise ValueError(
+                f"encode expects [B, {self.config.in_channels}, H, W], got {tuple(x.shape)}"
+            )
+        h, w = x.shape[-2:]
+        if h % self.spatial_factor or w % self.spatial_factor:
+            raise ValueError(
+                f"encode requires H and W divisible by {self.spatial_factor}, got {h}x{w}"
+            )
+        moments = self.encoder(x)
+        mean = torch.chunk(moments, 2, dim=1)[0]
+        return self.normalize(self.pack(mean))
+
+    @torch.no_grad()
+    def decode(self, z: Tensor) -> Tensor:
+        """Decode a packed, normalised latent ``[B, 128, h, w]`` to ``[B, 3, 16h, 16w]``.
+
+        The output is approximately ``[-1, 1]`` and is **not** clamped here; the
+        pipeline applies ``(x * 0.5 + 0.5).clamp(0, 1)``.
+        """
+        if z.ndim != 4 or z.shape[1] != self.latent_channels:
+            raise ValueError(
+                f"decode expects [B, {self.latent_channels}, h, w], got {tuple(z.shape)}"
+            )
+        return self.decoder(self.unpack(self.inv_normalize(z)))
+
+    # ---- construction / freezing -------------------------------------------
+
+    @classmethod
+    def from_single_file(
+        cls,
+        path: str | os.PathLike,
+        torch_dtype: torch.dtype = torch.float32,
+    ) -> "AutoencoderFlux2":
+        """Load the single-file ``ae.safetensors`` (BFL key names) with ``strict=True``.
+
+        The released file is the Apache-2.0 FLUX.2-klein-base-4B autoencoder
+        re-keyed to this layout; it is stored in bfloat16 and is upcast to
+        ``torch_dtype``.  ReFlowSET runs the autoencoder in float32.
+        """
+        from safetensors.torch import load_file
+
+        path = os.fspath(path)
+        if not os.path.isfile(path):
+            raise FileNotFoundError(
+                f"FLUX.2 autoencoder weights not found at: {path}. Expected the "
+                "single-file 'ae.safetensors' shipped with ReFlowSET."
+            )
+        model = cls()
+        model.load_state_dict(load_file(path, device="cpu"), strict=True)
+        model.to(dtype=torch_dtype)
+        model.eval()
+        model.requires_grad_(False)
+        return model
+
+    def train(self, mode: bool = True) -> "AutoencoderFlux2":
+        """The autoencoder is frozen: never leave eval mode (flux2_ae.py:437-439)."""
+        return super().train(False)

sar2opt/model_index.json

@@ -0,0 +1,16 @@
+{
+  "_class_name": "ReFlowSETPipeline",
+  "_diffusers_version": "0.37.1",
+  "transformer": [
+    "transformer_reflowset",
+    "ReFlowSETTransformer2DModel"
+  ],
+  "vae": [
+    "autoencoder_flux2",
+    "AutoencoderFlux2"
+  ],
+  "scheduler": [
+    "scheduler_flow_bridge",
+    "FlowBridgeScheduler"
+  ]
+}

sar2opt/pipeline.py

@@ -0,0 +1,267 @@
+"""ReFlowSET SAR -> EO translation pipeline."""
+
+from __future__ import annotations
+
+import warnings
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from PIL import Image
+
+from .autoencoder_flux2 import AutoencoderFlux2
+from .scheduler_flow_bridge import FlowBridgeScheduler
+from .transformer_reflowset import ReFlowSETTransformer2DModel
+
+#: PIL modes the SAR loader accepts.  The reference loader calls ``np.array(im)``
+#: with no ``convert()`` (datasets.py:454, 574), so a 16-bit (``I;16``) or
+#: palette (``P``) raster would flow straight into ``x / 127.5 - 1`` and be
+#: badly out of range.  That is an unguarded trap upstream; it is guarded here.
+_ACCEPTED_SAR_MODES = ("L", "RGB", "RGBA")
+
+
+class ReFlowSETPipeline(DiffusionPipeline):
+    """Generate an EO image from a SAR image with ReFlowSET's flow bridge.
+
+    Args:
+        transformer: The velocity transformer.
+        vae: The frozen FLUX.2 autoencoder that defines the latent space.
+        scheduler: The Design-B flow-bridge Euler solver.
+
+    To reproduce the paper's numbers, sample at ``num_inference_steps=50``,
+    ``guidance_scale=1.5``, float32, one image per call, with a generator freshly
+    seeded to 2024 on the compute device before each call — every test image in
+    the reported evaluation starts from the same seeded noise draw, and CPU-drawn
+    noise does not reproduce a CUDA draw.
+    """
+
+    model_cpu_offload_seq = "transformer->vae"
+
+    def __init__(
+        self,
+        transformer: ReFlowSETTransformer2DModel,
+        vae: AutoencoderFlux2,
+        scheduler: FlowBridgeScheduler,
+    ) -> None:
+        super().__init__()
+        self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)
+
+    # ---- preprocessing (datasets.py:124-205, 452-474, 572-589) --------------
+
+    @staticmethod
+    def _sar_hwc(raster: Union[Image.Image, np.ndarray]) -> np.ndarray:
+        """SAR raster -> ``[H, W, C]`` float32 in ``[0, 255]``, collapsed to 1 channel."""
+        if isinstance(raster, Image.Image):
+            if raster.mode not in _ACCEPTED_SAR_MODES:
+                raise ValueError(
+                    f"SAR image mode {raster.mode!r} is not an 8-bit display raster; expected "
+                    f"one of {_ACCEPTED_SAR_MODES}. ReFlowSET was trained on 8-bit display "
+                    "quicklooks (sar_value_domain='display_png'); convert with .convert('L') "
+                    "and be aware that the contrast stretch you choose is part of the input."
+                )
+            # No .convert() on the SAR side, matching datasets.py:454, 574.
+            arr = np.array(raster)
+        else:
+            arr = np.asarray(raster)
+        if arr.ndim == 2:  # PIL mode "L" (datasets.py:171-172)
+            arr = arr[:, :, None]
+        if arr.shape[-1] == 4:  # drop a container alpha channel (datasets.py:173-174)
+            arr = arr[..., :3]
+        arr = arr.astype(np.float32)
+        if arr.shape[-1] > 1:
+            # Exact-equality test, tol=0.0 (datasets.py:100-111): a display RGB
+            # quicklook collapses to its single amplitude channel.
+            if np.abs(arr - arr[..., :1]).max() == 0.0:
+                arr = arr[..., :1]
+            else:
+                warnings.warn(
+                    "SAR raster has non-identical colour channels; feeding all 3 to the "
+                    "frozen encoder. The released arms were trained on single-channel "
+                    "amplitude quicklooks, so this is an undeclared input.",
+                    RuntimeWarning,
+                    stacklevel=3,
+                )
+        return arr
+
+    @staticmethod
+    def _center_crop(arr: np.ndarray, crop: int) -> np.ndarray:
+        """Center-crop ``[H, W, C]`` to ``crop`` — **never** resize (datasets.py:145-165).
+
+        The offsets are albumentations' ``CenterCrop`` arithmetic ``(n - c) // 2``:
+        the SAR2Opt protocol takes the central 512 of 600 at offset 44.
+        """
+        h, w = arr.shape[:2]
+        if h < crop or w < crop:
+            raise ValueError(
+                f"image {h}x{w} is smaller than the requested crop {crop}; ReFlowSET never "
+                "upscales an input"
+            )
+        top, left = (h - crop) // 2, (w - crop) // 2
+        return arr[top : top + crop, left : left + crop]
+
+    def preprocess(
+        self,
+        sar: Union[Image.Image, np.ndarray, torch.Tensor, list],
+        crop: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Build the model-boundary SAR tensor ``[B, 3, H, W]`` in ``[-1, 1]``.
+
+        Args:
+            sar: A PIL image, a list of PIL images, an ``[H, W]`` / ``[H, W, C]``
+                uint8 array, or a float tensor already in ``[-1, 1]`` shaped
+                ``[H, W]``, ``[C, H, W]`` or ``[B, C, H, W]``.
+            crop: Center-crop size applied before normalisation.  ``None``
+                center-crops to the arm's own training resolution when the
+                raster is larger and not already a multiple of the latent
+                stride -- which is exactly the SAR2Opt 600 -> 512 protocol the
+                reported numbers use.  Pass an explicit size to override, or
+                ``0`` to keep the native raster and fail loudly if it does not
+                fit.
+
+        Images are read as 8-bit display rasters and mapped to ``[-1, 1]`` by
+        ``x / 127.5 - 1`` (datasets.py:124-126) with no per-image statistics, no
+        percentile stretch and no resize.  The single SAR channel is then
+        replicated to 3 at the model boundary (evaluate.py:566-570), because the
+        frozen FLUX.2 encoder is the same one that encodes EO — ReFlowSET has no
+        separate SAR encoder.
+        """
+        if crop == 0:
+            crop = None
+        elif crop is None:
+            # Fall back to the resolution this arm was trained at. Cropping is
+            # the protocol (train.py random-crops, evaluate.py center-crops);
+            # ReFlowSET never resizes, so an un-croppable raster is an error
+            # rather than something to silently rescale.
+            crop = self.transformer.config.sample_size
+
+        if isinstance(sar, torch.Tensor):
+            x = sar.float()
+            if x.ndim == 2:
+                x = x[None, None]
+            elif x.ndim == 3:
+                x = x[None]
+            elif x.ndim != 4:
+                raise ValueError(f"SAR tensor must have 2, 3 or 4 dims, got {tuple(sar.shape)}")
+            h, w = x.shape[-2:]
+            if crop is not None and (h, w) != (crop, crop):
+                if h < crop or w < crop:
+                    raise ValueError(f"tensor {h}x{w} is smaller than the requested crop {crop}")
+                top, left = (h - crop) // 2, (w - crop) // 2
+                x = x[..., top : top + crop, left : left + crop]
+        else:
+            images = sar if isinstance(sar, list) else [sar]
+            arrays = []
+            for item in images:
+                if not isinstance(item, (Image.Image, np.ndarray)):
+                    raise TypeError(f"unsupported SAR input type {type(item)!r}")
+                arr = self._sar_hwc(item)
+                if crop is not None and arr.shape[:2] != (crop, crop):
+                    arr = self._center_crop(arr, crop)
+                arrays.append(np.ascontiguousarray(arr.transpose(2, 0, 1)))
+            x = torch.from_numpy(np.stack(arrays)) / 127.5 - 1.0
+
+        # Train-side clamp (train.py:770); a no-op on 8-bit input, which maps
+        # exactly onto [-1, 1].
+        x = x.clamp(-1.0, 1.0)
+        if x.shape[1] == 1:
+            x = x.repeat(1, 3, 1, 1)
+        elif x.shape[1] != 3:
+            raise ValueError(
+                f"the frozen FLUX.2 encoder takes 1 or 3 SAR channels, got {x.shape[1]}"
+            )
+        factor = self.vae.spatial_factor
+        if x.shape[-2] % factor or x.shape[-1] % factor:
+            raise ValueError(
+                f"SAR size {x.shape[-2]}x{x.shape[-1]} must be divisible by {factor}; pass "
+                "crop= to center-crop (ReFlowSET never resizes)"
+            )
+        return x
+
+    # ---- postprocessing -----------------------------------------------------
+
+    @staticmethod
+    def _to_pil(images: torch.Tensor) -> list[Image.Image]:
+        """``[B, 3, H, W]`` in ``[0, 1]`` -> PIL, quantised round-half-up.
+
+        ``255 * x + 0.5`` truncated is what ``torchvision.utils.save_image``
+        does and is therefore what the released PNGs contain; numpy's
+        ``round()`` is banker's rounding and would differ on exact halves.
+        """
+        arr = (images * 255 + 0.5).clamp(0, 255).to(torch.uint8)
+        arr = arr.permute(0, 2, 3, 1).cpu().numpy()
+        return [Image.fromarray(a) for a in arr]
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        sar: Union[Image.Image, np.ndarray, torch.Tensor, list],
+        num_inference_steps: int = 50,
+        guidance_scale: float = 1.5,
+        generator: Optional[Union[torch.Generator, list[torch.Generator]]] = None,
+        output_type: str = "pil",
+        crop: Optional[int] = None,
+        return_dict: bool = True,
+    ) -> Union[ImagePipelineOutput, tuple[list]]:
+        """Translate a SAR image into an EO image.
+
+        Args:
+            sar: SAR input; see :meth:`preprocess`.
+            num_inference_steps: NFE, the number of velocity evaluations.  The
+                paper's main results are NFE 50; NFE 4 is the efficiency
+                operating point and trades FID for PSNR/SSIM, so the two must
+                not be mixed in one comparison.
+            guidance_scale: Classifier-free guidance scale.  1.5 is the published
+                setting; 1.0 disables guidance and halves the cost.
+            generator: Generator for the initial noise.  Create it on the compute
+                device — CPU-drawn noise does not reproduce a CUDA draw.
+            output_type: ``"pil"``, ``"np"`` or ``"pt"``.
+            crop: Center-crop size applied to the SAR input before encoding.
+            return_dict: Return an ``ImagePipelineOutput`` instead of a tuple.
+
+        Returns:
+            The generated EO image(s) in ``[0, 1]`` (or as PIL).
+        """
+        if output_type not in ("pil", "np", "pt"):
+            raise ValueError(f"output_type must be 'pil', 'np' or 'pt', got {output_type!r}")
+
+        device = self._execution_device
+        dtype = self.transformer.dtype
+
+        sar_pm1 = self.preprocess(sar, crop=crop).to(device=device, dtype=self.vae.dtype)
+        # The SAR condition is encoded by the SAME frozen autoencoder that
+        # defines the EO latent space (evaluate.py:553-577).
+        z_s = self.vae.encode(sar_pm1).to(dtype)
+
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        # Design B: the bridge starts at t = 0 from pure Gaussian noise
+        # (bridge.py:409-433), NOT from the SAR latent.
+        latents = randn_tensor(z_s.shape, generator=generator, device=device, dtype=z_s.dtype)
+
+        for t in self.progress_bar(self.scheduler.timesteps):
+            timestep = t.expand(latents.shape[0])
+            velocity = self.transformer(latents, timestep, z_s, return_dict=False)[0]
+            if guidance_scale != 1.0:
+                # Two passes; the null branch is cond=None, which the transformer
+                # turns into an all-zero conditioning latent (bridge.py:530-535).
+                uncond = self.transformer(latents, timestep, None, return_dict=False)[0]
+                velocity = uncond + guidance_scale * (velocity - uncond)
+            latents = self.scheduler.step(velocity, t, latents, return_dict=False)[0]
+
+        image = self.vae.decode(latents.to(self.vae.dtype))
+        # `--denorm standard` (evaluate.py:292-295). The `legacy` C-DiffSET
+        # convention `(x + 0.5).clamp(0, 1)` is a 2x contrast stretch and must
+        # not be used with these numbers.
+        image = (image * 0.5 + 0.5).clamp(0.0, 1.0)
+
+        self.maybe_free_model_hooks()
+
+        if output_type == "pil":
+            image = self._to_pil(image)
+        elif output_type == "np":
+            image = image.permute(0, 2, 3, 1).float().cpu().numpy()
+
+        if not return_dict:
+            return (image,)
+        return ImagePipelineOutput(images=image)

sar2opt/pipeline_reflowset.py

@@ -0,0 +1,267 @@
+"""ReFlowSET SAR -> EO translation pipeline."""
+
+from __future__ import annotations
+
+import warnings
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+from PIL import Image
+
+from .autoencoder_flux2 import AutoencoderFlux2
+from .scheduler_flow_bridge import FlowBridgeScheduler
+from .transformer_reflowset import ReFlowSETTransformer2DModel
+
+#: PIL modes the SAR loader accepts.  The reference loader calls ``np.array(im)``
+#: with no ``convert()`` (datasets.py:454, 574), so a 16-bit (``I;16``) or
+#: palette (``P``) raster would flow straight into ``x / 127.5 - 1`` and be
+#: badly out of range.  That is an unguarded trap upstream; it is guarded here.
+_ACCEPTED_SAR_MODES = ("L", "RGB", "RGBA")
+
+
+class ReFlowSETPipeline(DiffusionPipeline):
+    """Generate an EO image from a SAR image with ReFlowSET's flow bridge.
+
+    Args:
+        transformer: The velocity transformer.
+        vae: The frozen FLUX.2 autoencoder that defines the latent space.
+        scheduler: The Design-B flow-bridge Euler solver.
+
+    To reproduce the paper's numbers, sample at ``num_inference_steps=50``,
+    ``guidance_scale=1.5``, float32, one image per call, with a generator freshly
+    seeded to 2024 on the compute device before each call — every test image in
+    the reported evaluation starts from the same seeded noise draw, and CPU-drawn
+    noise does not reproduce a CUDA draw.
+    """
+
+    model_cpu_offload_seq = "transformer->vae"
+
+    def __init__(
+        self,
+        transformer: ReFlowSETTransformer2DModel,
+        vae: AutoencoderFlux2,
+        scheduler: FlowBridgeScheduler,
+    ) -> None:
+        super().__init__()
+        self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)
+
+    # ---- preprocessing (datasets.py:124-205, 452-474, 572-589) --------------
+
+    @staticmethod
+    def _sar_hwc(raster: Union[Image.Image, np.ndarray]) -> np.ndarray:
+        """SAR raster -> ``[H, W, C]`` float32 in ``[0, 255]``, collapsed to 1 channel."""
+        if isinstance(raster, Image.Image):
+            if raster.mode not in _ACCEPTED_SAR_MODES:
+                raise ValueError(
+                    f"SAR image mode {raster.mode!r} is not an 8-bit display raster; expected "
+                    f"one of {_ACCEPTED_SAR_MODES}. ReFlowSET was trained on 8-bit display "
+                    "quicklooks (sar_value_domain='display_png'); convert with .convert('L') "
+                    "and be aware that the contrast stretch you choose is part of the input."
+                )
+            # No .convert() on the SAR side, matching datasets.py:454, 574.
+            arr = np.array(raster)
+        else:
+            arr = np.asarray(raster)
+        if arr.ndim == 2:  # PIL mode "L" (datasets.py:171-172)
+            arr = arr[:, :, None]
+        if arr.shape[-1] == 4:  # drop a container alpha channel (datasets.py:173-174)
+            arr = arr[..., :3]
+        arr = arr.astype(np.float32)
+        if arr.shape[-1] > 1:
+            # Exact-equality test, tol=0.0 (datasets.py:100-111): a display RGB
+            # quicklook collapses to its single amplitude channel.
+            if np.abs(arr - arr[..., :1]).max() == 0.0:
+                arr = arr[..., :1]
+            else:
+                warnings.warn(
+                    "SAR raster has non-identical colour channels; feeding all 3 to the "
+                    "frozen encoder. The released arms were trained on single-channel "
+                    "amplitude quicklooks, so this is an undeclared input.",
+                    RuntimeWarning,
+                    stacklevel=3,
+                )
+        return arr
+
+    @staticmethod
+    def _center_crop(arr: np.ndarray, crop: int) -> np.ndarray:
+        """Center-crop ``[H, W, C]`` to ``crop`` — **never** resize (datasets.py:145-165).
+
+        The offsets are albumentations' ``CenterCrop`` arithmetic ``(n - c) // 2``:
+        the SAR2Opt protocol takes the central 512 of 600 at offset 44.
+        """
+        h, w = arr.shape[:2]
+        if h < crop or w < crop:
+            raise ValueError(
+                f"image {h}x{w} is smaller than the requested crop {crop}; ReFlowSET never "
+                "upscales an input"
+            )
+        top, left = (h - crop) // 2, (w - crop) // 2
+        return arr[top : top + crop, left : left + crop]
+
+    def preprocess(
+        self,
+        sar: Union[Image.Image, np.ndarray, torch.Tensor, list],
+        crop: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Build the model-boundary SAR tensor ``[B, 3, H, W]`` in ``[-1, 1]``.
+
+        Args:
+            sar: A PIL image, a list of PIL images, an ``[H, W]`` / ``[H, W, C]``
+                uint8 array, or a float tensor already in ``[-1, 1]`` shaped
+                ``[H, W]``, ``[C, H, W]`` or ``[B, C, H, W]``.
+            crop: Center-crop size applied before normalisation.  ``None``
+                center-crops to the arm's own training resolution when the
+                raster is larger and not already a multiple of the latent
+                stride -- which is exactly the SAR2Opt 600 -> 512 protocol the
+                reported numbers use.  Pass an explicit size to override, or
+                ``0`` to keep the native raster and fail loudly if it does not
+                fit.
+
+        Images are read as 8-bit display rasters and mapped to ``[-1, 1]`` by
+        ``x / 127.5 - 1`` (datasets.py:124-126) with no per-image statistics, no
+        percentile stretch and no resize.  The single SAR channel is then
+        replicated to 3 at the model boundary (evaluate.py:566-570), because the
+        frozen FLUX.2 encoder is the same one that encodes EO — ReFlowSET has no
+        separate SAR encoder.
+        """
+        if crop == 0:
+            crop = None
+        elif crop is None:
+            # Fall back to the resolution this arm was trained at. Cropping is
+            # the protocol (train.py random-crops, evaluate.py center-crops);
+            # ReFlowSET never resizes, so an un-croppable raster is an error
+            # rather than something to silently rescale.
+            crop = self.transformer.config.sample_size
+
+        if isinstance(sar, torch.Tensor):
+            x = sar.float()
+            if x.ndim == 2:
+                x = x[None, None]
+            elif x.ndim == 3:
+                x = x[None]
+            elif x.ndim != 4:
+                raise ValueError(f"SAR tensor must have 2, 3 or 4 dims, got {tuple(sar.shape)}")
+            h, w = x.shape[-2:]
+            if crop is not None and (h, w) != (crop, crop):
+                if h < crop or w < crop:
+                    raise ValueError(f"tensor {h}x{w} is smaller than the requested crop {crop}")
+                top, left = (h - crop) // 2, (w - crop) // 2
+                x = x[..., top : top + crop, left : left + crop]
+        else:
+            images = sar if isinstance(sar, list) else [sar]
+            arrays = []
+            for item in images:
+                if not isinstance(item, (Image.Image, np.ndarray)):
+                    raise TypeError(f"unsupported SAR input type {type(item)!r}")
+                arr = self._sar_hwc(item)
+                if crop is not None and arr.shape[:2] != (crop, crop):
+                    arr = self._center_crop(arr, crop)
+                arrays.append(np.ascontiguousarray(arr.transpose(2, 0, 1)))
+            x = torch.from_numpy(np.stack(arrays)) / 127.5 - 1.0
+
+        # Train-side clamp (train.py:770); a no-op on 8-bit input, which maps
+        # exactly onto [-1, 1].
+        x = x.clamp(-1.0, 1.0)
+        if x.shape[1] == 1:
+            x = x.repeat(1, 3, 1, 1)
+        elif x.shape[1] != 3:
+            raise ValueError(
+                f"the frozen FLUX.2 encoder takes 1 or 3 SAR channels, got {x.shape[1]}"
+            )
+        factor = self.vae.spatial_factor
+        if x.shape[-2] % factor or x.shape[-1] % factor:
+            raise ValueError(
+                f"SAR size {x.shape[-2]}x{x.shape[-1]} must be divisible by {factor}; pass "
+                "crop= to center-crop (ReFlowSET never resizes)"
+            )
+        return x
+
+    # ---- postprocessing -----------------------------------------------------
+
+    @staticmethod
+    def _to_pil(images: torch.Tensor) -> list[Image.Image]:
+        """``[B, 3, H, W]`` in ``[0, 1]`` -> PIL, quantised round-half-up.
+
+        ``255 * x + 0.5`` truncated is what ``torchvision.utils.save_image``
+        does and is therefore what the released PNGs contain; numpy's
+        ``round()`` is banker's rounding and would differ on exact halves.
+        """
+        arr = (images * 255 + 0.5).clamp(0, 255).to(torch.uint8)
+        arr = arr.permute(0, 2, 3, 1).cpu().numpy()
+        return [Image.fromarray(a) for a in arr]
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        sar: Union[Image.Image, np.ndarray, torch.Tensor, list],
+        num_inference_steps: int = 50,
+        guidance_scale: float = 1.5,
+        generator: Optional[Union[torch.Generator, list[torch.Generator]]] = None,
+        output_type: str = "pil",
+        crop: Optional[int] = None,
+        return_dict: bool = True,
+    ) -> Union[ImagePipelineOutput, tuple[list]]:
+        """Translate a SAR image into an EO image.
+
+        Args:
+            sar: SAR input; see :meth:`preprocess`.
+            num_inference_steps: NFE, the number of velocity evaluations.  The
+                paper's main results are NFE 50; NFE 4 is the efficiency
+                operating point and trades FID for PSNR/SSIM, so the two must
+                not be mixed in one comparison.
+            guidance_scale: Classifier-free guidance scale.  1.5 is the published
+                setting; 1.0 disables guidance and halves the cost.
+            generator: Generator for the initial noise.  Create it on the compute
+                device — CPU-drawn noise does not reproduce a CUDA draw.
+            output_type: ``"pil"``, ``"np"`` or ``"pt"``.
+            crop: Center-crop size applied to the SAR input before encoding.
+            return_dict: Return an ``ImagePipelineOutput`` instead of a tuple.
+
+        Returns:
+            The generated EO image(s) in ``[0, 1]`` (or as PIL).
+        """
+        if output_type not in ("pil", "np", "pt"):
+            raise ValueError(f"output_type must be 'pil', 'np' or 'pt', got {output_type!r}")
+
+        device = self._execution_device
+        dtype = self.transformer.dtype
+
+        sar_pm1 = self.preprocess(sar, crop=crop).to(device=device, dtype=self.vae.dtype)
+        # The SAR condition is encoded by the SAME frozen autoencoder that
+        # defines the EO latent space (evaluate.py:553-577).
+        z_s = self.vae.encode(sar_pm1).to(dtype)
+
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        # Design B: the bridge starts at t = 0 from pure Gaussian noise
+        # (bridge.py:409-433), NOT from the SAR latent.
+        latents = randn_tensor(z_s.shape, generator=generator, device=device, dtype=z_s.dtype)
+
+        for t in self.progress_bar(self.scheduler.timesteps):
+            timestep = t.expand(latents.shape[0])
+            velocity = self.transformer(latents, timestep, z_s, return_dict=False)[0]
+            if guidance_scale != 1.0:
+                # Two passes; the null branch is cond=None, which the transformer
+                # turns into an all-zero conditioning latent (bridge.py:530-535).
+                uncond = self.transformer(latents, timestep, None, return_dict=False)[0]
+                velocity = uncond + guidance_scale * (velocity - uncond)
+            latents = self.scheduler.step(velocity, t, latents, return_dict=False)[0]
+
+        image = self.vae.decode(latents.to(self.vae.dtype))
+        # `--denorm standard` (evaluate.py:292-295). The `legacy` C-DiffSET
+        # convention `(x + 0.5).clamp(0, 1)` is a 2x contrast stretch and must
+        # not be used with these numbers.
+        image = (image * 0.5 + 0.5).clamp(0.0, 1.0)
+
+        self.maybe_free_model_hooks()
+
+        if output_type == "pil":
+            image = self._to_pil(image)
+        elif output_type == "np":
+            image = image.permute(0, 2, 3, 1).float().cpu().numpy()
+
+        if not return_dict:
+            return (image,)
+        return ImagePipelineOutput(images=image)

sar2opt/scheduler/scheduler_config.json

@@ -0,0 +1,5 @@
+{
+  "_class_name": "FlowBridgeScheduler",
+  "_diffusers_version": "0.37.1",
+  "t_end": 1.0
+}

sar2opt/scheduler/scheduler_flow_bridge.py

@@ -0,0 +1,176 @@
+"""ReFlowSET's Design-B flow bridge and its explicit-Euler solver.
+
+Forward (training) process, with ``eps ~ N(0, I)`` and ``z_e`` the EO latent::
+
+    z_t = (1 - t) * eps + t * z_e          (bridge.py:311, sigma_b = 0)
+    u*  = z_e - eps                        (bridge.py:328 at sigma_b = 0)
+
+Sampling starts from ``z_0 ~ N(0, I)`` and integrates the predicted velocity
+with explicit Euler on a uniform grid ``linspace(0, t_end, nfe + 1)``
+(bridge.py:519, 536).  The bridge is deterministic: ``sigma_b = 0``, so no
+stochastic term ever executes, and the only randomness in a sample is the
+initial noise draw.
+
+**Time direction.**  ``t = 0`` is NOISE and ``t = 1`` is DATA, and the solver
+integrates ``t`` **ascending** (bridge.py:86-88).  That is the opposite of
+`diffusers`' ``sigma`` convention: setting ``sigma := 1 - t`` recovers
+``FlowMatchEulerDiscreteScheduler``'s interpolation, but then this bridge's
+velocity is the **negative** of the diffusers flow-matching target and the
+network must still be fed ``1 - sigma``.  This scheduler keeps ReFlowSET's own
+sign and direction so neither flip is needed; ``timesteps`` therefore *increase*
+from 0 towards 1, unlike every noise-schedule scheduler in `diffusers`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class FlowBridgeSchedulerOutput(BaseOutput):
+    """Output of :meth:`FlowBridgeScheduler.step`.
+
+    Args:
+        prev_sample: The bridge state at the next time on the grid.
+    """
+
+    prev_sample: torch.Tensor
+
+
+class FlowBridgeScheduler(SchedulerMixin, ConfigMixin):
+    """Explicit-Euler solver for ReFlowSET's Design-B flow bridge.
+
+    Args:
+        t_end: End time of the integration grid (1.0 — the EO endpoint).  The
+            model is evaluated at ``linspace(0, t_end, nfe + 1)[:-1]`` and the
+            final Euler step lands on ``t_end``; the network is never queried at
+            ``t = t_end``.
+    """
+
+    order = 1
+
+    @register_to_config
+    def __init__(self, t_end: float = 1.0) -> None:
+        if not 0.0 < t_end <= 1.0:
+            raise ValueError(f"t_end must lie in (0, 1], got {t_end}")
+        self._grid: Optional[torch.Tensor] = None
+        self._step_index: Optional[int] = None
+        self.num_inference_steps: Optional[int] = None
+
+    @property
+    def timesteps(self) -> torch.Tensor:
+        """The ``nfe`` bridge times at which the model is evaluated, ascending."""
+        if self._grid is None:
+            raise ValueError("call set_timesteps() before reading timesteps")
+        return self._grid[:-1]
+
+    @property
+    def step_index(self) -> Optional[int]:
+        """Index of the next grid interval; ``None`` until the first :meth:`step`."""
+        return self._step_index
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        device: Optional[Union[str, torch.device]] = None,
+    ) -> None:
+        """Build the uniform grid ``linspace(0, t_end, num_inference_steps + 1)``.
+
+        Args:
+            num_inference_steps: NFE — the number of velocity evaluations.
+                50 reproduces the paper's main results; 4 is the efficiency
+                operating point.
+            device: Device the grid is built on.
+
+        There is no shift, no dynamic shifting, no Karras or exponential
+        spacing, and no timestep-spacing option: the reference solver uses a
+        plain uniform grid (bridge.py:519).
+        """
+        if num_inference_steps < 1:
+            raise ValueError(f"num_inference_steps must be >= 1, got {num_inference_steps}")
+        self.num_inference_steps = num_inference_steps
+        self._grid = torch.linspace(
+            0.0, self.config.t_end, num_inference_steps + 1, device=device, dtype=torch.float32
+        )
+        self._step_index = 0
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[float, torch.Tensor],
+        sample: torch.Tensor,
+        return_dict: bool = True,
+    ) -> Union[FlowBridgeSchedulerOutput, tuple[torch.Tensor]]:
+        """One explicit-Euler step: ``z + (t_next - t_cur) * v`` (bridge.py:536).
+
+        Args:
+            model_output: The predicted velocity ``dz/dt`` at ``timestep``,
+                already classifier-free-guided by the caller.
+            timestep: The current bridge time.  Present for API compatibility and
+                checked against the grid; the step size comes from the grid.
+            sample: The current bridge state.
+            return_dict: Return a :class:`FlowBridgeSchedulerOutput` instead of a
+                tuple.
+
+        Steps must be taken in order, starting from the first entry of
+        :attr:`timesteps`.
+        """
+        if self._grid is None or self._step_index is None:
+            raise ValueError("call set_timesteps() before step()")
+        if self._step_index >= self.num_inference_steps:
+            raise ValueError(
+                f"already took {self.num_inference_steps} steps; call set_timesteps() again"
+            )
+        t_cur, t_next = self._grid[self._step_index], self._grid[self._step_index + 1]
+        if not torch.isclose(torch.as_tensor(timestep, dtype=torch.float32).to(t_cur.device), t_cur):
+            raise ValueError(
+                f"step {self._step_index} expects timestep {t_cur.item()}, got {float(timestep)}; "
+                "the flow bridge must be integrated in ascending grid order"
+            )
+
+        # The state is carried in float32 even if the model ran lower (bridge.py:515-517).
+        dtype = sample.dtype if sample.dtype in (torch.float32, torch.float64) else torch.float32
+        prev_sample = sample.to(dtype) + (t_next - t_cur) * model_output.to(dtype)
+        prev_sample = prev_sample.to(sample.dtype)
+
+        self._step_index += 1
+        if not return_dict:
+            return (prev_sample,)
+        return FlowBridgeSchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        """The training-side bridge state ``z_t = (1 - t) * eps + t * z_e`` (bridge.py:311).
+
+        Args:
+            original_samples: The EO latent ``z_e`` (the ``t = 1`` endpoint).
+            noise: ``eps ~ N(0, I)`` (the ``t = 0`` endpoint).
+            timesteps: Bridge times in ``[0, 1]``, broadcastable over the batch.
+        """
+        t = timesteps.to(original_samples.device, original_samples.dtype)
+        t = t.view(-1, *([1] * (original_samples.ndim - 1)))
+        return (1.0 - t) * noise + t * original_samples
+
+    def get_velocity(
+        self,
+        sample: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        """The training target ``u* = z_e - eps`` (bridge.py:328 at ``sigma_b = 0``).
+
+        Constant along the path, hence independent of ``timesteps``; the argument
+        is kept for `diffusers` API compatibility.
+        """
+        del timesteps
+        return sample - noise

sar2opt/scheduler_flow_bridge.py

@@ -0,0 +1,176 @@
+"""ReFlowSET's Design-B flow bridge and its explicit-Euler solver.
+
+Forward (training) process, with ``eps ~ N(0, I)`` and ``z_e`` the EO latent::
+
+    z_t = (1 - t) * eps + t * z_e          (bridge.py:311, sigma_b = 0)
+    u*  = z_e - eps                        (bridge.py:328 at sigma_b = 0)
+
+Sampling starts from ``z_0 ~ N(0, I)`` and integrates the predicted velocity
+with explicit Euler on a uniform grid ``linspace(0, t_end, nfe + 1)``
+(bridge.py:519, 536).  The bridge is deterministic: ``sigma_b = 0``, so no
+stochastic term ever executes, and the only randomness in a sample is the
+initial noise draw.
+
+**Time direction.**  ``t = 0`` is NOISE and ``t = 1`` is DATA, and the solver
+integrates ``t`` **ascending** (bridge.py:86-88).  That is the opposite of
+`diffusers`' ``sigma`` convention: setting ``sigma := 1 - t`` recovers
+``FlowMatchEulerDiscreteScheduler``'s interpolation, but then this bridge's
+velocity is the **negative** of the diffusers flow-matching target and the
+network must still be fed ``1 - sigma``.  This scheduler keeps ReFlowSET's own
+sign and direction so neither flip is needed; ``timesteps`` therefore *increase*
+from 0 towards 1, unlike every noise-schedule scheduler in `diffusers`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class FlowBridgeSchedulerOutput(BaseOutput):
+    """Output of :meth:`FlowBridgeScheduler.step`.
+
+    Args:
+        prev_sample: The bridge state at the next time on the grid.
+    """
+
+    prev_sample: torch.Tensor
+
+
+class FlowBridgeScheduler(SchedulerMixin, ConfigMixin):
+    """Explicit-Euler solver for ReFlowSET's Design-B flow bridge.
+
+    Args:
+        t_end: End time of the integration grid (1.0 — the EO endpoint).  The
+            model is evaluated at ``linspace(0, t_end, nfe + 1)[:-1]`` and the
+            final Euler step lands on ``t_end``; the network is never queried at
+            ``t = t_end``.
+    """
+
+    order = 1
+
+    @register_to_config
+    def __init__(self, t_end: float = 1.0) -> None:
+        if not 0.0 < t_end <= 1.0:
+            raise ValueError(f"t_end must lie in (0, 1], got {t_end}")
+        self._grid: Optional[torch.Tensor] = None
+        self._step_index: Optional[int] = None
+        self.num_inference_steps: Optional[int] = None
+
+    @property
+    def timesteps(self) -> torch.Tensor:
+        """The ``nfe`` bridge times at which the model is evaluated, ascending."""
+        if self._grid is None:
+            raise ValueError("call set_timesteps() before reading timesteps")
+        return self._grid[:-1]
+
+    @property
+    def step_index(self) -> Optional[int]:
+        """Index of the next grid interval; ``None`` until the first :meth:`step`."""
+        return self._step_index
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        device: Optional[Union[str, torch.device]] = None,
+    ) -> None:
+        """Build the uniform grid ``linspace(0, t_end, num_inference_steps + 1)``.
+
+        Args:
+            num_inference_steps: NFE — the number of velocity evaluations.
+                50 reproduces the paper's main results; 4 is the efficiency
+                operating point.
+            device: Device the grid is built on.
+
+        There is no shift, no dynamic shifting, no Karras or exponential
+        spacing, and no timestep-spacing option: the reference solver uses a
+        plain uniform grid (bridge.py:519).
+        """
+        if num_inference_steps < 1:
+            raise ValueError(f"num_inference_steps must be >= 1, got {num_inference_steps}")
+        self.num_inference_steps = num_inference_steps
+        self._grid = torch.linspace(
+            0.0, self.config.t_end, num_inference_steps + 1, device=device, dtype=torch.float32
+        )
+        self._step_index = 0
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[float, torch.Tensor],
+        sample: torch.Tensor,
+        return_dict: bool = True,
+    ) -> Union[FlowBridgeSchedulerOutput, tuple[torch.Tensor]]:
+        """One explicit-Euler step: ``z + (t_next - t_cur) * v`` (bridge.py:536).
+
+        Args:
+            model_output: The predicted velocity ``dz/dt`` at ``timestep``,
+                already classifier-free-guided by the caller.
+            timestep: The current bridge time.  Present for API compatibility and
+                checked against the grid; the step size comes from the grid.
+            sample: The current bridge state.
+            return_dict: Return a :class:`FlowBridgeSchedulerOutput` instead of a
+                tuple.
+
+        Steps must be taken in order, starting from the first entry of
+        :attr:`timesteps`.
+        """
+        if self._grid is None or self._step_index is None:
+            raise ValueError("call set_timesteps() before step()")
+        if self._step_index >= self.num_inference_steps:
+            raise ValueError(
+                f"already took {self.num_inference_steps} steps; call set_timesteps() again"
+            )
+        t_cur, t_next = self._grid[self._step_index], self._grid[self._step_index + 1]
+        if not torch.isclose(torch.as_tensor(timestep, dtype=torch.float32).to(t_cur.device), t_cur):
+            raise ValueError(
+                f"step {self._step_index} expects timestep {t_cur.item()}, got {float(timestep)}; "
+                "the flow bridge must be integrated in ascending grid order"
+            )
+
+        # The state is carried in float32 even if the model ran lower (bridge.py:515-517).
+        dtype = sample.dtype if sample.dtype in (torch.float32, torch.float64) else torch.float32
+        prev_sample = sample.to(dtype) + (t_next - t_cur) * model_output.to(dtype)
+        prev_sample = prev_sample.to(sample.dtype)
+
+        self._step_index += 1
+        if not return_dict:
+            return (prev_sample,)
+        return FlowBridgeSchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        """The training-side bridge state ``z_t = (1 - t) * eps + t * z_e`` (bridge.py:311).
+
+        Args:
+            original_samples: The EO latent ``z_e`` (the ``t = 1`` endpoint).
+            noise: ``eps ~ N(0, I)`` (the ``t = 0`` endpoint).
+            timesteps: Bridge times in ``[0, 1]``, broadcastable over the batch.
+        """
+        t = timesteps.to(original_samples.device, original_samples.dtype)
+        t = t.view(-1, *([1] * (original_samples.ndim - 1)))
+        return (1.0 - t) * noise + t * original_samples
+
+    def get_velocity(
+        self,
+        sample: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        """The training target ``u* = z_e - eps`` (bridge.py:328 at ``sigma_b = 0``).
+
+        Constant along the path, hence independent of ``timesteps``; the argument
+        is kept for `diffusers` API compatibility.
+        """
+        del timesteps
+        return sample - noise

sar2opt/transformer/config.json

@@ -0,0 +1,18 @@
+{
+  "_class_name": "ReFlowSETTransformer2DModel",
+  "_diffusers_version": "0.37.1",
+  "axes_dim": [
+    32,
+    32
+  ],
+  "depth": 24,
+  "double_blocks": 8,
+  "double_merge": "token",
+  "hidden_size": 1024,
+  "in_channels": 128,
+  "mlp_ratio": 4.0,
+  "num_heads": 16,
+  "out_channels": 128,
+  "sample_size": 512,
+  "theta": 10000
+}

sar2opt/transformer/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:14d5a1aaf41a32d2077c33f9aa19d356ac07611bb4d9d1f7c4ea61b9f6630481
+size 2037319452

sar2opt/transformer/transformer_reflowset.py

@@ -0,0 +1,471 @@
+"""ReFlowSET velocity transformer — a latent DiT with an EO/SAR double stream.
+
+The network predicts the flow-bridge velocity ``dz/dt`` in the frozen FLUX.2
+latent space.  It takes the noisy EO latent ``[B, 128, h, w]``, a scalar bridge
+time ``t`` in ``[0, 1]``, and the SAR conditioning latent of the same shape; the
+first 8 of its 24 blocks are double-stream (one EO tower and one SAR tower over
+a single joint attention), the remaining 16 are single-stream over the
+concatenated ``[EO | SAR]`` sequence, and only the EO half is decoded.
+
+This is an inference-only port.  The training-only REPA projection head
+(``repa_proj``) is a separate module in the reference implementation and is
+deliberately absent here.
+
+Deviations from `diffusers`' FLUX blocks that this file has to keep — each one
+is silent if you get it wrong:
+
+* ``FinalLayer`` unpacks ``shift, scale`` (dit.py:360), the **opposite** order of
+  ``AdaLayerNormContinuous``.
+* The single-stream MLP is **SwiGLU** of width 2752, not a 4x GELU of width 4096.
+* ``linear1``/``linear2`` are **bias-free**, and the QK-norm parameter is called
+  ``scale``, not ``weight``.
+* The timestep is multiplied by 1000 *inside* the model and the sinusoid is
+  **cos first, then sin**.
+* RoPE runs on **two** axes of **centred half-integer** coordinates, not on
+  FLUX's three axes of integers starting at 0.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional, Union
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+#: Width of the sinusoidal timestep embedding fed to ``time_in`` (dit.py:44).
+#: A module constant, deliberately independent of ``hidden_size``.
+TIME_EMBED_DIM = 256
+
+
+def swiglu_hidden_dim(hidden_size: int, mlp_ratio: float) -> int:
+    """SwiGLU intermediate width (dit.py:120-127).
+
+    The canonical 2/3 rule rounded to a multiple of 64, so a gated MLP at
+    ``mlp_ratio=4.0`` costs the same parameters as a plain 4x GELU MLP.
+    ``hidden_size=1024, mlp_ratio=4.0 -> 2752``.
+    """
+    return int(round(hidden_size * mlp_ratio * 2 / 3 / 64)) * 64
+
+
+class SwiGLU(nn.Module):
+    """``silu(first half) * second half`` — gate first, value second (dit.py:130-133)."""
+
+    def forward(self, x: Tensor) -> Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        return F.silu(x1) * x2
+
+
+class RMSNorm(nn.Module):
+    """RMS norm computed in float32.  The parameter is named ``scale`` (dit.py:136-145)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        x_dtype = x.dtype
+        x = x.float()
+        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
+        return (x * rrms).to(dtype=x_dtype) * self.scale
+
+
+class QKNorm(nn.Module):
+    """Per-head query/key RMS norm, applied **before** RoPE (dit.py:148-155)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.query_norm = RMSNorm(dim)
+        self.key_norm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        return self.query_norm(q).to(v), self.key_norm(k).to(v)
+
+
+class MLPEmbedder(nn.Module):
+    """``Linear -> SiLU -> Linear`` time-embedding MLP (dit.py:158-166).
+
+    Checkpoint keys are ``time_in.in_layer.*`` / ``time_in.out_layer.*``, not
+    diffusers' ``time_text_embed.timestep_embedder.linear_{1,2}``.
+    """
+
+    def __init__(self, in_dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+        self.silu = nn.SiLU()
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class Modulation(nn.Module):
+    """AdaLN-Zero triple.  Order is ``shift, scale, gate`` (dit.py:169-181)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.lin = nn.Linear(dim, 3 * dim, bias=True)
+
+    def forward(self, vec: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        out = self.lin(F.silu(vec))
+        if out.ndim == 2:
+            out = out[:, None, :]
+        shift, scale, gate = out.chunk(3, dim=-1)
+        return shift, scale, gate
+
+
+def timestep_embedding(
+    t: Tensor, dim: int, max_period: int = 10000, time_factor: float = 1000.0
+) -> Tensor:
+    """Sinusoidal embedding of a fractional bridge time (dit.py:184-201).
+
+    Two things differ from `diffusers`' ``get_timestep_embedding`` defaults:
+    ``t`` is a fraction in ``[0, 1]`` that is scaled by ``time_factor = 1000``
+    **here**, and the concatenation order is ``[cos, sin]`` (FLUX's ordering,
+    i.e. ``flip_sin_to_cos=True``).
+    """
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period)
+        * torch.arange(start=0, end=half, device=t.device, dtype=torch.float32)
+        / half
+    )
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    """Per-axis rotation matrices ``[..., L, dim/2, 2, 2]`` (dit.py:204-211)."""
+    if dim % 2:
+        raise ValueError(f"RoPE axis dim must be even, got {dim}")
+    scale = torch.arange(0, dim, 2, dtype=pos.dtype, device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
+    return out.reshape(*out.shape[:-1], 2, 2).float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+    """Rotate consecutive dimension pairs — the interleaved (FLUX) convention.
+
+    ``(x0, x1) -> (cos*x0 - sin*x1, sin*x0 + cos*x1)`` on ``(x[2k], x[2k+1])``
+    (dit.py:214-219).  Equivalent to diffusers' ``apply_rotary_emb(...,
+    use_real_unbind_dim=-1)``; ``-2`` is the split-halves convention and is wrong
+    for these weights.
+    """
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+
+
+class EmbedND(nn.Module):
+    """Concatenates the per-axis RoPE ladders and inserts the head axis (dit.py:222-234).
+
+    Holds no parameters and no buffers: the grid is rebuilt on every forward,
+    which is what lets one checkpoint serve 256 and 512 inputs.
+    """
+
+    def __init__(self, theta: int, axes_dim: list[int]) -> None:
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        emb = torch.cat(
+            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(len(self.axes_dim))],
+            dim=-3,
+        )
+        return emb.unsqueeze(1)
+
+
+def latent_image_ids(h: int, w: int, device, dtype=torch.float32) -> Tensor:
+    """Centred ``(y, x)`` coordinates for an ``h x w`` latent grid, ``[h*w, 2]``.
+
+    ``arange(n) - (n - 1) / 2`` with unit spacing (dit.py:237-252), so for even
+    ``n`` the coordinates are half-integers and the central 16x16 region of a
+    32x32 grid carries exactly the coordinates a 256-trained model saw — RoPE
+    only extrapolates outwards, it never rescales.  Row-major, so token
+    ``p = y * w + x``.  This is **not** FLUX's 3-axis integer id grid.
+    """
+    y = torch.arange(h, device=device, dtype=dtype) - (h - 1) / 2
+    x = torch.arange(w, device=device, dtype=dtype) - (w - 1) / 2
+    ids = torch.zeros(h, w, 2, device=device, dtype=dtype)
+    ids[..., 0] = y[:, None]
+    ids[..., 1] = x[None, :]
+    return ids.reshape(h * w, 2)
+
+
+class SingleStreamBlock(nn.Module):
+    """Fused attention + SwiGLU MLP under one modulation and one residual.
+
+    ``linear1`` emits ``[q | k | v | mlp_gate | mlp_value]`` in that order; the
+    qkv slab is K-major (``(K H D)``).  Both linears are bias-free
+    (dit.py:255-300).
+    """
+
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.mlp_hidden_dim = swiglu_hidden_dim(hidden_size, mlp_ratio)
+
+        self.linear1 = nn.Linear(hidden_size, 3 * hidden_size + 2 * self.mlp_hidden_dim, bias=False)
+        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size, bias=False)
+        self.norm = QKNorm(head_dim)
+        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp_act = SwiGLU()
+        self.modulation = Modulation(hidden_size)
+
+    def pre_attention(self, x: Tensor, vec: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+        """Everything up to (not including) RoPE and attention (dit.py:272-288)."""
+        shift, scale, gate = self.modulation(vec)
+        x_mod = (1 + scale) * self.pre_norm(x) + shift
+
+        qkv, mlp = torch.split(
+            self.linear1(x_mod), [3 * self.hidden_size, 2 * self.mlp_hidden_dim], dim=-1
+        )
+        b, length, _ = qkv.shape
+        # "B L (K H D) -> K B H L D" with K=3, H=num_heads.
+        q, k, v = qkv.reshape(b, length, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        q, k = self.norm(q, k, v)
+        return q, k, v, mlp, gate
+
+    def post_attention(self, x: Tensor, attn: Tensor, mlp: Tensor, gate: Tensor) -> Tensor:
+        """Output projection and the single gated residual (dit.py:290-294)."""
+        b, heads, length, head_dim = attn.shape
+        attn = attn.transpose(1, 2).reshape(b, length, heads * head_dim)
+        out = self.linear2(torch.cat((attn, self.mlp_act(mlp)), dim=-1))
+        return x + gate * out
+
+    def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
+        q, k, v, mlp, gate = self.pre_attention(x, vec)
+        q, k = apply_rope(q, k, pe)
+        attn = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+        return self.post_attention(x, attn, mlp, gate)
+
+
+class DoubleStreamBlock(nn.Module):
+    """Two independent towers over **one** joint attention across ``[EO | SAR]``.
+
+    The towers have completely separate weights but share the modulation vector
+    ``vec`` and the RoPE grid, so an EO token and the SAR token at the same
+    ground position carry an identical phase (dit.py:303-343).  SAR plays the
+    structural role text plays in FLUX.
+    """
+
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0) -> None:
+        super().__init__()
+        self.eo = SingleStreamBlock(hidden_size, num_heads, mlp_ratio)
+        self.sar = SingleStreamBlock(hidden_size, num_heads, mlp_ratio)
+
+    def forward(self, eo: Tensor, sar: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]:
+        """``pe`` must already cover the joint 2P-token sequence."""
+        q_e, k_e, v_e, mlp_e, gate_e = self.eo.pre_attention(eo, vec)
+        q_s, k_s, v_s, mlp_s, gate_s = self.sar.pre_attention(sar, vec)
+
+        q = torch.cat((q_e, q_s), dim=2)
+        k = torch.cat((k_e, k_s), dim=2)
+        v = torch.cat((v_e, v_s), dim=2)
+        q, k = apply_rope(q, k, pe)
+
+        attn = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+        attn_e, attn_s = attn.split([q_e.shape[2], q_s.shape[2]], dim=2)
+        return (
+            self.eo.post_attention(eo, attn_e, mlp_e, gate_e),
+            self.sar.post_attention(sar, attn_s, mlp_s, gate_s),
+        )
+
+
+class FinalLayer(nn.Module):
+    """AdaLN output layer.
+
+    ``adaLN`` unpacks ``shift, scale`` — the **opposite** order of diffusers'
+    ``AdaLayerNormContinuous`` (dit.py:346-362).  ``logvar_proj`` belongs to a
+    beta-NLL loss that was never enabled (``loss.flow = mse``); its weights are
+    kept so the published checkpoint loads with ``strict=True``, but inference
+    never evaluates it — the sampler reads only the velocity (bridge.py:531).
+    """
+
+    def __init__(self, hidden_size: int, out_channels: int) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.adaLN = nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+        self.proj = nn.Linear(hidden_size, out_channels, bias=True)
+        self.logvar_proj = nn.Linear(hidden_size, 1, bias=True)
+
+    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+        mod = self.adaLN(F.silu(vec))
+        if mod.ndim == 2:
+            mod = mod[:, None, :]
+        shift, scale = mod.chunk(2, dim=-1)
+        return self.proj((1 + scale) * self.norm(x) + shift)
+
+
+class ReFlowSETTransformer2DModel(ModelMixin, ConfigMixin):
+    """ReFlowSET's flow-velocity transformer (509.32 M parameters as configured).
+
+    Args:
+        in_channels: Channels of the packed FLUX.2 latent (128).
+        out_channels: Channels of the predicted velocity (128).
+        hidden_size: Residual width (1024).
+        depth: **Total** blocks, double plus single (24).
+        num_heads: Attention heads (16), so ``head_dim = 64``.
+        mlp_ratio: Nominal MLP ratio; the SwiGLU width is derived from it.
+        axes_dim: RoPE dims for the ``(y, x)`` axes; must sum to ``head_dim``.
+        theta: RoPE base period (10000).
+        sample_size: Input image resolution the released arm was trained at
+            (256 for QXS-SAROPT, 512 for SAR2Opt).  Recorded for provenance
+            only: the forward pass derives every shape from its input and the
+            RoPE grid is rebuilt per call, so one checkpoint serves any size
+            divisible by 16.
+        double_blocks: Leading double-stream blocks (8); the remaining
+            ``depth - double_blocks`` are single-stream.
+        double_merge: How the two streams become one.  ``"token"`` (the released
+            setting) concatenates on the sequence axis, so the single stack runs
+            over 2P tokens and the SAR half is dropped only at the very end;
+            ``"channel"`` fuses per position and keeps P tokens.
+
+    Forward contract: ``forward(hidden_states, timestep, condition)`` where
+    ``hidden_states`` is the bridge state ``[B, 128, h, w]``, ``timestep`` is the
+    bridge time in ``[0, 1]`` (**not** an integer diffusion step), and
+    ``condition`` is the SAR latent of the same shape or ``None``.  ``None`` is
+    the classifier-free-guidance null branch and is turned into an all-zero
+    latent inside the model — there is no learned null token.
+    """
+
+    _supports_gradient_checkpointing = False
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 128,
+        out_channels: int = 128,
+        hidden_size: int = 1024,
+        depth: int = 24,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        axes_dim: tuple[int, ...] = (32, 32),
+        theta: int = 10000,
+        sample_size: Optional[int] = None,
+        double_blocks: int = 8,
+        double_merge: str = "token",
+    ) -> None:
+        super().__init__()
+        if hidden_size % num_heads != 0:
+            raise ValueError(f"hidden_size {hidden_size} must be divisible by num_heads {num_heads}")
+        pe_dim = hidden_size // num_heads
+        if sum(axes_dim) != pe_dim:
+            raise ValueError(f"axes_dim {list(axes_dim)} must sum to the per-head dim {pe_dim}")
+        if not 0 <= double_blocks < depth:
+            raise ValueError(f"double_blocks {double_blocks} must be in [0, depth={depth})")
+        if double_merge not in ("token", "channel"):
+            raise ValueError(f"double_merge must be 'token' or 'channel', got {double_merge!r}")
+
+        self.pe_embedder = EmbedND(theta=theta, axes_dim=list(axes_dim))
+        if double_blocks:
+            # Each stream gets its own 1x1 "patchify": they are two token
+            # sequences now, not two halves of one channel stack.
+            self.in_proj_eo = nn.Linear(in_channels, hidden_size, bias=True)
+            self.in_proj_sar = nn.Linear(in_channels, hidden_size, bias=True)
+            if double_merge == "channel":
+                self.merge = nn.Linear(2 * hidden_size, hidden_size, bias=True)
+        else:
+            self.in_proj = nn.Linear(2 * in_channels, hidden_size, bias=True)
+        self.time_in = MLPEmbedder(TIME_EMBED_DIM, hidden_size)
+        self.double_stream = nn.ModuleList(
+            [DoubleStreamBlock(hidden_size, num_heads, mlp_ratio) for _ in range(double_blocks)]
+        )
+        self.blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(hidden_size, num_heads, mlp_ratio)
+                for _ in range(depth - double_blocks)
+            ]
+        )
+        self.final_layer = FinalLayer(hidden_size, out_channels)
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        timestep: Tensor,
+        condition: Optional[Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[Transformer2DModelOutput, tuple[Tensor]]:
+        """Predict the flow velocity ``dz/dt``.
+
+        Args:
+            hidden_states: ``[B, in_channels, h, w]`` bridge state.
+            timestep: Bridge time in ``[0, 1]``; a scalar or ``[B]``.
+            condition: ``[B, in_channels, h, w]`` SAR latent, or ``None`` for the
+                null branch (an all-zero conditioning latent, dit.py:531-532).
+            return_dict: Return a ``Transformer2DModelOutput`` instead of a tuple.
+
+        Returns:
+            The velocity ``[B, out_channels, h, w]``.  This is a flow velocity,
+            not ``epsilon`` and not diffusers' ``v_prediction``.
+        """
+        if hidden_states.ndim != 4:
+            raise ValueError(f"hidden_states must be [B, C, h, w], got {tuple(hidden_states.shape)}")
+        batch, _, h, w = hidden_states.shape
+        if condition is None:
+            condition = torch.zeros_like(hidden_states)
+        elif condition.shape != hidden_states.shape:
+            raise ValueError(
+                f"condition shape {tuple(condition.shape)} must match "
+                f"hidden_states shape {tuple(hidden_states.shape)}"
+            )
+        if timestep.ndim == 0:
+            timestep = timestep.expand(batch)
+
+        n_double = self.config.double_blocks
+        if n_double:
+            eo = self.in_proj_eo(hidden_states.flatten(2).transpose(1, 2))  # [B, P, D]
+            sar = self.in_proj_sar(condition.flatten(2).transpose(1, 2))  # [B, P, D]
+            ref = eo
+        else:
+            x = torch.cat([hidden_states, condition], dim=1).flatten(2).transpose(1, 2)
+            x = self.in_proj(x)
+            ref = x
+
+        vec = self.time_in(timestep_embedding(timestep, TIME_EMBED_DIM).to(ref.dtype))
+
+        ids = latent_image_ids(h, w, device=hidden_states.device, dtype=torch.float32)
+        pe = self.pe_embedder(ids[None].expand(batch, -1, -1))
+
+        num_tokens = ref.shape[1]
+        pe_single = pe
+        if n_double:
+            # Token axis of pe is dim 2 ([B, 1, L, head_dim/2, 2, 2]); repeating
+            # the same P coordinates gives EO and SAR one shared grid.
+            pe_joint = torch.cat((pe, pe), dim=2)
+            for block in self.double_stream:
+                eo, sar = block(eo, sar, vec, pe_joint)
+            if self.config.double_merge == "token":
+                x = torch.cat((eo, sar), dim=1)  # [B, 2P, D]
+                pe_single = pe_joint
+            else:
+                x = self.merge(torch.cat((eo, sar), dim=-1))  # [B, P, D]
+
+        for block in self.blocks:
+            x = block(x, vec, pe_single)
+
+        if n_double and self.config.double_merge == "token":
+            x = x[:, :num_tokens]  # drop the SAR half: only EO is decoded
+
+        v = self.final_layer(x, vec)
+        v = v.transpose(1, 2).reshape(batch, self.config.out_channels, h, w)
+        if not return_dict:
+            return (v,)
+        return Transformer2DModelOutput(sample=v)

sar2opt/transformer_reflowset.py

@@ -0,0 +1,471 @@
+"""ReFlowSET velocity transformer — a latent DiT with an EO/SAR double stream.
+
+The network predicts the flow-bridge velocity ``dz/dt`` in the frozen FLUX.2
+latent space.  It takes the noisy EO latent ``[B, 128, h, w]``, a scalar bridge
+time ``t`` in ``[0, 1]``, and the SAR conditioning latent of the same shape; the
+first 8 of its 24 blocks are double-stream (one EO tower and one SAR tower over
+a single joint attention), the remaining 16 are single-stream over the
+concatenated ``[EO | SAR]`` sequence, and only the EO half is decoded.
+
+This is an inference-only port.  The training-only REPA projection head
+(``repa_proj``) is a separate module in the reference implementation and is
+deliberately absent here.
+
+Deviations from `diffusers`' FLUX blocks that this file has to keep — each one
+is silent if you get it wrong:
+
+* ``FinalLayer`` unpacks ``shift, scale`` (dit.py:360), the **opposite** order of
+  ``AdaLayerNormContinuous``.
+* The single-stream MLP is **SwiGLU** of width 2752, not a 4x GELU of width 4096.
+* ``linear1``/``linear2`` are **bias-free**, and the QK-norm parameter is called
+  ``scale``, not ``weight``.
+* The timestep is multiplied by 1000 *inside* the model and the sinusoid is
+  **cos first, then sin**.
+* RoPE runs on **two** axes of **centred half-integer** coordinates, not on
+  FLUX's three axes of integers starting at 0.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional, Union
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+#: Width of the sinusoidal timestep embedding fed to ``time_in`` (dit.py:44).
+#: A module constant, deliberately independent of ``hidden_size``.
+TIME_EMBED_DIM = 256
+
+
+def swiglu_hidden_dim(hidden_size: int, mlp_ratio: float) -> int:
+    """SwiGLU intermediate width (dit.py:120-127).
+
+    The canonical 2/3 rule rounded to a multiple of 64, so a gated MLP at
+    ``mlp_ratio=4.0`` costs the same parameters as a plain 4x GELU MLP.
+    ``hidden_size=1024, mlp_ratio=4.0 -> 2752``.
+    """
+    return int(round(hidden_size * mlp_ratio * 2 / 3 / 64)) * 64
+
+
+class SwiGLU(nn.Module):
+    """``silu(first half) * second half`` — gate first, value second (dit.py:130-133)."""
+
+    def forward(self, x: Tensor) -> Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        return F.silu(x1) * x2
+
+
+class RMSNorm(nn.Module):
+    """RMS norm computed in float32.  The parameter is named ``scale`` (dit.py:136-145)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        x_dtype = x.dtype
+        x = x.float()
+        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
+        return (x * rrms).to(dtype=x_dtype) * self.scale
+
+
+class QKNorm(nn.Module):
+    """Per-head query/key RMS norm, applied **before** RoPE (dit.py:148-155)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.query_norm = RMSNorm(dim)
+        self.key_norm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        return self.query_norm(q).to(v), self.key_norm(k).to(v)
+
+
+class MLPEmbedder(nn.Module):
+    """``Linear -> SiLU -> Linear`` time-embedding MLP (dit.py:158-166).
+
+    Checkpoint keys are ``time_in.in_layer.*`` / ``time_in.out_layer.*``, not
+    diffusers' ``time_text_embed.timestep_embedder.linear_{1,2}``.
+    """
+
+    def __init__(self, in_dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+        self.silu = nn.SiLU()
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class Modulation(nn.Module):
+    """AdaLN-Zero triple.  Order is ``shift, scale, gate`` (dit.py:169-181)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.lin = nn.Linear(dim, 3 * dim, bias=True)
+
+    def forward(self, vec: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        out = self.lin(F.silu(vec))
+        if out.ndim == 2:
+            out = out[:, None, :]
+        shift, scale, gate = out.chunk(3, dim=-1)
+        return shift, scale, gate
+
+
+def timestep_embedding(
+    t: Tensor, dim: int, max_period: int = 10000, time_factor: float = 1000.0
+) -> Tensor:
+    """Sinusoidal embedding of a fractional bridge time (dit.py:184-201).
+
+    Two things differ from `diffusers`' ``get_timestep_embedding`` defaults:
+    ``t`` is a fraction in ``[0, 1]`` that is scaled by ``time_factor = 1000``
+    **here**, and the concatenation order is ``[cos, sin]`` (FLUX's ordering,
+    i.e. ``flip_sin_to_cos=True``).
+    """
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period)
+        * torch.arange(start=0, end=half, device=t.device, dtype=torch.float32)
+        / half
+    )
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    """Per-axis rotation matrices ``[..., L, dim/2, 2, 2]`` (dit.py:204-211)."""
+    if dim % 2:
+        raise ValueError(f"RoPE axis dim must be even, got {dim}")
+    scale = torch.arange(0, dim, 2, dtype=pos.dtype, device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
+    return out.reshape(*out.shape[:-1], 2, 2).float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+    """Rotate consecutive dimension pairs — the interleaved (FLUX) convention.
+
+    ``(x0, x1) -> (cos*x0 - sin*x1, sin*x0 + cos*x1)`` on ``(x[2k], x[2k+1])``
+    (dit.py:214-219).  Equivalent to diffusers' ``apply_rotary_emb(...,
+    use_real_unbind_dim=-1)``; ``-2`` is the split-halves convention and is wrong
+    for these weights.
+    """
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+
+
+class EmbedND(nn.Module):
+    """Concatenates the per-axis RoPE ladders and inserts the head axis (dit.py:222-234).
+
+    Holds no parameters and no buffers: the grid is rebuilt on every forward,
+    which is what lets one checkpoint serve 256 and 512 inputs.
+    """
+
+    def __init__(self, theta: int, axes_dim: list[int]) -> None:
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        emb = torch.cat(
+            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(len(self.axes_dim))],
+            dim=-3,
+        )
+        return emb.unsqueeze(1)
+
+
+def latent_image_ids(h: int, w: int, device, dtype=torch.float32) -> Tensor:
+    """Centred ``(y, x)`` coordinates for an ``h x w`` latent grid, ``[h*w, 2]``.
+
+    ``arange(n) - (n - 1) / 2`` with unit spacing (dit.py:237-252), so for even
+    ``n`` the coordinates are half-integers and the central 16x16 region of a
+    32x32 grid carries exactly the coordinates a 256-trained model saw — RoPE
+    only extrapolates outwards, it never rescales.  Row-major, so token
+    ``p = y * w + x``.  This is **not** FLUX's 3-axis integer id grid.
+    """
+    y = torch.arange(h, device=device, dtype=dtype) - (h - 1) / 2
+    x = torch.arange(w, device=device, dtype=dtype) - (w - 1) / 2
+    ids = torch.zeros(h, w, 2, device=device, dtype=dtype)
+    ids[..., 0] = y[:, None]
+    ids[..., 1] = x[None, :]
+    return ids.reshape(h * w, 2)
+
+
+class SingleStreamBlock(nn.Module):
+    """Fused attention + SwiGLU MLP under one modulation and one residual.
+
+    ``linear1`` emits ``[q | k | v | mlp_gate | mlp_value]`` in that order; the
+    qkv slab is K-major (``(K H D)``).  Both linears are bias-free
+    (dit.py:255-300).
+    """
+
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.mlp_hidden_dim = swiglu_hidden_dim(hidden_size, mlp_ratio)
+
+        self.linear1 = nn.Linear(hidden_size, 3 * hidden_size + 2 * self.mlp_hidden_dim, bias=False)
+        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size, bias=False)
+        self.norm = QKNorm(head_dim)
+        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp_act = SwiGLU()
+        self.modulation = Modulation(hidden_size)
+
+    def pre_attention(self, x: Tensor, vec: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+        """Everything up to (not including) RoPE and attention (dit.py:272-288)."""
+        shift, scale, gate = self.modulation(vec)
+        x_mod = (1 + scale) * self.pre_norm(x) + shift
+
+        qkv, mlp = torch.split(
+            self.linear1(x_mod), [3 * self.hidden_size, 2 * self.mlp_hidden_dim], dim=-1
+        )
+        b, length, _ = qkv.shape
+        # "B L (K H D) -> K B H L D" with K=3, H=num_heads.
+        q, k, v = qkv.reshape(b, length, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        q, k = self.norm(q, k, v)
+        return q, k, v, mlp, gate
+
+    def post_attention(self, x: Tensor, attn: Tensor, mlp: Tensor, gate: Tensor) -> Tensor:
+        """Output projection and the single gated residual (dit.py:290-294)."""
+        b, heads, length, head_dim = attn.shape
+        attn = attn.transpose(1, 2).reshape(b, length, heads * head_dim)
+        out = self.linear2(torch.cat((attn, self.mlp_act(mlp)), dim=-1))
+        return x + gate * out
+
+    def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
+        q, k, v, mlp, gate = self.pre_attention(x, vec)
+        q, k = apply_rope(q, k, pe)
+        attn = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+        return self.post_attention(x, attn, mlp, gate)
+
+
+class DoubleStreamBlock(nn.Module):
+    """Two independent towers over **one** joint attention across ``[EO | SAR]``.
+
+    The towers have completely separate weights but share the modulation vector
+    ``vec`` and the RoPE grid, so an EO token and the SAR token at the same
+    ground position carry an identical phase (dit.py:303-343).  SAR plays the
+    structural role text plays in FLUX.
+    """
+
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0) -> None:
+        super().__init__()
+        self.eo = SingleStreamBlock(hidden_size, num_heads, mlp_ratio)
+        self.sar = SingleStreamBlock(hidden_size, num_heads, mlp_ratio)
+
+    def forward(self, eo: Tensor, sar: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]:
+        """``pe`` must already cover the joint 2P-token sequence."""
+        q_e, k_e, v_e, mlp_e, gate_e = self.eo.pre_attention(eo, vec)
+        q_s, k_s, v_s, mlp_s, gate_s = self.sar.pre_attention(sar, vec)
+
+        q = torch.cat((q_e, q_s), dim=2)
+        k = torch.cat((k_e, k_s), dim=2)
+        v = torch.cat((v_e, v_s), dim=2)
+        q, k = apply_rope(q, k, pe)
+
+        attn = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+        attn_e, attn_s = attn.split([q_e.shape[2], q_s.shape[2]], dim=2)
+        return (
+            self.eo.post_attention(eo, attn_e, mlp_e, gate_e),
+            self.sar.post_attention(sar, attn_s, mlp_s, gate_s),
+        )
+
+
+class FinalLayer(nn.Module):
+    """AdaLN output layer.
+
+    ``adaLN`` unpacks ``shift, scale`` — the **opposite** order of diffusers'
+    ``AdaLayerNormContinuous`` (dit.py:346-362).  ``logvar_proj`` belongs to a
+    beta-NLL loss that was never enabled (``loss.flow = mse``); its weights are
+    kept so the published checkpoint loads with ``strict=True``, but inference
+    never evaluates it — the sampler reads only the velocity (bridge.py:531).
+    """
+
+    def __init__(self, hidden_size: int, out_channels: int) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.adaLN = nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+        self.proj = nn.Linear(hidden_size, out_channels, bias=True)
+        self.logvar_proj = nn.Linear(hidden_size, 1, bias=True)
+
+    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+        mod = self.adaLN(F.silu(vec))
+        if mod.ndim == 2:
+            mod = mod[:, None, :]
+        shift, scale = mod.chunk(2, dim=-1)
+        return self.proj((1 + scale) * self.norm(x) + shift)
+
+
+class ReFlowSETTransformer2DModel(ModelMixin, ConfigMixin):
+    """ReFlowSET's flow-velocity transformer (509.32 M parameters as configured).
+
+    Args:
+        in_channels: Channels of the packed FLUX.2 latent (128).
+        out_channels: Channels of the predicted velocity (128).
+        hidden_size: Residual width (1024).
+        depth: **Total** blocks, double plus single (24).
+        num_heads: Attention heads (16), so ``head_dim = 64``.
+        mlp_ratio: Nominal MLP ratio; the SwiGLU width is derived from it.
+        axes_dim: RoPE dims for the ``(y, x)`` axes; must sum to ``head_dim``.
+        theta: RoPE base period (10000).
+        sample_size: Input image resolution the released arm was trained at
+            (256 for QXS-SAROPT, 512 for SAR2Opt).  Recorded for provenance
+            only: the forward pass derives every shape from its input and the
+            RoPE grid is rebuilt per call, so one checkpoint serves any size
+            divisible by 16.
+        double_blocks: Leading double-stream blocks (8); the remaining
+            ``depth - double_blocks`` are single-stream.
+        double_merge: How the two streams become one.  ``"token"`` (the released
+            setting) concatenates on the sequence axis, so the single stack runs
+            over 2P tokens and the SAR half is dropped only at the very end;
+            ``"channel"`` fuses per position and keeps P tokens.
+
+    Forward contract: ``forward(hidden_states, timestep, condition)`` where
+    ``hidden_states`` is the bridge state ``[B, 128, h, w]``, ``timestep`` is the
+    bridge time in ``[0, 1]`` (**not** an integer diffusion step), and
+    ``condition`` is the SAR latent of the same shape or ``None``.  ``None`` is
+    the classifier-free-guidance null branch and is turned into an all-zero
+    latent inside the model — there is no learned null token.
+    """
+
+    _supports_gradient_checkpointing = False
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 128,
+        out_channels: int = 128,
+        hidden_size: int = 1024,
+        depth: int = 24,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        axes_dim: tuple[int, ...] = (32, 32),
+        theta: int = 10000,
+        sample_size: Optional[int] = None,
+        double_blocks: int = 8,
+        double_merge: str = "token",
+    ) -> None:
+        super().__init__()
+        if hidden_size % num_heads != 0:
+            raise ValueError(f"hidden_size {hidden_size} must be divisible by num_heads {num_heads}")
+        pe_dim = hidden_size // num_heads
+        if sum(axes_dim) != pe_dim:
+            raise ValueError(f"axes_dim {list(axes_dim)} must sum to the per-head dim {pe_dim}")
+        if not 0 <= double_blocks < depth:
+            raise ValueError(f"double_blocks {double_blocks} must be in [0, depth={depth})")
+        if double_merge not in ("token", "channel"):
+            raise ValueError(f"double_merge must be 'token' or 'channel', got {double_merge!r}")
+
+        self.pe_embedder = EmbedND(theta=theta, axes_dim=list(axes_dim))
+        if double_blocks:
+            # Each stream gets its own 1x1 "patchify": they are two token
+            # sequences now, not two halves of one channel stack.
+            self.in_proj_eo = nn.Linear(in_channels, hidden_size, bias=True)
+            self.in_proj_sar = nn.Linear(in_channels, hidden_size, bias=True)
+            if double_merge == "channel":
+                self.merge = nn.Linear(2 * hidden_size, hidden_size, bias=True)
+        else:
+            self.in_proj = nn.Linear(2 * in_channels, hidden_size, bias=True)
+        self.time_in = MLPEmbedder(TIME_EMBED_DIM, hidden_size)
+        self.double_stream = nn.ModuleList(
+            [DoubleStreamBlock(hidden_size, num_heads, mlp_ratio) for _ in range(double_blocks)]
+        )
+        self.blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(hidden_size, num_heads, mlp_ratio)
+                for _ in range(depth - double_blocks)
+            ]
+        )
+        self.final_layer = FinalLayer(hidden_size, out_channels)
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        timestep: Tensor,
+        condition: Optional[Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[Transformer2DModelOutput, tuple[Tensor]]:
+        """Predict the flow velocity ``dz/dt``.
+
+        Args:
+            hidden_states: ``[B, in_channels, h, w]`` bridge state.
+            timestep: Bridge time in ``[0, 1]``; a scalar or ``[B]``.
+            condition: ``[B, in_channels, h, w]`` SAR latent, or ``None`` for the
+                null branch (an all-zero conditioning latent, dit.py:531-532).
+            return_dict: Return a ``Transformer2DModelOutput`` instead of a tuple.
+
+        Returns:
+            The velocity ``[B, out_channels, h, w]``.  This is a flow velocity,
+            not ``epsilon`` and not diffusers' ``v_prediction``.
+        """
+        if hidden_states.ndim != 4:
+            raise ValueError(f"hidden_states must be [B, C, h, w], got {tuple(hidden_states.shape)}")
+        batch, _, h, w = hidden_states.shape
+        if condition is None:
+            condition = torch.zeros_like(hidden_states)
+        elif condition.shape != hidden_states.shape:
+            raise ValueError(
+                f"condition shape {tuple(condition.shape)} must match "
+                f"hidden_states shape {tuple(hidden_states.shape)}"
+            )
+        if timestep.ndim == 0:
+            timestep = timestep.expand(batch)
+
+        n_double = self.config.double_blocks
+        if n_double:
+            eo = self.in_proj_eo(hidden_states.flatten(2).transpose(1, 2))  # [B, P, D]
+            sar = self.in_proj_sar(condition.flatten(2).transpose(1, 2))  # [B, P, D]
+            ref = eo
+        else:
+            x = torch.cat([hidden_states, condition], dim=1).flatten(2).transpose(1, 2)
+            x = self.in_proj(x)
+            ref = x
+
+        vec = self.time_in(timestep_embedding(timestep, TIME_EMBED_DIM).to(ref.dtype))
+
+        ids = latent_image_ids(h, w, device=hidden_states.device, dtype=torch.float32)
+        pe = self.pe_embedder(ids[None].expand(batch, -1, -1))
+
+        num_tokens = ref.shape[1]
+        pe_single = pe
+        if n_double:
+            # Token axis of pe is dim 2 ([B, 1, L, head_dim/2, 2, 2]); repeating
+            # the same P coordinates gives EO and SAR one shared grid.
+            pe_joint = torch.cat((pe, pe), dim=2)
+            for block in self.double_stream:
+                eo, sar = block(eo, sar, vec, pe_joint)
+            if self.config.double_merge == "token":
+                x = torch.cat((eo, sar), dim=1)  # [B, 2P, D]
+                pe_single = pe_joint
+            else:
+                x = self.merge(torch.cat((eo, sar), dim=-1))  # [B, P, D]
+
+        for block in self.blocks:
+            x = block(x, vec, pe_single)
+
+        if n_double and self.config.double_merge == "token":
+            x = x[:, :num_tokens]  # drop the SAR half: only EO is decoded
+
+        v = self.final_layer(x, vec)
+        v = v.transpose(1, 2).reshape(batch, self.config.out_channels, h, w)
+        if not return_dict:
+            return (v,)
+        return Transformer2DModelOutput(sample=v)

sar2opt/vae/autoencoder_flux2.py

@@ -0,0 +1,426 @@
+"""Frozen FLUX.2 autoencoder — the latent endpoint of ReFlowSET.
+
+ReFlowSET never fine-tunes this module: it is loaded once, frozen, and used to
+encode the SAR condition and to decode the sampled EO latent.  The released
+weights are the **Apache-2.0** FLUX.2-klein-base-4B copy of the autoencoder,
+re-keyed to the layout below (see ``scripts/convert_flux2_ae.py``).
+
+Three details of the checkpoint are non-standard for `diffusers` and are
+preserved exactly, because the file must load with ``strict=True``:
+
+* ``quant_conv`` lives **inside** ``encoder.*`` and is the last op of the
+  encoder forward; ``post_quant_conv`` lives **inside** ``decoder.*`` and is the
+  first op of the decoder forward.  `diffusers`' ``AutoencoderKL`` makes both
+  siblings of the encoder/decoder.
+* The latent normaliser is a real ``BatchNorm2d(128, affine=False)`` whose
+  running statistics ship in the checkpoint under ``bn.*`` — a per-channel mean
+  **and** variance, not a scalar ``scaling_factor``/``shift_factor``.  Its
+  epsilon is ``1e-4``, not torch's ``1e-5``.
+* ``encode`` returns the posterior **mean**; the log-variance chunk of the
+  encoder's moments is discarded, so encoding is deterministic and there is no
+  ``DiagonalGaussianDistribution`` and no ``.sample()``.
+
+The public latent is ``[B, 128, H/16, W/16]``: an 8x convolutional stride
+followed by a 2x2 space-to-depth pack that is part of the *autoencoder*, not of
+the transformer.
+"""
+
+from __future__ import annotations
+
+import os
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+
+def swish(x: Tensor) -> Tensor:
+    """``x * sigmoid(x)`` — the activation used throughout the FLUX.2 AE."""
+    return x * torch.sigmoid(x)
+
+
+class AttnBlock(nn.Module):
+    """Single-head self-attention over the spatial grid (head dim == channels)."""
+
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q, k, v = self.q(h_), self.k(h_), self.v(h_)
+        b, c, h, w = q.shape
+        # "b c h w -> b 1 (h w) c": ONE head whose head-dim is the full channel
+        # count (flux2_ae.py:70-73).
+        q = q.reshape(b, c, h * w).transpose(1, 2).unsqueeze(1).contiguous()
+        k = k.reshape(b, c, h * w).transpose(1, 2).unsqueeze(1).contiguous()
+        v = v.reshape(b, c, h * w).transpose(1, 2).unsqueeze(1).contiguous()
+        h_ = F.scaled_dot_product_attention(q, k, v)
+        return h_.squeeze(1).transpose(1, 2).reshape(b, c, h, w)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if in_channels != out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x: Tensor) -> Tensor:
+        h = self.conv1(swish(self.norm1(x)))
+        h = self.conv2(swish(self.norm2(h)))
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+        return x + h
+
+
+class Downsample(nn.Module):
+    """Stride-2 conv with FLUX's asymmetric ``(0, 1, 0, 1)`` pad (flux2_ae.py:111-121)."""
+
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.conv(F.pad(x, (0, 1, 0, 1), mode="constant", value=0))
+
+
+class Upsample(nn.Module):
+    """Nearest-neighbour 2x followed by a 3x3 conv (flux2_ae.py:124-132)."""
+
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.conv(F.interpolate(x, scale_factor=2.0, mode="nearest"))
+
+
+class Encoder(nn.Module):
+    """FLUX.2 encoder.  Emits ``2 * z_channels`` moments; ``quant_conv`` is internal."""
+
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        z_channels: int,
+    ) -> None:
+        super().__init__()
+        # Declared first so the checkpoint key is `encoder.quant_conv.*`
+        # (flux2_ae.py:146) — diffusers keeps quant_conv outside the encoder.
+        self.quant_conv = nn.Conv2d(2 * z_channels, 2 * z_channels, 1)
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        self.conv_in = nn.Conv2d(in_channels, ch, kernel_size=3, stride=1, padding=1)
+
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.down = nn.ModuleList()
+        block_in = ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(num_res_blocks):
+                block.append(ResnetBlock(block_in, block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            # Empty at every level in this checkpoint: attention exists only in
+            # `mid` (flux2_ae.py:162).  Kept so the forward guard is meaningful.
+            down.attn = nn.ModuleList()
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in)
+            self.down.append(down)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(block_in, block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(block_in, block_in)
+
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        h = self.mid.block_2(self.mid.attn_1(self.mid.block_1(hs[-1])))
+        h = self.conv_out(swish(self.norm_out(h)))
+        return self.quant_conv(h)  # last op of the encoder (flux2_ae.py:207)
+
+
+class Decoder(nn.Module):
+    """FLUX.2 decoder.  ``post_quant_conv`` is internal and runs first."""
+
+    def __init__(
+        self,
+        ch: int,
+        out_ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+    ) -> None:
+        super().__init__()
+        # Checkpoint key `decoder.post_quant_conv.*` (flux2_ae.py:223).
+        self.post_quant_conv = nn.Conv2d(z_channels, z_channels, 1)
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(block_in, block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(block_in, block_in)
+
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(num_res_blocks + 1):
+                block.append(ResnetBlock(block_in, block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = nn.ModuleList()  # empty in this checkpoint (flux2_ae.py:249)
+            if i_level != 0:
+                up.upsample = Upsample(block_in)
+            self.up.insert(0, up)  # prepend so `up.<i>` indexes by resolution level
+
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: Tensor) -> Tensor:
+        z = self.post_quant_conv(z)  # first op of the decoder (flux2_ae.py:267)
+        upscale_dtype = next(self.up.parameters()).dtype
+
+        h = self.conv_in(z)
+        h = self.mid.block_2(self.mid.attn_1(self.mid.block_1(h)))
+        h = h.to(upscale_dtype)
+
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        return self.conv_out(swish(self.norm_out(h)))
+
+
+class AutoencoderFlux2(ModelMixin, ConfigMixin):
+    """Frozen FLUX.2 autoencoder with ReFlowSET's packed, BN-normalised latent.
+
+    ``encode`` maps ``[B, 3, H, W]`` in ``[-1, 1]`` to ``[B, 128, H/16, W/16]``
+    and ``decode`` inverts it.  The module is frozen: ``train()`` is a no-op that
+    always selects eval mode, and the latent BatchNorm is additionally forced to
+    eval on every call so no batch statistic can ever leak into the latent.
+
+    Args:
+        resolution: Nominal training resolution of the original autoencoder.
+            Only used to size bookkeeping attributes; any ``H``, ``W`` divisible
+            by 16 may be encoded.
+        in_channels: Input image channels (3).
+        ch: Base width.
+        out_ch: Output image channels (3).
+        ch_mult: Per-level width multipliers; ``len(ch_mult) - 1`` downsamples.
+        num_res_blocks: Residual blocks per level.
+        z_channels: Pre-pack latent channels (32).
+        patch_size: Space-to-depth factor applied after the encoder (2), which
+            takes the latent from 32 channels at ``H/8`` to 128 at ``H/16``.
+        bn_eps: Epsilon of the latent BatchNorm.  **1e-4**, not torch's 1e-5
+            (flux2_ae.py:331); using 1e-5 shifts the latent by up to 2.6e-5.
+    """
+
+    _supports_gradient_checkpointing = False
+
+    @register_to_config
+    def __init__(
+        self,
+        resolution: int = 256,
+        in_channels: int = 3,
+        ch: int = 128,
+        out_ch: int = 3,
+        ch_mult: tuple[int, ...] = (1, 2, 4, 4),
+        num_res_blocks: int = 2,
+        z_channels: int = 32,
+        patch_size: int = 2,
+        bn_eps: float = 1e-4,
+    ) -> None:
+        super().__init__()
+        ch_mult = list(ch_mult)
+        self.encoder = Encoder(
+            resolution=resolution,
+            in_channels=in_channels,
+            ch=ch,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            z_channels=z_channels,
+        )
+        self.decoder = Decoder(
+            ch=ch,
+            out_ch=out_ch,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            in_channels=in_channels,
+            resolution=resolution,
+            z_channels=z_channels,
+        )
+        # Per-channel latent normaliser with the checkpoint's running statistics.
+        # affine=False, so there is no weight/bias to load (flux2_ae.py:334-340).
+        self.bn = nn.BatchNorm2d(
+            patch_size * patch_size * z_channels,
+            eps=bn_eps,
+            momentum=0.1,
+            affine=False,
+            track_running_stats=True,
+        )
+
+    @property
+    def latent_channels(self) -> int:
+        """Channels of the public latent: ``patch_size**2 * z_channels`` = 128."""
+        return self.config.patch_size**2 * self.config.z_channels
+
+    @property
+    def spatial_factor(self) -> int:
+        """Total stride: 8x convolutional times ``patch_size`` packing = 16."""
+        return 2 ** (len(self.config.ch_mult) - 1) * self.config.patch_size
+
+    # ---- 2x2 space-to-depth pack / unpack -----------------------------------
+
+    def pack(self, z: Tensor) -> Tensor:
+        """``[B, C, H, W] -> [B, C*p*p, H/p, W/p]``, channel-major.
+
+        Bit-identical to the reference ``rearrange("... c (i pi) (j pj) -> ...
+        (c pi pj) i j")`` (flux2_ae.py:349-357).  Note this is **not** diffusers'
+        ``_pack_latents``, whose channel grouping is transposed.
+        """
+        return F.pixel_unshuffle(z, self.config.patch_size)
+
+    def unpack(self, z: Tensor) -> Tensor:
+        """Exact inverse of :meth:`pack` (flux2_ae.py:359-367)."""
+        return F.pixel_shuffle(z, self.config.patch_size)
+
+    # ---- latent normalisation ----------------------------------------------
+
+    def normalize(self, z: Tensor) -> Tensor:
+        """``(z - running_mean) / sqrt(running_var + bn_eps)``, per channel."""
+        self.bn.eval()  # forced every call (flux2_ae.py:372); train mode shifts z by ~1.67
+        return self.bn(z)
+
+    def inv_normalize(self, z: Tensor) -> Tensor:
+        """Exact inverse of :meth:`normalize` — same ``bn_eps`` (flux2_ae.py:375-379)."""
+        self.bn.eval()
+        s = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.config.bn_eps)
+        m = self.bn.running_mean.view(1, -1, 1, 1)
+        return z * s + m
+
+    # ---- public API ---------------------------------------------------------
+
+    @torch.no_grad()
+    def encode(self, x: Tensor) -> Tensor:
+        """Encode an image to the packed, normalised latent.
+
+        Args:
+            x: ``[B, 3, H, W]`` in ``[-1, 1]``; ``H`` and ``W`` divisible by 16.
+
+        Returns:
+            ``[B, 128, H/16, W/16]`` — the posterior **mean**, packed and
+            BN-normalised.  The encoder's log-variance chunk is discarded
+            (flux2_ae.py:396), so this is deterministic: there is no posterior
+            distribution object and nothing to sample.
+        """
+        if x.ndim != 4 or x.shape[1] != self.config.in_channels:
+            raise ValueError(
+                f"encode expects [B, {self.config.in_channels}, H, W], got {tuple(x.shape)}"
+            )
+        h, w = x.shape[-2:]
+        if h % self.spatial_factor or w % self.spatial_factor:
+            raise ValueError(
+                f"encode requires H and W divisible by {self.spatial_factor}, got {h}x{w}"
+            )
+        moments = self.encoder(x)
+        mean = torch.chunk(moments, 2, dim=1)[0]
+        return self.normalize(self.pack(mean))
+
+    @torch.no_grad()
+    def decode(self, z: Tensor) -> Tensor:
+        """Decode a packed, normalised latent ``[B, 128, h, w]`` to ``[B, 3, 16h, 16w]``.
+
+        The output is approximately ``[-1, 1]`` and is **not** clamped here; the
+        pipeline applies ``(x * 0.5 + 0.5).clamp(0, 1)``.
+        """
+        if z.ndim != 4 or z.shape[1] != self.latent_channels:
+            raise ValueError(
+                f"decode expects [B, {self.latent_channels}, h, w], got {tuple(z.shape)}"
+            )
+        return self.decoder(self.unpack(self.inv_normalize(z)))
+
+    # ---- construction / freezing -------------------------------------------
+
+    @classmethod
+    def from_single_file(
+        cls,
+        path: str | os.PathLike,
+        torch_dtype: torch.dtype = torch.float32,
+    ) -> "AutoencoderFlux2":
+        """Load the single-file ``ae.safetensors`` (BFL key names) with ``strict=True``.
+
+        The released file is the Apache-2.0 FLUX.2-klein-base-4B autoencoder
+        re-keyed to this layout; it is stored in bfloat16 and is upcast to
+        ``torch_dtype``.  ReFlowSET runs the autoencoder in float32.
+        """
+        from safetensors.torch import load_file
+
+        path = os.fspath(path)
+        if not os.path.isfile(path):
+            raise FileNotFoundError(
+                f"FLUX.2 autoencoder weights not found at: {path}. Expected the "
+                "single-file 'ae.safetensors' shipped with ReFlowSET."
+            )
+        model = cls()
+        model.load_state_dict(load_file(path, device="cpu"), strict=True)
+        model.to(dtype=torch_dtype)
+        model.eval()
+        model.requires_grad_(False)
+        return model
+
+    def train(self, mode: bool = True) -> "AutoencoderFlux2":
+        """The autoencoder is frozen: never leave eval mode (flux2_ae.py:437-439)."""
+        return super().train(False)

sar2opt/vae/config.json

@@ -0,0 +1,18 @@
+{
+  "_class_name": "AutoencoderFlux2",
+  "_diffusers_version": "0.37.1",
+  "bn_eps": 0.0001,
+  "ch": 128,
+  "ch_mult": [
+    1,
+    2,
+    4,
+    4
+  ],
+  "in_channels": 3,
+  "num_res_blocks": 2,
+  "out_ch": 3,
+  "patch_size": 2,
+  "resolution": 256,
+  "z_channels": 32
+}

sar2opt/vae/diffusion_pytorch_model.safetensors

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

scheduler_flow_bridge.py

@@ -0,0 +1,176 @@
+"""ReFlowSET's Design-B flow bridge and its explicit-Euler solver.
+
+Forward (training) process, with ``eps ~ N(0, I)`` and ``z_e`` the EO latent::
+
+    z_t = (1 - t) * eps + t * z_e          (bridge.py:311, sigma_b = 0)
+    u*  = z_e - eps                        (bridge.py:328 at sigma_b = 0)
+
+Sampling starts from ``z_0 ~ N(0, I)`` and integrates the predicted velocity
+with explicit Euler on a uniform grid ``linspace(0, t_end, nfe + 1)``
+(bridge.py:519, 536).  The bridge is deterministic: ``sigma_b = 0``, so no
+stochastic term ever executes, and the only randomness in a sample is the
+initial noise draw.
+
+**Time direction.**  ``t = 0`` is NOISE and ``t = 1`` is DATA, and the solver
+integrates ``t`` **ascending** (bridge.py:86-88).  That is the opposite of
+`diffusers`' ``sigma`` convention: setting ``sigma := 1 - t`` recovers
+``FlowMatchEulerDiscreteScheduler``'s interpolation, but then this bridge's
+velocity is the **negative** of the diffusers flow-matching target and the
+network must still be fed ``1 - sigma``.  This scheduler keeps ReFlowSET's own
+sign and direction so neither flip is needed; ``timesteps`` therefore *increase*
+from 0 towards 1, unlike every noise-schedule scheduler in `diffusers`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class FlowBridgeSchedulerOutput(BaseOutput):
+    """Output of :meth:`FlowBridgeScheduler.step`.
+
+    Args:
+        prev_sample: The bridge state at the next time on the grid.
+    """
+
+    prev_sample: torch.Tensor
+
+
+class FlowBridgeScheduler(SchedulerMixin, ConfigMixin):
+    """Explicit-Euler solver for ReFlowSET's Design-B flow bridge.
+
+    Args:
+        t_end: End time of the integration grid (1.0 — the EO endpoint).  The
+            model is evaluated at ``linspace(0, t_end, nfe + 1)[:-1]`` and the
+            final Euler step lands on ``t_end``; the network is never queried at
+            ``t = t_end``.
+    """
+
+    order = 1
+
+    @register_to_config
+    def __init__(self, t_end: float = 1.0) -> None:
+        if not 0.0 < t_end <= 1.0:
+            raise ValueError(f"t_end must lie in (0, 1], got {t_end}")
+        self._grid: Optional[torch.Tensor] = None
+        self._step_index: Optional[int] = None
+        self.num_inference_steps: Optional[int] = None
+
+    @property
+    def timesteps(self) -> torch.Tensor:
+        """The ``nfe`` bridge times at which the model is evaluated, ascending."""
+        if self._grid is None:
+            raise ValueError("call set_timesteps() before reading timesteps")
+        return self._grid[:-1]
+
+    @property
+    def step_index(self) -> Optional[int]:
+        """Index of the next grid interval; ``None`` until the first :meth:`step`."""
+        return self._step_index
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        device: Optional[Union[str, torch.device]] = None,
+    ) -> None:
+        """Build the uniform grid ``linspace(0, t_end, num_inference_steps + 1)``.
+
+        Args:
+            num_inference_steps: NFE — the number of velocity evaluations.
+                50 reproduces the paper's main results; 4 is the efficiency
+                operating point.
+            device: Device the grid is built on.
+
+        There is no shift, no dynamic shifting, no Karras or exponential
+        spacing, and no timestep-spacing option: the reference solver uses a
+        plain uniform grid (bridge.py:519).
+        """
+        if num_inference_steps < 1:
+            raise ValueError(f"num_inference_steps must be >= 1, got {num_inference_steps}")
+        self.num_inference_steps = num_inference_steps
+        self._grid = torch.linspace(
+            0.0, self.config.t_end, num_inference_steps + 1, device=device, dtype=torch.float32
+        )
+        self._step_index = 0
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[float, torch.Tensor],
+        sample: torch.Tensor,
+        return_dict: bool = True,
+    ) -> Union[FlowBridgeSchedulerOutput, tuple[torch.Tensor]]:
+        """One explicit-Euler step: ``z + (t_next - t_cur) * v`` (bridge.py:536).
+
+        Args:
+            model_output: The predicted velocity ``dz/dt`` at ``timestep``,
+                already classifier-free-guided by the caller.
+            timestep: The current bridge time.  Present for API compatibility and
+                checked against the grid; the step size comes from the grid.
+            sample: The current bridge state.
+            return_dict: Return a :class:`FlowBridgeSchedulerOutput` instead of a
+                tuple.
+
+        Steps must be taken in order, starting from the first entry of
+        :attr:`timesteps`.
+        """
+        if self._grid is None or self._step_index is None:
+            raise ValueError("call set_timesteps() before step()")
+        if self._step_index >= self.num_inference_steps:
+            raise ValueError(
+                f"already took {self.num_inference_steps} steps; call set_timesteps() again"
+            )
+        t_cur, t_next = self._grid[self._step_index], self._grid[self._step_index + 1]
+        if not torch.isclose(torch.as_tensor(timestep, dtype=torch.float32).to(t_cur.device), t_cur):
+            raise ValueError(
+                f"step {self._step_index} expects timestep {t_cur.item()}, got {float(timestep)}; "
+                "the flow bridge must be integrated in ascending grid order"
+            )
+
+        # The state is carried in float32 even if the model ran lower (bridge.py:515-517).
+        dtype = sample.dtype if sample.dtype in (torch.float32, torch.float64) else torch.float32
+        prev_sample = sample.to(dtype) + (t_next - t_cur) * model_output.to(dtype)
+        prev_sample = prev_sample.to(sample.dtype)
+
+        self._step_index += 1
+        if not return_dict:
+            return (prev_sample,)
+        return FlowBridgeSchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        """The training-side bridge state ``z_t = (1 - t) * eps + t * z_e`` (bridge.py:311).
+
+        Args:
+            original_samples: The EO latent ``z_e`` (the ``t = 1`` endpoint).
+            noise: ``eps ~ N(0, I)`` (the ``t = 0`` endpoint).
+            timesteps: Bridge times in ``[0, 1]``, broadcastable over the batch.
+        """
+        t = timesteps.to(original_samples.device, original_samples.dtype)
+        t = t.view(-1, *([1] * (original_samples.ndim - 1)))
+        return (1.0 - t) * noise + t * original_samples
+
+    def get_velocity(
+        self,
+        sample: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        """The training target ``u* = z_e - eps`` (bridge.py:328 at ``sigma_b = 0``).
+
+        Constant along the path, hence independent of ``timesteps``; the argument
+        is kept for `diffusers` API compatibility.
+        """
+        del timesteps
+        return sample - noise

transformer_reflowset.py

@@ -0,0 +1,471 @@
+"""ReFlowSET velocity transformer — a latent DiT with an EO/SAR double stream.
+
+The network predicts the flow-bridge velocity ``dz/dt`` in the frozen FLUX.2
+latent space.  It takes the noisy EO latent ``[B, 128, h, w]``, a scalar bridge
+time ``t`` in ``[0, 1]``, and the SAR conditioning latent of the same shape; the
+first 8 of its 24 blocks are double-stream (one EO tower and one SAR tower over
+a single joint attention), the remaining 16 are single-stream over the
+concatenated ``[EO | SAR]`` sequence, and only the EO half is decoded.
+
+This is an inference-only port.  The training-only REPA projection head
+(``repa_proj``) is a separate module in the reference implementation and is
+deliberately absent here.
+
+Deviations from `diffusers`' FLUX blocks that this file has to keep — each one
+is silent if you get it wrong:
+
+* ``FinalLayer`` unpacks ``shift, scale`` (dit.py:360), the **opposite** order of
+  ``AdaLayerNormContinuous``.
+* The single-stream MLP is **SwiGLU** of width 2752, not a 4x GELU of width 4096.
+* ``linear1``/``linear2`` are **bias-free**, and the QK-norm parameter is called
+  ``scale``, not ``weight``.
+* The timestep is multiplied by 1000 *inside* the model and the sinusoid is
+  **cos first, then sin**.
+* RoPE runs on **two** axes of **centred half-integer** coordinates, not on
+  FLUX's three axes of integers starting at 0.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional, Union
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+#: Width of the sinusoidal timestep embedding fed to ``time_in`` (dit.py:44).
+#: A module constant, deliberately independent of ``hidden_size``.
+TIME_EMBED_DIM = 256
+
+
+def swiglu_hidden_dim(hidden_size: int, mlp_ratio: float) -> int:
+    """SwiGLU intermediate width (dit.py:120-127).
+
+    The canonical 2/3 rule rounded to a multiple of 64, so a gated MLP at
+    ``mlp_ratio=4.0`` costs the same parameters as a plain 4x GELU MLP.
+    ``hidden_size=1024, mlp_ratio=4.0 -> 2752``.
+    """
+    return int(round(hidden_size * mlp_ratio * 2 / 3 / 64)) * 64
+
+
+class SwiGLU(nn.Module):
+    """``silu(first half) * second half`` — gate first, value second (dit.py:130-133)."""
+
+    def forward(self, x: Tensor) -> Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        return F.silu(x1) * x2
+
+
+class RMSNorm(nn.Module):
+    """RMS norm computed in float32.  The parameter is named ``scale`` (dit.py:136-145)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        x_dtype = x.dtype
+        x = x.float()
+        rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
+        return (x * rrms).to(dtype=x_dtype) * self.scale
+
+
+class QKNorm(nn.Module):
+    """Per-head query/key RMS norm, applied **before** RoPE (dit.py:148-155)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.query_norm = RMSNorm(dim)
+        self.key_norm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        return self.query_norm(q).to(v), self.key_norm(k).to(v)
+
+
+class MLPEmbedder(nn.Module):
+    """``Linear -> SiLU -> Linear`` time-embedding MLP (dit.py:158-166).
+
+    Checkpoint keys are ``time_in.in_layer.*`` / ``time_in.out_layer.*``, not
+    diffusers' ``time_text_embed.timestep_embedder.linear_{1,2}``.
+    """
+
+    def __init__(self, in_dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+        self.silu = nn.SiLU()
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class Modulation(nn.Module):
+    """AdaLN-Zero triple.  Order is ``shift, scale, gate`` (dit.py:169-181)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.lin = nn.Linear(dim, 3 * dim, bias=True)
+
+    def forward(self, vec: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        out = self.lin(F.silu(vec))
+        if out.ndim == 2:
+            out = out[:, None, :]
+        shift, scale, gate = out.chunk(3, dim=-1)
+        return shift, scale, gate
+
+
+def timestep_embedding(
+    t: Tensor, dim: int, max_period: int = 10000, time_factor: float = 1000.0
+) -> Tensor:
+    """Sinusoidal embedding of a fractional bridge time (dit.py:184-201).
+
+    Two things differ from `diffusers`' ``get_timestep_embedding`` defaults:
+    ``t`` is a fraction in ``[0, 1]`` that is scaled by ``time_factor = 1000``
+    **here**, and the concatenation order is ``[cos, sin]`` (FLUX's ordering,
+    i.e. ``flip_sin_to_cos=True``).
+    """
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period)
+        * torch.arange(start=0, end=half, device=t.device, dtype=torch.float32)
+        / half
+    )
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    """Per-axis rotation matrices ``[..., L, dim/2, 2, 2]`` (dit.py:204-211)."""
+    if dim % 2:
+        raise ValueError(f"RoPE axis dim must be even, got {dim}")
+    scale = torch.arange(0, dim, 2, dtype=pos.dtype, device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
+    return out.reshape(*out.shape[:-1], 2, 2).float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+    """Rotate consecutive dimension pairs — the interleaved (FLUX) convention.
+
+    ``(x0, x1) -> (cos*x0 - sin*x1, sin*x0 + cos*x1)`` on ``(x[2k], x[2k+1])``
+    (dit.py:214-219).  Equivalent to diffusers' ``apply_rotary_emb(...,
+    use_real_unbind_dim=-1)``; ``-2`` is the split-halves convention and is wrong
+    for these weights.
+    """
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+
+
+class EmbedND(nn.Module):
+    """Concatenates the per-axis RoPE ladders and inserts the head axis (dit.py:222-234).
+
+    Holds no parameters and no buffers: the grid is rebuilt on every forward,
+    which is what lets one checkpoint serve 256 and 512 inputs.
+    """
+
+    def __init__(self, theta: int, axes_dim: list[int]) -> None:
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        emb = torch.cat(
+            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(len(self.axes_dim))],
+            dim=-3,
+        )
+        return emb.unsqueeze(1)
+
+
+def latent_image_ids(h: int, w: int, device, dtype=torch.float32) -> Tensor:
+    """Centred ``(y, x)`` coordinates for an ``h x w`` latent grid, ``[h*w, 2]``.
+
+    ``arange(n) - (n - 1) / 2`` with unit spacing (dit.py:237-252), so for even
+    ``n`` the coordinates are half-integers and the central 16x16 region of a
+    32x32 grid carries exactly the coordinates a 256-trained model saw — RoPE
+    only extrapolates outwards, it never rescales.  Row-major, so token
+    ``p = y * w + x``.  This is **not** FLUX's 3-axis integer id grid.
+    """
+    y = torch.arange(h, device=device, dtype=dtype) - (h - 1) / 2
+    x = torch.arange(w, device=device, dtype=dtype) - (w - 1) / 2
+    ids = torch.zeros(h, w, 2, device=device, dtype=dtype)
+    ids[..., 0] = y[:, None]
+    ids[..., 1] = x[None, :]
+    return ids.reshape(h * w, 2)
+
+
+class SingleStreamBlock(nn.Module):
+    """Fused attention + SwiGLU MLP under one modulation and one residual.
+
+    ``linear1`` emits ``[q | k | v | mlp_gate | mlp_value]`` in that order; the
+    qkv slab is K-major (``(K H D)``).  Both linears are bias-free
+    (dit.py:255-300).
+    """
+
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.mlp_hidden_dim = swiglu_hidden_dim(hidden_size, mlp_ratio)
+
+        self.linear1 = nn.Linear(hidden_size, 3 * hidden_size + 2 * self.mlp_hidden_dim, bias=False)
+        self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size, bias=False)
+        self.norm = QKNorm(head_dim)
+        self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp_act = SwiGLU()
+        self.modulation = Modulation(hidden_size)
+
+    def pre_attention(self, x: Tensor, vec: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+        """Everything up to (not including) RoPE and attention (dit.py:272-288)."""
+        shift, scale, gate = self.modulation(vec)
+        x_mod = (1 + scale) * self.pre_norm(x) + shift
+
+        qkv, mlp = torch.split(
+            self.linear1(x_mod), [3 * self.hidden_size, 2 * self.mlp_hidden_dim], dim=-1
+        )
+        b, length, _ = qkv.shape
+        # "B L (K H D) -> K B H L D" with K=3, H=num_heads.
+        q, k, v = qkv.reshape(b, length, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        q, k = self.norm(q, k, v)
+        return q, k, v, mlp, gate
+
+    def post_attention(self, x: Tensor, attn: Tensor, mlp: Tensor, gate: Tensor) -> Tensor:
+        """Output projection and the single gated residual (dit.py:290-294)."""
+        b, heads, length, head_dim = attn.shape
+        attn = attn.transpose(1, 2).reshape(b, length, heads * head_dim)
+        out = self.linear2(torch.cat((attn, self.mlp_act(mlp)), dim=-1))
+        return x + gate * out
+
+    def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
+        q, k, v, mlp, gate = self.pre_attention(x, vec)
+        q, k = apply_rope(q, k, pe)
+        attn = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+        return self.post_attention(x, attn, mlp, gate)
+
+
+class DoubleStreamBlock(nn.Module):
+    """Two independent towers over **one** joint attention across ``[EO | SAR]``.
+
+    The towers have completely separate weights but share the modulation vector
+    ``vec`` and the RoPE grid, so an EO token and the SAR token at the same
+    ground position carry an identical phase (dit.py:303-343).  SAR plays the
+    structural role text plays in FLUX.
+    """
+
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0) -> None:
+        super().__init__()
+        self.eo = SingleStreamBlock(hidden_size, num_heads, mlp_ratio)
+        self.sar = SingleStreamBlock(hidden_size, num_heads, mlp_ratio)
+
+    def forward(self, eo: Tensor, sar: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]:
+        """``pe`` must already cover the joint 2P-token sequence."""
+        q_e, k_e, v_e, mlp_e, gate_e = self.eo.pre_attention(eo, vec)
+        q_s, k_s, v_s, mlp_s, gate_s = self.sar.pre_attention(sar, vec)
+
+        q = torch.cat((q_e, q_s), dim=2)
+        k = torch.cat((k_e, k_s), dim=2)
+        v = torch.cat((v_e, v_s), dim=2)
+        q, k = apply_rope(q, k, pe)
+
+        attn = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+        attn_e, attn_s = attn.split([q_e.shape[2], q_s.shape[2]], dim=2)
+        return (
+            self.eo.post_attention(eo, attn_e, mlp_e, gate_e),
+            self.sar.post_attention(sar, attn_s, mlp_s, gate_s),
+        )
+
+
+class FinalLayer(nn.Module):
+    """AdaLN output layer.
+
+    ``adaLN`` unpacks ``shift, scale`` — the **opposite** order of diffusers'
+    ``AdaLayerNormContinuous`` (dit.py:346-362).  ``logvar_proj`` belongs to a
+    beta-NLL loss that was never enabled (``loss.flow = mse``); its weights are
+    kept so the published checkpoint loads with ``strict=True``, but inference
+    never evaluates it — the sampler reads only the velocity (bridge.py:531).
+    """
+
+    def __init__(self, hidden_size: int, out_channels: int) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.adaLN = nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+        self.proj = nn.Linear(hidden_size, out_channels, bias=True)
+        self.logvar_proj = nn.Linear(hidden_size, 1, bias=True)
+
+    def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+        mod = self.adaLN(F.silu(vec))
+        if mod.ndim == 2:
+            mod = mod[:, None, :]
+        shift, scale = mod.chunk(2, dim=-1)
+        return self.proj((1 + scale) * self.norm(x) + shift)
+
+
+class ReFlowSETTransformer2DModel(ModelMixin, ConfigMixin):
+    """ReFlowSET's flow-velocity transformer (509.32 M parameters as configured).
+
+    Args:
+        in_channels: Channels of the packed FLUX.2 latent (128).
+        out_channels: Channels of the predicted velocity (128).
+        hidden_size: Residual width (1024).
+        depth: **Total** blocks, double plus single (24).
+        num_heads: Attention heads (16), so ``head_dim = 64``.
+        mlp_ratio: Nominal MLP ratio; the SwiGLU width is derived from it.
+        axes_dim: RoPE dims for the ``(y, x)`` axes; must sum to ``head_dim``.
+        theta: RoPE base period (10000).
+        sample_size: Input image resolution the released arm was trained at
+            (256 for QXS-SAROPT, 512 for SAR2Opt).  Recorded for provenance
+            only: the forward pass derives every shape from its input and the
+            RoPE grid is rebuilt per call, so one checkpoint serves any size
+            divisible by 16.
+        double_blocks: Leading double-stream blocks (8); the remaining
+            ``depth - double_blocks`` are single-stream.
+        double_merge: How the two streams become one.  ``"token"`` (the released
+            setting) concatenates on the sequence axis, so the single stack runs
+            over 2P tokens and the SAR half is dropped only at the very end;
+            ``"channel"`` fuses per position and keeps P tokens.
+
+    Forward contract: ``forward(hidden_states, timestep, condition)`` where
+    ``hidden_states`` is the bridge state ``[B, 128, h, w]``, ``timestep`` is the
+    bridge time in ``[0, 1]`` (**not** an integer diffusion step), and
+    ``condition`` is the SAR latent of the same shape or ``None``.  ``None`` is
+    the classifier-free-guidance null branch and is turned into an all-zero
+    latent inside the model — there is no learned null token.
+    """
+
+    _supports_gradient_checkpointing = False
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 128,
+        out_channels: int = 128,
+        hidden_size: int = 1024,
+        depth: int = 24,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        axes_dim: tuple[int, ...] = (32, 32),
+        theta: int = 10000,
+        sample_size: Optional[int] = None,
+        double_blocks: int = 8,
+        double_merge: str = "token",
+    ) -> None:
+        super().__init__()
+        if hidden_size % num_heads != 0:
+            raise ValueError(f"hidden_size {hidden_size} must be divisible by num_heads {num_heads}")
+        pe_dim = hidden_size // num_heads
+        if sum(axes_dim) != pe_dim:
+            raise ValueError(f"axes_dim {list(axes_dim)} must sum to the per-head dim {pe_dim}")
+        if not 0 <= double_blocks < depth:
+            raise ValueError(f"double_blocks {double_blocks} must be in [0, depth={depth})")
+        if double_merge not in ("token", "channel"):
+            raise ValueError(f"double_merge must be 'token' or 'channel', got {double_merge!r}")
+
+        self.pe_embedder = EmbedND(theta=theta, axes_dim=list(axes_dim))
+        if double_blocks:
+            # Each stream gets its own 1x1 "patchify": they are two token
+            # sequences now, not two halves of one channel stack.
+            self.in_proj_eo = nn.Linear(in_channels, hidden_size, bias=True)
+            self.in_proj_sar = nn.Linear(in_channels, hidden_size, bias=True)
+            if double_merge == "channel":
+                self.merge = nn.Linear(2 * hidden_size, hidden_size, bias=True)
+        else:
+            self.in_proj = nn.Linear(2 * in_channels, hidden_size, bias=True)
+        self.time_in = MLPEmbedder(TIME_EMBED_DIM, hidden_size)
+        self.double_stream = nn.ModuleList(
+            [DoubleStreamBlock(hidden_size, num_heads, mlp_ratio) for _ in range(double_blocks)]
+        )
+        self.blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(hidden_size, num_heads, mlp_ratio)
+                for _ in range(depth - double_blocks)
+            ]
+        )
+        self.final_layer = FinalLayer(hidden_size, out_channels)
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        timestep: Tensor,
+        condition: Optional[Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[Transformer2DModelOutput, tuple[Tensor]]:
+        """Predict the flow velocity ``dz/dt``.
+
+        Args:
+            hidden_states: ``[B, in_channels, h, w]`` bridge state.
+            timestep: Bridge time in ``[0, 1]``; a scalar or ``[B]``.
+            condition: ``[B, in_channels, h, w]`` SAR latent, or ``None`` for the
+                null branch (an all-zero conditioning latent, dit.py:531-532).
+            return_dict: Return a ``Transformer2DModelOutput`` instead of a tuple.
+
+        Returns:
+            The velocity ``[B, out_channels, h, w]``.  This is a flow velocity,
+            not ``epsilon`` and not diffusers' ``v_prediction``.
+        """
+        if hidden_states.ndim != 4:
+            raise ValueError(f"hidden_states must be [B, C, h, w], got {tuple(hidden_states.shape)}")
+        batch, _, h, w = hidden_states.shape
+        if condition is None:
+            condition = torch.zeros_like(hidden_states)
+        elif condition.shape != hidden_states.shape:
+            raise ValueError(
+                f"condition shape {tuple(condition.shape)} must match "
+                f"hidden_states shape {tuple(hidden_states.shape)}"
+            )
+        if timestep.ndim == 0:
+            timestep = timestep.expand(batch)
+
+        n_double = self.config.double_blocks
+        if n_double:
+            eo = self.in_proj_eo(hidden_states.flatten(2).transpose(1, 2))  # [B, P, D]
+            sar = self.in_proj_sar(condition.flatten(2).transpose(1, 2))  # [B, P, D]
+            ref = eo
+        else:
+            x = torch.cat([hidden_states, condition], dim=1).flatten(2).transpose(1, 2)
+            x = self.in_proj(x)
+            ref = x
+
+        vec = self.time_in(timestep_embedding(timestep, TIME_EMBED_DIM).to(ref.dtype))
+
+        ids = latent_image_ids(h, w, device=hidden_states.device, dtype=torch.float32)
+        pe = self.pe_embedder(ids[None].expand(batch, -1, -1))
+
+        num_tokens = ref.shape[1]
+        pe_single = pe
+        if n_double:
+            # Token axis of pe is dim 2 ([B, 1, L, head_dim/2, 2, 2]); repeating
+            # the same P coordinates gives EO and SAR one shared grid.
+            pe_joint = torch.cat((pe, pe), dim=2)
+            for block in self.double_stream:
+                eo, sar = block(eo, sar, vec, pe_joint)
+            if self.config.double_merge == "token":
+                x = torch.cat((eo, sar), dim=1)  # [B, 2P, D]
+                pe_single = pe_joint
+            else:
+                x = self.merge(torch.cat((eo, sar), dim=-1))  # [B, P, D]
+
+        for block in self.blocks:
+            x = block(x, vec, pe_single)
+
+        if n_double and self.config.double_merge == "token":
+            x = x[:, :num_tokens]  # drop the SAR half: only EO is decoded
+
+        v = self.final_layer(x, vec)
+        v = v.transpose(1, 2).reshape(batch, self.config.out_channels, h, w)
+        if not return_dict:
+            return (v,)
+        return Transformer2DModelOutput(sample=v)