huggingface / huggingface/diffusers

bria model/pipeline review

Open
#13,625 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

# `bria` model/pipeline review

Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`

Review performed against the repository review rules.

Files/categories reviewed: target pipeline/model files, public/lazy imports, docs, fast/slow tests, dtype/device/offload hooks, scheduler/custom timestep paths, attention processor paths, and related Flux precedent.

Duplicate search: checked GitHub issues/PRs for `bria`, `BriaPipeline`, `BriaTransformer2DModel`, `attention_kwargs`, `controlnet_single_block_samples`, `BriaPipelineOutput`, `guidance_embeds`, `timesteps sigmas`, `normalize`, and LoRA. Found broad integration/refactor items like `#11978`, `#12214`, `#13341`, and the separate `bria_fibo` review `#13618`, but no exact duplicates for the issues below.

Test coverage status: fast transformer and pipeline tests exist, and a slow Bria pipeline test exists. The current fast tests do not cover the failure modes below. Targeted pytest collection under `.venv` failed before running tests because this Windows torch build lacks `torch._C._distributed_c10d`.

## Issue 1: `attention_kwargs` are ignored during normal transformer inference

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_bria.py#L649-L666
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_bria.py#L674-L691

Problem:
`BriaTransformer2DModel.forward()` passes `attention_kwargs` to blocks only in the gradient-checkpointing branch. Normal inference drops them, so masks and processor kwargs supplied by `BriaPipeline(..., attention_kwargs=...)` have no effect.

Impact:
Attention masks, custom processor kwargs, and related conditioning paths silently do nothing in the default inference path.

Reproduction:
```python
import torch
from diffusers import BriaTransformer2DModel

torch.manual_seed(0)
model = BriaTransformer2DModel(
patch_size=1, in_channels=4, num_layers=1, num_single_layers=1,
attention_head_dim=8, num_attention_heads=2, joint_attention_dim=32,
axes_dims_rope=[0, 4, 4],
).eval()

inputs = dict(
hidden_states=torch.randn(1, 4, 4),
encoder_hidden_states=torch.randn(1, 4, 32),
img_ids=torch.randn(4, 3),
txt_ids=torch.randn(4, 3),
timestep=torch.ones(1),
return_dict=False,
)
mask = torch.ones(1, 8, dtype=torch.bool)
mask[:, -1] = False

with torch.no_grad():
a = model(**inputs)[0]
b = model(**inputs, attention_kwargs={"attention_mask": mask})[0]

print((a - b).abs().max().item()) # 0.0: mask was dropped
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L714-L764

Suggested fix:
```python
encoder_hidden_states, hidden_states = block(
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
temb=temb,
image_rotary_emb=image_rotary_emb,
attention_kwargs=attention_kwargs,
)
```

## Issue 2: `controlnet_single_block_samples` residuals crash

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_bria.py#L694-L699

Problem:
After each `BriaSingleTransformerBlock`, `hidden_states` already contains image tokens only. The code slices it with `encoder_hidden_states.shape[1]`, producing an empty or wrong-sized slice before adding ControlNet residuals.

Impact:
Any ControlNet path using `controlnet_single_block_samples` fails at runtime.

Reproduction:
```python
import torch
from diffusers import BriaTransformer2DModel

model = BriaTransformer2DModel(
patch_size=1, in_channels=4, num_layers=1, num_single_layers=1,
attention_head_dim=8, num_attention_heads=2, joint_attention_dim=32,
axes_dims_rope=[0, 4, 4],
).eval()

control = [torch.zeros(1, 4, 16)]
model(
hidden_states=torch.randn(1, 4, 4),
encoder_hidden_states=torch.randn(1, 4, 32),
img_ids=torch.randn(4, 3),
txt_ids=torch.randn(4, 3),
timestep=torch.ones(1),
controlnet_single_block_samples=control,
)
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L766-L770

Suggested fix:
```python
hidden_states = hidden_states + controlnet_single_block_samples[index_block // interval_control]
```

## Issue 3: Custom `timesteps` are unusable

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/bria/pipeline_bria.py#L614-L652

Problem:
`timesteps` is public, but the pipeline still builds `sigmas` and passes both `timesteps` and `sigmas` to `retrieve_timesteps()`, which rejects that combination. In the DDIM/Euler ancestral branch, custom timesteps are dropped by passing `None`.

Impact:
Users cannot run custom timestep schedules despite the documented argument.

Reproduction:
```python
import numpy as np
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.pipelines.flux.pipeline_flux import retrieve_timesteps

scheduler = FlowMatchEulerDiscreteScheduler()
sigmas = np.linspace(1.0, 0.5, 2)
retrieve_timesteps(scheduler, 2, "cpu", timesteps=[999, 500], sigmas=sigmas)
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux/pipeline_flux.py#L119-L143

Suggested fix:
```python
sigmas = None if timesteps is not None else get_original_sigmas(
num_train_timesteps=self.scheduler.config.num_train_timesteps,
num_inference_steps=num_inference_steps,
)
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler,
num_inference_steps,
device,
timesteps=timesteps,
sigmas=sigmas,
)
```

## Issue 4: `normalize=True` crashes when CFG is disabled

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/bria/pipeline_bria.py#L687-L695

Problem:
`cfg_noise_pred_text` is defined only inside the classifier-free guidance branch, but `normalize=True` reads it unconditionally.

Impact:
`BriaPipeline(..., guidance_scale<=1, normalize=True)` fails with `UnboundLocalError`.

Reproduction:
```python
import torch
from diffusers import AutoencoderKL, BriaPipeline, BriaTransformer2DModel, FlowMatchEulerDiscreteScheduler

transformer = BriaTransformer2DModel(
patch_size=1, in_channels=16, num_layers=1, num_single_layers=1,
attention_head_dim=8, num_attention_heads=2, joint_attention_dim=32,
axes_dims_rope=[0, 4, 4],
)
vae = AutoencoderKL(
block_out_channels=(32,), in_channels=3, out_channels=3,
down_block_types=["DownEncoderBlock2D"], up_block_types=["UpDecoderBlock2D"],
latent_channels=4, sample_size=32, shift_factor=0, scaling_factor=0.13025,
)
pipe = BriaPipeline(transformer, FlowMatchEulerDiscreteScheduler(), vae, None, None)
pipe(prompt_embeds=torch.randn(1, 4, 32), height=16, width=16, num_inference_steps=1,
guidance_scale=1.0, normalize=True, output_type="latent")
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py#L1253-L1256

Suggested fix:
```python
if normalize:
if not self.do_classifier_free_guidance:
raise ValueError("`normalize=True` requires `guidance_scale > 1`.")
noise_pred = noise_pred * (0.7 * (cfg_noise_pred_text / noise_pred.std())) + 0.3 * noise_pred
```

## Issue 5: Negative prompt embedding validation is not wired into `__call__`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/bria/pipeline_bria.py#L304-L316
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/bria/pipeline_bria.py#L560-L568

Problem:
`check_inputs()` can validate `negative_prompt` and `negative_prompt_embeds`, but `__call__()` does not pass them. Shape mismatches reach `torch.cat()` and fail with a low-level tensor error.

Impact:
Invalid user inputs are not caught with the intended clear `ValueError`.

Reproduction:
```python
import torch
from diffusers import AutoencoderKL, BriaPipeline, BriaTransformer2DModel, FlowMatchEulerDiscreteScheduler

transformer = BriaTransformer2DModel(
patch_size=1, in_channels=16, num_layers=1, num_single_layers=1,
attention_head_dim=8, num_attention_heads=2, joint_attention_dim=32,
axes_dims_rope=[0, 4, 4],
)
vae = AutoencoderKL(
block_out_channels=(32,), in_channels=3, out_channels=3,
down_block_types=["DownEncoderBlock2D"], up_block_types=["UpDecoderBlock2D"],
latent_channels=4, sample_size=32, shift_factor=0, scaling_factor=0.13025,
)
pipe = BriaPipeline(transformer, FlowMatchEulerDiscreteScheduler(), vae, None, None)
pipe(prompt_embeds=torch.randn(1, 4, 32), negative_prompt_embeds=torch.randn(1, 5, 32),
height=16, width=16, num_inference_steps=1, guidance_scale=5.0, output_type="latent")
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux/pipeline_flux.py#L494-L500

Suggested fix:
```python
self.check_inputs(
prompt=prompt,
height=height,
width=width,
negative_prompt=negative_prompt,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
max_sequence_length=max_sequence_length,
)
```

## Issue 6: `guidance_embeds=True` cannot construct or run correctly

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_bria.py#L551-L552
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_bria.py#L627-L636

Problem:
`BriaTimestepProjEmbeddings` requires `time_theta`, but `guidance_embed` is constructed without it. The forward path then uses `if guidance:` on a tensor, which is ambiguous for batch sizes above 1.

Impact:
Any config or future checkpoint with `guidance_embeds=True` fails.

Reproduction:
```python
from diffusers import BriaTransformer2DModel

BriaTransformer2DModel(
patch_size=1, in_channels=4, num_layers=1, num_single_layers=1,
attention_head_dim=8, num_attention_heads=2, joint_attention_dim=32,
axes_dims_rope=[0, 4, 4], guidance_embeds=True,
)
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L682-L690

Suggested fix:
```python
if guidance_embeds:
self.guidance_embed = BriaTimestepProjEmbeddings(embedding_dim=self.inner_dim, time_theta=time_theta)

...
if guidance is not None:
temb += self.guidance_embed(guidance, dtype=hidden_states.dtype)
```

## Issue 7: Pipeline has dead LoRA code but no LoRA loader mixin

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/bria/pipeline_bria.py#L13
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/bria/pipeline_bria.py#L89
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/bria/pipeline_bria.py#L183-L236

Problem:
`BriaPipeline` imports and checks `FluxLoraLoaderMixin`, but the class does not inherit it. Therefore `isinstance(self, FluxLoraLoaderMixin)` is always false and public LoRA loader methods are absent.

Impact:
Users cannot call `pipe.load_lora_weights()`, and `attention_kwargs={"scale": ...}` cannot scale text-encoder LoRA layers through this pipeline path.

Reproduction:
```python
from diffusers import BriaPipeline, FluxPipeline

print(hasattr(BriaPipeline, "load_lora_weights")) # False
print(hasattr(FluxPipeline, "load_lora_weights")) # True
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux/pipeline_flux.py#L147-L153

Suggested fix:
```python
class BriaPipeline(DiffusionPipeline, FluxLoraLoaderMixin):
...
```
If Bria LoRA loading is intentionally unsupported, remove the unused import, `isinstance` branches, and `lora_scale` plumbing instead.

## Issue 8: `BriaPipelineOutput` is not exported by the lazy package

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/bria/__init__.py#L13-L25
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/bria/pipeline_output.py#L9-L20

Problem:
`pipeline_output.py` defines `BriaPipelineOutput`, but `src/diffusers/pipelines/bria/__init__.py` does not include it in `_import_structure`.

Impact:
The output type is public in the pipeline return docs but cannot be imported from `diffusers.pipelines.bria`.

Reproduction:
```python
from diffusers.pipelines.bria.pipeline_output import BriaPipelineOutput
print(BriaPipelineOutput.__name__)

from diffusers.pipelines.bria import BriaPipelineOutput
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux/__init__.py#L15

Suggested fix:
```python
_import_structure = {"pipeline_output": ["BriaPipelineOutput"]}

...
else:
from .pipeline_bria import BriaPipeline
from .pipeline_output import BriaPipelineOutput
```

## Issue 9: Transformer is missing attention/offload integration metadata

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_bria.py#L506-L527

Problem:
`BriaTransformer2DModel` does not inherit `AttentionMixin` and does not define `_no_split_modules`. Comparable transformer models expose model-level attention processor APIs and protect transformer blocks from unsafe `device_map` splitting.

Impact:
Users cannot call standard APIs like `set_attn_processor()` / `fuse_qkv_projections()`, and automatic placement/offload can split residual transformer blocks.

Reproduction:
```python
from diffusers import BriaTransformer2DModel, FluxTransformer2DModel

print(hasattr(BriaTransformer2DModel, "set_attn_processor")) # False
print(getattr(BriaTransformer2DModel, "_no_split_modules", None)) # None
print(FluxTransformer2DModel._no_split_modules)
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L525-L568

Suggested fix:
```python
from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward

class BriaTransformer2DModel(..., CacheMixin, AttentionMixin):
_no_split_modules = ["BriaTransformerBlock", "BriaSingleTransformerBlock"]
_skip_layerwise_casting_patterns = ["pos_embed", "norm"]
_repeated_blocks = ["BriaTransformerBlock", "BriaSingleTransformerBlock"]
```

## Issue 10: RoPE still selects float64 on NPU

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_bria.py#L278-L280
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_bria.py#L346-L348

Problem:
The RoPE helpers gate float64 only for MPS. The review rules call out NPU as another backend where float64 is unsupported.

Impact:
NPU users can hit unsupported dtype failures in positional embedding creation.

Reproduction:
```python
import torch
from diffusers.models.transformers.transformer_bria import BriaEmbedND

if not hasattr(torch, "npu") or not torch.npu.is_available():
print("Requires an NPU runtime")
else:
ids = torch.zeros(4, 3, device="npu")
BriaEmbedND(theta=10000, axes_dim=[0, 4, 4])(ids)
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L505-L508

Suggested fix:
```python
is_mps = ids.device.type == "mps"
is_npu = ids.device.type == "npu"
freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64
```

## Issue 11: Public `__call__` docstring contains unresolved conflict markers

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/bria/pipeline_bria.py#L493-L508

Problem:
The public `BriaPipeline.__call__` docstring still contains `<<<<<<< HEAD`, `=======`, and `>>>>>>> main`.

Impact:
Generated API docs expose merge-conflict text, and the file violates the review rule against ephemeral PR-context artifacts.

Reproduction:
```python
from diffusers import BriaPipeline

doc = BriaPipeline.__call__.__doc__ or ""
print("<<<<<<<" in doc, "=======" in doc, ">>>>>>>" in doc)
```

Relevant precedent:
N/A.

Suggested fix:
```python
# Delete the conflict markers and keep one resolved guidance_scale paragraph.
```

Contributor guide

Open the contributing guide

Research direction

Start with src/diffusers/models/transformers/transformer_bria.py and src/diffusers/pipelines/bria/pipeline_bria.py, then compare the cited Flux implementations and existing fast transformer/pipeline tests. Reproduce the listed failures with a working torch environment, add regression coverage for each affected path, and confirm the Bria tests pass with the documented behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, testing-qa
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.