Upload 3 files

pa_src/attn_processor.py

@@ -0,0 +1,443 @@
+from typing import Callable, Optional
+import torch
+import torch.nn.functional as F
+from diffusers.models.attention_processor import *
+from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel
+
+def default_set_attn_proc_func(
+    name: str,
+    hidden_size: int,
+    cross_attention_dim: Optional[int],
+    ori_attn_proc: object,
+) -> object:
+    return ori_attn_proc
+
+def set_flux_transformer_attn_processor(
+    transformer: FluxTransformer2DModel,
+    set_attn_proc_func: Callable = default_set_attn_proc_func,
+    set_attn_module_names: Optional[list[str]] = None,
+) -> None:
+    do_set_processor = lambda name, module_names: (
+        any([name.startswith(module_name) for module_name in module_names])
+        if module_names is not None
+        else True
+    )  # prefix match
+
+    attn_procs = {}
+    for name, attn_processor in transformer.attn_processors.items():
+        dim_head = transformer.config.attention_head_dim
+        num_heads = transformer.config.num_attention_heads
+        if name.endswith("attn.processor"):
+            attn_procs[name] = (
+                set_attn_proc_func(name, dim_head, num_heads, attn_processor)
+                if do_set_processor(name, set_attn_module_names)
+                else attn_processor
+            )
+
+    transformer.set_attn_processor(attn_procs)
+
+class PersonalizeAnythingAttnProcessor:
+    
+    def __init__(self, name, mask, device, tau=0.98, concept_process=False, shift_mask = None, img_dims=4096):
+        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.name = name
+        self.mask = mask.view(img_dims).bool().to(device)
+        self.device = device
+        self.tau = tau
+        self.concept_process = concept_process
+        self.img_dims = img_dims
+        
+        if shift_mask is None:
+            self.shift_mask = self.mask
+        else:
+            self.shift_mask = shift_mask.view(img_dims).bool().to(device)
+
+    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,
+        timestep = None,
+    ) -> torch.FloatTensor:
+        batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        
+        ###################################################################################
+        if timestep is not None:
+            timestep = timestep
+            
+        concept_process = self.concept_process # token concatenation
+        c_q = concept_process and True # if token concatenation is applied to q
+        c_kv = concept_process and True  # if token concatenation is applied to kv
+        
+        t_flag = timestep > self.tau  # token replacement
+        r_q = True and t_flag # if token concatenation is applied to q
+        r_k = True and t_flag # if token concatenation is applied to k
+        r_v = True and t_flag # if token concatenation is applied to v    
+        
+        if encoder_hidden_states is not None:
+            concept_feature_ = hidden_states[0, self.mask, :] 
+        else:
+            concept_feature_ = hidden_states[0, 512:, :][self.mask, :] 
+        
+        if r_k or r_q or r_v:
+            r_hidden_states = hidden_states
+            if encoder_hidden_states is not None:
+                r_hidden_states[1, self.shift_mask, :] = concept_feature_
+            else:
+                text_hidden_states = hidden_states[1, :512, :]
+                image_hidden_states = hidden_states[1, 512:, :]
+                image_hidden_states[self.shift_mask, :] = concept_feature_
+                
+                r_hidden_states[1] = torch.cat([text_hidden_states, image_hidden_states], dim=0)
+        ###################################################################################
+
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+        query = attn.to_q(hidden_states)
+        
+        ###################################################################################        
+        if r_k:
+            key = attn.to_k(r_hidden_states)
+        if r_q:
+            query = attn.to_q(r_hidden_states)
+        if r_v:
+            value = attn.to_v(r_hidden_states)
+
+        if concept_process:    
+            if c_q:
+                c_query = attn.to_q(concept_feature_)
+                c_query = c_query.repeat(query.shape[0], 1, 1)
+                query = torch.cat([query, c_query], dim=1)
+            if c_kv:
+                c_key = attn.to_k(concept_feature_)
+                c_key = c_key.repeat(key.shape[0], 1, 1)
+                c_value = attn.to_v(concept_feature_)
+                c_value = c_value.repeat(value.shape[0], 1, 1)
+                key = torch.cat([key, c_key], dim=1)
+                value = torch.cat([value, c_value], dim=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)
+
+            # 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 diffusers.models.embeddings import apply_rotary_emb
+            
+            # use original position emb or text emb
+            if not c_q:
+                query = apply_rotary_emb(query, image_rotary_emb)
+            if not c_kv:
+                key = apply_rotary_emb(key, image_rotary_emb)
+            
+            ###################################################################################
+            # get original position emb
+            def get_concept_rotary_emb(ori_rotary_emb, mask):
+                enc_emb = ori_rotary_emb[:512, :]
+                hid_emb = ori_rotary_emb[512:, :]
+                concept_emb = hid_emb[mask, :]
+                
+                image_rotary_emb = torch.cat([enc_emb, hid_emb, concept_emb], dim=0)
+                return image_rotary_emb
+                
+            if concept_process:
+                # 1. use original position emb
+                image_rotary_emb_0 = get_concept_rotary_emb(image_rotary_emb[0], self.shift_mask)
+                image_rotary_emb_1 = get_concept_rotary_emb(image_rotary_emb[1], self.shift_mask)
+                image_rotary_emb = (image_rotary_emb_0, image_rotary_emb_1)
+                
+                # 2. use text emb
+                # dims = (self.mask == 1).sum().item()
+                # concept_rotary_emb_0 = torch.ones((dims, 128)).to(self.device)
+                # concept_rotary_emb_1 = torch.zeros((dims, 128)).to(self.device)
+                # image_rotary_emb = (
+                #     torch.cat([image_rotary_emb[0], concept_rotary_emb_0], dim=0), 
+                #     torch.cat([image_rotary_emb[1], concept_rotary_emb_1], dim=0))      
+                
+                if c_q:
+                    query = apply_rotary_emb(query, image_rotary_emb)
+                if c_kv:
+                    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)
+            
+            ################################################################
+            # restore after token concatenation
+            hidden_states = hidden_states[:, :self.img_dims, :]
+            ################################################################
+
+            return hidden_states, encoder_hidden_states
+        else:
+            ################################################################
+            dims = self.img_dims + 512
+            hidden_states = hidden_states[:, :dims, :]
+            ################################################################
+            
+            return hidden_states
+        
+class MultiPersonalizeAnythingAttnProcessor:
+    
+    def __init__(self, name, masks, device, tau=0.98, concept_process=False, shift_masks = None, img_dims=4096):
+        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.name = name
+        self.device = device
+        self.tau = tau
+        self.concept_process = concept_process
+        self.img_dims = img_dims
+        
+        for i in range(len(masks)):
+            masks[i] = masks[i].view(img_dims).bool().to(device)
+        self.masks = masks
+        
+        if shift_masks is None:
+            self.shift_masks = self.masks
+        else:
+            for i in range(len(shift_masks)):
+                shift_masks[i] = shift_masks[i].view(img_dims).bool().to(device)
+            self.shift_masks = shift_masks
+
+    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,
+        timestep = None,
+    ) -> torch.FloatTensor:
+        batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        
+        ###################################################################################
+        if timestep is not None:
+            timestep = timestep
+            
+        concept_process = self.concept_process # token concatenation
+        c_q = concept_process and True # if token concatenation is applied to q
+        c_kv = concept_process and True  # if token concatenation is applied to kv
+        
+        t_flag = timestep > self.tau  # token replacement
+        r_q = True and t_flag # if token concatenation is applied to q
+        r_k = True and t_flag # if token concatenation is applied to k
+        r_v = True and t_flag # if token concatenation is applied to v    
+        
+        concept_features = []
+        r_hidden_states = hidden_states
+        for id, mask in enumerate(self.masks):
+            if encoder_hidden_states is not None:
+                concept_feature_ = hidden_states[id, mask, :] 
+            else:
+                concept_feature_ = hidden_states[id, 512:, :][mask, :]
+                
+            shift_mask = self.shift_masks[id]
+            concept_features.append(concept_feature_)
+            
+            if r_k or r_q or r_v:
+                if encoder_hidden_states is not None:
+                    r_hidden_states[-1, shift_mask, :] = concept_feature_
+                else:
+                    text_hidden_states = r_hidden_states[-1, :512, :]
+                    image_hidden_states = r_hidden_states[-1, 512:, :]
+                    image_hidden_states[shift_mask, :] = concept_feature_
+                    r_hidden_states[-1] = torch.cat([text_hidden_states, image_hidden_states], dim=0)
+        ###################################################################################
+
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+        query = attn.to_q(hidden_states)
+        
+        ###################################################################################        
+        if r_k:
+            key = attn.to_k(r_hidden_states)
+        if r_q:
+            query = attn.to_q(r_hidden_states)
+        if r_v:
+            value = attn.to_v(r_hidden_states)
+
+        if concept_process:  
+            for concept_feature_ in concept_features:  
+                if c_q:
+                    c_query = attn.to_q(concept_feature_)
+                    c_query = c_query.repeat(query.shape[0], 1, 1)
+                    query = torch.cat([query, c_query], dim=1)
+                if c_kv:
+                    c_key = attn.to_k(concept_feature_)
+                    c_key = c_key.repeat(key.shape[0], 1, 1)
+                    
+                    c_value = attn.to_v(concept_feature_)
+                    c_value = c_value.repeat(value.shape[0], 1, 1)
+                    
+                    key = torch.cat([key, c_key], dim=1)
+                    value = torch.cat([value, c_value], dim=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)
+
+            # 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 diffusers.models.embeddings import apply_rotary_emb
+            
+            # use original position emb or text emb
+            if not c_q:
+                query = apply_rotary_emb(query, image_rotary_emb)
+            if not c_kv:
+                key = apply_rotary_emb(key, image_rotary_emb)
+            
+            ###################################################################################
+            def get_concept_rotary_emb(ori_rotary_emb, shift_masks):
+                enc_emb = ori_rotary_emb[:512, :]
+                hid_emb = ori_rotary_emb[512:, :]
+                
+                concept_embs = []
+                for mask in shift_masks:
+                    concept_embs.append(hid_emb[mask, :])
+                concept_emb = torch.cat(concept_embs, dim=0) if len(concept_embs) > 0 else torch.zeros(0, hid_emb.shape[1], device=hid_emb.device)
+                image_rotary_emb = torch.cat([enc_emb, hid_emb, concept_emb], dim=0)
+                return image_rotary_emb
+            
+            if concept_process:
+                # 1. use original position emb with plural masks
+                image_rotary_emb_0 = get_concept_rotary_emb(image_rotary_emb[0], self.shift_masks)
+                image_rotary_emb_1 = get_concept_rotary_emb(image_rotary_emb[1], self.shift_masks)
+                image_rotary_emb = (image_rotary_emb_0, image_rotary_emb_1)
+                
+                # 2. use text emb with plural masks
+                # total_dims = sum((mask == 1).sum().item() for mask in self.masks)
+                # concept_rotary_emb_0 = torch.ones((total_dims, 128)).to(self.device)
+                # concept_rotary_emb_1 = torch.zeros((total_dims, 128)).to(self.device)
+                # image_rotary_emb = (
+                #     torch.cat([image_rotary_emb[0], concept_rotary_emb_0], dim=0), 
+                #     torch.cat([image_rotary_emb[1], concept_rotary_emb_1], dim=0)
+                # )
+                
+                if c_q:
+                    query = apply_rotary_emb(query, image_rotary_emb)
+                if c_kv:
+                    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)
+            
+            ################################################################
+            # restore after token concatenation
+            hidden_states = hidden_states[:, :self.img_dims, :]
+            ################################################################
+
+            return hidden_states, encoder_hidden_states
+        else:
+            ################################################################
+            dims = self.img_dims + 512
+            hidden_states = hidden_states[:, :dims, :]
+            ################################################################
+            
+            return hidden_states

pa_src/pipeline.py

@@ -0,0 +1,1272 @@
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import (
+    FluxLoraLoaderMixin,
+    FromSingleFileMixin,
+    TextualInversionLoaderMixin,
+)
+from diffusers.models.autoencoders import AutoencoderKL
+from diffusers.models.transformers import FluxTransformer2DModel
+from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    is_torch_xla_available,
+    logging,
+    replace_example_docstring,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from diffusers.utils.torch_utils import randn_tensor
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> import requests
+        >>> import PIL
+        >>> from io import BytesIO
+        >>> from diffusers import DiffusionPipeline
+
+        >>> pipe = DiffusionPipeline.from_pretrained(
+        ...    "black-forest-labs/FLUX.1-dev",
+        ...    torch_dtype=torch.bfloat16,
+        ...    custom_pipeline="pipeline_flux_rf_inversion")
+        >>> pipe.to("cuda")
+
+         >>> def download_image(url):
+        ...     response = requests.get(url)
+        ...     return PIL.Image.open(BytesIO(response.content)).convert("RGB")
+
+
+        >>> img_url = "https://www.aiml.informatik.tu-darmstadt.de/people/mbrack/tennis.jpg"
+        >>> image = download_image(img_url)
+
+        >>> inverted_latents, image_latents, latent_image_ids = pipe.invert(image=image, num_inversion_steps=28, gamma=0.5)
+
+        >>> edited_image = pipe(
+        ...     prompt="a tomato",
+        ...     inverted_latents=inverted_latents,
+        ...     image_latents=image_latents,
+        ...     latent_image_ids=latent_image_ids,
+        ...     start_timestep=0,
+        ...     stop_timestep=.25,
+        ...     num_inference_steps=28,
+        ...     eta=0.9,
+        ... ).images[0]
+        ```
+"""
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.16,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError(
+            "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values"
+        )
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(
+            inspect.signature(scheduler.set_timesteps).parameters.keys()
+        )
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(
+            inspect.signature(scheduler.set_timesteps).parameters.keys()
+        )
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class RFPanoInversionParallelFluxPipeline(
+    DiffusionPipeline,
+    FluxLoraLoaderMixin,
+    FromSingleFileMixin,
+    TextualInversionLoaderMixin,
+):
+    r"""
+    The Flux pipeline for text-to-image generation.
+
+    Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+
+    Args:
+        transformer ([`FluxTransformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModel`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+        text_encoder_2 ([`T5EncoderModel`]):
+            [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
+            the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_2 (`T5TokenizerFast`):
+            Second Tokenizer of class
+            [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast).
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
+    _optional_components = []
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        text_encoder_2: T5EncoderModel,
+        tokenizer_2: T5TokenizerFast,
+        transformer: FluxTransformer2DModel,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+        self.vae_scale_factor = (
+            2 ** (len(self.vae.config.block_out_channels) - 1)
+            if hasattr(self, "vae") and self.vae is not None
+            else 8
+        )
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.tokenizer_max_length = (
+            self.tokenizer.model_max_length
+            if hasattr(self, "tokenizer") and self.tokenizer is not None
+            else 77
+        )
+        self.default_sample_size = 128
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]] = None,
+        num_images_per_prompt: int = 1,
+        max_sequence_length: int = 512,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        if isinstance(self, TextualInversionLoaderMixin):
+            prompt = self.maybe_convert_prompt(prompt, self.tokenizer_2)
+
+        text_inputs = self.tokenizer_2(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_length=False,
+            return_overflowing_tokens=False,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer_2(
+            prompt, padding="longest", return_tensors="pt"
+        ).input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+            text_input_ids, untruncated_ids
+        ):
+            removed_text = self.tokenizer_2.batch_decode(
+                untruncated_ids[:, self.tokenizer_max_length - 1 : -1]
+            )
+            logger.warning(
+                "The following part of your input was truncated because `max_sequence_length` is set to "
+                f" {max_sequence_length} tokens: {removed_text}"
+            )
+
+        prompt_embeds = self.text_encoder_2(
+            text_input_ids.to(device), output_hidden_states=False
+        )[0]
+
+        dtype = self.text_encoder_2.dtype
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        _, seq_len, _ = prompt_embeds.shape
+
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(
+            batch_size * num_images_per_prompt, seq_len, -1
+        )
+
+        return prompt_embeds
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds
+    def _get_clip_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]],
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+    ):
+        device = device or self._execution_device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        if isinstance(self, TextualInversionLoaderMixin):
+            prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer_max_length,
+            truncation=True,
+            return_overflowing_tokens=False,
+            return_length=False,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer(
+            prompt, padding="longest", return_tensors="pt"
+        ).input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+            text_input_ids, untruncated_ids
+        ):
+            removed_text = self.tokenizer.batch_decode(
+                untruncated_ids[:, self.tokenizer_max_length - 1 : -1]
+            )
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {self.tokenizer_max_length} tokens: {removed_text}"
+            )
+        prompt_embeds = self.text_encoder(
+            text_input_ids.to(device), output_hidden_states=False
+        )
+
+        # Use pooled output of CLIPTextModel
+        prompt_embeds = prompt_embeds.pooler_output
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+        return prompt_embeds
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        prompt_2: Union[str, List[str]],
+        device: Optional[torch.device] = None,
+        num_images_per_prompt: int = 1,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        max_sequence_length: int = 512,
+        lora_scale: Optional[float] = None,
+    ):
+        r"""
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in all text-encoders
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            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.
+            lora_scale (`float`, *optional*):
+                A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+        """
+        device = device or self._execution_device
+
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, FluxLoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            if self.text_encoder is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder, lora_scale)
+            if self.text_encoder_2 is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder_2, lora_scale)
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        if prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            # We only use the pooled prompt output from the CLIPTextModel
+            pooled_prompt_embeds = self._get_clip_prompt_embeds(
+                prompt=prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+            )
+            prompt_embeds = self._get_t5_prompt_embeds(
+                prompt=prompt_2,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+            )
+
+        if self.text_encoder is not None:
+            if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder, lora_scale)
+
+        if self.text_encoder_2 is not None:
+            if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder_2, lora_scale)
+
+        dtype = (
+            self.text_encoder.dtype
+            if self.text_encoder is not None
+            else self.transformer.dtype
+        )
+        text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=dtype)
+
+        return prompt_embeds, pooled_prompt_embeds, text_ids
+
+    @torch.no_grad()
+    # Modified from diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion.LEditsPPPipelineStableDiffusion.encode_image
+    def encode_image(
+        self,
+        image,
+        dtype=None,
+        height=None,
+        width=None,
+        resize_mode="default",
+        crops_coords=None,
+    ):
+        image = self.image_processor.preprocess(
+            image=image,
+            height=height,
+            width=width,
+            resize_mode=resize_mode,
+            crops_coords=crops_coords,
+        )
+        resized = self.image_processor.postprocess(image=image, output_type="pil")
+
+        if max(image.shape[-2:]) > self.vae.config["sample_size"] * 1.5:
+            logger.warning(
+                "Your input images far exceed the default resolution of the underlying diffusion model. "
+                "The output images may contain severe artifacts! "
+                "Consider down-sampling the input using the `height` and `width` parameters"
+            )
+        image = image.to(dtype)
+
+        x0 = self.vae.encode(image.to(self.device)).latent_dist.sample()
+        x0 = (x0 - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+        x0 = x0.to(dtype)
+        return x0, resized
+
+    def check_inputs(
+        self,
+        prompt,
+        prompt_2,
+        inverted_latents,
+        image_latents,
+        latent_image_ids,
+        height,
+        width,
+        start_timestep,
+        stop_timestep,
+        prompt_embeds=None,
+        pooled_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+        max_sequence_length=None,
+    ):
+        if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0:
+            raise ValueError(
+                f"`height` and `width` have to be divisible by {self.vae_scale_factor} but are {height} and {width}."
+            )
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs
+            for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt_2 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (
+            not isinstance(prompt, str) and not isinstance(prompt, list)
+        ):
+            raise ValueError(
+                f"`prompt` has to be of type `str` or `list` but is {type(prompt)}"
+            )
+        elif prompt_2 is not None and (
+            not isinstance(prompt_2, str) and not isinstance(prompt_2, list)
+        ):
+            raise ValueError(
+                f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}"
+            )
+
+        if prompt_embeds is not None and pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
+            )
+
+        if max_sequence_length is not None and max_sequence_length > 512:
+            raise ValueError(
+                f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}"
+            )
+
+        if inverted_latents is not None and (
+            image_latents is None or latent_image_ids is None
+        ):
+            raise ValueError(
+                "If `inverted_latents` are provided, `image_latents` and `latent_image_ids` also have to be passed. "
+            )
+        # check start_timestep and stop_timestep
+        if start_timestep < 0 or start_timestep > stop_timestep:
+            raise ValueError(
+                f"`start_timestep` should be in [0, stop_timestep] but is {start_timestep}"
+            )
+
+    @staticmethod
+    def _prepare_latent_image_ids_offset(
+        batch_size, height, width, device, dtype, x_offset=0, y_offset=32
+    ):
+        latent_image_ids = torch.zeros(height, width, 3)
+        latent_image_ids[..., 1] = (
+            latent_image_ids[..., 1] + torch.arange(height)[:, None]
+        )
+        latent_image_ids[..., 2] = (
+            latent_image_ids[..., 2] + torch.arange(width)[None, :]
+        )
+
+        latent_image_ids[..., 1] += x_offset
+        latent_image_ids[..., 2] += y_offset
+
+        latent_image_id_height, latent_image_id_width, latent_image_id_channels = (
+            latent_image_ids.shape
+        )
+
+        latent_image_ids = latent_image_ids.reshape(
+            latent_image_id_height * latent_image_id_width, latent_image_id_channels
+        )
+
+        return latent_image_ids.to(device=device, dtype=dtype)
+
+    @staticmethod
+    def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
+        latent_image_ids = torch.zeros(height, width, 3)
+        latent_image_ids[..., 1] = (
+            latent_image_ids[..., 1] + torch.arange(height)[:, None]
+        )
+        latent_image_ids[..., 2] = (
+            latent_image_ids[..., 2] + torch.arange(width)[None, :]
+        )
+
+        latent_image_id_height, latent_image_id_width, latent_image_id_channels = (
+            latent_image_ids.shape
+        )
+
+        latent_image_ids = latent_image_ids.reshape(
+            latent_image_id_height * latent_image_id_width, latent_image_id_channels
+        )
+
+        return latent_image_ids.to(device=device, dtype=dtype)
+
+    @staticmethod
+    def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+        latents = latents.view(
+            batch_size, num_channels_latents, height // 2, 2, width // 2, 2
+        )
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(
+            batch_size, (height // 2) * (width // 2), num_channels_latents * 4
+        )
+
+        return latents
+
+    @staticmethod
+    def _unpack_latents(latents, height, width, vae_scale_factor):
+        batch_size, num_patches, channels = latents.shape
+
+        height = height // vae_scale_factor
+        width = width // vae_scale_factor
+
+        latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+        latents = latents.reshape(batch_size, channels // (2 * 2), height, width)
+
+        return latents
+
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        self.vae.enable_tiling()
+
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_tiling()
+
+    def prepare_latents_inversion(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        image_latents,
+    ):
+        height = int(height) // self.vae_scale_factor
+        width = int(width) // self.vae_scale_factor
+        latents = self._pack_latents(
+            image_latents, batch_size, num_channels_latents, height, width
+        )
+
+        latent_image_ids = self._prepare_latent_image_ids(
+            batch_size, height // 2, width // 2, device, dtype
+        )
+
+        return latents, latent_image_ids
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.prepare_latents
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        shape = (batch_size, num_channels_latents, height, width)
+
+        if latents is not None:
+            latent_image_ids = self._prepare_latent_image_ids(
+                batch_size, height // 2, width // 2, device, dtype
+            )
+            return latents.to(device=device, dtype=dtype), latent_image_ids
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = self._pack_latents(
+            latents, batch_size, num_channels_latents, height, width
+        )
+
+        latent_image_ids = self._prepare_latent_image_ids(
+            batch_size, height // 2, width // 2, device, dtype
+        )
+
+        return latents, latent_image_ids
+
+    # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_img2img.StableDiffusion3Img2ImgPipeline.get_timesteps
+    def get_timesteps(self, num_inference_steps, strength=1.0):
+        # get the original timestep using init_timestep
+        init_timestep = min(num_inference_steps * strength, num_inference_steps)
+
+        t_start = int(max(num_inference_steps - init_timestep, 0))
+        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
+        sigmas = self.scheduler.sigmas[t_start * self.scheduler.order :]
+        if hasattr(self.scheduler, "set_begin_index"):
+            self.scheduler.set_begin_index(t_start * self.scheduler.order)
+
+        return timesteps, sigmas, num_inference_steps - t_start
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        inverted_latents: Optional[torch.FloatTensor] = None,
+        image_latents: Optional[torch.FloatTensor] = None,
+        latent_image_ids: Optional[torch.FloatTensor] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        eta: float = 1.0,
+        decay_eta: Optional[bool] = False,
+        eta_decay_power: Optional[float] = 1.0,
+        strength: float = 1.0,
+        start_timestep: float = 0,
+        stop_timestep: float = 0.25,
+        num_inference_steps: int = 28,
+        sigmas: Optional[List[float]] = None,
+        timesteps: List[int] = 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,
+        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,
+        ###################################
+        mask = None,
+        use_timestep = False
+    ):
+        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
+            inverted_latents (`torch.Tensor`, *optional*):
+                The inverted latents from `pipe.invert`.
+            image_latents (`torch.Tensor`, *optional*):
+                The image latents from `pipe.invert`.
+            latent_image_ids (`torch.Tensor`, *optional*):
+                The latent image ids from `pipe.invert`.
+            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.
+            eta (`float`, *optional*, defaults to 1.0):
+                The controller guidance, balancing faithfulness & editability:
+                higher eta - better faithfullness, less editability. For more significant edits, lower the value of eta.
+            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.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be used. Must be in descending order.
+            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.
+            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 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,
+            inverted_latents,
+            image_latents,
+            latent_image_ids,
+            height,
+            width,
+            start_timestep,
+            stop_timestep,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=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._interrupt = False
+        do_rf_inversion = inverted_latents is not None
+
+        # 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
+        )
+        (
+            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,
+        )
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        if do_rf_inversion:
+            latents = inverted_latents
+            new_latents, _ = self.prepare_latents(
+                1,
+                num_channels_latents,
+                height,
+                width,
+                prompt_embeds.dtype,
+                device,
+                generator,
+                latents=None,
+            )
+            ###############################################
+            n_h = height // 16
+            n_w = width // 16
+            bsz, _, dim = latents.shape
+
+            new_latents = new_latents.reshape(bsz, n_h, n_w, dim)
+            first_col = new_latents[:, :, 0:1, :]
+            last_col = new_latents[:, :, -1:, :]
+            new_latents = torch.cat([last_col, new_latents, first_col], dim=2)
+            new_latents = new_latents.reshape(bsz, -1, dim)
+            ###############################################
+
+            latents = torch.cat((latents, new_latents), dim=0)
+            bsz, _, dim = latents.shape
+        else:
+            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 = (int(height) // self.vae_scale_factor // 2) * (
+            int(width) // self.vae_scale_factor // 2
+        )
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.base_image_seq_len,
+            self.scheduler.config.max_image_seq_len,
+            self.scheduler.config.base_shift,
+            self.scheduler.config.max_shift,
+        )
+        # if timesteps is not None:
+        #     sigmas = None
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            timesteps,
+            sigmas,
+            mu=mu,
+        )
+        if do_rf_inversion:
+            start_timestep = int(start_timestep * num_inference_steps)
+            stop_timestep = min(
+                int(stop_timestep * num_inference_steps), num_inference_steps
+            )
+            timesteps, sigmas, num_inference_steps = self.get_timesteps(
+                num_inference_steps, strength
+            )
+        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(latents.shape[0])
+        else:
+            guidance = None
+
+        if do_rf_inversion:
+            y_0 = image_latents.clone()
+
+        # 6. Denoising loop / Controlled Reverse ODE, Algorithm 2 from: https://arxiv.org/pdf/2410.10792
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if do_rf_inversion:
+                    # ti (current timestep) as annotated in algorithm 2 - i/num_inference_steps.
+                    t_i = 1 - t / 1000
+                    dt = torch.tensor(1 / (len(timesteps) - 1), device=device)
+
+                if self.interrupt:
+                    continue
+
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+                # 添加额外的 timestep
+                if use_timestep:
+                    self._joint_attention_kwargs = (
+                        {}
+                        if self._joint_attention_kwargs is None
+                        else self._joint_attention_kwargs
+                    )
+                    self._joint_attention_kwargs["timestep"] = timestep[0].item() / 1000
+
+                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]
+
+                latents_dtype = latents.dtype
+
+                # noise_pred.shape: torch.Size([2, 4096, 64])
+                if do_rf_inversion:
+                    v_t = -noise_pred[:-1]
+                    v_t_cond = (y_0 - latents[:-1]) / (1 - t_i)
+                    eta_t = eta if start_timestep <= i < stop_timestep else 0.0
+                    if decay_eta:
+                        eta_t = (
+                            eta_t * (1 - i / num_inference_steps) ** eta_decay_power
+                        )  # Decay eta over the loop
+                    v_hat_t = v_t + eta_t * (v_t_cond - v_t)
+                    # SDE Eq: 17 from https://arxiv.org/pdf/2410.10792
+                    latents[:-1] = latents[:-1] + v_hat_t * (sigmas[i] - sigmas[i + 1])
+
+                    latents[-1] = self.scheduler.step(
+                        noise_pred[-1], t, latents[-1], return_dict=False
+                    )[0]
+
+                else:
+                    # compute the previous noisy sample x_t -> x_t-1
+                    latents = self.scheduler.step(
+                        noise_pred, t, latents, return_dict=False
+                    )[0]
+
+                ############################################
+                # enhance image consistency
+                t_l = 0.5   # smaller values enforce stronger consistency but may cause visual discontinuities
+                if mask is not None and timestep[1].item() / 1000 >= t_l:
+                    mask = mask.to(device)
+                    latents[1] = latents[1] * (1.0-mask) + latents[0] * mask
+                ############################################
+
+                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()
+
+        ######################################################
+        latents = latents.reshape(bsz, n_h, n_w+2, dim)
+        latents = latents[:, :, 1:-1, :]
+        latents = latents.reshape(bsz, -1, dim)
+        ######################################################
+
+        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)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return FluxPipelineOutput(images=image)
+
+    @torch.no_grad()
+    def invert(
+        self,
+        image: PipelineImageInput,
+        source_prompt: str = "",
+        source_guidance_scale=0.0,
+        num_inversion_steps: int = 28,
+        strength: float = 1.0,
+        gamma: float = 0.5,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        timesteps: List[int] = None,
+        dtype: Optional[torch.dtype] = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        generator: Optional[torch.Generator] = None,
+    ):
+        r"""
+        Performs Algorithm 1: Controlled Forward ODE from https://arxiv.org/pdf/2410.10792
+        Args:
+            image (`PipelineImageInput`):
+                Input for the image(s) that are to be edited. Multiple input images have to default to the same aspect
+                ratio.
+            source_prompt (`str` or `List[str]`, *optional* defaults to an empty prompt as done in the original paper):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            source_guidance_scale (`float`, *optional*, defaults to 0.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). For this algorithm, it's better to keep it 0.
+            num_inversion_steps (`int`, *optional*, defaults to 28):
+                The number of discretization steps.
+            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.
+            gamma (`float`, *optional*, defaults to 0.5):
+                The controller guidance for the forward ODE, balancing faithfulness & editability:
+                higher eta - better faithfullness, less editability. For more significant edits, lower the value of eta.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be used. Must be in descending order.
+        """
+        dtype = dtype or self.text_encoder.dtype
+        batch_size = 1
+        self._joint_attention_kwargs = joint_attention_kwargs
+        num_channels_latents = self.transformer.config.in_channels // 4
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+        device = self._execution_device
+
+        # 1. prepare image
+        image_latents, _ = self.encode_image(
+            image, height=height, width=width, dtype=dtype
+        )
+        image_latents, latent_image_ids = self.prepare_latents_inversion(
+            batch_size,
+            num_channels_latents,
+            height,
+            width,
+            dtype,
+            device,
+            image_latents,
+        )
+        #####################################################
+        n_h = height // 16
+        n_w = width // 16
+        bsz, _, dim = image_latents.shape
+
+        image_latents = image_latents.reshape(bsz, n_h, n_w, dim)
+        first_col = image_latents[:, :, 0:1, :]
+        last_col = image_latents[:, :, -1:, :]
+        image_latents = torch.cat([last_col, image_latents, first_col], dim=2)
+        image_latents = image_latents.reshape(bsz, -1, dim)
+
+        latent_image_ids = latent_image_ids.reshape(n_h, n_w, 3)
+        first_col = latent_image_ids[:, 0:1, :]
+        last_col = latent_image_ids[:, -1:, :]
+        latent_image_ids = torch.cat([last_col, latent_image_ids, first_col], dim=1)
+        latent_image_ids = latent_image_ids.reshape(-1, 3)
+        #####################################################
+
+        # 2. prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inversion_steps, num_inversion_steps)
+        image_seq_len = (int(height) // self.vae_scale_factor // 2) * (
+            int(width) // self.vae_scale_factor // 2
+        )
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.base_image_seq_len,
+            self.scheduler.config.max_image_seq_len,
+            self.scheduler.config.base_shift,
+            self.scheduler.config.max_shift,
+        )
+        timesteps, num_inversion_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inversion_steps,
+            device,
+            timesteps,
+            sigmas,
+            mu=mu,
+        )
+        timesteps, sigmas, num_inversion_steps = self.get_timesteps(
+            num_inversion_steps, strength
+        )
+
+        # 3. prepare text embeddings
+        (
+            prompt_embeds,
+            pooled_prompt_embeds,
+            text_ids,
+        ) = self.encode_prompt(
+            prompt=source_prompt,
+            prompt_2=source_prompt,
+            device=device,
+        )
+        # 4. handle guidance
+        if self.transformer.config.guidance_embeds:
+            guidance = torch.full(
+                [1], source_guidance_scale, device=device, dtype=torch.float32
+            )
+        else:
+            guidance = None
+
+        # Eq 8 dY_t = [u_t(Y_t) + γ(u_t(Y_t|y_1) - u_t(Y_t))]dt
+        Y_t = image_latents
+        y_1 = torch.randn(
+            Y_t.shape, device=Y_t.device, dtype=Y_t.dtype, generator=generator
+        )
+        N = len(sigmas)
+
+        # forward ODE loop
+        with self.progress_bar(total=N - 1) as progress_bar:
+            for i in range(N - 1):
+                t_i = torch.tensor(i / (N), dtype=Y_t.dtype, device=device)
+                timestep = torch.tensor(t_i, dtype=Y_t.dtype, device=device).repeat(
+                    batch_size
+                )
+                # invert_timestep_list.append(timestep.item() * 1000)
+                # get the unconditional vector field
+                u_t_i = self.transformer(
+                    hidden_states=Y_t,
+                    timestep=timestep,
+                    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]
+
+                # get the conditional vector field
+                u_t_i_cond = (y_1 - Y_t) / (1 - t_i)
+
+                # controlled vector field
+                # Eq 8 dY_t = [u_t(Y_t) + γ(u_t(Y_t|y_1) - u_t(Y_t))]dt
+                u_hat_t_i = u_t_i + gamma * (u_t_i_cond - u_t_i)
+                Y_t = Y_t + u_hat_t_i * (sigmas[i] - sigmas[i + 1])
+                progress_bar.update()
+
+        # return the inverted latents (start point for the denoising loop), encoded image & latent image ids
+        return Y_t, image_latents, latent_image_ids

pa_src/utils.py

@@ -0,0 +1,106 @@
+import torch
+import numpy as np
+from PIL import Image
+import cv2
+from typing import Optional
+
+
+def shift_tensor(tensor, x):
+    shifted_tensor = torch.zeros_like(tensor)
+
+    if x > 0:
+        shifted_tensor[:, x:] = tensor[:, :-x]
+    elif x < 0:
+        shifted_tensor[:, :x] = tensor[:, -x:]
+    else:
+        shifted_tensor = tensor  # No shift for x == 0
+
+    return shifted_tensor
+
+
+def create_mask(input_image_path, w=64, h=64):
+    img = (
+        Image.open(input_image_path)
+        .resize((w, h), Image.Resampling.NEAREST)
+        .convert("L")
+    )
+    img_array = np.array(img)
+    mask = np.where(img_array == 255, 1, 0)
+    mask_tensor = torch.tensor(mask).int()
+
+    return mask_tensor
+
+
+def save_array_as_png(array, path):
+    if array.dtype != np.uint8:
+        array = (array * 255).clip(0, 255).astype(np.uint8)
+    image = Image.fromarray(array, "RGBA")
+    image.save(path)
+
+
+def convert_to_mask_inpainting(image_array, mask_path):
+    if image_array.shape[2] != 4:
+        raise ValueError("输入数组必须是 RGBA 格式")
+    mask = np.ones(image_array.shape[:2], dtype=np.uint8) * 255
+    alpha_channel = image_array[:, :, 3]
+    mask[alpha_channel != 0] = 0
+    mask_image = Image.fromarray(mask, mode="L")
+    mask_image.save(mask_path)
+
+    return mask_image
+
+
+# mask for Subject Customiztion
+def composite_images(background_path: str, mask_path: str) -> Image.Image:
+    background = Image.open(background_path).convert("RGBA")
+    mask = Image.open(mask_path).convert("L")
+
+    if background.size != mask.size:
+        mask = mask.resize(background.size)
+
+    mask_array = np.array(mask) > 128
+
+    if background.mode == "RGBA":
+        white_canvas = Image.new("RGBA", background.size, (255, 255, 255, 255))
+    else:
+        white_canvas = Image.new("RGB", background.size, (255, 255, 255))
+
+    composite = Image.composite(background, white_canvas, Image.fromarray(mask_array))
+
+    return composite.convert("RGB")
+
+
+def process_mask_array(mask_array: np.ndarray) -> Image.Image:
+    alpha = mask_array[..., 3]
+    gray_array = np.where(alpha > 0, 0, 255).astype(np.uint8)
+    mask_image = Image.fromarray(gray_array, mode="L")
+    return mask_image.convert("1")
+
+
+def process_mask(mask: Image.Image) -> Image.Image:
+    if mask.mode != "L":
+        mask = mask.convert("L")
+    return mask.point(lambda x: 1 if x > 128 else 0, mode="1")
+
+
+def merge_masks(mask1: Image.Image, mask2: Image.Image) -> Image.Image:
+    arr1 = np.array(mask1, dtype=bool)
+    arr2 = np.array(mask2, dtype=bool)
+    merged = np.logical_and(arr1, arr2)
+    return Image.fromarray(merged).convert("1")
+
+
+def save_merged_mask(
+    mask_array: np.ndarray, mask: Optional[Image.Image], output_path: str
+) -> None:
+    mask1 = process_mask_array(mask_array)
+
+    if mask is not None:
+        mask2 = process_mask(mask)
+        if mask1.size != mask2.size:
+            mask2 = mask2.resize(mask1.size, Image.NEAREST)
+        merged = merge_masks(mask1, mask2)
+    else:
+        merged = mask1
+
+    merged.save(output_path)