New quant types, activation quantization

modules/modelLoader/GenericEmbeddingModelLoader.py

@@ -2,10 +2,10 @@
 from modules.modelLoader.BaseModelLoader import BaseModelLoader
 from modules.modelLoader.mixin.InternalModelLoaderMixin import InternalModelLoaderMixin
 from modules.modelLoader.mixin.ModelSpecModelLoaderMixin import ModelSpecModelLoaderMixin
+from modules.util.config.TrainConfig import QuantizationConfig
 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
 
 
 def make_embedding_model_loader(
@@ -33,7 +33,7 @@ def load(
                 model_type: ModelType,
                 model_names: ModelNames,
                 weight_dtypes: ModelWeightDtypes,
-                quant_filters: list[ModuleFilter] | None = None,
+                quantization: QuantizationConfig,
         ) -> model_class | None:
             base_model_loader = model_loader_class()
             embedding_loader = embedding_loader_class()
@@ -43,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, quant_filters)
+                base_model_loader.load(model, model_type, model_names, weight_dtypes, quantization)
             embedding_loader.load(model, model_names.embedding.model_name, model_names)
 
             return model

modules/modelLoader/GenericFineTuneModelLoader.py

@@ -2,10 +2,10 @@
 from modules.modelLoader.BaseModelLoader import BaseModelLoader
 from modules.modelLoader.mixin.InternalModelLoaderMixin import InternalModelLoaderMixin
 from modules.modelLoader.mixin.ModelSpecModelLoaderMixin import ModelSpecModelLoaderMixin
+from modules.util.config.TrainConfig import QuantizationConfig
 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
 
 
 def make_fine_tune_model_loader(
@@ -33,7 +33,7 @@ def load(
                 model_type: ModelType,
                 model_names: ModelNames,
                 weight_dtypes: ModelWeightDtypes,
-                quant_filters: list[ModuleFilter] | None = None,
+                quantization: QuantizationConfig,
         ) -> model_class | None:
             base_model_loader = model_loader_class()
             if embedding_loader_class is not None:
@@ -44,7 +44,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, quant_filters)
+            base_model_loader.load(model, model_type, model_names, weight_dtypes, quantization)
             if embedding_loader_class is not None:
                 embedding_loader.load(model, model_names.base_model, model_names)
 

modules/modelLoader/GenericLoRAModelLoader.py

@@ -2,10 +2,10 @@
 from modules.modelLoader.BaseModelLoader import BaseModelLoader
 from modules.modelLoader.mixin.InternalModelLoaderMixin import InternalModelLoaderMixin
 from modules.modelLoader.mixin.ModelSpecModelLoaderMixin import ModelSpecModelLoaderMixin
+from modules.util.config.TrainConfig import QuantizationConfig
 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
 
 
 def make_lora_model_loader(
@@ -34,7 +34,7 @@ def load(
                 model_type: ModelType,
                 model_names: ModelNames,
                 weight_dtypes: ModelWeightDtypes,
-                quant_filters: list[ModuleFilter] | None = None,
+                quantization: QuantizationConfig,
         ) -> model_class | None:
             base_model_loader = model_loader_class()
             lora_model_loader = lora_loader_class()
@@ -46,7 +46,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, quant_filters)
+                base_model_loader.load(model, model_type, model_names, weight_dtypes, quantization)
             lora_model_loader.load(model, model_names)
             if embedding_loader_class is not None:
                 embedding_loader.load(model, model_names.lora, model_names)

modules/modelLoader/chroma/ChromaModelLoader.py

@@ -3,11 +3,10 @@
 
 from modules.model.ChromaModel import ChromaModel
 from modules.modelLoader.mixin.HFModelLoaderMixin import HFModelLoaderMixin
-from modules.util.enum.DataType import DataType
+from modules.util.config.TrainConfig import QuantizationConfig
 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
 
 import torch
 
@@ -34,11 +33,11 @@ def __load_internal(
             base_model_name: str,
             transformer_model_name: str,
             vae_model_name: str,
-            quant_filters: list[ModuleFilter],
+            quantization: QuantizationConfig,
     ):
         if os.path.isfile(os.path.join(base_model_name, "meta.json")):
             self.__load_diffusers(
-                model, model_type, weight_dtypes, base_model_name, transformer_model_name, vae_model_name, quant_filters,
+                model, model_type, weight_dtypes, base_model_name, transformer_model_name, vae_model_name, quantization,
             )
         else:
             raise Exception("not an internal model")
@@ -51,7 +50,7 @@ def __load_diffusers(
             base_model_name: str,
             transformer_model_name: str,
             vae_model_name: str,
-            quant_filters: list[ModuleFilter],
+            quantization: QuantizationConfig,
     ):
         diffusers_sub = []
         if not transformer_model_name:
@@ -104,10 +103,10 @@ def __load_diffusers(
                 transformer_model_name,
                 #avoid loading the transformer in float32:
                 torch_dtype = torch.bfloat16 if weight_dtypes.transformer.torch_dtype() is None else weight_dtypes.transformer.torch_dtype(),
-                quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16) if weight_dtypes.transformer == DataType.GGUF else None,
+                quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16) if weight_dtypes.transformer.is_gguf() else None,
             )
             transformer = self._convert_diffusers_sub_module_to_dtype(
-                transformer, weight_dtypes.transformer, weight_dtypes.train_dtype, quant_filters,
+                transformer, weight_dtypes.transformer, weight_dtypes.train_dtype, quantization,
             )
         else:
             transformer = self._load_diffusers_sub_module(
@@ -116,7 +115,7 @@ def __load_diffusers(
                 weight_dtypes.train_dtype,
                 base_model_name,
                 "transformer",
-                quant_filters,
+                quantization,
             )
 
         model.model_type = model_type
@@ -134,7 +133,7 @@ def __load_safetensors(
             base_model_name: str,
             transformer_model_name: str,
             vae_model_name: str,
-            quant_filters: list[ModuleFilter],
+            quantization: QuantizationConfig,
     ):
         #no single file .safetensors for Chroma available at the time of writing this code
         raise NotImplementedError("Loading of single file Chroma models not supported. Use the diffusers model instead. Optionally, transformer-only safetensor files can be loaded by overriding the transformer.")
@@ -145,29 +144,29 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
-            quant_filters: list[ModuleFilter] | None = None,
+            quantization: QuantizationConfig,
     ):
         stacktraces = []
 
         try:
             self.__load_internal(
-                model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quant_filters,
+                model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization,
             )
             return
         except Exception:
             stacktraces.append(traceback.format_exc())
 
         try:
             self.__load_diffusers(
-                model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quant_filters,
+                model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization,
             )
             return
         except Exception:
             stacktraces.append(traceback.format_exc())
 
         try:
             self.__load_safetensors(
-                model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quant_filters,
+                model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization,
             )
             return
         except Exception:

modules/modelLoader/flux/FluxModelLoader.py

@@ -3,11 +3,10 @@
 
 from modules.model.FluxModel import FluxModel
 from modules.modelLoader.mixin.HFModelLoaderMixin import HFModelLoaderMixin
-from modules.util.enum.DataType import DataType
+from modules.util.config.TrainConfig import QuantizationConfig
 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
 
 import torch
 
@@ -37,12 +36,12 @@ def __load_internal(
             vae_model_name: str,
             include_text_encoder_1: bool,
             include_text_encoder_2: bool,
-            quant_filters: list[ModuleFilter],
+            quantization: QuantizationConfig,
     ):
         if os.path.isfile(os.path.join(base_model_name, "meta.json")):
             self.__load_diffusers(
                 model, model_type, weight_dtypes, base_model_name, transformer_model_name, vae_model_name,
-                include_text_encoder_1, include_text_encoder_2, quant_filters,
+                include_text_encoder_1, include_text_encoder_2, quantization,
             )
         else:
             raise Exception("not an internal model")
@@ -57,7 +56,7 @@ def __load_diffusers(
             vae_model_name: str,
             include_text_encoder_1: bool,
             include_text_encoder_2: bool,
-            quant_filters: list[ModuleFilter],
+            quantization: QuantizationConfig,
     ):
         diffusers_sub = []
         transformers_sub = []
@@ -140,10 +139,10 @@ def __load_diffusers(
                 transformer_model_name,
                 #avoid loading the transformer in float32:
                 torch_dtype = torch.bfloat16 if weight_dtypes.transformer.torch_dtype() is None else weight_dtypes.transformer.torch_dtype(),
-                quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16) if weight_dtypes.transformer == DataType.GGUF else None,
+                quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16) if weight_dtypes.transformer.is_gguf() else None,
             )
             transformer = self._convert_diffusers_sub_module_to_dtype(
-                transformer, weight_dtypes.transformer, weight_dtypes.train_dtype, quant_filters,
+                transformer, weight_dtypes.transformer, weight_dtypes.train_dtype, quantization,
             )
         else:
             transformer = self._load_diffusers_sub_module(
@@ -152,7 +151,7 @@ def __load_diffusers(
                 weight_dtypes.train_dtype,
                 base_model_name,
                 "transformer",
-                quant_filters,
+                quantization,
             )
 
         model.model_type = model_type
@@ -174,7 +173,7 @@ def __load_safetensors(
             vae_model_name: str,
             include_text_encoder_1: bool,
             include_text_encoder_2: bool,
-            quant_filters: list[ModuleFilter],
+            quantization: QuantizationConfig,
     ):
         transformer = FluxTransformer2DModel.from_single_file(
             #always load transformer separately even though FluxPipeLine.from_single_file() could load it, to avoid loading in float32:
@@ -227,7 +226,7 @@ def __load_safetensors(
             print("text encoder 2 (t5) not loaded, continuing without it")
 
         transformer = self._convert_diffusers_sub_module_to_dtype(
-            pipeline.transformer, weight_dtypes.transformer, weight_dtypes.train_dtype, quant_filters,
+            pipeline.transformer, weight_dtypes.transformer, weight_dtypes.train_dtype, quantization,
         )
 
         model.model_type = model_type
@@ -245,14 +244,14 @@ def load(
             model_type: ModelType,
             model_names: ModelNames,
             weight_dtypes: ModelWeightDtypes,
-            quant_filters: list[ModuleFilter] | None = None,
+            quantization: QuantizationConfig,
     ):
         stacktraces = []
 
         try:
             self.__load_internal(
                 model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model,
-                model_names.include_text_encoder, model_names.include_text_encoder_2, quant_filters,
+                model_names.include_text_encoder, model_names.include_text_encoder_2, quantization,
             )
             return
         except Exception:
@@ -261,7 +260,7 @@ def load(
         try:
             self.__load_diffusers(
                 model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model,
-                model_names.include_text_encoder, model_names.include_text_encoder_2, quant_filters,
+                model_names.include_text_encoder, model_names.include_text_encoder_2, quantization,
             )
             return
         except Exception:
@@ -270,7 +269,7 @@ def load(
         try:
             self.__load_safetensors(
                 model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model,
-                model_names.include_text_encoder, model_names.include_text_encoder_2, quant_filters,
+                model_names.include_text_encoder, model_names.include_text_encoder_2, quantization,
             )
             return
         except Exception: