Coverage for wrapper/vibevoice/schedule/timestep_sampler.py: 100%

1# mypy: ignore-errors 

2# Copy from https://github.com/microsoft/VibeVoice/blob/main/vibevoice/schedule/timestep_sampler.py 

3 

4import math 

5import torch 

6 

7 

8class UniformSampler: 

9 def __init__( 

10 self, 

11 timesteps: int = 1000 

12 ) -> None: 

13 self.timesteps = timesteps 

14 

15 def sample( 

16 self, 

17 batch_size: int, 

18 device: torch.device 

19 ) -> torch.Tensor: 

20 return torch.randint(0, self.timesteps, (batch_size,), device=device) 

21 

22 

23class LogitNormalSampler: 

24 def __init__( 

25 self, 

26 timesteps: int = 1000, 

27 m: int = 0, 

28 s: int = 1 

29 ) -> None: 

30 self.timesteps = timesteps 

31 timesteps = torch.linspace(0, 1, timesteps) 

32 logit = torch.log(timesteps / (1 - timesteps)) 

33 self.prob = torch.exp(-0.5 * (logit - m) ** 2 / s ** 2) / (s * math.sqrt(2 * math.pi)) 

34 

35 def sample( 

36 self, 

37 batch_size: int, 

38 device: torch.device 

39 ) -> torch.Tensor: 

40 return torch.multinomial(self.prob, batch_size, replacement=True).to(device)