mirror of
https://github.com/modelscope/DiffSynth-Studio.git
synced 2026-03-24 18:28:10 +00:00
update flux training
This commit is contained in:
@@ -16,17 +16,13 @@ class FlowMatchScheduler():
|
|||||||
sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength
|
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 = torch.linspace(sigma_start, self.sigma_min, num_inference_steps)
|
||||||
self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)
|
self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)
|
||||||
|
self.timesteps = self.sigmas * self.num_train_timesteps
|
||||||
if training:
|
if training:
|
||||||
self.timesteps = torch.linspace(1000, 0, num_inference_steps)
|
x = self.timesteps
|
||||||
|
|
||||||
# prepare timestep weights
|
|
||||||
x = torch.arange(num_inference_steps, dtype=torch.float32)
|
|
||||||
y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2)
|
y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2)
|
||||||
y_shifted = y - y.min()
|
y_shifted = y - y.min()
|
||||||
bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum())
|
bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum())
|
||||||
self.linear_timesteps_weights = bsmntw_weighing
|
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):
|
def step(self, model_output, timestep, sample, to_final=False):
|
||||||
|
|||||||
Reference in New Issue
Block a user