diffusers/docs/source/en/tutorials/custom_pipelines_components.md

Working with fully custom pipelines and components

Diffusers supports the use custom pipelines letting the users add any additional features on top of the [DiffusionPipeline]. However, it can get cumbersome if you're dealing with a custom pipeline where its components (such as the UNet, VAE, scheduler) are also custom.

We allow loading of such pipelines by exposing a trust_remote_code argument inside [DiffusionPipeline]. The advantage of trust_remote_code lies in its flexibility. You can have different levels of customizations for a pipeline. Following are a few examples:

• Only UNet is custom
• UNet and VAE both are custom
• Pipeline is custom
• UNet, VAE, scheduler, and pipeline are custom

With trust_remote_code=True, you can achieve perform of the above!

This tutorial covers how to author your pipeline repository so that it becomes compatible with trust_remote_code. You'll use a custom UNet, a custom scheduler, and a custom pipeline for this purpose.

You should use trust_remote_code=True only when you fully trust the code and have verified its usage.

Pipeline components

In the interest of brevity, you'll use the custom UNet, scheduler, and pipeline classes that we've already authored:

# Custom UNet
wget https://huggingface.co/sayakpaul/custom_pipeline_remote_code/raw/main/unet/my_unet_model.py
# Custom scheduler
wget https://huggingface.co/sayakpaul/custom_pipeline_remote_code/raw/main/scheduler/my_scheduler.py
# Custom pipeline
wget https://huggingface.co/sayakpaul/custom_pipeline_remote_code/raw/main/my_pipeline.py

The above classes are just for references. We encourage you to experiment with these classes for desired customizations.

Load the individual components, starting with the UNet:

from my_unet_model import MyUNetModel

pretrained_id = "hf-internal-testing/tiny-sdxl-custom-all"
unet = MyUNetModel.from_pretrained(pretrained_id, subfolder="unet")

Then go for the scheduler:

from my_scheduler import MyUNetModel

scheduler = MyScheduler.from_pretrained(pretrained_id, subfolder="scheduler")

Finally, the VAE and the text encoders:

from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
from diffusers import AutoencoderKL

text_encoder = CLIPTextModel.from_pretrained(pretrained_id, subfolder="text_encoder")
text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(pretrained_id, subfolder="text_encoder_2")
tokenizer = CLIPTokenizer.from_pretrained(pretrained_id, subfolder="tokenizer")
tokenizer_2 = CLIPTokenizer.from_pretrained(pretrained_id, subfolder="tokenizer_2")

vae = AutoencoderKL.from_pretrained(pretrained_id, subfolder="vae")

MyUNetModel, MyScheduler, and MyPipeline use blocks that are already supported by Diffusers. If you are using any custom blocks make sure to put them in the module files themselves.

Pipeline initialization and serialization

With all the components, you can now initialize the custom pipeline:

pipeline = MyPipeline(
    vae=vae, 
    unet=unet, 
    text_encoder=text_encoder, 
    text_encoder_2=text_encoder_2,
    tokenizer=tokenizer, 
    tokenizer_2=tokenizer_2, 
    scheduler=scheduler,
)

Now, push the pipeline to the Hub:

pipeline.push_to_hub("custom_pipeline_remote_code")

Since the pipeline itself is a custom pipeline, its corresponding Python module will also be pushed (example). If the pipeline has any other custom components, they will be pushed as well (UNet, scheduler).

If you want to keep the pipeline local, then use the [PushToHubMixin.save_pretrained] method.

Pipeline loading

You can load this pipeline from the Hub by specifying trust_remote_code=True:

from diffusers import DiffusionPipeline

reloaded_pipeline = DiffusionPipeline.from_pretrained(
    "sayakpaul/custom_pipeline_remote_code", 
    torch_dtype=torch.float16,
    trust_remote_code=True,
).to("cuda")

And then perform inference:

prompt = "hey"
num_inference_steps = 2

_ = reloaded_pipeline(prompt=prompt, num_inference_steps=num_inference_steps)[0]

For more complex pipelines, readers are welcome to check out this comment on GitHub.