| 206 | |
| 207 | |
| 208 | def add_noise(self, original_samples, noise, timestep): |
| 209 | if isinstance(timestep, torch.Tensor): |
| 210 | timestep = timestep.cpu() |
| 211 | # timestep: [batch_size], self.timesteps: [num_inference_steps] |
| 212 | timestep_id = torch.argmin((self.timesteps - timestep |
| 213 | .unsqueeze(-1).repeat(1, self.timesteps.shape[0]) |
| 214 | .to(self.timesteps.device)).abs(), dim=1) |
| 215 | sigma = self.sigmas[timestep_id].to(noise.device) # [batch_size] |
| 216 | # unsqueeze to match the shape of original_samples |
| 217 | shape_diff = len(original_samples.shape) - len(sigma.shape) |
| 218 | for _ in range(shape_diff): |
| 219 | sigma = sigma.unsqueeze(-1) |
| 220 | sample = (1 - sigma) * original_samples + sigma * noise |
| 221 | return sample |
| 222 | |
| 223 | |
| 224 | def training_target(self, sample, noise, timestep): |