Add clip skip for diffusion pipeline

[NormXU]

Introduction

clip skip is a trick to feed the early-stopped features encoded by CLIPTextModel into the cross-attention. If clip_skip = 2, it means that we want to use the features from the layer before the last of the clip text encoder to guide our image generation. And our current diffusion pipeline can be regarded as clip_skip = 1, which means that we just use the feature from the last layer of clip text encoder.

Here is a brief introduction to clip skip webui-wiki and related discussion link

A dominant majority of models need clip_skip=2 to reach a more aesthetic generation. I think adding this feature can give people more choices to optimize their generation.

Implementation

Adding clip_skip into diffusers is both simple and difficult. The main idea of clip_skip is simple, however, since our text encoder is imported from transformers, it is not easy to hack the CLIPTexModel in diffusers.

To do so, we need to overwrite CLIPTextModel and CLIPTextTransformer. Here is my implementation:

class MyCLIPTextTransformer(CLIPTextTransformer):

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        clip_skip: Optional[int] = 1, # <-- newly added: take the last N layer of encoder as output
    ) -> Union[Tuple, BaseModelOutputWithPooling]:

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if input_ids is None:
            raise ValueError("You have to specify input_ids")

        input_shape = input_ids.size()
        input_ids = input_ids.view(-1, input_shape[-1])

        hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)

        bsz, seq_len = input_shape
        # CLIP's text model uses causal mask, prepare it here.
        # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324
        causal_attention_mask = self._build_causal_attention_mask(bsz, seq_len, hidden_states.dtype).to(
            hidden_states.device
        )
        # expand attention_mask
        if attention_mask is not None:
            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
            attention_mask = _expand_mask(attention_mask, hidden_states.dtype)


        encoder_outputs = self.encoder(
            inputs_embeds=hidden_states,
            attention_mask=attention_mask,
            causal_attention_mask=causal_attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=True,
            return_dict=return_dict,
        )

        last_hidden_state = encoder_outputs.hidden_states[-clip_skip] # <-- newly added:  take the last N layer of encoder as output
        last_hidden_state = self.final_layer_norm(last_hidden_state)

        # text_embeds.shape = [batch_size, sequence_length, transformer.width]
        # take features from the eot embedding (eot_token is the highest number in each sequence)
        # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14
        pooled_output = last_hidden_state[
            torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
            input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1),
        ]

        if not return_dict:
            return (last_hidden_state, pooled_output) + encoder_outputs[1:]

        return BaseModelOutputWithPooling(
            last_hidden_state=last_hidden_state,
            pooler_output=pooled_output,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )

    def _build_causal_attention_mask(self, bsz, seq_len, dtype):
        # lazily create causal attention mask, with full attention between the vision tokens
        # pytorch uses additive attention mask; fill with -inf
        mask = torch.empty(bsz, seq_len, seq_len, dtype=dtype)
        mask.fill_(torch.tensor(torch.finfo(dtype).min))
        mask.triu_(1)  # zero out the lower diagonal
        mask = mask.unsqueeze(1)  # expand mask
        return mask


class MyCLIPTextModel(CLIPTextModel):
    config_class = CLIPTextConfig

    _no_split_modules = ["CLIPEncoderLayer"]

    def __init__(self, config: CLIPTextConfig):
        super().__init__(config)
        self.text_model = MyCLIPTextTransformer(config) # <-- newly added:  use the overrided clip_text_transformer
        # Initialize weights and apply final processing
        self.post_init()

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        clip_skip: Optional[int] = 1, # <-- newly added:  take the last N layer of encoder as output
    ) -> Union[Tuple, BaseModelOutputWithPooling]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        return self.text_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            clip_skip=clip_skip  # <-- newly added:  take the last N layer of encoder as output
        )

We can then use the overwritten clip_text_encoder in any_encode_prompt function of the diffuser pipeline. For example,
in pipeline_stable_diffusion.py

    def _encode_prompt(
        self,
        prompt,
        device,
        num_images_per_prompt,
        do_classifier_free_guidance,
        negative_prompt=None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
       clip_skip: Optional[int] = 1,
    ):
        # ......  # Omit Unchanged Codes

        prompt_embeds = self.text_encoder(
            text_input_ids.to(device),
            clip_skip=clip_skip,  
            attention_mask=attention_mask,
        )
        prompt_embeds = prompt_embeds[0]

        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)

        # ......  # Omit Unchanged Codes
        # this trick often applies on prompt embedding instead of negative prompt embedding
        negative_prompt_embeds = self.text_encoder(
            uncond_input.input_ids.to(device),
            attention_mask=attention_mask,
        )
        negative_prompt_embeds = negative_prompt_embeds[0]

       # ......  # Omit Unchanged Codes

        return prompt_embeds

ImplementingCLIPTextTransformer and CLIPTextModel to support clip_skip cleanly and nicely is difficult for me. I'd like to leave this issue to the diffusers team.

[patrickvonplaten]

Hey @NormXU, you could also just do the following no:

# we skip one layer of the encoder
text_encoder = CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder", num_hidden_layers=11, torch_dtype=torch.float16)

controlnet = ControlNetModel.from_pretrained(checkpoint, torch_dtype=torch.float16)

By loading the text encoder only with 11 layers you are skipping the final layer.

[NormXU]

Hey @patrickvonplaten, thank you for your quick reply.

Setting num_hidden_layers while initializing the text encoder is a good choice. However, I think it will be easier to use if we can set clip_skip as an editable parameter for diffusion forward inference, such as:

@torch.no_grad()
    def __call__(
        self,
        clip_skip: Optional[int] = 1,
    ):

Or, we need to initialize a text encoder every time we want to change the value of clip_skip.

[patrickvonplaten]

@NormXU, can you explain a bit more when people would want to change clip_skip during runtime? Usually a model performs best with one specific layer output not with multiple IMO so I don't think one would really change this during runtime. Can you explain when this is necessary?

[NormXU]

@patrickvonplaten I usually change the clip_skip to see whether I can get a better-looking generation, especially when I download a new model and want to see which layer output from the text encoder can best fit my text description. I also like to compare how the clip text encoder can influence the generation by setting clip_skip from 1 to 3,

But you are right, in my use cases, it seems that it is not necessary to introduce a new parameter into the runtime. 😂

CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder", num_hidden_layers=11, torch_dtype=torch.float16)

is good and convenient enough to use.

[sandeshrajbhandari]

@NormXU, can you explain a bit more when people would want to change clip_skip during runtime? Usually a model performs best with one specific layer output not with multiple IMO so I don't think one would really change this during runtime. Can you explain when this is necessary?

clip_skip 2 is used generally with models like Anything. using clip_skip 2 has shown positive results during image generation compared to clip_skip 1 for anime-related models. However, most users tend to use either clip_skip 1 and 2 to ideate, and using the num_hidden_layers=11 should suffice.

[DaXu999]

Any guidance on how to use this when loading models from .safetensor files:

pipeline = StableDiffusionPipeline.from_single_file( "/.safetensors",torch_dtype=torch.float16,clip_model=clip_model )

the clip_model=clip_model argument is being ignored

[patrickvonplaten]

Where do you see clip_model as being an input? I think it should be text_encoder=clip_model no?

[DaXu999]

Oh yeah, that seems to have worked. I've got a follow up question on this though. I'm running into this error when combining this with compel, I've disabled the truncation of the prompts but am running into this error on this line:
negative_conditioning = compel.build_conditioning_tensor(negative_prompt)
'index out of range in self'

[zetyquickly]

@patrickvonplaten I usually change the clip_skip to see whether I can get a better-looking generation, especially when I download a new model and want to see which layer output from the text encoder can best fit my text description. I also like to compare how the clip text encoder can influence the generation by setting clip_skip from 1 to 3,

But you are right, in my use cases, it seems that it is not necessary to introduce a new parameter into the runtime. 😂

CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder", num_hidden_layers=11, torch_dtype=torch.float16)

is good and convenient enough to use.

Hello. Given the recent introduction of new functionality, our existing approach may no longer be viable. Consider the implementation of LoraLoaderMixin:

class LoraLoaderMixin:
....
def load_lora_weights(...)
....

Should there be a mismatch between the keys of LoRA weights and the CLIPTextModel, the loading process could be disrupted. For instance, if the LoRA weights have 12 layers, but the text encoder only encompasses 11, an error similar to the following would be encountered:

...
File "/opt/conda/envs/backend-new/lib/python3.7/site-packages/diffusers/loaders.py", line 882, in load_lora_weights
    lora_scale=self.lora_scale,
  File "/opt/conda/envs/backend-new/lib/python3.7/site-packages/diffusers/loaders.py", line 1148, in load_lora_into_text_encoder
    f"failed to load text encoder state dict, unexpected keys: {load_state_dict_results.unexpected_keys}"
ValueError: failed to load text encoder state dict, unexpected keys: ['text_model.encoder.layers.11.mlp.fc1.lora_linear_layer.down.weight', ...

Given this scenario, I propose two potential solutions:

1. Incorporate clip_skip as an argument in __call__ and modify the inference flow accordingly.
2. Implement a failsafe loading mechanism for LoRA-like mixin extensions.

[Gynjn]

How can i define CLIPTextModel?

[patrickvonplaten]

Hmm is clip_skip really being that important as a feature? I still haven't seen a use case where a model produces better results with clip_skip different to what it was trained on in inference. Would love to see some concrete examples.

Also cc @sayakpaul and @yiyixuxu FYI

[neggles]

Almost every SD1.5 anime model is based off the NAI model, which was trained for a huge number of steps at CLIP Skip 2. If you run them at CLIP Skip 1 (which AIUI is a bit of a misnomer since it doesn't skip any layers, but whatever), you'll often get weird glitches in fine details or strange fractals in the background; it's been the source of quite some annoyance with some merges/finetunes.

A couple of not particularly great examples:

The second one is a better example; at CLIP skip 1 the model failed to pick up on the arms behind back in the prompt, but it followed it correctly at CLIP skip 2.

[sayakpaul]

Is this specific to anime models only? How does this approach generalize to other kinds of models?

The solution shown in the OP seems like the best approach as we cannot really change things at the end of transformers.

If we can see some diverse and varied examples where using CLIP Skip has been truly crucial, I think the original solution proposed in the OP could be directly incorporated.