Step Distillation Scheduler for accelerated inference. This scheduler works with step-distilled LoRA models to enable 4-step inference instead of the standard 30+ steps. Key differences from standard Flow Matching scheduler: 1. Uses fixed denoising step list (e.g., [10
| 20 | |
| 21 | |
| 22 | class StepDistillScheduler(SchedulerMixin, ConfigMixin): |
| 23 | """ |
| 24 | Step Distillation Scheduler for accelerated inference. |
| 25 | |
| 26 | This scheduler works with step-distilled LoRA models to enable |
| 27 | 4-step inference instead of the standard 30+ steps. |
| 28 | |
| 29 | Key differences from standard Flow Matching scheduler: |
| 30 | 1. Uses fixed denoising step list (e.g., [1000, 750, 500, 250]) |
| 31 | 2. Each step predicts a larger "jump" in the denoising process |
| 32 | 3. Works with distilled LoRA weights that learned these larger jumps |
| 33 | |
| 34 | Args: |
| 35 | num_train_timesteps: Number of training timesteps (default: 1000) |
| 36 | shift: Noise schedule shift parameter (default: 5.0 for 720p) |
| 37 | denoising_step_list: List of timesteps to denoise at |
| 38 | """ |
| 39 | |
| 40 | @register_to_config |
| 41 | def __init__( |
| 42 | self, |
| 43 | num_train_timesteps: int = 1000, |
| 44 | shift: float = 5.0, |
| 45 | denoising_step_list: List[int] = None, |
| 46 | base_image_seq_len: int = 256, |
| 47 | max_image_seq_len: int = 4096, |
| 48 | base_shift: float = 0.5, |
| 49 | max_shift: float = 1.16, |
| 50 | ): |
| 51 | if denoising_step_list is None: |
| 52 | # Default 4-step distillation schedule |
| 53 | denoising_step_list = [1000, 750, 500, 250] |
| 54 | |
| 55 | self.num_train_timesteps = num_train_timesteps |
| 56 | self.shift = shift |
| 57 | self.denoising_step_list = denoising_step_list |
| 58 | self.infer_steps = len(denoising_step_list) |
| 59 | |
| 60 | self.sigma_max = 1.0 |
| 61 | self.sigma_min = 0.0 |
| 62 | |
| 63 | # Initialize timesteps and sigmas |
| 64 | self.timesteps = None |
| 65 | self.sigmas = None |
| 66 | self._step_index = None # Use _step_index for compatibility with pipeline |
| 67 | |
| 68 | self.base_image_seq_len = base_image_seq_len |
| 69 | self.max_image_seq_len = max_image_seq_len |
| 70 | self.base_shift = base_shift |
| 71 | self.max_shift = max_shift |
| 72 | |
| 73 | # Required by diffusers pipeline |
| 74 | self.order = 1 # First-order method |
| 75 | |
| 76 | @property |
| 77 | def step_index(self): |
| 78 | """The current step index.""" |
| 79 | return self._step_index |
no outgoing calls
no test coverage detected