Fibo Edit: multi-reference conditioning and batching

[ShivamShrirao]

What does this PR do?

Adds multi-reference image conditioning to BriaFiboEditPipeline and fixes the two remaining items from the bria_fibo review issue.

Fixes the remaining items (issues 2 and 5) of #13618, issues 1/3/4/6 were fixed by #13981.

Multi-reference conditioning

image now accepts a PIL.Image.Image or a list of them. A list means multiple references, not a batch:

• Each reference is VAE-encoded at its own aspect ratio (area capped at 1024², dims rounded to /16), packed, and appended to the sequence as context tokens.
• Reference i is placed on RoPE time plane i (generated tokens stay on plane 0).
• Output resolution defaults to the first reference's size; explicit height/width still override it.
• Masks remain supported with exactly one reference.

Batching

• Prompt lists are now accepted (references repeat across the batch).
• num_images_per_prompt > 1 now works, upstream crashed at a latent reshape because reference latents were never duplicated (fixes issue 5 of bria_fibo model/pipeline review #13618). test_num_images_per_prompt is un-skipped and passes.

Other fixes

• Custom timesteps are now honored in both BriaFiboPipeline and BriaFiboEditPipeline, they were accepted and documented but hard-coded to None in the retrieve_timesteps call (fixes issue 2 of bria_fibo model/pipeline review #13618).
• The attention mask is now a bool key-padding mask (batch, 1, 1, seq) instead of a float outer-product matrix, and is dropped entirely when nothing is padded, this keeps backends without mask support (flash-attn 2/3) usable at guidance_scale <= 1.
• BriaFiboTransformer2DModel declares _repeated_blocks for regional compilation.
• do_patching=True previously produced mismatched token widths between generated and reference latents (references were always packed unpatched); reference packing now follows do_patching.

Intentional behavior changes (please review)

• torch.Tensor image input now raises. References are PIL-only, following the Flux2 / QwenImage Edit Plus convention. Tensor input only became functional recently (fix(bria_fibo): fix guidance_embeds, prompt_embeds, tensor-image and multi-image crashes #13981 fixed a crash); np.ndarray was already rejected.
• List semantics changed: a list of images previously went into preprocess as a batch; it now means multiple references. image=[img] behaves identically to before.
• A single reference is no longer resized to the output resolution, it is encoded at its own aspect ratio like any other reference. With default height/width (taken from the reference) behavior is equivalent; it only diverges when an explicit output size differs from the reference.
• Reference-latent normalization runs in the latent's dtype (repo convention, as in Flux2/Wan/Qwen). Verified against fp32 normalization on real weights: MAE 0.95/255, visually identical.

Validation

• Generation verified end-to-end on real Fibo-Edit-1.5 weights (single and multi-reference, batched prompts).
• Reference indexing/RoPE layout verified against the internal training code (reference i → time id i, per-reference variable resolution, target-first sequence order).
• Fast suites: all bria_fibo + bria_fibo_edit tests pass. The pre-existing test_inference_batch_single_identical failure in the base bria_fibo suite reproduces on current main without this PR, not introduced here.
• New tests: distinct RoPE time planes per reference, batched prompts × multiple references, mask-requires-single-reference.
• make style, make fix-copies clean.

Docs: added a "Multiple reference images" section to bria_fibo_edit.md with a runnable two-reference example; fixed and extended EXAMPLE_DOC_STRING.

Self-review notes (final round)

Rubric: .ai/review-rules.md (+ .ai/AGENTS.md, .ai/pipelines.md, .ai/testing.md, .ai/models.md, model-integration pitfalls). Scope: full diff vs upstream base.

Verdict: READY

No blocking findings. Reference preparation is a single public lifecycle method (prepare_reference_latents) called from __call__, one code path handles one or N references, and tests assert on public API and end-to-end outputs only.

Checks: full bria_fibo + bria_fibo_edit fast suites pass (aside from the pre-existing failure noted below), make style clean, check_copies clean.

Behavior changes (disclosed intentionally)

• torch.Tensor image input now raises (was supported upstream since fix(bria_fibo): fix guidance_embeds, prompt_embeds, tensor-image and multi-image crashes #13981). References are PIL-only, following the Flux2 / QwenImage Edit Plus convention. np.ndarray was already rejected upstream, unchanged.
• List semantics changed: upstream fed a list into preprocess as a batch; a list now means multiple references. image=[img] behaves identically to before; image=[a, b] is a two-reference edit.
• Single reference is no longer resized to the output resolution, it is encoded at its own aspect ratio (1024² area cap), same as any other reference; output resolution defaults from the first reference.
• Prompt lists accepted (batched prompts); reference latents repeat across the batch.
• Custom timesteps are now honored in both pipelines (fixes issue 2 of bria_fibo model/pipeline review #13618): they were accepted and documented but hard-coded to None in the retrieve_timesteps call, silently falling back to the default schedule.
• num_images_per_prompt > 1 now works in the edit pipeline (fixes issue 5 of bria_fibo model/pipeline review #13618): upstream crashed at a latent reshape because reference latents were never duplicated. Everything now keys off the encoded prompt batch (batch_size * num_images_per_prompt); test_num_images_per_prompt is un-skipped and passes.

Left for the reviewer

• _vae_safe_dims / _vae_safe_size carry base_resolution=1024 / multiple=16 defaults no caller overrides, named constants vs. unused knobs is a judgment call.
• Pre-existing test_inference_batch_single_identical failure in the base bria_fibo pipeline reproduces on main without this PR, not introduced here.

Dead-code analysis (advisory)

Traced __call__ end to end: paste_mask_on_image, is_valid_mask, get_mask_size, and both do_patching branches are all reachable. No likely-dead code beyond the default params noted above.

Docs

docs/source/en/api/pipelines/bria_fibo_edit.md gained a "Multiple reference images" section (semantics + runnable two-reference example); EXAMPLE_DOC_STRING fixed (imports) and extended with a multi-reference call.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? bria_fibo model/pipeline review #13618
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

@yiyixuxu @sayakpaul