From a55e3f84dc4c85e862827d150a089fa39730b3ba Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 24 Feb 2025 02:07:48 -0500 Subject: [PATCH] notes --- .../desktop/src-tauri/src/ml/lcm_scheduler.rs | 195 ++++++++++-------- 1 file changed, 107 insertions(+), 88 deletions(-) diff --git a/crates/desktop/src-tauri/src/ml/lcm_scheduler.rs b/crates/desktop/src-tauri/src/ml/lcm_scheduler.rs index 721dc2af19..c27965ac23 100644 --- a/crates/desktop/src-tauri/src/ml/lcm_scheduler.rs +++ b/crates/desktop/src-tauri/src/ml/lcm_scheduler.rs @@ -38,6 +38,11 @@ struct LcmScheduler { // Standard deviation of the initial noise distribution // NB: Only in diffusers, and seems to only hold the value "1.0". Swift implementation omits it. init_noise_sigma: f64, + + // TODO: This is set by the set_timesteps() method after construction in python. + // We might want to handle this upfront in the constructor ourselves + num_inference_steps: Option, + step_index: Option, } impl LcmScheduler { @@ -163,8 +168,8 @@ impl LcmScheduler { [x] self.betas [ ] self.custom_timesteps [x] self.final_alpha_cumprod - [ ] self.init_noise_sigma - [ ] self.num_inference_steps + [x] self.init_noise_sigma + [x] self.num_inference_steps [x] self.timesteps = None # Swift Implementation @@ -188,6 +193,8 @@ impl LcmScheduler { betas, init_noise_sigma, timesteps, + num_inference_steps: None, + step_index: None, }) } } @@ -268,113 +275,125 @@ impl Scheduler for LcmScheduler { self.init_noise_sigma } - /// Ensures interchangeability with schedulers that need to scale the denoising model input + /// Ensures interchangeability with schedulers that need to scale the denoising model input /// depending on the current timestep. (For LCM, this is a pass through no-op!) fn scale_model_input(&self, sample: Tensor, _timestep: usize) -> candle_core::Result { Ok(sample) // NB: The diffusers scheduling_lcm.py implementation is a pass through no-op ! } fn step(&mut self, model_output: &Tensor, timestep: usize, sample: &Tensor) -> candle_core::Result { + + if self.num_inference_steps.is_none() { + return Err(candle_core::Error::Msg( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + .to_string())); + } + + if self.step_index.is_none() { + //self.init_step_index(); // TODO + } + /* - # NB: This is from scheduling_lcm.py: + # NB: This is from scheduling_lcm.py: - def step( - self, - model_output: torch.Tensor, - timestep: int, - sample: torch.Tensor, - generator: Optional[torch.Generator] = None, - return_dict: bool = True, - ) -> Union[LCMSchedulerOutput, Tuple]: - """ - Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion - process from the learned model outputs (most often the predicted noise). + def step( + self, + model_output: torch.Tensor, + timestep: int, + sample: torch.Tensor, + generator: Optional[torch.Generator] = None, + return_dict: bool = True, + ) -> Union[LCMSchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). - Args: - model_output (`torch.Tensor`): - The direct output from learned diffusion model. - timestep (`float`): - The current discrete timestep in the diffusion chain. - sample (`torch.Tensor`): - A current instance of a sample created by the diffusion process. - generator (`torch.Generator`, *optional*): - A random number generator. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`. - Returns: - [`~schedulers.scheduling_utils.LCMSchedulerOutput`] or `tuple`: - If return_dict is `True`, [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] is returned, otherwise a - tuple is returned where the first element is the sample tensor. - """ - if self.num_inference_steps is None: - raise ValueError( - "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" - ) + Args: + model_output (`torch.Tensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + generator (`torch.Generator`, *optional*): + A random number generator. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`. + Returns: + [`~schedulers.scheduling_utils.LCMSchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] is returned, otherwise a + tuple is returned where the first element is the sample tensor. + """ + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) - if self.step_index is None: - self._init_step_index(timestep) + if self.step_index is None: + self._init_step_index(timestep) - # 1. get previous step value - prev_step_index = self.step_index + 1 - if prev_step_index < len(self.timesteps): - prev_timestep = self.timesteps[prev_step_index] - else: - prev_timestep = timestep + # 1. get previous step value + prev_step_index = self.step_index + 1 + if prev_step_index < len(self.timesteps): + prev_timestep = self.timesteps[prev_step_index] + else: + prev_timestep = timestep - # 2. compute alphas, betas - alpha_prod_t = self.alphas_cumprod[timestep] - alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod + # 2. compute alphas, betas + alpha_prod_t = self.alphas_cumprod[timestep] + alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod - beta_prod_t = 1 - alpha_prod_t - beta_prod_t_prev = 1 - alpha_prod_t_prev + beta_prod_t = 1 - alpha_prod_t + beta_prod_t_prev = 1 - alpha_prod_t_prev - # 3. Get scalings for boundary conditions - c_skip, c_out = self.get_scalings_for_boundary_condition_discrete(timestep) + # 3. Get scalings for boundary conditions + c_skip, c_out = self.get_scalings_for_boundary_condition_discrete(timestep) - # 4. Compute the predicted original sample x_0 based on the model parameterization - if self.config.prediction_type == "epsilon": # noise-prediction - predicted_original_sample = (sample - beta_prod_t.sqrt() * model_output) / alpha_prod_t.sqrt() - elif self.config.prediction_type == "sample": # x-prediction - predicted_original_sample = model_output - elif self.config.prediction_type == "v_prediction": # v-prediction - predicted_original_sample = alpha_prod_t.sqrt() * sample - beta_prod_t.sqrt() * model_output - else: - raise ValueError( - f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or" - " `v_prediction` for `LCMScheduler`." - ) + # 4. Compute the predicted original sample x_0 based on the model parameterization + if self.config.prediction_type == "epsilon": # noise-prediction + predicted_original_sample = (sample - beta_prod_t.sqrt() * model_output) / alpha_prod_t.sqrt() + elif self.config.prediction_type == "sample": # x-prediction + predicted_original_sample = model_output + elif self.config.prediction_type == "v_prediction": # v-prediction + predicted_original_sample = alpha_prod_t.sqrt() * sample - beta_prod_t.sqrt() * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or" + " `v_prediction` for `LCMScheduler`." + ) - # 5. Clip or threshold "predicted x_0" - if self.config.thresholding: - predicted_original_sample = self._threshold_sample(predicted_original_sample) - elif self.config.clip_sample: - predicted_original_sample = predicted_original_sample.clamp( - -self.config.clip_sample_range, self.config.clip_sample_range - ) + # 5. Clip or threshold "predicted x_0" + if self.config.thresholding: + predicted_original_sample = self._threshold_sample(predicted_original_sample) + elif self.config.clip_sample: + predicted_original_sample = predicted_original_sample.clamp( + -self.config.clip_sample_range, self.config.clip_sample_range + ) - # 6. Denoise model output using boundary conditions - denoised = c_out * predicted_original_sample + c_skip * sample + # 6. Denoise model output using boundary conditions + denoised = c_out * predicted_original_sample + c_skip * sample - # 7. Sample and inject noise z ~ N(0, I) for MultiStep Inference - # Noise is not used on the final timestep of the timestep schedule. - # This also means that noise is not used for one-step sampling. - if self.step_index != self.num_inference_steps - 1: - noise = randn_tensor( - model_output.shape, generator=generator, device=model_output.device, dtype=denoised.dtype - ) - prev_sample = alpha_prod_t_prev.sqrt() * denoised + beta_prod_t_prev.sqrt() * noise - else: - prev_sample = denoised + # 7. Sample and inject noise z ~ N(0, I) for MultiStep Inference + # Noise is not used on the final timestep of the timestep schedule. + # This also means that noise is not used for one-step sampling. + if self.step_index != self.num_inference_steps - 1: + noise = randn_tensor( + model_output.shape, generator=generator, device=model_output.device, dtype=denoised.dtype + ) + prev_sample = alpha_prod_t_prev.sqrt() * denoised + beta_prod_t_prev.sqrt() * noise + else: + prev_sample = denoised - # upon completion increase step index by one - self._step_index += 1 + # upon completion increase step index by one + self._step_index += 1 - if not return_dict: - return (prev_sample, denoised) + if not return_dict: + return (prev_sample, denoised) - return LCMSchedulerOutput(prev_sample=prev_sample, denoised=denoised) - */ - todo!() + return LCMSchedulerOutput(prev_sample=prev_sample, denoised=denoised) + */ + todo!() + } } }