New model: Microsoft Lens

modules/dataLoader/LensBaseDataLoader.py

@@ -0,0 +1,172 @@
+import os
+
+from modules.dataLoader.BaseDataLoader import BaseDataLoader
+from modules.dataLoader.mixin.DataLoaderText2ImageMixin import DataLoaderText2ImageMixin
+from modules.model.LensModel import (
+    PROMPT_MAX_LENGTH,
+    PROMPT_TEMPLATE_CROP_START,
+    LensModel,
+    make_lens_conversation,
+)
+from modules.modelSetup.BaseLensSetup import BaseLensSetup
+from modules.util import factory
+from modules.util.config.TrainConfig import TrainConfig
+from modules.util.enum.ModelType import ModelType
+from modules.util.TrainProgress import TrainProgress
+
+from mgds.pipelineModules.DecodeTokens import DecodeTokens
+from mgds.pipelineModules.DecodeVAE import DecodeVAE
+from mgds.pipelineModules.EncodeLensText import EncodeLensText
+from mgds.pipelineModules.EncodeVAE import EncodeVAE
+from mgds.pipelineModules.PadMaskedTokens import PadMaskedTokens
+from mgds.pipelineModules.PruneMaskedTokens import PruneMaskedTokens
+from mgds.pipelineModules.RescaleImageChannels import RescaleImageChannels
+from mgds.pipelineModules.SampleVAEDistribution import SampleVAEDistribution
+from mgds.pipelineModules.SaveImage import SaveImage
+from mgds.pipelineModules.SaveText import SaveText
+from mgds.pipelineModules.ScaleImage import ScaleImage
+from mgds.pipelineModules.Tokenize import Tokenize
+
+
+@factory.register(BaseDataLoader, ModelType.LENS)
+class LensBaseDataLoader(
+    BaseDataLoader,
+    DataLoaderText2ImageMixin,
+):
+    def _preparation_modules(self, config: TrainConfig, model: LensModel):
+        rescale_image = RescaleImageChannels(image_in_name='image', image_out_name='image', in_range_min=0, in_range_max=1, out_range_min=-1, out_range_max=1)
+        encode_image = EncodeVAE(in_name='image', out_name='latent_image_distribution', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype())
+        image_sample = SampleVAEDistribution(in_name='latent_image_distribution', out_name='latent_image', mode='mean')
+        downscale_mask = ScaleImage(in_name='mask', out_name='latent_mask', factor=0.125)
+
+        tokenize_prompt = Tokenize(
+            in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask',
+            tokenizer=model.tokenizer,
+            max_token_length=PROMPT_MAX_LENGTH + PROMPT_TEMPLATE_CROP_START,
+            apply_chat_template=make_lens_conversation,
+            apply_chat_template_kwargs={'add_generation_prompt': False},
+            apply_chat_template_post_process=lambda t: t.split("<|return|>")[0],
+        )
+
+        encode_prompt = EncodeLensText(
+            tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask',
+            hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask',
+            text_encoder=model.text_encoder,
+            crop_start=PROMPT_TEMPLATE_CROP_START,
+            autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype(),
+        )
+        prune_masked_tokens = PruneMaskedTokens(tokens_name='tokens', tokens_mask_name='tokens_mask', hidden_state_name='text_encoder_hidden_state')
+
+        modules = [rescale_image, encode_image, image_sample]
+        if config.masked_training or config.model_type.has_mask_input():
+            modules.append(downscale_mask)
+
+        modules += [tokenize_prompt, encode_prompt]
+
+        if config.latent_caching:
+            modules.append(prune_masked_tokens)
+
+        return modules
+
+    def _cache_modules(self, config: TrainConfig, model: LensModel, model_setup: BaseLensSetup):
+        image_split_names = ['latent_image', 'original_resolution', 'crop_offset']
+
+        if config.masked_training or config.model_type.has_mask_input():
+            image_split_names.append('latent_mask')
+
+        image_aggregate_names = ['crop_resolution', 'image_path']
+
+        text_split_names = []
+
+        sort_names = image_aggregate_names + image_split_names + [
+            'prompt', 'tokens', 'tokens_mask', 'text_encoder_hidden_state',
+            'concept'
+        ]
+
+        text_split_names += ['tokens', 'tokens_mask', 'text_encoder_hidden_state']
+
+        return self._cache_modules_from_names(
+            model, model_setup,
+            image_split_names=image_split_names,
+            image_aggregate_names=image_aggregate_names,
+            text_split_names=text_split_names,
+            sort_names=sort_names,
+            config=config,
+            text_caching=True,
+        )
+
+    def _output_modules(self, config: TrainConfig, model: LensModel, model_setup: BaseLensSetup):
+        pad_masked_tokens = PadMaskedTokens(tokens_name='tokens', tokens_mask_name='tokens_mask', hidden_state_name='text_encoder_hidden_state', max_length=PROMPT_MAX_LENGTH)
+
+        output_names = [
+            'image_path', 'latent_image',
+            'prompt',
+            'tokens',
+            'tokens_mask',
+            'original_resolution', 'crop_resolution', 'crop_offset',
+        ]
+
+        if config.masked_training or config.model_type.has_mask_input():
+            output_names.append('latent_mask')
+
+        output_names.append('text_encoder_hidden_state')
+
+        output_module_list = self._output_modules_from_out_names(
+            model, model_setup,
+            output_names=output_names,
+            config=config,
+            use_conditioning_image=False,
+            vae=model.vae,
+            autocast_context=[model.autocast_context],
+            train_dtype=model.train_dtype,
+        )
+
+        if config.latent_caching:
+            output_module_list = [pad_masked_tokens] + output_module_list
+
+        return output_module_list
+
+    def _debug_modules(self, config: TrainConfig, model: LensModel):
+        debug_dir = os.path.join(config.debug_dir, "dataloader")
+
+        def before_save_fun():
+            model.vae_to(self.train_device)
+
+        decode_image = DecodeVAE(in_name='latent_image', out_name='decoded_image', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype())
+        upscale_mask = ScaleImage(in_name='latent_mask', out_name='decoded_mask', factor=8)
+        decode_prompt = DecodeTokens(in_name='tokens', out_name='decoded_prompt', tokenizer=model.tokenizer)
+        save_image = SaveImage(image_in_name='decoded_image', original_path_in_name='image_path', path=debug_dir, in_range_min=-1, in_range_max=1, before_save_fun=before_save_fun)
+        # SaveImage(image_in_name='latent_mask', original_path_in_name='image_path', path=debug_dir, in_range_min=0, in_range_max=1, before_save_fun=before_save_fun)
+        save_mask = SaveImage(image_in_name='decoded_mask', original_path_in_name='image_path', path=debug_dir, in_range_min=0, in_range_max=1, before_save_fun=before_save_fun)
+        save_prompt = SaveText(text_in_name='decoded_prompt', original_path_in_name='image_path', path=debug_dir, before_save_fun=before_save_fun)
+
+        # These modules don't really work, since they are inserted after a sorting operation that does not include this data
+        # SaveImage(image_in_name='mask', original_path_in_name='image_path', path=debug_dir, in_range_min=0, in_range_max=1),
+        # SaveImage(image_in_name='image', original_path_in_name='image_path', path=debug_dir, in_range_min=-1, in_range_max=1),
+
+        modules = []
+
+        modules.append(decode_image)
+        modules.append(save_image)
+
+        if config.masked_training or config.model_type.has_mask_input():
+            modules.append(upscale_mask)
+            modules.append(save_mask)
+
+        modules.append(decode_prompt)
+        modules.append(save_prompt)
+
+        return modules
+
+    def _create_dataset(
+            self,
+            config: TrainConfig,
+            model: LensModel,
+            model_setup: BaseLensSetup,
+            train_progress: TrainProgress,
+            is_validation: bool = False,
+    ):
+        return DataLoaderText2ImageMixin._create_dataset(self,
+            config, model, model_setup, train_progress, is_validation,
+            aspect_bucketing_quantization=64,
+        )

modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py

@@ -55,7 +55,7 @@ def _setup_cache_device(
             temp_device: torch.device,
             config: TrainConfig,
     ):
-        model.to(self.temp_device)
+        model.release()
 
         model.vae_to(train_device)
 

modules/dataLoader/WuerstchenBaseDataLoader.py

@@ -74,7 +74,7 @@ def _cache_modules(self, config: TrainConfig, model: WuerstchenModel, model_setu
         ]
 
         def before_cache_image_fun():
-            model.to(self.temp_device)
+            model.release()
             model.effnet_encoder_to(self.train_device)
             model.eval()
             torch_gc()
@@ -109,7 +109,7 @@ def _output_modules(self, config: TrainConfig, model: WuerstchenModel, model_set
                 output_names.append('pooled_text_encoder_output')
 
         def before_cache_image_fun():
-            model.to(self.temp_device)
+            model.release()
             model.effnet_encoder_to(self.train_device)
             model.eval()
             torch_gc()

modules/dataLoader/mixin/DataLoaderText2ImageMixin.py

@@ -274,7 +274,7 @@ def _output_modules_from_out_names(
     ):
         if before_cache_image_fun is None:
             def prepare_vae():
-                model.to(self.temp_device)
+                model.release()
                 model.vae_to(self.train_device)
                 model.eval()
                 torch_gc()
@@ -340,7 +340,7 @@ def _cache_modules_from_names(
 
         if before_cache_image_fun is None:
             def prepare_vae():
-                model.to(self.temp_device)
+                model.release()
                 model.vae_to(self.train_device)
                 model.eval()
                 torch_gc()

modules/model/BaseModel.py

@@ -94,8 +94,11 @@ def __init__(
         self.autocast_context = nullcontext()
         self.train_dtype = DataType.FLOAT_32
 
+    #park the whole model on the temp device to free VRAM. Models with on-demand components
+    #(which cannot be parked, only discarded and rebuilt) override this to free those components
+    #instead of moving them.
     @abstractmethod
-    def to(self, device: torch.device):
+    def release(self):
         pass
 
     @abstractmethod

modules/model/ChromaModel.py

@@ -113,8 +113,7 @@ def vae_to(self, device: torch.device):
 
     def text_encoder_to(self, device: torch.device):
         if self.text_encoder is not None:
-            if self.text_encoder_offload_conductor is not None and \
-                    self.text_encoder_offload_conductor.layer_offload_activated():
+            if self.text_encoder_offload_conductor is not None:
                 self.text_encoder_offload_conductor.to(device)
             else:
                 self.text_encoder.to(device=device)
@@ -123,19 +122,18 @@ def text_encoder_to(self, device: torch.device):
             self.text_encoder_lora.to(device)
 
     def transformer_to(self, device: torch.device):
-        if self.transformer_offload_conductor is not None and \
-                self.transformer_offload_conductor.layer_offload_activated():
+        if self.transformer_offload_conductor is not None:
             self.transformer_offload_conductor.to(device)
         else:
             self.transformer.to(device=device)
 
         if self.transformer_lora is not None:
             self.transformer_lora.to(device)
 
-    def to(self, device: torch.device):
-        self.vae_to(device)
-        self.text_encoder_to(device)
-        self.transformer_to(device)
+    def release(self):
+        self.vae_to(self.train_config.temp_device)
+        self.text_encoder_to(self.train_config.temp_device)
+        self.transformer_to(self.train_config.temp_device)
 
     def eval(self):
         self.vae.eval()

modules/model/ErnieModel.py

@@ -72,26 +72,24 @@ def vae_to(self, device: torch.device):
 
     def text_encoder_to(self, device: torch.device):
         if self.text_encoder is not None:
-            if self.text_encoder_offload_conductor is not None and \
-                    self.text_encoder_offload_conductor.layer_offload_activated():
+            if self.text_encoder_offload_conductor is not None:
                 self.text_encoder_offload_conductor.to(device)
             else:
                 self.text_encoder.to(device=device)
 
     def transformer_to(self, device: torch.device):
-        if self.transformer_offload_conductor is not None and \
-                self.transformer_offload_conductor.layer_offload_activated():
+        if self.transformer_offload_conductor is not None:
             self.transformer_offload_conductor.to(device)
         else:
             self.transformer.to(device=device)
 
         if self.transformer_lora is not None:
             self.transformer_lora.to(device)
 
-    def to(self, device: torch.device):
-        self.vae_to(device)
-        self.text_encoder_to(device)
-        self.transformer_to(device)
+    def release(self):
+        self.vae_to(self.train_config.temp_device)
+        self.text_encoder_to(self.train_config.temp_device)
+        self.transformer_to(self.train_config.temp_device)
 
     def eval(self):
         self.vae.eval()

modules/model/Flux2Model.py

@@ -122,26 +122,24 @@ def vae_to(self, device: torch.device):
 
     def text_encoder_to(self, device: torch.device):
         if self.text_encoder is not None:
-            if self.text_encoder_offload_conductor is not None and \
-                    self.text_encoder_offload_conductor.layer_offload_activated():
+            if self.text_encoder_offload_conductor is not None:
                 self.text_encoder_offload_conductor.to(device)
             else:
                 self.text_encoder.to(device=device)
 
     def transformer_to(self, device: torch.device):
-        if self.transformer_offload_conductor is not None and \
-                self.transformer_offload_conductor.layer_offload_activated():
+        if self.transformer_offload_conductor is not None:
             self.transformer_offload_conductor.to(device)
         else:
             self.transformer.to(device=device)
 
         if self.transformer_lora is not None:
             self.transformer_lora.to(device)
 
-    def to(self, device: torch.device):
-        self.vae_to(device)
-        self.text_encoder_to(device)
-        self.transformer_to(device)
+    def release(self):
+        self.vae_to(self.train_config.temp_device)
+        self.text_encoder_to(self.train_config.temp_device)
+        self.transformer_to(self.train_config.temp_device)
 
     def eval(self):
         self.vae.eval()
@@ -164,7 +162,7 @@ def encode_text(
             train_device: torch.device,
             batch_size: int = 1, #TODO unused
             rand: Random | None = None,
-            text: str = None,
+            text: str | list[str] = None,
             tokens: Tensor = None,
             tokens_mask: Tensor = None,
             text_encoder_sequence_length: int | None = None,

modules/model/FluxModel.py

@@ -149,8 +149,7 @@ def text_encoder_1_to(self, device: torch.device):
 
     def text_encoder_2_to(self, device: torch.device):
         if self.text_encoder_2 is not None:
-            if self.text_encoder_2_offload_conductor is not None and \
-                    self.text_encoder_2_offload_conductor.layer_offload_activated():
+            if self.text_encoder_2_offload_conductor is not None:
                 self.text_encoder_2_offload_conductor.to(device)
             else:
                 self.text_encoder_2.to(device=device)
@@ -159,19 +158,18 @@ def text_encoder_2_to(self, device: torch.device):
             self.text_encoder_2_lora.to(device)
 
     def transformer_to(self, device: torch.device):
-        if self.transformer_offload_conductor is not None and \
-                self.transformer_offload_conductor.layer_offload_activated():
+        if self.transformer_offload_conductor is not None:
             self.transformer_offload_conductor.to(device)
         else:
             self.transformer.to(device=device)
 
         if self.transformer_lora is not None:
             self.transformer_lora.to(device)
 
-    def to(self, device: torch.device):
-        self.vae_to(device)
-        self.text_encoder_to(device)
-        self.transformer_to(device)
+    def release(self):
+        self.vae_to(self.train_config.temp_device)
+        self.text_encoder_to(self.train_config.temp_device)
+        self.transformer_to(self.train_config.temp_device)
 
     def eval(self):
         self.vae.eval()