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

Merge remote-tracking branch 'origin/main' into accelerate-multi-gpu

Browse Source
This commit is contained in:
Jaret Burkett
2025-01-26 11:19:34 -07:00
parent 5e663746b8 1188cf1e8a
commit 2141c6e06c
1 changed files with 13 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
toolkit/sampler.py
View File

@@ -76,6 +76,18 @@ pixart_config = {
"variance_type": None
}
flux_config = {
"_class_name": "FlowMatchEulerDiscreteScheduler",
"_diffusers_version": "0.30.0.dev0",
"base_image_seq_len": 256,
"base_shift": 0.5,
"max_image_seq_len": 4096,
"max_shift": 1.15,
"num_train_timesteps": 1000,
"shift": 3.0,
"use_dynamic_shifting": True
}
def get_sampler(
sampler: str,
@@ -120,12 +132,7 @@ def get_sampler(
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_to_use = copy.deepcopy(flux_config)
else:
raise ValueError(f"Sampler {sampler} not supported")