Prior model preservation

modules/dataLoader/FluxBaseDataLoader.py

@@ -201,15 +201,6 @@ def _output_modules(self, config: TrainConfig, model: FluxModel):
         if not config.train_text_encoder_2_or_embedding():
             output_names.append('text_encoder_2_hidden_state')
 
-        sort_names = output_names + ['concept']
-        output_names = output_names + [('concept.loss_weight', 'loss_weight')]
-
-        # add for calculating loss per concept
-        if config.validation:
-            output_names.append(('concept.name', 'concept_name'))
-            output_names.append(('concept.path', 'concept_path'))
-            output_names.append(('concept.seed', 'concept_seed'))
-
         def before_cache_image_fun():
             model.to(self.temp_device)
             model.vae_to(self.train_device)
@@ -218,7 +209,6 @@ def before_cache_image_fun():
 
         return self._output_modules_from_out_names(
             output_names=output_names,
-            sort_names=sort_names,
             config=config,
             before_cache_image_fun=before_cache_image_fun,
             use_conditioning_image=True,

modules/dataLoader/HunyuanVideoBaseDataLoader.py

@@ -193,15 +193,6 @@ def _output_modules(self, config: TrainConfig, model: HunyuanVideoModel):
         if not config.train_text_encoder_2_or_embedding():
             output_names.append('text_encoder_2_pooled_state')
 
-        sort_names = output_names + ['concept']
-        output_names = output_names + [('concept.loss_weight', 'loss_weight')]
-
-        # add for calculating loss per concept
-        if config.validation:
-            output_names.append(('concept.name', 'concept_name'))
-            output_names.append(('concept.path', 'concept_path'))
-            output_names.append(('concept.seed', 'concept_seed'))
-
         def before_cache_image_fun():
             model.to(self.temp_device)
             model.vae_to(self.train_device)
@@ -210,7 +201,6 @@ def before_cache_image_fun():
 
         return self._output_modules_from_out_names(
             output_names=output_names,
-            sort_names=sort_names,
             config=config,
             before_cache_image_fun=before_cache_image_fun,
             use_conditioning_image=True,

modules/dataLoader/PixArtAlphaBaseDataLoader.py

@@ -166,15 +166,6 @@ def _output_modules(self, config: TrainConfig, model: PixArtAlphaModel):
         if not config.train_text_encoder_or_embedding():
             output_names.append('text_encoder_hidden_state')
 
-        sort_names = output_names + ['concept']
-        output_names = output_names + [('concept.loss_weight', 'loss_weight')]
-
-        # add for calculating loss per concept
-        if config.validation:
-            output_names.append(('concept.name', 'concept_name'))
-            output_names.append(('concept.path', 'concept_path'))
-            output_names.append(('concept.seed', 'concept_seed'))
-
         def before_cache_image_fun():
             model.to(self.temp_device)
             model.vae_to(self.train_device)
@@ -183,7 +174,6 @@ def before_cache_image_fun():
 
         return self._output_modules_from_out_names(
             output_names=output_names,
-            sort_names=sort_names,
             config=config,
             before_cache_image_fun=before_cache_image_fun,
             use_conditioning_image=True,

modules/dataLoader/SanaBaseDataLoader.py

@@ -159,15 +159,6 @@ def _output_modules(self, config: TrainConfig, model: SanaModel):
         if not config.train_text_encoder_or_embedding():
             output_names.append('text_encoder_hidden_state')
 
-        sort_names = output_names + ['concept']
-        output_names = output_names + [('concept.loss_weight', 'loss_weight')]
-
-        # add for calculating loss per concept
-        if config.validation:
-            output_names.append(('concept.name', 'concept_name'))
-            output_names.append(('concept.path', 'concept_path'))
-            output_names.append(('concept.seed', 'concept_seed'))
-
         def before_cache_image_fun():
             model.to(self.temp_device)
             model.vae_to(self.train_device)
@@ -176,7 +167,6 @@ def before_cache_image_fun():
 
         return self._output_modules_from_out_names(
             output_names=output_names,
-            sort_names=sort_names,
             config=config,
             before_cache_image_fun=before_cache_image_fun,
             use_conditioning_image=True,

modules/dataLoader/StableDiffusion3BaseDataLoader.py

@@ -224,15 +224,6 @@ def _output_modules(self, config: TrainConfig, model: StableDiffusion3Model):
         if not config.train_text_encoder_3_or_embedding():
             output_names.append('text_encoder_3_hidden_state')
 
-        sort_names = output_names + ['concept']
-        output_names = output_names + [('concept.loss_weight', 'loss_weight')]
-
-        # add for calculating loss per concept
-        if config.validation:
-            output_names.append(('concept.name', 'concept_name'))
-            output_names.append(('concept.path', 'concept_path'))
-            output_names.append(('concept.seed', 'concept_seed'))
-
         def before_cache_image_fun():
             model.to(self.temp_device)
             model.vae_to(self.train_device)
@@ -241,7 +232,6 @@ def before_cache_image_fun():
 
         return self._output_modules_from_out_names(
             output_names=output_names,
-            sort_names=sort_names,
             config=config,
             before_cache_image_fun=before_cache_image_fun,
             use_conditioning_image=True,

modules/dataLoader/StableDiffusionBaseDataLoader.py

@@ -170,15 +170,6 @@ def _output_modules(self, config: TrainConfig, model: StableDiffusionModel):
         if not config.train_text_encoder_or_embedding():
             output_names.append('text_encoder_hidden_state')
 
-        sort_names = output_names + ['concept']
-        output_names = output_names + [('concept.loss_weight', 'loss_weight')]
-
-        # add for calculating loss per concept
-        if config.validation:
-            output_names.append(('concept.name', 'concept_name'))
-            output_names.append(('concept.path', 'concept_path'))
-            output_names.append(('concept.seed', 'concept_seed'))
-
         def before_cache_image_fun():
             model.to(self.temp_device)
             model.vae_to(self.train_device)
@@ -187,7 +178,6 @@ def before_cache_image_fun():
 
         return self._output_modules_from_out_names(
             output_names=output_names,
-            sort_names=sort_names,
             config=config,
             before_cache_image_fun=before_cache_image_fun,
             use_conditioning_image=True,

modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py

@@ -238,6 +238,7 @@ def __output_modules(self, config: TrainConfig):
 
         sort_names = output_names + ['concept']
         output_names = output_names + [('concept.loss_weight', 'loss_weight')]
+        output_names = output_names + [('concept.type', 'concept_type')]
 
         # add for calculating loss per concept
         if config.validation:

modules/dataLoader/StableDiffusionXLBaseDataLoader.py

@@ -192,15 +192,6 @@ def _output_modules(self, config: TrainConfig, model: StableDiffusionXLModel):
             output_names.append('text_encoder_2_hidden_state')
             output_names.append('text_encoder_2_pooled_state')
 
-        sort_names = output_names + ['concept']
-        output_names = output_names + [('concept.loss_weight', 'loss_weight')]
-
-        # add for calculating loss per concept
-        if config.validation:
-            output_names.append(('concept.name', 'concept_name'))
-            output_names.append(('concept.path', 'concept_path'))
-            output_names.append(('concept.seed', 'concept_seed'))
-
         def before_cache_image_fun():
             model.to(self.temp_device)
             model.vae_to(self.train_device)
@@ -209,7 +200,6 @@ def before_cache_image_fun():
 
         return self._output_modules_from_out_names(
             output_names=output_names,
-            sort_names=sort_names,
             config=config,
             before_cache_image_fun=before_cache_image_fun,
             use_conditioning_image=True,

modules/dataLoader/WuerstchenBaseDataLoader.py

@@ -160,15 +160,6 @@ def _output_modules(self, config: TrainConfig, model: WuerstchenModel):
             if model.model_type.is_stable_cascade():
                 output_names.append('pooled_text_encoder_output')
 
-        sort_names = output_names + ['concept']
-        output_names = output_names + [('concept.loss_weight', 'loss_weight')]
-
-        # add for calculating loss per concept
-        if config.validation:
-            output_names.append(('concept.name', 'concept_name'))
-            output_names.append(('concept.path', 'concept_path'))
-            output_names.append(('concept.seed', 'concept_seed'))
-
         def before_cache_image_fun():
             model.to(self.temp_device)
             model.effnet_encoder_to(self.train_device)
@@ -177,7 +168,6 @@ def before_cache_image_fun():
 
         return self._output_modules_from_out_names(
             output_names=output_names,
-            sort_names=sort_names,
             config=config,
             before_cache_image_fun=before_cache_image_fun,
             autocast_context=[model.autocast_context],

modules/dataLoader/mixin/DataLoaderMgdsMixin.py

@@ -3,6 +3,7 @@
 
 from modules.util.config.ConceptConfig import ConceptConfig
 from modules.util.config.TrainConfig import TrainConfig
+from modules.util.enum.ConceptType import ConceptType
 from modules.util.TrainProgress import TrainProgress
 
 from mgds.MGDS import MGDS
@@ -26,7 +27,7 @@ def _create_mgds(
                 concepts = [ConceptConfig.default_values().from_dict(c) for c in json.load(f)]
 
         # choose all validation concepts, or none of them, depending on is_validation
-        concepts = [concept for concept in concepts if concept.validation_concept == is_validation]
+        concepts = [concept for concept in concepts if (ConceptType(concept.type) == ConceptType.VALIDATION) == is_validation]
 
         # convert before passing to MGDS
         concepts = [c.to_dict() for c in concepts]

modules/dataLoader/mixin/DataLoaderText2ImageMixin.py

@@ -225,15 +225,26 @@ def _inpainting_modules(self, config: TrainConfig):
 
     def _output_modules_from_out_names(
             self,
-            output_names: list[str],
-            sort_names: list[str],
+            output_names: list[str | tuple[str, str]],
             config: TrainConfig,
             before_cache_image_fun: Callable[[], None] | None = None,
             use_conditioning_image: bool = False,
             vae: AutoencoderKL | None = None,
             autocast_context: list[torch.autocast | None] = None,
             train_dtype: DataType | None = None,
     ):
+        sort_names = output_names + ['concept']
+
+        output_names = output_names + [
+            ('concept.loss_weight', 'loss_weight'),
+            ('concept.type', 'concept_type'),
+        ]
+
+        if config.validation:
+            output_names.append(('concept.name', 'concept_name'))
+            output_names.append(('concept.path', 'concept_path'))
+            output_names.append(('concept.seed', 'concept_seed'))
+
         mask_remove = RandomLatentMaskRemove(
             latent_mask_name='latent_mask', latent_conditioning_image_name='latent_conditioning_image' if use_conditioning_image else None,
             replace_probability=config.unmasked_probability, vae=vae,

modules/model/BaseModel.py

@@ -3,6 +3,7 @@
 from uuid import uuid4
 
 from modules.module.EMAModule import EMAModuleWrapper
+from modules.module.LoRAModule import LoRAModuleWrapper
 from modules.util.config.TrainConfig import TrainConfig
 from modules.util.enum.DataType import DataType
 from modules.util.enum.ModelType import ModelType
@@ -101,6 +102,10 @@ def to(self, device: torch.device):
     def eval(self):
         pass
 
+    @abstractmethod
+    def adapters(self) -> list[LoRAModuleWrapper]:
+        pass
+
     @staticmethod
     def _add_embeddings_to_prompt(
             additional_embeddings: list[BaseModelEmbedding],

modules/model/FluxModel.py

@@ -112,6 +112,13 @@ def __init__(
         self.transformer_lora = None
         self.lora_state_dict = None
 
+    def adapters(self) -> list[LoRAModuleWrapper]:
+        return [a for a in [
+            self.text_encoder_1_lora,
+            self.text_encoder_2_lora,
+            self.transformer_lora,
+        ] if a is not None]
+
     def all_embeddings(self) -> list[FluxModelEmbedding]:
         return self.additional_embeddings \
                + ([self.embedding] if self.embedding is not None else [])

modules/model/HunyuanVideoModel.py

@@ -130,6 +130,13 @@ def __init__(
         self.transformer_lora = None
         self.lora_state_dict = None
 
+    def adapters(self) -> list[LoRAModuleWrapper]:
+        return [a for a in [
+            self.text_encoder_1_lora,
+            self.text_encoder_2_lora,
+            self.transformer_lora,
+        ] if a is not None]
+
     def all_embeddings(self) -> list[HunyuanVideoModelEmbedding]:
         return self.additional_embeddings \
                + ([self.embedding] if self.embedding is not None else [])

modules/model/PixArtAlphaModel.py

@@ -94,6 +94,12 @@ def __init__(
         self.transformer_lora = None
         self.lora_state_dict = None
 
+    def adapters(self) -> list[LoRAModuleWrapper]:
+        return [a for a in [
+            self.text_encoder_lora,
+            self.transformer_lora,
+        ] if a is not None]
+
     def all_embeddings(self) -> list[PixArtAlphaModelEmbedding]:
         return self.additional_embeddings \
                + ([self.embedding] if self.embedding is not None else [])

modules/model/SanaModel.py

@@ -96,6 +96,12 @@ def __init__(
         self.transformer_lora = None
         self.lora_state_dict = None
 
+    def adapters(self) -> list[LoRAModuleWrapper]:
+        return [a for a in [
+            self.text_encoder_lora,
+            self.transformer_lora,
+        ] if a is not None]
+
     def all_embeddings(self) -> list[SanaModelEmbedding]:
         return self.additional_embeddings \
                + ([self.embedding] if self.embedding is not None else [])

modules/model/StableDiffusion3Model.py

@@ -129,6 +129,14 @@ def __init__(
         self.transformer_lora = None
         self.lora_state_dict = None
 
+    def adapters(self) -> list[LoRAModuleWrapper]:
+        return [a for a in [
+            self.text_encoder_1_lora,
+            self.text_encoder_2_lora,
+            self.text_encoder_3_lora,
+            self.transformer_lora,
+        ] if a is not None]
+
     def all_embeddings(self) -> list[StableDiffusion3ModelEmbedding]:
         return self.additional_embeddings \
                + ([self.embedding] if self.embedding is not None else [])

modules/model/StableDiffusionModel.py

@@ -90,6 +90,12 @@ def __init__(
         self.sd_config = None
         self.sd_config_filename = None
 
+    def adapters(self) -> list[LoRAModuleWrapper]:
+        return [a for a in [
+            self.text_encoder_lora,
+            self.unet_lora,
+        ] if a is not None]
+
     def all_embeddings(self) -> list[StableDiffusionModelEmbedding]:
         return self.additional_embeddings \
                + ([self.embedding] if self.embedding is not None else [])

modules/model/StableDiffusionXLModel.py

@@ -105,6 +105,13 @@ def __init__(
         self.sd_config = None
         self.sd_config_filename = None
 
+    def adapters(self) -> list[LoRAModuleWrapper]:
+        return [a for a in [
+            self.text_encoder_1_lora,
+            self.text_encoder_2_lora,
+            self.unet_lora,
+        ] if a is not None]
+
     def all_embeddings(self) -> list[StableDiffusionXLModelEmbedding]:
         return self.additional_embeddings \
                + ([self.embedding] if self.embedding is not None else [])

modules/model/WuerstchenModel.py

@@ -123,6 +123,12 @@ def __init__(
         self.prior_prior_lora = None
         self.lora_state_dict = None
 
+    def adapters(self) -> list[LoRAModuleWrapper]:
+        return [a for a in [
+            self.prior_text_encoder_lora,
+            self.prior_prior_lora,
+        ] if a is not None]
+
     def all_embeddings(self) -> list[WuerstchenModelEmbedding]:
         return self.additional_embeddings \
                + ([self.embedding] if self.embedding is not None else [])