theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-19 06:48:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

update flux training

Browse Source
This commit is contained in:
Artiprocher
2024-10-14 10:00:32 +08:00
parent 22e4ae99e8
commit 4e3a184199
1 changed files with 2 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
diffsynth/schedulers/flow_match.py
View File

@@ -16,17 +16,13 @@ class FlowMatchScheduler():
sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength
self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps)
self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)
self.timesteps = self.sigmas * self.num_train_timesteps
if training:
self.timesteps = torch.linspace(1000, 0, num_inference_steps)
# prepare timestep weights
x = torch.arange(num_inference_steps, dtype=torch.float32)
x = self.timesteps
y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2)
y_shifted = y - y.min()
bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum())
self.linear_timesteps_weights = bsmntw_weighing
else:
self.timesteps = self.sigmas * self.num_train_timesteps
def step(self, model_output, timestep, sample, to_final=False):