upload

.gitignore

@@ -0,0 +1,6 @@
+.idea/
+
+__pycache__/
+*.py[cod]
+*$py.class
+

README.md

@@ -1,12 +1,15 @@
 ---
 title: NAG FLUX.1-dev
-emoji: 🏢
-colorFrom: green
-colorTo: purple
+emoji: 🚀
+colorFrom: purple
+colorTo: gray
 sdk: gradio
 sdk_version: 5.32.0
 app_file: app.py
 pinned: false
+license: mit
+short_description: Demo of Normalized Attention Guidance for FLUX.1-dev
 ---
 
+[[arXiv Paper]](https://arxiv.org/abs/2505.21179) [[Project Page]](https://chendaryen.github.io/NAG.github.io/)
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,161 @@
+import random
+import os
+
+import spaces
+import torch
+from PIL import Image
+import huggingface_hub
+import gradio as gr
+
+from src.pipeline_flux_nag import NAGFluxPipeline
+from src.transformer_flux import NAGFluxTransformer2DModel
+
+
+theme = gr.themes.Base(
+    font=[gr.themes.GoogleFont('Libre Franklin'), gr.themes.GoogleFont('Public Sans'), 'system-ui', 'sans-serif'],
+)
+
+transformer = NAGFluxTransformer2DModel.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    subfolder="transformer",
+    torch_dtype=torch.bfloat16,
+)
+pipe = NAGFluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    transformer=transformer,
+    torch_dtype=torch.bfloat16,
+)
+
+device = "cuda"
+pipe = pipe.to(device)
+
+examples = [
+    ["Portrait of AI researcher.", "Glasses.", 5],
+    ["A beautiful cyborg.", "Robot.", 5],
+    ["Minimalist abstract line drawing: face portrait of a girl with long hair.", "Complex, detail.", 5],
+    ["A baby phoenix made of fire and flames is born from the smoking ashes.", "Low resolution, blurry, lack of details, illustration, cartoon, painting.", 5],
+    ["A tiny astronaut hatching from an egg on the moon.", "Low resolution, blurry, lack of details, illustration, cartoon, painting.", 5]
+]
+
+@spaces.GPU
+def sample(
+        prompt,
+        negative_prompt=None, guidance_scale=3.5,
+        nag_negative_prompt=None, nag_scale=5.0,
+        num_inference_steps=25,
+        seed=2025, randomize_seed=False,
+        compare=True,
+):
+    prompt = prompt.strip()
+    negative_prompt = negative_prompt.strip() if negative_prompt and negative_prompt.strip() else None
+    guidance_scale = float(guidance_scale)
+    num_inference_steps = int(num_inference_steps)
+    
+    if (randomize_seed):
+        seed = random.randint(0, 9007199254740991)
+    else:
+        seed = int(seed)
+
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    image_nag = pipe(
+        prompt,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        nag_negative_prompt=nag_negative_prompt,
+        nag_scale=nag_scale,
+        generator=generator,
+        num_inference_steps=num_inference_steps,
+    ).images[0]
+
+    if compare:
+        generator = torch.Generator(device="cuda").manual_seed(seed)
+        image_normal = pipe(
+            prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            generator=generator,
+            num_inference_steps=num_inference_steps,
+        ).images[0]
+    else:
+        image_normal = Image.new("RGB", image_nag.size, color=(0, 0, 0))
+
+    return (image_normal, image_nag), seed
+
+
+def sample_example(
+        prompt,
+        nag_negative_prompt,
+        nag_scale,
+):
+    outputs, seed = sample(
+        prompt=prompt,
+        nag_negative_prompt=nag_negative_prompt,
+        nag_scale=nag_scale,
+    )
+    return outputs, 3.5, 25, seed, True
+
+
+css = '''
+.gradio-container{
+max-width: 768px !important;
+margin: 0 auto;
+}
+'''
+
+with gr.Blocks(css=css, theme=theme) as demo:
+    gr.Markdown('''# Normalized Attention Guidance (NAG) Flux-Dev
+    Implementation of [Normalized Attention Guidance](https://chendaryen.github.io/NAG.github.io/)
+    ''')
+    with gr.Group():
+        prompt = gr.Textbox(
+            label="Prompt",
+            max_lines=1,
+            placeholder="Enter your prompt",
+        )
+        nag_negative_prompt = gr.Textbox(
+            label="Negative Prompt for NAG",
+            value="Low resolution, blurry, lack of details, illustration, cartoon, painting.",
+            max_lines=1,
+        )
+        nag_scale = gr.Slider(label="NAG Scale", minimum=1., maximum=20., step=0.25, value=5.)
+        compare = gr.Checkbox(label="Compare with baseline", info="If unchecked, only sample with NAG will be generated.", value=True)
+        button = gr.Button("Generate", min_width=120)
+        output = gr.ImageSlider(label="Left: Baseline, Right: With NAG", interactive=False)
+        with gr.Accordion("Advanced Settings", open=False):
+            negative_prompt = gr.Textbox(label="Negative Prompt", value=None, visible=False)
+            guidance_scale = gr.Slider(label="Guidance Scale", minimum=1., maximum=15., step=0.1, value=3.5)
+            num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=25)
+            seed = gr.Slider(label="Seed", minimum=1, maximum=9007199254740991, step=1, randomize=True)
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+        gr.Examples(
+            examples=examples,
+            fn=sample_example,
+            inputs=[
+                prompt,
+                nag_negative_prompt,
+                nag_scale,
+            ],
+            outputs=[output, guidance_scale, num_inference_steps, seed, compare],
+            cache_examples="lazy",
+        )
+
+    gr.on(
+        triggers=[
+            button.click,
+            prompt.submit
+        ],
+        fn=sample,
+        inputs=[
+            prompt,
+            negative_prompt, guidance_scale,
+            nag_negative_prompt, nag_scale,
+            num_inference_steps,
+            seed, randomize_seed,
+            compare,
+        ],
+        outputs=[output, seed],
+    )
+if __name__ == "__main__":
+    huggingface_hub.login(os.getenv('HF_TOKEN'))
+    demo.launch(share=True)

requirements.txt

@@ -0,0 +1,5 @@
+accelerate
+diffusers
+torch
+transformers
+sentencepiece

src/attention_flux_nag.py

@@ -0,0 +1,176 @@
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+
+from diffusers.models.attention_processor import Attention
+from diffusers.models.embeddings import apply_rotary_emb
+
+
+class NAGFluxAttnProcessor2_0:
+    """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+    def __init__(
+            self,
+            nag_scale: float = 1.0,
+            nag_tau=2.5,
+            nag_alpha=0.25,
+            encoder_hidden_states_length: int = None,
+    ):
+        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.")
+        self.nag_scale = nag_scale
+        self.nag_tau = nag_tau
+        self.nag_alpha = nag_alpha
+        self.encoder_hidden_states_length = encoder_hidden_states_length
+
+    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
+
+        if self.nag_scale > 1.:
+            if encoder_hidden_states is not None:
+                assert len(hidden_states) == batch_size * 0.5
+            apply_guidance = True
+        else:
+            apply_guidance = False
+
+        # `sample` projections.
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        # attention
+        if apply_guidance and encoder_hidden_states is not None:
+            query = query.tile(2, 1, 1)
+            key = key.tile(2, 1, 1)
+            value = value.tile(2, 1, 1)
+
+        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)
+
+            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)
+
+            encoder_hidden_states_length = encoder_hidden_states.shape[1]
+
+        else:
+            assert self.encoder_hidden_states_length is not None
+            encoder_hidden_states_length = self.encoder_hidden_states_length
+
+        if image_rotary_emb is not None:
+            query = apply_rotary_emb(query, image_rotary_emb)
+            key = apply_rotary_emb(key, image_rotary_emb)
+
+        if not apply_guidance:
+            hidden_states = F.scaled_dot_product_attention(query, key, value, 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)
+
+        else:
+            origin_batch_size = batch_size // 2
+            query, query_negative = torch.chunk(query, 2, dim=0)
+            key, key_negative = torch.chunk(key, 2, dim=0)
+            value, value_negative = torch.chunk(value, 2, dim=0)
+
+            hidden_states_negative = F.scaled_dot_product_attention(query_negative, key_negative, value_negative, dropout_p=0.0, is_causal=False)
+            hidden_states_negative = hidden_states_negative.transpose(1, 2).reshape(origin_batch_size, -1, attn.heads * head_dim)
+            hidden_states_negative = hidden_states_negative.to(query.dtype)
+
+            hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+            hidden_states = hidden_states.transpose(1, 2).reshape(origin_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] :],
+            )
+
+            if apply_guidance:
+                encoder_hidden_states_negative, hidden_states_negative = (
+                    hidden_states_negative[:, : encoder_hidden_states.shape[1]],
+                    hidden_states_negative[:, encoder_hidden_states.shape[1]:],
+                )
+                hidden_states_positive = hidden_states
+                hidden_states_guidance = hidden_states_positive * self.nag_scale - hidden_states_negative * (self.nag_scale - 1)
+                norm_positive = torch.norm(hidden_states_positive, p=2, dim=-1, keepdim=True).expand(*hidden_states_positive.shape)
+                norm_guidance = torch.norm(hidden_states_guidance, p=2, dim=-1, keepdim=True).expand(*hidden_states_positive.shape)
+
+                scale = norm_guidance / norm_positive
+                hidden_states_guidance = hidden_states_guidance * torch.minimum(scale, scale.new_ones(1) * self.nag_tau) / scale
+
+                hidden_states = hidden_states_guidance * self.nag_alpha + hidden_states_positive * (1 - self.nag_alpha)
+
+                encoder_hidden_states = torch.cat((encoder_hidden_states, encoder_hidden_states_negative), dim=0)
+
+            # 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:
+            if apply_guidance:
+                image_hidden_states_negative = hidden_states_negative[:, encoder_hidden_states_length:]
+                image_hidden_states = hidden_states[:, encoder_hidden_states_length:]
+
+                image_hidden_states_positive = image_hidden_states
+                image_hidden_states_guidance = image_hidden_states_positive * self.nag_scale - image_hidden_states_negative * (self.nag_scale - 1)
+                norm_positive = torch.norm(image_hidden_states_positive, p=2, dim=-1, keepdim=True).expand(*image_hidden_states_positive.shape)
+                norm_guidance = torch.norm(image_hidden_states_guidance, p=2, dim=-1, keepdim=True).expand(*image_hidden_states_positive.shape)
+
+                scale = norm_guidance / norm_positive
+                image_hidden_states_guidance = image_hidden_states_guidance * torch.minimum(scale, scale.new_ones(1) * self.nag_tau) / scale
+                # scale = torch.nan_to_num(scale, 10)
+                # image_hidden_states_guidance[scale > self.nag_tau] = image_hidden_states_guidance[scale > self.nag_tau] / (norm_guidance[scale > self.nag_tau] + 1e-7) * norm_positive[scale > self.nag_tau] * self.nag_tau
+
+                image_hidden_states = image_hidden_states_guidance * self.nag_alpha + image_hidden_states_positive * (1 - self.nag_alpha)
+
+                hidden_states_negative[:, encoder_hidden_states_length:] = image_hidden_states
+                hidden_states[:, encoder_hidden_states_length:] = image_hidden_states
+                hidden_states = torch.cat((hidden_states, hidden_states_negative), dim=0)
+
+            return hidden_states

src/normalization.py

@@ -0,0 +1,54 @@
+from typing import Optional, Tuple
+
+import torch
+from diffusers.models.normalization import AdaLayerNorm, AdaLayerNormContinuous, AdaLayerNormZero, SD35AdaLayerNormZeroX
+
+
+class TruncAdaLayerNorm(AdaLayerNorm):
+    def forward(
+            self, x: torch.Tensor, timestep: Optional[torch.Tensor] = None, temb: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        batch_size = x.shape[0]
+        return self.forward_old(
+            x, 
+            temb[:batch_size] if temb is not None else None,
+        )
+
+
+class TruncAdaLayerNormContinuous(AdaLayerNormContinuous):
+    def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor:
+        batch_size = x.shape[0]
+        return self.forward_old(x, conditioning_embedding[:batch_size])
+
+
+class TruncAdaLayerNormZero(AdaLayerNormZero):
+    def forward(
+            self,
+            x: torch.Tensor,
+            timestep: Optional[torch.Tensor] = None,
+            class_labels: Optional[torch.LongTensor] = None,
+            hidden_dtype: Optional[torch.dtype] = None,
+            emb: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        batch_size = x.shape[0]
+        return self.forward_old(
+            x,
+            timestep[:batch_size] if timestep is not None else None,
+            class_labels[:batch_size] if class_labels is not None else None,
+            hidden_dtype,
+            emb[:batch_size] if emb is not None else None,
+        )
+
+
+class TruncSD35AdaLayerNormZeroX(SD35AdaLayerNormZeroX):
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            emb: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, ...]:
+        batch_size = hidden_states.shape[0]
+        return self.forward_old(
+            hidden_states,
+            emb[:batch_size] if emb is not None else None,
+        )
+        

src/pipeline_flux_nag.py

@@ -0,0 +1,461 @@
+import math
+from typing import Union, List, Optional, Dict, Any, Callable
+import types
+
+import numpy as np
+import torch
+
+from diffusers import FluxPipeline
+from diffusers.pipelines.flux.pipeline_flux import calculate_shift, retrieve_timesteps
+from diffusers.image_processor import PipelineImageInput
+from diffusers.utils import is_torch_xla_available
+from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
+from diffusers.models.normalization import AdaLayerNormZero, AdaLayerNormContinuous
+
+from src.attention_flux_nag import NAGFluxAttnProcessor2_0
+from src.normalization import TruncAdaLayerNormZero, TruncAdaLayerNormContinuous
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+class NAGFluxPipeline(FluxPipeline):
+    @property
+    def do_normalized_attention_guidance(self):
+        return self._nag_scale > 1
+
+    def _set_nag_attn_processor(
+            self,
+            nag_scale,
+            encoder_hidden_states_length,
+            nag_tau=2.5,
+            nag_alpha=0.25,
+    ):
+        attn_procs = {}
+        for name in self.transformer.attn_processors.keys():
+            attn_procs[name] = NAGFluxAttnProcessor2_0(
+                nag_scale=nag_scale,
+                nag_tau=nag_tau,
+                nag_alpha=nag_alpha,
+                encoder_hidden_states_length=encoder_hidden_states_length,
+            )
+        self.transformer.set_attn_processor(attn_procs)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        negative_prompt: Union[str, List[str]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        true_cfg_scale: float = 1.0,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 28,
+        sigmas: Optional[List[float]] = None,
+        guidance_scale: float = 3.5,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        negative_ip_adapter_image: Optional[PipelineImageInput] = None,
+        negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+
+        nag_scale: float = 1.0,
+        nag_tau: float = 2.5,
+        nag_alpha: float = 0.25,
+        nag_end: float = 1.0,
+        nag_negative_prompt: str = None,
+        nag_negative_prompt_2: str = None,
+        nag_negative_prompt_embeds: Optional[torch.Tensor] = None,
+        nag_negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                will be used instead.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
+                not greater than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
+            true_cfg_scale (`float`, *optional*, defaults to 1.0):
+                When > 1.0 and a provided `negative_prompt`, enables true classifier-free guidance.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            guidance_scale (`float`, *optional*, defaults to 7.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
+            ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
+                IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
+                provided, embeddings are computed from the `ip_adapter_image` input argument.
+            negative_ip_adapter_image:
+                (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
+            negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
+                IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
+                provided, embeddings are computed from the `ip_adapter_image` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                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 512): Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
+            is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
+            images.
+        """
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+        self._nag_scale = nag_scale
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        lora_scale = (
+            self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
+        )
+        has_neg_prompt = negative_prompt is not None or (
+            negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None
+        )
+        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+        # do_true_cfg = do_true_cfg or self.do_normalized_attention_guidance
+
+        (
+            prompt_embeds,
+            pooled_prompt_embeds,
+            text_ids,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+        if do_true_cfg:
+            (
+                negative_prompt_embeds,
+                negative_pooled_prompt_embeds,
+                _,
+            ) = self.encode_prompt(
+                prompt=negative_prompt,
+                prompt_2=negative_prompt_2,
+                prompt_embeds=negative_prompt_embeds,
+                pooled_prompt_embeds=negative_pooled_prompt_embeds,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                lora_scale=lora_scale,
+            )
+
+        if self.do_normalized_attention_guidance:
+            if nag_negative_prompt_embeds is None or nag_negative_pooled_prompt_embeds is None:
+                if nag_negative_prompt is None:
+                    if negative_prompt is not None:
+                        if do_true_cfg:
+                            nag_negative_prompt_embeds = negative_prompt_embeds
+                            nag_negative_pooled_prompt_embeds = negative_pooled_prompt_embeds
+                        else:
+                            nag_negative_prompt = negative_prompt
+                            nag_negative_prompt_2 = negative_prompt_2
+                    else:
+                        nag_negative_prompt = ""
+
+                if nag_negative_prompt is not None:
+                    nag_negative_prompt_embeds, nag_negative_pooled_prompt_embeds = self.encode_prompt(
+                        prompt=nag_negative_prompt,
+                        prompt_2=nag_negative_prompt_2,
+                        device=device,
+                        num_images_per_prompt=num_images_per_prompt,
+                        max_sequence_length=max_sequence_length,
+                        lora_scale=lora_scale,
+                    )[:2]
+
+        if self.do_normalized_attention_guidance:
+            pooled_prompt_embeds = torch.cat([pooled_prompt_embeds, nag_negative_pooled_prompt_embeds], dim=0)
+            prompt_embeds = torch.cat([prompt_embeds, nag_negative_prompt_embeds], dim=0)
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents, latent_image_ids = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.5),
+            self.scheduler.config.get("max_shift", 1.16),
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # handle guidance
+        if self.transformer.config.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+            guidance = guidance.expand(prompt_embeds.shape[0])
+        else:
+            guidance = None
+
+        if (ip_adapter_image is not None or ip_adapter_image_embeds is not None) and (
+            negative_ip_adapter_image is None and negative_ip_adapter_image_embeds is None
+        ):
+            negative_ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
+        elif (ip_adapter_image is None and ip_adapter_image_embeds is None) and (
+            negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None
+        ):
+            ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
+
+        if self.joint_attention_kwargs is None:
+            self._joint_attention_kwargs = {}
+
+        image_embeds = None
+        negative_image_embeds = None
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+            )
+        if negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None:
+            negative_image_embeds = self.prepare_ip_adapter_image_embeds(
+                negative_ip_adapter_image,
+                negative_ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+            )
+
+        origin_attn_procs = self.transformer.attn_processors
+        if self.do_normalized_attention_guidance:
+            self._set_nag_attn_processor(nag_scale, prompt_embeds.shape[1], nag_tau, nag_alpha)
+            attn_procs_recovered = False
+
+        for sub_mod in self.transformer.modules():
+            if not hasattr(sub_mod, "forward_old") :
+                sub_mod.forward_old = sub_mod.forward
+                if isinstance(sub_mod, AdaLayerNormZero):
+                    sub_mod.forward = types.MethodType(TruncAdaLayerNormZero.forward, sub_mod)
+                elif isinstance(sub_mod, AdaLayerNormContinuous):
+                    sub_mod.forward = types.MethodType(TruncAdaLayerNormContinuous.forward, sub_mod)
+
+        # 6. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                if t < (1 - nag_end) * 1000 and self.do_normalized_attention_guidance and not attn_procs_recovered:
+                    self.transformer.set_attn_processor(origin_attn_procs)
+                    if guidance is not None:
+                        guidance = guidance[:len(latents)]
+                    pooled_prompt_embeds = pooled_prompt_embeds[:len(latents)]
+                    prompt_embeds = prompt_embeds[:len(latents)]
+                    attn_procs_recovered = True
+
+                self._current_timestep = t
+                if image_embeds is not None:
+                    self._joint_attention_kwargs["ip_adapter_image_embeds"] = image_embeds
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(prompt_embeds.shape[0]).to(latents.dtype)
+
+                noise_pred = self.transformer(
+                    hidden_states=latents,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    pooled_projections=pooled_prompt_embeds,
+                    encoder_hidden_states=prompt_embeds,
+                    txt_ids=text_ids,
+                    img_ids=latent_image_ids,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    return_dict=False,
+                )[0]
+
+                if do_true_cfg:
+                    if negative_image_embeds is not None:
+                        self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds
+                    neg_noise_pred = self.transformer(
+                        hidden_states=latents,
+                        timestep=timestep / 1000,
+                        guidance=guidance,
+                        pooled_projections=negative_pooled_prompt_embeds,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        txt_ids=text_ids,
+                        img_ids=latent_image_ids,
+                        joint_attention_kwargs=self.joint_attention_kwargs,
+                        return_dict=False,
+                    )[0]
+                    noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        self._current_timestep = None
+
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        if self.do_normalized_attention_guidance and not attn_procs_recovered:
+            self.transformer.set_attn_processor(origin_attn_procs)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return FluxPipelineOutput(images=image)

src/transformer_flux.py

@@ -0,0 +1,187 @@
+from typing import Any, Dict, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from diffusers.utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class NAGFluxTransformer2DModel(FluxTransformer2DModel):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        pooled_projections: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
+        guidance: torch.Tensor = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        controlnet_block_samples=None,
+        controlnet_single_block_samples=None,
+        return_dict: bool = True,
+        controlnet_blocks_repeat: bool = False,
+    ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+        """
+        The [`FluxTransformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
+                Input `hidden_states`.
+            encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
+                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+            pooled_projections (`torch.Tensor` of shape `(batch_size, projection_dim)`): Embeddings projected
+                from the embeddings of input conditions.
+            timestep ( `torch.LongTensor`):
+                Used to indicate denoising step.
+            block_controlnet_hidden_states: (`list` of `torch.Tensor`):
+                A list of tensors that if specified are added to the residuals of transformer blocks.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+        if joint_attention_kwargs is not None:
+            joint_attention_kwargs = joint_attention_kwargs.copy()
+            lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+        else:
+            if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
+                logger.warning(
+                    "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+                )
+
+        do_nag = hidden_states.shape[0] != encoder_hidden_states.shape[0]
+
+        hidden_states = self.x_embedder(hidden_states)
+
+        timestep = timestep.to(hidden_states.dtype) * 1000
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype) * 1000
+        else:
+            guidance = None
+
+        temb = (
+            self.time_text_embed(timestep, pooled_projections)
+            if guidance is None
+            else self.time_text_embed(timestep, guidance, pooled_projections)
+        )
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        if txt_ids.ndim == 3:
+            logger.warning(
+                "Passing `txt_ids` 3d torch.Tensor is deprecated."
+                "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            )
+            txt_ids = txt_ids[0]
+        if img_ids.ndim == 3:
+            logger.warning(
+                "Passing `img_ids` 3d torch.Tensor is deprecated."
+                "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            )
+            img_ids = img_ids[0]
+
+        ids = torch.cat((txt_ids, img_ids), dim=0)
+        image_rotary_emb = self.pos_embed(ids)
+
+        if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs:
+            ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds")
+            ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds)
+            joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states})
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                )
+
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if controlnet_block_samples is not None:
+                interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                # For Xlabs ControlNet.
+                if controlnet_blocks_repeat:
+                    hidden_states = (
+                        hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
+                    )
+                else:
+                    hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+
+        if do_nag:
+            hidden_states = hidden_states.tile(2, 1, 1)
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    temb,
+                    image_rotary_emb,
+                )
+
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if controlnet_single_block_samples is not None:
+                interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                controlnet_single_block_sample = controlnet_single_block_samples[index_block // interval_control]
+                if do_nag:
+                    controlnet_single_block_sample = controlnet_single_block_sample.tile(2, 1, 1)
+                hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                    hidden_states[:, encoder_hidden_states.shape[1] :, ...] + controlnet_single_block_sample
+                )
+
+        hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+
+        if do_nag:
+            hidden_states = torch.chunk(hidden_states, 2, dim=0)[0]
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)