How to train a "Variance Exploding" diffusion model from scratch

[modantailleur]

I'm trying to train a "Variance Exploding" diffusion model from scratch with the diffusers library. In the train loop in the training tutorial, the model is conditioned on a simple timestep:

            # Sample a random timestep for each image

            timesteps = torch.randint(

                0, noise_scheduler.config.num_train_timesteps, (bs,), device=clean_images.device,

                dtype=torch.int64

            )

But on the ScoreSdeVePipeline , it's a sigma that is used as input of the Unet:

        for i, t in enumerate(self.progress_bar(self.scheduler.timesteps)):
            sigma_t = self.scheduler.sigmas[i] * torch.ones(shape[0], device=self.device)

            # correction step
            for _ in range(self.scheduler.config.correct_steps):
                model_output = self.unet(sample, sigma_t).sample
                sample = self.scheduler.step_correct(model_output, sample, generator=generator).prev_sample

            # prediction step
            model_output = model(sample, sigma_t).sample
            output = self.scheduler.step_pred(model_output, t, sample, generator=generator)

            sample, sample_mean = output.prev_sample, output.prev_sample_mean

How can I modify the training code to adapt it to SdeVe ? Is there examples of implementation ?