external_query argument mismatch between flux_kontext.py and MochiAttnProcessor2_0 in diffusers

[dyk-npu]

FluxAttnProcessor.call() got an unexpected keyword argument external_query

Hi, thanks for the great work!

I encountered an error when running the project. It seems that the FluxAttnProcessor in my installed version of diffusers does not support the external_query argument.

---

Error

The error occurs at the following call:

return self.processor(
    self,
    hidden_states,
    external_query=external_query,
    encoder_hidden_states=encoder_hidden_states,
    attention_mask=attention_mask,
    wtext=wtext,
    **cross_attention_kwargs,
)

Runtime error:

FluxAttnProcessor.__call__() got an unexpected keyword argument 'external_query'

---

My Investigation

I checked the implementation of FluxAttnProcessor in:

diffusers/models/transformers/transformer_flux.py

The definition in my environment is:

class FluxAttnProcessor:
    _attention_backend = None

    def __init__(self):
        if not hasattr(F, "scaled_dot_product_attention"):
            raise ImportError(
                f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version."
            )

    def __call__(
        self,
        attn: "FluxAttention",
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        image_rotary_emb: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:

As shown above, external_query is not included in the function arguments, which leads to the error.

---

Environment

• diffusers: 0.35.0
• k-diffusion: 0.1.1.post1

---

Question

Has this issue been encountered before?

Is the project expecting a custom or modified version of diffusers, or a specific diffusers version/commit where FluxAttnProcessor supports external_query?

Any guidance would be appreciated. Thanks!

[RohanU97]

I have fixed this by installing diffusers dev0 however the output looks very uninspired. Just the same as my input. This repository has not released the correct model it has used for benchmarking and showcasing the transformation

[HilaManor]

Hi! Sorry for getting to you back so late, I did not receive notifications for issues in the repo.

So apparently the diffusers version from when we developed the code and when publishing it changed in a non-backward compatible manner. FluxAttnProcessor2_0 was deprecated in huggingface/diffusers@18c8f10 without a deprecation warning. beforehand.

A simple fix is to introduce the original processor back via updating extensions_built_in/flux_kontext.py:

-from diffusers.models.attention_processor import FluxAttnProcessor2_0, Attention
+from diffusers.models.attention_processor import Attention

...

+class FluxAttnProcessor2_0:
+    """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("FluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.FloatTensor:
+        batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+        # `sample` projections.
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+        if encoder_hidden_states is not None:
+            # `context` projections.
+            encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+            encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+            encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+            encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+
+            if attn.norm_added_q is not None:
+                encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+            if attn.norm_added_k is not None:
+                encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+            # attention
+            query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+            key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+            value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+        if image_rotary_emb is not None:
+            from .embeddings import apply_rotary_emb
+
+            query = apply_rotary_emb(query, image_rotary_emb)
+            key = apply_rotary_emb(key, image_rotary_emb)
+
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        if encoder_hidden_states is not None:
+            encoder_hidden_states, hidden_states = (
+                hidden_states[:, : encoder_hidden_states.shape[1]],
+                hidden_states[:, encoder_hidden_states.shape[1] :],
+            )
+
+            # linear proj
+            hidden_states = attn.to_out[0](hidden_states)
+            # dropout
+            hidden_states = attn.to_out[1](hidden_states)
+
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+            return hidden_states, encoder_hidden_states
+        else:
+            return hidden_states

This worked for me with diffusers 0.37.0.dev0 (which I updated from the github repo a month-ish ago).

Regarding @RohanU97, as mentioned in the README this model was re-trained on non-NVIDIA resources, so it does not produce the exact images from the paper, however it should perform equivalently.
For example, running this model with the following commands (with the eval data decoded) produces the images in the samples directory.

python inference.py -w "lorweb_model/lorweb_model.safetensors" -c "config.yaml" -a
"data/original/unsplash_animals/alvan-nee-brFsZ7qszSY-unsplash.jpg" -t "data/decoded_random_inference_set/db_not_in_trainset/animals/Give_this_animal_a_fantastical_set_of_armor/alvan-nee-brFsZ7qszSY-unsplash_2.png" -b "data/original/unsplash_animals/ryan-walton-AbNO2iejoXA-unsplash.jpg" -o "samples/armor_doggy.jpg" -p "Give this animal a fantastical set of armor"
python inference.py -w "lorweb_model/lorweb_model.safetensors" -c "config.yaml" -a "data/original/unsplash_persons_cropped/masoud-mostafaei-YZfxKBCze_c-unsplash.jpg" -t "data/decoded_random_inference_set/db_not_in_trainset_loras/persons/clay_toy/masoud-mostafaei-YZfxKBCze_c-unsplash_2.png" -b "data/original/unsplash_persons_cropped/oana-valeria-serban-NNSz7P1ECT4-unsplash.jpg" -o "samples/clay_woman_man.jpg" -p "Turn this image into the Clay_Toy style."
python inference.py -w "lorweb_model/lorweb_model.safetensors" -c "config.yaml" -a "data/original/unsplash_animals/aarn-giri-3lGi0BXJ1W0-unsplash.jpg" -t "data/decoded_random_inference_set/db_not_in_trainset_loras/animals/ghibli/aarn-giri-3lGi0BXJ1W0-unsplash_2.png" -b "data/original/unsplash_animals/timo-volz-ZlFKIG6dApg-unsplash.jpg" -o "samples/ghibli_animals_8.jpg" -p "Turn this image into the Ghibli style"
python inference.py -w "lorweb_model/lorweb_model.safetensors" -c "config.yaml" -a "data/original/unsplash_objects/johann-siemens-EPy0gBJzzZU-unsplash.jpg" -t "data/decoded_random_inference_set/db_not_in_trainset/objects/Turn_this_photo_into_an_architectural_rendering/johann-siemens-EPy0gBJzzZU-unsplash_2.png" -b "data/original/unsplash_objects/gustavo-S-W9vDL5whU-unsplash.jpg" -o "samples/architectural_tree_car.jpg" -p "Turn this photo into an architectural rendering"

I'm attaching here the input images and the results to reproduce these samples (just the first 2 for brevitiy).
Please let me know if you still have trouble reproducing the results :)

data/original/unsplash_animals/alvan-nee-brFsZ7qszSY-unsplash.jpg

"data/decoded_random_inference_set/db_not_in_trainset/animals/Give_this_animal_a_fantastical_set_of_armor/alvan-nee-brFsZ7qszSY-unsplash_2.png"

"data/original/unsplash_animals/ryan-walton-AbNO2iejoXA-unsplash.jpg"

results

"data/original/unsplash_persons_cropped/masoud-mostafaei-YZfxKBCze_c-unsplash.jpg"

"data/decoded_random_inference_set/db_not_in_trainset_loras/persons/clay_toy/masoud-mostafaei-YZfxKBCze_c-unsplash_2.png"

"data/original/unsplash_persons_cropped/oana-valeria-serban-NNSz7P1ECT4-unsplash.jpg"

results

[RohanU97]

Thanks for the prompt response.
I was able to execute this command (python inference.py -w "lorweb_model/lorweb_model.safetensors" -c "config.yaml" -a
"data/original/unsplash_animals/alvan-nee-brFsZ7qszSY-unsplash.jpg" -t "data/decoded_random_inference_set/db_not_in_trainset/animals/Give_this_animal_a_fantastical_set_of_armor/alvan-nee-brFsZ7qszSY-unsplash_2.png" -b "data/original/unsplash_animals/ryan-walton-AbNO2iejoXA-unsplash.jpg" -o "samples/armor_doggy.jpg" -p "Give this animal a fantastical set of armor"
) and get desired transformed output .

However, when changing the input image -b to alvan-nee-ZCHj_2lJP00-unsplash.jpg or any other animal image from the dataset folder "unsplash_animals" seems like there is no transformation happening.

[bethharmoning1]

I also encountered the problem. After trying several examples, I found that only the demo example produced good results.

[HilaManor]

Sorry for the delay again, it takes time to get access to the NVIDIA data (since my internship is over) and verify what was going on.
Due to licensing issues, we couldn't release the NVIDIA-trained model, and I tried retraining on the university resources.
It seems that something indeed failed during this retraining, and I didn't catch it before uploading.
I retrained again and the newer model (already up on Hugging Face) seems to perform much better and more similar to the paper version.

Some samples of the corresponding b' generations (using the default seed 50 and cfg 2.5):

Please let me know if the issues persist now.
Sorry again!

[RohanU97]

Great. Thanks for the model update. It does work for all prompts and examples in eval set db_not_in_trainset, however when tested on db_not_in_trainset_loras for example

"a_path": "inference_set/unsplash_animals/ryan-walton-AbNO2iejoXA-unsplash.jpg",
"atag_path": "random_inference_set/db_not_in_trainset_loras/animals/3d_game_asset/ryan-walton-AbNO2iejoXA-unsplash_2.jpg",
"b_path": "inference_set/unsplash_animals/alan-king-KZv7w34tluA-unsplash.jpg",
"b_tag_output_path": "db_not_in_trainset_loras/animals/3d_game_asset/alan-king-KZv7w34tluA-unsplash.jpg",
"category": "animals",
"subset": "db_not_in_trainset_loras",
"prompt": "3d game asset"

It fails to produce expected output.

[HilaManor]

Hi,
Indeed, not all analogies work (it's also reflected by the reported metrics). Specifically, we expect LoRA samples to be harder, since these are "styles" the base model has a rough time generating from scratch. Indeed, I verified this example doesn't work with the NVIDIA model as well