Spaces:
Running on Zero
Running on Zero
| # Copyright (c) 2026 Bytedance Ltd. and/or its affiliate | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| """Flow-matching scheduler for Wan sampling (inference only).""" | |
| import torch | |
| class FlowMatchScheduler: | |
| def __init__( | |
| self, | |
| num_inference_steps: int = 100, | |
| num_train_timesteps: int = 1000, | |
| shift: float = 3.0, | |
| sigma_max: float = 1.0, | |
| sigma_min: float = 0.003 / 1.002, | |
| inverse_timesteps: bool = False, | |
| extra_one_step: bool = False, | |
| reverse_sigmas: bool = False, | |
| ): | |
| self.num_train_timesteps = num_train_timesteps | |
| self.shift = shift | |
| self.sigma_max = sigma_max | |
| self.sigma_min = sigma_min | |
| self.inverse_timesteps = inverse_timesteps | |
| self.extra_one_step = extra_one_step | |
| self.reverse_sigmas = reverse_sigmas | |
| self.set_timesteps(num_inference_steps) | |
| def set_timesteps( | |
| self, | |
| num_inference_steps: int = 100, | |
| denoising_strength: float = 1.0, | |
| shift: float = None, | |
| device: str = None, | |
| dtype: torch.dtype = torch.bfloat16, | |
| training: bool = False, | |
| ): | |
| if shift is not None: | |
| self.shift = shift | |
| if device is None: | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength | |
| if self.extra_one_step: | |
| self.sigmas = torch.linspace( | |
| sigma_start, self.sigma_min, num_inference_steps + 1, device=device, dtype=dtype | |
| )[:-1] | |
| else: | |
| self.sigmas = torch.linspace( | |
| sigma_start, self.sigma_min, num_inference_steps, device=device, dtype=dtype | |
| ) | |
| if self.inverse_timesteps: | |
| self.sigmas = torch.flip(self.sigmas, dims=[0]) | |
| self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) | |
| if self.reverse_sigmas: | |
| self.sigmas = 1 - self.sigmas | |
| self.timesteps = self.sigmas * self.num_train_timesteps | |
| self.training = training | |
| def get_noise_sigma(self, timestep): | |
| timestep = timestep.to(self.timesteps.device) if isinstance(timestep, torch.Tensor) else torch.tensor(timestep, device=self.timesteps.device) | |
| timestep_id = torch.argmin((self.timesteps.unsqueeze(-1) - timestep.reshape(1, -1)).abs(), dim=0) | |
| return self.sigmas[timestep_id].to(timestep.device) | |
| def step(self, model_output, timestep, sample, to_final: bool = False, **kwargs): | |
| if isinstance(timestep, torch.Tensor): | |
| timestep = timestep.cuda(non_blocking=True) | |
| timestep_id = torch.argmin((self.timesteps - timestep).abs()) | |
| sigma = self.sigmas[timestep_id] | |
| if to_final or timestep_id + 1 >= len(self.timesteps): | |
| sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0 | |
| else: | |
| sigma_ = self.sigmas[timestep_id + 1] | |
| return sample + model_output * (sigma_ - sigma) | |