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

Added flux training. Still a WIP. Wont train right without rectified flow working right

Browse Source
This commit is contained in:
Jaret Burkett
2024-08-02 15:00:30 -06:00
parent 03613c523f
commit 87ba867fdc
6 changed files with 292 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
toolkit/config_modules.py
View File

@@ -367,6 +367,7 @@ class ModelConfig:
self.is_pixart_sigma: bool = kwargs.get('is_pixart_sigma', False)
self.is_auraflow: bool = kwargs.get('is_auraflow', False)
self.is_v3: bool = kwargs.get('is_v3', False)
self.is_flux: bool = kwargs.get('is_flux', False)
if self.is_pixart_sigma:
self.is_pixart = True
self.is_ssd: bool = kwargs.get('is_ssd', False)
@@ -404,6 +405,9 @@ class ModelConfig:
self.vae_dtype = kwargs.get("vae_dtype", self.dtype)
self.te_device = kwargs.get("te_device", None)
self.te_dtype = kwargs.get("te_dtype", self.dtype)
# only for flux for now
self.quantize = kwargs.get("quantize", False)
pass