| |
| 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 |
|
|
