You should also add the `register_modules` function to ensure your pipeline and its components can be saved with [`~DiffusionPipeline.save_pretrained`].  
```py
from diffusers import DiffusionPipeline
import torch