ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-01-26 16:39:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

WIP on mean flow loss. Still a WIP.

Browse Source
This commit is contained in:
Jaret Burkett
2025-06-12 08:00:51 -06:00
parent cf11f128b9
commit fc83eb7691
6 changed files with 465 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
toolkit/sampler.py
View File

@@ -16,6 +16,7 @@ from diffusers import (
LCMScheduler,
FlowMatchEulerDiscreteScheduler,
)
from toolkit.samplers.mean_flow_scheduler import MeanFlowScheduler
from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
@@ -159,6 +160,8 @@ def get_sampler(
scheduler_cls = LCMScheduler
elif sampler == "custom_lcm":
scheduler_cls = CustomLCMScheduler
elif sampler == "mean_flow":
scheduler_cls = MeanFlowScheduler
elif sampler == "flowmatch":
scheduler_cls = CustomFlowMatchEulerDiscreteScheduler
config_to_use = copy.deepcopy(flux_config)