| | |
| | 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.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size), |
| | ) |
| | timestep = 1 |
| |
|
| | model_output = self.unet(image, timestep).sample |
| | scheduler_output = self.scheduler.step(model_output, timestep, image).prev_sample |
| |
|
| | result = scheduler_output - scheduler_output + torch.ones_like(scheduler_output) |
| |
|
| | return result |
| |
|
