RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix custom_flowmatch_sampler (#783)

This commit is contained in:
Asaf Agami 2026-04-13 18:42:05 +03:00 committed by GitHub
parent 233e292256
commit e868fca562
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
toolkit/samplers/custom_flowmatch_sampler.py
View File

@ -68,7 +68,7 @@ class CustomFlowMatchEulerDiscreteScheduler(FlowMatchEulerDiscreteScheduler):
device=timesteps.device, device=timesteps.device,
dtype=timesteps.dtype dtype=timesteps.dtype
) )
if v2: elif v2:
weights = self.linear_timesteps_weights2[step_indices].flatten() weights = self.linear_timesteps_weights2[step_indices].flatten()
else: else:
weights = self.linear_timesteps_weights[step_indices].flatten() weights = self.linear_timesteps_weights[step_indices].flatten()