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

Some work on sd3 training. Not working

Browse Source
This commit is contained in:
Jaret Burkett
2024-06-13 12:19:16 -06:00
parent cb5d28cba9
commit bd10d2d668
12 changed files with 306 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
toolkit/sampler.py
View File

@@ -13,9 +13,12 @@ from diffusers import (
HeunDiscreteScheduler,
KDPM2DiscreteScheduler,
KDPM2AncestralDiscreteScheduler,
LCMScheduler
LCMScheduler,
FlowMatchEulerDiscreteScheduler,
)
from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler
from k_diffusion.external import CompVisDenoiser
from toolkit.samplers.custom_lcm_scheduler import CustomLCMScheduler
@@ -112,6 +115,15 @@ def get_sampler(
scheduler_cls = LCMScheduler
elif sampler == "custom_lcm":
scheduler_cls = CustomLCMScheduler
elif sampler == "flowmatch":
scheduler_cls = CustomFlowMatchEulerDiscreteScheduler
config_to_use = {
"_class_name": "FlowMatchEulerDiscreteScheduler",
"_diffusers_version": "0.29.0.dev0",
"num_train_timesteps": 1000,
"shift": 3.0
}
config = copy.deepcopy(config_to_use)
config.update(sched_init_args)