[yiyixuxu]

This PR prompt upsampler block into official pipeline

"""FLUX.3 prompt upsampler examples, as they appear in the docs: the pipeline runs the upsampler on `prompt` by default,
for text-to-image, an edit, and the two layout forms (boxes on a blank canvas; a box over a reference); the upsampler
on its own; and the pipeline without it, for a caption to encode as given. The requests, references and boxes are the
ones of BFL's examples (github.com/black-forest-labs/flux-3-image-klein).

Standalone: writes the captions and five PNGs under `out/`.

    python flux3_upsampler_examples.py
"""

import json
import os

import torch

from diffusers import ModularPipeline
from diffusers.utils import load_image


REPO = "black-forest-labs/flux-3-image-klein-base"  # or black-forest-labs/flux-3-image-klein
REVISION = "refs/pr/17"  # the index with the upsampler components, until the PR is merged (klein: refs/pr/22)
INPUTS = "https://raw.githubusercontent.com/black-forest-labs/flux-3-image-klein/main/examples/inputs"
OUT = "out"
os.makedirs(OUT, exist_ok=True)
captions = {}

# the pipeline runs the prompt upsampler on `prompt` first; the task follows the inputs, like the rest of the pipeline
pipe = ModularPipeline.from_pretrained(REPO, revision=REVISION)
pipe.load_components(dtype={"vae": torch.float32, "default": torch.bfloat16})  # the image VAE runs in fp32
pipe.to("cuda")

# text-to-image
out = pipe(
    prompt="A man walking with his dog, afternoon sun",
    generator=torch.Generator().manual_seed(7),
    output=["images", "prompt"],
)
out["images"][0].save(os.path.join(OUT, "t2i.png"))
captions["t2i"] = out["prompt"]  # the caption the image model received

# an edit: the reference goes to the upsampler (which writes the caption from it) and to the image model
dog = load_image(f"{INPUTS}/3.png")
out = pipe(
    prompt="put a red bandana around the dog's neck",
    image=dog,
    generator=torch.Generator().manual_seed(22),
    output=["images", "prompt"],
)
out["images"][0].save(os.path.join(OUT, "i2i.png"))
captions["i2i"] = out["prompt"]

# layout: boxes are [y_min, x_min, y_max, x_max] in 0-1; the caption comes back with the box table appended and
# the given boxes kept exactly
boxes = [
    {"id": "obj_1", "desc": "a red and yellow striped hot-air balloon", "bbox": [0.08, 0.06, 0.48, 0.38]},
    {"id": "obj_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,
    generator=torch.Generator().manual_seed(1),
    output=["images", "prompt"],
)
out["images"][0].save(os.path.join(OUT, "bbox_t2i.png"))
captions["bbox_t2i"] = out["prompt"]

# layout edit: a box over the reference places a new element ('from': 'target'); 'from': 'ref_image_0' with
# `src_bbox` and `tgt_bbox` would move an existing one
mug = load_image(f"{INPUTS}/2.png")
boxes = [{"id": "obj_1", "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,
    generator=torch.Generator().manual_seed(1),
    output=["images", "prompt"],
)
out["images"][0].save(os.path.join(OUT, "bbox_i2i.png"))
captions["bbox_i2i"] = out["prompt"]

# the upsampler on its own: its step of the pipeline, as a pipeline that writes captions (sharing the loaded components)
pup = pipe.blocks.sub_blocks["prompt_upsampler"].init_pipeline(REPO)
pup.update_components(prompt_upsampler=pipe.prompt_upsampler, prompt_upsampler_processor=pipe.prompt_upsampler_processor)
captions["standalone"] = pup(prompt="A man walking with his dog, afternoon sun", output="prompt")

# a caption to encode as given: the pipeline without the step
blocks = ModularPipeline.from_pretrained(REPO, revision=REVISION).blocks
blocks.sub_blocks.pop("prompt_upsampler")
plain = blocks.init_pipeline(REPO)
plain.load_components(dtype={"vae": torch.float32, "default": torch.bfloat16})
plain.to("cuda")
image = plain(prompt=captions["t2i"], generator=torch.Generator().manual_seed(7), output="images")[0]
image.save(os.path.join(OUT, "t2i_from_caption.png"))

json.dump(captions, open(os.path.join(OUT, "captions.json"), "w"), indent=1, ensure_ascii=False)
print("wrote", sorted(os.listdir(OUT)))

text to image

edit

layout

layout edit

text to image (from expanded prompt)

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[vladmandic]

prompt-upsampling may be strongly recommended, but this pr makes it mandatory as pipeline FAILS if prompt upsampler is not loaded. and making 9b upsampling model mandatory for model that is intended to run on lower-end gpus doesn't make sense. also, loader does not load upsampler automatically which makes upsampling mandatory, but must-be-loaded-by-user-separately - also not in line with any other model.

[yiyixuxu]

it's by default, but not mandatory

you can pop out the block if yon't want it, the doc covers the use case https://huggingface.co/docs/diffusers/main/en/api/pipelines/flux3#prompt-upsampling

blocks.sub_blocks.pop("prompt_upsampler")
pipe = blocks.init_pipeline("")

and it's gone, the block and its components

this is the user experience BFL team preferred. like we explained in the doc, the model needs to take full scene captions and gives noticeably worse results without the upsampler - that's something specific to how this model is trained and we don't think it make any sense to overlook that simply to be in line with other models

[yiyixuxu]

also, loader does not load upsampler automatically which makes upsampling mandatory, but must-be-loaded-by-user-separately

also this, not sure what you're talking about
it loads automatically, so feel free to open an issue with reproducible script if not for you