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

Added chroma model to the ui. Added logic to easily pull latest, use local, or use a specific version of chroma. Allow ustom name or path in the ui for custom models

Browse Source
This commit is contained in:
Jaret Burkett
2025-05-07 12:06:30 -06:00
parent d9700bdb99
commit 43cb5603ad
11 changed files with 830 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
toolkit/config_modules.py
View File

@@ -463,7 +463,7 @@ class TrainConfig:
self.blended_blur_noise = kwargs.get('blended_blur_noise', False)
ModelArch = Literal['sd1', 'sd2', 'sd3', 'sdxl', 'pixart', 'pixart_sigma', 'auraflow', 'flux', 'flex2', 'lumina2', 'vega', 'ssd', 'wan21']
ModelArch = Literal['sd1', 'sd2', 'sd3', 'sdxl', 'pixart', 'pixart_sigma', 'auraflow', 'flux', 'flex1', 'flex2', 'lumina2', 'vega', 'ssd', 'wan21']
class ModelConfig:
@@ -553,6 +553,15 @@ class ModelConfig:
# kwargs to pass to the model
self.model_kwargs = kwargs.get("model_kwargs", {})
# allow frontend to pass arch with a color like arch:tag
# but remove the tag
if self.arch is not None:
if ':' in self.arch:
self.arch = self.arch.split(':')[0]
if self.arch == "flex1":
self.arch = "flux"
# handle migrating to new model arch
if self.arch is not None:
# reverse the arch to the old style