[Community] One step unet (#840)
This commit is contained in:
committed by
GitHub
parent
e48ca0f0a2
commit
d3eb3b35be
@@ -5,3 +5,4 @@
|
|||||||
| Example | Description | Author | Colab |
|
| 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/) | [](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/CLIP_Guided_Stable_diffusion_with_diffusers.ipynb) |
|
| CLIP Guided Stable Diffusion | Doing CLIP guidance for text to image generation with Stable Diffusion| [Suraj Patil](https://github.com/patil-suraj/) | [](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/) | - |
|
||||||
|
|||||||
Executable
+22
@@ -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
|
||||||
Reference in New Issue
Block a user