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

Adjust flux sample sampler to handle some new breaking changes in diffusers.

Browse Source
This commit is contained in:
Jaret Burkett
2025-01-26 18:09:21 +00:00
parent 441474e81f
commit 1188cf1e8a
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 "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( def get_sampler(
sampler: str, sampler: str,
@@ -120,12 +132,7 @@ def get_sampler(
scheduler_cls = CustomLCMScheduler scheduler_cls = CustomLCMScheduler
elif sampler == "flowmatch": elif sampler == "flowmatch":
scheduler_cls = CustomFlowMatchEulerDiscreteScheduler scheduler_cls = CustomFlowMatchEulerDiscreteScheduler
config_to_use = { config_to_use = copy.deepcopy(flux_config)
"_class_name": "FlowMatchEulerDiscreteScheduler",
"_diffusers_version": "0.29.0.dev0",
"num_train_timesteps": 1000,
"shift": 3.0
}
else: else:
raise ValueError(f"Sampler {sampler} not supported") raise ValueError(f"Sampler {sampler} not supported")