FLUX.3 prompt upsampler block

[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.