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

Added ability to set timesteps to linear for flowmatching schedule

Browse Source
This commit is contained in:
Jaret Burkett
2024-08-11 13:06:08 -06:00
parent f8f0657b68
commit a6aa4b2c7d
3 changed files with 19 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
jobs/process/BaseSDTrainProcess.py
View File

@@ -908,7 +908,9 @@ class BaseSDTrainProcess(BaseTrainProcess):
) )
elif self.train_config.noise_scheduler == 'flowmatch': elif self.train_config.noise_scheduler == 'flowmatch':
self.sd.noise_scheduler.set_train_timesteps( self.sd.noise_scheduler.set_train_timesteps(
num_train_timesteps, device=self.device_torch num_train_timesteps,
device=self.device_torch,
linear=self.train_config.linear_timesteps
) )
else: else:
self.sd.noise_scheduler.set_timesteps( self.sd.noise_scheduler.set_timesteps(

1
toolkit/config_modules.py
View File

@@ -356,6 +356,7 @@ class TrainConfig:
# adds an additional loss to the network to encourage it output a normalized standard deviation # adds an additional loss to the network to encourage it output a normalized standard deviation
self.target_norm_std = kwargs.get('target_norm_std', None) self.target_norm_std = kwargs.get('target_norm_std', None)
self.target_norm_std_value = kwargs.get('target_norm_std_value', 1.0) self.target_norm_std_value = kwargs.get('target_norm_std_value', 1.0)
self.linear_timesteps = kwargs.get('linear_timesteps', False)
class ModelConfig: class ModelConfig:

7
toolkit/samplers/custom_flowmatch_sampler.py
View File

@@ -44,7 +44,12 @@ class CustomFlowMatchEulerDiscreteScheduler(FlowMatchEulerDiscreteScheduler):
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): def set_train_timesteps(self, num_timesteps, device, linear=False):
if linear:
timesteps = torch.linspace(1000, 0, num_timesteps, device=device)
self.timesteps = timesteps
return timesteps
else:
# distribute them closer to center. Inference distributes them as a bias toward first # distribute them closer to center. Inference distributes them as a bias toward first
# Generate values from 0 to 1 # Generate values from 0 to 1
t = torch.sigmoid(torch.randn((num_timesteps,), device=device)) t = torch.sigmoid(torch.randn((num_timesteps,), device=device))