feat(diffusers): add Bria model and pipeline to diffusers #1384
Conversation
Summary of ChangesHello @Cui-yshoho, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request integrates the Bria 3.2 text-to-image model and its associated pipeline into the Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the Bria model and pipeline, which is a significant feature addition. The implementation looks solid and follows the existing structure of the library. I've found a few issues, including a critical runtime error due to a missing import, a logic bug in prompt handling, and some code duplication. I've also pointed out several documentation typos and areas for minor code improvements to enhance robustness and maintainability. Once these points are addressed, the PR should be in great shape.
| FlowMatchEulerDiscreteScheduler, | ||
| KarrasDiffusionSchedulers, | ||
| ) | ||
| from ...utils import logging, scale_lora_layers, unscale_lora_layers |
There was a problem hiding this comment.
The name USE_PEFT_BACKEND used on line 169 is not defined, which will cause a NameError at runtime. You need to import it from ...utils.
| from ...utils import logging, scale_lora_layers, unscale_lora_layers | |
| from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers |
| class BriaPosEmbed(nn.Cell): | ||
| # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11 | ||
| def __init__(self, theta: int, axes_dim: List[int]): | ||
| super().__init__() | ||
| self.theta = theta | ||
| self.axes_dim = axes_dim | ||
|
|
||
| def construct(self, ids: ms.Tensor) -> ms.Tensor: | ||
| n_axes = ids.shape[-1] | ||
| cos_out = [] | ||
| sin_out = [] | ||
| pos = ids.float() | ||
| freqs_dtype = ms.float64 | ||
| for i in range(n_axes): | ||
| cos, sin = get_1d_rotary_pos_embed( | ||
| self.axes_dim[i], | ||
| pos[:, i], | ||
| theta=self.theta, | ||
| repeat_interleave_real=True, | ||
| use_real=True, | ||
| freqs_dtype=freqs_dtype, | ||
| ) | ||
| cos_out.append(cos) | ||
| sin_out.append(sin) | ||
| freqs_cos = mint.cat(cos_out, dim=-1) | ||
| freqs_sin = mint.cat(sin_out, dim=-1) | ||
| return freqs_cos, freqs_sin | ||
|
|
There was a problem hiding this comment.
The class BriaPosEmbed is an exact duplicate of BriaEmbedND defined earlier in this file (lines 348-375). This code duplication increases maintenance overhead. You should remove BriaPosEmbed and use BriaEmbedND in its place.
| else: | ||
| attn_output, context_attn_output, ip_attn_output = None, None, None | ||
|
|
There was a problem hiding this comment.
The else block here silently sets attn_output, context_attn_output, and ip_attn_output to None if attention_outputs does not have a length of 2 or 3. This could hide bugs if self.attn returns an unexpected number of outputs. It would be safer to raise a ValueError to make such issues immediately apparent.
| else: | |
| attn_output, context_attn_output, ip_attn_output = None, None, None | |
| else: | |
| raise ValueError(f"Unexpected number of attention outputs: {len(attention_outputs)}") |
| return_tensors="np", | ||
| ) | ||
| text_input_ids = text_inputs.input_ids | ||
| untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="np").input_ids |
There was a problem hiding this comment.
The logic for detecting and warning about prompt truncation appears to be incorrect. untruncated_ids is calculated using the entire prompt list, while text_input_ids is based on a single prompt p from the list within the loop. This comparison is invalid. You should calculate untruncated_ids for the single prompt p inside the loop.
| untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="np").input_ids | |
| untruncated_ids = tokenizer(p, padding="longest", return_tensors="np").input_ids |
| Original model checkpoints for Bria 3.2 can be found [here](https://huggingface.co/briaai/BRIA-3.2). | ||
| Github repo for Bria 3.2 can be found [here](https://github.com/Bria-AI/BRIA-3.2). | ||
|
|
||
| If you want to learn more about the Bria platform, and get free traril access, please visit [bria.ai](https://bria.ai). |
There was a problem hiding this comment.
There's a typo in the word 'traril'. It should be 'trial'.
| If you want to learn more about the Bria platform, and get free traril access, please visit [bria.ai](https://bria.ai). | |
| If you want to learn more about the Bria platform, and get free trial access, please visit [bria.ai](https://bria.ai). |
| def _get_fused_projections(attn: "BriaAttention", hidden_states, encoder_hidden_states=None): | ||
| query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1) | ||
|
|
||
| encoder_query = encoder_key = encoder_value = (None,) | ||
| if encoder_hidden_states is not None and hasattr(attn, "to_added_qkv"): | ||
| encoder_query, encoder_key, encoder_value = attn.to_added_qkv(encoder_hidden_states).chunk(3, dim=-1) | ||
|
|
||
| return query, key, value, encoder_query, encoder_key, encoder_value | ||
|
|
||
|
|
||
| def _get_qkv_projections(attn: "BriaAttention", hidden_states, encoder_hidden_states=None): | ||
| if attn.fused_projections: | ||
| return _get_fused_projections(attn, hidden_states, encoder_hidden_states) | ||
| return _get_projections(attn, hidden_states, encoder_hidden_states) |
There was a problem hiding this comment.
The function _get_fused_projections and the check for attn.fused_projections in _get_qkv_projections appear to be unused, as the BriaAttention class does not have a fused_projections attribute. This looks like dead code that could be removed to improve clarity and reduce maintenance overhead.
| The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list | ||
| will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the | ||
| `._callback_tensor_inputs` attribute of your pipeline class. | ||
| max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`. |
There was a problem hiding this comment.
The docstring indicates that max_sequence_length defaults to 256, but the function signature on line 439 sets the default to 128. Please update the docstring to match the implementation for consistency.
| max_sequence_length (`int` defaults to 256): Maximum sequence length to use with the `prompt`. | |
| max_sequence_length (`int` defaults to 128): Maximum sequence length to use with the `prompt`. |
|
|
||
| # expand the latents if we are doing classifier free guidance | ||
| latent_model_input = mint.cat([latents] * 2) if self.do_classifier_free_guidance else latents | ||
| if type(self.scheduler) != FlowMatchEulerDiscreteScheduler: |
There was a problem hiding this comment.
Using type() for type checking is generally discouraged in favor of isinstance(), as isinstance() correctly handles inheritance. It's more robust to check if not isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):.
| if type(self.scheduler) != FlowMatchEulerDiscreteScheduler: | |
| if not isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler): |
| # While this behavior is consistent with HF Diffusers for now, | ||
| # it may still be a potential bug source worth validating. | ||
| if attn_mask is not None and 1.0 in attn_mask: | ||
| if attn_mask is not None and attn_mask.dtype != ms.bool_ and 1.0 in attn_mask: |
There was a problem hiding this comment.
这个改动是否是对齐原实现的
嗯嗯，发现的一个新问题，这里修改一下，主要是为了解决sdpa输入如果是float和ms的fa无法对齐的问题
4 participants
What does this PR do?
Add
mindone.diffusers.BriaTransformer2DModelmindone.diffusers.BriaPipelineUsage
Performance
Experiments are tested on Ascend Atlas 800T A2 machines with mindspore 2.7.0.
Before submitting
What's New. Here are thedocumentation guidelines
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@xxx