ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-30 19:21:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Wan22 14b training is working, still need tons of testing and some bug fixes

Browse Source
This commit is contained in:
Jaret Burkett
2025-08-14 13:03:27 -06:00
parent be71cc75ce
commit 3413fa537f
8 changed files with 554 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
jobs/process/BaseSDTrainProcess.py
View File

@@ -437,19 +437,24 @@ class BaseSDTrainProcess(BaseTrainProcess):
# Combine and sort the lists
combined_items = safetensors_files + directories + pt_files
combined_items.sort(key=os.path.getctime)
num_saves_to_keep = self.save_config.max_step_saves_to_keep
if hasattr(self.sd, 'max_step_saves_to_keep_multiplier'):
num_saves_to_keep *= self.sd.max_step_saves_to_keep_multiplier
# Use slicing with a check to avoid 'NoneType' error
safetensors_to_remove = safetensors_files[
:-self.save_config.max_step_saves_to_keep] if safetensors_files else []
pt_files_to_remove = pt_files[:-self.save_config.max_step_saves_to_keep] if pt_files else []
directories_to_remove = directories[:-self.save_config.max_step_saves_to_keep] if directories else []
embeddings_to_remove = embed_files[:-self.save_config.max_step_saves_to_keep] if embed_files else []
critic_to_remove = critic_items[:-self.save_config.max_step_saves_to_keep] if critic_items else []
:-num_saves_to_keep] if safetensors_files else []
pt_files_to_remove = pt_files[:-num_saves_to_keep] if pt_files else []
directories_to_remove = directories[:-num_saves_to_keep] if directories else []
embeddings_to_remove = embed_files[:-num_saves_to_keep] if embed_files else []
critic_to_remove = critic_items[:-num_saves_to_keep] if critic_items else []
items_to_remove = safetensors_to_remove + pt_files_to_remove + directories_to_remove + embeddings_to_remove + critic_to_remove
# remove all but the latest max_step_saves_to_keep
# items_to_remove = combined_items[:-self.save_config.max_step_saves_to_keep]
# items_to_remove = combined_items[:-num_saves_to_keep]
# remove duplicates
items_to_remove = list(dict.fromkeys(items_to_remove))