diff --git a/examples/community/README.md b/examples/community/README.md index aada72163b2e..795fbe7efa86 100644 --- a/examples/community/README.md +++ b/examples/community/README.md @@ -5,3 +5,4 @@ | Example | Description | Author | Colab | |:----------|:----------------------|:-----------------|----------:| | CLIP Guided Stable Diffusion | Doing CLIP guidance for text to image generation with Stable Diffusion| [Suraj Patil](https://github.com/patil-suraj/) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/CLIP_Guided_Stable_diffusion_with_diffusers.ipynb) | +| One Step U-Net (Dummy) | [Patrick von Platen](https://github.com/patrickvonplaten/) | - | diff --git a/examples/community/one_step_unet.py b/examples/community/one_step_unet.py new file mode 100755 index 000000000000..425d70fe2350 --- /dev/null +++ b/examples/community/one_step_unet.py @@ -0,0 +1,22 @@ +#!/usr/bin/env python3 +import torch + +from diffusers import DiffusionPipeline + + +class UnetSchedulerOneForwardPipeline(DiffusionPipeline): + def __init__(self, unet, scheduler): + super().__init__() + + self.register_modules(unet=unet, scheduler=scheduler) + + def __call__(self): + image = torch.randn( + (1, self.unet.in_channels, self.unet.sample_size, self.unet.sample_size), + ) + timestep = 1 + + model_output = self.unet(image, timestep).sample + scheduler_output = self.scheduler.step(model_output, timestep, image).prev_sample + + return scheduler_output