github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Reworked timestep distribution on flowmatch sampler when training.

Browse Source
This commit is contained in:
Jaret Burkett
2024-08-08 06:01:45 -06:00
parent acafe9984f
commit e69a520616
2 changed files with 27 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
jobs/process/BaseSDTrainProcess.py
View File

@@ -906,6 +906,10 @@ class BaseSDTrainProcess(BaseTrainProcess):
self.sd.noise_scheduler.set_timesteps( self.sd.noise_scheduler.set_timesteps(
num_train_timesteps, device=self.device_torch, original_inference_steps=num_train_timesteps num_train_timesteps, device=self.device_torch, original_inference_steps=num_train_timesteps
) )
elif self.train_config.noise_scheduler == 'flowmatch':
self.sd.noise_scheduler.set_train_timesteps(
num_train_timesteps, device=self.device_torch
)
else: else:
self.sd.noise_scheduler.set_timesteps( self.sd.noise_scheduler.set_timesteps(
num_train_timesteps, device=self.device_torch num_train_timesteps, device=self.device_torch

27
toolkit/samplers/custom_flowmatch_sampler.py
View File

@@ -3,6 +3,7 @@ from typing import Union
from diffusers import FlowMatchEulerDiscreteScheduler from diffusers import FlowMatchEulerDiscreteScheduler
import torch import torch
class CustomFlowMatchEulerDiscreteScheduler(FlowMatchEulerDiscreteScheduler): class CustomFlowMatchEulerDiscreteScheduler(FlowMatchEulerDiscreteScheduler):
def get_sigmas(self, timesteps: torch.Tensor, n_dim, dtype, device) -> torch.Tensor: def get_sigmas(self, timesteps: torch.Tensor, n_dim, dtype, device) -> torch.Tensor:
@@ -30,10 +31,28 @@ class CustomFlowMatchEulerDiscreteScheduler(FlowMatchEulerDiscreteScheduler):
# sigmas = get_sigmas(timesteps, n_dim=model_input.ndim, dtype=model_input.dtype) # sigmas = get_sigmas(timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
# noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise # noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise
n_dim = original_samples.ndim # timestep needs to be in [0, 1], we store them in [0, 1000]
sigmas = self.get_sigmas(timesteps, n_dim, original_samples.dtype, original_samples.device) # noisy_sample = (1 - timestep) * latent + timestep * noise
noisy_model_input = (1.0 - sigmas) * original_samples + sigmas * noise t_01 = (timesteps / 1000).to(original_samples.device)
noisy_model_input = (1 - t_01) * original_samples + t_01 * noise
# n_dim = original_samples.ndim
# sigmas = self.get_sigmas(timesteps, n_dim, original_samples.dtype, original_samples.device)
# noisy_model_input = (1.0 - sigmas) * original_samples + sigmas * noise
return noisy_model_input return noisy_model_input
def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor: def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
return sample return sample
def set_train_timesteps(self, num_timesteps, device):
# distribute them closer to center. Inference distributes them as a bias toward first
# Generate values from 0 to 1
t = torch.sigmoid(torch.randn((num_timesteps,), device=device))
# Scale and reverse the values to go from 1000 to 0
timesteps = ((1 - t) * 1000)
# Sort the timesteps in descending order
timesteps, _ = torch.sort(timesteps, descending=True)
return timesteps