huggingface / huggingface/diffusers

bria model/pipeline review

Abierto
#13,625 2 comentarios 0 reacciones 0 asignados Ver en GitHub
Lenguaje dominante
Python
Estrellas
34.5k
Forks
7.3k
Merge medio
3 d 3 h
PR fusionados (30 d)
91

Descripción

# `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.
```

Guía de contribución

Abrir la guía de contribución

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Evaluación

Stack tecnológico
python, pytorch
Área
machine-learning, testing-qa
Tipo de issue
Error
Dificultad
4/5
Tiempo estimado
3-5 días
Estado de actividad
Tranquilo
Claridad
Bastante claro
Aptitud para principiantes
48/100

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.