diff --git a/modules/dataLoader/LensBaseDataLoader.py b/modules/dataLoader/LensBaseDataLoader.py new file mode 100644 index 000000000..ade6ccfb0 --- /dev/null +++ b/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, + ) diff --git a/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py b/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py index ad0c890b0..02d7829cd 100644 --- a/modules/dataLoader/StableDiffusionFineTuneVaeDataLoader.py +++ b/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) diff --git a/modules/dataLoader/WuerstchenBaseDataLoader.py b/modules/dataLoader/WuerstchenBaseDataLoader.py index 4689b09a4..29ada84c2 100644 --- a/modules/dataLoader/WuerstchenBaseDataLoader.py +++ b/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() diff --git a/modules/dataLoader/mixin/DataLoaderText2ImageMixin.py b/modules/dataLoader/mixin/DataLoaderText2ImageMixin.py index 2654bdd19..af79459d6 100644 --- a/modules/dataLoader/mixin/DataLoaderText2ImageMixin.py +++ b/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() diff --git a/modules/model/BaseModel.py b/modules/model/BaseModel.py index fef1fa9e1..e1caca247 100644 --- a/modules/model/BaseModel.py +++ b/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 diff --git a/modules/model/ChromaModel.py b/modules/model/ChromaModel.py index 59967c8dc..bb4017f13 100644 --- a/modules/model/ChromaModel.py +++ b/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,8 +122,7 @@ 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) @@ -132,10 +130,10 @@ def transformer_to(self, device: torch.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() diff --git a/modules/model/ErnieModel.py b/modules/model/ErnieModel.py index fb44425e8..8a78d5b78 100644 --- a/modules/model/ErnieModel.py +++ b/modules/model/ErnieModel.py @@ -72,15 +72,13 @@ 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) @@ -88,10 +86,10 @@ def transformer_to(self, device: torch.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() diff --git a/modules/model/Flux2Model.py b/modules/model/Flux2Model.py index 004c70bee..58b67d3f4 100644 --- a/modules/model/Flux2Model.py +++ b/modules/model/Flux2Model.py @@ -122,15 +122,13 @@ 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) @@ -138,10 +136,10 @@ def transformer_to(self, device: torch.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, diff --git a/modules/model/FluxModel.py b/modules/model/FluxModel.py index b981865c4..d0b0c1002 100644 --- a/modules/model/FluxModel.py +++ b/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,8 +158,7 @@ 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) @@ -168,10 +166,10 @@ def transformer_to(self, device: torch.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() diff --git a/modules/model/HiDreamModel.py b/modules/model/HiDreamModel.py index 049b0e6d3..7b0b5c1d6 100644 --- a/modules/model/HiDreamModel.py +++ b/modules/model/HiDreamModel.py @@ -220,8 +220,7 @@ def text_encoder_2_to(self, device: torch.device): def text_encoder_3_to(self, device: torch.device): if self.text_encoder_3 is not None: - if self.text_encoder_3_offload_conductor is not None and \ - self.text_encoder_3_offload_conductor.layer_offload_activated(): + if self.text_encoder_3_offload_conductor is not None: self.text_encoder_3_offload_conductor.to(device) else: self.text_encoder_3.to(device=device) @@ -231,8 +230,7 @@ def text_encoder_3_to(self, device: torch.device): def text_encoder_4_to(self, device: torch.device): if self.text_encoder_4 is not None: - if self.text_encoder_4_offload_conductor is not None and \ - self.text_encoder_4_offload_conductor.layer_offload_activated(): + if self.text_encoder_4_offload_conductor is not None: self.text_encoder_4_offload_conductor.to(device) else: self.text_encoder_4.to(device=device) @@ -241,8 +239,7 @@ def text_encoder_4_to(self, device: torch.device): self.text_encoder_4_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) @@ -250,10 +247,10 @@ def transformer_to(self, device: torch.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() diff --git a/modules/model/HunyuanVideoModel.py b/modules/model/HunyuanVideoModel.py index 2107e55d7..9c0e2dc0f 100644 --- a/modules/model/HunyuanVideoModel.py +++ b/modules/model/HunyuanVideoModel.py @@ -157,8 +157,7 @@ def text_encoder_to(self, device: torch.device): def text_encoder_1_to(self, device: torch.device): if self.text_encoder_1 is not None: - if self.text_encoder_1_offload_conductor is not None and \ - self.text_encoder_1_offload_conductor.layer_offload_activated(): + if self.text_encoder_1_offload_conductor is not None: self.text_encoder_1_offload_conductor.to(device) else: self.text_encoder_1.to(device=device) @@ -174,8 +173,7 @@ 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) @@ -183,10 +181,10 @@ def transformer_to(self, device: torch.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() diff --git a/modules/model/LensModel.py b/modules/model/LensModel.py new file mode 100644 index 000000000..525aebf4f --- /dev/null +++ b/modules/model/LensModel.py @@ -0,0 +1,297 @@ +import math +from contextlib import nullcontext +from random import Random + +from modules.model.BaseModel import BaseModel +from modules.module.LoRAModule import LoRAModuleWrapper +from modules.util.enum.ModelType import ModelType +from modules.util.LayerOffloadConductor import LayerOffloadConductor +from modules.util.OnDemandModule import OnDemandModule +from modules.util.torch_util import torch_gc + +import torch +from torch import Tensor + +from diffusers import ( + AutoencoderKLFlux2, + DiffusionPipeline, + FlowMatchEulerDiscreteScheduler, +) +from transformers import PreTrainedTokenizerFast + +from lens.pipeline import LensPipeline +from lens.text_encoder import LensGptOssEncoder +from lens.transformer import LensTransformer2DModel + +# Chat template constants, matching lens/pipeline.py +CHAT_SYSTEM = ( + "Describe the image by detailing the color, shape, size, texture, " + "quantity, text, spatial relationships of the objects and background." +) +CHAT_ASSISTANT_THINKING = "Need to generate one image according to the description." +PROMPT_TEMPLATE_CROP_START = 97 # tokens consumed by the chat template prefix +PROMPT_MAX_LENGTH = 512 # caption token budget (chat template tokens are added on top) + + +def make_lens_conversation(caption: str) -> list[dict]: + return [ + {"role": "system", "content": CHAT_SYSTEM, "thinking": None}, + {"role": "user", "content": caption, "thinking": None}, + {"role": "assistant", "thinking": CHAT_ASSISTANT_THINKING, "content": ""}, + ] + + +class LensModel(BaseModel): + # base model data + tokenizer: PreTrainedTokenizerFast | None + noise_scheduler: FlowMatchEulerDiscreteScheduler | None + text_encoder: LensGptOssEncoder | OnDemandModule | None + text_encoder_hidden_size: int | None # cached so encode_text() works after encoder is deleted + vae: AutoencoderKLFlux2 | None + transformer: LensTransformer2DModel | None + + # autocast context + text_encoder_autocast_context: torch.autocast | nullcontext + + transformer_offload_conductor: LayerOffloadConductor | None + + transformer_lora: LoRAModuleWrapper | None + lora_state_dict: dict | None + + def __init__( + self, + model_type: ModelType, + ): + super().__init__( + model_type=model_type, + ) + + self.tokenizer = None + self.noise_scheduler = None + self.text_encoder = None + self.text_encoder_hidden_size = None + self.vae = None + self.transformer = None + + self.text_encoder_autocast_context = nullcontext() + + self.transformer_offload_conductor = None + + self.transformer_lora = None + self.lora_state_dict = None + + def adapters(self) -> list[LoRAModuleWrapper]: + return [a for a in [ + self.transformer_lora, + ] if a is not None] + + def vae_to(self, device: torch.device): + self.vae.to(device=device) + + def text_encoder_to(self, device: torch.device): + # only moves a resident/materialized encoder. The on-demand proxy's to() is a no-op; + # its placement is owned by materialize_text_encoder / release_text_encoder. + if self.text_encoder is not None: + self.text_encoder.to(device=device) + + # Lens always loads on demand (see TrainConfig.text_encoder_on_demand): the GPT-OSS encoder is + # MXFP4-quantized and cannot be parked on the CPU temp device, so it is rebuilt straight onto the + # accelerator when needed and discarded afterwards. MXFP4 quantization happens inside + # from_pretrained (the loader lambda), so no quantize_layers call is needed here. + def materialize_text_encoder(self, device: torch.device): + if isinstance(self.text_encoder, OnDemandModule): + self.text_encoder.materialize() + self.text_encoder.inner.to(device) + else: + self.text_encoder_to(device) # resident: just move + + # load an on-demand encoder for saving. resident encoders already hold their weights. + def materialize_text_encoder_for_save(self): + if isinstance(self.text_encoder, OnDemandModule): + self.text_encoder.materialize() + + def release_text_encoder(self): + if isinstance(self.text_encoder, OnDemandModule): + self.text_encoder.discard() # free the weights + torch_gc() + else: + self.text_encoder_to(torch.device(self.train_config.temp_device)) # resident: park on temp + + def transformer_to(self, device: torch.device): + if self.transformer_offload_conductor is not None and \ + self.transformer_offload_conductor.layer_offload_activated(): + 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) + + # park resident components on the temp device and discard the on-demand text encoder + def release(self): + temp_device = torch.device(self.train_config.temp_device) + self.vae_to(temp_device) + self.release_text_encoder() + self.transformer_to(temp_device) + + def eval(self): + self.vae.eval() + if self.text_encoder is not None: + self.text_encoder.eval() + self.transformer.eval() + + # the real module behind the (possibly on-demand) text_encoder: the inner module when + # materialized, None when discarded, or the resident module. The pipeline and save path + # need the unwrapped module, not the proxy. + def _resolved_text_encoder(self) -> torch.nn.Module | None: + if isinstance(self.text_encoder, OnDemandModule): + return self.text_encoder.inner + return self.text_encoder + + def create_pipeline(self) -> DiffusionPipeline: + return LensPipeline( + scheduler=self.noise_scheduler, + vae=self.vae, + text_encoder=self._resolved_text_encoder(), + tokenizer=self.tokenizer, + transformer=self.transformer, + ) + + def encode_text( + self, + train_device: torch.device, + batch_size: int = 1, + rand: Random | None = None, + text: str | list[str] = None, + tokens: Tensor = None, + tokens_mask: Tensor = None, + text_encoder_dropout_probability: float | None = None, + text_encoder_output: list[Tensor] | None = None, + ) -> tuple[list[Tensor] | None, Tensor | None]: + + # the on-demand proxy is always truthy, so resolve the real (materialized) encoder here and + # gate the fresh-encode path on it being present rather than on `self.text_encoder`. + text_encoder = self._resolved_text_encoder() + + if tokens is None and text is not None: + if isinstance(text, str): + text = [text] + + rendered = [] + for prompt in text: + t = self.tokenizer.apply_chat_template( + make_lens_conversation(prompt), tokenize=False, add_generation_prompt=False + ) + rendered.append(t.split("<|return|>")[0]) + + tokenizer_output = self.tokenizer( + rendered, + max_length=PROMPT_MAX_LENGTH + PROMPT_TEMPLATE_CROP_START, + padding='max_length', + truncation=True, + return_tensors="pt", + add_special_tokens=True, + ) + tokens = tokenizer_output.input_ids.to(text_encoder.device) + tokens_mask = tokenizer_output.attention_mask.to(text_encoder.device) + + if text_encoder_output is None and text_encoder is not None: + with self.text_encoder_autocast_context: + # encode_layers() is used instead of the standard output_hidden_states=True API + # because transformers' @capture_outputs decorator applies tie_last_hidden_states, + # which silently replaces hidden_states[-1] with the norm-applied final output. + # GPT-OSS has 24 layers and selected_layer_index ends at layer 23 (the last), + # so output_hidden_states=True returns the normed version for that layer while + # encode_layers() correctly returns the pre-norm hidden state. + layer_outputs = text_encoder.encode_layers(tokens, tokens_mask) + if tokens.shape[1] > PROMPT_TEMPLATE_CROP_START: + text_encoder_output = [feat[:, PROMPT_TEMPLATE_CROP_START:, :].contiguous() for feat in layer_outputs] + tokens_mask = tokens_mask[:, PROMPT_TEMPLATE_CROP_START:] + else: + #TODO can this ever happen? max_length=PROMPT_MAX_LENGTH+PROMPT_TEMPLATE_CROP_START, + #so the tokenizer output should always be longer than PROMPT_TEMPLATE_CROP_START. + #upstream pipeline has the same guard, presumably as a safety net. + zero_shape = (tokens.shape[0], 0, layer_outputs[0].shape[-1]) + text_encoder_output = [layer_outputs[0].new_zeros(zero_shape) for _ in layer_outputs] + tokens_mask = torch.zeros((tokens.shape[0], 0), dtype=torch.bool, device=tokens.device) + + elif text_encoder_output is not None: + # Cached: EncodeLensText concatenates layers along dim=-1; split back into per-layer list. + hidden_dim = self.text_encoder_hidden_size + text_encoder_output = list(text_encoder_output.split(hidden_dim, dim=-1)) + + if text_encoder_dropout_probability is not None and text_encoder_dropout_probability > 0.0: + raise NotImplementedError # https://github.com/Nerogar/OneTrainer/issues/957 + + #prune tokens that are masked in all batch samples: + seq_lengths = tokens_mask.sum(dim=1) + max_seq_length = int(seq_lengths.max().item()) + + #pad to 16 because attention processors and/or torch.compile can have issues with uneven sequence lengths, but only pad if an attention mask has to be used anyway: + if max_seq_length % 16 > 0 and (seq_lengths != max_seq_length).any(): + max_seq_length += (16 - max_seq_length % 16) + + text_encoder_output = [feat[:, :max_seq_length, :] for feat in text_encoder_output] + tokens_mask = tokens_mask[:, :max_seq_length].bool() + + return text_encoder_output, tokens_mask + + def calculate_timestep_shift(self, latent_height: int, latent_width: int) -> float: + base_seq_len = self.noise_scheduler.config.base_image_seq_len + max_seq_len = self.noise_scheduler.config.max_image_seq_len + base_shift = self.noise_scheduler.config.base_shift + max_shift = self.noise_scheduler.config.max_shift + patch_size = 2 + + image_seq_len = (latent_width // patch_size) * (latent_height // patch_size) + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return math.exp(mu) + + #packing and unpacking on patchified latents + @staticmethod + def pack_latents(latents) -> Tensor: + batch_size, num_channels, height, width = latents.shape + return latents.reshape(batch_size, num_channels, height * width).permute(0, 2, 1) + + @staticmethod + def unpack_latents(latents, height: int, width: int) -> Tensor: + batch_size, seq_len, num_channels = latents.shape + return latents.reshape(batch_size, height, width, num_channels).permute(0, 3, 1, 2) + + @staticmethod + def patchify_latents(latents: torch.Tensor) -> torch.Tensor: + batch_size, num_channels_latents, height, width = latents.shape + latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 1, 3, 5, 2, 4) + latents = latents.reshape(batch_size, num_channels_latents * 4, height // 2, width // 2) + return latents + + @staticmethod + def unpatchify_latents(latents: torch.Tensor) -> torch.Tensor: + batch_size, num_channels_latents, height, width = latents.shape + latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), 2, 2, height, width) + latents = latents.permute(0, 1, 4, 2, 5, 3) + latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), height * 2, width * 2) + return latents + + #scaling on patchified latents + def scale_latents(self, latents: Tensor) -> Tensor: + latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) + latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to( + latents.device, latents.dtype + ) + return (latents - latents_bn_mean) / latents_bn_std + + def unscale_latents(self, latents: Tensor) -> Tensor: + latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) + latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to( + latents.device, latents.dtype + ) + return latents * latents_bn_std + latents_bn_mean diff --git a/modules/model/PixArtAlphaModel.py b/modules/model/PixArtAlphaModel.py index 466cc61f9..92f2c798b 100644 --- a/modules/model/PixArtAlphaModel.py +++ b/modules/model/PixArtAlphaModel.py @@ -114,8 +114,7 @@ def vae_to(self, device: torch.device): self.vae.to(device=device) def text_encoder_to(self, device: torch.device): - 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) @@ -124,8 +123,7 @@ 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) @@ -133,10 +131,10 @@ def transformer_to(self, device: torch.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() diff --git a/modules/model/QwenModel.py b/modules/model/QwenModel.py index afa6c24fe..361d7ee2f 100644 --- a/modules/model/QwenModel.py +++ b/modules/model/QwenModel.py @@ -81,8 +81,7 @@ def vae_to(self, device: torch.device): def text_encoder_to(self, device: torch.device): #TODO share more code between models 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) @@ -91,8 +90,7 @@ def text_encoder_to(self, device: torch.device): #TODO share more code between m 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) @@ -100,10 +98,10 @@ def transformer_to(self, device: torch.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() diff --git a/modules/model/SanaModel.py b/modules/model/SanaModel.py index 9e8008219..8e96c82a5 100644 --- a/modules/model/SanaModel.py +++ b/modules/model/SanaModel.py @@ -116,8 +116,7 @@ def vae_to(self, device: torch.device): self.vae.to(device=device) def text_encoder_to(self, device: torch.device): - 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) @@ -126,8 +125,7 @@ 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) @@ -135,10 +133,10 @@ def transformer_to(self, device: torch.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() diff --git a/modules/model/StableDiffusion3Model.py b/modules/model/StableDiffusion3Model.py index 8f6cf5818..b3710a619 100644 --- a/modules/model/StableDiffusion3Model.py +++ b/modules/model/StableDiffusion3Model.py @@ -180,8 +180,7 @@ def text_encoder_2_to(self, device: torch.device): def text_encoder_3_to(self, device: torch.device): if self.text_encoder_3 is not None: - if self.text_encoder_3_offload_conductor is not None and \ - self.text_encoder_3_offload_conductor.layer_offload_activated(): + if self.text_encoder_3_offload_conductor is not None: self.text_encoder_3_offload_conductor.to(device) else: self.text_encoder_3.to(device=device) @@ -190,8 +189,7 @@ def text_encoder_3_to(self, device: torch.device): self.text_encoder_3_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) @@ -199,10 +197,10 @@ def transformer_to(self, device: torch.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() diff --git a/modules/model/StableDiffusionModel.py b/modules/model/StableDiffusionModel.py index 4b1b11a2c..6b741764a 100644 --- a/modules/model/StableDiffusionModel.py +++ b/modules/model/StableDiffusionModel.py @@ -125,11 +125,11 @@ def unet_to(self, device: torch.device): if self.unet_lora is not None: self.unet_lora.to(device) - def to(self, device: torch.device): - self.vae_to(device) - self.depth_estimator_to(device) - self.text_encoder_to(device) - self.unet_to(device) + def release(self): + self.vae_to(self.train_config.temp_device) + self.depth_estimator_to(self.train_config.temp_device) + self.text_encoder_to(self.train_config.temp_device) + self.unet_to(self.train_config.temp_device) def eval(self): self.vae.eval() diff --git a/modules/model/StableDiffusionXLModel.py b/modules/model/StableDiffusionXLModel.py index 43edead8e..b5fa046aa 100644 --- a/modules/model/StableDiffusionXLModel.py +++ b/modules/model/StableDiffusionXLModel.py @@ -159,10 +159,10 @@ def unet_to(self, device: torch.device): if self.unet_lora is not None: self.unet_lora.to(device) - def to(self, device: torch.device): - self.vae_to(device) - self.text_encoder_to(device) - self.unet_to(device) + def release(self): + self.vae_to(self.train_config.temp_device) + self.text_encoder_to(self.train_config.temp_device) + self.unet_to(self.train_config.temp_device) def eval(self): self.vae.eval() diff --git a/modules/model/WuerstchenModel.py b/modules/model/WuerstchenModel.py index 7b50d81ed..c12491245 100644 --- a/modules/model/WuerstchenModel.py +++ b/modules/model/WuerstchenModel.py @@ -163,14 +163,14 @@ def prior_prior_to(self, device: torch.device): if self.prior_prior_lora is not None: self.prior_prior_lora.to(device) - def to(self, device: torch.device): + def release(self): if self.model_type.is_wuerstchen_v2(): - self.decoder_text_encoder_to(device) - self.decoder_decoder_to(device) - self.decoder_vqgan_to(device) - self.effnet_encoder_to(device) - self.prior_text_encoder_to(device) - self.prior_prior_to(device) + self.decoder_text_encoder_to(self.train_config.temp_device) + self.decoder_decoder_to(self.train_config.temp_device) + self.decoder_vqgan_to(self.train_config.temp_device) + self.effnet_encoder_to(self.train_config.temp_device) + self.prior_text_encoder_to(self.train_config.temp_device) + self.prior_prior_to(self.train_config.temp_device) def eval(self): if self.model_type.is_wuerstchen_v2(): diff --git a/modules/model/ZImageModel.py b/modules/model/ZImageModel.py index 7fd9e52cb..33e09ffd7 100644 --- a/modules/model/ZImageModel.py +++ b/modules/model/ZImageModel.py @@ -83,15 +83,13 @@ def vae_to(self, device: torch.device): def text_encoder_to(self, device: torch.device): #TODO share more code between models 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) @@ -99,10 +97,10 @@ def transformer_to(self, device: torch.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() diff --git a/modules/modelLoader/Flux2ModelLoader.py b/modules/modelLoader/Flux2ModelLoader.py index 55bbeebd4..c50f60714 100644 --- a/modules/modelLoader/Flux2ModelLoader.py +++ b/modules/modelLoader/Flux2ModelLoader.py @@ -198,8 +198,7 @@ def __init__(self): super().__init__() def _get_convert_key_sets(self, model: BaseModel) -> list[LoraConversionKeySet] | None: - return None #TODO - #return convert_flux_lora_key_sets() + return None def load( self, diff --git a/modules/modelLoader/LensModelLoader.py b/modules/modelLoader/LensModelLoader.py new file mode 100644 index 000000000..d1dba3561 --- /dev/null +++ b/modules/modelLoader/LensModelLoader.py @@ -0,0 +1,212 @@ +import os +import traceback + +from modules.model.BaseModel import BaseModel +from modules.model.LensModel import LensModel +from modules.modelLoader.GenericFineTuneModelLoader import make_fine_tune_model_loader +from modules.modelLoader.GenericLoRAModelLoader import make_lora_model_loader +from modules.modelLoader.mixin.HFModelLoaderMixin import HFModelLoaderMixin +from modules.modelLoader.mixin.LoRALoaderMixin import LoRALoaderMixin +from modules.util.config.TrainConfig import QuantizationConfig +from modules.util.convert.lora.convert_lora_util import LoraConversionKeySet +from modules.util.enum.ModelType import ModelType +from modules.util.ModelNames import ModelNames +from modules.util.ModelWeightDtypes import ModelWeightDtypes +from modules.util.OnDemandModule import OnDemandModule + +import torch + +from diffusers import ( + AutoencoderKLFlux2, + FlowMatchEulerDiscreteScheduler, + GGUFQuantizationConfig, +) +from transformers import PreTrainedTokenizerFast + +from lens.text_encoder import LensGptOssEncoder +from lens.transformer import LensTransformer2DModel + + +class LensModelLoader( + HFModelLoaderMixin, +): + def __init__(self): + super().__init__() + + def __load_internal( + self, + model: LensModel, + model_type: ModelType, + weight_dtypes: ModelWeightDtypes, + base_model_name: str, + transformer_model_name: str, + vae_model_name: str, + quantization: QuantizationConfig, + text_encoder_on_demand: bool, + ): + 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, quantization, text_encoder_on_demand, + ) + else: + raise Exception("not an internal model") + + def __load_diffusers( + self, + model: LensModel, + model_type: ModelType, + weight_dtypes: ModelWeightDtypes, + base_model_name: str, + transformer_model_name: str, + vae_model_name: str, + quantization: QuantizationConfig, + text_encoder_on_demand: bool, + ): + if transformer_model_name: + transformer = LensTransformer2DModel.from_single_file( + transformer_model_name, + config=base_model_name, + subfolder="transformer", + #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.is_gguf() else None, + ) + transformer = self._convert_diffusers_sub_module_to_dtype( + transformer, weight_dtypes.transformer, weight_dtypes.train_dtype, quantization, + ) + else: + transformer = self._load_diffusers_sub_module( + LensTransformer2DModel, + weight_dtypes.transformer, + weight_dtypes.train_dtype, + base_model_name, + "transformer", + quantization, + ) + + #TODO verify whether the TokenizersBackend warning actually appears for Lens; if so, uncomment the log suppression below (see ErnieModelLoader for the pattern) + #tokenization_logger = logging.getLogger("transformers.tokenization_utils_base") + #prev_level = tokenization_logger.level + #tokenization_logger.setLevel(logging.ERROR) + tokenizer = PreTrainedTokenizerFast.from_pretrained( + base_model_name, + subfolder="tokenizer", + ) + #tokenization_logger.setLevel(prev_level) + + selected_layer_index = transformer.config.selected_layer_index + + def load_text_encoder(): + text_encoder = LensGptOssEncoder.from_pretrained( + base_model_name, + subfolder="text_encoder", + ) + # set_selected_layers must be called before encode_layers(); the upstream does this in + # LensPipeline.__init__ — we do it here since OneTrainer loads components separately. + text_encoder.set_selected_layers(selected_layer_index) + return text_encoder + + # Lens always loads on demand: its MXFP4 encoder cannot be parked on the CPU temp device, so it + # is built straight onto the accelerator when needed and discarded afterwards (see + # TrainConfig.text_encoder_on_demand / LensModel.materialize_text_encoder). + text_encoder = OnDemandModule(load_text_encoder) if text_encoder_on_demand else load_text_encoder() + + noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + base_model_name, + subfolder="scheduler", + ) + + if vae_model_name: + vae = self._load_diffusers_sub_module( + AutoencoderKLFlux2, + weight_dtypes.vae, + weight_dtypes.train_dtype, + vae_model_name, + ) + else: + vae = self._load_diffusers_sub_module( + AutoencoderKLFlux2, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + "vae", + ) + + model.model_type = model_type + model.tokenizer = tokenizer + model.noise_scheduler = noise_scheduler + model.text_encoder = text_encoder + # read hidden_size from the config only (no weights), so an on-demand encoder stays unmaterialized + model.text_encoder_hidden_size = LensGptOssEncoder.config_class.from_pretrained( + base_model_name, subfolder="text_encoder", + ).hidden_size + model.vae = vae + model.transformer = transformer + + def load( + self, + model: LensModel, + model_type: ModelType, + model_names: ModelNames, + weight_dtypes: ModelWeightDtypes, + quantization: QuantizationConfig, + ): + stacktraces = [] + + try: + self.__load_internal( + model, model_type, weight_dtypes, model_names.base_model, model_names.transformer_model, model_names.vae_model, quantization, model_names.text_encoder_on_demand, + ) + 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, quantization, model_names.text_encoder_on_demand, + ) + return + except Exception: + stacktraces.append(traceback.format_exc()) + + for stacktrace in stacktraces: + print(stacktrace) + raise Exception("could not load model: " + model_names.base_model) + + + +class LensLoRALoader( + LoRALoaderMixin +): + def __init__(self): + super().__init__() + + def _get_convert_key_sets(self, model: BaseModel) -> list[LoraConversionKeySet] | None: + return None #TODO + + def load( + self, + model: LensModel, + model_names: ModelNames, + ): + return self._load(model, model_names) + + +LensLoRAModelLoader = make_lora_model_loader( + model_spec_map={ + ModelType.LENS: "resources/sd_model_spec/lens-lora.json", + }, + model_class=LensModel, + model_loader_class=LensModelLoader, + lora_loader_class=LensLoRALoader, + embedding_loader_class=None, +) + +LensFineTuneModelLoader = make_fine_tune_model_loader( + model_spec_map={ + ModelType.LENS: "resources/sd_model_spec/lens.json", + }, + model_class=LensModel, + model_loader_class=LensModelLoader, + embedding_loader_class=None, +) diff --git a/modules/modelSampler/LensSampler.py b/modules/modelSampler/LensSampler.py new file mode 100644 index 000000000..2d4ccb96b --- /dev/null +++ b/modules/modelSampler/LensSampler.py @@ -0,0 +1,188 @@ +import copy +import inspect +from collections.abc import Callable + +from modules.model.LensModel import LensModel +from modules.modelSampler.BaseModelSampler import BaseModelSampler, ModelSamplerOutput +from modules.util import factory +from modules.util.config.SampleConfig import SampleConfig +from modules.util.enum.AudioFormat import AudioFormat +from modules.util.enum.FileType import FileType +from modules.util.enum.ImageFormat import ImageFormat +from modules.util.enum.ModelType import ModelType +from modules.util.enum.NoiseScheduler import NoiseScheduler +from modules.util.enum.VideoFormat import VideoFormat +from modules.util.torch_util import torch_gc + +import torch + +import numpy as np +from lens.pipeline import LensPipeline, compute_empirical_mu +from tqdm import tqdm + + +@factory.register(BaseModelSampler, ModelType.LENS) +class LensSampler(BaseModelSampler): + def __init__( + self, + train_device: torch.device, + temp_device: torch.device, + model: LensModel, + model_type: ModelType, + ): + super().__init__(train_device, temp_device) + + self.model = model + self.model_type = model_type + + @torch.no_grad() + def __sample_base( + self, + prompt: str, + negative_prompt: str, + height: int, + width: int, + seed: int, + random_seed: bool, + diffusion_steps: int, + cfg_scale: float, + noise_scheduler: NoiseScheduler, + on_update_progress: Callable[[int, int], None] = lambda _, __: None, + ) -> ModelSamplerOutput: + with self.model.autocast_context: + generator = torch.Generator(device=self.train_device) + if random_seed: + generator.seed() + else: + generator.manual_seed(seed) + + noise_scheduler = copy.deepcopy(self.model.noise_scheduler) + transformer = self.model.transformer + vae = self.model.vae + + vae_scale_factor = 8 + num_latent_channels = 32 + patch_size = 2 + + # prepare prompt + self.model.materialize_text_encoder(self.train_device) + + batch_size = 2 if cfg_scale > 1.0 else 1 + (prompt_features, prompt_mask) = self.model.encode_text( + text=[prompt, negative_prompt] if batch_size == 2 else [prompt], + train_device=self.train_device, + ) + + self.model.release_text_encoder() + torch_gc() + + # prepare latent image + latent_image = torch.randn( + size=(1, num_latent_channels, height // vae_scale_factor, width // vae_scale_factor), + generator=generator, + device=self.train_device, + dtype=torch.float32, + ) + + latent_image = self.model.patchify_latents(latent_image) + + latent_image = self.model.pack_latents(latent_image) + image_seq_len = latent_image.shape[1] + mu = compute_empirical_mu(image_seq_len, diffusion_steps) + + # prepare timesteps + #TODO for other models, too? This is different than with sigmas=None + sigmas = np.linspace(1.0, 1 / diffusion_steps, diffusion_steps) + noise_scheduler.set_timesteps(diffusion_steps, device=self.train_device, mu=mu, sigmas=sigmas) + timesteps = noise_scheduler.timesteps + + # denoising loop + extra_step_kwargs = {} #TODO remove + if "generator" in set(inspect.signature(noise_scheduler.step).parameters.keys()): + extra_step_kwargs["generator"] = generator + + # img_shapes: (frame, h_lat, w_lat) for RoPE + img_shapes = [(1, height // vae_scale_factor // patch_size, width // vae_scale_factor // patch_size)] + + self.model.transformer_to(self.train_device) + for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): + latent_model_input = torch.cat([latent_image] * batch_size) + expanded_timestep = timestep.expand(latent_model_input.shape[0]) + + + noise_pred = transformer( + hidden_states=latent_model_input.to(dtype=self.model.train_dtype.torch_dtype()), + encoder_hidden_states=[f.to(dtype=self.model.train_dtype.torch_dtype()) for f in prompt_features], # list of 4 per-GPT-OSS-layer tensors, each [B, S_txt, dim] + encoder_hidden_states_mask=prompt_mask, + timestep=expanded_timestep / 1000, + img_shapes=img_shapes, + ) + + if batch_size == 2: + cond, uncond = noise_pred.chunk(2) + # Norm-rescaled CFG: rescale the combined prediction to match ||cond||, + # preventing magnitude blowup at high guidance scales. Used by the upstream + # pipeline even though not documented in the Lens paper (arxiv:2605.21573). + comb = uncond + cfg_scale * (cond - uncond) + cond_norm = torch.norm(cond, dim=-1, keepdim=True) + comb_norm = torch.norm(comb, dim=-1, keepdim=True) + scale = torch.where(comb_norm > 0, cond_norm / comb_norm.clamp_min(1e-12), torch.ones_like(comb_norm)) + noise_pred = comb * scale + + latent_image = noise_scheduler.step(noise_pred, timestep, latent_image, return_dict=False, **extra_step_kwargs)[0] + + on_update_progress(i + 1, len(timesteps)) + + self.model.transformer_to(self.temp_device) + torch_gc() + self.model.vae_to(self.train_device) + + latent_image = self.model.unpack_latents( + latent_image, + height // vae_scale_factor // patch_size, + width // vae_scale_factor // patch_size, + ) + latents = self.model.unscale_latents(latent_image) + latents = self.model.unpatchify_latents(latents) + + decoded = vae.decode(latents, return_dict=False)[0] + + image = LensPipeline._to_pil(decoded) + + self.model.vae_to(self.temp_device) + torch_gc() + + return ModelSamplerOutput( + file_type=FileType.IMAGE, + data=image[0], + ) + + def sample( + self, + sample_config: SampleConfig, + destination: str, + image_format: ImageFormat | None = None, + video_format: VideoFormat | None = None, + audio_format: AudioFormat | None = None, + on_sample: Callable[[ModelSamplerOutput], None] = lambda _: None, + on_update_progress: Callable[[int, int], None] = lambda _, __: None, + ): + sampler_output = self.__sample_base( + prompt=sample_config.prompt, + negative_prompt=sample_config.negative_prompt, + height=self.quantize_resolution(sample_config.height, 64), + width=self.quantize_resolution(sample_config.width, 64), + seed=sample_config.seed, + random_seed=sample_config.random_seed, + diffusion_steps=sample_config.diffusion_steps, + cfg_scale=sample_config.cfg_scale, + noise_scheduler=sample_config.noise_scheduler, + on_update_progress=on_update_progress, + ) + + self.save_sampler_output( + sampler_output, destination, + image_format, video_format, audio_format, + ) + + on_sample(sampler_output) diff --git a/modules/modelSaver/LensFineTuneModelSaver.py b/modules/modelSaver/LensFineTuneModelSaver.py new file mode 100644 index 000000000..300c18f3a --- /dev/null +++ b/modules/modelSaver/LensFineTuneModelSaver.py @@ -0,0 +1,11 @@ +from modules.model.LensModel import LensModel +from modules.modelSaver.GenericFineTuneModelSaver import make_fine_tune_model_saver +from modules.modelSaver.lens.LensModelSaver import LensModelSaver +from modules.util.enum.ModelType import ModelType + +LensFineTuneModelSaver = make_fine_tune_model_saver( + ModelType.LENS, + model_class=LensModel, + model_saver_class=LensModelSaver, + embedding_saver_class=None, +) diff --git a/modules/modelSaver/LensLoRAModelSaver.py b/modules/modelSaver/LensLoRAModelSaver.py new file mode 100644 index 000000000..54656d48c --- /dev/null +++ b/modules/modelSaver/LensLoRAModelSaver.py @@ -0,0 +1,11 @@ +from modules.model.LensModel import LensModel +from modules.modelSaver.GenericLoRAModelSaver import make_lora_model_saver +from modules.modelSaver.lens.LensLoRASaver import LensLoRASaver +from modules.util.enum.ModelType import ModelType + +LensLoRAModelSaver = make_lora_model_saver( + ModelType.LENS, + model_class=LensModel, + lora_saver_class=LensLoRASaver, + embedding_saver_class=None, +) diff --git a/modules/modelSaver/lens/LensLoRASaver.py b/modules/modelSaver/lens/LensLoRASaver.py new file mode 100644 index 000000000..97846bed3 --- /dev/null +++ b/modules/modelSaver/lens/LensLoRASaver.py @@ -0,0 +1,38 @@ +from modules.model.LensModel import LensModel +from modules.modelSaver.mixin.LoRASaverMixin import LoRASaverMixin +from modules.util.convert.lora.convert_lora_util import LoraConversionKeySet +from modules.util.enum.ModelFormat import ModelFormat + +import torch +from torch import Tensor + + +class LensLoRASaver( + LoRASaverMixin, +): + def __init__(self): + super().__init__() + + def _get_convert_key_sets(self, model: LensModel) -> list[LoraConversionKeySet] | None: + return None + + def _get_state_dict( + self, + model: LensModel, + ) -> dict[str, Tensor]: + state_dict = {} + if model.transformer_lora is not None: + state_dict |= model.transformer_lora.state_dict() + if model.lora_state_dict is not None: + state_dict |= model.lora_state_dict + + return state_dict + + def save( + self, + model: LensModel, + output_model_format: ModelFormat, + output_model_destination: str, + dtype: torch.dtype | None, + ): + self._save(model, output_model_format, output_model_destination, dtype) diff --git a/modules/modelSaver/lens/LensModelSaver.py b/modules/modelSaver/lens/LensModelSaver.py new file mode 100644 index 000000000..26bfdabd7 --- /dev/null +++ b/modules/modelSaver/lens/LensModelSaver.py @@ -0,0 +1,87 @@ +import copy +import os.path +from pathlib import Path + +from modules.model.LensModel import LensModel +from modules.modelSaver.mixin.DtypeModelSaverMixin import DtypeModelSaverMixin +from modules.util.enum.ModelFormat import ModelFormat + +import torch + +from safetensors.torch import save_file + + +class LensModelSaver( + DtypeModelSaverMixin, +): + def __init__(self): + super().__init__() + + def __save_diffusers( + self, + model: LensModel, + destination: str, + dtype: torch.dtype | None, + ): + # this is the only save path that embeds the base text encoder, so the on-demand encoder is + # loaded here (and freed again afterwards). resident encoders are untouched. + model.materialize_text_encoder_for_save() + try: + # Copy the model to cpu by first moving the original model to cpu. This preserves some VRAM. + pipeline = model.create_pipeline() + pipeline.to("cpu") + if dtype is not None: #TODO necessary? + # replace the tokenizers __deepcopy__ before calling deepcopy, to prevent a copy being made. + # the tokenizer tries to reload from the file system otherwise + tokenizer = pipeline.tokenizer + tokenizer.__deepcopy__ = lambda memo: tokenizer + + save_pipeline = copy.deepcopy(pipeline) + save_pipeline.to(device="cpu", dtype=dtype, silence_dtype_warnings=True) + + delattr(tokenizer, '__deepcopy__') + else: + save_pipeline = pipeline + + os.makedirs(Path(destination).absolute(), exist_ok=True) + save_pipeline.save_pretrained(destination) + + if dtype is not None: + del save_pipeline + finally: + model.release_text_encoder() + + def __save_safetensors( + self, + model: LensModel, + destination: str, + dtype: torch.dtype | None, + ): + # Lens transformer uses diffusers-format keys; no key conversion needed. + state_dict = model.transformer.state_dict() + save_state_dict = self._convert_state_dict_dtype(state_dict, dtype) + self._convert_state_dict_to_contiguous(save_state_dict) + os.makedirs(Path(destination).parent.absolute(), exist_ok=True) + save_file(save_state_dict, destination, self._create_safetensors_header(model, save_state_dict)) + + def __save_internal( + self, + model: LensModel, + destination: str, + ): + self.__save_diffusers(model, destination, None) + + def save( + self, + model: LensModel, + output_model_format: ModelFormat, + output_model_destination: str, + dtype: torch.dtype | None, + ): + match output_model_format: + case ModelFormat.DIFFUSERS: + self.__save_diffusers(model, output_model_destination, dtype) + case ModelFormat.SAFETENSORS: + self.__save_safetensors(model, output_model_destination, dtype) + case ModelFormat.INTERNAL: + self.__save_internal(model, output_model_destination) diff --git a/modules/modelSetup/BaseChromaSetup.py b/modules/modelSetup/BaseChromaSetup.py index 0eb623399..c6fbb0f14 100644 --- a/modules/modelSetup/BaseChromaSetup.py +++ b/modules/modelSetup/BaseChromaSetup.py @@ -49,12 +49,9 @@ def setup_optimizations( model: ChromaModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_chroma_transformer(model.transformer, config) - if model.text_encoder is not None: - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config) + model.transformer_offload_conductor = enable_checkpointing_for_chroma_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None: + model.text_encoder_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, @@ -279,7 +276,7 @@ def calculate_loss( def prepare_text_caching(self, model: ChromaModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseErnieSetup.py b/modules/modelSetup/BaseErnieSetup.py index 3d630b8cf..469604318 100644 --- a/modules/modelSetup/BaseErnieSetup.py +++ b/modules/modelSetup/BaseErnieSetup.py @@ -45,11 +45,10 @@ def setup_optimizations( model: ErnieModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_ernie_transformer(model.transformer, config) - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_mistral_encoder_layers(model.text_encoder, config) + if config.transformer.checkpointing_or_offloading_enabled(): + model.transformer_offload_conductor = enable_checkpointing_for_ernie_transformer(model.transformer, config, config.transformer) + if config.text_encoder.checkpointing_or_offloading_enabled(): + model.text_encoder_offload_conductor = enable_checkpointing_for_mistral_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, @@ -170,7 +169,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: ErnieModel, config: TrainConfig): - model.to(self.temp_device) + model.release() model.text_encoder_to(self.train_device) model.eval() torch_gc() diff --git a/modules/modelSetup/BaseFlux2Setup.py b/modules/modelSetup/BaseFlux2Setup.py index b91cd8af8..1d98f4aef 100644 --- a/modules/modelSetup/BaseFlux2Setup.py +++ b/modules/modelSetup/BaseFlux2Setup.py @@ -45,16 +45,12 @@ def setup_optimizations( model: Flux2Model, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_flux2_transformer(model.transformer, config) - if model.text_encoder is not None: - if model.is_dev(): - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_mistral_encoder_layers(model.text_encoder, config) - else: - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config) + model.transformer_offload_conductor = enable_checkpointing_for_flux2_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None: + if model.is_dev(): + model.text_encoder_offload_conductor = enable_checkpointing_for_mistral_encoder_layers(model.text_encoder, config, config.text_encoder) + else: + model.text_encoder_offload_conductor = enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, @@ -194,7 +190,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: FluxModel, config: TrainConfig): - model.to(self.temp_device) + model.release() model.text_encoder_to(self.train_device) model.eval() torch_gc() diff --git a/modules/modelSetup/BaseFluxSetup.py b/modules/modelSetup/BaseFluxSetup.py index 9bae83cde..8f095018a 100644 --- a/modules/modelSetup/BaseFluxSetup.py +++ b/modules/modelSetup/BaseFluxSetup.py @@ -49,14 +49,11 @@ def setup_optimizations( model: FluxModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_flux_transformer(model.transformer, config) - if model.text_encoder_1 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config) - if model.text_encoder_2 is not None: - model.text_encoder_2_offload_conductor = \ - enable_checkpointing_for_t5_encoder_layers(model.text_encoder_2, config) + model.transformer_offload_conductor = enable_checkpointing_for_flux_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None: + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) + if model.text_encoder_2 is not None: + model.text_encoder_2_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, @@ -339,7 +336,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: FluxModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseHiDreamSetup.py b/modules/modelSetup/BaseHiDreamSetup.py index 17fbcc0d6..e6bce8d14 100644 --- a/modules/modelSetup/BaseHiDreamSetup.py +++ b/modules/modelSetup/BaseHiDreamSetup.py @@ -49,19 +49,15 @@ def setup_optimizations( model: HiDreamModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_hi_dream_transformer(model.transformer, config) - if model.text_encoder_1 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config) - if model.text_encoder_2 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config) - if model.text_encoder_3 is not None: - model.text_encoder_3_offload_conductor = \ - enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config) - if model.text_encoder_4 is not None: - model.text_encoder_4_offload_conductor = \ - enable_checkpointing_for_llama_encoder_layers(model.text_encoder_4, config) + model.transformer_offload_conductor = enable_checkpointing_for_hi_dream_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None: + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) + if model.text_encoder_2 is not None: + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) + if model.text_encoder_3 is not None: + model.text_encoder_3_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config, config.text_encoder_3) + if model.text_encoder_4 is not None: + model.text_encoder_4_offload_conductor = enable_checkpointing_for_llama_encoder_layers(model.text_encoder_4, config, config.text_encoder_4) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, @@ -432,7 +428,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: HiDreamModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseHunyuanVideoSetup.py b/modules/modelSetup/BaseHunyuanVideoSetup.py index b072bf4ba..5301e4e38 100644 --- a/modules/modelSetup/BaseHunyuanVideoSetup.py +++ b/modules/modelSetup/BaseHunyuanVideoSetup.py @@ -49,14 +49,11 @@ def setup_optimizations( model: HunyuanVideoModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_hunyuan_video_transformer(model.transformer, config) - if model.text_encoder_1 is not None: - model.text_encoder_1_offload_conductor = \ - enable_checkpointing_for_llama_encoder_layers(model.text_encoder_1, config) - if model.text_encoder_2 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config) + model.transformer_offload_conductor = enable_checkpointing_for_hunyuan_video_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None: + model.text_encoder_1_offload_conductor = enable_checkpointing_for_llama_encoder_layers(model.text_encoder_1, config, config.text_encoder) + if model.text_encoder_2 is not None: + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, @@ -308,7 +305,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: HunyuanVideoModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseLensSetup.py b/modules/modelSetup/BaseLensSetup.py new file mode 100644 index 000000000..1f884e3f0 --- /dev/null +++ b/modules/modelSetup/BaseLensSetup.py @@ -0,0 +1,179 @@ +from abc import ABCMeta +from random import Random + +import modules.util.multi_gpu_util as multi +from modules.model.LensModel import LensModel +from modules.modelSetup.BaseModelSetup import BaseModelSetup +from modules.modelSetup.mixin.ModelSetupDebugMixin import ModelSetupDebugMixin +from modules.modelSetup.mixin.ModelSetupDiffusionLossMixin import ModelSetupDiffusionLossMixin +from modules.modelSetup.mixin.ModelSetupEmbeddingMixin import ModelSetupEmbeddingMixin +from modules.modelSetup.mixin.ModelSetupFlowMatchingMixin import ModelSetupFlowMatchingMixin +from modules.modelSetup.mixin.ModelSetupNoiseMixin import ModelSetupNoiseMixin +from modules.util.checkpointing_util import ( + enable_checkpointing_for_lens_transformer, +) +from modules.util.config.TrainConfig import TrainConfig +from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.quantization_util import quantize_layers +from modules.util.torch_util import torch_gc +from modules.util.TrainProgress import TrainProgress + +import torch +from torch import Tensor + + +class BaseLensSetup( + BaseModelSetup, + ModelSetupDiffusionLossMixin, + ModelSetupDebugMixin, + ModelSetupNoiseMixin, + ModelSetupFlowMatchingMixin, + ModelSetupEmbeddingMixin, + metaclass=ABCMeta +): + LAYER_PRESETS = { + "attn-mlp": ["attn", "mlp"], + "attn-only": ["attn"], + "blocks": ["transformer_block"], + "full": [], + } + + def setup_optimizations( + self, + model: LensModel, + config: TrainConfig, + ): + model.transformer_offload_conductor = \ + enable_checkpointing_for_lens_transformer(model.transformer, config, config.transformer) + + model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ + config.weight_dtypes().transformer, + config.weight_dtypes().text_encoder, + config.weight_dtypes().vae, + config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, + ], config.enable_autocast_cache) + + model.text_encoder_autocast_context, model.text_encoder_train_dtype = \ + disable_fp16_autocast_context( + self.train_device, + config.train_dtype, + config.fallback_train_dtype, + [ + config.weight_dtypes().text_encoder, + config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, + ], + config.enable_autocast_cache, + ) + + quantize_layers(model.vae, self.train_device, model.train_dtype, config) + quantize_layers(model.transformer, self.train_device, model.train_dtype, config) + #model.text_encoder is not quantized - GPT OSS is loaded in MXFP4 + + def predict( + self, + model: LensModel, + batch: dict, + config: TrainConfig, + train_progress: TrainProgress, + *, + deterministic: bool = False, + ) -> dict: + with model.autocast_context: + batch_seed = 0 if deterministic else train_progress.global_step * multi.world_size() + multi.rank() + generator = torch.Generator(device=config.train_device) + generator.manual_seed(batch_seed) + rand = Random(batch_seed) + + text_encoder_output, tokens_mask = model.encode_text( + train_device=self.train_device, + batch_size=batch['latent_image'].shape[0], + rand=rand, + tokens=batch.get("tokens"), + tokens_mask=batch.get("tokens_mask"), + text_encoder_output=batch.get('text_encoder_hidden_state'), + text_encoder_dropout_probability=config.text_encoder.dropout_probability if not deterministic else None, + ) + latent_image = model.patchify_latents(batch['latent_image'].float()) + latent_height = latent_image.shape[-2] + latent_width = latent_image.shape[-1] + scaled_latent_image = model.scale_latents(latent_image) + + latent_noise = self._create_noise(scaled_latent_image, config, generator) + + shift = model.calculate_timestep_shift(latent_height, latent_width) + timestep = self._get_timestep_discrete( + model.noise_scheduler.config['num_train_timesteps'], + deterministic, + generator, + scaled_latent_image.shape[0], + config, + shift = shift if config.dynamic_timestep_shifting else config.timestep_shift, + ) + + scaled_noisy_latent_image, sigma = self._add_noise_discrete( + scaled_latent_image, + latent_noise, + timestep, + model.noise_scheduler.timesteps, + ) + + img_shapes = [(1, latent_height, latent_width)] + packed_latent_input = model.pack_latents(scaled_noisy_latent_image) + + packed_predicted_flow = model.transformer( + hidden_states=packed_latent_input.to(dtype=model.train_dtype.torch_dtype()), + encoder_hidden_states=[f.to(dtype=model.train_dtype.torch_dtype()) for f in text_encoder_output], + encoder_hidden_states_mask=tokens_mask, + timestep=timestep / 1000, + img_shapes=img_shapes, + ) + + predicted_flow = model.unpack_latents( + packed_predicted_flow, + latent_height, + latent_width, + ) + + flow = latent_noise - scaled_latent_image + model_output_data = { + 'loss_type': 'target', + 'timestep': timestep, + #unpatchify, to make the shape match the mask shape of masked training: + 'predicted': model.unpatchify_latents(predicted_flow), + 'target': model.unpatchify_latents(flow), + } + + if config.debug_mode: + with torch.no_grad(): + predicted_scaled_latent_image = scaled_noisy_latent_image - predicted_flow * sigma + self._save_tokens("7-prompt", batch['tokens'], model.tokenizer, config, train_progress) + self._save_latent("1-noise", latent_noise, config, train_progress) + self._save_latent("2-noisy_image", scaled_noisy_latent_image, config, train_progress) + self._save_latent("3-predicted_flow", predicted_flow, config, train_progress) + self._save_latent("4-flow", flow, config, train_progress) + self._save_latent("5-predicted_image", predicted_scaled_latent_image, config, train_progress) + self._save_latent("6-image", scaled_latent_image, config, train_progress) + + return model_output_data + + def calculate_loss( + self, + model: LensModel, + batch: dict, + data: dict, + config: TrainConfig, + ) -> Tensor: + return self._flow_matching_losses( + batch=batch, + data=data, + config=config, + train_device=self.train_device, + sigmas=model.noise_scheduler.sigmas, + ).mean() + + def prepare_text_caching(self, model: LensModel, config: TrainConfig): + model.release() + model.materialize_text_encoder(self.train_device) + model.eval() + torch_gc() diff --git a/modules/modelSetup/BasePixArtAlphaSetup.py b/modules/modelSetup/BasePixArtAlphaSetup.py index 3069b4884..e7b21cd6e 100644 --- a/modules/modelSetup/BasePixArtAlphaSetup.py +++ b/modules/modelSetup/BasePixArtAlphaSetup.py @@ -51,12 +51,8 @@ def setup_optimizations( model: PixArtAlphaModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.vae.enable_gradient_checkpointing() - model.transformer_offload_conductor = \ - enable_checkpointing_for_basic_transformer_blocks(model.transformer, config, offload_enabled=True) - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config) + model.transformer_offload_conductor = enable_checkpointing_for_basic_transformer_blocks(model.transformer, config, config.transformer, offload_enabled=True) + model.text_encoder_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, @@ -340,7 +336,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: PixArtAlphaModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseQwenSetup.py b/modules/modelSetup/BaseQwenSetup.py index a8a7be8f6..68d4acc2a 100644 --- a/modules/modelSetup/BaseQwenSetup.py +++ b/modules/modelSetup/BaseQwenSetup.py @@ -46,12 +46,9 @@ def setup_optimizations( model: QwenModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_qwen_transformer(model.transformer, config) - if model.text_encoder is not None: - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_qwen25vl_encoder_layers(model.text_encoder, config) + model.transformer_offload_conductor = enable_checkpointing_for_qwen_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None: + model.text_encoder_offload_conductor = enable_checkpointing_for_qwen25vl_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, @@ -189,7 +186,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: QwenModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseSanaSetup.py b/modules/modelSetup/BaseSanaSetup.py index 0c9ea6da0..d73086d01 100644 --- a/modules/modelSetup/BaseSanaSetup.py +++ b/modules/modelSetup/BaseSanaSetup.py @@ -52,12 +52,8 @@ def setup_optimizations( config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - # model.vae.enable_gradient_checkpointing() - model.transformer_offload_conductor = \ - enable_checkpointing_for_sana_transformer(model.transformer, config) - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_gemma_layers(model.text_encoder, config) + model.transformer_offload_conductor = enable_checkpointing_for_sana_transformer(model.transformer, config, config.transformer) + model.text_encoder_offload_conductor = enable_checkpointing_for_gemma_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, @@ -263,7 +259,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: SanaModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseStableDiffusion3Setup.py b/modules/modelSetup/BaseStableDiffusion3Setup.py index de5dc04e8..37f8acfcb 100644 --- a/modules/modelSetup/BaseStableDiffusion3Setup.py +++ b/modules/modelSetup/BaseStableDiffusion3Setup.py @@ -48,16 +48,13 @@ def setup_optimizations( model: StableDiffusion3Model, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_stable_diffusion_3_transformer(model.transformer, config) - if model.text_encoder_1 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config) - if model.text_encoder_2 is not None: - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config) - if model.text_encoder_3 is not None: - model.text_encoder_3_offload_conductor = \ - enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config) + model.transformer_offload_conductor = enable_checkpointing_for_stable_diffusion_3_transformer(model.transformer, config, config.transformer) + if model.text_encoder_1 is not None: + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) + if model.text_encoder_2 is not None: + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) + if model.text_encoder_3 is not None: + model.text_encoder_3_offload_conductor = enable_checkpointing_for_t5_encoder_layers(model.text_encoder_3, config, config.text_encoder_3) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, @@ -361,7 +358,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: StableDiffusion3Model, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseStableDiffusionSetup.py b/modules/modelSetup/BaseStableDiffusionSetup.py index 8cf63ac07..90e5da8e0 100644 --- a/modules/modelSetup/BaseStableDiffusionSetup.py +++ b/modules/modelSetup/BaseStableDiffusionSetup.py @@ -51,11 +51,12 @@ def setup_optimizations( model: StableDiffusionModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.vae.enable_gradient_checkpointing() + if config.unet.checkpointing_or_offloading_enabled(): model.unet.enable_gradient_checkpointing() - enable_checkpointing_for_basic_transformer_blocks(model.unet, config, offload_enabled=False) - enable_checkpointing_for_clip_encoder_layers(model.text_encoder, config) + enable_checkpointing_for_basic_transformer_blocks(model.unet, config, config.unet, offload_enabled=False) + if config.vae.checkpointing_enabled(): + model.vae.enable_gradient_checkpointing() + enable_checkpointing_for_clip_encoder_layers(model.text_encoder, config, config.text_encoder) if config.force_circular_padding: apply_circular_padding_to_conv2d(model.vae) @@ -340,7 +341,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: StableDiffusionModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseStableDiffusionXLSetup.py b/modules/modelSetup/BaseStableDiffusionXLSetup.py index 61cb1e457..5903c2a31 100644 --- a/modules/modelSetup/BaseStableDiffusionXLSetup.py +++ b/modules/modelSetup/BaseStableDiffusionXLSetup.py @@ -48,11 +48,11 @@ def setup_optimizations( model: StableDiffusionXLModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): + if config.unet.checkpointing_or_offloading_enabled(): model.unet.enable_gradient_checkpointing() - enable_checkpointing_for_basic_transformer_blocks(model.unet, config, offload_enabled=False) - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config) - enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config) + enable_checkpointing_for_basic_transformer_blocks(model.unet, config, config.unet, offload_enabled=False) + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_1, config, config.text_encoder) + enable_checkpointing_for_clip_encoder_layers(model.text_encoder_2, config, config.text_encoder_2) if config.force_circular_padding: apply_circular_padding_to_conv2d(model.vae) @@ -388,7 +388,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: StableDiffusionXLModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseWuerstchenSetup.py b/modules/modelSetup/BaseWuerstchenSetup.py index e1a4c39d2..a5aac3b4b 100644 --- a/modules/modelSetup/BaseWuerstchenSetup.py +++ b/modules/modelSetup/BaseWuerstchenSetup.py @@ -54,9 +54,10 @@ def setup_optimizations( model: WuerstchenModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): + if config.prior.checkpointing_enabled(): model.prior_prior.enable_gradient_checkpointing() - enable_checkpointing_for_clip_encoder_layers(model.prior_text_encoder, config) + if config.text_encoder.checkpointing_or_offloading_enabled(): + enable_checkpointing_for_clip_encoder_layers(model.prior_text_encoder, config, config.text_encoder) if config.force_circular_padding: apply_circular_padding_to_conv2d(model.decoder_vqgan) @@ -355,7 +356,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: WuerstchenModel, config: TrainConfig): - model.to(self.temp_device) + model.release() if not config.train_text_encoder_or_embedding(): model.text_encoder_to(self.train_device) diff --git a/modules/modelSetup/BaseZImageSetup.py b/modules/modelSetup/BaseZImageSetup.py index b822d7304..dda618df9 100644 --- a/modules/modelSetup/BaseZImageSetup.py +++ b/modules/modelSetup/BaseZImageSetup.py @@ -47,12 +47,9 @@ def setup_optimizations( model: ZImageModel, config: TrainConfig, ): - if config.gradient_checkpointing.enabled(): - model.transformer_offload_conductor = \ - enable_checkpointing_for_z_image_transformer(model.transformer, config) - if model.text_encoder is not None: - model.text_encoder_offload_conductor = \ - enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config) + model.transformer_offload_conductor = enable_checkpointing_for_z_image_transformer(model.transformer, config, config.transformer) + if model.text_encoder is not None: + model.text_encoder_offload_conductor = enable_checkpointing_for_qwen3_encoder_layers(model.text_encoder, config, config.text_encoder) model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ config.weight_dtypes().transformer, @@ -172,7 +169,7 @@ def calculate_loss( ).mean() def prepare_text_caching(self, model: ZImageModel, config: TrainConfig): - model.to(self.temp_device) + model.release() model.text_encoder_to(self.train_device) model.eval() diff --git a/modules/modelSetup/LensFineTuneSetup.py b/modules/modelSetup/LensFineTuneSetup.py new file mode 100644 index 000000000..d2b08cc0b --- /dev/null +++ b/modules/modelSetup/LensFineTuneSetup.py @@ -0,0 +1,95 @@ +from modules.model.LensModel import LensModel +from modules.modelSetup.BaseLensSetup import BaseLensSetup +from modules.modelSetup.BaseModelSetup import BaseModelSetup +from modules.util import factory +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.ModuleFilter import ModuleFilter +from modules.util.NamedParameterGroup import NamedParameterGroupCollection +from modules.util.optimizer_util import init_model_parameters +from modules.util.TrainProgress import TrainProgress + +import torch + + +@factory.register(BaseModelSetup, ModelType.LENS, TrainingMethod.FINE_TUNE) +class LensFineTuneSetup( + BaseLensSetup, +): + def __init__( + self, + train_device: torch.device, + temp_device: torch.device, + debug_mode: bool, + ): + super().__init__( + train_device=train_device, + temp_device=temp_device, + debug_mode=debug_mode, + ) + + def create_parameters( + self, + model: LensModel, + config: TrainConfig, + ) -> NamedParameterGroupCollection: + parameter_group_collection = NamedParameterGroupCollection() + + self._create_model_part_parameters(parameter_group_collection, "transformer", model.transformer, config.transformer, + freeze=ModuleFilter.create(config), debug=config.debug_mode) + return parameter_group_collection + + def __setup_requires_grad( + self, + model: LensModel, + config: TrainConfig, + ): + self._setup_model_part_requires_grad("transformer", model.transformer, config.transformer, model.train_progress) + model.vae.requires_grad_(False) + if model.text_encoder is not None: + model.text_encoder.requires_grad_(False) + + + def setup_model( + self, + model: LensModel, + config: TrainConfig, + ): + params = self.create_parameters(model, config) + self.__setup_requires_grad(model, config) + init_model_parameters(model, params, self.train_device) + + def setup_train_device( + self, + model: LensModel, + config: TrainConfig, + ): + vae_on_train_device = not config.latent_caching + text_encoder_on_train_device = not config.latent_caching + + # the encoder is always on-demand for Lens, so materialize it when training needs it and + # discard it (freeing VRAM) otherwise -- text_encoder_to would be a no-op on the proxy. + if text_encoder_on_train_device: + model.materialize_text_encoder(self.train_device) + else: + model.release_text_encoder() + model.vae_to(self.train_device if vae_on_train_device else self.temp_device) + model.transformer_to(self.train_device) + + if model.text_encoder is not None: + model.text_encoder.eval() + model.vae.eval() + + if config.transformer.train: + model.transformer.train() + else: + model.transformer.eval() + + def after_optimizer_step( + self, + model: LensModel, + config: TrainConfig, + train_progress: TrainProgress + ): + self.__setup_requires_grad(model, config) diff --git a/modules/modelSetup/LensLoRASetup.py b/modules/modelSetup/LensLoRASetup.py new file mode 100644 index 000000000..760307fad --- /dev/null +++ b/modules/modelSetup/LensLoRASetup.py @@ -0,0 +1,107 @@ +from modules.model.LensModel import LensModel +from modules.modelSetup.BaseLensSetup import BaseLensSetup +from modules.modelSetup.BaseModelSetup import BaseModelSetup +from modules.module.LoRAModule import LoRAModuleWrapper +from modules.util import factory +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.NamedParameterGroup import NamedParameterGroupCollection +from modules.util.optimizer_util import init_model_parameters +from modules.util.TrainProgress import TrainProgress + +import torch + + +@factory.register(BaseModelSetup, ModelType.LENS, TrainingMethod.LORA) +class LensLoRASetup( + BaseLensSetup, +): + def __init__( + self, + train_device: torch.device, + temp_device: torch.device, + debug_mode: bool, + ): + super().__init__( + train_device=train_device, + temp_device=temp_device, + debug_mode=debug_mode, + ) + + def create_parameters( + self, + model: LensModel, + config: TrainConfig, + ) -> NamedParameterGroupCollection: + parameter_group_collection = NamedParameterGroupCollection() + + self._create_model_part_parameters(parameter_group_collection, "transformer", model.transformer_lora, config.transformer) + return parameter_group_collection + + def __setup_requires_grad( + self, + model: LensModel, + config: TrainConfig, + ): + if model.text_encoder is not None: + model.text_encoder.requires_grad_(False) + model.transformer.requires_grad_(False) + model.vae.requires_grad_(False) + + self._setup_model_part_requires_grad("transformer", model.transformer_lora, config.transformer, model.train_progress) + + def setup_model( + self, + model: LensModel, + config: TrainConfig, + ): + model.transformer_lora = LoRAModuleWrapper( + model.transformer, "transformer", config, config.layer_filter.split(",") + ) + + if model.lora_state_dict: + model.transformer_lora.load_state_dict(model.lora_state_dict) + model.lora_state_dict = None + + model.transformer_lora.set_dropout(config.dropout_probability) + model.transformer_lora.to(dtype=config.lora_weight_dtype.torch_dtype()) + model.transformer_lora.hook_to_module() + + params = self.create_parameters(model, config) + self.__setup_requires_grad(model, config) + init_model_parameters(model, params, self.train_device) + + def setup_train_device( + self, + model: LensModel, + config: TrainConfig, + ): + vae_on_train_device = not config.latent_caching + text_encoder_on_train_device = not config.latent_caching + + # the encoder is always on-demand for Lens, so materialize it when training needs it and + # discard it (freeing VRAM) otherwise -- text_encoder_to would be a no-op on the proxy. + if text_encoder_on_train_device: + model.materialize_text_encoder(self.train_device) + else: + model.release_text_encoder() + model.vae_to(self.train_device if vae_on_train_device else self.temp_device) + model.transformer_to(self.train_device) + + if model.text_encoder is not None: + model.text_encoder.eval() + model.vae.eval() + + if config.transformer.train: + model.transformer.train() + else: + model.transformer.eval() + + def after_optimizer_step( + self, + model: LensModel, + config: TrainConfig, + train_progress: TrainProgress + ): + self.__setup_requires_grad(model, config) diff --git a/modules/trainer/GenericTrainer.py b/modules/trainer/GenericTrainer.py index dd17ad76e..990ee6ab3 100644 --- a/modules/trainer/GenericTrainer.py +++ b/modules/trainer/GenericTrainer.py @@ -139,8 +139,7 @@ def start(self): self.model_setup.setup_optimizations(self.model, self.config) self.model_setup.setup_train_device(self.model, self.config) self.model_setup.setup_model(self.model, self.config) - self.model.to(self.temp_device) - self.model.eval() + self.model.release() torch_gc() self.callbacks.on_update_status("creating the data loader/caching") @@ -253,8 +252,7 @@ def on_sample_custom(sampler_output: ModelSamplerOutput): on_sample = on_sample_custom if is_custom_sample else on_sample_default on_update_progress = self.callbacks.on_update_sample_custom_progress if is_custom_sample else self.callbacks.on_update_sample_default_progress - self.model.to(self.temp_device) - self.model.eval() + self.model.release() sample_config = copy.copy(sample_config) sample_config.from_train_config(self.config) @@ -717,7 +715,7 @@ def sample_commands_fun(): backup = self.commands.get_and_reset_backup_command() save = self.commands.get_and_reset_save_command() if multi.is_master() and (backup or save): - self.model.to(self.temp_device) + self.model.release() if backup: self.__backup(train_progress, True, step_tqdm.write) if save: @@ -842,7 +840,7 @@ def sample_commands_fun(): def end(self): if self.one_step_trained: - self.model.to(self.temp_device) + self.model.release() if self.config.backup_before_save and multi.is_master(): self.__backup(self.model.train_progress) @@ -875,7 +873,7 @@ def end(self): ) if self.model is not None: - self.model.to(self.temp_device) + self.model.release() if multi.is_master(): self.tensorboard.close() diff --git a/modules/ui/ConvertModelUI.py b/modules/ui/ConvertModelUI.py index 6cb1b507a..a32987de4 100644 --- a/modules/ui/ConvertModelUI.py +++ b/modules/ui/ConvertModelUI.py @@ -67,6 +67,7 @@ def main_frame(self, master): ("Flux Dev", ModelType.FLUX_DEV_1), ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), ("Flux 2", ModelType.FLUX_2), + ("Lens", ModelType.LENS), ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), ("Chroma1", ModelType.CHROMA_1), #TODO does this just work? HiDream is not here ("QwenImage", ModelType.QWEN), #TODO does this just work? HiDream is not here diff --git a/modules/ui/ModelTab.py b/modules/ui/ModelTab.py index ff17ea3ba..73ce614a3 100644 --- a/modules/ui/ModelTab.py +++ b/modules/ui/ModelTab.py @@ -45,307 +45,68 @@ def refresh_ui(self): base_frame.grid_columnconfigure(3, weight=0) base_frame.grid_columnconfigure(4, weight=1) - if self.train_config.model_type.is_stable_diffusion(): #TODO simplify - self.__setup_stable_diffusion_ui(base_frame) - if self.train_config.model_type.is_stable_diffusion_3(): - self.__setup_stable_diffusion_3_ui(base_frame) - elif self.train_config.model_type.is_stable_diffusion_xl(): - self.__setup_stable_diffusion_xl_ui(base_frame) - elif self.train_config.model_type.is_wuerstchen(): - self.__setup_wuerstchen_ui(base_frame) - elif self.train_config.model_type.is_pixart(): - self.__setup_pixart_alpha_ui(base_frame) - elif self.train_config.model_type.is_flux_1(): - self.__setup_flux_ui(base_frame) - elif self.train_config.model_type.is_flux_2(): - self.__setup_flux_2_ui(base_frame) - elif self.train_config.model_type.is_z_image(): - self.__setup_z_image_ui(base_frame) - elif self.train_config.model_type.is_chroma(): - self.__setup_chroma_ui(base_frame) - elif self.train_config.model_type.is_qwen(): - self.__setup_qwen_ui(base_frame) - elif self.train_config.model_type.is_sana(): - self.__setup_sana_ui(base_frame) - elif self.train_config.model_type.is_hunyuan_video(): - self.__setup_hunyuan_video_ui(base_frame) - elif self.train_config.model_type.is_hi_dream(): - self.__setup_hi_dream_ui(base_frame) - elif self.train_config.model_type.is_ernie(): - self.__setup_ernie_ui(base_frame) - - def __setup_stable_diffusion_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_unet=True, - has_text_encoder=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method in [ - TrainingMethod.FINE_TUNE, - TrainingMethod.FINE_TUNE_VAE, - ], - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_stable_diffusion_3_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_text_encoder_3=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_flux_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_flux_2_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) + self.__setup_ui(base_frame) - def __setup_z_image_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) + def __setup_ui(self, frame): + model_type = self.train_config.model_type + training_method = self.train_config.training_method + parts = model_type.model_parts() - def __setup_ernie_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + # The transformer override path exists only for these architectures; SD3, PixArt, Sana + # and HiDream have a transformer but expose no override field. + allow_override_transformer = ( + model_type.is_flux() + or model_type.is_z_image() + or model_type.is_ernie() + or model_type.is_chroma() + or model_type.is_qwen() + or model_type.is_hunyuan_video() + or model_type.is_lens() ) - def __setup_chroma_ui(self, frame): row = 0 row = self.__create_base_dtype_components(frame, row) row = self.__create_base_components( frame, row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) + has_unet="unet" in parts, + has_prior="prior" in parts, + allow_override_prior=model_type.is_stable_cascade(), + has_transformer="transformer" in parts, + allow_override_transformer=allow_override_transformer, + # Lens: GPT-OSS dtype is fixed by MXFP4 + checkpoint; not user-configurable. + has_text_encoder=not model_type.has_multiple_text_encoders() and not model_type.is_lens(), + has_text_encoder_1=model_type.has_multiple_text_encoders(), + has_text_encoder_2="text_encoder_2" in parts, + has_text_encoder_3="text_encoder_3" in parts, + has_text_encoder_4="text_encoder_4" in parts, + allow_override_text_encoder_4="text_encoder_4" in parts, + has_vae="vae" in parts, + ) + if "effnet_encoder" in parts: + row = self.__create_effnet_encoder_components(frame, row) + if "decoder" in parts: + row = self.__create_decoder_components(frame, row, "decoder_text_encoder" in parts) + + if model_type.is_sana(): + allow_safetensors = training_method != TrainingMethod.FINE_TUNE + elif model_type.is_wuerstchen(): + allow_safetensors = training_method != TrainingMethod.FINE_TUNE \ + or model_type.is_stable_cascade() + else: + allow_safetensors = True + + if model_type.is_stable_diffusion(): + allow_diffusers = training_method in [TrainingMethod.FINE_TUNE, TrainingMethod.FINE_TUNE_VAE] + else: + allow_diffusers = training_method == TrainingMethod.FINE_TUNE - def __setup_qwen_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_stable_diffusion_xl_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_unet=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_wuerstchen_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_prior=True, - allow_override_prior=self.train_config.model_type.is_stable_cascade(), - has_text_encoder=True, - ) - row = self.__create_effnet_encoder_components(frame, row) - row = self.__create_decoder_components(frame, row, self.train_config.model_type.is_wuerstchen_v2()) - row = self.__create_output_components( - frame, - row, - allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE - or self.train_config.model_type.is_stable_cascade(), - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_pixart_alpha_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_sana_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=self.train_config.training_method != TrainingMethod.FINE_TUNE, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_hunyuan_video_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - allow_override_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_vae=True, - ) - row = self.__create_output_components( - frame, - row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, - ) - - def __setup_hi_dream_ui(self, frame): - row = 0 - row = self.__create_base_dtype_components(frame, row) - row = self.__create_base_components( - frame, - row, - has_transformer=True, - has_text_encoder_1=True, - has_text_encoder_2=True, - has_text_encoder_3=True, - has_text_encoder_4=True, - allow_override_text_encoder_4=True, - has_vae=True, - ) row = self.__create_output_components( frame, row, - allow_safetensors=True, - allow_diffusers=self.train_config.training_method == TrainingMethod.FINE_TUNE, - allow_legacy_safetensors=self.train_config.training_method == TrainingMethod.LORA, + allow_safetensors=allow_safetensors, + allow_diffusers=allow_diffusers, + allow_legacy_safetensors=training_method == TrainingMethod.LORA, ) def __create_dtype_options(self, include_gguf: bool=False, include_a8: bool=False) -> list[tuple[str, DataType]]: diff --git a/modules/ui/OffloadingWindow.py b/modules/ui/OffloadingWindow.py deleted file mode 100644 index 54035e121..000000000 --- a/modules/ui/OffloadingWindow.py +++ /dev/null @@ -1,75 +0,0 @@ -from modules.util.config.TrainConfig import TrainConfig -from modules.util.enum.GradientCheckpointingMethod import ( - GradientCheckpointingMethod, -) -from modules.util.ui import components -from modules.util.ui.ui_utils import set_window_icon -from modules.util.ui.UIState import UIState - -import customtkinter as ctk - - -class OffloadingWindow(ctk.CTkToplevel): - def __init__( - self, - parent, - config: TrainConfig, - ui_state: UIState, - *args, **kwargs, - ): - super().__init__(parent, *args, **kwargs) - - self.config = config - self.ui_state = ui_state - self.image_preview_file_index = 0 - self.ax = None - self.canvas = None - - self.title("Offloading") - self.geometry("800x400") - self.resizable(True, True) - - self.grid_rowconfigure(0, weight=1) - self.grid_columnconfigure(0, weight=1) - - frame = self.__content_frame(self) - frame.grid(row=0, column=0, sticky='nsew') - components.button(self, 1, 0, "ok", self.__ok) - - self.wait_visibility() - self.grab_set() - self.focus_set() - self.after(200, lambda: set_window_icon(self)) - - - def __content_frame(self, master): - frame = ctk.CTkScrollableFrame(master, fg_color="transparent") - frame.grid_columnconfigure(0, weight=1) - frame.grid_columnconfigure(1, weight=1) - - # timestep distribution - components.label(frame, 0, 0, "Gradient checkpointing", - tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") - components.options(frame, 0, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, - "gradient_checkpointing") - - # gradient checkpointing layer offloading - components.label(frame, 1, 0, "Async Offloading", - tooltip="Enables Asynchronous offloading.") - components.switch(frame, 1, 1, self.ui_state, "enable_async_offloading") - - # gradient checkpointing layer offloading - components.label(frame, 2, 0, "Offload Activations", - tooltip="Enables Activation Offloading") - components.switch(frame, 2, 1, self.ui_state, "enable_activation_offloading") - - # gradient checkpointing layer offloading - components.label(frame, 3, 0, "Layer offload fraction", - tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") - components.entry(frame, 3, 1, self.ui_state, "layer_offload_fraction") - - frame.pack(fill="both", expand=1) - return frame - - def __ok(self): - self.destroy() diff --git a/modules/ui/SampleWindow.py b/modules/ui/SampleWindow.py index 0f91ad2fa..bdd664d84 100644 --- a/modules/ui/SampleWindow.py +++ b/modules/ui/SampleWindow.py @@ -144,7 +144,7 @@ def __load_model(self) -> BaseModel: model_setup.setup_optimizations(model, self.initial_train_config) model_setup.setup_train_device(model, self.initial_train_config) model_setup.setup_model(model, self.initial_train_config) - model.to(torch.device(self.initial_train_config.temp_device)) + model.release() return model diff --git a/modules/ui/TopBar.py b/modules/ui/TopBar.py index 820fdb71a..e416ca9dd 100644 --- a/modules/ui/TopBar.py +++ b/modules/ui/TopBar.py @@ -95,6 +95,7 @@ def __init__( ("Flux Dev.1", ModelType.FLUX_DEV_1), ("Flux Fill Dev", ModelType.FLUX_FILL_DEV_1), ("Flux 2 [Dev, Klein]", ModelType.FLUX_2), + ("Lens", ModelType.LENS), ("Sana", ModelType.SANA), ("Hunyuan Video", ModelType.HUNYUAN_VIDEO), ("HiDream Full", ModelType.HI_DREAM_FULL), @@ -112,37 +113,13 @@ def __create_training_method(self): if self.training_method: self.training_method.destroy() - values = [] - #TODO simplify - if self.train_config.model_type.is_stable_diffusion(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ("Fine Tune VAE", TrainingMethod.FINE_TUNE_VAE), - ] - elif self.train_config.model_type.is_stable_diffusion_3() \ - or self.train_config.model_type.is_stable_diffusion_xl() \ - or self.train_config.model_type.is_wuerstchen() \ - or self.train_config.model_type.is_pixart() \ - or self.train_config.model_type.is_flux_1() \ - or self.train_config.model_type.is_sana() \ - or self.train_config.model_type.is_hunyuan_video() \ - or self.train_config.model_type.is_hi_dream() \ - or self.train_config.model_type.is_chroma(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ("Embedding", TrainingMethod.EMBEDDING), - ] - elif self.train_config.model_type.is_qwen() \ - or self.train_config.model_type.is_z_image() \ - or self.train_config.model_type.is_flux_2() \ - or self.train_config.model_type.is_ernie(): - values = [ - ("Fine Tune", TrainingMethod.FINE_TUNE), - ("LoRA", TrainingMethod.LORA), - ] + labels = { + TrainingMethod.FINE_TUNE: "Fine Tune", + TrainingMethod.LORA: "LoRA", + TrainingMethod.EMBEDDING: "Embedding", + TrainingMethod.FINE_TUNE_VAE: "Fine Tune VAE", + } + values = [(labels[m], m) for m in self.train_config.model_type.supported_training_methods()] # training method self.training_method = components.options_kv( diff --git a/modules/ui/TrainUI.py b/modules/ui/TrainUI.py index b9fa0c04a..1217c06b4 100644 --- a/modules/ui/TrainUI.py +++ b/modules/ui/TrainUI.py @@ -311,6 +311,10 @@ def create_general_tab(self, master): tooltip="The device used for training. Can be \"cuda\", \"cuda:0\", \"cuda:1\" etc. Default:\"cuda\". Must be \"cuda\" for multi-GPU training.") components.entry(frame, 11, 1, self.ui_state, "train_device", required=True) + components.label(frame, 11, 2, "Async Offloading", + tooltip="Overlaps CPU<->GPU transfers with computation using CUDA streams. Applies to every offloaded component") + components.switch(frame, 11, 3, self.ui_state, "async_offloading") + components.label(frame, 12, 0, "Multi-GPU", tooltip="Enable multi-GPU training") components.switch(frame, 12, 1, self.ui_state, "multi_gpu") diff --git a/modules/ui/TrainingTab.py b/modules/ui/TrainingTab.py index f897bb8ce..aeffbc056 100644 --- a/modules/ui/TrainingTab.py +++ b/modules/ui/TrainingTab.py @@ -1,4 +1,3 @@ -from modules.ui.OffloadingWindow import OffloadingWindow from modules.ui.OptimizerParamsWindow import OptimizerParamsWindow from modules.ui.SchedulerParamsWindow import SchedulerParamsWindow from modules.ui.TimestepDistributionWindow import TimestepDistributionWindow @@ -6,13 +5,13 @@ from modules.util.config.TrainConfig import TrainConfig from modules.util.enum.DataType import DataType from modules.util.enum.EMAMode import EMAMode -from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod from modules.util.enum.LearningRateScaler import LearningRateScaler from modules.util.enum.LearningRateScheduler import LearningRateScheduler from modules.util.enum.LossScaler import LossScaler from modules.util.enum.LossWeight import LossWeight from modules.util.enum.Optimizer import Optimizer from modules.util.enum.TimestepDistribution import TimestepDistribution +from modules.util.enum.TrainingMethod import TrainingMethod from modules.util.optimizer_util import change_optimizer from modules.util.ui import components from modules.util.ui.UIState import UIState @@ -74,6 +73,8 @@ def refresh_ui(self): self.__setup_flux_ui(column_0, column_1, column_2) elif self.train_config.model_type.is_flux_2(): self.__setup_flux_2_ui(column_0, column_1, column_2) + elif self.train_config.model_type.is_lens(): + self.__setup_lens_ui(column_0, column_1, column_2) elif self.train_config.model_type.is_chroma(): self.__setup_chroma_ui(column_0, column_1, column_2) elif self.train_config.model_type.is_qwen(): @@ -99,6 +100,8 @@ def __setup_stable_diffusion_ui(self, column_0, column_1, column_2): self.__create_unet_frame(column_1, 1) self.__create_noise_frame(column_1, 2, supports_generalized_offset_noise=True) + if self.train_config.training_method == TrainingMethod.FINE_TUNE_VAE: + self.__create_vae_frame(column_2, 0) self.__create_masked_frame(column_2, 1) self.__create_loss_frame(column_2, 2) self.__create_layer_frame(column_2, 3) @@ -184,6 +187,19 @@ def __setup_flux_2_ui(self, column_0, column_1, column_2): self.__create_loss_frame(column_2, 2) self.__create_layer_frame(column_2, 3) + def __setup_lens_ui(self, column_0, column_1, column_2): + self.__create_base_frame(column_0, 0) + # the GPT-OSS encoder is always on-demand (MXFP4); layer offloading is incompatible with that, so hide it + self.__create_text_encoder_frame(column_0, 1, supports_clip_skip=False, supports_training=False, supports_sequence_length=False, supports_offloading=False) + + self.__create_base2_frame(column_1, 0) + self.__create_transformer_frame(column_1, 1, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_noise_frame(column_1, 2, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1) + self.__create_loss_frame(column_2, 2) + self.__create_layer_frame(column_2, 3) + def __setup_chroma_ui(self, column_0, column_1, column_2): self.__create_base_frame(column_0, 0) self.__create_text_encoder_frame(column_0, 1) @@ -375,19 +391,6 @@ def __create_base2_frame(self, master, row, video_training_enabled: bool=False, components.entry(frame, row, 1, self.ui_state, "ema_update_step_interval") row += 1 - # gradient checkpointing - components.label(frame, row, 0, "Gradient checkpointing", - tooltip="Enables gradient checkpointing. This reduces memory usage, but increases training time") - components.options_adv(frame, row, 1, [str(x) for x in list(GradientCheckpointingMethod)], self.ui_state, - "gradient_checkpointing", adv_command=self.__open_offloading_window) - row += 1 - - # gradient checkpointing layer offloading - components.label(frame, row, 0, "Layer offload fraction", - tooltip="Enables offloading of individual layers during training to reduce VRAM usage. Increases training time and uses more RAM. Only available if checkpointing is set to CPU_OFFLOADED. values between 0 and 1, 0=disabled") - components.entry(frame, row, 1, self.ui_state, "layer_offload_fraction") - row += 1 - # train dtype components.label(frame, row, 0, "Train Data Type", tooltip="The mixed precision data type used for training. This can increase training speed, but reduces precision") @@ -434,7 +437,29 @@ def __create_base2_frame(self, master, row, video_training_enabled: bool=False, tooltip="Enables circular padding for all conv layers to better train seamless images") components.switch(frame, row, 1, self.ui_state, "force_circular_padding") - def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supports_training=True, supports_sequence_length=False): + def __create_offloading_widgets(self, frame, row, part, supports_checkpointing=True, supports_activation_offloading=False, supports_layer_offloading=True): + # per-component offloading / checkpointing controls (bound to the model part config) + if supports_checkpointing: + components.label(frame, row, 0, "Gradient Checkpointing", + tooltip="Enables gradient checkpointing for this component. Reduces VRAM usage at the cost of training speed") + components.switch(frame, row, 1, self.ui_state, f"{part}.gradient_checkpointing") + row += 1 + + if supports_layer_offloading: + components.label(frame, row, 0, "Layer Offload Fraction", + tooltip="Fraction of this component's layers to offload to CPU to reduce VRAM usage. Increases training time and RAM usage. 0=disabled, 1=all layers") + components.entry(frame, row, 1, self.ui_state, f"{part}.offload_fraction") + row += 1 + + if supports_activation_offloading: + components.label(frame, row, 0, "Offload Activations", + tooltip="Offloads this component's activations to CPU during training to reduce VRAM usage") + components.switch(frame, row, 1, self.ui_state, f"{part}.activation_offloading") + row += 1 + + return row + + def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supports_training=True, supports_sequence_length=False, supports_offloading=True): frame = ctk.CTkFrame(master=master, corner_radius=5) frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) @@ -445,6 +470,12 @@ def __create_text_encoder_frame(self, master, row, supports_clip_skip=True, supp tooltip="Enables training the text encoder model") components.switch(frame, row, 1, self.ui_state, "text_encoder.train") row += 1 + else: + # no Train switch to act as the frame's header, so add an explicit one + components.label(frame, row, 0, "Text Encoder") + row += 1 + + row = self.__create_offloading_widgets(frame, row, "text_encoder", supports_checkpointing=supports_training, supports_layer_offloading=supports_offloading) # dropout components.label(frame, row, 0, "Caption Dropout Probability", @@ -509,6 +540,8 @@ def __create_text_encoder_n_frame( components.switch(frame, row, 1, self.ui_state, f"text_encoder{suffix}.train") row += 1 + row = self.__create_offloading_widgets(frame, row, f"text_encoder{suffix}") + # train text encoder embedding components.label(frame, row, 0, f"Train Text Encoder {i} Embedding", tooltip=f"Enables training embeddings for the text encoder {i} model") @@ -566,82 +599,119 @@ def __create_unet_frame(self, master, row): frame = ctk.CTkFrame(master=master, corner_radius=5) frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) + row = 0 # train unet - components.label(frame, 0, 0, "Train UNet", + components.label(frame, row, 0, "Train UNet", tooltip="Enables training the UNet model") - components.switch(frame, 0, 1, self.ui_state, "unet.train") + components.switch(frame, row, 1, self.ui_state, "unet.train") + row += 1 + + row = self.__create_offloading_widgets(frame, row, "unet", supports_activation_offloading=True) # train unet epochs - components.label(frame, 1, 0, "Stop Training After", + components.label(frame, row, 0, "Stop Training After", tooltip="When to stop training the UNet") - components.time_entry(frame, 1, 1, self.ui_state, "unet.stop_training_after", "unet.stop_training_after_unit", + components.time_entry(frame, row, 1, self.ui_state, "unet.stop_training_after", "unet.stop_training_after_unit", supports_time_units=False) + row += 1 # unet learning rate - components.label(frame, 2, 0, "UNet Learning Rate", + components.label(frame, row, 0, "UNet Learning Rate", tooltip="The learning rate of the UNet. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "unet.learning_rate") + components.entry(frame, row, 1, self.ui_state, "unet.learning_rate") + row += 1 # rescale noise scheduler to zero terminal SNR - rescale_label = components.label(frame, 3, 0, "Rescale Noise Scheduler + V-pred", + rescale_label = components.label(frame, row, 0, "Rescale Noise Scheduler + V-pred", tooltip="Rescales the noise scheduler to a zero terminal signal to noise ratio and switches the model to a v-prediction target") rescale_label.configure(wraplength=130, justify="left") - components.switch(frame, 3, 1, self.ui_state, "rescale_noise_scheduler_to_zero_terminal_snr") + components.switch(frame, row, 1, self.ui_state, "rescale_noise_scheduler_to_zero_terminal_snr") + row += 1 + + def __create_vae_frame(self, master, row): + frame = ctk.CTkFrame(master=master, corner_radius=5) + frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") + frame.grid_columnconfigure(0, weight=1) + row = 0 + + components.label(frame, row, 0, "Train VAE", + tooltip="Enables training the VAE model") + components.switch(frame, row, 1, self.ui_state, "vae.train") + row += 1 + + components.label(frame, row, 0, "Gradient Checkpointing", + tooltip="Enables gradient checkpointing for the VAE. Reduces VRAM usage at the cost of training speed") + components.switch(frame, row, 1, self.ui_state, "vae.gradient_checkpointing") + row += 1 def __create_prior_frame(self, master, row): frame = ctk.CTkFrame(master=master, corner_radius=5) frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) + row = 0 # train prior - components.label(frame, 0, 0, "Train Prior", + components.label(frame, row, 0, "Train Prior", tooltip="Enables training the Prior model") - components.switch(frame, 0, 1, self.ui_state, "prior.train") + components.switch(frame, row, 1, self.ui_state, "prior.train") + row += 1 + + row = self.__create_offloading_widgets(frame, row, "prior", supports_activation_offloading=True) # train prior epochs - components.label(frame, 1, 0, "Stop Training After", + components.label(frame, row, 0, "Stop Training After", tooltip="When to stop training the Prior") - components.time_entry(frame, 1, 1, self.ui_state, "prior.stop_training_after", "prior.stop_training_after_unit", + components.time_entry(frame, row, 1, self.ui_state, "prior.stop_training_after", "prior.stop_training_after_unit", supports_time_units=False) + row += 1 # prior learning rate - components.label(frame, 2, 0, "Prior Learning Rate", + components.label(frame, row, 0, "Prior Learning Rate", tooltip="The learning rate of the Prior. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "prior.learning_rate") + components.entry(frame, row, 1, self.ui_state, "prior.learning_rate") + row += 1 def __create_transformer_frame(self, master, row, supports_guidance_scale: bool = False, supports_force_attention_mask: bool = True): frame = ctk.CTkFrame(master=master, corner_radius=5) frame.grid(row=row, column=0, padx=5, pady=5, sticky="nsew") frame.grid_columnconfigure(0, weight=1) + row = 0 # train transformer - components.label(frame, 0, 0, "Train Transformer", + components.label(frame, row, 0, "Train Transformer", tooltip="Enables training the Transformer model") - components.switch(frame, 0, 1, self.ui_state, "transformer.train") + components.switch(frame, row, 1, self.ui_state, "transformer.train") + row += 1 + + row = self.__create_offloading_widgets(frame, row, "transformer", supports_activation_offloading=True) # train transformer epochs - components.label(frame, 1, 0, "Stop Training After", + components.label(frame, row, 0, "Stop Training After", tooltip="When to stop training the Transformer") - components.time_entry(frame, 1, 1, self.ui_state, "transformer.stop_training_after", "transformer.stop_training_after_unit", + components.time_entry(frame, row, 1, self.ui_state, "transformer.stop_training_after", "transformer.stop_training_after_unit", supports_time_units=False) + row += 1 # transformer learning rate - components.label(frame, 2, 0, "Transformer Learning Rate", + components.label(frame, row, 0, "Transformer Learning Rate", tooltip="The learning rate of the Transformer. Overrides the base learning rate") - components.entry(frame, 2, 1, self.ui_state, "transformer.learning_rate") + components.entry(frame, row, 1, self.ui_state, "transformer.learning_rate") + row += 1 if supports_force_attention_mask: # transformer learning rate - components.label(frame, 3, 0, "Force Attention Mask", + components.label(frame, row, 0, "Force Attention Mask", tooltip="Force enables passing of a text embedding attention mask to the transformer. This can improve training on shorter captions.") - components.switch(frame, 3, 1, self.ui_state, "transformer.attention_mask") + components.switch(frame, row, 1, self.ui_state, "transformer.attention_mask") + row += 1 if supports_guidance_scale: # guidance scale - components.label(frame, 4, 0, "Guidance Scale", + components.label(frame, row, 0, "Guidance Scale", tooltip="The guidance scale of guidance distilled models passed to the transformer during training.") - components.entry(frame, 4, 1, self.ui_state, "transformer.guidance_scale") + components.entry(frame, row, 1, self.ui_state, "transformer.guidance_scale") + row += 1 def __create_noise_frame(self, master, row, supports_generalized_offset_noise: bool = False, supports_dynamic_timestep_shifting: bool = False): frame = ctk.CTkFrame(master=master, corner_radius=5) @@ -844,10 +914,6 @@ def __open_timestep_distribution_window(self): window = TimestepDistributionWindow(self.master, self.train_config, self.ui_state) self.master.wait_window(window) - def __open_offloading_window(self): - window = OffloadingWindow(self.master, self.train_config, self.ui_state) - self.master.wait_window(window) - def __restore_optimizer_config(self, *args): optimizer_config = change_optimizer(self.train_config) self.ui_state.get_var("optimizer").update(optimizer_config) diff --git a/modules/util/LayerOffloadConductor.py b/modules/util/LayerOffloadConductor.py index 0b7e6c7ca..062d1391e 100644 --- a/modules/util/LayerOffloadConductor.py +++ b/modules/util/LayerOffloadConductor.py @@ -2,7 +2,7 @@ import random from typing import Any -from modules.util.config.TrainConfig import TrainConfig +from modules.util.config.TrainConfig import TrainConfig, TrainModelPartConfig from modules.util.quantization_util import get_offload_tensor_bytes, offload_quantized from modules.util.torch_util import ( create_stream_context, @@ -566,6 +566,7 @@ def __init__( self, module: nn.Module, config: TrainConfig, + part: TrainModelPartConfig, ): super().__init__() @@ -573,16 +574,16 @@ def __init__( self.__layers = [] self.__layer_device_map = [] - self.__layer_offload_fraction = config.layer_offload_fraction + self.__layer_offload_fraction = part.offload_fraction self.__layer_activations_included_offload_param_indices_map = [] self.__train_device = torch.device(config.train_device) self.__temp_device = torch.device(config.temp_device) - self.__offload_activations = config.gradient_checkpointing.offload() and config.enable_activation_offloading - self.__offload_layers = config.gradient_checkpointing.offload() and config.layer_offload_fraction > 0 - self.__async_transfer = self.__train_device.type == "cuda" and config.enable_async_offloading + self.__offload_activations = part.activation_offloading + self.__offload_layers = part.offload_fraction > 0 + self.__async_transfer = self.__train_device.type == "cuda" and config.async_offloading if self.__async_transfer: self.__train_stream = torch.cuda.default_stream(self.__train_device) diff --git a/modules/util/ModelNames.py b/modules/util/ModelNames.py index 8dec9a9bf..c58a2d284 100644 --- a/modules/util/ModelNames.py +++ b/modules/util/ModelNames.py @@ -25,6 +25,10 @@ def __init__( include_text_encoder_2: bool = True, include_text_encoder_3: bool = True, include_text_encoder_4: bool = True, + text_encoder_on_demand: bool = False, + text_encoder_2_on_demand: bool = False, + text_encoder_3_on_demand: bool = False, + text_encoder_4_on_demand: bool = False, ): self.base_model = base_model self.prior_model = prior_model @@ -40,6 +44,10 @@ def __init__( self.include_text_encoder_2 = include_text_encoder_2 self.include_text_encoder_3 = include_text_encoder_3 self.include_text_encoder_4 = include_text_encoder_4 + self.text_encoder_on_demand = text_encoder_on_demand + self.text_encoder_2_on_demand = text_encoder_2_on_demand + self.text_encoder_3_on_demand = text_encoder_3_on_demand + self.text_encoder_4_on_demand = text_encoder_4_on_demand def all_embedding(self): if self.embedding is not None: diff --git a/modules/util/OnDemandModule.py b/modules/util/OnDemandModule.py new file mode 100644 index 000000000..e8eba8cbf --- /dev/null +++ b/modules/util/OnDemandModule.py @@ -0,0 +1,67 @@ +from collections.abc import Callable + +from modules.util.torch_util import torch_gc + +from torch import nn + + +# A persistent delegating proxy for a module loaded on demand and discarded +# (weights freed) after use. Always truthy, so presence gates and captured +# references (e.g. MGDS Encode*Text nodes) stay valid while no weights are loaded. +# On-demand modules are frozen and inference-only by definition, so the proxy +# enforces that invariant: eval/requires_grad_(False) are accepted, while +# train(True)/requires_grad_(True) raise. Other attribute access delegates to the +# inner. The loader loads to cpu; the caller (model) is responsible for any +# quantization and for moving the materialized module to the accelerator. +# +# Intentionally a plain object, not an nn.Module, to avoid submodule/parameter +# registration fighting the materialize/discard swap. +class OnDemandModule: + def __init__(self, loader: Callable[[], nn.Module]): + self._inner: nn.Module | None = None + self._loader = loader + + def __bool__(self): + return True + + @property + def inner(self) -> nn.Module | None: + return self._inner + + def __call__(self, *args, **kwargs): + return self._inner(*args, **kwargs) + + def __getattr__(self, name): + if self._inner is None: + raise AttributeError( + f"OnDemandModule has no materialized inner module; cannot access '{name}'" + ) + return getattr(self._inner, name) + + def eval(self): + return self.train(False) + + def train(self, mode: bool = True): + if mode: + raise RuntimeError("OnDemandModule is inference-only; train mode is not allowed") + return self + + def requires_grad_(self, requires_grad: bool = True): + if requires_grad: + raise RuntimeError("OnDemandModule is frozen; requires_grad_(True) is not allowed") + return self + + # No-op: the materialized module is placed on the accelerator by the caller and + # freed by discard(). A blanket model.to(temp_device) must never move or reload it. + def to(self, *args, **kwargs): + return self + + def materialize(self): + if self._inner is None: + self._inner = self._loader() + self._inner.eval() + self._inner.requires_grad_(False) + + def discard(self): + self._inner = None + torch_gc() diff --git a/modules/util/checkpointing_util.py b/modules/util/checkpointing_util.py index 4e6ee5529..9e602ac00 100644 --- a/modules/util/checkpointing_util.py +++ b/modules/util/checkpointing_util.py @@ -3,7 +3,7 @@ from typing import Any from modules.util.compile_util import init_compile -from modules.util.config.TrainConfig import TrainConfig +from modules.util.config.TrainConfig import TrainConfig, TrainModelPartConfig from modules.util.LayerOffloadConductor import LayerOffloadConductor from modules.util.torch_util import add_dummy_grad_fn_, has_grad_fn @@ -19,6 +19,7 @@ ) from transformers.models.clip.modeling_clip import CLIPEncoderLayer from transformers.models.gemma2.modeling_gemma2 import Gemma2DecoderLayer +from transformers.models.gpt_oss.modeling_gpt_oss import GptOssDecoderLayer from transformers.models.llama.modeling_llama import LlamaDecoderLayer from transformers.models.mistral.modeling_mistral import MistralDecoderLayer from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLDecoderLayer @@ -69,21 +70,25 @@ def __init__(self, *args, **kwargs): class CheckpointLayer(BaseCheckpointLayer): - def __init__(self, orig_module: nn.Module, orig_forward, train_device: torch.device): + def __init__(self, orig_module: nn.Module, orig_forward, train_device: torch.device, checkpointing: bool = True): super().__init__() assert (orig_module is None or orig_forward is None) and not (orig_module is None and orig_forward is None) self.checkpoint = orig_module self.orig_forward = orig_forward + self.checkpointing = checkpointing # dummy tensor that requires grad is needed for checkpointing to work when training a LoRA self.dummy = torch.zeros((1,), device=train_device, requires_grad=True) - def __checkpointing_forward(self, dummy: torch.Tensor, *args, **kwargs): + def __orig(self, *args, **kwargs): return self.orig_forward(*args, **kwargs) if self.checkpoint is None else self.checkpoint(*args, **kwargs) + def __checkpointing_forward(self, dummy: torch.Tensor, *args, **kwargs): + return self.__orig(*args, **kwargs) + def forward(self, *args, **kwargs): - if torch.is_grad_enabled(): + if self.checkpointing and torch.is_grad_enabled(): return torch.utils.checkpoint.checkpoint( self.__checkpointing_forward, self.dummy, @@ -92,7 +97,7 @@ def forward(self, *args, **kwargs): use_reentrant=False ) else: - return self.orig_forward(*args, **kwargs) if self.checkpoint is None else self.checkpoint(*args, **kwargs) + return self.__orig(*args, **kwargs) class OffloadCheckpointLayer(BaseCheckpointLayer): def __init__(self, orig_module: nn.Module, orig_forward, train_device: torch.device, conductor: LayerOffloadConductor, layer_index: int): @@ -149,6 +154,7 @@ def create_checkpoint( train_device: torch.device, include_from_offload_param_names: list[str] = None, conductor: LayerOffloadConductor | None = None, + checkpointing: bool = True, layer_index: int = 0, compile: bool = False, ) -> Callable: @@ -160,6 +166,9 @@ def create_checkpoint( conductor.add_layer(orig_module, included_offload_param_indices) if conductor is not None and conductor.offload_activated(): + # offloading is structurally coupled to use_reentrant=True checkpointing during the back pass + # (the recompute is what fires before_layer/after_layer in the backward direction), so the offload + # layer always checkpoints when grad is enabled, regardless of the part's gradient_checkpointing flag. if compile: layer = OffloadCheckpointLayer(orig_module=orig_module, orig_forward=None, train_device=train_device, conductor=conductor, layer_index=layer_index) #don't compile the checkpointing layer - offloading cannot be compiled: @@ -172,12 +181,12 @@ def create_checkpoint( return orig_module else: if compile: - layer = CheckpointLayer(orig_module=orig_module, orig_forward=None, train_device=train_device) + layer = CheckpointLayer(orig_module=orig_module, orig_forward=None, train_device=train_device, checkpointing=checkpointing) #do compile the checkpointing layer - slightly faster layer.compile(fullgraph=True) return layer else: - layer = CheckpointLayer(orig_module=None, orig_forward=orig_module.forward, train_device=train_device) + layer = CheckpointLayer(orig_module=None, orig_forward=orig_module.forward, train_device=train_device, checkpointing=checkpointing) orig_module.forward = layer.forward return orig_module @@ -185,6 +194,7 @@ def _create_checkpoints_for_module_list( module_list: nn.ModuleList, include_from_offload_param_names: list[str], conductor: LayerOffloadConductor, + checkpointing: bool, train_device: torch.device, layer_index: int, compile: bool, @@ -196,7 +206,7 @@ def _create_checkpoints_for_module_list( module_list[i] = create_checkpoint( layer, train_device, include_from_offload_param_names, - conductor, layer_index, compile=compile, + conductor, checkpointing, layer_index, compile=compile, ) layer_index += 1 return layer_index @@ -209,11 +219,18 @@ def _remove_checkpoint_keys(module, state_dict, prefix, local_metadata): def enable_checkpointing( model: nn.Module, config: TrainConfig, + part: TrainModelPartConfig, compile: bool, lists, # if there are multiple entries in this list, they must be in the exact order they are executed - otherwise offloading fails offload_enabled: bool = True, -) -> LayerOffloadConductor: - conductor = LayerOffloadConductor(model, config) +) -> LayerOffloadConductor | None: + if not part.checkpointing_or_offloading_enabled() and not compile: + return None + + # a conductor exists iff this part actually offloads (and the component supports conductor offloading) + offload = offload_enabled and part.offloading_enabled() + conductor = LayerOffloadConductor(model, config, part) if offload else None + checkpointing = part.checkpointing_enabled() layer_index = 0 for type_or_list, param_names in lists: @@ -224,7 +241,8 @@ def enable_checkpointing( layer_index = _create_checkpoints_for_module_list( module_list, param_names, - conductor if offload_enabled else None, + conductor, + checkpointing, torch.device(config.train_device), layer_index, compile = compile, @@ -238,7 +256,8 @@ def enable_checkpointing( layer_index = _create_checkpoints_for_module_list( module_list, param_names, - conductor if offload_enabled else None, + conductor, + checkpointing, torch.device(config.train_device), layer_index, compile = compile, @@ -249,9 +268,10 @@ def enable_checkpointing( def enable_checkpointing_for_basic_transformer_blocks( model: nn.Module, config: TrainConfig, + part: TrainModelPartConfig, offload_enabled: bool, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (BasicTransformerBlock , []), ], offload_enabled = offload_enabled, @@ -260,16 +280,18 @@ def enable_checkpointing_for_basic_transformer_blocks( def enable_checkpointing_for_clip_encoder_layers( model: nn.Module, config: TrainConfig, + part: TrainModelPartConfig, ): - return enable_checkpointing(model, config, False, [ + return enable_checkpointing(model, config, part, False, [ (CLIPEncoderLayer, []), # No activation offloading for text encoders, because the output might be taken from the middle of the network ]) def enable_checkpointing_for_t5_encoder_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ (T5Block, []), ]) @@ -277,8 +299,9 @@ def enable_checkpointing_for_t5_encoder_layers( def enable_checkpointing_for_gemma_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ (Gemma2DecoderLayer, []), ]) @@ -286,17 +309,19 @@ def enable_checkpointing_for_gemma_layers( def enable_checkpointing_for_llama_encoder_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ (LlamaDecoderLayer, []), ]) def enable_checkpointing_for_mistral_encoder_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ - (MistralDecoderLayer, []), + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ + (MistralDecoderLayer, []), # no activation offloading: this encoder is never trained ]) @@ -304,32 +329,36 @@ def enable_checkpointing_for_mistral_encoder_layers( def enable_checkpointing_for_qwen25vl_encoder_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ (Qwen2_5_VLDecoderLayer, []), # TODO No activation offloading for other encoders, see above. But clip skip is not implemented for QwenVL. Then do activation offloading? ]) def enable_checkpointing_for_qwen3_encoder_layers( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, False, [ - (Qwen3DecoderLayer, []), # No activation offloading, because hidden states are taken from the middle of the network by Flux2 + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ + (Qwen3DecoderLayer, []), # no activation offloading: this encoder is never trained ]) def enable_checkpointing_for_stable_diffusion_3_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (JointTransformerBlock, ["hidden_states", "encoder_hidden_states"]), ]) def enable_checkpointing_for_flux_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.transformer_blocks, ["hidden_states", "encoder_hidden_states"]), (model.single_transformer_blocks, ["hidden_states" ]), ]) @@ -337,8 +366,9 @@ def enable_checkpointing_for_flux_transformer( def enable_checkpointing_for_flux2_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.transformer_blocks, ["hidden_states", "encoder_hidden_states"]), (model.single_transformer_blocks, ["hidden_states" ]), ]) @@ -347,8 +377,9 @@ def enable_checkpointing_for_flux2_transformer( def enable_checkpointing_for_chroma_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.transformer_blocks, ["hidden_states", "encoder_hidden_states"]), (model.single_transformer_blocks, ["hidden_states" ]), ]) @@ -357,16 +388,18 @@ def enable_checkpointing_for_chroma_transformer( def enable_checkpointing_for_qwen_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.transformer_blocks, ["hidden_states", "encoder_hidden_states"]), ]) def enable_checkpointing_for_z_image_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.noise_refiner, ["x"]), (model.context_refiner, ["x"]), (model.layers, ["x"]), @@ -376,16 +409,18 @@ def enable_checkpointing_for_z_image_transformer( def enable_checkpointing_for_sana_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (SanaTransformerBlock, ["hidden_states"]), ]) def enable_checkpointing_for_hunyuan_video_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (HunyuanVideoIndividualTokenRefinerBlock, ["hidden_states" ]), (HunyuanVideoTransformerBlock, ["hidden_states", "encoder_hidden_states"]), (HunyuanVideoSingleTransformerBlock, ["hidden_states" ]), @@ -394,8 +429,9 @@ def enable_checkpointing_for_hunyuan_video_transformer( def enable_checkpointing_for_hi_dream_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.double_stream_blocks, ["hidden_states", "encoder_hidden_states"]), (model.single_stream_blocks, ["hidden_states" ]), ]) @@ -403,7 +439,28 @@ def enable_checkpointing_for_hi_dream_transformer( def enable_checkpointing_for_ernie_transformer( model: nn.Module, config: TrainConfig, -) -> LayerOffloadConductor: - return enable_checkpointing(model, config, config.compile, [ + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ (model.layers, ["x"]), ]) + + +def enable_checkpointing_for_lens_transformer( + model: nn.Module, + config: TrainConfig, + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ + (model.transformer_blocks, ["hidden_states", "encoder_hidden_states"]), + ]) + + +def enable_checkpointing_for_gpt_oss_encoder_layers( + model: nn.Module, + config: TrainConfig, + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, False, [ + (GptOssDecoderLayer, []), # No activation offloading: hidden states are taken from intermediate layers by encode_layers() + ]) diff --git a/modules/util/config/SampleConfig.py b/modules/util/config/SampleConfig.py index 38e99e182..9e079f4e7 100644 --- a/modules/util/config/SampleConfig.py +++ b/modules/util/config/SampleConfig.py @@ -79,6 +79,13 @@ def _get_model_defaults(model_type) -> dict: "diffusion_steps": 30, "cfg_scale": 4.0, }) + elif model_type.is_lens(): + defaults.update({ + "width": 1024, + "height": 1024, + "diffusion_steps": 30, + "cfg_scale": 5.0, + }) elif model_type.is_qwen(): defaults.update({ "width": 1024, diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index f52988502..d9e801a6f 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -13,7 +13,6 @@ from modules.util.enum.ConfigPart import ConfigPart from modules.util.enum.DataType import DataType from modules.util.enum.EMAMode import EMAMode -from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod from modules.util.enum.GradientReducePrecision import GradientReducePrecision from modules.util.enum.ImageFormat import ImageFormat from modules.util.enum.LearningRateScaler import LearningRateScaler @@ -267,10 +266,28 @@ class TrainModelPartConfig(BaseConfig): train_embedding: bool attention_mask: bool guidance_scale: float + gradient_checkpointing: bool + offload_fraction: float + activation_offloading: bool + load_on_demand: bool def __init__(self, data: list[(str, Any, type, bool)]): super().__init__(data) + def offloading_enabled(self) -> bool: + # a conductor should exist iff this is True. Layer offloading applies even to frozen parts (to fit + # them in VRAM), but activation offloading only does work during a backward pass, so it only applies + # when the part is trained -- even if activation_offloading is True in the config. + return self.offload_fraction > 0 or (self.activation_offloading and self.train) + + def checkpointing_enabled(self) -> bool: + # the inner torch checkpoint() should run iff this is True + return self.gradient_checkpointing and self.train + + def checkpointing_or_offloading_enabled(self) -> bool: + # whether the checkpoint layer wrapper needs to be installed for this part at all + return self.checkpointing_enabled() or self.offloading_enabled() + @staticmethod def default_values(): data = [] @@ -287,6 +304,10 @@ def default_values(): data.append(("train_embedding", True, bool, False)) data.append(("attention_mask", False, bool, False)) data.append(("guidance_scale", 1.0, float, False)) + data.append(("gradient_checkpointing", True, bool, False)) + data.append(("offload_fraction", 0.0, float, False)) + data.append(("activation_offloading", True, bool, False)) + data.append(("load_on_demand", False, bool, False)) return TrainModelPartConfig(data) @@ -374,10 +395,7 @@ class TrainConfig(BaseConfig): output_dtype: DataType output_model_format: ModelFormat output_model_destination: str - gradient_checkpointing: GradientCheckpointingMethod - enable_async_offloading: bool - enable_activation_offloading: bool - layer_offload_fraction: float + async_offloading: bool force_circular_padding: bool compile: bool @@ -569,7 +587,7 @@ class TrainConfig(BaseConfig): def __init__(self, data: list[(str, Any, type, bool)]): super().__init__( data, - config_version=10, + config_version=11, config_migrations={ 0: self.__migration_0, 1: self.__migration_1, @@ -581,6 +599,7 @@ def __init__(self, data: list[(str, Any, type, bool)]): 7: self.__migration_7, 8: self.__migration_8, 9: self.__migration_9, + 10: self.__migration_10, } ) @@ -727,12 +746,14 @@ def __migration_3(self, data: dict) -> dict: def __migration_4(self, data: dict) -> dict: migrated_data = data.copy() + # Translate the old bool form of gradient_checkpointing into the v5..v10 + # string/enum form. __migration_10 later fans this out per-component. gradient_checkpointing = migrated_data.pop("gradient_checkpointing", True) if gradient_checkpointing: - migrated_data["gradient_checkpointing"] = GradientCheckpointingMethod.ON + migrated_data["gradient_checkpointing"] = "ON" else: - migrated_data["gradient_checkpointing"] = GradientCheckpointingMethod.OFF + migrated_data["gradient_checkpointing"] = "OFF" return migrated_data @@ -800,6 +821,42 @@ def replace_dtype(part: str): return migrated_data + def __migration_10(self, data: dict) -> dict: + migrated_data = data.copy() + + # Fan the four old global offload/checkpointing settings out per-component. + # After __migration_4 gradient_checkpointing is a string "OFF"/"ON"/"CPU_OFFLOADED". + gc = migrated_data.pop("gradient_checkpointing", "ON") + act = migrated_data.pop("enable_activation_offloading", True) + frac = migrated_data.pop("layer_offload_fraction", 0.0) + migrated_data["async_offloading"] = migrated_data.pop("enable_async_offloading", True) + + def fan_out(part: str): + if part in migrated_data: + migrated_data[part]["gradient_checkpointing"] = gc != "OFF" + migrated_data[part]["activation_offloading"] = (gc == "CPU_OFFLOADED") and act + migrated_data[part]["offload_fraction"] = frac if gc == "CPU_OFFLOADED" else 0.0 + + fan_out("unet") + fan_out("prior") + fan_out("transformer") + fan_out("text_encoder") + fan_out("text_encoder_2") + fan_out("text_encoder_3") + fan_out("text_encoder_4") + fan_out("vae") + fan_out("effnet_encoder") + fan_out("decoder") + fan_out("decoder_text_encoder") + fan_out("decoder_vqgan") + + return migrated_data + + def model_part_configs(self) -> list[TrainModelPartConfig]: + # the per-part configs for the components this model_type actually has. Avoids "phantom" parts whose + # fields keep their defaults (train=True) or migrated offload values but don't exist in the model. + return [getattr(self, name) for name in self.model_type.model_parts()] + def weight_dtypes(self) -> ModelWeightDtypes: return ModelWeightDtypes( self.train_dtype, @@ -838,6 +895,10 @@ def model_names(self) -> ModelNames: include_text_encoder_2=self.text_encoder_2.include, include_text_encoder_3=self.text_encoder_3.include, include_text_encoder_4=self.text_encoder_4.include, + text_encoder_on_demand=self.text_encoder_on_demand(), + text_encoder_2_on_demand=self.text_encoder_2_on_demand(), + text_encoder_3_on_demand=self.text_encoder_3_on_demand(), + text_encoder_4_on_demand=self.text_encoder_4_on_demand(), ) def train_any_embedding(self) -> bool: @@ -872,6 +933,30 @@ def train_text_encoder_4_or_embedding(self) -> bool: or ((self.text_encoder_4.train_embedding or not self.model_type.has_multiple_text_encoders()) and self.train_any_embedding()) + #an encoder is loaded on demand only when it is requested, frozen (not trained, no embedding + #training) and its conditioning is cached -- otherwise it is needed resident every step. + def text_encoder_on_demand(self) -> bool: + if self.model_type.is_lens(): + return True + return self.text_encoder.load_on_demand \ + and not self.train_text_encoder_or_embedding() \ + and self.latent_caching + + def text_encoder_2_on_demand(self) -> bool: + return self.text_encoder_2.load_on_demand \ + and not self.train_text_encoder_2_or_embedding() \ + and self.latent_caching + + def text_encoder_3_on_demand(self) -> bool: + return self.text_encoder_3.load_on_demand \ + and not self.train_text_encoder_3_or_embedding() \ + and self.latent_caching + + def text_encoder_4_on_demand(self) -> bool: + return self.text_encoder_4.load_on_demand \ + and not self.train_text_encoder_4_or_embedding() \ + and self.latent_caching + def all_embedding_configs(self): if self.training_method == TrainingMethod.EMBEDDING: return self.additional_embeddings + [self.embedding] @@ -969,10 +1054,7 @@ def default_values() -> 'TrainConfig': data.append(("output_dtype", DataType.FLOAT_32, DataType, False)) data.append(("output_model_format", ModelFormat.SAFETENSORS, ModelFormat, False)) data.append(("output_model_destination", "models/model.safetensors", str, False)) - data.append(("gradient_checkpointing", GradientCheckpointingMethod.ON, GradientCheckpointingMethod, False)) - data.append(("enable_async_offloading", True, bool, False)) - data.append(("enable_activation_offloading", True, bool, False)) - data.append(("layer_offload_fraction", 0.0, float, False)) + data.append(("async_offloading", True, bool, False)) data.append(("force_circular_padding", False, bool, False)) data.append(("compile", False, bool, False)) diff --git a/modules/util/create.py b/modules/util/create.py index 7c0194da8..0c0adba16 100644 --- a/modules/util/create.py +++ b/modules/util/create.py @@ -110,7 +110,11 @@ def create_data_loader( train_progress: TrainProgress | None = None, is_validation: bool = False ) -> BaseDataLoader | None: - if config.gradient_checkpointing.offload() and config.layer_offload_fraction > 0 and config.dataloader_threads > 1: + # Layer offloading uses a non-thread-safe conductor. This check is too broad: it trips whenever any model + # part does layer offloading, even though only a component that is actually cached really runs in the + # dataloader worker threads. + # TODO: narrow this to the cached components only. + if config.dataloader_threads > 1 and any(part.offload_fraction > 0 for part in config.model_part_configs()): raise RuntimeError('layer offloading can not be activated if "dataloader_threads" > 1') if train_progress is None: @@ -133,7 +137,8 @@ def create_optimizer( if optimizer_config.optimizer is None: return None - if config.gradient_checkpointing.offload() and config.layer_offload_fraction > 0: + # a trained, layer-offloaded part has its params evicted during the back pass, so it needs fused_back_pass + if any(part.offload_fraction > 0 and part.train for part in config.model_part_configs()): if (not optimizer_config.optimizer.supports_fused_back_pass() or not optimizer_config.fused_back_pass) \ and config.training_method == TrainingMethod.FINE_TUNE: raise RuntimeError('layer offloading can only be used for fine tuning when using an optimizer that supports "fused_back_pass"') diff --git a/modules/util/enum/GradientCheckpointingMethod.py b/modules/util/enum/GradientCheckpointingMethod.py deleted file mode 100644 index d3f05666a..000000000 --- a/modules/util/enum/GradientCheckpointingMethod.py +++ /dev/null @@ -1,17 +0,0 @@ -from enum import Enum - - -class GradientCheckpointingMethod(Enum): - OFF = 'OFF' - ON = 'ON' - CPU_OFFLOADED = 'CPU_OFFLOADED' - - def __str__(self): - return self.value - - def enabled(self): - return self == GradientCheckpointingMethod.ON \ - or self == GradientCheckpointingMethod.CPU_OFFLOADED - - def offload(self): - return self == GradientCheckpointingMethod.CPU_OFFLOADED diff --git a/modules/util/enum/ModelType.py b/modules/util/enum/ModelType.py index a3ad940ec..8f1e61bde 100644 --- a/modules/util/enum/ModelType.py +++ b/modules/util/enum/ModelType.py @@ -1,5 +1,7 @@ from enum import Enum +from modules.util.enum.TrainingMethod import TrainingMethod + class ModelType(Enum): STABLE_DIFFUSION_15 = 'STABLE_DIFFUSION_15' @@ -41,6 +43,8 @@ class ModelType(Enum): ERNIE = 'ERNIE' + LENS = 'LENS' + def __str__(self): return self.value @@ -112,6 +116,9 @@ def is_z_image(self): def is_ernie(self): return self == ModelType.ERNIE + def is_lens(self): + return self == ModelType.LENS + def has_mask_input(self) -> bool: return self == ModelType.STABLE_DIFFUSION_15_INPAINTING \ or self == ModelType.STABLE_DIFFUSION_20_INPAINTING \ @@ -161,11 +168,66 @@ def is_flow_matching(self) -> bool: or self.is_hunyuan_video() \ or self.is_hi_dream() \ or self.is_z_image() \ - or self.is_ernie() + or self.is_ernie() \ + or self.is_lens() def is_video_model(self) -> bool: return self.is_hunyuan_video() #incase we add more video models in the future + def model_parts(self) -> tuple[str, ...]: + return _MODEL_PARTS[self] + + def supported_training_methods(self) -> tuple[TrainingMethod, ...]: + if self.is_stable_diffusion(): + return (TrainingMethod.FINE_TUNE, TrainingMethod.LORA, TrainingMethod.EMBEDDING, TrainingMethod.FINE_TUNE_VAE) + if self.is_stable_diffusion_3() \ + or self.is_stable_diffusion_xl() \ + or self.is_wuerstchen() \ + or self.is_pixart() \ + or self.is_flux_1() \ + or self.is_sana() \ + or self.is_hunyuan_video() \ + or self.is_hi_dream() \ + or self.is_chroma(): + return (TrainingMethod.FINE_TUNE, TrainingMethod.LORA, TrainingMethod.EMBEDDING) + if self.is_qwen() or self.is_z_image() or self.is_flux_2() or self.is_ernie() or self.is_lens(): + return (TrainingMethod.FINE_TUNE, TrainingMethod.LORA) + raise ValueError(f"No supported training methods defined for model type {self}") + + +# The first text encoder is always "text_encoder" here (matching the config field), even for +# multi-encoder models that refer to it as "text_encoder_1" elsewhere in the code. +_MODEL_PARTS: dict[ModelType, tuple[str, ...]] = { + ModelType.STABLE_DIFFUSION_15: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_15_INPAINTING: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_20: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_20_BASE: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_20_INPAINTING: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_20_DEPTH: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_21: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_21_BASE: ("text_encoder", "unet", "vae"), + ModelType.STABLE_DIFFUSION_3: ("text_encoder", "text_encoder_2", "text_encoder_3", "transformer", "vae"), + ModelType.STABLE_DIFFUSION_35: ("text_encoder", "text_encoder_2", "text_encoder_3", "transformer", "vae"), + ModelType.STABLE_DIFFUSION_XL_10_BASE: ("text_encoder", "text_encoder_2", "unet", "vae"), + ModelType.STABLE_DIFFUSION_XL_10_BASE_INPAINTING: ("text_encoder", "text_encoder_2", "unet", "vae"), + # Only Würstchen v2's decoder has its own text encoder; Stable Cascade's decoder does not. + ModelType.WUERSTCHEN_2: ("text_encoder", "prior", "effnet_encoder", "decoder", "decoder_text_encoder", "decoder_vqgan"), + ModelType.STABLE_CASCADE_1: ("text_encoder", "prior", "effnet_encoder", "decoder", "decoder_vqgan"), + ModelType.PIXART_ALPHA: ("text_encoder", "transformer", "vae"), + ModelType.PIXART_SIGMA: ("text_encoder", "transformer", "vae"), + ModelType.FLUX_DEV_1: ("text_encoder", "text_encoder_2", "transformer", "vae"), + ModelType.FLUX_FILL_DEV_1: ("text_encoder", "text_encoder_2", "transformer", "vae"), + ModelType.FLUX_2: ("text_encoder", "transformer", "vae"), + ModelType.SANA: ("text_encoder", "transformer", "vae"), + ModelType.HUNYUAN_VIDEO: ("text_encoder", "text_encoder_2", "transformer", "vae"), + ModelType.HI_DREAM_FULL: ("text_encoder", "text_encoder_2", "text_encoder_3", "text_encoder_4", "transformer", "vae"), + ModelType.CHROMA_1: ("text_encoder", "transformer", "vae"), + ModelType.QWEN: ("text_encoder", "transformer", "vae"), + ModelType.Z_IMAGE: ("text_encoder", "transformer", "vae"), + ModelType.ERNIE: ("text_encoder", "transformer", "vae"), + ModelType.LENS: ("text_encoder", "transformer", "vae"), +} + class PeftType(Enum): LORA = 'LORA' diff --git a/modules/util/optimizer/muon_util.py b/modules/util/optimizer/muon_util.py index b2630c16b..cd39c06bb 100644 --- a/modules/util/optimizer/muon_util.py +++ b/modules/util/optimizer/muon_util.py @@ -35,7 +35,7 @@ def build_muon_adam_key_fn( 'block', # UNet 'text_model.encoder.layers', # TEs (CLIPs) ] - case ModelType.STABLE_DIFFUSION_3 | ModelType.STABLE_DIFFUSION_35 | ModelType.SANA | ModelType.FLUX_DEV_1 | ModelType.FLUX_2 | ModelType.CHROMA_1 | ModelType.QWEN | ModelType.PIXART_ALPHA | ModelType.PIXART_SIGMA: + case ModelType.STABLE_DIFFUSION_3 | ModelType.STABLE_DIFFUSION_35 | ModelType.SANA | ModelType.FLUX_DEV_1 | ModelType.FLUX_2 | ModelType.CHROMA_1 | ModelType.QWEN | ModelType.PIXART_ALPHA | ModelType.PIXART_SIGMA | ModelType.LENS: default_patterns = [ 'transformer_blocks', 'encoder.block', # TE (T5) diff --git a/requirements-global.txt b/requirements-global.txt index ba3d9f476..a3a8ee857 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -28,11 +28,14 @@ omegaconf==2.3.0 # needed to load stable diffusion from single ckpt files invisible-watermark==0.2.0 # needed for the SDXL pipeline # other models +git+https://github.com/dxqb/Lens.git@daba3a8#egg=lens +kernels==0.14.1 # Triton kernel loader for Lens GPT-OSS MXFP4 text encoder +einops==0.8.2 pooch==1.8.2 open-clip-torch==2.32.0 # data loader --e git+https://github.com/Nerogar/mgds.git@9320a69#egg=mgds +-e git+https://github.com/dxqb/mgds.git@f3b1bd1#egg=mgds # optimizers dadaptation==3.2 # dadaptation optimizers diff --git a/resources/sd_model_spec/lens-lora.json b/resources/sd_model_spec/lens-lora.json new file mode 100644 index 000000000..31fede12a --- /dev/null +++ b/resources/sd_model_spec/lens-lora.json @@ -0,0 +1,6 @@ +{ + "modelspec.sai_model_spec": "1.0.0", + "modelspec.architecture": "Lens/lora", + "modelspec.implementation": "https://github.com/microsoft/Lens", + "modelspec.title": "Lens LoRA" +} diff --git a/resources/sd_model_spec/lens.json b/resources/sd_model_spec/lens.json new file mode 100644 index 000000000..57c2681aa --- /dev/null +++ b/resources/sd_model_spec/lens.json @@ -0,0 +1,6 @@ +{ + "modelspec.sai_model_spec": "1.0.0", + "modelspec.architecture": "Lens", + "modelspec.implementation": "https://github.com/microsoft/Lens", + "modelspec.title": "Lens" +} diff --git a/training_presets/#chroma Finetune 16GB.json b/training_presets/#chroma Finetune 16GB.json index 2dacbee20..145ab5b04 100644 --- a/training_presets/#chroma Finetune 16GB.json +++ b/training_presets/#chroma Finetune 16GB.json @@ -4,16 +4,15 @@ "learning_rate": 1e-5, "model_type": "CHROMA_1", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.4, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.4 }, "text_encoder": { "train": false, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "FLOAT_8" }, "training_method": "FINE_TUNE", "vae": { diff --git a/training_presets/#chroma Finetune 8GB.json b/training_presets/#chroma Finetune 8GB.json index 508410995..29b36c84b 100644 --- a/training_presets/#chroma Finetune 8GB.json +++ b/training_presets/#chroma Finetune 8GB.json @@ -4,16 +4,15 @@ "learning_rate": 1e-5, "model_type": "CHROMA_1", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.85, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.85 }, "text_encoder": { "train": false, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "FLOAT_8" }, "training_method": "FINE_TUNE", "vae": { diff --git a/training_presets/#chroma LoRA 8GB.json b/training_presets/#chroma LoRA 8GB.json index 78027aac4..437ac71d7 100644 --- a/training_presets/#chroma LoRA 8GB.json +++ b/training_presets/#chroma LoRA 8GB.json @@ -4,16 +4,15 @@ "learning_rate": 0.0003, "model_type": "CHROMA_1", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.6, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.6 }, "text_encoder": { "train": false, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "FLOAT_8" }, "training_method": "LORA", "vae": { diff --git a/training_presets/#ernie LoRA 8GB.json b/training_presets/#ernie LoRA 8GB.json index bc1e4a82b..a63d34ece 100644 --- a/training_presets/#ernie LoRA 8GB.json +++ b/training_presets/#ernie LoRA 8GB.json @@ -7,7 +7,8 @@ "compile": true, "transformer": { "train": true, - "weight_dtype": "INT_W8A8" + "weight_dtype": "INT_W8A8", + "offload_fraction": 0.7 }, "text_encoder": { "train": false, @@ -26,7 +27,5 @@ "layer_filter_preset": "blocks" }, "timestep_distribution": "LOGIT_NORMAL", - "dataloader_threads": 1, - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.7 + "dataloader_threads": 1 } diff --git a/training_presets/#flux2 Finetune 16GB.json b/training_presets/#flux2 Finetune 16GB.json index ac07a501e..00094da12 100644 --- a/training_presets/#flux2 Finetune 16GB.json +++ b/training_presets/#flux2 Finetune 16GB.json @@ -7,11 +7,12 @@ "compile": true, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.6 }, "text_encoder": { "train": false, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "FLOAT_8" }, "training_method": "FINE_TUNE", "vae": { @@ -26,8 +27,6 @@ "timestep_distribution": "LOGIT_NORMAL", "dynamic_timestep_shifting": true, "dataloader_threads": 1, - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.6, "optimizer": { "optimizer": "ADAFACTOR" }, diff --git a/training_presets/#flux2 LoRA 8GB.json b/training_presets/#flux2 LoRA 8GB.json index 0bd10d116..160a999ef 100644 --- a/training_presets/#flux2 LoRA 8GB.json +++ b/training_presets/#flux2 LoRA 8GB.json @@ -7,11 +7,13 @@ "compile": true, "transformer": { "train": true, - "weight_dtype": "INT_W8A8" + "weight_dtype": "INT_W8A8", + "offload_fraction": 0.7 }, "text_encoder": { "train": false, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.7 }, "training_method": "LORA", "vae": { @@ -27,7 +29,5 @@ }, "timestep_distribution": "LOGIT_NORMAL", "dynamic_timestep_shifting": true, - "dataloader_threads": 1, - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.7 + "dataloader_threads": 1 } diff --git a/training_presets/#hidream LoRA.json b/training_presets/#hidream LoRA.json index 2eb588763..8b025a628 100644 --- a/training_presets/#hidream LoRA.json +++ b/training_presets/#hidream LoRA.json @@ -2,8 +2,6 @@ "backup_after": 10, "base_model_name": "HiDream-ai/HiDream-I1-Full", "batch_size": 4, - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.5, "dataloader_threads": 1, "learning_rate": 0.0003, "model_type": "HI_DREAM_FULL", @@ -16,7 +14,8 @@ "training_method": "LORA", "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 }, "text_encoder": { "train": false, @@ -28,11 +27,13 @@ }, "text_encoder_3": { "train": false, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 }, "text_encoder_4": { "model_name": "meta-llama/Llama-3.1-8B-Instruct", "train": false, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 } } diff --git a/training_presets/#hunyuan video LoRA.json b/training_presets/#hunyuan video LoRA.json index 754550155..96bbb80e7 100644 --- a/training_presets/#hunyuan video LoRA.json +++ b/training_presets/#hunyuan video LoRA.json @@ -2,8 +2,6 @@ "backup_after": 10, "base_model_name": "hunyuanvideo-community/HunyuanVideo", "batch_size": 4, - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.5, "dataloader_threads": 1, "learning_rate": 0.0003, "model_type": "HUNYUAN_VIDEO", @@ -16,11 +14,13 @@ "training_method": "LORA", "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 }, "text_encoder": { "train": false, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 }, "text_encoder_2": { "train": false, diff --git a/training_presets/#lens Finetune 16GB.json b/training_presets/#lens Finetune 16GB.json new file mode 100644 index 000000000..d1dd235ab --- /dev/null +++ b/training_presets/#lens Finetune 16GB.json @@ -0,0 +1,47 @@ +{ + "base_model_name": "microsoft/Lens", + "batch_size": 2, + "learning_rate": 1e-6, + "model_type": "LENS", + "resolution": "512", + "compile": true, + "transformer": { + "train": true, + "weight_dtype": "BFLOAT_16" + }, + "text_encoder": { + "train": false, + "weight_dtype": "BFLOAT_16" + }, + "training_method": "FINE_TUNE", + "vae": { + "weight_dtype": "FLOAT_32" + }, + "train_dtype": "BFLOAT_16", + "output_dtype": "BFLOAT_16", + "quantization": { + "layer_filter": "transformer_block", + "layer_filter_preset": "blocks" + }, + "timestep_distribution": "LOGIT_NORMAL", + "dataloader_threads": 1, + "optimizer": { + "optimizer": "ADAFACTOR" + }, + "optimizer_defaults": { + "ADAFACTOR": { + "optimizer": "ADAFACTOR", + "fused_back_pass": true, + "beta1": null, + "clip_threshold": 1.0, + "decay_rate": -0.8, + "eps": 1e-30, + "eps2": 0.001, + "relative_step": false, + "scale_parameter": false, + "stochastic_rounding": true, + "warmup_init": false, + "weight_decay": 0.0 + } + } +} diff --git a/training_presets/#lens LoRA 16GB.json b/training_presets/#lens LoRA 16GB.json new file mode 100644 index 000000000..5400f22df --- /dev/null +++ b/training_presets/#lens LoRA 16GB.json @@ -0,0 +1,30 @@ +{ + "base_model_name": "microsoft/Lens", + "batch_size": 2, + "learning_rate": 3e-4, + "model_type": "LENS", + "resolution": "512", + "compile": true, + "transformer": { + "train": true, + "weight_dtype": "INT_W8A8" + }, + "text_encoder": { + "train": false, + "weight_dtype": "BFLOAT_16" + }, + "training_method": "LORA", + "vae": { + "weight_dtype": "FLOAT_32" + }, + "train_dtype": "BFLOAT_16", + "output_dtype": "BFLOAT_16", + "layer_filter": "attn,mlp", + "layer_filter_preset": "attn-mlp", + "quantization": { + "layer_filter": "attn,mlp", + "layer_filter_preset": "attn-mlp" + }, + "timestep_distribution": "LOGIT_NORMAL", + "dataloader_threads": 1 +} diff --git a/training_presets/#qwen Finetune 16GB.json b/training_presets/#qwen Finetune 16GB.json index 811d7e0b1..2224993f4 100644 --- a/training_presets/#qwen Finetune 16GB.json +++ b/training_presets/#qwen Finetune 16GB.json @@ -4,12 +4,11 @@ "learning_rate": 1e-5, "model_type": "QWEN", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.75, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.75 }, "text_encoder": { "train": false, diff --git a/training_presets/#qwen Finetune 24GB.json b/training_presets/#qwen Finetune 24GB.json index 8bee3cd3f..1cf9dc09f 100644 --- a/training_presets/#qwen Finetune 24GB.json +++ b/training_presets/#qwen Finetune 24GB.json @@ -4,12 +4,11 @@ "learning_rate": 1e-5, "model_type": "QWEN", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.55, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.55 }, "text_encoder": { "train": false, diff --git a/training_presets/#qwen LoRA 16GB.json b/training_presets/#qwen LoRA 16GB.json index b4e0d7e88..0eda34d5d 100644 --- a/training_presets/#qwen LoRA 16GB.json +++ b/training_presets/#qwen LoRA 16GB.json @@ -4,12 +4,11 @@ "learning_rate": 0.0003, "model_type": "QWEN", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.5, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.5 }, "text_encoder": { "train": false, diff --git a/training_presets/#qwen LoRA 24GB.json b/training_presets/#qwen LoRA 24GB.json index 696648a42..cd7b7216e 100644 --- a/training_presets/#qwen LoRA 24GB.json +++ b/training_presets/#qwen LoRA 24GB.json @@ -4,12 +4,11 @@ "learning_rate": 0.0003, "model_type": "QWEN", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.1, "dataloader_threads": 1, "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.1 }, "text_encoder": { "train": false, diff --git a/training_presets/#z-image DeTurbo LoRA 8GB.json b/training_presets/#z-image DeTurbo LoRA 8GB.json index cc38e60eb..957f08b80 100644 --- a/training_presets/#z-image DeTurbo LoRA 8GB.json +++ b/training_presets/#z-image DeTurbo LoRA 8GB.json @@ -4,13 +4,12 @@ "learning_rate": 0.0003, "model_type": "Z_IMAGE", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.6, "compile": true, "transformer": { "train": true, "weight_dtype": "INT_W8A8", - "model_name": "https://huggingface.co/ostris/Z-Image-De-Turbo/blob/main/z_image_de_turbo_v1_bf16.safetensors" + "model_name": "https://huggingface.co/ostris/Z-Image-De-Turbo/blob/main/z_image_de_turbo_v1_bf16.safetensors", + "offload_fraction": 0.6 }, "text_encoder": { "train": false, diff --git a/training_presets/#z-image Finetune 16GB.json b/training_presets/#z-image Finetune 16GB.json index 0d23d3992..3911c866c 100644 --- a/training_presets/#z-image Finetune 16GB.json +++ b/training_presets/#z-image Finetune 16GB.json @@ -4,12 +4,11 @@ "learning_rate": 1e-5, "model_type": "Z_IMAGE", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.1, "compile": true, "transformer": { "train": true, - "weight_dtype": "BFLOAT_16" + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.1 }, "text_encoder": { "train": false, diff --git a/training_presets/#z-image LoRA 8GB.json b/training_presets/#z-image LoRA 8GB.json index 78b4b05cc..8dbce7b21 100644 --- a/training_presets/#z-image LoRA 8GB.json +++ b/training_presets/#z-image LoRA 8GB.json @@ -4,12 +4,11 @@ "learning_rate": 0.0003, "model_type": "Z_IMAGE", "resolution": "512", - "gradient_checkpointing": "CPU_OFFLOADED", - "layer_offload_fraction": 0.6, "compile": true, "transformer": { "train": true, - "weight_dtype": "FLOAT_8" + "weight_dtype": "FLOAT_8", + "offload_fraction": 0.6 }, "text_encoder": { "train": false,