From f99768b48b191cb5d8df485798285ab1591cca67 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Sun, 18 Oct 2026 11:41:40 +0200 Subject: [PATCH 1/3] FLUX.3: match the released BFL reference image sizing A single reference is kept whole: downscaled so its padded canvas fits under 4 MP, reflect-padded to the token grid (`resize_mode="fill"` of the FLUX.3 image processor), and cropped back on output. Several references are fitted under 1 MP each. The default output size follows the first reference. Fixes #15068 Co-Authored-By: Claude Fable 5.1 --- docs/source/en/api/pipelines/flux3.md | 2 +- .../modular_pipelines/flux3/decoders.py | 9 +- .../modular_pipelines/flux3/encoders.py | 60 ++++++++--- .../flux3/image_processor.py | 102 +++++++++++++++++- .../modular_pipelines/flux3/inputs.py | 43 ++++---- .../flux3/modular_blocks_flux3.py | 40 ++++--- .../flux3/modular_blocks_flux3_distilled.py | 24 ++--- .../flux3/test_modular_pipeline_flux3.py | 39 ++++--- 8 files changed, 240 insertions(+), 79 deletions(-) diff --git a/docs/source/en/api/pipelines/flux3.md b/docs/source/en/api/pipelines/flux3.md index b99b4eea4d92..274ab023f5a1 100644 --- a/docs/source/en/api/pipelines/flux3.md +++ b/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( diff --git a/src/diffusers/modular_pipelines/flux3/decoders.py b/src/diffusers/modular_pipelines/flux3/decoders.py index 32f3d5a72e17..c0800b1ed914 100644 --- a/src/diffusers/modular_pipelines/flux3/decoders.py +++ b/src/diffusers/modular_pipelines/flux3/decoders.py @@ -125,6 +125,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 +168,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) diff --git a/src/diffusers/modular_pipelines/flux3/encoders.py b/src/diffusers/modular_pipelines/flux3/encoders.py index 18c6a454a63f..bbfc41a01b70 100644 --- a/src/diffusers/modular_pipelines/flux3/encoders.py +++ b/src/diffusers/modular_pipelines/flux3/encoders.py @@ -36,9 +36,11 @@ 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 @@ -57,8 +59,9 @@ 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 +72,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 +89,41 @@ 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 + image_processor = components.image_processor + multiple_of = image_processor.config.vae_scale_factor # one token: the VAE downsampling and the 2 x 2 packing + if len(images) == 1: + # one reference is kept whole: downscaled so that its canvas (the size rounded up to the token grid) fits + # under 4 megapixels, then reflect-padded to the canvas (`resize_mode="fill"` for this processor); the + # generation defaults to the canvas and the decode step crops back to the reference + image = images[0] + max_area, max_side, min_side = 4 * 1024 * 1024, 4096, 64 + scale = min(1.0, math.sqrt(max_area / (image.width * image.height)), max_side / max(image.size)) + while True: + reference_height, reference_width = ( + max(1, round(side * scale)) for side in (image.height, image.width) + ) + height, width = ( + min(max_side, max(min_side, -(-side // multiple_of) * multiple_of)) + for side in (reference_height, reference_width) + ) + if height * width <= max_area: + break + scale *= 0.995 + images = [image_processor.resize(image, height=reference_height, width=reference_width)] + preprocess_kwargs = {"height": height, "width": width, "resize_mode": "fill"} + else: + # several references are downscaled to fit under 1 megapixel each and center-cropped to the token grid + # (the default size rounds down to the grid; `resize_mode="crop"` is a pure crop for this processor) + images = [image_processor._resize_if_exceeds_area(image, 1024 * 1024) for image in images] + preprocess_kwargs = {"resize_mode": "crop"} + condition_images, postprocessing_kwargs = zip( + *(image_processor.preprocess(image, **preprocess_kwargs) 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 diff --git a/src/diffusers/modular_pipelines/flux3/image_processor.py b/src/diffusers/modular_pipelines/flux3/image_processor.py index 08b8df17fdff..7c82875b71f2 100644 --- a/src/diffusers/modular_pipelines/flux3/image_processor.py +++ b/src/diffusers/modular_pipelines/flux3/image_processor.py @@ -14,15 +14,19 @@ import math +import numpy as np import PIL.Image +import torch -from ...image_processor import VaeImageProcessor +from ...image_processor import PIL_INTERPOLATION, PipelineImageInput, 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"` pads the image to the given size with its reflection, and `preprocess` also + returns the `reference_crop` box of the image inside the padded canvas, which `postprocess` takes to crop the + generated image back to it. Args: do_resize (`bool`, *optional*, defaults to `True`): @@ -114,3 +118,97 @@ def _resize_and_crop( bottom = top + height return image.crop((left, top, right, bottom)) + + 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"`: downscaled to fit within the dimensions, maintaining the aspect ratio, and centered. + """ + image_width, image_height = image.size + if image_width > width or image_height > height: + scale = min(width / image_width, height / image_height) + image_width, image_height = max(1, round(image_width * scale)), max(1, round(image_height * scale)) + 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""" + Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center + the image within the dimensions, filling the rest with its reflection. + + Args: + image (`PIL.Image.Image`): + The image to resize and fill. + width (`int`): + The width to resize the image to. + height (`int`): + The height to resize the image to. + + Returns: + `PIL.Image.Image`: + The resized and filled image. + """ + left, top, right, bottom = self._get_fill_region(image, height, width) + if (right - left, bottom - top) != image.size: + image = image.resize((right - left, bottom - top), resample=PIL_INTERPOLATION[self.config.resample]) + 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: PipelineImageInput, + height: int | None = None, + width: int | None = None, + resize_mode: str = "default", # "default", "fill", "crop" + crops_coords: tuple[int, int, int, int] | None = None, + ) -> tuple[torch.Tensor, dict]: + """ + Preprocess the image input as [`~image_processor.VaeImageProcessor.preprocess`] does. + + 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`. + """ + reference_crop = None + if resize_mode == "fill": + first = image[0] if isinstance(image, list) else image + if not isinstance(first, PIL.Image.Image): + raise ValueError(f"Only PIL image input is supported for resize_mode {resize_mode}") + height, width = self.get_default_height_width(first, height, width) + reference_crop = self._get_fill_region(first, height, width) + image = super().preprocess( + image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords + ) + 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) diff --git a/src/diffusers/modular_pipelines/flux3/inputs.py b/src/diffusers/modular_pipelines/flux3/inputs.py index f4ab66553a6e..bc28e12d4740 100644 --- a/src/diffusers/modular_pipelines/flux3/inputs.py +++ b/src/diffusers/modular_pipelines/flux3/inputs.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import math import torch @@ -243,19 +242,12 @@ def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> Pi class Flux3ImageCondInputStep(ModularPipelineBlocks): model_name = "flux3" - def __init__(self, default_height: int = 1024, default_width: int = 1024): - # the generation size defaults to this area at the first reference image's aspect ratio - self.default_height = default_height - self.default_width = default_width - super().__init__() - @property def description(self) -> str: return ( "Step that expands the reference image latents to the final batch size and sets the generation " - f"`height`/`width` when not provided: {self.default_height} x {self.default_width} pixels of area at the " - "first reference image's aspect ratio, on the 16-pixel grid (a given side overrides its default). The " - "reference images keep their own resolution, the generation size is independent of them." + "`height`/`width` when not provided: the first reference image's own size (its padded canvas for a " + "single reference, which the decode step crops back). A given size is used as is, without the crop-back." ) @property @@ -271,13 +263,19 @@ def inputs(self) -> list[InputParam]: "height", type_hint=int, description="The height in pixels of the generated image; defaults to the first reference image's " - "aspect ratio at the default area", + "own height", ), InputParam( "width", type_hint=int, description="The width in pixels of the generated image; defaults to the first reference image's " - "aspect ratio at the default area", + "own width", + ), + InputParam( + "reference_crop", + type_hint=list[tuple[int, int, int, int] | None] | None, + description="Per reference, the box of the reference inside its padded canvas or None, from the " + "preprocess step", ), InputParam.template("num_images_per_prompt"), InputParam("batch_size", required=True, type_hint=int), @@ -293,20 +291,25 @@ def intermediate_outputs(self) -> list[OutputParam]: ), OutputParam("height", type_hint=int, description="The height in pixels of the generated image"), OutputParam("width", type_hint=int, description="The width in pixels of the generated image"), + OutputParam( + "reference_crop", + type_hint=list[tuple[int, int, int, int] | None] | None, + description="The crop-back boxes for the decode step; None when the generation size was given", + ), ] @torch.no_grad() def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - # the generation size defaults to the default area at the first reference's aspect ratio, rounded to a - # multiple of 16 pixels (one token); a given `height` or `width` overrides its side - if block_state.height is None or block_state.width is None: - multiple_of = components.vae_scale_factor * 2 - latent_height, latent_width = block_state.image_cond_latents[0].shape[-2:] - scale = math.sqrt(self.default_height * self.default_width / (latent_height * latent_width)) / multiple_of - block_state.height = block_state.height or max(multiple_of, round(latent_height * scale) * multiple_of) - block_state.width = block_state.width or max(multiple_of, round(latent_width * scale) * multiple_of) + # the generation size defaults to the first (prepared) reference's own size; a given size is used as is and + # disables the crop-back of a padded reference + multiple_of = components.vae_scale_factor * 2 + latent_height, latent_width = block_state.image_cond_latents[0].shape[-2:] + if block_state.height is not None or block_state.width is not None: + block_state.reference_crop = None + block_state.height = block_state.height or latent_height * multiple_of + block_state.width = block_state.width or latent_width * multiple_of block_state.image_cond_latents = [ repeat_tensor_to_batch_size( diff --git a/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py b/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py index ddaf7f485c0a..31dcc62f7b0d 100644 --- a/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py +++ b/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py @@ -47,12 +47,16 @@ class Flux3VaeEncoderSequentialStep(SequentialPipelineBlocks): Inputs: image (`Image | list`): - The reference image or images to condition the generation on; each keeps its own resolution - (center-cropped to the 16-pixel grid) + 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 Outputs: condition_images (`list`): The preprocessed reference images, one tensor of shape (1, 3, height_i, width_i) in [-1, 1] per image + reference_crop (`list`): + 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 image_cond_latents (`list`): Normalized packed latents of the reference images, one tensor of shape (1, 128, height_i//16, width_i//16) per image @@ -79,12 +83,16 @@ class Flux3AutoVaeEncoderStep(AutoPipelineBlocks): Inputs: image (`Image | list`, *optional*): - The reference image or images to condition the generation on; each keeps its own resolution - (center-cropped to the 16-pixel grid) + 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 Outputs: condition_images (`list`): The preprocessed reference images, one tensor of shape (1, 3, height_i, width_i) in [-1, 1] per image + reference_crop (`list`): + 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 image_cond_latents (`list`): Normalized packed latents of the reference images, one tensor of shape (1, 128, height_i//16, width_i//16) per image @@ -209,11 +217,11 @@ class Flux3ImageConditionedCoreDenoiseStep(SequentialPipelineBlocks): Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16) per image height (`int`, *optional*): - The height in pixels of the generated image; defaults to the first reference image's aspect ratio at the - default area + The height in pixels of the generated image; defaults to the first reference image's own height width (`int`, *optional*): - The width in pixels of the generated image; defaults to the first reference image's aspect ratio at the - default area + The width in pixels of the generated image; defaults to the first reference image's own width + reference_crop (`list | NoneType`, *optional*): + Per reference, the box of the reference inside its padded canvas or None, from the preprocess step latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. generator (`Generator`, *optional*): @@ -273,11 +281,11 @@ class Flux3AutoCoreDenoiseStep(AutoPipelineBlocks): Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16) per image height (`int`, *optional*, defaults to None or 1024, depending on the workflow): - The height in pixels of the generated image; defaults to the first reference image's aspect ratio at the - default area + The height in pixels of the generated image; defaults to the first reference image's own height width (`int`, *optional*, defaults to None or 1024, depending on the workflow): - The width in pixels of the generated image; defaults to the first reference image's aspect ratio at the - default area + The width in pixels of the generated image; defaults to the first reference image's own width + reference_crop (`list | NoneType`, *optional*): + Per reference, the box of the reference inside its padded canvas or None, from the preprocess step latents (`Tensor`): Pre-generated noisy latents for image generation. generator (`Generator`, *optional*): @@ -366,11 +374,11 @@ class Flux3AutoBlocks(SequentialPipelineBlocks): Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16) per image height (`int`, *optional*, defaults to None or 1024, depending on the workflow): - The height in pixels of the generated image; defaults to the first reference image's aspect ratio at the - default area + The height in pixels of the generated image; defaults to the first reference image's own height width (`int`, *optional*, defaults to None or 1024, depending on the workflow): - The width in pixels of the generated image; defaults to the first reference image's aspect ratio at the - default area + The width in pixels of the generated image; defaults to the first reference image's own width + reference_crop (`list | NoneType`, *optional*): + Per reference, the box of the reference inside its padded canvas or None, from the preprocess step latents (`Tensor`): Pre-generated noisy latents for image generation. generator (`Generator`, *optional*): diff --git a/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py b/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py index 307ef2d70454..d23f894c6df7 100644 --- a/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py +++ b/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py @@ -140,11 +140,11 @@ class Flux3DistilledImageConditionedCoreDenoiseStep(SequentialPipelineBlocks): Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16) per image height (`int`, *optional*): - The height in pixels of the generated image; defaults to the first reference image's aspect ratio at the - default area + The height in pixels of the generated image; defaults to the first reference image's own height width (`int`, *optional*): - The width in pixels of the generated image; defaults to the first reference image's aspect ratio at the - default area + The width in pixels of the generated image; defaults to the first reference image's own width + reference_crop (`list | NoneType`, *optional*): + Per reference, the box of the reference inside its padded canvas or None, from the preprocess step latents (`Tensor`, *optional*): Pre-generated noisy latents for image generation. generator (`Generator`, *optional*): @@ -201,11 +201,11 @@ class Flux3DistilledAutoCoreDenoiseStep(AutoPipelineBlocks): Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16) per image height (`int`, *optional*, defaults to None or 1024, depending on the workflow): - The height in pixels of the generated image; defaults to the first reference image's aspect ratio at the - default area + The height in pixels of the generated image; defaults to the first reference image's own height width (`int`, *optional*, defaults to None or 1024, depending on the workflow): - The width in pixels of the generated image; defaults to the first reference image's aspect ratio at the - default area + The width in pixels of the generated image; defaults to the first reference image's own width + reference_crop (`list | NoneType`, *optional*): + Per reference, the box of the reference inside its padded canvas or None, from the preprocess step latents (`Tensor`): Pre-generated noisy latents for image generation. generator (`Generator`, *optional*): @@ -292,11 +292,11 @@ class Flux3DistilledAutoBlocks(SequentialPipelineBlocks): Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16) per image height (`int`, *optional*, defaults to None or 1024, depending on the workflow): - The height in pixels of the generated image; defaults to the first reference image's aspect ratio at the - default area + The height in pixels of the generated image; defaults to the first reference image's own height width (`int`, *optional*, defaults to None or 1024, depending on the workflow): - The width in pixels of the generated image; defaults to the first reference image's aspect ratio at the - default area + The width in pixels of the generated image; defaults to the first reference image's own width + reference_crop (`list | NoneType`, *optional*): + Per reference, the box of the reference inside its padded canvas or None, from the preprocess step latents (`Tensor`): Pre-generated noisy latents for image generation. generator (`Generator`, *optional*): diff --git a/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py b/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py index f024128022e8..e5b6c07b117c 100644 --- a/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py +++ b/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py @@ -26,7 +26,6 @@ Flux3ModularPipeline, ) from diffusers.modular_pipelines.flux3.encoders import Flux3TextEncoderStep -from diffusers.modular_pipelines.flux3.inputs import Flux3ImageCondInputStep from diffusers.utils import logging from ...testing_utils import CaptureLogger, torch_device @@ -104,7 +103,8 @@ "sigmas": None, "output_type": "pil", } -FLUX3_TEXT2IMAGE_INPUTS = {**FLUX3_INPUTS, "height": 1024, "width": 1024} +# the decode step's crop-back box is produced by the image branch; text2image exposes it as an optional input +FLUX3_TEXT2IMAGE_INPUTS = {**FLUX3_INPUTS, "height": 1024, "width": 1024, "reference_crop": None} FLUX3_IMAGE_CONDITIONED_INPUTS = {**FLUX3_INPUTS, "height": None, "width": None} FLUX3_WORKFLOW_DEFAULTS = { "text2image": { @@ -162,7 +162,7 @@ "text2image": { "components": FLUX3_DISTILLED_COMPONENTS, "required_inputs": ["prompt"], - "inputs": {**FLUX3_DISTILLED_INPUTS, "height": 1024, "width": 1024}, + "inputs": {**FLUX3_DISTILLED_INPUTS, "height": 1024, "width": 1024, "reference_crop": None}, "component_configs": FLUX3_DISTILLED_SAMPLING, }, "image_conditioned": { @@ -250,21 +250,30 @@ def test_prompt_upsampling(self): assert not torch.allclose(out["images"], direct["images"]) def test_image_to_image_default_size(self): - # without a size, image-to-image generates at the default area (1024 x 1024 for the real model) at the first - # reference's aspect ratio, rounded to the token grid; a given side overrides its default. The default area is - # shrunk to 64 x 64 for the tiny model. - blocks = self.pipeline_blocks_class().get_workflow("image_conditioned") - blocks.sub_blocks["denoise.cond_input"] = Flux3ImageCondInputStep(default_height=64, default_width=64) - pipe = blocks.init_pipeline(self.pretrained_model_name_or_path) - pipe.load_components(dtype=torch.float32) - pipe.to(torch_device) + # without a size, image-to-image generates at the first reference's own size; a given side overrides it. + # A single reference is padded to the image processor's 16-pixel grid and the image is cropped back + pipe = self.get_pipeline().to(torch_device) inputs = {k: v for k, v in self.get_dummy_inputs().items() if k not in ("height", "width")} inputs["num_inference_steps"] = 1 - inputs["image"] = self.get_reference_images(sizes=((64, 48),)) # 64 x 48 pixels, 4 pixels per token - # 64 x 64 of area at the 4:3 aspect ratio: 72 x 56 on the 4-pixel grid - assert pipe(**inputs, output="images").shape[-2:] == (56, 72) - assert pipe(**inputs, height=32, width=32, output="images").shape[-2:] == (32, 32) + assert pipe(**inputs, image=self.get_reference_images(sizes=((64, 48),)), output="images").shape[-2:] == ( + 48, + 64, + ) + # 62 x 46 is padded to the 64 x 48 canvas, generated there and cropped back to 62 x 46 + assert pipe(**inputs, image=self.get_reference_images(sizes=((62, 46),)), output="images").shape[-2:] == ( + 46, + 62, + ) + # a given size is used as is, without the crop-back + image = pipe( + **inputs, image=self.get_reference_images(sizes=((62, 46),)), height=32, width=32, output="images" + ) + assert image.shape[-2:] == (32, 32) + # several references: the first one's size, each cropped to the grid + assert pipe(**inputs, image=self.get_reference_images(sizes=((62, 46), (32, 32))), output="images").shape[ + -2: + ] == (32, 48) def test_text_window_pads_to_multiple(self): # the prompt is right-padded to the block's pad multiple (80 for the real model), capped at `max_sequence_length` From a87c6182be814f60c336dd77bcd35d38b8fd9b93 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Mon, 19 Oct 2026 01:18:04 +0200 Subject: [PATCH 2/3] FLUX.3 image processor: own the reference sizing rules `preprocess` takes only the image and `resize_mode`: "fill" downscales the reference so that its canvas fits under 4 megapixels and reflect-pads it to the canvas, "crop" fits it under 1 megapixel and center-crops it to the grid. The process-images block only picks the mode. Fill is a pure pad, like crop is a pure crop, and the crop-back box is read from the preprocessed tensor. The flags of the base processor are fixed; only the VAE scale factor and latent channels remain configurable. Co-Authored-By: Claude Fable 5.1 --- .../modular_pipelines/flux3/decoders.py | 8 +- .../modular_pipelines/flux3/encoders.py | 40 +----- .../flux3/image_processor.py | 133 ++++++++++-------- 3 files changed, 80 insertions(+), 101 deletions(-) diff --git a/src/diffusers/modular_pipelines/flux3/decoders.py b/src/diffusers/modular_pipelines/flux3/decoders.py index c0800b1ed914..ce30073bc544 100644 --- a/src/diffusers/modular_pipelines/flux3/decoders.py +++ b/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 diff --git a/src/diffusers/modular_pipelines/flux3/encoders.py b/src/diffusers/modular_pipelines/flux3/encoders.py index bbfc41a01b70..bd407a35be42 100644 --- a/src/diffusers/modular_pipelines/flux3/encoders.py +++ b/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 @@ -46,12 +45,7 @@ def description(self) -> str: @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 @@ -92,35 +86,11 @@ def __call__( 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)}") - image_processor = components.image_processor - multiple_of = image_processor.config.vae_scale_factor # one token: the VAE downsampling and the 2 x 2 packing - if len(images) == 1: - # one reference is kept whole: downscaled so that its canvas (the size rounded up to the token grid) fits - # under 4 megapixels, then reflect-padded to the canvas (`resize_mode="fill"` for this processor); the - # generation defaults to the canvas and the decode step crops back to the reference - image = images[0] - max_area, max_side, min_side = 4 * 1024 * 1024, 4096, 64 - scale = min(1.0, math.sqrt(max_area / (image.width * image.height)), max_side / max(image.size)) - while True: - reference_height, reference_width = ( - max(1, round(side * scale)) for side in (image.height, image.width) - ) - height, width = ( - min(max_side, max(min_side, -(-side // multiple_of) * multiple_of)) - for side in (reference_height, reference_width) - ) - if height * width <= max_area: - break - scale *= 0.995 - images = [image_processor.resize(image, height=reference_height, width=reference_width)] - preprocess_kwargs = {"height": height, "width": width, "resize_mode": "fill"} - else: - # several references are downscaled to fit under 1 megapixel each and center-cropped to the token grid - # (the default size rounds down to the grid; `resize_mode="crop"` is a pure crop for this processor) - images = [image_processor._resize_if_exceeds_area(image, 1024 * 1024) for image in images] - preprocess_kwargs = {"resize_mode": "crop"} + # 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( - *(image_processor.preprocess(image, **preprocess_kwargs) for image in images) + *(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] diff --git a/src/diffusers/modular_pipelines/flux3/image_processor.py b/src/diffusers/modular_pipelines/flux3/image_processor.py index 7c82875b71f2..891098a0fcd0 100644 --- a/src/diffusers/modular_pipelines/flux3/image_processor.py +++ b/src/diffusers/modular_pipelines/flux3/image_processor.py @@ -18,54 +18,36 @@ import PIL.Image import torch -from ...image_processor import PIL_INTERPOLATION, PipelineImageInput, VaeImageProcessor +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. `resize_mode="fill"` pads the image to the given size with its reflection, and `preprocess` also - returns the `reference_crop` box of the image inside the padded canvas, which `postprocess` takes to crop the - generated image back to it. + 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. 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 @@ -119,38 +101,68 @@ def _resize_and_crop( 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`]. + """ + 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 fits under `target_area` and its sides under + `max_side`. + """ + 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"`: downscaled to fit within the dimensions, maintaining the aspect ratio, and centered. + `resize_mode="fill"`: centered. """ image_width, image_height = image.size if image_width > width or image_height > height: - scale = min(width / image_width, height / image_height) - image_width, image_height = max(1, round(image_width * scale)), max(1, round(image_height * scale)) + 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""" - Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center - the image within the dimensions, filling the rest with its reflection. + 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 resize and fill. + The image to fill. width (`int`): - The width to resize the image to. + The width to pad the image to. height (`int`): - The height to resize the image to. + The height to pad the image to. Returns: `PIL.Image.Image`: - The resized and filled image. + The filled image. """ left, top, right, bottom = self._get_fill_region(image, height, width) - if (right - left, bottom - top) != image.size: - image = image.resize((right - left, bottom - top), resample=PIL_INTERPOLATION[self.config.resample]) if image.size == (width, height): return image @@ -161,14 +173,14 @@ def _resize_and_fill(self, image: PIL.Image.Image, width: int, height: int) -> P def preprocess( self, - image: PipelineImageInput, - height: int | None = None, - width: int | None = None, - resize_mode: str = "default", # "default", "fill", "crop" - crops_coords: tuple[int, int, int, int] | None = None, + 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 image input as [`~image_processor.VaeImageProcessor.preprocess`] does. + 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]`: @@ -176,16 +188,19 @@ def preprocess( (left, top, right, bottom) box of the (first) image inside its canvas when it was padded (`resize_mode="fill"`), else `None`. """ - reference_crop = 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": - first = image[0] if isinstance(image, list) else image - if not isinstance(first, PIL.Image.Image): - raise ValueError(f"Only PIL image input is supported for resize_mode {resize_mode}") - height, width = self.get_default_height_width(first, height, width) - reference_crop = self._get_fill_region(first, height, width) - image = super().preprocess( - image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords - ) + 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( From 48f1f5a0d9b1368b5b177136e9a7f7f53a4d0d30 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Mon, 19 Oct 2026 08:47:38 +0200 Subject: [PATCH 3/3] FLUX.3 image processor: reference the BFL sizing code Co-Authored-By: Claude Fable 5.1 --- .../modular_pipelines/flux3/image_processor.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/src/diffusers/modular_pipelines/flux3/image_processor.py b/src/diffusers/modular_pipelines/flux3/image_processor.py index 891098a0fcd0..f52df6243e32 100644 --- a/src/diffusers/modular_pipelines/flux3/image_processor.py +++ b/src/diffusers/modular_pipelines/flux3/image_processor.py @@ -28,7 +28,9 @@ class Flux3ImageProcessor(VaeImageProcessor): 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. + 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: vae_scale_factor (`int`, *optional*, defaults to `16`): @@ -105,7 +107,8 @@ 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 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( @@ -116,8 +119,12 @@ 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 fits under `target_area` and its sides under - `max_side`. + 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))