arlaz/modular-flux123-multidiffusion

FLUX.1 / FLUX.2 / FLUX 3 MultiDiffusion Modular Blocks

Hugging Face Modular Diffusers blocks for MultiDiffusion with FLUX.1 dev and schnell, FLUX.2 dev and Klein, and FLUX 3 Image Klein (base and distilled): windowed denoising of large canvases, batched windows, regional prompts with masks, seamless panoramas, img2img and reference images.

Path	Contents
block.py	Hub remote-code entry point; it imports and re-exports the blocks of multidiffusion/
multidiffusion/geometry.py	Canvas layout, window grid and weights, reference crops, the window and panorama inputs
multidiffusion/window_loop.py	Per-step window preprocessing and the mixin that runs an upstream denoiser per window
multidiffusion/steps.py	Shared img2img and regional prompt split steps
multidiffusion/panorama.py	Circular-padding VAE decode
multidiffusion/composition.py	The top-level blocks base class and the helpers that edit a copy of the upstream blocks
multidiffusion/flux1.py, flux2.py, flux3.py	Each family's blocks, its builder and its top-level blocks
repository.py, cli.py	Family selection, Hub export helpers and the multidiff-modular command
model/<family>/	Generated metadata JSON for every family
examples/	Local and remote generation runners, see examples/README.md

Install

Python 3.14 (.python-version pins 3.14.10) with Diffusers 0.42.1:

uv sync --all-extras

or python -m pip install -e ".[dev,examples,quantization]". Extras:

• dev: pytest and ruff.
• examples: Pillow for images and masks, PEFT for --lora-path.
• quantization: TorchAO for --quantization.

export-hf copies block.py and multidiffusion/ from the checkout, so it needs a checkout (uv sync or pip install -e .), not a plain pip install ..

Remote Loading

The quickest test needs no checkout: it loads the blocks from the root block.py with trust_remote_code=True. You need diffusers>=0.42.1,<0.43, transformers>=5.2, torch, accelerate, and torchvision>=0.22 for FLUX 3, plus access to the gated upstream weights (accept their licenses on the Hub). The root default family is flux3-klein.

import torch
from diffusers import ComponentsManager, ModularPipelineBlocks

repo_id = "arlaz/modular-flux123-multidiffusion"

manager = ComponentsManager()
manager.enable_auto_cpu_offload(device="cuda")

blocks = ModularPipelineBlocks.from_pretrained(repo_id, trust_remote_code=True)
blocks.sub_blocks.pop("prompt_upsampler")  # optional: skip the 9B FLUX 3 prompt upsampler
pipe = blocks.init_pipeline(repo_id, components_manager=manager)
pipe.load_components(dtype=torch.bfloat16)

image = pipe(
    prompt="a dense renaissance fresco",
    height=1024,
    width=2048,
    height_generation=1024,
    width_generation=1024,
    window_stride_height=512,
    window_stride_width=512,
    generator=torch.Generator().manual_seed(42),
).images[0]
image.save("remote_smoke.png")

Load another family from the same block.py, then point it at its upstream model repo:

from diffusers.utils.dynamic_modules_utils import get_class_from_dynamic_module

blocks_cls = get_class_from_dynamic_module(
    repo_id, "block.py", class_name="Flux1MultiDiffusionAutoBlocks", trust_remote_code=True
)
pipe = blocks_cls().init_pipeline("black-forest-labs/FLUX.1-schnell", components_manager=manager)
pipe.load_components(dtype=torch.bfloat16)

The runner equivalent, examples/example_remote.py --repo-id ... [--family <family>], needs this checkout installed; see examples/README.md.

Local Python Use

From a checkout, import the blocks from the multidiffusion package (root block.py is only the Hub entry point), or build a family's pipeline on its upstream repo with repository.build_pipeline:

import torch
from diffusers import ComponentsManager

from repository import build_pipeline

manager = ComponentsManager()
manager.enable_auto_cpu_offload(device="cuda")

pipe = build_pipeline("flux2-klein-4b", components_manager=manager)
pipe.load_components(dtype=torch.bfloat16)

image = pipe(
    prompt="a dense renaissance fresco",
    height=2048,
    width=2048,
    height_generation=1024,
    width_generation=1024,
    window_stride_height=512,
    window_stride_width=512,
    num_inference_steps=4,
).images[0]

Pipeline Inputs

Input	Meaning
height, width	The output canvas in pixels.
height_generation, width_generation	The denoiser window; defaults to the full canvas.
window_stride_height, window_stride_width	Distance between windows; defaults to half the window.
window_stride_height_offset, window_stride_width_offset	Shift of the window grid added at every step.
weighting_type, weighting_range	Window blending: "cosine" (default), "linear" or "none"; the ramp length as a fraction of the window, defaulting to stride / window.
window_batch_size	Windows sent to the transformer at once (default 1).
panorama_width, panorama_height	Wrap the canvas so its opposite edges join seamlessly.
prompt=["", p1, ..., pN] + regional_masks	Regional prompts: a fallback prompt plus one prompt per mask; regional_masks is (N, height // 16, width // 16) and must cover every cell.
image_img2img + strength	Img2img source as packed, normalized latents shaped like latents (the runners build them from an image with examples.example.prepare_img2img_latents); strength below 1.0 keeps part of the source.
image	FLUX.2 and FLUX 3 reference images, cropped per window.

Window sizes, strides and offsets must be multiples of vae_scale_factor * 2 (16 px for every family); sizes and strides must be positive, offsets non-negative, the window at most the canvas and the stride at most the window.

Families

Family	Upstream model repo	Blocks	Default steps	Guidance	Upstream license
flux1-dev	black-forest-labs/FLUX.1-dev	Flux1MultiDiffusionAutoBlocks	50	embedded, 3.5	FLUX.1 [dev] Non-Commercial, gated
flux1-schnell	black-forest-labs/FLUX.1-schnell	Flux1MultiDiffusionAutoBlocks	50 (runners: 4)	none	Apache-2.0, gated
flux2-dev	black-forest-labs/FLUX.2-dev	Flux2DevMultiDiffusionAutoBlocks	50	embedded, 4.0	FLUX Non-Commercial, gated
flux2-klein-4b	black-forest-labs/FLUX.2-klein-4B	Flux2KleinMultiDiffusionAutoBlocks	50 (runners: 4)	none	Apache-2.0
flux2-klein-base-4b	black-forest-labs/FLUX.2-klein-base-4B	Flux2KleinBaseMultiDiffusionAutoBlocks	50	CFG, 4.0	Apache-2.0
flux2-klein-9b	black-forest-labs/FLUX.2-klein-9B	Flux2KleinMultiDiffusionAutoBlocks	50 (runners: 4)	none	FLUX Non-Commercial, gated
flux2-klein-base-9b	black-forest-labs/FLUX.2-klein-base-9B	Flux2KleinBaseMultiDiffusionAutoBlocks	50	CFG, 4.0	FLUX Non-Commercial, gated
flux3-klein	black-forest-labs/flux-3-image-klein	Flux3KleinMultiDiffusionAutoBlocks	16	none	FLUX Kommunity, gated
flux3-klein-base	black-forest-labs/flux-3-image-klein-base	Flux3KleinBaseMultiDiffusionAutoBlocks	50	CFG, 3.5	FLUX Kommunity, gated

• The step counts are the block defaults. The runners run FLUX.1 schnell and the distilled FLUX.2 Klein in 4 steps, and schnell with max_sequence_length=256.
• In Python, pass guidance_scale to the embedded-guidance families (FLUX.1 dev, FLUX.2 dev). Swap the guider of the CFG families with pipe.update_components(guider=pipe.guider.new(guidance_scale=4.0)). The runners do both for --guidance-scale. The other families ignore it.
• Both FLUX 3 families also load the prompt upsampler from black-forest-labs/flux-3-image-klein-pup (Apache-2.0).
• The root metadata is the flux3-klein export, and model/<family>/ holds the metadata of every family. Those folders are not standalone remote-code repos.
• The code in this repo is Apache-2.0. The upstream model licenses apply to the weights.

Feature Matrix

Feature	FLUX.1	FLUX.2	FLUX 3
Windows, batched windows, weighting, per-step offsets	yes	yes	yes
Schedule shift from the window size	dev (schnell is static)	yes	not needed (constant shift)
Regional prompts with masks	yes	yes	yes, each regional prompt encoded on its own
Panorama wrap with circular VAE decode	yes	yes	yes
Img2img (image_img2img + strength)	yes; also the native image workflow in Python	yes	yes
Reference image conditioning	no (Kontext is not covered)	yes, cropped per window	yes, cropped per window

FLUX 3 Notes

• Prompt upsampler:
  • By default it runs first, as Black Forest Labs recommends. It is a Qwen3.5-9B fine-tune.
  • The Diffusers docs list about 39 GB in bf16 with it and about 20 GB without it.
  • Skip it with blocks.sub_blocks.pop("prompt_upsampler") before init_pipeline, or with --no-prompt-upsampler in the runners.
• Regional prompts (prompt=["", p1, ..., pN]): each regional prompt is upsampled and encoded on its own. Windows are batched together only when their prompts have the same token length. The upsampler's layout boxes describe the whole canvas, so they cannot be combined with regional prompts.
• Reference images: a single reference image is cropped back to its own box before it is mapped onto the canvas, so give references the canvas aspect ratio.
• Sampling defaults:
  • The blocks keep the Diffusers 0.42.1 defaults: a constant time shift of 6.93, and CFG 3.5 for the base model.
  • Change the CFG scale with the guider swap above.
  • Change the shift by replacing the set_timesteps step before init_pipeline:

from copy import deepcopy

from diffusers.modular_pipelines.flux3.before_denoise import Flux3SetTimestepsStep

from multidiffusion import Flux3KleinBaseMultiDiffusionAutoBlocks

# top-level blocks share their sub-blocks across instances: edit a copy
blocks = deepcopy(Flux3KleinBaseMultiDiffusionAutoBlocks())
for workflow in blocks.sub_blocks["denoise"].sub_blocks.values():
    workflow.sub_blocks["set_timesteps"] = Flux3SetTimestepsStep(shift=16.0)
pipe = blocks.init_pipeline("black-forest-labs/flux-3-image-klein-base")

For the distilled model, keep its 16-step default with Flux3SetTimestepsStep(shift=16.0, num_inference_steps=16).

Export Metadata

Refresh the root default metadata, or a family folder:

uv run multidiff-modular export-hf . --family flux3-klein
uv run multidiff-modular export-hf model/flux2-klein-base-9b --family flux2-klein-base-9b

To create a standalone single-family Hub repo folder, run the command below. It copies root block.py and multidiffusion/ beside the generated JSON. With --push-to-hub, the folder is uploaded once all its files are written. The checkout root and its model/ folders are published with git instead.

uv run multidiff-modular export-hf hf_export --family flux2-klein-4b --repo-id <user>/<repo> --push-to-hub

The root README is copied as-is. Edit its default-family text in the export before publishing. scripts/export_hf_repo.py runs the same command without installing the project.

Hugging Face Setup

No model weights are stored here: the components load from the upstream FLUX repos named in the metadata. Authenticate before using gated models, and prefill the cache for offline runs:

hf auth login
hf download black-forest-labs/FLUX.2-klein-4B

--local-files-only needs every repo the run reads already cached:

• the upstream model repo;
• for FLUX 3, also black-forest-labs/flux-3-image-klein-pup (unless you pass --no-prompt-upsampler);
• for example_remote.py, also arlaz/modular-flux123-multidiffusion itself.

Development

Commands, conventions and the Modular Diffusers rules of this repo are in AGENTS.md. tests/test_denoiser_equivalence.py checks every MultiDiffusion denoise step against its upstream step, on tiny random transformers on CPU.

The real-model smoke test is opt-in and reads models from the local cache. It loads the repo root, or the family named by MULTIDIFF_MODULAR_SMOKE_FAMILY:

MULTIDIFF_MODULAR_RUN_REAL_INFERENCE=1 uv run pytest -ra

The Ubuntu-based dev container (VS Code and PyCharm Dev Containers) mounts the workspace, exposes NVIDIA GPUs and shares the host Hugging Face cache, ~/.cache/huggingface, which must exist on the host. Run uv sync --all-extras after opening it.