Full finetuning with activation quantization

modules/modelLoader/ChromaEmbeddingModelLoader.py

@@ -7,6 +7,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class ChromaEmbeddingModelLoader(
@@ -32,6 +33,7 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> ChromaModel | None:
         base_model_loader = ChromaModelLoader()
         embedding_loader = ChromaEmbeddingLoader()
@@ -41,7 +43,7 @@ def load(
         model.model_spec = self._load_default_model_spec(model_type)
 
         if model_names.base_model is not None:
-            base_model_loader.load(model, model_type, model_names, weight_dtypes)
+            base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
         embedding_loader.load(model, model_names.embedding.model_name, model_names)
 
         return model

modules/modelLoader/ChromaFineTuneModelLoader.py

@@ -7,6 +7,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class ChromaFineTuneModelLoader(
@@ -32,6 +33,7 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> ChromaModel | None:
         base_model_loader = ChromaModelLoader()
         embedding_loader = ChromaEmbeddingLoader()
@@ -41,7 +43,7 @@ def load(
         self._load_internal_data(model, model_names.base_model)
         model.model_spec = self._load_default_model_spec(model_type)
 
-        base_model_loader.load(model, model_type, model_names, weight_dtypes)
+        base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
         embedding_loader.load(model, model_names.base_model, model_names)
 
         return model

modules/modelLoader/ChromaLoRAModelLoader.py

@@ -8,6 +8,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class ChromaLoRAModelLoader(
@@ -33,6 +34,7 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> ChromaModel | None:
         base_model_loader = ChromaModelLoader()
         lora_model_loader = ChromaLoRALoader()
@@ -43,7 +45,7 @@ def load(
         model.model_spec = self._load_default_model_spec(model_type)
 
         if model_names.base_model is not None:
-            base_model_loader.load(model, model_type, model_names, weight_dtypes)
+            base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
         lora_model_loader.load(model, model_names)
         embedding_loader.load(model, model_names.lora, model_names)
 

modules/modelLoader/FluxEmbeddingModelLoader.py

@@ -7,6 +7,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class FluxEmbeddingModelLoader(
@@ -34,6 +35,7 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> FluxModel | None:
         base_model_loader = FluxModelLoader()
         embedding_loader = FluxEmbeddingLoader()
@@ -43,7 +45,7 @@ def load(
         model.model_spec = self._load_default_model_spec(model_type)
 
         if model_names.base_model is not None:
-            base_model_loader.load(model, model_type, model_names, weight_dtypes)
+            base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
         embedding_loader.load(model, model_names.embedding.model_name, model_names)
 
         return model

modules/modelLoader/FluxFineTuneModelLoader.py

@@ -7,6 +7,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class FluxFineTuneModelLoader(
@@ -34,6 +35,7 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> FluxModel | None:
         base_model_loader = FluxModelLoader()
         embedding_loader = FluxEmbeddingLoader()
@@ -43,7 +45,7 @@ def load(
         self._load_internal_data(model, model_names.base_model)
         model.model_spec = self._load_default_model_spec(model_type)
 
-        base_model_loader.load(model, model_type, model_names, weight_dtypes)
+        base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
         embedding_loader.load(model, model_names.base_model, model_names)
 
         return model

modules/modelLoader/FluxLoRAModelLoader.py

@@ -8,6 +8,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class FluxLoRAModelLoader(
@@ -35,6 +36,7 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> FluxModel | None:
         base_model_loader = FluxModelLoader()
         lora_model_loader = FluxLoRALoader()
@@ -45,7 +47,7 @@ def load(
         model.model_spec = self._load_default_model_spec(model_type)
 
         if model_names.base_model is not None:
-            base_model_loader.load(model, model_type, model_names, weight_dtypes)
+            base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
         lora_model_loader.load(model, model_names)
         embedding_loader.load(model, model_names.lora, model_names)
 

modules/modelLoader/HiDreamEmbeddingModelLoader.py

@@ -7,6 +7,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class HiDreamEmbeddingModelLoader(
@@ -32,6 +33,7 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> HiDreamModel | None:
         base_model_loader = HiDreamModelLoader()
         embedding_loader = HiDreamEmbeddingLoader()
@@ -41,7 +43,7 @@ def load(
         model.model_spec = self._load_default_model_spec(model_type)
 
         if model_names.base_model is not None:
-            base_model_loader.load(model, model_type, model_names, weight_dtypes)
+            base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
         embedding_loader.load(model, model_names.embedding.model_name, model_names)
 
         return model

modules/modelLoader/HiDreamFineTuneModelLoader.py

@@ -7,6 +7,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class HiDreamFineTuneModelLoader(
@@ -32,6 +33,7 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> HiDreamModel | None:
         base_model_loader = HiDreamModelLoader()
         embedding_loader = HiDreamEmbeddingLoader()
@@ -41,7 +43,7 @@ def load(
         self._load_internal_data(model, model_names.base_model)
         model.model_spec = self._load_default_model_spec(model_type)
 
-        base_model_loader.load(model, model_type, model_names, weight_dtypes)
+        base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
         embedding_loader.load(model, model_names.base_model, model_names)
 
         return model

modules/modelLoader/HiDreamLoRAModelLoader.py

@@ -8,6 +8,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class HiDreamLoRAModelLoader(
@@ -33,6 +34,7 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> HiDreamModel | None:
         base_model_loader = HiDreamModelLoader()
         lora_model_loader = HiDreamLoRALoader()
@@ -43,7 +45,7 @@ def load(
         model.model_spec = self._load_default_model_spec(model_type)
 
         if model_names.base_model is not None:
-            base_model_loader.load(model, model_type, model_names, weight_dtypes)
+            base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
         lora_model_loader.load(model, model_names)
         embedding_loader.load(model, model_names.lora, model_names)
 

modules/modelLoader/HunyuanVideoEmbeddingModelLoader.py

@@ -7,6 +7,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class HunyuanVideoEmbeddingModelLoader(
@@ -32,6 +33,7 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> HunyuanVideoModel | None:
         base_model_loader = HunyuanVideoModelLoader()
         embedding_loader = HunyuanVideoEmbeddingLoader()
@@ -41,7 +43,7 @@ def load(
         model.model_spec = self._load_default_model_spec(model_type)
 
         if model_names.base_model is not None:
-            base_model_loader.load(model, model_type, model_names, weight_dtypes)
+            base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
         embedding_loader.load(model, model_names.embedding.model_name, model_names)
 
         return model

modules/modelLoader/HunyuanVideoFineTuneModelLoader.py

@@ -7,6 +7,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class HunyuanVideoFineTuneModelLoader(
@@ -32,6 +33,7 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> HunyuanVideoModel | None:
         base_model_loader = HunyuanVideoModelLoader()
         embedding_loader = HunyuanVideoEmbeddingLoader()
@@ -41,7 +43,7 @@ def load(
         self._load_internal_data(model, model_names.base_model)
         model.model_spec = self._load_default_model_spec(model_type)
 
-        base_model_loader.load(model, model_type, model_names, weight_dtypes)
+        base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
         embedding_loader.load(model, model_names.base_model, model_names)
 
         return model

modules/modelLoader/HunyuanVideoLoRAModelLoader.py

@@ -8,6 +8,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class HunyuanVideoLoRAModelLoader(
@@ -33,6 +34,7 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> HunyuanVideoModel | None:
         base_model_loader = HunyuanVideoModelLoader()
         lora_model_loader = HunyuanVideoLoRALoader()
@@ -43,7 +45,7 @@ def load(
         model.model_spec = self._load_default_model_spec(model_type)
 
         if model_names.base_model is not None:
-            base_model_loader.load(model, model_type, model_names, weight_dtypes)
+            base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
         lora_model_loader.load(model, model_names)
         embedding_loader.load(model, model_names.lora, model_names)
 

modules/modelLoader/PixArtAlphaEmbeddingModelLoader.py

@@ -7,6 +7,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class PixArtAlphaEmbeddingModelLoader(
@@ -34,6 +35,7 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> PixArtAlphaModel | None:
         base_model_loader = PixArtAlphaModelLoader()
         embedding_loader = PixArtAlphaEmbeddingLoader()
@@ -43,7 +45,7 @@ def load(
         model.model_spec = self._load_default_model_spec(model_type)
 
         if model_names.base_model:
-            base_model_loader.load(model, model_type, model_names, weight_dtypes)
+            base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
         embedding_loader.load(model, model_names.embedding.model_name, model_names)
 
         return model

modules/modelLoader/PixArtAlphaFineTuneModelLoader.py

@@ -7,6 +7,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class PixArtAlphaFineTuneModelLoader(
@@ -34,6 +35,7 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> PixArtAlphaModel | None:
         base_model_loader = PixArtAlphaModelLoader()
         embedding_loader = PixArtAlphaEmbeddingLoader()
@@ -43,7 +45,7 @@ def load(
         self._load_internal_data(model, model_names.base_model)
         model.model_spec = self._load_default_model_spec(model_type)
 
-        base_model_loader.load(model, model_type, model_names, weight_dtypes)
+        base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
         embedding_loader.load(model, model_names.base_model, model_names)
 
         return model

modules/modelLoader/PixArtAlphaLoRAModelLoader.py

@@ -8,6 +8,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class PixArtAlphaLoRAModelLoader(
@@ -35,6 +36,7 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> PixArtAlphaModel | None:
         base_model_loader = PixArtAlphaModelLoader()
         lora_model_loader = PixArtAlphaLoRALoader()
@@ -45,7 +47,7 @@ def load(
         model.model_spec = self._load_default_model_spec(model_type)
 
         if model_names.base_model is not None:
-            base_model_loader.load(model, model_type, model_names, weight_dtypes)
+            base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
         lora_model_loader.load(model, model_names)
         embedding_loader.load(model, model_names.lora, model_names)
 

modules/modelLoader/QwenFineTuneModelLoader.py

@@ -6,6 +6,7 @@
 from modules.util.enum.ModelType import ModelType
 from modules.util.ModelNames import ModelNames
 from modules.util.ModelWeightDtypes import ModelWeightDtypes
+from modules.util.ModuleFilter import ModuleFilter
 
 
 class QwenFineTuneModelLoader(
@@ -31,13 +32,14 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
+            quant_filters: list[ModuleFilter] | None = None,
     ) -> QwenModel | None:
         base_model_loader = QwenModelLoader()
         model = QwenModel(model_type=model_type)
 
         self._load_internal_data(model, model_names.base_model)
         model.model_spec = self._load_default_model_spec(model_type)
 
-        base_model_loader.load(model, model_type, model_names, weight_dtypes)
+        base_model_loader.load(model, model_type, model_names, weight_dtypes, quant_filters)
 
         return model