ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-28 02:01:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

Imitial lumina3 support

Browse Source
This commit is contained in:
Jaret Burkett
2025-02-08 10:59:53 -07:00
parent c6d8eedb94
commit d138f07365
8 changed files with 769 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
toolkit/sampler.py
View File

@@ -88,6 +88,23 @@ flux_config = {
"use_dynamic_shifting": True
}
lumina2_config = {
"_class_name": "FlowMatchEulerDiscreteScheduler",
"_diffusers_version": "0.33.0.dev0",
"base_image_seq_len": 256,
"base_shift": 0.5,
"invert_sigmas": False,
"max_image_seq_len": 4096,
"max_shift": 1.15,
"num_train_timesteps": 1000,
"shift": 6.0,
"shift_terminal": None,
"use_beta_sigmas": False,
"use_dynamic_shifting": False,
"use_exponential_sigmas": False,
"use_karras_sigmas": False
}
def get_sampler(
sampler: str,
@@ -132,7 +149,13 @@ def get_sampler(
scheduler_cls = CustomLCMScheduler
elif sampler == "flowmatch":
scheduler_cls = CustomFlowMatchEulerDiscreteScheduler
config_to_use = copy.deepcopy(flux_config)
if arch == "flux":
config_to_use = copy.deepcopy(flux_config)
elif arch == "lumina2":
config_to_use = copy.deepcopy(lumina2_config)
else:
# use flux by default
config_to_use = copy.deepcopy(flux_config)
else:
raise ValueError(f"Sampler {sampler} not supported")