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

fix custom_flowmatch_sampler (#783)

Browse Source
This commit is contained in:
Asaf Agami
2026-04-13 18:42:05 +03:00
committed by GitHub
parent 233e292256
commit e868fca562
1 changed files with 1 additions and 1 deletions
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()