modelscope/DiffSynth-Painter

1import torch2 3 4class ContinuousODEScheduler():5 6    def __init__(self, num_inference_steps=100, sigma_max=700.0, sigma_min=0.002, rho=7.0):7        self.sigma_max = sigma_max8        self.sigma_min = sigma_min9        self.rho = rho10        self.set_timesteps(num_inference_steps)11 12 13    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0):14        ramp = torch.linspace(1-denoising_strength, 1, num_inference_steps)15        min_inv_rho = torch.pow(torch.tensor((self.sigma_min,)), (1 / self.rho))16        max_inv_rho = torch.pow(torch.tensor((self.sigma_max,)), (1 / self.rho))17        self.sigmas = torch.pow(max_inv_rho + ramp * (min_inv_rho - max_inv_rho), self.rho)18        self.timesteps = torch.log(self.sigmas) * 0.2519 20 21    def step(self, model_output, timestep, sample, to_final=False):22        timestep_id = torch.argmin((self.timesteps - timestep).abs())23        sigma = self.sigmas[timestep_id]24        sample *= (sigma*sigma + 1).sqrt()25        estimated_sample = -sigma / (sigma*sigma + 1).sqrt() * model_output + 1 / (sigma*sigma + 1) * sample26        if to_final or timestep_id + 1 >= len(self.timesteps):27            prev_sample = estimated_sample28        else:29            sigma_ = self.sigmas[timestep_id + 1]30            derivative = 1 / sigma * (sample - estimated_sample)31            prev_sample = sample + derivative * (sigma_ - sigma)32            prev_sample /= (sigma_*sigma_ + 1).sqrt()33        return prev_sample34    35 36    def return_to_timestep(self, timestep, sample, sample_stablized):37        # This scheduler doesn't support this function.38        pass39    40    41    def add_noise(self, original_samples, noise, timestep):42        timestep_id = torch.argmin((self.timesteps - timestep).abs())43        sigma = self.sigmas[timestep_id]44        sample = (original_samples + noise * sigma) / (sigma*sigma + 1).sqrt()45        return sample46    47 48    def training_target(self, sample, noise, timestep):49        timestep_id = torch.argmin((self.timesteps - timestep).abs())50        sigma = self.sigmas[timestep_id]51        target = (-(sigma*sigma + 1).sqrt() / sigma + 1 / (sigma*sigma + 1).sqrt() / sigma) * sample + 1 / (sigma*sigma + 1).sqrt() * noise52        return target53    54 55    def training_weight(self, timestep):56        timestep_id = torch.argmin((self.timesteps - timestep).abs())57        sigma = self.sigmas[timestep_id]58        weight = (1 + sigma*sigma).sqrt() / sigma59        return weight60