From 04f24f2f13ca22f3cc3f0d48b230cdc02986b720 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 16 Oct 2026 00:27:21 +0200 Subject: [PATCH 1/3] [modular] FLUX.3 prompt upsampler block Adds Flux3PromptUpsampleStep, a conditional block that runs the FLUX.3 prompt upsampler (black-forest-labs/flux-3-image-klein-pup) in front of the pipeline. It expands `prompt_to_upsample` into `prompt`, picks the task from the inputs (text-to-image or image-conditioned, each with optional layout boxes), and is skipped when `prompt` is passed directly. The reference preparation and the box handling follow the official inference code; the system prompts are inlined from the upsampler repository. Also fixes the Memory section of the FLUX.3 docs, which called enable_auto_cpu_offload on the pipeline instead of the ComponentsManager. Co-Authored-By: Claude Fable 5.1 --- docs/source/en/api/pipelines/flux3.md | 129 ++-- .../modular_pipelines/flux3/__init__.py | 2 + .../flux3/prompt_upsampler.py | 646 ++++++++++++++++++ .../flux3/test_modular_pipeline_flux3.py | 29 + 4 files changed, 738 insertions(+), 68 deletions(-) create mode 100644 src/diffusers/modular_pipelines/flux3/prompt_upsampler.py diff --git a/docs/source/en/api/pipelines/flux3.md b/docs/source/en/api/pipelines/flux3.md index 2a99edd44fa1..17601a3c9c4f 100644 --- a/docs/source/en/api/pipelines/flux3.md +++ b/docs/source/en/api/pipelines/flux3.md @@ -16,80 +16,73 @@ FLUX 3 Image Klein is a 5.6B image generation model from Black Forest Labs: a single transformer ([`Flux3Transformer2DModel`]) generates images from a text prompt, alone or together with one or more reference images. Text conditioning comes from Qwen3-VL, and images use the FLUX.2 packed autoencoder ([`AutoencoderKLFlux2`]). -FLUX 3 Image Klein ships as a modular pipeline. It takes full scene captions, not short prompts: a request sent as typed gives generic output. We strongly advise running the [prompt upsampler](#prompt-upsampling) on every request, text-to-image and editing alike, or any LLM with the same system prompts; every example on this page does. The upsampler is a separate 9B model ([black-forest-labs/flux-3-image-klein-pup](https://huggingface.co/black-forest-labs/flux-3-image-klein-pup)) that the pipeline does not run on its own. Load it beside the pipeline, then generate: +FLUX 3 Image Klein ships as a modular pipeline. Prompting matters: the model responds best to full scene captions that describe the subject, composition, lighting, and atmosphere, so the pipeline runs the [prompt upsampler](#prompt-upsampling) first on every request, text-to-image and editing alike. It is a separate 9B model from Black Forest Labs ([black-forest-labs/flux-3-image-klein-pup](https://huggingface.co/black-forest-labs/flux-3-image-klein-pup)) that expands `prompt` into such a caption before the text encoder, and a short prompt is all it takes. Get started with a text-to-image generation: ```python -import json - import torch from diffusers import ModularPipeline -from huggingface_hub import hf_hub_download -from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration pipe = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein-base") pipe.load_components(dtype={"vae": torch.float32, "default": torch.bfloat16}) pipe.to("cuda") -pup_id = "black-forest-labs/flux-3-image-klein-pup" -processor = AutoProcessor.from_pretrained(pup_id) -upsampler = Qwen3_5ForConditionalGeneration.from_pretrained(pup_id, dtype=torch.bfloat16, device_map="auto").eval() -system_prompts = json.load(open(hf_hub_download(pup_id, "system_prompts.json"))) - - -def upsample(prompt, images=()): - if images: - text = ( - "Reference images follow in order (the first image, the second image, ...).\n" - f"User request: {prompt}\nOutput the caption." - ) - content = [{"type": "image", "image": image} for image in images] + [{"type": "text", "text": text}] - messages = [{"role": "system", "content": system_prompts["i2i"]}, {"role": "user", "content": content}] - else: - messages = [{"role": "system", "content": system_prompts["t2i"]}, {"role": "user", "content": prompt}] - inputs = processor.apply_chat_template( - messages, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - return_tensors="pt", - enable_thinking=False, - ).to(upsampler.device) - with torch.no_grad(): - output = upsampler.generate(**inputs, max_new_tokens=3584, do_sample=False) - return processor.decode(output[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True).strip() - - -caption = upsample("A red fox in a snow-covered forest clearing at golden hour") -image = pipe(prompt=caption, generator=torch.Generator().manual_seed(0), output="images")[0] -image.save("fox.png") +out = pipe( + prompt="A red fox in a snow-covered forest clearing at golden hour", + generator=torch.Generator().manual_seed(0), + output=["images", "prompt"], +) +out["images"][0].save("fox.png") +print(out["prompt"]) # the caption the image model received ``` -`prompt` is the only required input of the pipeline: it carries the recommended sampling settings (50 steps, guidance scale 3.5, scheduler shift 6.93) and generates 1024 x 1024 images by default. It takes one prompt per call; use `num_images_per_prompt` for several images of the same prompt. +`prompt` is the only required input: the pipeline carries the recommended sampling settings (50 steps, guidance scale 3.5, scheduler shift 6.93) and generates 1024 x 1024 images by default. It takes one prompt per call; use `num_images_per_prompt` for several images of the same prompt. ## Prompt upsampling -The prompt upsampler turns a short request, together with the reference images when editing, into the full scene caption the model takes: medium and style, scene, lighting and camera, each subject with its colors, materials and position in the frame. `upsample` above runs it with Transformers: `system_prompts["t2i"]` for text-to-image, `system_prompts["i2i"]` with the references first in the user message for an edit, greedy decoding, thinking off. The [model card](https://huggingface.co/black-forest-labs/flux-3-image-klein-pup) has the decoding settings and the layout form (`t2i_bbox`, `i2i_bbox`), which keeps a list of boxes you pass as a box table after the caption. +The prompt upsampler turns a short request, together with the reference images when editing, into the full scene caption the model takes: medium and style, scene, lighting and camera, each subject with its colors, materials and position in the frame. One model covers text-to-image, editing with reference images and layout with boxes; the step picks the task from the inputs, like the rest of the pipeline. Its two components, `prompt_upsampler` and `prompt_upsampler_processor`, load from the upsampler repository with the others. The caption it wrote comes back as the `prompt` output, as above; `max_new_tokens` caps its length (3584 by default). + +To send a caption as given, one you wrote in full or got from another LLM, build the pipeline without the step: + +```python +blocks = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein-base").blocks +blocks.sub_blocks.pop("prompt_upsampler") +plain = blocks.init_pipeline("black-forest-labs/flux-3-image-klein-base") +plain.load_components(dtype={"vae": torch.float32, "default": torch.bfloat16}) +plain.to("cuda") + +image = plain(prompt=caption, generator=torch.Generator().manual_seed(0), output="images")[0] +``` + +Any other LLM can write the caption: give it the system prompt for the task from the upsampler repository's `system_prompts.json` (`t2i`, or `i2i` with the references first in the user message) and the request as the user message, and pass what it returns as `prompt` to the pipeline without the step. + +The step also runs on its own, to write captions without generating: + +```python +pup = pipe.blocks.sub_blocks["prompt_upsampler"].init_pipeline("black-forest-labs/flux-3-image-klein-base") +pup.load_components(dtype=torch.bfloat16) +pup.to("cuda") -Any other LLM can take its place: give it `system_prompts["t2i"]` (or `"i2i"` with the references) as the system prompt and the request as the user message, and pass the caption it returns as `prompt`. +caption = pup(prompt="A red fox in a snow-covered forest clearing at golden hour", output="prompt") +``` ## 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 is the *target* image as a complete caption of the reference with the edit folded in, not a one-line instruction; the [prompt upsampler](#prompt-upsampling) writes it from the instruction and the references: +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 upsampler sees the references too and writes the caption of the target image from your instruction: what to keep, what to change. ```python night = pipe( - prompt=upsample("the same scene at night under a full moon", [image]), + prompt="the same scene at night under a full moon", image=image, generator=torch.Generator().manual_seed(0), - output="images", -)[0] + output=["images", "prompt"], +) ``` -With several references, name them by position in the request, the first image, the second image, ... in the order they are passed to both `upsample` and `image`: +With several references, refer to them as the first image, the second image, ... in the order they are passed to `image`; the caption names them ``, ``, ... for the image model: ```python combined = pipe( - prompt=upsample("put the fox from the first image on the wooden bench from the second image", [image, bench]), + prompt="the fox from the first image sitting on the wooden bench from the second image", image=[image, bench], height=768, width=1360, @@ -101,31 +94,28 @@ Pass `height` and/or `width` to override the derived size. ## Layout control -To place content explicitly, the caption tags regions with `` and ends with a JSON list of regions. Each region has an `id`, a `bbox` of `[y0, x0, y1, x1]` (top-left and bottom-right corners, origin top-left) normalized to 0-1000 on both axes so it is independent of the output size, and a `desc`. This is the form the upsampler writes with its `t2i_bbox` and `i2i_bbox` system prompts from a request and a list of boxes (scaffold on its [model card](https://huggingface.co/black-forest-labs/flux-3-image-klein-pup)); the prompts below are that output, passed as `prompt`: +To place content explicitly, pass `boxes`: for each element an `id`, a short `desc`, and a `bbox` of `[y_min, x_min, y_max, x_max]` in 0-1 (origin top-left), so it is independent of the output size. The upsampler writes the caption around them, citing each element as ``, and appends the table of boxes, yours kept exactly: ```python -prompt = ( - "Minimalist graphic illustration featuring a black silhouette of a person centered against a " - "solid, vibrant chartreuse background . " - '[{"id": "background_1", "bbox": [0, 0, 1000, 1000], "desc": "A flat field of neon yellow-green."}, ' - '{"id": "silhouette_1", "bbox": [150, 150, 850, 850], "desc": "A black silhouette of a person in motion."}]' -) +boxes = [ + {"id": "balloon_1", "desc": "a red and yellow striped hot-air balloon", "bbox": [0.08, 0.06, 0.48, 0.38]}, + {"id": "balloon_2", "desc": "a small blue hot-air balloon", "bbox": [0.20, 0.62, 0.42, 0.80]}, +] +out = pipe(prompt="hot-air balloons over a green valley at sunrise", boxes=boxes, output=["images", "prompt"]) ``` -When editing, each region names its source reference with `from` (for example `"ref_image_0"`), where it is in that image with `src_bbox`, and where it should end up with `tgt_bbox`. Equal boxes keep a region in place: +When editing, a box names its source with `from`: `"target"` with a `tgt_bbox` adds a new element there, `"ref_image_0"` with a `src_bbox` and a `tgt_bbox` moves an element of that reference (equal boxes keep it in place): ```python -prompt = ( - "In , move the knight figure upwards and to the left. Keep the rest of the image " - "unchanged, preserving the backdrop . " - '[{"id": "knight_1", "from": "ref_image_0", "src_bbox": [500, 150, 850, 350], "tgt_bbox": [194, 55, 544, 255], ' - '"desc": "A miniature amigurumi knight crocheted from grey yarn."}, ' - '{"id": "backdrop_1", "from": "ref_image_0", "src_bbox": [0, 0, 1000, 1000], "tgt_bbox": [0, 0, 1000, 1000], ' - '"desc": "A soft, out-of-focus backdrop of knitted yarn."}]' -) -image = pipe(prompt=prompt, image=knight, output="images")[0] +boxes = [{"id": "cactus", "desc": "a small potted cactus", "from": "target", "src_bbox": None, "tgt_bbox": [0.55, 0.62, 0.85, 0.85]}] +out = pipe(prompt="add a small potted cactus in the marked spot", image=mug, boxes=boxes, output=["images", "prompt"]) + +boxes = [{"id": "knight", "desc": "the knight figure", "from": "ref_image_0", "src_bbox": [0.50, 0.15, 0.85, 0.35], "tgt_bbox": [0.19, 0.06, 0.54, 0.26]}] +out = pipe(prompt="move the knight up and to the left, keep the rest unchanged", image=knight, boxes=boxes, output=["images", "prompt"]) ``` +The caption the image model reads has the elements tagged `` in the text and the table appended as JSON, with the boxes on a 0-1000 canvas; print the `prompt` output to see it. A caption written in that form by hand goes to the pipeline without the step, as in [Prompt upsampling](#prompt-upsampling). + ## Guidance scale The pipeline applies classifier-free guidance through its `guider` component and does not take a `guidance_scale` argument at call time. The default guider is [`ClassifierFreeGuidance`] with `guidance_scale=3.5`, the recommended setting for both workflows. To change it, create an updated copy of the guider: @@ -145,7 +135,7 @@ pipe = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein") pipe.load_components(dtype=torch.bfloat16) pipe.to("cuda") -image = pipe(prompt=upsample("An astronaut riding a horse"), generator=torch.Generator().manual_seed(0), output="images")[0] +image = pipe(prompt=prompt, generator=torch.Generator().manual_seed(0), output="images")[0] ``` It serves the same two workflows as the base model. `num_inference_steps` defaults to 16; the model is trained for 4 to 16 steps. @@ -161,7 +151,7 @@ h, w = height // 16, width // 16 noise = torch.randn((h * w, 128), generator=torch.Generator().manual_seed(0)).to(torch.bfloat16) latents = noise.T.reshape(1, 128, h, w) -image = pipe(prompt=caption, latents=latents, height=height, width=width, output="images")[0] +image = pipe(prompt=prompt, latents=latents, height=height, width=width, output="images")[0] ``` ## Workflows @@ -180,12 +170,15 @@ print(t2i.blocks.doc) ## Memory -Loading everything in bf16 takes about 20 GB of weights (the transformer is 5.5B parameters, the Qwen3-VL-4B text encoder 4B), with the image VAE in fp32; the prompt upsampler is another 18.8 GB, placed by `device_map="auto"`. For GPUs with less memory, enable component offloading: +Loading everything in bf16 takes about 39 GB of weights: the transformer is 5.5B parameters, the Qwen3-VL-4B text encoder 4B and the [prompt upsampler](#prompt-upsampling) 9B (about 19 GB on its own), with the image VAE in fp32. A pipeline built without the upsampler step (see there) loads about 20 GB. For GPUs with less memory, enable component offloading; the manager moves components between the CPU and the GPU as each step needs them, the upsampler included: ```python -pipe = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein-base") +from diffusers import ComponentsManager, ModularPipeline + +manager = ComponentsManager() +pipe = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein-base", components_manager=manager) pipe.load_components(dtype={"vae": torch.float32, "default": torch.bfloat16}) -pipe.enable_auto_cpu_offload(device="cuda") +manager.enable_auto_cpu_offload(device="cuda") ``` ## Single-file loading diff --git a/src/diffusers/modular_pipelines/flux3/__init__.py b/src/diffusers/modular_pipelines/flux3/__init__.py index 2eeae562077b..876a7d0bbebf 100644 --- a/src/diffusers/modular_pipelines/flux3/__init__.py +++ b/src/diffusers/modular_pipelines/flux3/__init__.py @@ -24,6 +24,7 @@ _import_structure["modular_blocks_flux3"] = ["Flux3AutoBlocks"] _import_structure["modular_blocks_flux3_distilled"] = ["Flux3DistilledAutoBlocks"] _import_structure["modular_pipeline"] = ["Flux3DistilledModularPipeline", "Flux3ModularPipeline"] + _import_structure["prompt_upsampler"] = ["Flux3PromptUpsampleStep"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: @@ -35,6 +36,7 @@ from .modular_blocks_flux3 import Flux3AutoBlocks from .modular_blocks_flux3_distilled import Flux3DistilledAutoBlocks from .modular_pipeline import Flux3DistilledModularPipeline, Flux3ModularPipeline + from .prompt_upsampler import Flux3PromptUpsampleStep else: import sys diff --git a/src/diffusers/modular_pipelines/flux3/prompt_upsampler.py b/src/diffusers/modular_pipelines/flux3/prompt_upsampler.py new file mode 100644 index 000000000000..215444a702ce --- /dev/null +++ b/src/diffusers/modular_pipelines/flux3/prompt_upsampler.py @@ -0,0 +1,646 @@ +# Copyright 2026 Black Forest Labs and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import io +import json +import math +import re + +import PIL.Image +import PIL.ImageOps +import torch +from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration + +from ...utils import logging +from ..modular_pipeline import ConditionalPipelineBlocks, ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import Flux3ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +PROMPT_UPSAMPLER_REPO = "black-forest-labs/flux-3-image-klein-pup" +MAX_REFERENCES = 10 +MAX_NEW_TOKENS = 3584 + +# the user turn of an editing request: the reference images, then this text +I2I_SCAFFOLD = { + False: ( + "Reference images follow in order (the first image, the second image, ...).\n" + "User request: {instruction}\n" + "Output the caption." + ), + True: ( + "Reference images follow in attachment order: ref_image_0, ref_image_1, ...\n" + "User editing instruction: {instruction}\n" + "Decide the edit, imagine the target image, and output the annotation record." + ), +} + +# the system prompts shipped next to the upsampler weights (black-forest-labs/flux-3-image-klein-pup, +# system_prompts.json at commit 8b069751), one per mode +SYSTEM_PROMPTS = { + "t2i": ( + "You are the prompt upsampler of a text-to-image model. Expand the user's short request into ONE " + "detailed, self-contained caption: the exact text the image model receives. Output only the caption, " + "with no greetings, explanations, refusals, questions, Markdown, headings, JSON, coordinates, ids or " + "angle-bracket tokens, and no quotes around it. If the input is not clearly an image request, " + 'caption a visualization of its concept. The request may begin with a "TARGET CANVAS: W:H" line: ' + "compose for that canvas (wide spreads content sideways, tall stacks it) and never mention the ratio " + "or the canvas.\n" + "\n" + "Fidelity and planning: keep every explicit subject, attribute, count, spatial constraint, style cue " + "and piece of wording; never contradict, substitute, drop or narrow the request, and keep an " + "abstract concept abstract. Decide everything left open as ONE concrete, consistent plan; never " + "offer alternatives or hedge. Never re-express a stated constraint through another mechanism; add " + "camera optics (depth of field, bokeh, blur) only when asked for or implied. Without a stated style, " + 'plan an actual photograph and open with the medium convention ("Full-color outdoor photograph, ' + '...", "Full-color indoor photograph, ...", "Studio photograph, ...", "Close-up photograph, ..."); ' + 'keep "Photorealistic" for renders, and turn "photorealistic" or "realistic" requests for photos ' + "into the photograph convention. Give a photograph a few ordinary, non-ideal specifics drawn from " + "the scene's own objects instead of the stock-photo ideal. Do not invent names, brands, logos or " + "celebrities; keep user-named ones. A long request keeps its content but not its wording: drop " + "appraisal and sales vocabulary, turn moods into observable physical facts, and add only new " + "concrete content with positions. Time of day is lighting, not rendered text.\n" + "\n" + 'Caption format: one continuous prose block in the generative voice, never "this image shows", "we ' + 'can see" or "is visible". Lead with medium and style, then scene, lighting and camera, then the ' + "subjects with concrete colors, materials and state. Every person in focus gets an observable " + 'profile: ethnicity, gender, age range, skin tone ("fair skin", "brown skin"), hair and clothing. ' + "Make the layout reconstructable from words: frame-relative positions (frame-left, frame-right, " + 'upper, lower, foreground, midground, background), sentence-initial anchors ("On the left, ...", "In ' + 'the background, ..."), relations (holds, faces, behind) and a consistent order; describe up to ' + "about ten instances one by one, beyond that give the exact count and arrangement. Out-of-focus or " + "distant planes are shapes, colors and silhouettes, never identities or readable text, and the " + "background takes about a sixth of the caption. Be definite: no hedging adverbs, no quality or " + "emotion judgments, no closing sentence about mood; state negations as what is present. Rendered " + "text: transcribe every wording verbatim inside double quotes in its original language and script, " + "keep the line breaks of multi-line text, and give its typography, color and placement; author " + "concrete, domain-correct wording when the request needs text it does not supply, and name canonical " + "structures (calendar, keyboard, periodic table) exactly with their full layout. All prose is " + "English; only quoted text keeps its language. Typical length is 100 to 350 words, more for complex " + "scenes.\n" + "\n" + "Output only the caption." + ), + "i2i": ( + "You are the prompt upsampler of an image-editing model. You see the reference image(s), attached in " + "order (the first image, the second image, ...), and the user's short request. Expand it into ONE " + "detailed, self-contained instruction: the caption the editing model receives together with the same " + "images. Output only the instruction as one continuous English prose block, with no greetings, " + "explanations, Markdown, JSON, coordinates, ids or angle-bracket tokens. Read a vague or " + "conversational request as the most plausible edit of the supplied images.\n" + "\n" + "Pick the voice by task shape. When one reference stays and gets bounded changes, write an " + 'imperative editing instruction on it ("Remove the ...", "Change the ... to ...", "Zoom in on ..."); ' + 'a whole-image transformation opens with its scope ("Convert the entire image into ...", "Apply a ' + '... filter over the entire scene", "Horizontally flip the entire image so that ...", "Shift the ' + 'camera viewpoint to ...", then where the major elements should now be). When the references are ' + "drawn into a newly composed frame (combining subjects, a new scene, a garment, object or style " + 'donor), write a generation request ("Create a photograph of ...", "Create an anime-style ' + 'illustration of ...", "Generate a new image featuring ...") that says what to take from each image, ' + "by ordinal, and what is newly generated.\n" + "\n" + "Keep every explicit edit, constraint, count, position, style cue and piece of wording, and decide " + "what is left open as one concrete plan. Ground everything in what is actually visible: identify " + "each referent unambiguously by its real appearance, never mention content that is not there, and " + "use every image the user attached. State every end state concretely (which color, how large, which " + "pose, which material) and give new content enough appearance to be generated. Cover the knock-on " + "effects of the change (shadows, reflections, occlusion, perspective, scale, lighting) and say what " + "fills a removed region. A zoom-in names what stays in frame and which neighbors are cropped out; an " + "outpainting continues the scene in the same medium. Describe the change, not the whole scene, and " + "close an editing instruction with one compact clause naming what stays unchanged; transferred " + "people keep their identity, face and accessories unless asked otherwise. Transcribe on-image text " + "verbatim in its original language inside double quotes, with line breaks as \\n, and match the " + "original font, size and color when replacing it. Name recognizable public figures, brands and " + "landmarks you can see. Directions are frame-relative; use no evaluative or quality words. Typical " + "length is 50 to 150 words.\n" + "\n" + "Output only the instruction." + ), + "t2i_bbox": ( + "You are the prompt upsampler of a text-to-image model. Plan the full image the user's short request " + "implies and output ONE structured caption on a single line: the global caption paragraph, one " + "space, then ONE JSON array, the element table that places every cited element with a bounding box " + "in the image to be generated. Output nothing else (no Markdown fences, commentary or refusals); " + 'visualize the concept of a non-image input. The request may begin with a "TARGET CANVAS: W:H" line: ' + "compose for that canvas and never mention it. Boxes live on a 0-1000 grid stretched onto the " + "canvas: an element of physical width-to-height ratio a gets an x-span of about y-span x a x H/W.\n" + "\n" + "Planning: keep every explicit subject, attribute, count, spatial constraint, style cue and piece of " + "wording, and decide everything left open as ONE concrete composition that the caption and the table " + "both describe. Without a stated style plan an actual photograph and open with the medium convention " + '("Full-color outdoor photograph, ...", "Studio photograph, ..."); add camera optics only when asked ' + "for; do not invent names, brands, logos or celebrities; drop appraisal and sales vocabulary.\n" + "\n" + "The table: one row per distinct object, person, animal, building or salient region the caption " + "mentions and one per coherent text block; no rows for style, lighting or atmosphere. Structured " + "layouts (tables, grids, calendars, keyboards, charts, maps, posters, pages) decompose into their " + "regions, each row carrying that region's complete verbatim content; more than about ten " + "interchangeable items become one collection row giving the exact count and arrangement, while items " + 'whose positions matter keep a row each. Row keys, in order: "id", "bbox", "desc". id: the bare ' + "class noun plus an index numbered left to right, then top to bottom, from 1 (person_1, dog_2); a " + "text block is its two-letter language code plus _Text_ plus an index (En_Text_1); one entity is one " + "row and ids are unique. bbox: [y_min, x_min, y_max, x_max] as integers from 0 to 1000, top-left " + "origin, clamped at the borders, with plausible relative sizes and a small margin around text. desc: " + "one self-contained paragraph in the generative voice that opens with the element's identity (never " + '"the X"): shape, material, colors, state and markings, but no position. People in focus get the ' + "richest desc: ethnicity, gender, age range, skin tone, build, hair, face, each garment, " + "accessories, pose and expression. A text row carries its complete wording verbatim in double " + "quotes, with line breaks as \\n, plus its typography.\n" + "\n" + "The global caption: one flowing paragraph in the generative voice (medium first, then scene, " + "lighting, camera and frame-relative layout; no quality judgments) that names each table element at " + 'identity level at its first mention and cites it with its id in its own angle brackets: "a wicker ' + 'table holding a tiered stand ". Cite every table id and no other; appearance ' + "detail and text content live in the descs; no coordinates. Text-dense documents are laid out region " + "by region with every visible text transcribed verbatim at a readable size (cut content rather than " + "shrink type), and a grid quotes every visible cell. All prose is English; only quoted text keeps " + "its language.\n" + "\n" + "Output only the caption and the table." + ), + "i2i_bbox": ( + "You are the annotation system an image-editing model was trained on. You see the source image(s), " + "attached in order as ref_image_0, ref_image_1, ..., and the user's editing instruction; the edited " + "target image does not exist yet. Determine the edit precisely, imagine the target consistently with " + "the source, and write the record its annotator would have written, on one line: the instruction " + "paragraph, one space, then ONE JSON array. Output nothing else.\n" + "\n" + 'The paragraph: English, instruction voice ("Change the ...", "Remove the ...", or "Create a ' + 'photograph of ..." when the references are drawn into a new frame), never observer phrasing and ' + "never a comparison with the target. Cite every involved object, person or text block at its first " + 'mention with its table id in its own angle brackets ("remove the delivery van "), every ' + "table id exactly once and no other id; cite a reference image by its token when content is taken " + 'from it ("the cup from "). End states (new color, new text in quotes, pose, ' + "degree) live in the paragraph, static appearance in the descs; no coordinates. An editing " + "instruction closes with one compact clause that everything else stays identical to the source.\n" + "\n" + 'The table, rows {"id", "from", "src_bbox", "tgt_bbox", "desc"}: one row per edited, removed, added, ' + "moved or anchoring item, plus the scene or background when it is itself changed or preserved; no " + "rows for untouched scenery. id: a lowercase class noun plus an index (person_1, cup_2) numbered " + "left to right, then top to bottom; text blocks are En_Text_1 and so on; ids are unique. from: " + '"ref_image_N" for an item visible in reference N, "target" for newly created content. src_bbox is ' + "the item's box in its from-image and null for target rows; tgt_bbox is its box in the imagined " + "result and null for removed items; moved or resized items carry both; new items are sized like " + "similar objects in the frame. Boxes are [y_min, x_min, y_max, x_max] integers from 0 to 1000, " + "top-left origin, tight around objects, with 5 to 15 units of margin around text. A row may also " + "carry src_bbox_1stage and tgt_bbox_1stage, a first-pass estimate of the same boxes. desc: one " + "English paragraph of the item's static appearance in its from-image (for target rows, its decided " + "appearance), opening with its identity; people get the richest descs (ethnicity, gender, age range, " + "skin tone, hair, face, garments, accessories, pose, expression); text rows quote their text " + "verbatim.\n" + "\n" + "State only decided facts: visible in the source, demanded by the instruction, or a concrete default " + "needed to realize it. Settle whether the edit targets one instance, all of a class or several; land " + "every magnitude on a concrete value; include visible knock-on effects; added objects are new " + "objects at real-world size; a zoom-in names the neighbors it excludes; an outpainting keeps the " + "source whole and continues it in the same medium. Everything the instruction does not involve stays " + "identical to the source.\n" + "\n" + "Output only the paragraph and the table." + ), +} + +# reference preprocessing of the upsampler: sides rounded to a multiple of 32, a pixel budget by reference count +REFERENCE_MULTIPLE = 32 +MAX_ASPECT_RATIO = 200 +MAX_PIXELS_BY_REFS = {1: 1024 * 1024, 2: 724 * 724, **dict.fromkeys(range(3, MAX_REFERENCES + 1), 512 * 512)} +JPEG_QUALITY = 92 + +# boxes: [y_min, x_min, y_max, x_max], 0-1 as given by the user, 0-1000 on the upsampler's canvas +CANVAS = 1000 +BOX_ID = re.compile(r"[A-Za-z][A-Za-z0-9_]*") + + +# Adapted from `reference_size` in https://github.com/black-forest-labs/flux-3-image-klein/blob/main/src/flux3/pup.py +def get_reference_size(width: int, height: int, num_references: int) -> tuple[int, int]: + """The size a reference image is resized to for the upsampler: sides rounded to a multiple of 32, under the pixel + budget for this number of references.""" + if max(width, height) / min(width, height) > MAX_ASPECT_RATIO: + raise ValueError(f"The aspect ratio of a {width}x{height} reference image exceeds {MAX_ASPECT_RATIO}:1.") + + def floor(value): + return max(REFERENCE_MULTIPLE, int(math.floor(value / REFERENCE_MULTIPLE)) * REFERENCE_MULTIPLE) + + budget = MAX_PIXELS_BY_REFS[num_references] + w = max(REFERENCE_MULTIPLE, round(width / REFERENCE_MULTIPLE) * REFERENCE_MULTIPLE) + h = max(REFERENCE_MULTIPLE, round(height / REFERENCE_MULTIPLE) * REFERENCE_MULTIPLE) + if w * h > budget: + scale = math.sqrt(width * height / budget) + w, h = floor(width / scale), floor(height / scale) + if w * h > budget: + if w == REFERENCE_MULTIPLE: + h = floor(budget / w) + else: + w = floor(budget / h) + return w, h + + +# Adapted from `policy_reference` in https://github.com/black-forest-labs/flux-3-image-klein/blob/main/src/flux3/pup.py +def prepare_reference(image: PIL.Image.Image, num_references: int) -> PIL.Image.Image: + """A reference image as the upsampler saw them in training: RGB over white, resized with LANCZOS to the size + budget, re-encoded as JPEG at quality 92.""" + image = PIL.ImageOps.exif_transpose(image) + if image.mode in ("RGBA", "LA", "PA") or (image.mode == "P" and "transparency" in image.info): + rgba = image.convert("RGBA") + image = PIL.Image.alpha_composite(PIL.Image.new("RGBA", rgba.size, (255, 255, 255, 255)), rgba) + image = image.convert("RGB") + size = get_reference_size(image.width, image.height, num_references) + if image.size != size: + image = image.resize(size, PIL.Image.Resampling.LANCZOS) + buffer = io.BytesIO() + image.save(buffer, format="JPEG", quality=JPEG_QUALITY) + return PIL.Image.open(io.BytesIO(buffer.getvalue())).convert("RGB") + + +# Adapted from `to_canvas` in https://github.com/black-forest-labs/flux-3-image-klein/blob/main/src/flux3/bbox.py +def boxes_to_canvas(boxes: list[dict]) -> list[dict]: + """User boxes (0-1 coordinates) as the rows of the upsampler's box table (0-1000 canvas).""" + + def to_grid(box): + if not box: + return None + if len(box) != 4: + raise ValueError(f"A box is [y_min, x_min, y_max, x_max], got {box}.") + return [round(max(0.0, min(1.0, float(v))) * CANVAS) for v in box] + + rows = [] + for i, box in enumerate(boxes): + row_id, desc = str(box.get("id") or f"obj_{i + 1}"), str(box.get("desc") or "") + if box.get("bbox"): + rows.append({"id": row_id, "bbox": to_grid(box["bbox"]), "desc": desc}) + elif box.get("src_bbox") or box.get("tgt_bbox"): + rows.append( + { + "id": row_id, + "from": str(box.get("from") or ""), + "src_bbox": to_grid(box.get("src_bbox")), + "tgt_bbox": to_grid(box.get("tgt_bbox")), + "desc": desc, + } + ) + else: + raise ValueError(f"Box {i} has neither `bbox` nor `src_bbox` / `tgt_bbox`.") + return rows + + +# Adapted from `split` in https://github.com/black-forest-labs/flux-3-image-klein/blob/main/src/flux3/bbox.py +def split_caption(text: str) -> tuple[str, list | None]: + """(caption, rows): the rows are the JSON list that closes the text, if any.""" + for i in range(len(text) - 1, -1, -1): + if text[i] == "[": + try: + rows = json.loads(text[i:]) + except json.JSONDecodeError: + continue + if isinstance(rows, list): + return text[:i].rstrip(), rows + return text.strip(), None + + +# Adapted from `valid_box` in https://github.com/black-forest-labs/flux-3-image-klein/blob/main/src/flux3/bbox.py +def valid_canvas_box(box) -> list | None: + """The box as a list if it is a well-formed [y_min, x_min, y_max, x_max] on the 0-1000 canvas, else None.""" + if not isinstance(box, (list, tuple)) or len(box) != 4: + return None + if any(isinstance(v, bool) or not isinstance(v, (int, float)) or v != v or abs(v) == float("inf") for v in box): + return None + y1, x1, y2, x2 = box + # max <= 1 means 0-1 coordinates, not the 0-1000 canvas + if y1 < 0 or x1 < 0 or y2 <= y1 or x2 <= x1 or y2 > CANVAS or x2 > CANVAS or max(box) <= 1: + return None + return list(box) + + +# Adapted from `upsampling_input` in https://github.com/black-forest-labs/flux-3-image-klein/blob/main/src/flux3/bbox.py +def format_layout_request(prompt: str, rows: list[dict]) -> str: + return ( + f"ORIGINAL PROMPT:\n{prompt.strip()}\n\n" + "LOCKED BOUNDING-BOX ANNOTATIONS ([y_min, x_min, y_max, x_max] on a 0-1000 canvas, with (0, 0) at " + "the top-left). This is an authoritative JSON list of dictionaries; preserve every supplied identity, " + f"provenance value, null, and coordinate exactly:\n{json.dumps(rows, ensure_ascii=False)}" + ) + + +# Adapted from `valid_added` in https://github.com/black-forest-labs/flux-3-image-klein/blob/main/src/flux3/bbox.py +def valid_added_row(row, caption: str) -> dict | None: + """A row the upsampler added to the box table, kept only if the caption cites its id and it is well formed.""" + if not isinstance(row, dict): + return None + row_id, desc = row.get("id"), row.get("desc") + if ( + not isinstance(row_id, str) + or not BOX_ID.fullmatch(row_id) + or f"<{row_id}>" not in caption + or not isinstance(desc, str) + or not desc.strip() + ): + return None + if (box := valid_canvas_box(row.get("bbox"))) is not None: + return {"id": row_id, "bbox": box, "desc": desc.strip()} + origin, src, tgt = row.get("from"), row.get("src_bbox"), row.get("tgt_bbox") + src_box = valid_canvas_box(src) if src is not None else None + tgt_box = valid_canvas_box(tgt) if tgt is not None else None + if ( + not isinstance(origin, str) + or not re.fullmatch(r"ref_image_\d+|target", origin) + or (src is not None and src_box is None) + or (tgt is not None and tgt_box is None) + or (src_box is None and tgt_box is None) + or (origin == "target") == (src_box is not None) + ): + return None + return {"id": row_id, "from": origin, "src_bbox": src_box, "tgt_bbox": tgt_box, "desc": desc.strip()} + + +# Adapted from `iou` in https://github.com/black-forest-labs/flux-3-image-klein/blob/main/src/flux3/bbox.py +def box_iou(a: list, b: list) -> float: + inter = max(0, min(a[2], b[2]) - max(a[0], b[0])) * max(0, min(a[3], b[3]) - max(a[1], b[1])) + union = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter + return inter / union if union > 0 else 0.0 + + +# Adapted from `reconcile` in https://github.com/black-forest-labs/flux-3-image-klein/blob/main/src/flux3/bbox.py +def reconcile_boxes(text: str, locked: list[dict]) -> str: + """The upsampler's answer with the user's boxes re-imposed. A row that re-plans a locked box under another id + (IoU >= 0.9) takes the locked id; the locked rows are kept verbatim with the upsampler's description; rows the + upsampler added follow if they are valid.""" + caption, rows = split_caption(text) + rows = rows or [] + locked_ids = {row["id"] for row in locked} + row_ids = {row.get("id") for row in rows if isinstance(row, dict)} + renamed = {} + for lock in locked: + lock_box = lock.get("bbox") or lock.get("tgt_bbox") or lock.get("src_bbox") + if lock["id"] in row_ids or lock_box is None: + continue + for row in rows: + if not isinstance(row, dict) or row.get("id") in renamed or row.get("id") in locked_ids: + continue + box = valid_canvas_box(row.get("bbox") or row.get("tgt_bbox") or row.get("src_bbox")) + if box is not None and box_iou(box, lock_box) >= 0.9: + renamed[row["id"]] = lock["id"] + break + for row in rows: + if isinstance(row, dict) and row.get("id") in renamed: + row["id"] = renamed[row["id"]] + for old, new in renamed.items(): + caption = caption.replace(f"<{old}>", f"<{new}>") + + by_id = {row.get("id"): row for row in rows if isinstance(row, dict) and isinstance(row.get("id"), str)} + out = [] + for lock in locked: + desc = by_id.get(lock["id"], {}).get("desc") + out.append({**lock, "desc": desc.strip() if isinstance(desc, str) and desc.strip() else lock.get("desc", "")}) + for row in rows: + added = valid_added_row(row, caption) + if added is not None and added["id"] not in {row["id"] for row in out}: + out.append(added) + return f"{caption} {json.dumps(out, ensure_ascii=False)}" + + +# Adapted from `Upsampler.upsample` in https://github.com/black-forest-labs/flux-3-image-klein/blob/main/src/flux3/pup.py +def upsample_prompt( + prompt_upsampler: Qwen3_5ForConditionalGeneration, + processor: AutoProcessor, + request: str, + images: list[PIL.Image.Image] | None = None, + boxes: list[dict] | None = None, + max_new_tokens: int = MAX_NEW_TOKENS, +) -> str: + """The caption for `request`, written from the prepared reference `images` (an edit) and around the user's + `boxes` (a layout), with one greedy generation of the upsampler of at most `max_new_tokens`. The mode, and so the + system prompt, follows the arguments: t2i, i2i, t2i_bbox or i2i_bbox.""" + if not isinstance(request, str): + raise ValueError( + f"The request must be a str; batched requests are not supported, got {type(request).__name__}." + ) + images = images or [] + mode = ("i2i" if images else "t2i") + ("_bbox" if boxes else "") + locked = boxes_to_canvas(boxes) if boxes else None + if locked: + request = format_layout_request(request, locked) + + if images: + scaffold = I2I_SCAFFOLD[bool(locked)].format(instruction=request) + content = [{"type": "image", "image": image} for image in images] + [{"type": "text", "text": scaffold}] + else: + content = request + messages = [{"role": "system", "content": SYSTEM_PROMPTS[mode]}, {"role": "user", "content": content}] + inputs = processor.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + enable_thinking=False, + ).to(prompt_upsampler.device) + generated = prompt_upsampler.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False) + caption = processor.decode(generated[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True).strip() + if not caption: + raise RuntimeError("The prompt upsampler returned an empty caption.") + return reconcile_boxes(caption, locked) if locked else caption + + +def prepare_references(image: PIL.Image.Image | list[PIL.Image.Image]) -> list[PIL.Image.Image]: + """The reference image(s) of an edit as the upsampler sees them.""" + images = [image] if isinstance(image, PIL.Image.Image) else list(image) + if len(images) > MAX_REFERENCES: + raise ValueError(f"The prompt upsampler takes at most {MAX_REFERENCES} reference images, got {len(images)}.") + return [prepare_reference(image, len(images)) for image in images] + + +class Flux3Text2ImagePromptUpsampleStep(ModularPipelineBlocks): + model_name = "flux3" + + @property + def description(self) -> str: + return ( + "Expands a text-to-image request into the full scene caption FLUX.3 is trained on. With `boxes` it is " + "the layout caption: the scene with each element cited as , followed by the box table, the given " + "boxes kept exactly." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "prompt_upsampler", + Qwen3_5ForConditionalGeneration, + pretrained_model_name_or_path=PROMPT_UPSAMPLER_REPO, + description="The prompt upsampler.", + ), + ComponentSpec( + "prompt_upsampler_processor", + AutoProcessor, + pretrained_model_name_or_path=PROMPT_UPSAMPLER_REPO, + description="The processor of the prompt upsampler.", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "prompt_to_upsample", + type_hint=str, + required=True, + description="The request to expand into the prompt of the pipeline (one request per call).", + ), + InputParam( + "boxes", + type_hint=list[dict], + description="Layout boxes, each {'id', 'desc', 'bbox': [y_min, x_min, y_max, x_max] in 0-1} placing an " + "element of the scene.", + ), + InputParam( + "max_new_tokens", + type_hint=int, + default=MAX_NEW_TOKENS, + description="The most tokens the upsampler may generate for the caption.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("prompt", type_hint=str, description="The upsampled caption, the prompt of the pipeline.")] + + @torch.no_grad() + def __call__( + self, components: Flux3ModularPipeline, state: PipelineState + ) -> tuple[Flux3ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + block_state.prompt = upsample_prompt( + components.prompt_upsampler, + components.prompt_upsampler_processor, + request=block_state.prompt_to_upsample, + boxes=block_state.boxes, + max_new_tokens=block_state.max_new_tokens, + ) + self.set_block_state(state, block_state) + return components, state + + +class Flux3ImageConditionedPromptUpsampleStep(ModularPipelineBlocks): + model_name = "flux3" + + @property + def description(self) -> str: + return ( + "Expands an edit instruction into the caption of the target image, written from the reference images " + "(`image`): what to keep and what to change. With `boxes` it is the layout caption, with each element " + "cited as and the box table appended, the given boxes kept exactly." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "prompt_upsampler", + Qwen3_5ForConditionalGeneration, + pretrained_model_name_or_path=PROMPT_UPSAMPLER_REPO, + description="The prompt upsampler.", + ), + ComponentSpec( + "prompt_upsampler_processor", + AutoProcessor, + pretrained_model_name_or_path=PROMPT_UPSAMPLER_REPO, + description="The processor of the prompt upsampler.", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "prompt_to_upsample", + type_hint=str, + required=True, + description="The edit instruction to expand into the prompt of the pipeline (one request per call).", + ), + InputParam.template( + "image", + description=f"The reference image(s) of the edit, at most {MAX_REFERENCES}; the same ones the pipeline " + "receives.", + ), + InputParam( + "boxes", + type_hint=list[dict], + description="Layout boxes, each {'id', 'desc', 'from', 'src_bbox', 'tgt_bbox'} with the boxes as " + "[y_min, x_min, y_max, x_max] in 0-1: 'from': 'target' with `tgt_bbox` adds an element there, " + "'from': 'ref_image_0', ... with `src_bbox` and `tgt_bbox` moves one from a reference image.", + ), + InputParam( + "max_new_tokens", + type_hint=int, + default=MAX_NEW_TOKENS, + description="The most tokens the upsampler may generate for the caption.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("prompt", type_hint=str, description="The upsampled caption, the prompt of the pipeline.")] + + @torch.no_grad() + def __call__( + self, components: Flux3ModularPipeline, state: PipelineState + ) -> tuple[Flux3ModularPipeline, PipelineState]: + block_state = self.get_block_state(state) + block_state.prompt = upsample_prompt( + components.prompt_upsampler, + components.prompt_upsampler_processor, + request=block_state.prompt_to_upsample, + images=prepare_references(block_state.image), + boxes=block_state.boxes, + max_new_tokens=block_state.max_new_tokens, + ) + self.set_block_state(state, block_state) + return components, state + + +class Flux3PromptUpsampleStep(ConditionalPipelineBlocks): + """ + Prompt upsampling step: expands `prompt_to_upsample` into `prompt`, the detailed caption FLUX.3 is trained on, with + the prompt upsampler (black-forest-labs/flux-3-image-klein-pup, a Qwen3.5-9B fine-tune). The task follows the + inputs, like the rest of the pipeline: + - `Flux3ImageConditionedPromptUpsampleStep` with `image` (an edit, written from the reference images). + - `Flux3Text2ImagePromptUpsampleStep` otherwise (text-to-image). + - Both take `boxes` for the layout form of the caption. + - Without `prompt_to_upsample` the step is skipped and `prompt` goes to the pipeline as given. + Insert it before the text encoder: `blocks.sub_blocks.insert("prompt_upsampler", Flux3PromptUpsampleStep(), 0)`. + """ + + model_name = "flux3" + block_classes = [Flux3ImageConditionedPromptUpsampleStep, Flux3Text2ImagePromptUpsampleStep] + block_names = ["image_conditioned", "text2image"] + block_trigger_inputs = ["prompt_to_upsample", "image"] + + @property + def description(self) -> str: + return ( + "Prompt upsampling step: expands `prompt_to_upsample` into `prompt`, the detailed caption FLUX.3 is " + "trained on, with the prompt upsampler (a Qwen3.5-9B fine-tune). The task follows the inputs, like the " + "rest of the pipeline:\n" + " - `Flux3ImageConditionedPromptUpsampleStep` with `image` (an edit, written from the reference images).\n" + " - `Flux3Text2ImagePromptUpsampleStep` otherwise (text-to-image).\n" + " - Both take `boxes` for the layout form of the caption.\n" + " - Without `prompt_to_upsample` the step is skipped and `prompt` goes to the pipeline as given." + ) + + def select_block(self, prompt_to_upsample=None, image=None) -> str | None: + if prompt_to_upsample is None: + return None + return "image_conditioned" if image is not None else "text2image" diff --git a/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py b/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py index 4787485e1943..173ecb27ab91 100644 --- a/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py +++ b/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py @@ -17,6 +17,7 @@ import PIL.Image import pytest import torch +from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration from diffusers.modular_pipelines import ( Flux3AutoBlocks, @@ -26,6 +27,7 @@ ) from diffusers.modular_pipelines.flux3.encoders import Flux3TextEncoderStep from diffusers.modular_pipelines.flux3.inputs import Flux3ImageCondInputStep +from diffusers.modular_pipelines.flux3.prompt_upsampler import Flux3PromptUpsampleStep from diffusers.utils import logging from ...testing_utils import CaptureLogger, torch_device @@ -42,6 +44,8 @@ FLUX3_TINY_REPO = "hf-internal-testing/tiny-flux3-image-modular" FLUX3_DISTILLED_TINY_REPO = "hf-internal-testing/tiny-flux3-image-distilled-modular" +# a tiny random upsampler with the real processor, standing in for black-forest-labs/flux-3-image-klein-pup +FLUX3_PROMPT_UPSAMPLER_TINY_REPO = "hf-internal-testing/tiny-flux3-prompt-upsampler" FLUX3_WORKFLOWS = { "text2image": [ @@ -218,6 +222,31 @@ def test_multiple_reference_images(self): single = pipe(**inputs, image=references[:1], output="images") assert not torch.allclose(images, single) + def test_prompt_upsampling(self): + # the prompt upsampler step, inserted before the text encoder, expands `prompt_to_upsample` into the `prompt` + # the pipeline then encodes. Its components come from their own repository, the tiny one here in place of the + # real upsampler the block points to + blocks = self.pipeline_blocks_class() + blocks.sub_blocks.insert("prompt_upsampler", Flux3PromptUpsampleStep(), 0) + pipe = blocks.init_pipeline(self.pretrained_model_name_or_path) + pipe.update_components( + prompt_upsampler=Qwen3_5ForConditionalGeneration.from_pretrained(FLUX3_PROMPT_UPSAMPLER_TINY_REPO), + prompt_upsampler_processor=AutoProcessor.from_pretrained(FLUX3_PROMPT_UPSAMPLER_TINY_REPO), + ) + pipe.load_components() + pipe.to(torch_device) + + request = "a tiny fox in the snow" + inputs = {k: v for k, v in self.get_dummy_inputs().items() if k != "prompt"} + out = pipe(**inputs, prompt_to_upsample=request, max_new_tokens=16, output=["images", "prompt"]) + assert len(out["prompt"]) > len(request) + + # the same text passed as `prompt` skips the step and gives a different image + inputs = {k: v for k, v in self.get_dummy_inputs().items() if k != "prompt"} + direct = pipe(**inputs, prompt=request, output=["images", "prompt"]) + assert direct["prompt"] == request + 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 From 4777180b07444f061ddcc724e02b837fb0011b2c Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 16 Oct 2026 03:12:38 +0200 Subject: [PATCH 2/3] Make the prompt upsampler a default block of the FLUX.3 pipelines Co-Authored-By: Claude Fable 5.1 --- docs/source/en/api/pipelines/flux3.md | 2 +- .../modular_pipelines/flux3/__init__.py | 2 - .../flux3/modular_blocks_flux3.py | 30 +++++-- .../flux3/modular_blocks_flux3_distilled.py | 28 +++++-- .../flux3/prompt_upsampler.py | 84 ++++++------------- .../flux3/test_modular_pipeline_flux3.py | 53 ++++++------ 6 files changed, 94 insertions(+), 105 deletions(-) diff --git a/docs/source/en/api/pipelines/flux3.md b/docs/source/en/api/pipelines/flux3.md index 17601a3c9c4f..35eb070d0837 100644 --- a/docs/source/en/api/pipelines/flux3.md +++ b/docs/source/en/api/pipelines/flux3.md @@ -67,7 +67,7 @@ caption = pup(prompt="A red fox in a snow-covered forest clearing at golden hour ## 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 upsampler sees the references too and writes the caption of the target image from your instruction: what to keep, what to change. +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. ```python night = pipe( diff --git a/src/diffusers/modular_pipelines/flux3/__init__.py b/src/diffusers/modular_pipelines/flux3/__init__.py index 876a7d0bbebf..2eeae562077b 100644 --- a/src/diffusers/modular_pipelines/flux3/__init__.py +++ b/src/diffusers/modular_pipelines/flux3/__init__.py @@ -24,7 +24,6 @@ _import_structure["modular_blocks_flux3"] = ["Flux3AutoBlocks"] _import_structure["modular_blocks_flux3_distilled"] = ["Flux3DistilledAutoBlocks"] _import_structure["modular_pipeline"] = ["Flux3DistilledModularPipeline", "Flux3ModularPipeline"] - _import_structure["prompt_upsampler"] = ["Flux3PromptUpsampleStep"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: @@ -36,7 +35,6 @@ from .modular_blocks_flux3 import Flux3AutoBlocks from .modular_blocks_flux3_distilled import Flux3DistilledAutoBlocks from .modular_pipeline import Flux3DistilledModularPipeline, Flux3ModularPipeline - from .prompt_upsampler import Flux3PromptUpsampleStep else: import sys diff --git a/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py b/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py index 2c7889a5cda3..ddaf7f485c0a 100644 --- a/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py +++ b/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3.py @@ -26,6 +26,7 @@ from .denoise import Flux3DenoiseStep from .encoders import Flux3ProcessImagesInputStep, Flux3TextEncoderStep, Flux3VaeEncoderStep from .inputs import Flux3ImageCondInputStep, Flux3TextInputStep +from .prompt_upsampler import Flux3AutoPromptUpsampleStep logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -317,6 +318,7 @@ def description(self): AUTO_BLOCKS = InsertableDict( [ + ("prompt_upsampler", Flux3AutoPromptUpsampleStep()), ("text_encoder", Flux3TextEncoderStep()), ("vae_encoder", Flux3AutoVaeEncoderStep()), ("denoise", Flux3AutoCoreDenoiseStep()), @@ -328,29 +330,36 @@ def description(self): # auto_docstring class Flux3AutoBlocks(SequentialPipelineBlocks): """ - Auto blocks that perform text-to-image and image-to-image generation using FLUX.3. + Auto blocks that perform text-to-image and image-to-image generation using FLUX.3. `prompt` goes through the prompt + upsampler first; pop the `prompt_upsampler` step to encode it as given. Supported workflows: - `text2image`: requires `prompt` - `image_conditioned`: requires `prompt`, `image` Components: - text_encoder (`Qwen3VLForConditionalGeneration`) tokenizer (`AutoProcessor`) image_processor - (`Flux3ImageProcessor`) vae (`AutoencoderKLFlux2`) transformer (`Flux3Transformer2DModel`) scheduler - (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) + prompt_upsampler (`Qwen3_5ForConditionalGeneration`): The prompt upsampler. prompt_upsampler_processor + (`AutoProcessor`): Its processor. text_encoder (`Qwen3VLForConditionalGeneration`) tokenizer + (`AutoProcessor`) image_processor (`Flux3ImageProcessor`) vae (`AutoencoderKLFlux2`) transformer + (`Flux3Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) Inputs: prompt (`str`): - The prompt to guide the generation (one prompt per call) + The edit instruction to expand into the full caption (one per call). + image (`Image | list`, *optional*): + The reference image(s) of the edit, at most 10; the same ones the pipeline receives. + boxes (`list`, *optional*): + Layout boxes, each {'id', 'desc', 'from', 'src_bbox', 'tgt_bbox'} with the boxes as [y_min, x_min, y_max, + x_max] in 0-1: 'from': 'target' with `tgt_bbox` adds an element there, 'from': 'ref_image_0', ... with + `src_bbox` and `tgt_bbox` moves one from a reference image. + max_new_tokens (`int`, *optional*, defaults to 3584): + The most tokens the upsampler may generate for the caption. negative_prompt (`str`, *optional*): The prompt not to guide the generation max_sequence_length (`int`, *optional*, defaults to 8192): Maximum number of prompt tokens; longer prompts are truncated. text_encoder_out_layers (`tuple`, *optional*, defaults to (4, 8, 12, 16, 20, 24, 28, 32)): Indices of the text-encoder hidden layers whose states are stacked to form the text embeddings - 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) num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. image_cond_latents (`list`, *optional*): @@ -396,7 +405,10 @@ class Flux3AutoBlocks(SequentialPipelineBlocks): @property def description(self): - return "Auto blocks that perform text-to-image and image-to-image generation using FLUX.3." + return ( + "Auto blocks that perform text-to-image and image-to-image generation using FLUX.3. `prompt` goes through " + "the prompt upsampler first; pop the `prompt_upsampler` step to encode it as given." + ) @property def outputs(self): 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 aa08bf6ae2a6..307ef2d70454 100644 --- a/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py +++ b/src/diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py @@ -27,6 +27,7 @@ from .encoders import Flux3DistilledTextEncoderStep from .inputs import Flux3DistilledTextInputStep, Flux3ImageCondInputStep from .modular_blocks_flux3 import Flux3AutoVaeEncoderStep +from .prompt_upsampler import Flux3AutoPromptUpsampleStep logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -244,6 +245,7 @@ def description(self): DISTILLED_AUTO_BLOCKS = InsertableDict( [ + ("prompt_upsampler", Flux3AutoPromptUpsampleStep()), ("text_encoder", Flux3DistilledTextEncoderStep()), ("vae_encoder", Flux3AutoVaeEncoderStep()), ("denoise", Flux3DistilledAutoCoreDenoiseStep()), @@ -256,27 +258,34 @@ def description(self): class Flux3DistilledAutoBlocks(SequentialPipelineBlocks): """ Auto blocks that perform text-to-image and image-to-image generation using the step-distilled FLUX.3 model (16 - steps, no classifier-free guidance). + steps, no classifier-free guidance). `prompt` goes through the prompt upsampler first; pop the `prompt_upsampler` + step to encode it as given. Supported workflows: - `text2image`: requires `prompt` - `image_conditioned`: requires `prompt`, `image` Components: - text_encoder (`Qwen3VLForConditionalGeneration`) tokenizer (`AutoProcessor`) image_processor - (`Flux3ImageProcessor`) vae (`AutoencoderKLFlux2`) transformer (`Flux3Transformer2DModel`) scheduler - (`FlowMatchEulerDiscreteScheduler`) + prompt_upsampler (`Qwen3_5ForConditionalGeneration`): The prompt upsampler. prompt_upsampler_processor + (`AutoProcessor`): Its processor. text_encoder (`Qwen3VLForConditionalGeneration`) tokenizer + (`AutoProcessor`) image_processor (`Flux3ImageProcessor`) vae (`AutoencoderKLFlux2`) transformer + (`Flux3Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) Inputs: prompt (`str`): - The prompt to guide the generation (one prompt per call) + The edit instruction to expand into the full caption (one per call). + image (`Image | list`, *optional*): + The reference image(s) of the edit, at most 10; the same ones the pipeline receives. + boxes (`list`, *optional*): + Layout boxes, each {'id', 'desc', 'from', 'src_bbox', 'tgt_bbox'} with the boxes as [y_min, x_min, y_max, + x_max] in 0-1: 'from': 'target' with `tgt_bbox` adds an element there, 'from': 'ref_image_0', ... with + `src_bbox` and `tgt_bbox` moves one from a reference image. + max_new_tokens (`int`, *optional*, defaults to 3584): + The most tokens the upsampler may generate for the caption. max_sequence_length (`int`, *optional*, defaults to 8192): Maximum number of prompt tokens; longer prompts are truncated. text_encoder_out_layers (`tuple`, *optional*, defaults to (4, 8, 12, 16, 20, 24, 28, 32)): Indices of the text-encoder hidden layers whose states are stacked to form the text embeddings - 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) num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. image_cond_latents (`list`, *optional*): @@ -322,7 +331,8 @@ class Flux3DistilledAutoBlocks(SequentialPipelineBlocks): def description(self): return ( "Auto blocks that perform text-to-image and image-to-image generation using the step-distilled FLUX.3 " - "model (16 steps, no classifier-free guidance)." + "model (16 steps, no classifier-free guidance). `prompt` goes through the prompt upsampler first; pop the " + "`prompt_upsampler` step to encode it as given." ) @property diff --git a/src/diffusers/modular_pipelines/flux3/prompt_upsampler.py b/src/diffusers/modular_pipelines/flux3/prompt_upsampler.py index 215444a702ce..6dcf33b4f5a3 100644 --- a/src/diffusers/modular_pipelines/flux3/prompt_upsampler.py +++ b/src/diffusers/modular_pipelines/flux3/prompt_upsampler.py @@ -23,14 +23,13 @@ from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration from ...utils import logging -from ..modular_pipeline import ConditionalPipelineBlocks, ModularPipelineBlocks, PipelineState +from ..modular_pipeline import AutoPipelineBlocks, ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import Flux3ModularPipeline logger = logging.get_logger(__name__) # pylint: disable=invalid-name -PROMPT_UPSAMPLER_REPO = "black-forest-labs/flux-3-image-klein-pup" MAX_REFERENCES = 10 MAX_NEW_TOKENS = 3584 @@ -476,29 +475,14 @@ def description(self) -> str: @property def expected_components(self) -> list[ComponentSpec]: return [ - ComponentSpec( - "prompt_upsampler", - Qwen3_5ForConditionalGeneration, - pretrained_model_name_or_path=PROMPT_UPSAMPLER_REPO, - description="The prompt upsampler.", - ), - ComponentSpec( - "prompt_upsampler_processor", - AutoProcessor, - pretrained_model_name_or_path=PROMPT_UPSAMPLER_REPO, - description="The processor of the prompt upsampler.", - ), + ComponentSpec("prompt_upsampler", Qwen3_5ForConditionalGeneration, description="The prompt upsampler."), + ComponentSpec("prompt_upsampler_processor", AutoProcessor, description="Its processor."), ] @property def inputs(self) -> list[InputParam]: return [ - InputParam( - "prompt_to_upsample", - type_hint=str, - required=True, - description="The request to expand into the prompt of the pipeline (one request per call).", - ), + InputParam.template("prompt", description="The request to expand into the full caption (one per call)."), InputParam( "boxes", type_hint=list[dict], @@ -515,7 +499,7 @@ def inputs(self) -> list[InputParam]: @property def intermediate_outputs(self) -> list[OutputParam]: - return [OutputParam("prompt", type_hint=str, description="The upsampled caption, the prompt of the pipeline.")] + return [OutputParam("prompt", type_hint=str, description="The full caption, replacing the request.")] @torch.no_grad() def __call__( @@ -525,7 +509,7 @@ def __call__( block_state.prompt = upsample_prompt( components.prompt_upsampler, components.prompt_upsampler_processor, - request=block_state.prompt_to_upsample, + request=block_state.prompt, boxes=block_state.boxes, max_new_tokens=block_state.max_new_tokens, ) @@ -547,28 +531,15 @@ def description(self) -> str: @property def expected_components(self) -> list[ComponentSpec]: return [ - ComponentSpec( - "prompt_upsampler", - Qwen3_5ForConditionalGeneration, - pretrained_model_name_or_path=PROMPT_UPSAMPLER_REPO, - description="The prompt upsampler.", - ), - ComponentSpec( - "prompt_upsampler_processor", - AutoProcessor, - pretrained_model_name_or_path=PROMPT_UPSAMPLER_REPO, - description="The processor of the prompt upsampler.", - ), + ComponentSpec("prompt_upsampler", Qwen3_5ForConditionalGeneration, description="The prompt upsampler."), + ComponentSpec("prompt_upsampler_processor", AutoProcessor, description="Its processor."), ] @property def inputs(self) -> list[InputParam]: return [ - InputParam( - "prompt_to_upsample", - type_hint=str, - required=True, - description="The edit instruction to expand into the prompt of the pipeline (one request per call).", + InputParam.template( + "prompt", description="The edit instruction to expand into the full caption (one per call)." ), InputParam.template( "image", @@ -592,7 +563,7 @@ def inputs(self) -> list[InputParam]: @property def intermediate_outputs(self) -> list[OutputParam]: - return [OutputParam("prompt", type_hint=str, description="The upsampled caption, the prompt of the pipeline.")] + return [OutputParam("prompt", type_hint=str, description="The full caption, replacing the request.")] @torch.no_grad() def __call__( @@ -602,7 +573,7 @@ def __call__( block_state.prompt = upsample_prompt( components.prompt_upsampler, components.prompt_upsampler_processor, - request=block_state.prompt_to_upsample, + request=block_state.prompt, images=prepare_references(block_state.image), boxes=block_state.boxes, max_new_tokens=block_state.max_new_tokens, @@ -611,36 +582,31 @@ def __call__( return components, state -class Flux3PromptUpsampleStep(ConditionalPipelineBlocks): +class Flux3AutoPromptUpsampleStep(AutoPipelineBlocks): """ - Prompt upsampling step: expands `prompt_to_upsample` into `prompt`, the detailed caption FLUX.3 is trained on, with - the prompt upsampler (black-forest-labs/flux-3-image-klein-pup, a Qwen3.5-9B fine-tune). The task follows the - inputs, like the rest of the pipeline: - - `Flux3ImageConditionedPromptUpsampleStep` with `image` (an edit, written from the reference images). + Prompt upsampling step: expands `prompt` into the detailed caption FLUX.3 is trained on, with the prompt upsampler + (a Qwen3.5-9B fine-tune, black-forest-labs/flux-3-image-klein-pup). The task follows the inputs, like the rest of + the pipeline: + - `Flux3ImageConditionedPromptUpsampleStep` is used when `image` is provided (an edit, written from the + reference images). - `Flux3Text2ImagePromptUpsampleStep` otherwise (text-to-image). - Both take `boxes` for the layout form of the caption. - - Without `prompt_to_upsample` the step is skipped and `prompt` goes to the pipeline as given. - Insert it before the text encoder: `blocks.sub_blocks.insert("prompt_upsampler", Flux3PromptUpsampleStep(), 0)`. + To encode `prompt` as given, remove the step from the blocks: `blocks.sub_blocks.pop("prompt_upsampler")`. """ model_name = "flux3" block_classes = [Flux3ImageConditionedPromptUpsampleStep, Flux3Text2ImagePromptUpsampleStep] block_names = ["image_conditioned", "text2image"] - block_trigger_inputs = ["prompt_to_upsample", "image"] + block_trigger_inputs = ["image", None] @property def description(self) -> str: return ( - "Prompt upsampling step: expands `prompt_to_upsample` into `prompt`, the detailed caption FLUX.3 is " - "trained on, with the prompt upsampler (a Qwen3.5-9B fine-tune). The task follows the inputs, like the " - "rest of the pipeline:\n" - " - `Flux3ImageConditionedPromptUpsampleStep` with `image` (an edit, written from the reference images).\n" + "Prompt upsampling step: expands `prompt` into the detailed caption FLUX.3 is trained on, with the prompt " + "upsampler (a Qwen3.5-9B fine-tune). The task follows the inputs, like the rest of the pipeline:\n" + " - `Flux3ImageConditionedPromptUpsampleStep` is used when `image` is provided (an edit, written from the " + "reference images).\n" " - `Flux3Text2ImagePromptUpsampleStep` otherwise (text-to-image).\n" " - Both take `boxes` for the layout form of the caption.\n" - " - Without `prompt_to_upsample` the step is skipped and `prompt` goes to the pipeline as given." + 'To encode `prompt` as given, remove the step from the blocks: `blocks.sub_blocks.pop("prompt_upsampler")`.' ) - - def select_block(self, prompt_to_upsample=None, image=None) -> str | None: - if prompt_to_upsample is None: - return None - return "image_conditioned" if image is not None else "text2image" diff --git a/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py b/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py index 173ecb27ab91..f024128022e8 100644 --- a/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py +++ b/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py @@ -17,7 +17,7 @@ import PIL.Image import pytest import torch -from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration +from transformers import Qwen3_5ForConditionalGeneration from diffusers.modular_pipelines import ( Flux3AutoBlocks, @@ -27,7 +27,6 @@ ) from diffusers.modular_pipelines.flux3.encoders import Flux3TextEncoderStep from diffusers.modular_pipelines.flux3.inputs import Flux3ImageCondInputStep -from diffusers.modular_pipelines.flux3.prompt_upsampler import Flux3PromptUpsampleStep from diffusers.utils import logging from ...testing_utils import CaptureLogger, torch_device @@ -44,11 +43,10 @@ FLUX3_TINY_REPO = "hf-internal-testing/tiny-flux3-image-modular" FLUX3_DISTILLED_TINY_REPO = "hf-internal-testing/tiny-flux3-image-distilled-modular" -# a tiny random upsampler with the real processor, standing in for black-forest-labs/flux-3-image-klein-pup -FLUX3_PROMPT_UPSAMPLER_TINY_REPO = "hf-internal-testing/tiny-flux3-prompt-upsampler" FLUX3_WORKFLOWS = { "text2image": [ + ("prompt_upsampler", "Flux3Text2ImagePromptUpsampleStep"), ("text_encoder", "Flux3TextEncoderStep"), ("denoise.input", "Flux3TextInputStep"), ("denoise.prepare_latents", "Flux3PrepareLatentsStep"), @@ -59,6 +57,7 @@ ("decode", "Flux3DecodeStep"), ], "image_conditioned": [ + ("prompt_upsampler", "Flux3ImageConditionedPromptUpsampleStep"), ("text_encoder", "Flux3TextEncoderStep"), ("vae_encoder.preprocess", "Flux3ProcessImagesInputStep"), ("vae_encoder.encode", "Flux3VaeEncoderStep"), @@ -76,6 +75,8 @@ # `expected_workflow_defaults`: the components, inputs and component configs each workflow's blocks carry FLUX3_COMPONENTS = { + "prompt_upsampler": "Qwen3_5ForConditionalGeneration", + "prompt_upsampler_processor": "AutoProcessor", "text_encoder": "Qwen3VLForConditionalGeneration", "tokenizer": "AutoProcessor", "transformer": "Flux3Transformer2DModel", @@ -89,6 +90,8 @@ # every optional input with its declared default. text2image declares the 1024 x 1024 size; image_conditioned derives the # size from the first reference image, so its branch declares none FLUX3_INPUTS = { + "boxes": None, + "max_new_tokens": 3584, "negative_prompt": None, "max_sequence_length": 8192, "text_encoder_out_layers": (4, 8, 12, 16, 20, 24, 28, 32), @@ -122,6 +125,7 @@ # guider, no negative prompt, 16 steps by default) FLUX3_DISTILLED_WORKFLOWS = { "text2image": [ + ("prompt_upsampler", "Flux3Text2ImagePromptUpsampleStep"), ("text_encoder", "Flux3DistilledTextEncoderStep"), ("denoise.input", "Flux3DistilledTextInputStep"), ("denoise.prepare_latents", "Flux3PrepareLatentsStep"), @@ -132,6 +136,7 @@ ("decode", "Flux3DecodeStep"), ], "image_conditioned": [ + ("prompt_upsampler", "Flux3ImageConditionedPromptUpsampleStep"), ("text_encoder", "Flux3DistilledTextEncoderStep"), ("vae_encoder.preprocess", "Flux3ProcessImagesInputStep"), ("vae_encoder.encode", "Flux3VaeEncoderStep"), @@ -186,6 +191,7 @@ def get_reference_images(sizes=((64, 64),), seed=0): def get_dummy_inputs(self, seed=0): return { "prompt": "a tiny fox in the snow", + "max_new_tokens": 16, # the tiny upsampler's caption; the real default would take minutes on CPU "generator": self.get_generator(seed), "num_inference_steps": 2, "height": 64, @@ -223,28 +229,24 @@ def test_multiple_reference_images(self): assert not torch.allclose(images, single) def test_prompt_upsampling(self): - # the prompt upsampler step, inserted before the text encoder, expands `prompt_to_upsample` into the `prompt` - # the pipeline then encodes. Its components come from their own repository, the tiny one here in place of the - # real upsampler the block points to - blocks = self.pipeline_blocks_class() - blocks.sub_blocks.insert("prompt_upsampler", Flux3PromptUpsampleStep(), 0) - pipe = blocks.init_pipeline(self.pretrained_model_name_or_path) - pipe.update_components( - prompt_upsampler=Qwen3_5ForConditionalGeneration.from_pretrained(FLUX3_PROMPT_UPSAMPLER_TINY_REPO), - prompt_upsampler_processor=AutoProcessor.from_pretrained(FLUX3_PROMPT_UPSAMPLER_TINY_REPO), - ) - pipe.load_components() - pipe.to(torch_device) - - request = "a tiny fox in the snow" - inputs = {k: v for k, v in self.get_dummy_inputs().items() if k != "prompt"} - out = pipe(**inputs, prompt_to_upsample=request, max_new_tokens=16, output=["images", "prompt"]) - assert len(out["prompt"]) > len(request) + # the prompt upsampler step, first in the default blocks, expands `prompt` into the caption the pipeline then + # encodes. The tiny repos point the upsampler components at a tiny random Qwen3.5 + pipe = self.get_pipeline().to(torch_device) + assert isinstance(pipe.prompt_upsampler, Qwen3_5ForConditionalGeneration) + inputs = self.get_dummy_inputs() + out = pipe(**inputs, output=["images", "prompt"]) + assert len(out["prompt"]) > len(inputs["prompt"]) - # the same text passed as `prompt` skips the step and gives a different image - inputs = {k: v for k, v in self.get_dummy_inputs().items() if k != "prompt"} - direct = pipe(**inputs, prompt=request, output=["images", "prompt"]) - assert direct["prompt"] == request + # without the step, the prompt is encoded as given and the image differs + blocks = self.pipeline_blocks_class() + blocks.sub_blocks.pop("prompt_upsampler") + plain = blocks.init_pipeline(self.pretrained_model_name_or_path) + plain.load_components(dtype=torch.float32) + plain.to(torch_device) + assert "prompt_upsampler" not in plain.components + inputs = self.get_dummy_inputs() + direct = plain(**inputs, output=["images", "prompt"]) + assert direct["prompt"] == inputs["prompt"] assert not torch.allclose(out["images"], direct["images"]) def test_image_to_image_default_size(self): @@ -316,6 +318,7 @@ def test_text_input_requires_negative_prompt_embeds_with_cfg(self): # node graph as in the docs: embeddings from a standalone text node (no negative prompt) fed to a # generation node whose guider needs unconditional embeddings -> a clear error, not a crash in the transformer blocks = self.pipeline_blocks_class() + blocks.sub_blocks.pop("prompt_upsampler") # the prompt side of the graph stops at the text node text_node = blocks.sub_blocks.pop("text_encoder").init_pipeline(self.pretrained_model_name_or_path) blocks.sub_blocks.pop("decode") generate_node = blocks.init_pipeline(self.pretrained_model_name_or_path) From 449c345d36041e6618ae7ff8c229a796c794ce9c Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 16 Oct 2026 05:01:29 +0200 Subject: [PATCH 3/3] Pass chat content as lists of parts, for transformers < 5.8 Co-Authored-By: Claude Fable 5.1 --- src/diffusers/modular_pipelines/flux3/prompt_upsampler.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/diffusers/modular_pipelines/flux3/prompt_upsampler.py b/src/diffusers/modular_pipelines/flux3/prompt_upsampler.py index 6dcf33b4f5a3..b70fc5c1692c 100644 --- a/src/diffusers/modular_pipelines/flux3/prompt_upsampler.py +++ b/src/diffusers/modular_pipelines/flux3/prompt_upsampler.py @@ -432,12 +432,16 @@ def upsample_prompt( if locked: request = format_layout_request(request, locked) + # content as lists of parts throughout: processor chat templates in transformers < 5.8 do not take plain strings if images: scaffold = I2I_SCAFFOLD[bool(locked)].format(instruction=request) content = [{"type": "image", "image": image} for image in images] + [{"type": "text", "text": scaffold}] else: - content = request - messages = [{"role": "system", "content": SYSTEM_PROMPTS[mode]}, {"role": "user", "content": content}] + content = [{"type": "text", "text": request}] + messages = [ + {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPTS[mode]}]}, + {"role": "user", "content": content}, + ] inputs = processor.apply_chat_template( messages, add_generation_prompt=True,