Coverage for supercat/models.py: 100.00%

1import torch 

2import torch.nn as nn 

3 

4from widitapp.diffusion import create_diffusion 

5 

6 

7class DiffusionPredictionModel(nn.Module): 

8 def __init__(self, diffusion_model: nn.Module, num_sampling_steps: int = 250): 

9 super().__init__() 

10 self.diffusion_model = diffusion_model 

11 self.diffusion = create_diffusion(str(num_sampling_steps)) 

12 

13 def forward(self, x): 

14 device = x.device 

15 z = torch.randn(x.shape, device=device) 

16 conditioned = x 

17 model_kwargs = dict(conditioned=conditioned) 

18 samples = self.diffusion.p_sample_loop( 

19 self.diffusion_model.forward, 

20 z.shape, 

21 z, 

22 clip_denoised=True, 

23 model_kwargs=model_kwargs, 

24 progress=True, 

25 device=device, 

26 ) 

27 return samples