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