`_encode_prompt` function

[patrickvonplaten]

At the moment, we have a bit of an inconsistent design with the _encode_prompt / encode_prompt functions.

We essentially have three different designs:

• 1.) The original stable diffusion 1.5 design, which means we have a private _encode_prompt function that does not use by default use torch.no_grad() and returns a single tensor
  

  diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

  Line 302 in 6427aa9

  def _encode_prompt(

• 2.) For IF we made the method public and wrapeed it into a torch.no_grad() mainly because one more or less has to run this function independently because of the size of the text encoder and the multi-stage inference process
  https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/deepfloyd_if/pipeline_if.py#L246

• 3.) For SD-XL for now we also made the method public, but did not add a torch.no_grad() decorator. The reason being here that there might be use cases where the user actually wants to compute the gradients when calling encode_prompt - e.g. when training LoRA with text encoder one could just call encode_prompt
  https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py#L212

What are our options to improve the API / unify the API here? cc @yiyixuxu @sayakpaul @williamberman @pcuenca

[sayakpaul]

I like the idea.

There's a clear inconsistency between

The original stable diffusion 1.5 design, which means we have a private _encode_prompt function that does not use by default use torch.no_grad() and returns a single tensor

and

For SD-XL for now we also made the method public, but did not add a torch.no_grad() decorator. The reason being here that there might be use cases where the user actually wants to compute the gradients when calling encode_prompt - e.g. when training LoRA with text encoder one could just call encode_prompt

as both SD and SDXL can be used for LoRA training. And FWIW, during LoRA training, we repurpose the encode_prompt() method from SDXL source. So, it's absolutely okay to decorate encode_prompt() of SDXL torch.no_grad IMO.

Maybe we can have text_utils.py and start unifying the APIs there? Does that make sense? A few things I would want to have:

• A dedicated class for handling all the methods like encode_prompt() and tokenize().
• We initialize this class from the pipeline classes and make use of the encode_prompt() method.

[yiyixuxu]

I think it makes sense to make encode_prompt a public method and wrapped inside @torch_no_grad() across all our pipelines - it might make it less convenient for training but I think it's ok to be a little bit more verbose in our training scripts.

[eliphatfs]

Though this may be a bit off-topic, I actually want the different parts in pipeline to be callable (and without no_grad) for pipelines so in training scripts we can reuse them and we don't have that many duplicate logic and worries that the training is actually off against the inference.

[sayakpaul]

Yeah I can understand that concern. This is why I think the utils what we are brainstorming here could make sense.

[patrickvonplaten]

So, first it seems like we should make a public encode method that does not have a torch.no_grad() dependency so that it can be used for training (so essentially what we're doing right now for SD-XL) => this makes sense to me.

Regarding a utils encoder class, I have the following concerns:

• encoding can be very pipeline-specific. E.g. IF and Stable Diffusion XL have their very own way of encoding. It would not make sense to build a class for them IMO. Also even classes like Stable Diffusion Instruct Pix2Pix don't have the same encoding class as Stable Diffusion
• Creating a "encoding" class would entail having to save the text encoder and tokenizer as attributes of this class which means that we would have both pipe.text_encoder and pipe.encoder.text_encoder assuming that encoder is such a utils class => this is not great as the user now has a more complex way of accessing tokenizer and encoder

I do admit that there is quite some boiler plate encoding code between SD pipeline classes though.

Do you have any design ideas how we could:

• a) reduce boiler plate code without sacrificing simplicity for the reader
• b) have a design that does not wrap text_encoder and tokenizer into a new class

maybe also cc @williamberman

[sayakpaul]

Hmm, fair enough.

I think this relates to the philosophy of Transformers. Duplication of code to promote readability. The subtle differences in some of the text encoding schemes also make it hard to approach this refactor.

[williamberman]

For the gradient question, this is the code I'd recommend. If we want it to be used in the training scripts, I'd also recommend moving encode_prompt_with_grad to a class function so we don't have to instantiate a pipeline. I think the exposed encode_prompt function should always be wrapped in no_grad because the pipelines are primarily used for inference.

class Pipeline
    @no_grad
    def encode_prompt():
        return self.encode_prompt_with_grad(...)

    @classmethod
    def encode_prompt_with_grad():
        ....

For boiler plate reduction, I don't know off the top of my head, would have to do some thought :)

[patrickvonplaten]

I don't think encode_prompt_with_grad() should be a class method - I don't want to pass the encoder etc... to it.
Encode prompt is too tied to the components of a pipeline.

I'd prefer to just stick to not adding torch.no_grad() & also to remove it from IF. This is also very much in line with PyTorch and keeping this library flexible so that encode_prompt can be used for training.

[patrickvonplaten]

So I'd be in favor of aligning everything to 3.) - happy to open a PR for this

[sayakpaul]

To summarize, it's hard to completely disentangle encode_prompt() from the pipelines. So, I think it's best to give each pipeline its encode_prompt() with clear # Copied from ... statements to help with maintenance.

Removal of torch.no_grad() should be useful. This way we also stay explicit similar to how we deal with the placement of modules on a requested device.

Am I on the same page as you?

[patrickvonplaten]

Yes, I'd be in favor of this!

[williamberman]

How many docs do we have already released that recommend using encode_prompt to re-use encode prompts? My understanding was it's considered a public api to be used in inference. I'd worry about removing no grad for it now.

I don't think I necessarily agree that encode_prompt is very dependent on the pipeline class. It doesn't mutate any state on the pipeline class, and it only requires reading two instance variables (the text encoder and the tokenizer).

If the recommended api to use the encoding from the training script is to instantiate a pipeline, we then have to manually pass None to all the other constructor arguments or call from_pretrained. We've already had success in simplifying the lora loader script where we moved instance methods to class methods where we are unecessarily calling from_pretrained and downloading weights just to call a function.

i.e.

tokenizer = ...
text_encoder = ...

pipe = Pipeline(tokenizer=tokenizer, text_encoder=text_encoder, unet=None, safety_checker=None, <... other arguments might be none here>)
pipe.encode_prompt("foo bar")

compare this to

tokenizer = ...
text_encoder = ...

pipe = Pipeline.encode_prompt_no_grad("foo bar", tokenizer, text_encoder) # could call the method something else like `encode_prompt_training` to emphasize its meant for training scripts

I'm pro moving as many things as possible away from instance methods if they don't mutate instance state.