[Community] Implement `prompt-to-prompt` pipelines

[apolinario]

Describe the solution you'd like
Now that we have an official way to tweak cross attention #1639 , would be great to have a pipeline (be it official or community) for prompt-to-prompt and further implementations of the technique (such as EDICT).

Describe alternatives you've considered
@amirhertz official Prompt-to-Prompt implementation is built on top of diffusers 0.3.0 with their own cross attention manipulation function.
@bloc97 community prompt-to-prompt implementation already uses diffusers, but it is pinned to version 0.4.1, also with a cross attention control of their own.
@bram-w / Salesforce EDICT , that adds inversion to prompt-to-prompt (allowing you to edit real images) also uses the above as a base with some modifications for double precision for inversion.

So while alternatives exist, they require users to pin old versions of diffusers and not enjoy the latest advancements. Given this technique is very useful, having it on a pipeline within diffusers could be really great. Also could potentially leverage the technique to other models (Karlo, IF, etc.)

Additional context
InstructPix2Pix and Imagic have shown how editing real and edited images is a trend. Prompt-to-prompt is a nice tool to have on that belt for practitioners, artists and professionals.

[patrickvonplaten]

+100 - just lacking the time at the moment. I wonder whether we should do a community sprint in a week or so trying to add the most important "tweak your text prompts" pipelines.

[patrickvonplaten]

Actually taking this as an opportunity to turn the feature request into a more precise explanation of how it can be added.

In short we have now all the necessary tools to add a Pipeline like Prompt-2-prompt in a nice & clean way.
Prompt-2-prompt is an official pipeline with paper release and 1k+ stars, so IMO we should put it in src/diffusers/pipelines

What you'll need to do:

• Add a new pipeline Prompt2PromptPipeline that can be more or less copied from StableDiffusionPipeline
• As @apolinario said Refactor cross attention and allow mechanism to tweak cross attention function #1639 now allows one to add a prompt-2-prompt specific attention processor which can also accept new arguments to the __call__ function. Those can be passed via cross_attention_kwargs={"swapped_out_embeds"=..., swapped_in_embeds=...} to:

  diffusers/src/diffusers/models/unet_2d_condition.py

  Line 403 in 7436e30

  cross_attention_kwargs: Optional[Dict[str, Any]] = None,

  which will then passed to your customized prompt-2-prompt attention processor.
• The prompt-2-prompt attention processor class should reside inside the src/pipelines/prompt2prompt/pipeline_prompt2prompt.py` script and can be loaded directly into the unet via:

  diffusers/src/diffusers/models/unet_2d_condition.py

  Line 297 in 7436e30

  def set_attn_processor(self, processor: Union[AttnProcessor, Dict[str, AttnProcessor]]):

  
  => The super cool thing about the pipeline is then that it'll be compatible with all SD checkpoints on the Hub.

Very keen on guiding someone from the community through a PR, but currently don't find the time to do it

[keturn]

You may also reference InvokeAI's update for the diffusers 0.12 attention API: invoke-ai/InvokeAI#2385

A few caveats:

• we've only implemented the replacement (swap) control so far. Refinement and re-weight are yet to do.
• the attention code is only a part of what's necessary for a full pipeline; you also need a way to identify which controls to apply where. Those other parts are probably easier to take from the other implementations than InvokeAI.

[kashif]

So I had the following attention processors in mind for this variant of the prompt-to-prompt: https://github.com/cccntu/efficient-prompt-to-prompt

class CrossAttnKVProcessor:
    def __call__(
        self, attn: CrossAttention, hidden_states, key_hidden_states=None, value_hidden_state=None, attention_mask=None
    ):
        _, sequence_length, _ = hidden_states.shape
        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length)

        query = attn.to_q(hidden_states)
        query = attn.head_to_batch_dim(query)

        key_hidden_states = key_hidden_states if key_hidden_states is not None else hidden_states
        value_hidden_state = value_hidden_state if value_hidden_state is not None else hidden_states
        key = attn.to_k(key_hidden_states)
        value = attn.to_v(value_hidden_state)
        key = attn.head_to_batch_dim(key)
        value = attn.head_to_batch_dim(value)

        attention_probs = attn.get_attention_scores(query, key, attention_mask)
        hidden_states = torch.bmm(attention_probs, value)
        hidden_states = attn.batch_to_head_dim(hidden_states)

        # linear proj
        hidden_states = attn.to_out[0](hidden_states)
        # dropout
        hidden_states = attn.to_out[1](hidden_states)

        return hidden_states

class XFormersCrossAttnKVProcessor:
    def __call__(
        self, attn: CrossAttention, hidden_states, key_hidden_states=None, value_hidden_state=None, attention_mask=None
    ):
        _, sequence_length, _ = hidden_states.shape

        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length)

        query = attn.to_q(hidden_states)

        key_hidden_states = key_hidden_states if key_hidden_states is not None else hidden_states
        value_hidden_state = value_hidden_state if value_hidden_state is not None else hidden_states
        key = attn.to_k(key_hidden_states)
        value = attn.to_v(value_hidden_state)

        query = attn.head_to_batch_dim(query).contiguous()
        key = attn.head_to_batch_dim(key).contiguous()
        value = attn.head_to_batch_dim(value).contiguous()

        hidden_states = xformers.ops.memory_efficient_attention(query, key, value, attn_bias=attention_mask)
        hidden_states = hidden_states.to(query.dtype)
        hidden_states = attn.batch_to_head_dim(hidden_states)

        # linear proj
        hidden_states = attn.to_out[0](hidden_states)
        # dropout
        hidden_states = attn.to_out[1](hidden_states)

        return hidden_states


class SlicedAttnKVProcessor:
    def __init__(self, slice_size):
        self.slice_size = slice_size

    def __call__(
        self, attn: CrossAttention, hidden_states, key_hidden_states=None, value_hidden_state=None, attention_mask=None
    ):
        _, sequence_length, _ = hidden_states.shape

        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length)

        query = attn.to_q(hidden_states)
        dim = query.shape[-1]
        query = attn.head_to_batch_dim(query)

        key_hidden_states = key_hidden_states if key_hidden_states is not None else hidden_states
        value_hidden_state = value_hidden_state if value_hidden_state is not None else hidden_states
        key = attn.to_k(key_hidden_states)
        value = attn.to_v(value_hidden_state)
        key = attn.head_to_batch_dim(key)
        value = attn.head_to_batch_dim(value)

        batch_size_attention = query.shape[0]
        hidden_states = torch.zeros(
            (batch_size_attention, sequence_length, dim // attn.heads), device=query.device, dtype=query.dtype
        )

        for i in range(hidden_states.shape[0] // self.slice_size):
            start_idx = i * self.slice_size
            end_idx = (i + 1) * self.slice_size

            query_slice = query[start_idx:end_idx]
            key_slice = key[start_idx:end_idx]
            attn_mask_slice = attention_mask[start_idx:end_idx] if attention_mask is not None else None

            attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)

            attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx])

            hidden_states[start_idx:end_idx] = attn_slice

        hidden_states = attn.batch_to_head_dim(hidden_states)

        # linear proj
        hidden_states = attn.to_out[0](hidden_states)
        # dropout
        hidden_states = attn.to_out[1](hidden_states)

        return hidden_states

[patrickvonplaten]

        self, attn: CrossAttention, hidden_states, key_hidden_states=None, value_hidd

Sure, this seems reasonable, guess would be great to see it in a pipeline class directly :-)

[unography]

Is this open? Would be happy to take it up!

[kashif]

@unography yes it's open, please feel free to contribute!

[unography]

@kashif sure, will add a draft PR soon

[evinpinar]

This looks plausible thanks! Furthermore, with the xformers implementation, how can we retrieve softmaxed k*q attention map (before applying to values)? See here: https://github.com/facebookresearch/xformers/blob/5df1f0b682a5b246577f0cf40dd3b15c1a04ce50/xformers/ops/fmha/__init__.py#L149

class XFormersCrossAttnKVProcessor:
    def __call__(
        self, attn: CrossAttention, hidden_states, key_hidden_states=None, value_hidden_state=None, attention_mask=None
    ):
        _, sequence_length, _ = hidden_states.shape

        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length)

        query = attn.to_q(hidden_states)

        key_hidden_states = key_hidden_states if key_hidden_states is not None else hidden_states
        value_hidden_state = value_hidden_state if value_hidden_state is not None else hidden_states
        key = attn.to_k(key_hidden_states)
        value = attn.to_v(value_hidden_state)

        query = attn.head_to_batch_dim(query).contiguous()
        key = attn.head_to_batch_dim(key).contiguous()
        value = attn.head_to_batch_dim(value).contiguous()

        hidden_states = xformers.ops.memory_efficient_attention(query, key, value, attn_bias=attention_mask)
        hidden_states = hidden_states.to(query.dtype)
        hidden_states = attn.batch_to_head_dim(hidden_states)