Unverified Commit d3eb3b35 authored by Patrick von Platen's avatar Patrick von Platen Committed by GitHub
Browse files

[Community] One step unet (#840)

parent e48ca0f0
......@@ -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/) | - |
#!/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
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment