Work to omprove pixart training

2026-04-26 17:29:27 +00:00 · 2024-06-23 20:46:48 +00:00
parent 5d47244c57
commit 7165f2d25a
6 changed files with 65 additions and 18 deletions
--- a/toolkit/config_modules.py
+++ b/toolkit/config_modules.py
@@ -372,6 +372,8 @@ class ModelConfig:

        # for text encoder quant. Only works with pixart currently
        self.text_encoder_bits = kwargs.get('text_encoder_bits', 8)  # 16, 8, 4
+        self.unet_path = kwargs.get("unet_path", None)
+        self.unet_sample_size = kwargs.get("unet_sample_size", None)


 class ReferenceDatasetConfig: