FLUX.3: match the released BFL reference image sizing

docs/source/en/api/pipelines/flux3.md

@@ -106,7 +106,7 @@ This way it can be served as a separate stage, on its own GPU or behind its own
 
 ## Image-to-image
 
-Pass one or more reference images with `image`. Every reference keeps its own resolution (center-cropped to the 16-pixel grid; references above 16 megapixels are downscaled first), and the generation size defaults to one megapixel at the first reference's aspect ratio. The prompt the image model takes is the *target* image as a complete caption of the reference with the edit folded in, not a one-line instruction; the upsampler sees the references too and writes it from your instruction: what to keep, what to change.
+Pass one or more reference images with `image`. A single reference is kept whole: it is downscaled only above 4 megapixels, padded to the 16-pixel grid, and the generation defaults to its size, so the edit comes back at the size of the input. Several references are each fitted under 1 megapixel and cropped to the grid, and the generation defaults to the first one's size. A given `height`/`width` overrides the default. Describe the *target* image as a complete caption of the reference with your edit folded in, rather than a one-line instruction, or let the [prompt upsampler](#prompt-upsampling) write it from the instruction and the references:
 
 ```python
 night = pipe(

src/diffusers/modular_pipelines/flux3/decoders.py

@@ -20,7 +20,6 @@
 import PIL
 import torch
 
-from ...configuration_utils import FrozenDict
 from ...models import AutoencoderKLFlux2
 from ...utils import logging
 from ..modular_pipeline import ModularPipelineBlocks, PipelineState
@@ -100,12 +99,7 @@ class Flux3DecodeStep(ModularPipelineBlocks):
     def expected_components(self) -> list[ComponentSpec]:
         return [
             ComponentSpec("vae", AutoencoderKLFlux2),
-            ComponentSpec(
-                "image_processor",
-                Flux3ImageProcessor,
-                config=FrozenDict({"vae_scale_factor": 16, "vae_latent_channels": 32}),
-                default_creation_method="from_config",
-            ),
+            ComponentSpec("image_processor", Flux3ImageProcessor, default_creation_method="from_config"),
         ]
 
     @property
@@ -125,6 +119,12 @@ def inputs(self) -> list[tuple[str, Any]]:
                 type_hint=torch.Tensor,
                 description="The denoised latents from the denoising step",
             ),
+            InputParam(
+                "reference_crop",
+                type_hint=list[tuple[int, int, int, int] | None] | None,
+                description="For image-to-image, per reference, the (left, top, right, bottom) box to crop the "
+                "decoded image back to (a single padded reference) or None",
+            ),
         ]
 
     @property
@@ -162,8 +162,9 @@ def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> Pi
         latents = self._unpatchify_latents(latents)
 
         block_state.images = vae.decode(latents.to(vae.dtype), return_dict=False)[0]
+        # a single padded reference is cropped back to its own size when the generation size was not given
         block_state.images = components.image_processor.postprocess(
-            block_state.images, output_type=block_state.output_type
+            block_state.images, output_type=block_state.output_type, reference_crop=block_state.reference_crop
         )
 
         self.set_block_state(state, block_state)

src/diffusers/modular_pipelines/flux3/encoders.py

@@ -18,7 +18,6 @@
 import torch
 from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
 
-from ...configuration_utils import FrozenDict
 from ...models import AutoencoderKLFlux2
 from ...utils import logging
 from ..modular_pipeline import ModularPipelineBlocks, PipelineState
@@ -36,29 +35,27 @@ class Flux3ProcessImagesInputStep(ModularPipelineBlocks):
     @property
     def description(self) -> str:
         return (
-            "Image preprocess step for FLUX.3 image-to-image generation. Each reference image keeps its own "
-            "resolution: it is converted to RGB, downscaled only when it exceeds 4096 x 4096 pixels, and "
-            "center-cropped to a multiple of 16 pixels."
+            "Image preprocess step for FLUX.3 image-to-image generation. A single reference is downscaled to fit "
+            "under 4 megapixels (sides at most 4096) and reflect-padded to the 16-pixel grid, so nothing is cut; the "
+            "generation defaults to that canvas and the decoded image is cropped back to the reference "
+            "(`reference_crop`). Each reference of a multi-reference edit is downscaled to fit under 1 megapixel "
+            "and center-cropped to the grid."
         )
 
     @property
     def expected_components(self) -> list[ComponentSpec]:
         return [
-            ComponentSpec(
-                "image_processor",
-                Flux3ImageProcessor,
-                config=FrozenDict({"vae_scale_factor": 16, "vae_latent_channels": 32}),
-                default_creation_method="from_config",
-            ),
+            ComponentSpec("image_processor", Flux3ImageProcessor, default_creation_method="from_config"),
         ]
 
     @property
     def inputs(self) -> list[InputParam]:
         return [
             InputParam.template(
                 "image",
-                description="The reference image or images to condition the generation on; each keeps its own "
-                "resolution (center-cropped to the 16-pixel grid)",
+                description="The reference image or images to condition the generation on. One reference is kept "
+                "whole (padded to the 16-pixel grid, up to 4 megapixels); several are fitted under 1 megapixel each "
+                "and cropped to the grid",
             ),
         ]
 
@@ -69,7 +66,14 @@ def intermediate_outputs(self) -> list[OutputParam]:
                 "condition_images",
                 type_hint=list[torch.Tensor],
                 description="The preprocessed reference images, one tensor of shape (1, 3, height_i, width_i) in [-1, 1] per image",
-            )
+            ),
+            OutputParam(
+                "reference_crop",
+                type_hint=list[tuple[int, int, int, int] | None],
+                description="Per reference, the (left, top, right, bottom) box of the reference inside its padded "
+                "canvas (a single reference) or None (several, cropped); the decoded image is cropped back to the box "
+                "when the generation size was not given",
+            ),
         ]
 
     @torch.no_grad()
@@ -79,15 +83,17 @@ def __call__(
         block_state = self.get_block_state(state)
 
         images = block_state.image if isinstance(block_state.image, list) else [block_state.image]
-        condition_images = []
         for image in images:
             if not isinstance(image, PIL.Image.Image):
                 raise ValueError(f"`image` must be a `PIL.Image.Image` or a list of them, got {type(image)}")
-            # references keep their own resolution: downscaled only above 4096 x 4096 pixels, then center-cropped to
-            # the 16-pixel grid (`resize_mode="crop"` is a pure crop for this processor)
-            image = components.image_processor._resize_if_exceeds_area(image, 4096 * 4096)
-            condition_images.append(components.image_processor.preprocess(image, resize_mode="crop"))
-        block_state.condition_images = condition_images
+        # one reference is kept whole (`resize_mode="fill"`: downscaled so that its canvas fits under 4 megapixels and
+        # reflect-padded to it); several are each downscaled to fit under 1 megapixel and center-cropped to the grid
+        resize_mode = "fill" if len(images) == 1 else "crop"
+        condition_images, postprocessing_kwargs = zip(
+            *(components.image_processor.preprocess(image, resize_mode=resize_mode) for image in images)
+        )
+        block_state.condition_images = list(condition_images)
+        block_state.reference_crop = [kwargs["reference_crop"] for kwargs in postprocessing_kwargs]
 
         self.set_block_state(state, block_state)
         return components, state

src/diffusers/modular_pipelines/flux3/image_processor.py

@@ -14,54 +14,42 @@
 
 import math
 
+import numpy as np
 import PIL.Image
+import torch
 
 from ...image_processor import VaeImageProcessor
 
 
 class Flux3ImageProcessor(VaeImageProcessor):
     r"""
     Image processor for FLUX.3: preprocesses the reference images of image-to-image generation and postprocesses the
-    generated images.
+    generated images. `resize_mode="fill"` keeps a reference whole: it is downscaled so that its canvas (its size
+    rounded up to the token grid) fits under 4 megapixels, then reflect-padded to the canvas, and `preprocess` also
+    returns the `reference_crop` box of the image inside the canvas, which `postprocess` takes to crop the generated
+    image back to it. `resize_mode="crop"` downscales a reference to fit under 1 megapixel and center-crops it to the
+    token grid. The sizing follows the released BFL code, `reference_canvas` / `prepare_reference` for one reference
+    and `prepare_multi_reference` for several, in
+    https://github.com/black-forest-labs/flux-3-image-klein/blob/main/src/flux3/pipeline.py.
 
     Args:
-        do_resize (`bool`, *optional*, defaults to `True`):
-            Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`. Can accept
-            `height` and `width` arguments from [`image_processor.VaeImageProcessor.preprocess`] method.
         vae_scale_factor (`int`, *optional*, defaults to `16`):
-            VAE (spatial) scale factor, including the 2x2 packing of the latents. If `do_resize` is `True`, the image
-            is automatically resized to multiples of this factor.
+            VAE (spatial) scale factor, including the 2x2 packing of the latents: one token. Images are resized to
+            multiples of this factor.
         vae_latent_channels (`int`, *optional*, defaults to `32`):
             VAE latent channels.
-        do_normalize (`bool`, *optional*, defaults to `True`):
-            Whether to normalize the image to [-1,1].
-        do_binarize (`bool`, *optional*, defaults to `False`):
-            Whether to binarize the image to 0/1.
-        do_convert_rgb (`bool`, *optional*, defaults to be `True`):
-            Whether to convert the images to RGB format.
-        do_convert_grayscale (`bool`, *optional*, defaults to be `False`):
-            Whether to convert the images to grayscale format.
     """
 
-    # Copied from diffusers.pipelines.flux2.image_processor.Flux2ImageProcessor.__init__
-    def __init__(
-        self,
-        do_resize: bool = True,
-        vae_scale_factor: int = 16,
-        vae_latent_channels: int = 32,
-        do_normalize: bool = True,
-        do_binarize: bool = False,
-        do_convert_rgb: bool = True,
-        do_convert_grayscale: bool = False,
-    ):
+    def __init__(self, vae_scale_factor: int = 16, vae_latent_channels: int = 32):
+        # `preprocess` always passes the size to resize to, and the VAE takes RGB images in [-1, 1]
         super().__init__(
-            do_resize=do_resize,
+            do_resize=True,
             vae_scale_factor=vae_scale_factor,
             vae_latent_channels=vae_latent_channels,
-            do_normalize=do_normalize,
-            do_binarize=do_binarize,
-            do_convert_rgb=do_convert_rgb,
-            do_convert_grayscale=do_convert_grayscale,
+            do_normalize=True,
+            do_binarize=False,
+            do_convert_rgb=True,
+            do_convert_grayscale=False,
         )
 
     @staticmethod
@@ -114,3 +102,135 @@ def _resize_and_crop(
         bottom = top + height
 
         return image.crop((left, top, right, bottom))
+
+    def _get_canvas_height_width(
+        self, image_height: int, image_width: int, max_side: int = 4096, min_side: int = 64
+    ) -> tuple[int, int]:
+        r"""
+        The canvas of an image kept whole: its size rounded up to the token grid, clamped to [`min_side`, `max_side`]
+        (the `canvas` of BFL's `reference_canvas`).
+        """
+        multiple_of = self.config.vae_scale_factor
+        return tuple(
+            min(max_side, max(min_side, -(-side // multiple_of) * multiple_of)) for side in (image_height, image_width)
+        )
+
+    def _resize_if_canvas_exceeds_area(
+        self, image: PIL.Image.Image, target_area: int = 4 * 1024 * 1024, max_side: int = 4096
+    ) -> PIL.Image.Image:
+        r"""
+        Downscale the image, keeping its aspect ratio, so that its canvas (its size rounded up to the token grid) fits
+        under `target_area` and its sides under `max_side`, as BFL's `reference_canvas` does.
+
+        The scale that fits the image itself under `target_area` can still leave the canvas over it once the sides are
+        rounded up, so it is backed off by 0.5% until the canvas fits. The first pass already fits unless the image is
+        within a token of the cap.
+        """
+        image_width, image_height = image.size
+        scale = min(1.0, math.sqrt(target_area / (image_width * image_height)), max_side / max(image.size))
+        while True:
+            height, width = (max(1, round(side * scale)) for side in (image_height, image_width))
+            canvas_height, canvas_width = self._get_canvas_height_width(height, width, max_side=max_side)
+            if canvas_height * canvas_width <= target_area:
+                break
+            scale *= 0.995
+        if (width, height) == image.size:
+            return image
+        return self.resize(image, height=height, width=width)
+
+    def _get_fill_region(self, image: PIL.Image.Image, height: int, width: int) -> tuple[int, int, int, int]:
+        r"""
+        The (left, top, right, bottom) box the image occupies once padded to `height` x `width` with
+        `resize_mode="fill"`: centered.
+        """
+        image_width, image_height = image.size
+        if image_width > width or image_height > height:
+            raise ValueError(
+                f"resize_mode='fill' pads the image to the given size, got a {image_width} x {image_height} image for "
+                f"a {width} x {height} canvas"
+            )
+        left, top = (width - image_width) // 2, (height - image_height) // 2
+        return left, top, left + image_width, top + image_height
+
+    def _resize_and_fill(self, image: PIL.Image.Image, width: int, height: int) -> PIL.Image.Image:
+        r"""
+        Center the image within the specified width and height, filling the rest with its reflection. The image is not
+        resized: it must fit within the dimensions.
+
+        Args:
+            image (`PIL.Image.Image`):
+                The image to fill.
+            width (`int`):
+                The width to pad the image to.
+            height (`int`):
+                The height to pad the image to.
+
+        Returns:
+            `PIL.Image.Image`:
+                The filled image.
+        """
+        left, top, right, bottom = self._get_fill_region(image, height, width)
+        if image.size == (width, height):
+            return image
+
+        padding = ((top, height - bottom), (left, width - right), (0, 0))
+        # reflection excludes the edge pixel; an image smaller than one token is mirrored instead
+        mode = "reflect" if min(image.size) > self.config.vae_scale_factor else "symmetric"
+        return PIL.Image.fromarray(np.pad(np.asarray(image.convert("RGB")), padding, mode=mode))
+
+    def preprocess(
+        self,
+        image: PIL.Image.Image | list[PIL.Image.Image],
+        resize_mode: str = "fill",  # "fill" (one reference, kept whole), "crop" (several references)
+    ) -> tuple[torch.Tensor, dict]:
+        """
+        Preprocess the reference image(s) as [`~image_processor.VaeImageProcessor.preprocess`] does. With
+        `resize_mode="fill"`, the image is downscaled so that its canvas fits under 4 megapixels and reflect-padded to
+        the canvas; with `resize_mode="crop"`, it is downscaled to fit under 1 megapixel and center-cropped to the
+        token grid.
+
+        Returns:
+            `tuple[torch.Tensor, dict]`:
+                The preprocessed image and the kwargs for [`~Flux3ImageProcessor.postprocess`]: `reference_crop`, the
+                (left, top, right, bottom) box of the (first) image inside its canvas when it was padded
+                (`resize_mode="fill"`), else `None`.
+        """
+        images = image if isinstance(image, list) else [image]
+        if not all(isinstance(image, PIL.Image.Image) for image in images):
+            raise ValueError(f"Only PIL image input is supported, got {[type(image) for image in images]}")
+        height = width = None
+        if resize_mode == "fill":
+            images = [self._resize_if_canvas_exceeds_area(image) for image in images]
+            height, width = self._get_canvas_height_width(images[0].height, images[0].width)
+        elif resize_mode == "crop":
+            images = [self._resize_if_exceeds_area(image, 1024 * 1024) for image in images]
+        else:
+            raise ValueError(f"resize_mode must be 'fill' or 'crop', got {resize_mode}")
+        image = super().preprocess(images, height=height, width=width, resize_mode=resize_mode)
+        reference_crop = self._get_fill_region(images[0], *image.shape[-2:]) if resize_mode == "fill" else None
+        return image, {"reference_crop": reference_crop}
+
+    def postprocess(
+        self,
+        image: torch.Tensor,
+        output_type: str = "pil",
+        do_denormalize: list[bool] | None = None,
+        reference_crop: tuple[int, int, int, int] | list[tuple[int, int, int, int] | None] | None = None,
+    ) -> PIL.Image.Image | np.ndarray | torch.Tensor:
+        """
+        Postprocess the image output as [`~image_processor.VaeImageProcessor.postprocess`] does, cropped to the
+        `reference_crop` box returned by [`~Flux3ImageProcessor.preprocess`] when given. For several references, pass
+        the list of their boxes: at most one may be a box (the padded reference), the others `None`.
+        """
+        if isinstance(reference_crop, list):
+            boxes = [box for box in reference_crop if box is not None]
+            if len(boxes) > 1:
+                raise ValueError(
+                    f"The generated image can be cropped back to one reference only, got {len(boxes)} boxes in "
+                    f"`reference_crop`: {reference_crop}"
+                )
+            reference_crop = boxes[0] if boxes else None
+        if reference_crop is not None:
+            left, top, right, bottom = reference_crop
+            image = image[..., top:bottom, left:right]
+        return super().postprocess(image, output_type=output_type, do_denormalize=do_denormalize)