still doesn't work

flux_train_comfy.py

@@ -276,7 +276,7 @@ class FluxTrainer:
         flux = flux_utils.load_flow_model(name, args.pretrained_model_name_or_path, weight_dtype, "cpu")
 
         if args.gradient_checkpointing:
-            flux.enable_gradient_checkpointing(cpu_offload=args.cpu_offload_checkpointing)
+            flux.enable_gradient_checkpointing(args.cpu_offload_checkpointing)
 
         flux.requires_grad_(True)
 
@@ -680,7 +680,7 @@ class FluxTrainer:
                     else:
                         with torch.no_grad():
                             # encode images to latents. images are [-1, 1]
-                            latents = ae.encode(batch["images"].to(ae.dtype)).to(accelerator.device, dtype=weight_dtype)
+                            latents = ae.encode(batch["images"])
 
                         # NaNが含まれていれば警告を表示し0に置き換える
                         if torch.any(torch.isnan(latents)):

flux_train_network_comfy.py

@@ -4,7 +4,7 @@ import math
 from typing import Any
 import argparse
 from .library import flux_models, flux_train_utils, flux_utils, sd3_train_utils, strategy_base, strategy_flux, train_util
-from .train_network import NetworkTrainer, clean_memory_on_device, setup_parser
+from .train_network import NetworkTrainer, clean_memory_on_device
 
 from accelerate import Accelerator
 
@@ -20,89 +20,56 @@ class FluxNetworkTrainer(NetworkTrainer):
 
     def assert_extra_args(self, args, train_dataset_group):
         super().assert_extra_args(args, train_dataset_group)
-        # sdxl_train_util.verify_sdxl_training_args(args)
-
-        if args.fp8_base_unet:
-            args.fp8_base = True  # if fp8_base_unet is enabled, fp8_base is also enabled for FLUX.1
-
-        if args.cache_text_encoder_outputs_to_disk and not args.cache_text_encoder_outputs:
-            logger.warning(
-                "cache_text_encoder_outputs_to_disk is enabled, so cache_text_encoder_outputs is also enabled / cache_text_encoder_outputs_to_diskが有効になっているため、cache_text_encoder_outputsも有効になります"
-            )
-            args.cache_text_encoder_outputs = True
 
         if args.cache_text_encoder_outputs:
             assert (
                 train_dataset_group.is_text_encoder_output_cacheable()
-            ), "when caching Text Encoder output, either caption_dropout_rate, shuffle_caption, token_warmup_step or caption_tag_dropout_rate cannot be used / Text Encoderの出力をキャッシュするときはcaption_dropout_rate, shuffle_caption, token_warmup_step, caption_tag_dropout_rateは使えません"
+            ), "when caching Text Encoder output, either caption_dropout_rate, shuffle_caption, token_warmup_step or caption_tag_dropout_rate cannot be used"
 
-        # prepare CLIP-L/T5XXL training flags
-        self.train_clip_l = not args.network_train_unet_only
-        self.train_t5xxl = False  # default is False even if args.network_train_unet_only is False
+        #assert (
+        #    args.network_train_unet_only or not args.cache_text_encoder_outputs
+        #), "network for Text Encoder cannot be trained with caching Text Encoder outputs"
+        if not args.network_train_unet_only:
+            logger.info(
+                "network for CLIP-L only will be trained. T5XXL will not be trained / CLIP-Lのネットワークのみが学習されます。T5XXLは学習されません"
+            )
 
         if args.max_token_length is not None:
             logger.warning("max_token_length is not used in Flux training")
 
-        assert not args.split_mode or not args.cpu_offload_checkpointing, (
-            "split_mode and cpu_offload_checkpointing cannot be used together"
-        )
-
         train_dataset_group.verify_bucket_reso_steps(32)  # TODO check this
 
     def get_flux_model_name(self, args):
         return "schnell" if "schnell" in args.pretrained_model_name_or_path else "dev"
-
+    
     def load_target_model(self, args, weight_dtype, accelerator):
         # currently offload to cpu for some models
         name = self.get_flux_model_name(args)
-
-        # if the file is fp8 and we are using fp8_base, we can load it as is (fp8)
-        loading_dtype = None if args.fp8_base else weight_dtype
-
-        # if we load to cpu, flux.to(fp8) takes a long time, so we should load to gpu in future
-        model = flux_utils.load_flow_model(
-            name, args.pretrained_model_name_or_path, loading_dtype, "cpu", disable_mmap=args.disable_mmap_load_safetensors
-        )
-        if args.fp8_base:
-            # check dtype of model
-            if model.dtype == torch.float8_e4m3fnuz or model.dtype == torch.float8_e5m2fnuz:
-                raise ValueError(f"Unsupported fp8 model dtype: {model.dtype}")
-            elif model.dtype == torch.float8_e4m3fn or model.dtype == torch.float8_e5m2:
-                logger.info(f"Loaded {model.dtype} FLUX model")
+        # if we load to cpu, flux.to(fp8) takes a long time
+        model = flux_utils.load_flow_model(name, args.pretrained_model_name_or_path, weight_dtype, "cpu")
 
         if args.split_mode:
-            model = self.prepare_split_model(model, args, weight_dtype, accelerator)
+            model = self.prepare_split_model(model, weight_dtype, accelerator, args)
 
-        clip_l = flux_utils.load_clip_l(args.clip_l, weight_dtype, "cpu", disable_mmap=args.disable_mmap_load_safetensors)
+        clip_l = flux_utils.load_clip_l(args.clip_l, weight_dtype, "cpu")
         clip_l.eval()
 
-        # if the file is fp8 and we are using fp8_base (not unet), we can load it as is (fp8)
-        if args.fp8_base and not args.fp8_base_unet:
-            loading_dtype = None  # as is
-        else:
-            loading_dtype = weight_dtype
-
         # loading t5xxl to cpu takes a long time, so we should load to gpu in future
-        t5xxl = flux_utils.load_t5xxl(args.t5xxl, loading_dtype, "cpu", disable_mmap=args.disable_mmap_load_safetensors)
+        t5xxl = flux_utils.load_t5xxl(args.t5xxl, weight_dtype, "cpu")
         t5xxl.eval()
-        if args.fp8_base and not args.fp8_base_unet:
-            # check dtype of model
-            if t5xxl.dtype == torch.float8_e4m3fnuz or t5xxl.dtype == torch.float8_e5m2 or t5xxl.dtype == torch.float8_e5m2fnuz:
-                raise ValueError(f"Unsupported fp8 model dtype: {t5xxl.dtype}")
-            elif t5xxl.dtype == torch.float8_e4m3fn:
-                logger.info("Loaded fp8 T5XXL model")
 
-        ae = flux_utils.load_ae(name, args.ae, weight_dtype, "cpu", disable_mmap=args.disable_mmap_load_safetensors)
+        ae = flux_utils.load_ae(name, args.ae, weight_dtype, "cpu")
 
         return flux_utils.MODEL_VERSION_FLUX_V1, [clip_l, t5xxl], ae, model
 
-    def prepare_split_model(self, model, args, weight_dtype, accelerator):
+    def prepare_split_model(self, model, weight_dtype, accelerator, args):
         from accelerate import init_empty_weights
 
         logger.info("prepare split model")
         with init_empty_weights():
-            flux_upper = flux_models.FluxUpper(model.params)
             flux_lower = flux_models.FluxLower(model.params)
+            flux_upper = flux_models.FluxUpper(model.params, flux_lower)
+            
         sd = model.state_dict()
 
         # lower (trainable)
@@ -115,13 +82,7 @@ class FluxNetworkTrainer(NetworkTrainer):
         flux_upper.load_state_dict(sd, strict=False, assign=True)
 
         logger.info("prepare upper model")
-        if args.fp8_base:
-            if args.fp8_dtype and args.fp8_dtype.lower() == "e5m2":
-                target_dtype = torch.float8_e5m2
-            else:
-                target_dtype = torch.float8_e4m3fn
-        else:
-            target_dtype =weight_dtype
+        target_dtype = torch.float8_e4m3fn if args.fp8_base else weight_dtype
         flux_upper.to(accelerator.device, dtype=target_dtype)
         flux_upper.eval()
 
@@ -163,35 +124,25 @@ class FluxNetworkTrainer(NetworkTrainer):
     def get_text_encoding_strategy(self, args):
         return strategy_flux.FluxTextEncodingStrategy(apply_t5_attn_mask=args.apply_t5_attn_mask)
 
-    def post_process_network(self, args, accelerator, network, text_encoders, unet):
-        # check t5xxl is trained or not
-        self.train_t5xxl = network.train_t5xxl
-
-        if self.train_t5xxl and args.cache_text_encoder_outputs:
-            raise ValueError(
-                "T5XXL is trained, so cache_text_encoder_outputs cannot be used / T5XXL学習時はcache_text_encoder_outputsは使用できません"
-            )
-
     def get_models_for_text_encoding(self, args, accelerator, text_encoders):
         if args.cache_text_encoder_outputs:
-            if self.train_clip_l and not self.train_t5xxl:
+            if self.is_train_text_encoder(args):
                 return text_encoders[0:1]  # only CLIP-L is needed for encoding because T5XXL is cached
             else:
-                return None  # no text encoders are needed for encoding because both are cached
+                return text_encoders  # ignored
         else:
             return text_encoders  # both CLIP-L and T5XXL are needed for encoding
 
     def get_text_encoders_train_flags(self, args, text_encoders):
-        return [self.train_clip_l, self.train_t5xxl]
+        return [True, False] if self.is_train_text_encoder(args) else [False, False]
 
     def get_text_encoder_outputs_caching_strategy(self, args):
         if args.cache_text_encoder_outputs:
-            # if the text encoders is trained, we need tokenization, so is_partial is True
             return strategy_flux.FluxTextEncoderOutputsCachingStrategy(
                 args.cache_text_encoder_outputs_to_disk,
                 None,
                 False,
-                is_partial=self.train_clip_l or self.train_t5xxl,
+                is_partial=self.is_train_text_encoder(args),
                 apply_t5_attn_mask=args.apply_t5_attn_mask,
             )
         else:
@@ -212,20 +163,13 @@ class FluxNetworkTrainer(NetworkTrainer):
 
             # When TE is not be trained, it will not be prepared so we need to use explicit autocast
             logger.info("move text encoders to gpu")
-            text_encoders[0].to(accelerator.device, dtype=weight_dtype)  # always not fp8
-            text_encoders[1].to(accelerator.device)
-
-            if text_encoders[1].dtype == torch.float8_e4m3fn:
-                # if we load fp8 weights, the model is already fp8, so we use it as is
-                self.prepare_text_encoder_fp8(1, text_encoders[1], text_encoders[1].dtype, weight_dtype)
-            else:
-                # otherwise, we need to convert it to target dtype
-                text_encoders[1].to(weight_dtype)
-
+            text_encoders[0].to(accelerator.device, dtype=weight_dtype)
+            text_encoders[1].to(accelerator.device, dtype=weight_dtype)
             with accelerator.autocast():
                 dataset.new_cache_text_encoder_outputs(text_encoders, accelerator.is_main_process)
 
             # cache sample prompts
+
             if args.sample_prompts is not None:
                 logger.info(f"cache Text Encoder outputs for sample prompt: {args.sample_prompts}")
 
@@ -263,10 +207,8 @@ class FluxNetworkTrainer(NetworkTrainer):
                                     tokenize_strategy, text_encoders, tokens_and_masks, args.apply_t5_attn_mask
                                 )
                 self.sample_prompts_te_outputs = sample_prompts_te_outputs
-
             accelerator.wait_for_everyone()
 
-            # move back to cpu
             if not self.is_train_text_encoder(args):
                 logger.info("move CLIP-L back to cpu")
                 text_encoders[0].to("cpu")
@@ -281,7 +223,7 @@ class FluxNetworkTrainer(NetworkTrainer):
         else:
             # Text Encoder
             text_encoders[0].to(accelerator.device, dtype=weight_dtype)
-            text_encoders[1].to(accelerator.device)
+            text_encoders[1].to(accelerator.device, dtype=weight_dtype)
 
     def sample_images_split_mode(self, accelerator, args, epoch, global_step, flux, ae, text_encoder, sample_prompts_te_outputs, validation_settings):
        
@@ -293,12 +235,12 @@ class FluxNetworkTrainer(NetworkTrainer):
                 self.target_device = device
 
             def forward(self, img, img_ids, txt, txt_ids, timesteps, y, guidance=None, txt_attention_mask=None):
-                self.flux_lower.to("cpu")
-                clean_memory_on_device(self.target_device)
+                #self.flux_lower.to("cpu")
+                #clean_memory_on_device(self.target_device)
                 self.flux_upper.to(self.target_device)
                 img, txt, vec, pe = self.flux_upper(img, img_ids, txt, txt_ids, timesteps, y, guidance, txt_attention_mask)
-                self.flux_upper.to("cpu")
-                clean_memory_on_device(self.target_device)
+                #self.flux_upper.to("cpu")
+                #clean_memory_on_device(self.target_device)
                 self.flux_lower.to(self.target_device)
                 return self.flux_lower(img, txt, vec, pe, txt_attention_mask)
 
@@ -314,6 +256,9 @@ class FluxNetworkTrainer(NetworkTrainer):
         noise_scheduler = sd3_train_utils.FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=args.discrete_flow_shift)
         self.noise_scheduler_copy = copy.deepcopy(noise_scheduler)
         return noise_scheduler
+    
+    def is_text_encoder_not_needed_for_training(self, args):
+        return args.cache_text_encoder_outputs and not self.is_train_text_encoder(args)
 
     def encode_images_to_latents(self, args, accelerator, vae, images):
         return vae.encode(images)
@@ -334,6 +279,55 @@ class FluxNetworkTrainer(NetworkTrainer):
         weight_dtype,
         train_unet,
     ):
+        # copy from sd3_train.py and modified
+
+        def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
+            sigmas = self.noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype)
+            schedule_timesteps = self.noise_scheduler_copy.timesteps.to(accelerator.device)
+            timesteps = timesteps.to(accelerator.device)
+            step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+
+            sigma = sigmas[step_indices].flatten()
+            while len(sigma.shape) < n_dim:
+                sigma = sigma.unsqueeze(-1)
+            return sigma
+
+        def compute_density_for_timestep_sampling(
+            weighting_scheme: str, batch_size: int, logit_mean: float = None, logit_std: float = None, mode_scale: float = None
+        ):
+            """Compute the density for sampling the timesteps when doing SD3 training.
+
+            Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528.
+
+            SD3 paper reference: https://arxiv.org/abs/2403.03206v1.
+            """
+            if weighting_scheme == "logit_normal":
+                # See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$).
+                u = torch.normal(mean=logit_mean, std=logit_std, size=(batch_size,), device="cpu")
+                u = torch.nn.functional.sigmoid(u)
+            elif weighting_scheme == "mode":
+                u = torch.rand(size=(batch_size,), device="cpu")
+                u = 1 - u - mode_scale * (torch.cos(math.pi * u / 2) ** 2 - 1 + u)
+            else:
+                u = torch.rand(size=(batch_size,), device="cpu")
+            return u
+
+        def compute_loss_weighting_for_sd3(weighting_scheme: str, sigmas=None):
+            """Computes loss weighting scheme for SD3 training.
+
+            Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528.
+
+            SD3 paper reference: https://arxiv.org/abs/2403.03206v1.
+            """
+            if weighting_scheme == "sigma_sqrt":
+                weighting = (sigmas**-2.0).float()
+            elif weighting_scheme == "cosmap":
+                bot = 1 - 2 * sigmas + 2 * sigmas**2
+                weighting = 2 / (math.pi * bot)
+            else:
+                weighting = torch.ones_like(sigmas)
+            return weighting
+
         # Sample noise that we'll add to the latents
         noise = torch.randn_like(latents)
         bsz = latents.shape[0]
@@ -349,8 +343,7 @@ class FluxNetworkTrainer(NetworkTrainer):
         img_ids = flux_utils.prepare_img_ids(bsz, packed_latent_height, packed_latent_width).to(device=accelerator.device)
 
         # get guidance
-        # ensure guidance_scale in args is float
-        guidance_vec = torch.full((bsz,), float(args.guidance_scale), device=accelerator.device)
+        guidance_vec = torch.full((bsz,), args.guidance_scale, device=accelerator.device)
 
         # ensure the hidden state will require grad
         if args.gradient_checkpointing:
@@ -382,16 +375,16 @@ class FluxNetworkTrainer(NetworkTrainer):
                 )
         else:
             # split forward to reduce memory usage
-            assert network.train_blocks == "single", "train_blocks must be single for split mode"
+            #assert network.train_blocks == "single", "train_blocks must be single for split mode"
             with accelerator.autocast():
                 # move flux lower to cpu, and then move flux upper to gpu
-                unet.to("cpu")
-                clean_memory_on_device(accelerator.device)
+                #unet.to("cpu")
+                #clean_memory_on_device(accelerator.device)
                 self.flux_upper.to(accelerator.device)
 
                 # upper model does not require grad
                 with torch.no_grad():
-                    intermediate_img, intermediate_txt, vec, pe = self.flux_upper(
+                    model_pred = self.flux_upper(
                         img=packed_noisy_model_input,
                         img_ids=img_ids,
                         txt=t5_out,
@@ -400,19 +393,20 @@ class FluxNetworkTrainer(NetworkTrainer):
                         timesteps=timesteps / 1000,
                         guidance=guidance_vec,
                         txt_attention_mask=t5_attn_mask,
+                        train_lower=True,
                     )
-
+                model_pred.requires_grad_(True)
                 # move flux upper back to cpu, and then move flux lower to gpu
-                self.flux_upper.to("cpu")
-                clean_memory_on_device(accelerator.device)
-                unet.to(accelerator.device)
+                #self.flux_upper.to("cpu")
+                #clean_memory_on_device(accelerator.device)
+                #unet.to(accelerator.device)
 
                 # lower model requires grad
-                intermediate_img.requires_grad_(True)
-                intermediate_txt.requires_grad_(True)
-                vec.requires_grad_(True)
-                pe.requires_grad_(True)
-                model_pred = unet(img=intermediate_img, txt=intermediate_txt, vec=vec, pe=pe, txt_attention_mask=t5_attn_mask)
+                # intermediate_img.requires_grad_(True)
+                # intermediate_txt.requires_grad_(True)
+                # vec.requires_grad_(True)
+                # pe.requires_grad_(True)
+                #model_pred = unet(img=intermediate_img, txt=intermediate_txt, vec=vec, pe=pe, txt_attention_mask=t5_attn_mask)
 
         # unpack latents
         model_pred = flux_utils.unpack_latents(model_pred, packed_latent_height, packed_latent_width)
@@ -442,70 +436,3 @@ class FluxNetworkTrainer(NetworkTrainer):
         metadata["ss_sigmoid_scale"] = args.sigmoid_scale
         metadata["ss_model_prediction_type"] = args.model_prediction_type
         metadata["ss_discrete_flow_shift"] = args.discrete_flow_shift
-
-    def is_text_encoder_not_needed_for_training(self, args):
-        return args.cache_text_encoder_outputs and not self.is_train_text_encoder(args)
-
-    def prepare_text_encoder_grad_ckpt_workaround(self, index, text_encoder):
-        if index == 0:  # CLIP-L
-            return super().prepare_text_encoder_grad_ckpt_workaround(index, text_encoder)
-        else:  # T5XXL
-            text_encoder.encoder.embed_tokens.requires_grad_(True)
-
-    def prepare_text_encoder_fp8(self, index, text_encoder, te_weight_dtype, weight_dtype):
-        if index == 0:  # CLIP-L
-            logger.info(f"prepare CLIP-L for fp8: set to {te_weight_dtype}, set embeddings to {weight_dtype}")
-            text_encoder.to(te_weight_dtype)  # fp8
-            text_encoder.text_model.embeddings.to(dtype=weight_dtype)
-        else:  # T5XXL
-
-            def prepare_fp8(text_encoder, target_dtype):
-                def forward_hook(module):
-                    def forward(hidden_states):
-                        hidden_gelu = module.act(module.wi_0(hidden_states))
-                        hidden_linear = module.wi_1(hidden_states)
-                        hidden_states = hidden_gelu * hidden_linear
-                        hidden_states = module.dropout(hidden_states)
-
-                        hidden_states = module.wo(hidden_states)
-                        return hidden_states
-
-                    return forward
-
-                for module in text_encoder.modules():
-                    if module.__class__.__name__ in ["T5LayerNorm", "Embedding"]:
-                        # print("set", module.__class__.__name__, "to", target_dtype)
-                        module.to(target_dtype)
-                    if module.__class__.__name__ in ["T5DenseGatedActDense"]:
-                        # print("set", module.__class__.__name__, "hooks")
-                        module.forward = forward_hook(module)
-
-            if flux_utils.get_t5xxl_actual_dtype(text_encoder) == torch.float8_e4m3fn and text_encoder.dtype == weight_dtype:
-                logger.info(f"T5XXL already prepared for fp8")
-            else:
-                logger.info(f"prepare T5XXL for fp8: set to {te_weight_dtype}, set embeddings to {weight_dtype}, add hooks")
-                text_encoder.to(te_weight_dtype)  # fp8
-                prepare_fp8(text_encoder, weight_dtype)
-
-
-def setup_parser() -> argparse.ArgumentParser:
-    parser = setup_parser()
-    flux_train_utils.add_flux_train_arguments(parser)
-
-    parser.add_argument(
-        "--split_mode",
-        action="store_true",
-        help="[EXPERIMENTAL] use split mode for Flux model, network arg `train_blocks=single` is required"
-    )
-    return parser
-
-
-if __name__ == "__main__":
-    parser = setup_parser()
-
-    args = parser.parse_args()
-    train_util.verify_command_line_training_args(args)
-    args = train_util.read_config_from_file(args, parser)
-
-    trainer = FluxNetworkTrainer()
-    trainer.train(args)

library/flux_models.py

@@ -1095,9 +1095,9 @@ class FluxUpper(nn.Module):
     Transformer model for flow matching on sequences.
     """
 
-    def __init__(self, params: FluxParams):
+    def __init__(self, params: FluxParams, lower_model):
         super().__init__()
-
+        self.lower_model = lower_model
         self.params = params
         self.in_channels = params.in_channels
         self.out_channels = self.in_channels
@@ -1127,6 +1127,19 @@ class FluxUpper(nn.Module):
             ]
         )
 
+        self.excluded_blocks = [7]
+        if self.excluded_blocks is None:
+            self.excluded_blocks = []  # default to no blocks excluded
+
+        self.single_blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio)
+                for i in range(params.depth_single_blocks) if i not in self.excluded_blocks
+            ]
+        )
+        print("UPPER: Single blocks: ", self.single_blocks)
+
+        self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
         self.gradient_checkpointing = False
 
     @property
@@ -1173,6 +1186,7 @@ class FluxUpper(nn.Module):
         y: Tensor,
         guidance: Tensor | None = None,
         txt_attention_mask: Tensor | None = None,
+        train_lower=False
     ) -> Tensor:
         if img.ndim != 3 or txt.ndim != 3:
             raise ValueError("Input img and txt tensors must have 3 dimensions.")
@@ -1193,7 +1207,20 @@ class FluxUpper(nn.Module):
         for block in self.double_blocks:
             img, txt = block(img=img, txt=txt, vec=vec, pe=pe, txt_attention_mask=txt_attention_mask)
 
-        return img, txt, vec, pe
+        img = torch.cat((txt, img), 1)
+
+        for i, block in enumerate(self.single_blocks):
+            if i in self.excluded_blocks:
+                img = self.lower_model(img, txt, vec=vec, pe=pe, txt_attention_mask=txt_attention_mask, train=train_lower)
+            else:
+                img = block(img, vec=vec, pe=pe, txt_attention_mask=txt_attention_mask)
+            print(img.shape)
+
+        img = img[:, txt.shape[1]:, ...]
+
+        img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
+
+        return img
 
 
 class FluxLower(nn.Module):
@@ -1207,14 +1234,23 @@ class FluxLower(nn.Module):
         self.num_heads = params.num_heads
         self.out_channels = params.in_channels
 
+        selected_blocks = [7]
+
+        if selected_blocks is None:
+            selected_blocks = range(params.depth_single_blocks)  # default to all blocks
+
         self.single_blocks = nn.ModuleList(
             [
                 SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio)
-                for _ in range(params.depth_single_blocks)
+                for i in selected_blocks
             ]
         )
+        
+        for i, block in enumerate(self.single_blocks):
+            print(f"LOWER: Single block {i}: {block.__class__.__name__}")
+        print("LOWER: Single blocks: ", self.single_blocks)
 
-        self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
+        #self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
 
         self.gradient_checkpointing = False
 
@@ -1249,11 +1285,20 @@ class FluxLower(nn.Module):
         vec: Tensor | None = None,
         pe: Tensor | None = None,
         txt_attention_mask: Tensor | None = None,
+        train: bool = False,
     ) -> Tensor:
-        img = torch.cat((txt, img), 1)
-        for block in self.single_blocks:
-            img = block(img, vec=vec, pe=pe, txt_attention_mask=txt_attention_mask)
-        img = img[:, txt.shape[1] :, ...]
+        
+        if train:
+            img.requires_grad_(True)
+            txt.requires_grad_(True)
+            vec.requires_grad_(True)
+            pe.requires_grad_(True)
+        #img = torch.cat((txt, img), 1)
+        print("img.shape to lower: ", img.shape)
+        with torch.enable_grad():
+            for block in self.single_blocks:
+                img = block(img, vec=vec, pe=pe, txt_attention_mask=txt_attention_mask)
+        #img = img[:, txt.shape[1] :, ...]
 
-        img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
+        #img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
         return img

library/flux_train_utils.py

@@ -83,7 +83,7 @@ def sample_images(
     except Exception:
         pass
 
-    with torch.no_grad(), accelerator.autocast():
+    with torch.no_grad():
         image_tensor_list = []
         for prompt_dict in prompts:
             image_tensor = sample_image_inference(
@@ -180,27 +180,13 @@ def sample_image_inference(
     tokenize_strategy = strategy_base.TokenizeStrategy.get_strategy()
     encoding_strategy = strategy_base.TextEncodingStrategy.get_strategy()
 
-    text_encoder_conds = []
     if sample_prompts_te_outputs and prompt in sample_prompts_te_outputs:
-        text_encoder_conds = sample_prompts_te_outputs[prompt]
-        print(f"Using cached text encoder outputs for prompt: {prompt}")
-    if text_encoders is not None:
-        print(f"Encoding prompt: {prompt}")
+        te_outputs = sample_prompts_te_outputs[prompt]
+    else:
         tokens_and_masks = tokenize_strategy.tokenize(prompt)
-        # strategy has apply_t5_attn_mask option
-        encoded_text_encoder_conds = encoding_strategy.encode_tokens(tokenize_strategy, text_encoders, tokens_and_masks)
-        print([x.shape if x is not None else None for x in encoded_text_encoder_conds])
+        te_outputs = encoding_strategy.encode_tokens(tokenize_strategy, text_encoders, tokens_and_masks)
 
-        # if text_encoder_conds is not cached, use encoded_text_encoder_conds
-        if len(text_encoder_conds) == 0:
-            text_encoder_conds = encoded_text_encoder_conds
-        else:
-            # if encoded_text_encoder_conds is not None, update cached text_encoder_conds
-            for i in range(len(encoded_text_encoder_conds)):
-                if encoded_text_encoder_conds[i] is not None:
-                    text_encoder_conds[i] = encoded_text_encoder_conds[i]
-
-    l_pooled, t5_out, txt_ids, t5_attn_mask = text_encoder_conds
+    l_pooled, t5_out, txt_ids, t5_attn_mask = te_outputs
 
     # sample image
     weight_dtype = ae.dtype  # TOFO give dtype as argument
@@ -536,7 +522,7 @@ def add_flux_train_arguments(parser: argparse.ArgumentParser):
     parser.add_argument(
         "--apply_t5_attn_mask",
         action="store_true",
-        help="apply attention mask to T5-XXL encode and FLUX double blocks / T5-XXLエンコードとFLUXダブルブロックにアテンションマスクを適用する",
+        help="apply attention mask (zero embs) to T5-XXL / T5-XXLにアテンションマスク（ゼロ埋め）を適用する",
     )
     parser.add_argument(
         "--cache_text_encoder_outputs", action="store_true", help="cache text encoder outputs / text encoderの出力をキャッシュする"
@@ -585,10 +571,9 @@ def add_flux_train_arguments(parser: argparse.ArgumentParser):
 
     parser.add_argument(
         "--timestep_sampling",
-        choices=["sigma", "uniform", "sigmoid", "shift", "flux_shift"],
+        choices=["sigma", "uniform", "sigmoid"],
         default="sigma",
-        help="Method to sample timesteps: sigma-based, uniform random, sigmoid of random normal, shift of sigmoid and FLUX.1 shifting."
-        " / タイムステップをサンプリングする方法：sigma、random uniform、random normalのsigmoid、sigmoidのシフト、FLUX.1のシフト。",
+        help="Method to sample timesteps: sigma-based, uniform random, or sigmoid of random normal. / タイムステップをサンプリングする方法：sigma、random uniform、またはrandom normalのsigmoid。",
     )
     parser.add_argument(
         "--sigmoid_scale",

library/flux_utils.py

@@ -1,5 +1,5 @@
 import json
-from typing import Optional, Union
+from typing import Union
 import einops
 import torch
 
@@ -8,7 +8,7 @@ from accelerate import init_empty_weights
 from transformers import CLIPTextModel, CLIPConfig, T5EncoderModel, T5Config
 
 from .flux_models import Flux, AutoEncoder, configs
-from .utils import setup_logging, MemoryEfficientSafeOpen
+from .utils import setup_logging
 
 setup_logging()
 import logging
@@ -18,67 +18,32 @@ logger = logging.getLogger(__name__)
 MODEL_VERSION_FLUX_V1 = "flux1"
 
 
-# temporary copy from sd3_utils TODO refactor
-def load_safetensors(
-    path: str, device: Union[str, torch.device], disable_mmap: bool = False, dtype: Optional[torch.dtype] = torch.float32
-):
-    if disable_mmap:
-        # return safetensors.torch.load(open(path, "rb").read())
-        # use experimental loader
-        logger.info(f"Loading without mmap (experimental)")
-        state_dict = {}
-        with MemoryEfficientSafeOpen(path) as f:
-            for key in f.keys():
-                state_dict[key] = f.get_tensor(key).to(device, dtype=dtype)
-        return state_dict
-    else:
-        try:
-            return load_file(path, device=device)
-        except:
-            return load_file(path)  # prevent device invalid Error
-
-
-def load_flow_model(
-    name: str, ckpt_path: str, dtype: Optional[torch.dtype], device: Union[str, torch.device], disable_mmap: bool = False
-) -> Flux:
+def load_flow_model(name: str, ckpt_path: str, dtype: torch.dtype, device: Union[str, torch.device]) -> Flux:
     logger.info(f"Building Flux model {name}")
     with torch.device("meta"):
-        model = Flux(configs[name].params)
-        if dtype is not None:
-            model = model.to(dtype)
+        model = Flux(configs[name].params).to(dtype)
 
     # load_sft doesn't support torch.device
     logger.info(f"Loading state dict from {ckpt_path}")
-    sd = load_safetensors(ckpt_path, device=str(device), disable_mmap=disable_mmap, dtype=dtype)
-
-     # Check the first key to see if it contains the prefix
-    first_key = next(iter(sd))
-    if first_key.startswith("model.diffusion_model."):
-        # Remove the 'model.diffusion_model.' prefix from keys if it exists
-        sd = {
-            key.replace("model.diffusion_model.", ""): value
-            for key, value in sd.items()
-        }
+    sd = load_file(ckpt_path, device=str(device))
     info = model.load_state_dict(sd, strict=False, assign=True)
     logger.info(f"Loaded Flux: {info}")
     return model
 
 
-def load_ae(
-    name: str, ckpt_path: str, dtype: torch.dtype, device: Union[str, torch.device], disable_mmap: bool = False
-) -> AutoEncoder:
+def load_ae(name: str, ckpt_path: str, dtype: torch.dtype, device: Union[str, torch.device]) -> AutoEncoder:
     logger.info("Building AutoEncoder")
     with torch.device("meta"):
         ae = AutoEncoder(configs[name].ae_params).to(dtype)
 
     logger.info(f"Loading state dict from {ckpt_path}")
-    sd = load_safetensors(ckpt_path, device=str(device), disable_mmap=disable_mmap, dtype=dtype)
+    sd = load_file(ckpt_path, device=str(device))
     info = ae.load_state_dict(sd, strict=False, assign=True)
     logger.info(f"Loaded AE: {info}")
     return ae
 
 
-def load_clip_l(ckpt_path: str, dtype: torch.dtype, device: Union[str, torch.device], disable_mmap: bool = False) -> CLIPTextModel:
+def load_clip_l(ckpt_path: str, dtype: torch.dtype, device: Union[str, torch.device]) -> CLIPTextModel:
     logger.info("Building CLIP")
     CLIPL_CONFIG = {
         "_name_or_path": "clip-vit-large-patch14/",
@@ -173,15 +138,13 @@ def load_clip_l(ckpt_path: str, dtype: torch.dtype, device: Union[str, torch.dev
         clip = CLIPTextModel._from_config(config)
 
     logger.info(f"Loading state dict from {ckpt_path}")
-    sd = load_safetensors(ckpt_path, device=str(device), disable_mmap=disable_mmap, dtype=dtype)
+    sd = load_file(ckpt_path, device=str(device))
     info = clip.load_state_dict(sd, strict=False, assign=True)
     logger.info(f"Loaded CLIP: {info}")
     return clip
 
 
-def load_t5xxl(
-    ckpt_path: str, dtype: Optional[torch.dtype], device: Union[str, torch.device], disable_mmap: bool = False
-) -> T5EncoderModel:
+def load_t5xxl(ckpt_path: str, dtype: torch.dtype, device: Union[str, torch.device]) -> T5EncoderModel:
     T5_CONFIG_JSON = """
 {
   "architectures": [
@@ -221,17 +184,12 @@ def load_t5xxl(
         t5xxl = T5EncoderModel._from_config(config)
 
     logger.info(f"Loading state dict from {ckpt_path}")
-    sd = load_safetensors(ckpt_path, device=str(device), disable_mmap=disable_mmap, dtype=dtype)
+    sd = load_file(ckpt_path, device=str(device))
     info = t5xxl.load_state_dict(sd, strict=False, assign=True)
     logger.info(f"Loaded T5xxl: {info}")
     return t5xxl
 
 
-def get_t5xxl_actual_dtype(t5xxl: T5EncoderModel) -> torch.dtype:
-    # nn.Embedding is the first layer, but it could be casted to bfloat16 or float32
-    return t5xxl.encoder.block[0].layer[0].SelfAttention.q.weight.dtype
-
-
 def prepare_img_ids(batch_size: int, packed_latent_height: int, packed_latent_width: int):
     img_ids = torch.zeros(packed_latent_height, packed_latent_width, 3)
     img_ids[..., 1] = img_ids[..., 1] + torch.arange(packed_latent_height)[:, None]

library/model_util.py

@@ -1004,15 +1004,6 @@ def load_models_from_stable_diffusion_checkpoint(v2, ckpt_path, device="cpu", dt
     unet_config = create_unet_diffusers_config(v2, unet_use_linear_projection_in_v2)
     converted_unet_checkpoint = convert_ldm_unet_checkpoint(v2, state_dict, unet_config)
 
-    # convert keys of comfy saved models
-    first_key = next(iter(converted_unet_checkpoint))
-    if first_key.startswith("model.diffusion_model."):
-        # Remove the 'model.diffusion_model.' prefix from keys if it exists
-        converted_unet_checkpoint = {
-            key.replace("model.diffusion_model.", ""): value
-            for key, value in converted_unet_checkpoint.items()
-        }
-
     unet = UNet2DConditionModel(**unet_config).to(device)
     info = unet.load_state_dict(converted_unet_checkpoint)
     logger.info(f"loading u-net: {info}")

library/strategy_flux.py

@@ -6,7 +6,6 @@ import numpy as np
 from transformers import CLIPTokenizer, T5TokenizerFast
 
 from . import train_util
-from .flux_utils import get_t5xxl_actual_dtype
 from .strategy_base import LatentsCachingStrategy, TextEncodingStrategy, TokenizeStrategy, TextEncoderOutputsCachingStrategy
 
 from .utils import setup_logging
@@ -81,7 +80,7 @@ class FluxTextEncodingStrategy(TextEncodingStrategy):
         else:
             t5_out = None
             txt_ids = None
-            t5_attn_mask = None  # caption may be dropped/shuffled, so t5_attn_mask should not be used to make sure the mask is same as the cached one
+            t5_attn_mask = None # caption may be dropped/shuffled, so t5_attn_mask should not be used to make sure the mask is same as the cached one
 
         return [l_pooled, t5_out, txt_ids, t5_attn_mask]  # returns t5_attn_mask for attention mask in transformer
 
@@ -100,8 +99,6 @@ class FluxTextEncoderOutputsCachingStrategy(TextEncoderOutputsCachingStrategy):
         super().__init__(cache_to_disk, batch_size, skip_disk_cache_validity_check, is_partial)
         self.apply_t5_attn_mask = apply_t5_attn_mask
 
-        self.warn_fp8_weights = False
-
     def get_outputs_npz_path(self, image_abs_path: str) -> str:
         return os.path.splitext(image_abs_path)[0] + FluxTextEncoderOutputsCachingStrategy.FLUX_TEXT_ENCODER_OUTPUTS_NPZ_SUFFIX
 
@@ -147,13 +144,6 @@ class FluxTextEncoderOutputsCachingStrategy(TextEncoderOutputsCachingStrategy):
     def cache_batch_outputs(
         self, tokenize_strategy: TokenizeStrategy, models: List[Any], text_encoding_strategy: TextEncodingStrategy, infos: List
     ):
-        if not self.warn_fp8_weights:
-            if get_t5xxl_actual_dtype(models[1]) == torch.float8_e4m3fn:
-                logger.warning(
-                    "T5 model is using fp8 weights for caching. This may affect the quality of the cached outputs."
-                )
-            self.warn_fp8_weights = True
-
         flux_text_encoding_strategy: FluxTextEncodingStrategy = text_encoding_strategy
         captions = [info.caption for info in infos]
 

library/train_util.py

@@ -3521,13 +3521,6 @@ def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth:
         "--full_bf16", action="store_true", help="bf16 training including gradients / 勾配も含めてbf16で学習する"
     )  # TODO move to SDXL training, because it is not supported by SD1/2
     parser.add_argument("--fp8_base", action="store_true", help="use fp8 for base model / base modelにfp8を使う")
-    parser.add_argument(
-        "--fp8_dtype",
-        type=str,
-        default="e4m3",
-        choices=["e4m3", "e5m2"],
-        help="fp8 dtype selection",
-    )
 
     parser.add_argument(
         "--ddp_timeout",
@@ -4789,8 +4782,6 @@ def prepare_dtype(args: argparse.Namespace):
         save_dtype = torch.float32
     elif args.save_precision == "fp8_e4m3fn":
         save_dtype = torch.float8_e4m3fn
-    elif args.save_precision == "fp8_e5m2":
-        save_dtype = torch.float8_e5m2
 
     return weight_dtype, save_dtype
 

networks/lora_flux.py

@@ -38,7 +38,6 @@ class LoRAModule(torch.nn.Module):
         dropout=None,
         rank_dropout=None,
         module_dropout=None,
-        split_dims: Optional[List[int]] = None,
     ):
         """if alpha == 0 or None, alpha is rank (no scaling)."""
         super().__init__()
@@ -52,34 +51,16 @@ class LoRAModule(torch.nn.Module):
             out_dim = org_module.out_features
 
         self.lora_dim = lora_dim
-        self.split_dims = split_dims
 
-        if split_dims is None:
-            if org_module.__class__.__name__ == "Conv2d":
-                kernel_size = org_module.kernel_size
-                stride = org_module.stride
-                padding = org_module.padding
-                self.lora_down = torch.nn.Conv2d(in_dim, self.lora_dim, kernel_size, stride, padding, bias=False)
-                self.lora_up = torch.nn.Conv2d(self.lora_dim, out_dim, (1, 1), (1, 1), bias=False)
-            else:
-                self.lora_down = torch.nn.Linear(in_dim, self.lora_dim, bias=False)
-                self.lora_up = torch.nn.Linear(self.lora_dim, out_dim, bias=False)
-
-            torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
-            torch.nn.init.zeros_(self.lora_up.weight)
+        if org_module.__class__.__name__ == "Conv2d":
+            kernel_size = org_module.kernel_size
+            stride = org_module.stride
+            padding = org_module.padding
+            self.lora_down = torch.nn.Conv2d(in_dim, self.lora_dim, kernel_size, stride, padding, bias=False)
+            self.lora_up = torch.nn.Conv2d(self.lora_dim, out_dim, (1, 1), (1, 1), bias=False)
         else:
-            # conv2d not supported
-            assert sum(split_dims) == out_dim, "sum of split_dims must be equal to out_dim"
-            assert org_module.__class__.__name__ == "Linear", "split_dims is only supported for Linear"
-            # print(f"split_dims: {split_dims}")
-            self.lora_down = torch.nn.ModuleList(
-                [torch.nn.Linear(in_dim, self.lora_dim, bias=False) for _ in range(len(split_dims))]
-            )
-            self.lora_up = torch.nn.ModuleList([torch.nn.Linear(self.lora_dim, split_dim, bias=False) for split_dim in split_dims])
-            for lora_down in self.lora_down:
-                torch.nn.init.kaiming_uniform_(lora_down.weight, a=math.sqrt(5))
-            for lora_up in self.lora_up:
-                torch.nn.init.zeros_(lora_up.weight)
+            self.lora_down = torch.nn.Linear(in_dim, self.lora_dim, bias=False)
+            self.lora_up = torch.nn.Linear(self.lora_dim, out_dim, bias=False)
 
         if type(alpha) == torch.Tensor:
             alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
@@ -88,6 +69,9 @@ class LoRAModule(torch.nn.Module):
         self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
 
         # same as microsoft's
+        torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
+        torch.nn.init.zeros_(self.lora_up.weight)
+
         self.multiplier = multiplier
         self.org_module = org_module  # remove in applying
         self.dropout = dropout
@@ -107,56 +91,30 @@ class LoRAModule(torch.nn.Module):
             if torch.rand(1) < self.module_dropout:
                 return org_forwarded
 
-        if self.split_dims is None:
-            lx = self.lora_down(x)
+        lx = self.lora_down(x)
 
-            # normal dropout
-            if self.dropout is not None and self.training:
-                lx = torch.nn.functional.dropout(lx, p=self.dropout)
+        # normal dropout
+        if self.dropout is not None and self.training:
+            lx = torch.nn.functional.dropout(lx, p=self.dropout)
 
-            # rank dropout
-            if self.rank_dropout is not None and self.training:
-                mask = torch.rand((lx.size(0), self.lora_dim), device=lx.device) > self.rank_dropout
-                if len(lx.size()) == 3:
-                    mask = mask.unsqueeze(1)  # for Text Encoder
-                elif len(lx.size()) == 4:
-                    mask = mask.unsqueeze(-1).unsqueeze(-1)  # for Conv2d
-                lx = lx * mask
+        # rank dropout
+        if self.rank_dropout is not None and self.training:
+            mask = torch.rand((lx.size(0), self.lora_dim), device=lx.device) > self.rank_dropout
+            if len(lx.size()) == 3:
+                mask = mask.unsqueeze(1)  # for Text Encoder
+            elif len(lx.size()) == 4:
+                mask = mask.unsqueeze(-1).unsqueeze(-1)  # for Conv2d
+            lx = lx * mask
 
-                # scaling for rank dropout: treat as if the rank is changed
-                # maskから計算することも考えられるが、augmentation的な効果を期待してrank_dropoutを用いる
-                scale = self.scale * (1.0 / (1.0 - self.rank_dropout))  # redundant for readability
-            else:
-                scale = self.scale
-
-            lx = self.lora_up(lx)
-
-            return org_forwarded + lx * self.multiplier * scale
+            # scaling for rank dropout: treat as if the rank is changed
+            # maskから計算することも考えられるが、augmentation的な効果を期待してrank_dropoutを用いる
+            scale = self.scale * (1.0 / (1.0 - self.rank_dropout))  # redundant for readability
         else:
-            lxs = [lora_down(x) for lora_down in self.lora_down]
+            scale = self.scale
 
-            # normal dropout
-            if self.dropout is not None and self.training:
-                lxs = [torch.nn.functional.dropout(lx, p=self.dropout) for lx in lxs]
+        lx = self.lora_up(lx)
 
-            # rank dropout
-            if self.rank_dropout is not None and self.training:
-                masks = [torch.rand((lx.size(0), self.lora_dim), device=lx.device) > self.rank_dropout for lx in lxs]
-                for i in range(len(lxs)):
-                    if len(lx.size()) == 3:
-                        masks[i] = masks[i].unsqueeze(1)
-                    elif len(lx.size()) == 4:
-                        masks[i] = masks[i].unsqueeze(-1).unsqueeze(-1)
-                    lxs[i] = lxs[i] * masks[i]
-
-                # scaling for rank dropout: treat as if the rank is changed
-                scale = self.scale * (1.0 / (1.0 - self.rank_dropout))  # redundant for readability
-            else:
-                scale = self.scale
-
-            lxs = [lora_up(lx) for lora_up, lx in zip(self.lora_up, lxs)]
-
-            return org_forwarded + torch.cat(lxs, dim=-1) * self.multiplier * scale
+        return org_forwarded + lx * self.multiplier * scale
 
 
 class LoRAInfModule(LoRAModule):
@@ -193,50 +151,31 @@ class LoRAInfModule(LoRAModule):
         if device is None:
             device = org_device
 
-        if self.split_dims is None:
-            # get up/down weight
-            down_weight = sd["lora_down.weight"].to(torch.float).to(device)
-            up_weight = sd["lora_up.weight"].to(torch.float).to(device)
+        # get up/down weight
+        up_weight = sd["lora_up.weight"].to(torch.float).to(device)
+        down_weight = sd["lora_down.weight"].to(torch.float).to(device)
 
-            # merge weight
-            if len(weight.size()) == 2:
-                # linear
-                weight = weight + self.multiplier * (up_weight @ down_weight) * self.scale
-            elif down_weight.size()[2:4] == (1, 1):
-                # conv2d 1x1
-                weight = (
-                    weight
-                    + self.multiplier
-                    * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3)
-                    * self.scale
-                )
-            else:
-                # conv2d 3x3
-                conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
-                # logger.info(conved.size(), weight.size(), module.stride, module.padding)
-                weight = weight + self.multiplier * conved * self.scale
-
-            # set weight to org_module
-            org_sd["weight"] = weight.to(dtype)
-            self.org_module.load_state_dict(org_sd)
+        # merge weight
+        if len(weight.size()) == 2:
+            # linear
+            weight = weight + self.multiplier * (up_weight @ down_weight) * self.scale
+        elif down_weight.size()[2:4] == (1, 1):
+            # conv2d 1x1
+            weight = (
+                weight
+                + self.multiplier
+                * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3)
+                * self.scale
+            )
         else:
-            # split_dims
-            total_dims = sum(self.split_dims)
-            for i in range(len(self.split_dims)):
-                # get up/down weight
-                down_weight = sd[f"lora_down.{i}.weight"].to(torch.float).to(device)  # (rank, in_dim)
-                up_weight = sd[f"lora_up.{i}.weight"].to(torch.float).to(device)  # (split dim, rank)
+            # conv2d 3x3
+            conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
+            # logger.info(conved.size(), weight.size(), module.stride, module.padding)
+            weight = weight + self.multiplier * conved * self.scale
 
-                # pad up_weight -> (total_dims, rank)
-                padded_up_weight = torch.zeros((total_dims, up_weight.size(0)), device=device, dtype=torch.float)
-                padded_up_weight[sum(self.split_dims[:i]) : sum(self.split_dims[: i + 1])] = up_weight
-
-                # merge weight
-                weight = weight + self.multiplier * (up_weight @ down_weight) * self.scale
-
-            # set weight to org_module
-            org_sd["weight"] = weight.to(dtype)
-            self.org_module.load_state_dict(org_sd)
+        # set weight to org_module
+        org_sd["weight"] = weight.to(dtype)
+        self.org_module.load_state_dict(org_sd)
 
     # 復元できるマージのため、このモジュールのweightを返す
     def get_weight(self, multiplier=None):
@@ -271,14 +210,7 @@ class LoRAInfModule(LoRAModule):
 
     def default_forward(self, x):
         # logger.info(f"default_forward {self.lora_name} {x.size()}")
-        if self.split_dims is None:
-            lx = self.lora_down(x)
-            lx = self.lora_up(lx)
-            return self.org_forward(x) + lx * self.multiplier * self.scale
-        else:
-            lxs = [lora_down(x) for lora_down in self.lora_down]
-            lxs = [lora_up(lx) for lora_up, lx in zip(self.lora_up, lxs)]
-            return self.org_forward(x) + torch.cat(lxs, dim=-1) * self.multiplier * self.scale
+        return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
 
     def forward(self, x):
         if not self.enabled:
@@ -321,22 +253,8 @@ def create_network(
 
     # single or double blocks
     train_blocks = kwargs.get("train_blocks", None)  # None (default), "all" (same as None), "single", "double"
-    if train_blocks is not None:
-        assert train_blocks in ["all", "single", "double"], f"invalid train_blocks: {train_blocks}"
-
-    only_if_contains = kwargs.get("only_if_contains", None)
-    if only_if_contains is not None:
-        only_if_contains = [word.strip() for word in only_if_contains.split(',')]
-
-    # split qkv
-    split_qkv = kwargs.get("split_qkv", False)
-    if split_qkv is not None:
-        split_qkv = True if split_qkv == "True" else False
-
-    # train T5XXL
-    train_t5xxl = kwargs.get("train_t5xxl", False)
-    if train_t5xxl is not None:
-        train_t5xxl = True if train_t5xxl == "True" else False
+    #if train_blocks is not None:
+    #    assert train_blocks in ["all", "single", "double"], f"invalid train_blocks: {train_blocks}"
 
     # すごく引数が多いな ( ^ω^)･･･
     network = LoRANetwork(
@@ -351,10 +269,7 @@ def create_network(
         conv_lora_dim=conv_dim,
         conv_alpha=conv_alpha,
         train_blocks=train_blocks,
-        split_qkv=split_qkv,
-        train_t5xxl=train_t5xxl,
         varbose=True,
-        only_if_contains=only_if_contains
     )
 
     loraplus_lr_ratio = kwargs.get("loraplus_lr_ratio", None)
@@ -380,10 +295,9 @@ def create_network_from_weights(multiplier, file, ae, text_encoders, flux, weigh
         else:
             weights_sd = torch.load(file, map_location="cpu")
 
-    # get dim/alpha mapping, and train t5xxl
+    # get dim/alpha mapping
     modules_dim = {}
     modules_alpha = {}
-    train_t5xxl = None
     for key, value in weights_sd.items():
         if "." not in key:
             continue
@@ -396,41 +310,10 @@ def create_network_from_weights(multiplier, file, ae, text_encoders, flux, weigh
             modules_dim[lora_name] = dim
             # logger.info(lora_name, value.size(), dim)
 
-        if train_t5xxl is None or train_t5xxl is False:
-            train_t5xxl = "lora_te3" in lora_name
-
-    if train_t5xxl is None:
-        train_t5xxl = False
-
-    # # split qkv
-    # double_qkv_rank = None
-    # single_qkv_rank = None
-    # rank = None
-    # for lora_name, dim in modules_dim.items():
-    #     if "double" in lora_name and "qkv" in lora_name:
-    #         double_qkv_rank = dim
-    #     elif "single" in lora_name and "linear1" in lora_name:
-    #         single_qkv_rank = dim
-    #     elif rank is None:
-    #         rank = dim
-    #     if double_qkv_rank is not None and single_qkv_rank is not None and rank is not None:
-    #         break
-    # split_qkv = (double_qkv_rank is not None and double_qkv_rank != rank) or (
-    #     single_qkv_rank is not None and single_qkv_rank != rank
-    # )
-    split_qkv = False  # split_qkv is not needed to care, because state_dict is qkv combined
-
     module_class = LoRAInfModule if for_inference else LoRAModule
 
     network = LoRANetwork(
-        text_encoders,
-        flux,
-        multiplier=multiplier,
-        modules_dim=modules_dim,
-        modules_alpha=modules_alpha,
-        module_class=module_class,
-        split_qkv=split_qkv,
-        train_t5xxl=train_t5xxl,
+        text_encoders, flux, multiplier=multiplier, modules_dim=modules_dim, modules_alpha=modules_alpha, module_class=module_class
     )
     return network, weights_sd
 
@@ -439,10 +322,10 @@ class LoRANetwork(torch.nn.Module):
     # FLUX_TARGET_REPLACE_MODULE = ["DoubleStreamBlock", "SingleStreamBlock"]
     FLUX_TARGET_REPLACE_MODULE_DOUBLE = ["DoubleStreamBlock"]
     FLUX_TARGET_REPLACE_MODULE_SINGLE = ["SingleStreamBlock"]
-    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPSdpaAttention", "CLIPMLP", "T5Attention", "T5DenseGatedActDense"]
+    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPSdpaAttention", "CLIPMLP"]
     LORA_PREFIX_FLUX = "lora_unet"  # make ComfyUI compatible
     LORA_PREFIX_TEXT_ENCODER_CLIP = "lora_te1"
-    LORA_PREFIX_TEXT_ENCODER_T5 = "lora_te3"  # make ComfyUI compatible
+    LORA_PREFIX_TEXT_ENCODER_T5 = "lora_te2"
 
     def __init__(
         self,
@@ -460,10 +343,7 @@ class LoRANetwork(torch.nn.Module):
         modules_dim: Optional[Dict[str, int]] = None,
         modules_alpha: Optional[Dict[str, int]] = None,
         train_blocks: Optional[str] = None,
-        split_qkv: bool = False,
-        train_t5xxl: bool = False,
         varbose: Optional[bool] = False,
-        only_if_contains: Optional[List[str]] = None,
     ) -> None:
         super().__init__()
         self.multiplier = multiplier
@@ -476,15 +356,11 @@ class LoRANetwork(torch.nn.Module):
         self.rank_dropout = rank_dropout
         self.module_dropout = module_dropout
         self.train_blocks = train_blocks if train_blocks is not None else "all"
-        self.split_qkv = split_qkv
-        self.train_t5xxl = train_t5xxl
 
         self.loraplus_lr_ratio = None
         self.loraplus_unet_lr_ratio = None
         self.loraplus_text_encoder_lr_ratio = None
 
-        self.only_if_contains = only_if_contains
-
         if modules_dim is not None:
             logger.info(f"create LoRA network from weights")
         else:
@@ -492,18 +368,10 @@ class LoRANetwork(torch.nn.Module):
             logger.info(
                 f"neuron dropout: p={self.dropout}, rank dropout: p={self.rank_dropout}, module dropout: p={self.module_dropout}"
             )
-            # if self.conv_lora_dim is not None:
-            #     logger.info(
-            #         f"apply LoRA to Conv2d with kernel size (3,3). dim (rank): {self.conv_lora_dim}, alpha: {self.conv_alpha}"
-            #     )
-        if self.split_qkv:
-            logger.info(f"split qkv for LoRA")
-        if self.train_blocks is not None:
-            logger.info(f"train {self.train_blocks} blocks only")
-        if train_t5xxl:
-            logger.info(f"train T5XXL as well")
-
-        #self.only_if_contains = ["lora_unet_single_blocks_20_linear2"]
+            if self.conv_lora_dim is not None:
+                logger.info(
+                    f"apply LoRA to Conv2d with kernel size (3,3). dim (rank): {self.conv_lora_dim}, alpha: {self.conv_alpha}"
+                )
 
         # create module instances
         def create_modules(
@@ -515,66 +383,77 @@ class LoRANetwork(torch.nn.Module):
                 else (self.LORA_PREFIX_TEXT_ENCODER_CLIP if text_encoder_idx == 0 else self.LORA_PREFIX_TEXT_ENCODER_T5)
             )
 
+            def process_module(name, module, prefix, modules_dim, modules_alpha, dropout, rank_dropout, module_dropout, target_replace_modules):
+                loras = []
+                skipped = []
+            
+                def process_child(child_name, child_module):
+                    is_linear = child_module.__class__.__name__ == "Linear"
+                    is_conv2d = child_module.__class__.__name__ == "Conv2d"
+                    is_conv2d_1x1 = is_conv2d and child_module.kernel_size == (1, 1)
+            
+                    if is_linear or is_conv2d:
+                        lora_name = prefix + "." + name + "." + child_name
+                        lora_name = lora_name.replace(".", "_")
+            
+                        dim = None
+                        alpha = None
+            
+                        if modules_dim is not None:
+                            # Module specified
+                            if lora_name in modules_dim:
+                                dim = modules_dim[lora_name]
+                                alpha = modules_alpha[lora_name]
+                        else:
+                            # Normally, target all
+                            if is_linear or is_conv2d_1x1:
+                                dim = self.lora_dim
+                                alpha = self.alpha
+                            elif self.conv_lora_dim is not None:
+                                dim = self.conv_lora_dim
+                                alpha = self.conv_alpha
+            
+                        if dim is None or dim == 0:
+                            # Output skipped information
+                            if is_linear or is_conv2d_1x1 or (self.conv_lora_dim is not None):
+                                skipped.append(lora_name)
+                            return
+            
+                        lora = module_class(
+                            lora_name,
+                            child_module,
+                            self.multiplier,
+                            dim,
+                            alpha,
+                            dropout=dropout,
+                            rank_dropout=rank_dropout,
+                            module_dropout=module_dropout,
+                        )
+                        loras.append(lora)
+            
+                for child_name, child_module in module.named_modules():
+                    process_child(child_name, child_module)
+            
+                return loras, skipped
+            
+            
             loras = []
             skipped = []
+
+            target_replace_modules = [module.strip() for module in target_replace_modules]
+            
             for name, module in root_module.named_modules():
-                if module.__class__.__name__ in target_replace_modules:
-                    for child_name, child_module in module.named_modules():
-                        is_linear = child_module.__class__.__name__ == "Linear"
-                        is_conv2d = child_module.__class__.__name__ == "Conv2d"
-                        is_conv2d_1x1 = is_conv2d and child_module.kernel_size == (1, 1)
-
-                        if is_linear or is_conv2d:
-                            lora_name = prefix + "." + name + "." + child_name
-                            lora_name = lora_name.replace(".", "_")
-                            #lora_unet_single_blocks_20_linear2
-
-                            if "unet" in lora_name and (self.only_if_contains is not None and not any(word in lora_name for word in self.only_if_contains)):
-                                continue
-
-                            dim = None
-                            alpha = None
-
-                            if modules_dim is not None:
-                                # モジュール指定あり
-                                if lora_name in modules_dim:
-                                    dim = modules_dim[lora_name]
-                                    alpha = modules_alpha[lora_name]
-                            else:
-                                # 通常、すべて対象とする
-                                if is_linear or is_conv2d_1x1:
-                                    dim = self.lora_dim
-                                    alpha = self.alpha
-                                elif self.conv_lora_dim is not None:
-                                    dim = self.conv_lora_dim
-                                    alpha = self.conv_alpha
-
-                            if dim is None or dim == 0:
-                                # skipした情報を出力
-                                if is_linear or is_conv2d_1x1 or (self.conv_lora_dim is not None):
-                                    skipped.append(lora_name)
-                                continue
-
-                            # qkv split
-                            split_dims = None
-                            if is_flux and split_qkv:
-                                if "double" in lora_name and "qkv" in lora_name:
-                                    split_dims = [3072] * 3
-                                elif "single" in lora_name and "linear1" in lora_name:
-                                    split_dims = [3072] * 3 + [12288]
-
-                            lora = module_class(
-                                lora_name,
-                                child_module,
-                                self.multiplier,
-                                dim,
-                                alpha,
-                                dropout=dropout,
-                                rank_dropout=rank_dropout,
-                                module_dropout=module_dropout,
-                                split_dims=split_dims,
-                            )
-                            loras.append(lora)
+                if any("blocks" in part for part in target_replace_modules) and name in target_replace_modules:
+                    print(module)
+                    module_loras, module_skipped = process_module(name, module, prefix, modules_dim, modules_alpha, dropout, rank_dropout, module_dropout, target_replace_modules)
+                    loras.extend(module_loras)
+                    skipped.extend(module_skipped)
+                elif module.__class__.__name__ in target_replace_modules:
+                    print(module)
+                    module_loras, module_skipped = process_module(name, module, prefix, modules_dim, modules_alpha, dropout, rank_dropout, module_dropout, target_replace_modules)
+                    loras.extend(module_loras)
+                    skipped.extend(module_skipped)
+            
             return loras, skipped
 
         # create LoRA for text encoder
@@ -583,18 +462,18 @@ class LoRANetwork(torch.nn.Module):
         skipped_te = []
         for i, text_encoder in enumerate(text_encoders):
             index = i
-            if not train_t5xxl and index > 0:  # 0: CLIP, 1: T5XXL, so we skip T5XXL if train_t5xxl is False
-                break
-
             logger.info(f"create LoRA for Text Encoder {index+1}:")
 
             text_encoder_loras, skipped = create_modules(False, index, text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)
-            logger.info(f"create LoRA for Text Encoder {index+1}: {len(text_encoder_loras)} modules.")
             self.text_encoder_loras.extend(text_encoder_loras)
             skipped_te += skipped
-
+        logger.info(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
+        print("TRAIN BLOCKS:", self.train_blocks)
         # create LoRA for U-Net
-        if self.train_blocks == "all":
+        if any("blocks" in part for part in self.train_blocks.split(',')):
+            target_replace_modules = self.train_blocks.split(',')
+            print("TARGET_REPLACE_MODULES:", target_replace_modules)
+        elif self.train_blocks == "all":
             target_replace_modules = LoRANetwork.FLUX_TARGET_REPLACE_MODULE_DOUBLE + LoRANetwork.FLUX_TARGET_REPLACE_MODULE_SINGLE
         elif self.train_blocks == "single":
             target_replace_modules = LoRANetwork.FLUX_TARGET_REPLACE_MODULE_SINGLE
@@ -604,7 +483,6 @@ class LoRANetwork(torch.nn.Module):
         self.unet_loras: List[Union[LoRAModule, LoRAInfModule]]
         self.unet_loras, skipped_un = create_modules(True, None, unet, target_replace_modules)
         logger.info(f"create LoRA for FLUX {self.train_blocks} blocks: {len(self.unet_loras)} modules.")
-        print(self.unet_loras)
 
         skipped = skipped_te + skipped_un
         if varbose and len(skipped) > 0:
@@ -640,111 +518,6 @@ class LoRANetwork(torch.nn.Module):
         info = self.load_state_dict(weights_sd, False)
         return info
 
-    def load_state_dict(self, state_dict, strict=True):
-        # override to convert original weight to split qkv
-        if not self.split_qkv:
-            return super().load_state_dict(state_dict, strict)
-
-        # split qkv
-        for key in list(state_dict.keys()):
-            if "double" in key and "qkv" in key:
-                split_dims = [3072] * 3
-            elif "single" in key and "linear1" in key:
-                split_dims = [3072] * 3 + [12288]
-            else:
-                continue
-
-            weight = state_dict[key]
-            lora_name = key.split(".")[0]
-            if "lora_down" in key and "weight" in key:
-                # dense weight (rank*3, in_dim)
-                split_weight = torch.chunk(weight, len(split_dims), dim=0)
-                for i, split_w in enumerate(split_weight):
-                    state_dict[f"{lora_name}.lora_down.{i}.weight"] = split_w
-
-                del state_dict[key]
-                # print(f"split {key}: {weight.shape} to {[w.shape for w in split_weight]}")
-            elif "lora_up" in key and "weight" in key:
-                # sparse weight (out_dim=sum(split_dims), rank*3)
-                rank = weight.size(1) // len(split_dims)
-                i = 0
-                for j in range(len(split_dims)):
-                    state_dict[f"{lora_name}.lora_up.{j}.weight"] = weight[i : i + split_dims[j], j * rank : (j + 1) * rank]
-                    i += split_dims[j]
-                del state_dict[key]
-
-                # # check is sparse
-                # i = 0
-                # is_zero = True
-                # for j in range(len(split_dims)):
-                #     for k in range(len(split_dims)):
-                #         if j == k:
-                #             continue
-                #         is_zero = is_zero and torch.all(weight[i : i + split_dims[j], k * rank : (k + 1) * rank] == 0)
-                #     i += split_dims[j]
-                # if not is_zero:
-                #     logger.warning(f"weight is not sparse: {key}")
-                # else:
-                #     logger.info(f"weight is sparse: {key}")
-
-                # print(
-                #     f"split {key}: {weight.shape} to {[state_dict[k].shape for k in [f'{lora_name}.lora_up.{j}.weight' for j in range(len(split_dims))]]}"
-                # )
-
-            # alpha is unchanged
-
-        return super().load_state_dict(state_dict, strict)
-
-    def state_dict(self, destination=None, prefix="", keep_vars=False):
-        if not self.split_qkv:
-            return super().state_dict(destination, prefix, keep_vars)
-
-        # merge qkv
-        state_dict = super().state_dict(destination, prefix, keep_vars)
-        new_state_dict = {}
-        for key in list(state_dict.keys()):
-            if "double" in key and "qkv" in key:
-                split_dims = [3072] * 3
-            elif "single" in key and "linear1" in key:
-                split_dims = [3072] * 3 + [12288]
-            else:
-                new_state_dict[key] = state_dict[key]
-                continue
-
-            if key not in state_dict:
-                continue  # already merged
-
-            lora_name = key.split(".")[0]
-
-            # (rank, in_dim) * 3
-            down_weights = [state_dict.pop(f"{lora_name}.lora_down.{i}.weight") for i in range(len(split_dims))]
-            # (split dim, rank) * 3
-            up_weights = [state_dict.pop(f"{lora_name}.lora_up.{i}.weight") for i in range(len(split_dims))]
-
-            alpha = state_dict.pop(f"{lora_name}.alpha")
-
-            # merge down weight
-            down_weight = torch.cat(down_weights, dim=0)  # (rank, split_dim) * 3 -> (rank*3, sum of split_dim)
-
-            # merge up weight (sum of split_dim, rank*3)
-            rank = up_weights[0].size(1)
-            up_weight = torch.zeros((sum(split_dims), down_weight.size(0)), device=down_weight.device, dtype=down_weight.dtype)
-            i = 0
-            for j in range(len(split_dims)):
-                up_weight[i : i + split_dims[j], j * rank : (j + 1) * rank] = up_weights[j]
-                i += split_dims[j]
-
-            new_state_dict[f"{lora_name}.lora_down.weight"] = down_weight
-            new_state_dict[f"{lora_name}.lora_up.weight"] = up_weight
-            new_state_dict[f"{lora_name}.alpha"] = alpha
-
-            # print(
-            #     f"merged {lora_name}: {lora_name}, {[w.shape for w in down_weights]}, {[w.shape for w in up_weights]} to {down_weight.shape}, {up_weight.shape}"
-            # )
-            print(f"new key: {lora_name}.lora_down.weight, {lora_name}.lora_up.weight, {lora_name}.alpha")
-
-        return new_state_dict
-
     def apply_to(self, text_encoders, flux, apply_text_encoder=True, apply_unet=True):
         if apply_text_encoder:
             logger.info(f"enable LoRA for text encoder: {len(self.text_encoder_loras)} modules")

nodes.py

@@ -9,8 +9,6 @@ import toml
 import json
 import time
 import shutil
-import shlex
-
 from pathlib import Path
 script_directory = os.path.dirname(os.path.abspath(__file__))
 
@@ -37,14 +35,11 @@ class FluxTrainModelSelect:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
-                    "transformer": (folder_paths.get_filename_list("unet"), ),
-                    "vae": (folder_paths.get_filename_list("vae"), ),
-                    "clip_l": (folder_paths.get_filename_list("clip"), ),
-                    "t5": (folder_paths.get_filename_list("clip"), ),
-                },
-                "optional": {
-                    "lora_path": ("STRING",{"multiline": True, "forceInput": True, "default": "", "tooltip": "pre-trained LoRA path to load (network_weights)"}),
-                }
+            "transformer": (folder_paths.get_filename_list("unet"), ),
+            "vae": (folder_paths.get_filename_list("vae"), ),
+            "clip_l": (folder_paths.get_filename_list("clip"), ),
+            "t5": (folder_paths.get_filename_list("clip"), ),
+           },
         }
 
     RETURN_TYPES = ("TRAIN_FLUX_MODELS",)
@@ -52,7 +47,7 @@ class FluxTrainModelSelect:
     FUNCTION = "loadmodel"
     CATEGORY = "FluxTrainer"
 
-    def loadmodel(self, transformer, vae, clip_l, t5, lora_path=""):
+    def loadmodel(self, transformer, vae, clip_l, t5):
         
         transformer_path = folder_paths.get_full_path("unet", transformer)
         vae_path = folder_paths.get_full_path("vae", vae)
@@ -63,20 +58,12 @@ class FluxTrainModelSelect:
             "transformer": transformer_path,
             "vae": vae_path,
             "clip_l": clip_path,
-            "t5": t5_path,
-            "lora_path": lora_path
+            "t5": t5_path
         }
         
         return (flux_models,)
 
 class TrainDatasetGeneralConfig:
-    queue_counter = 0
-    @classmethod
-    def IS_CHANGED(s, reset_on_queue=False, **kwargs):
-        if reset_on_queue:
-            s.queue_counter += 1
-        print(f"queue_counter: {s.queue_counter}")
-        return s.queue_counter
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
@@ -86,10 +73,6 @@ class TrainDatasetGeneralConfig:
             "caption_dropout_rate": ("FLOAT",{"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01,"tooltip": "tag dropout rate"}),
             "alpha_mask": ("BOOLEAN",{"default": False, "tooltip": "use alpha channel as mask for training"}),
             },
-            "optional": {
-                "reset_on_queue": ("BOOLEAN",{"default": False, "tooltip": "Force refresh of everything for cleaner queueing"}),
-                "reg_data_dir": ("STRING",{"multiline": True, "forceInput": True, "default": "", "tooltip": "reg data dir"}),
-            }
         }
 
     RETURN_TYPES = ("JSON",)
@@ -97,7 +80,7 @@ class TrainDatasetGeneralConfig:
     FUNCTION = "create_config"
     CATEGORY = "FluxTrainer"
 
-    def create_config(self, shuffle_caption, caption_dropout_rate, color_aug, flip_aug, alpha_mask, reset_on_queue=False, reg_data_dir=""):
+    def create_config(self, shuffle_caption, caption_dropout_rate, color_aug, flip_aug, alpha_mask):
         
         dataset = {
            "general": {
@@ -114,15 +97,13 @@ class TrainDatasetGeneralConfig:
         #print(dataset_json)
         dataset_config = {
             "datasets": dataset_json,
-            "alpha_mask": alpha_mask,
-            "reg_data_dir": reg_data_dir
+            "alpha_mask": alpha_mask
         }
         return (dataset_config,)
 
 class TrainDatasetAdd:
     def __init__(self):
         self.previous_dataset_signature = None
-        
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
@@ -137,6 +118,7 @@ class TrainDatasetAdd:
             "num_repeats": ("INT", {"default": 1, "min": 1, "tooltip": "number of times to repeat dataset for an epoch"}),
             "min_bucket_reso": ("INT", {"default": 256, "min": 64, "max": 4096, "step": 8, "tooltip": "min bucket resolution"}),
             "max_bucket_reso": ("INT", {"default": 1024, "min": 64, "max": 4096, "step": 8, "tooltip": "max bucket resolution"}),
+  
             },
         }
 
@@ -197,7 +179,7 @@ class OptimizerConfig:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": {
-            "optimizer_type": (["adamw8bit", "adamw","prodigy", "CAME", "Lion8bit", "Lion"], {"default": "adamw8bit", "tooltip": "optimizer type"}),
+            "optimizer_type": (["adamw8bit", "adamw","prodigy", "CAME"], {"default": "adamw8bit", "tooltip": "optimizer type"}),
             "max_grad_norm": ("FLOAT",{"default": 1.0, "min": 0.0, "tooltip": "gradient clipping"}),
             "lr_scheduler": (["constant", "cosine", "cosine_with_restarts", "polynomial", "constant_with_warmup"], {"default": "constant", "tooltip": "learning rate scheduler"}),
             "lr_warmup_steps": ("INT",{"default": 0, "min": 0, "tooltip": "learning rate warmup steps"}),
@@ -215,7 +197,7 @@ class OptimizerConfig:
 
     def create_config(self, min_snr_gamma, extra_optimizer_args, **kwargs):
         kwargs["min_snr_gamma"] = min_snr_gamma if min_snr_gamma != 0.0 else None
-        kwargs["optimizer_args"] = [arg.strip() for arg in extra_optimizer_args.strip().split('|') if arg.strip()]
+        kwargs["optimizer_args"] = [arg.strip() for arg in extra_optimizer_args.strip().split(',') if arg.strip()]
         return (kwargs,)
 
 class OptimizerConfigAdafactor:
@@ -243,7 +225,7 @@ class OptimizerConfigAdafactor:
 
     def create_config(self, relative_step, scale_parameter, warmup_init, clip_threshold, min_snr_gamma, extra_optimizer_args, **kwargs):
         kwargs["optimizer_type"] = "adafactor"
-        extra_args = [arg.strip() for arg in extra_optimizer_args.strip().split('|') if arg.strip()]
+        extra_args = [arg.strip() for arg in extra_optimizer_args.strip().split(',') if arg.strip()]
         node_args = [
                 f"relative_step={relative_step}",
                 f"scale_parameter={scale_parameter}",
@@ -279,7 +261,7 @@ class OptimizerConfigProdigy:
 
     def create_config(self, weight_decay, decouple, min_snr_gamma, use_bias_correction, extra_optimizer_args, **kwargs):
         kwargs["optimizer_type"] = "prodigy"
-        extra_args = [arg.strip() for arg in extra_optimizer_args.strip().split('|') if arg.strip()]
+        extra_args = [arg.strip() for arg in extra_optimizer_args.strip().split(',') if arg.strip()]
         node_args = [
                 f"weight_decay={weight_decay}",
                 f"decouple={decouple}",
@@ -288,7 +270,23 @@ class OptimizerConfigProdigy:
         kwargs["optimizer_args"] = node_args + extra_args
         kwargs["min_snr_gamma"] = min_snr_gamma if min_snr_gamma != 0.0 else None
         
-        return (kwargs,)    
+        return (kwargs,)
+    
+class FluxLoRATrainBlocks:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "train_blocks": ("STRING",{"default": "single_blocks.7", "multiline": True, "tooltip": "specify individual blocks to include in the training, for example 'single_blocks.7'"}),
+           },
+        }
+
+    RETURN_TYPES = ("BLOCKS",)
+    RETURN_NAMES = ("train_blocks",)
+    FUNCTION = "create_config"
+    CATEGORY = "FluxTrainer"
+
+    def create_config(self, train_blocks):
+        return (train_blocks,)    
 
 class InitFluxLoRATraining:
     @classmethod
@@ -302,8 +300,12 @@ class InitFluxLoRATraining:
             "network_dim": ("INT", {"default": 4, "min": 1, "max": 256, "step": 1, "tooltip": "network dim"}),
             "network_alpha": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 256.0, "step": 0.01, "tooltip": "network alpha"}),
             "learning_rate": ("FLOAT", {"default": 4e-4, "min": 0.0, "max": 10.0, "step": 0.000001, "tooltip": "learning rate"}),
+            #"unet_lr": ("FLOAT", {"default": 1e-4, "min": 0.0, "max": 10.0, "step": 0.00001, "tooltip": "unet learning rate"}),
+            #"max_train_epochs": ("INT", {"default": 4, "min": 1, "max": 1000, "step": 1, "tooltip": "max number of training epochs"}),
             "max_train_steps": ("INT", {"default": 1500, "min": 1, "max": 100000, "step": 1, "tooltip": "max number of training steps"}),
+            #"text_encoder_lr": ("FLOAT", {"default": 0, "min": 0.0, "max": 10.0, "step": 0.00001, "tooltip": "text encoder learning rate"}),
             "apply_t5_attn_mask": ("BOOLEAN", {"default": True, "tooltip": "apply t5 attention mask"}),
+            #"t5xxl_max_token_length": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 8, "tooltip": "dev uses 512, schnell 256"}),
             "cache_latents": (["disk", "memory", "disabled"], {"tooltip": "caches text encoder outputs"}),
             "cache_text_encoder_outputs": (["disk", "memory", "disabled"], {"tooltip": "caches text encoder outputs"}),
             "split_mode": ("BOOLEAN", {"default": False, "tooltip": "[EXPERIMENTAL] use split mode for Flux model, network arg `train_blocks=single` is required"}),
@@ -319,17 +321,17 @@ class InitFluxLoRATraining:
             "highvram": ("BOOLEAN", {"default": False, "tooltip": "memory mode"}),
             "fp8_base": ("BOOLEAN", {"default": True, "tooltip": "use fp8 for base model"}),
             "gradient_dtype": (["fp32", "fp16", "bf16"], {"default": "fp32", "tooltip": "the actual dtype training uses"}),
-            "save_dtype": (["fp32", "fp16", "bf16", "fp8_e4m3fn", "fp8_e5m2"], {"default": "bf16", "tooltip": "the dtype to save checkpoints as"}),
+            "save_dtype": (["fp32", "fp16", "bf16", "fp8_e4m3fn"], {"default": "bf16", "tooltip": "the dtype to save checkpoints as"}),
             "attention_mode": (["sdpa", "xformers", "disabled"], {"default": "sdpa", "tooltip": "memory efficient attention mode"}),
             "sample_prompts": ("STRING", {"multiline": True, "default": "illustration of a kitten | photograph of a turtle", "tooltip": "validation sample prompts, for multiple prompts, separate by `|`"}),
             },
             "optional": {
                 "additional_args": ("STRING", {"multiline": True, "default": "", "tooltip": "additional args to pass to the training command"}),
                 "resume_args": ("ARGS", {"default": "", "tooltip": "resume args to pass to the training command"}),
-                "train_text_encoder": (['disabled', 'clip_l', 'clip_l_fp8', 'clip_l+T5', 'clip_l+T5_fp8'], {"default": 'disabled', "tooltip": "also train the selected text encoders using specified dtype, T5 can not be trained without clip_l"}),
-                "text_encoder_lr": ("FLOAT", {"default": 0, "min": 0.0, "max": 10.0, "step": 0.000001, "tooltip": "text encoder learning rate"}),
-                "block_args": ("ARGS", {"default": "", "tooltip": "limit the blocks used in the LoRA"}),
-                "gradient_checkpointing": (["enabled", "enabled_with_cpu_offloading", "disabled"], {"default": "enabled", "tooltip": "use gradient checkpointing"}),
+                "train_clip_l": (['disabled', 'use_gradient_dtype', 'use_fp8'], {"default": 'disabled', "tooltip": "also train the clip_l text encoder using specified dtype"}),
+                "text_encoder_lr": ("FLOAT", {"default": 0, "min": 0.0, "max": 10.0, "step": 0.00001, "tooltip": "text encoder learning rate"}),
+                "train_blocks": ("BLOCKS", ),
+                "gradient_checkpointing": ("BOOLEAN", {"default": True, "tooltip": "use gradient checkpointing"}),
             },
         }
 
@@ -339,8 +341,7 @@ class InitFluxLoRATraining:
     CATEGORY = "FluxTrainer"
 
     def init_training(self, flux_models, dataset, optimizer_settings, sample_prompts, output_name, attention_mode, 
-                      gradient_dtype, save_dtype, split_mode, additional_args=None, resume_args=None, train_text_encoder='disabled', 
-                      block_args=None, gradient_checkpointing="enabled", **kwargs,):
+                      gradient_dtype, save_dtype, split_mode, additional_args=None, resume_args=None, train_clip_l='disabled', train_blocks=None, gradient_checkpointing=True, **kwargs,):
         mm.soft_empty_cache()
         
         output_dir = os.path.abspath(kwargs.get("output_dir"))
@@ -357,8 +358,7 @@ class InitFluxLoRATraining:
 
         parser = train_network_setup_parser()
         if additional_args is not None:
-            print(f"additional_args: {additional_args}")
-            args, _ = parser.parse_known_args(args=shlex.split(additional_args))
+            args, _ = parser.parse_known_args(args=[additional_args])
         else:
             args, _ = parser.parse_known_args()
         #print(args)
@@ -405,6 +405,7 @@ class InitFluxLoRATraining:
             "persistent_data_loader_workers": False,
             "max_data_loader_n_workers": 0,
             "seed": 42,
+            #"gradient_checkpointing": True,
             "network_module": ".networks.lora_flux",
             "dataset_config": dataset_toml,
             "output_name": f"{output_name}_rank{kwargs.get('network_dim')}_{save_dtype}",
@@ -412,11 +413,11 @@ class InitFluxLoRATraining:
             "text_encoder_lr": 0,
             "t5xxl_max_token_length": 512,
             "alpha_mask": dataset["alpha_mask"],
-            "network_train_unet_only": True if train_text_encoder == 'disabled' else False,
-            "fp8_base_unet": False if "fp8" in train_text_encoder else True,
-            "disable_mmap_load_safetensors": False,
-            "split_mode": split_mode,
+            "network_train_unet_only": True if train_clip_l == 'disabled' else False,
+            "fp8_base_unet": True if train_clip_l=='use_gradient_dtype' else False,
         }
+        if gradient_checkpointing:
+            config_dict["gradient_checkpointing"] = True
         attention_settings = {
             "sdpa": {"mem_eff_attn": True, "xformers": False, "spda": True},
             "xformers": {"mem_eff_attn": True, "xformers": True, "spda": False}
@@ -429,34 +430,20 @@ class InitFluxLoRATraining:
         }
         config_dict.update(gradient_dtype_settings.get(gradient_dtype, {}))
 
-        #network args
-        additional_network_args = []
-        
-        if "T5" in train_text_encoder:
-            additional_network_args.append("train_t5xxl=True")
-        if split_mode:
-            additional_network_args.append("train_blocks=single")
-        if block_args:
-            additional_network_args.append(block_args["include"])
-        
-        # Handle network_args in args Namespace
-        if hasattr(args, 'network_args') and isinstance(args.network_args, list):
-            args.network_args.extend(additional_network_args)
+        if train_blocks is None:
+            split_mode_settings = {
+                True: {"split_mode": True, "network_args": ["train_blocks=single"]},
+                False: {"split_mode": False, "network_args": ["train_blocks=all"]}
+            }
+            config_dict.update(split_mode_settings.get(split_mode, {}))
         else:
-            setattr(args, 'network_args', additional_network_args)
+            config_dict["split_mode"] = True
+            if "network_args" not in config_dict:
+                config_dict["network_args"] = []
+            config_dict["network_args"].append(f"train_blocks={train_blocks}")
 
-        if gradient_checkpointing == "disabled":
-            config_dict["gradient_checkpointing"] = False
-        elif gradient_checkpointing == "enabled_with_cpu_offloading":
-            config_dict["cpu_offload_checkpointing"] = True
-        else:
-            config_dict["gradient_checkpointing"] = True
-
-        if flux_models["lora_path"]:
-            config_dict["network_weights"] = flux_models["lora_path"]
-
-        if dataset["reg_data_dir"]:
-            config_dict["reg_data_dir"] = dataset["reg_data_dir"]
+        print("NETWORK ARGS: ", config_dict["network_args"])
+            
 
         config_dict.update(kwargs)
         config_dict.update(optimizer_settings)
@@ -596,7 +583,6 @@ class InitFluxTraining:
             "dataset_config": dataset_toml,
             "output_name": f"{output_name}_{save_dtype}",
             "mem_eff_save": True,
-            "disable_mmap_load_safetensors": True,
 
         }
         optimizer_fusing_settings = {
@@ -740,19 +726,18 @@ class FluxTrainLoop:
             initial_global_step = network_trainer.global_step
 
             target_global_step = network_trainer.global_step + steps
-            comfy_pbar = comfy.utils.ProgressBar(steps)
-            network_trainer.comfy_pbar = comfy_pbar
+            pbar = comfy.utils.ProgressBar(steps)
             while network_trainer.global_step < target_global_step:
                 steps_done = training_loop(
                     break_at_steps = target_global_step,
                     epoch = network_trainer.current_epoch.value,
                 )
-                #pbar.update(steps_done)
+                pbar.update(steps_done)
                
                 # Also break if the global steps have reached the max train steps
                 if network_trainer.global_step >= network_trainer.args.max_train_steps:
                     break
-            
+
             trainer = {
                 "network_trainer": network_trainer,
                 "training_loop": training_loop,
@@ -793,9 +778,7 @@ class FluxTrainSave:
 
             lora_path = os.path.join(trainer.args.output_dir, ckpt_name)
             if copy_to_comfy_lora_folder:
-                destination_dir = os.path.join(folder_paths.models_dir, "loras", "flux_trainer")
-                os.makedirs(destination_dir, exist_ok=True)
-                shutil.copy(lora_path, os.path.join(destination_dir, ckpt_name))
+                shutil.copy(lora_path, os.path.join(folder_paths.models_dir, "loras", "flux_trainer", ckpt_name))
         
             
         return (network_trainer, lora_path, global_step)
@@ -909,26 +892,6 @@ class FluxTrainResume:
             
         return (resume_args, )
     
-class FluxTrainBlockSelect:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-            "include": ("STRING", {"default": "lora_unet_single_blocks_20_linear2", "multiline": True, "tooltip": "blocks to include in the LoRA network"}),
-             },
-        }
-
-    RETURN_TYPES = ("ARGS", )
-    RETURN_NAMES = ("block_args", )
-    FUNCTION = "block_select"
-    CATEGORY = "FluxTrainer"
-
-    def block_select(self, include):
-        block_args ={
-            "include": f"only_if_contains={include}",
-        }
-            
-        return (block_args, )
-    
 class FluxTrainValidationSettings:
     @classmethod
     def INPUT_TYPES(s):
@@ -1073,7 +1036,6 @@ class FluxKohyaInferenceSampler:
             "guidance_scale": ("FLOAT", {"default": 3.5, "min": 1.0, "max": 32.0, "step": 0.05, "tooltip": "guidance scale"}),
             "seed": ("INT", {"default": 42,"min": 0, "max": 0xffffffffffffffff, "step": 1}),
             "use_fp8": ("BOOLEAN", {"default": True, "tooltip": "use fp8 weights"}),
-            "apply_t5_attn_mask": ("BOOLEAN", {"default": True, "tooltip": "use t5 attention mask"}),
             "prompt": ("STRING", {"multiline": True, "default": "illustration of a kitten", "tooltip": "prompt"}),
           
             },
@@ -1084,7 +1046,7 @@ class FluxKohyaInferenceSampler:
     FUNCTION = "sample"
     CATEGORY = "FluxTrainer"
 
-    def sample(self, flux_models, lora_name, steps, width, height, guidance_scale, seed, prompt, use_fp8, lora_method, apply_t5_attn_mask):
+    def sample(self, flux_models, lora_name, steps, width, height, guidance_scale, seed, prompt, use_fp8, lora_method):
 
         from .library import flux_utils as flux_utils
         from .library import strategy_flux as strategy_flux
@@ -1097,7 +1059,7 @@ class FluxKohyaInferenceSampler:
         import gc
 
         device = "cuda"
-        
+        apply_t5_attn_mask = True
 
         if use_fp8:
             accelerator = accelerate.Accelerator(mixed_precision="bf16")
@@ -1142,7 +1104,8 @@ class FluxKohyaInferenceSampler:
         # AE
         ae = flux_utils.load_ae("dev", ae, ae_dtype, loading_device)
         ae.eval()
-
+        #if is_fp8(ae_dtype):
+        #    ae = accelerator.prepare(ae)
 
         # LoRA
         lora_models: List[lora_flux.LoRANetwork] = []
@@ -1184,7 +1147,7 @@ class FluxKohyaInferenceSampler:
                 clip_l.to(ae_dtype)
                 t5xxl.to(ae_dtype)
                 with accelerator.autocast():
-                    l_pooled, t5_out, txt_ids, t5_attn_mask = encoding_strategy.encode_tokens(
+                    _, t5_out, txt_ids, t5_attn_mask = encoding_strategy.encode_tokens(
                         tokenize_strategy, [clip_l, t5xxl], tokens_and_masks, apply_t5_attn_mask
                     )
             else:
@@ -1290,7 +1253,6 @@ class FluxKohyaInferenceSampler:
             flux_dtype: torch.dtype,
         ):
             timesteps = get_schedule(num_steps, img.shape[1], shift=not is_schnell)
-            print(timesteps)
 
             # denoise initial noise
             if accelerator:
@@ -1299,11 +1261,9 @@ class FluxKohyaInferenceSampler:
                         model, img, img_ids, t5_out, txt_ids, l_pooled, timesteps=timesteps, guidance=guidance, t5_attn_mask=t5_attn_mask
                     )
             else:
-                with torch.autocast(device_type=device.type, dtype=flux_dtype):
-                    l_pooled, _, _, _ = encoding_strategy.encode_tokens(tokenize_strategy, [clip_l, None], tokens_and_masks)
-                with torch.autocast(device_type=device.type, dtype=flux_dtype):
-                    _, t5_out, txt_ids, t5_attn_mask = encoding_strategy.encode_tokens(
-                        tokenize_strategy, [None, t5xxl], tokens_and_masks, apply_t5_attn_mask
+                with torch.autocast(device_type=device.type, dtype=flux_dtype), torch.no_grad():
+                    x = denoise(
+                        model, img, img_ids, t5_out, txt_ids, l_pooled, timesteps=timesteps, guidance=guidance, t5_attn_mask=t5_attn_mask
                     )
 
             return x
@@ -1442,7 +1402,7 @@ class ExtractFluxLoRA:
                 "finetuned_model": (folder_paths.get_filename_list("unet"), ),
                 "output_path": ("STRING", {"default": f"{str(os.path.join(folder_paths.models_dir, 'loras', 'Flux'))}"}),
                 "dim": ("INT", {"default": 4, "min": 2, "max": 1024, "step": 2, "tooltip": "LoRA rank"}),
-                "save_dtype": (["fp32", "fp16", "bf16", "fp8_e4m3fn", "fp8_e5m2"], {"default": "bf16", "tooltip": "the dtype to save the LoRA as"}),
+                "save_dtype": (["fp32", "fp16", "bf16", "fp8_e4m3fn"], {"default": "bf16", "tooltip": "the dtype to save the LoRA as"}),
                 "load_device": (["cpu", "cuda"], {"default": "cuda", "tooltip": "the device to load the model to"}),
                 "store_device": (["cpu", "cuda"], {"default": "cpu", "tooltip": "the device to store the LoRA as"}),
                 "clamp_quantile": ("FLOAT", {"default": 0.99, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "clamp quantile"}),
@@ -1495,7 +1455,7 @@ NODE_CLASS_MAPPINGS = {
     "ExtractFluxLoRA": ExtractFluxLoRA,
     "OptimizerConfigProdigy": OptimizerConfigProdigy,
     "FluxTrainResume": FluxTrainResume,
-    "FluxTrainBlockSelect": FluxTrainBlockSelect
+    "FluxLoRATrainBlocks": FluxLoRATrainBlocks
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "InitFluxLoRATraining": "Init Flux LoRA Training",
@@ -1517,5 +1477,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ExtractFluxLoRA": "Extract Flux LoRA",
     "OptimizerConfigProdigy": "Optimizer Config Prodigy",
     "FluxTrainResume": "Flux Train Resume",
-    "FluxTrainBlockSelect": "Flux Train Block Select"
+    "FluxLoRATrainBlocks": "FluxLoRATrainBlocks"
 }

requirements.txt

@@ -19,5 +19,4 @@ rich>=13.7.0
 came_pytorch
 matplotlib
 # for T5XXL tokenizer (SD3/FLUX)
-sentencepiece>=0.2.0
-protobuf
+sentencepiece>=0.2.0

train_network.py

@@ -158,9 +158,6 @@ class NetworkTrainer:
 
     # region SD/SDXL
 
-    def post_process_network(self, args, accelerator, network, text_encoders, unet):
-        pass
-
     def get_noise_scheduler(self, args: argparse.Namespace, device: torch.device) -> Any:
         noise_scheduler = DDPMScheduler(
             beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
@@ -234,19 +231,12 @@ class NetworkTrainer:
 
     def get_sai_model_spec(self, args):
         return train_util.get_sai_model_spec(None, args, self.is_sdxl, True, False)
-
-    def update_metadata(self, metadata, args):
-        pass
-
+    
     def is_text_encoder_not_needed_for_training(self, args):
         return False  # use for sample images
 
-    def prepare_text_encoder_grad_ckpt_workaround(self, index, text_encoder):
-        # set top parameter requires_grad = True for gradient checkpointing works
-        text_encoder.text_model.embeddings.requires_grad_(True)
-
-    def prepare_text_encoder_fp8(self, index, text_encoder, te_weight_dtype, weight_dtype):
-        text_encoder.text_model.embeddings.to(dtype=weight_dtype)
+    def update_metadata(self, metadata, args):
+        pass
 
     # endregion
 
@@ -328,7 +318,7 @@ class NetworkTrainer:
         collator = train_util.collator_class(current_epoch, current_step, ds_for_collator)
 
         if args.debug_dataset:
-            train_dataset_group.set_current_strategies()  # dasaset needs to know the strategies explicitly
+            train_dataset_group.set_current_strategies()
             train_util.debug_dataset(train_dataset_group)
             return
         if len(train_dataset_group) == 0:
@@ -342,7 +332,7 @@ class NetworkTrainer:
                 train_dataset_group.is_latent_cacheable()
             ), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"
 
-        self.assert_extra_args(args, train_dataset_group)  # may change some args
+        self.assert_extra_args(args, train_dataset_group)
 
         # prepare accelerator
         logger.info("preparing accelerator")
@@ -444,24 +434,17 @@ class NetworkTrainer:
             )
             args.scale_weight_norms = False
 
-        self.post_process_network(args, accelerator, network, text_encoders, unet)
-
-        # apply network to unet and text_encoder
         train_unet = not args.network_train_text_encoder_only
         train_text_encoder = self.is_train_text_encoder(args)
         network.apply_to(text_encoder, unet, train_text_encoder, train_unet)
 
         if args.network_weights is not None:
-            # FIXME consider alpha of weights: this assumes that the alpha is not changed
+            # FIXME consider alpha of weights
             info = network.load_weights(args.network_weights)
             accelerator.print(f"load network weights from {args.network_weights}: {info}")
 
         if args.gradient_checkpointing:
-            if args.cpu_offload_checkpointing:
-                unet.enable_gradient_checkpointing(cpu_offload=True)
-            else:
-                unet.enable_gradient_checkpointing()
-
+            unet.enable_gradient_checkpointing()
             for t_enc, flag in zip(text_encoders, self.get_text_encoders_train_flags(args, text_encoders)):
                 if flag:
                     if t_enc.supports_gradient_checkpointing:
@@ -555,12 +538,11 @@ class NetworkTrainer:
                 args.mixed_precision != "no"
             ), "fp8_base requires mixed precision='fp16' or 'bf16'"
             accelerator.print("enable fp8 training for U-Net.")
-            unet_weight_dtype = torch.float8_e4m3fn if args.fp8_dtype == "e4m3" else torch.float8_e5m2
-            accelerator.print(f"unet_weight_dtype: {unet_weight_dtype}")
+            unet_weight_dtype = torch.float8_e4m3fn
 
             if not args.fp8_base_unet and not args.network_train_unet_only:
                 accelerator.print("enable fp8 training for Text Encoder.")
-            te_weight_dtype = torch.float8_e4m3fn if args.fp8_dtype == "e4m3" else torch.float8_e5m2
+            te_weight_dtype = weight_dtype if args.fp8_base_unet else torch.float8_e4m3fn
 
             # unet.to(accelerator.device)  # this makes faster `to(dtype)` below, but consumes 23 GB VRAM
             # unet.to(dtype=unet_weight_dtype)  # without moving to gpu, this takes a lot of time and main memory
@@ -569,16 +551,17 @@ class NetworkTrainer:
 
         unet.requires_grad_(False)
         unet.to(dtype=unet_weight_dtype)
-        for i, t_enc in enumerate(text_encoders):
+        for t_enc in text_encoders:
             t_enc.requires_grad_(False)
 
             # in case of cpu, dtype is already set to fp32 because cpu does not support fp8/fp16/bf16
             if t_enc.device.type != "cpu":
                 t_enc.to(dtype=te_weight_dtype)
-
-                # nn.Embedding not support FP8
-                if te_weight_dtype != weight_dtype:
-                    self.prepare_text_encoder_fp8(i, t_enc, te_weight_dtype, weight_dtype)
+                if hasattr(t_enc, "text_model") and hasattr(t_enc.text_model, "embeddings"):
+                    # nn.Embedding not support FP8
+                    t_enc.text_model.embeddings.to(dtype=(weight_dtype if te_weight_dtype != weight_dtype else te_weight_dtype))
+                elif hasattr(t_enc, "encoder") and hasattr(t_enc.encoder, "embeddings"):
+                    t_enc.encoder.embeddings.to(dtype=(weight_dtype if te_weight_dtype != weight_dtype else te_weight_dtype))
 
         # acceleratorがなんかよろしくやってくれるらしい / accelerator will do something good
         if args.deepspeed:
@@ -619,12 +602,12 @@ class NetworkTrainer:
         if args.gradient_checkpointing:
             # according to TI example in Diffusers, train is required
             unet.train()
-            for i, (t_enc, frag) in enumerate(zip(text_encoders, self.get_text_encoders_train_flags(args, text_encoders))):
+            for t_enc, frag in zip(text_encoders, self.get_text_encoders_train_flags(args, text_encoders)):
                 t_enc.train()
 
                 # set top parameter requires_grad = True for gradient checkpointing works
                 if frag:
-                    self.prepare_text_encoder_grad_ckpt_workaround(i, t_enc)
+                    t_enc.text_model.embeddings.requires_grad_(True)
 
         else:
             unet.eval()
@@ -1049,12 +1032,8 @@ class NetworkTrainer:
 
         # log device and dtype for each model
         logger.info(f"unet dtype: {unet_weight_dtype}, device: {unet.device}")
-        for i, t_enc in enumerate(text_encoders):
-            params_itr = t_enc.parameters()
-            params_itr.__next__()  # skip the first parameter
-            params_itr.__next__()  # skip the second parameter. because CLIP first two parameters are embeddings
-            param_3rd = params_itr.__next__()
-            logger.info(f"text_encoder [{i}] dtype: {param_3rd.dtype}, device: {t_enc.device}")
+        for t_enc in text_encoders:
+            logger.info(f"text_encoder dtype: {t_enc.dtype}, device: {t_enc.device}")
 
         clean_memory_on_device(accelerator.device)
 
@@ -1075,13 +1054,10 @@ class NetworkTrainer:
         self.lr_scheduler = lr_scheduler
         self.save_model = save_model
         self.remove_model = remove_model
-        self.comfy_pbar = None
         
         progress_bar = tqdm(range(args.max_train_steps - initial_step), smoothing=0, disable=False, desc="steps")
-        
         def training_loop(break_at_steps, epoch):
             steps_done = 0
-            
             #accelerator.print(f"\nepoch {epoch+1}/{num_train_epochs}")
             progress_bar.set_description(f"Epoch {epoch + 1}/{num_train_epochs} - steps")
 
@@ -1128,11 +1104,15 @@ class NetworkTrainer:
                         # print(f"set multiplier: {multipliers}")
                         accelerator.unwrap_model(network).set_multiplier(multipliers)
 
-                    text_encoder_conds = []
                     text_encoder_outputs_list = batch.get("text_encoder_outputs_list", None)
                     if text_encoder_outputs_list is not None:
                         text_encoder_conds = text_encoder_outputs_list  # List of text encoder outputs
-                    if len(text_encoder_conds) == 0 or text_encoder_conds[0] is None or train_text_encoder:
+                    if (
+                        text_encoder_conds is None
+                        or len(text_encoder_conds) == 0
+                        or text_encoder_conds[0] is None
+                        or train_text_encoder
+                    ):
                         with torch.set_grad_enabled(train_text_encoder), accelerator.autocast():
                             # Get the text embedding for conditioning
                             if args.weighted_captions:
@@ -1155,14 +1135,10 @@ class NetworkTrainer:
                                 if args.full_fp16:
                                     encoded_text_encoder_conds = [c.to(weight_dtype) for c in encoded_text_encoder_conds]
 
-                        # if text_encoder_conds is not cached, use encoded_text_encoder_conds
-                        if len(text_encoder_conds) == 0:
-                            text_encoder_conds = encoded_text_encoder_conds
-                        else:
-                            # if encoded_text_encoder_conds is not None, update cached text_encoder_conds
-                            for i in range(len(encoded_text_encoder_conds)):
-                                if encoded_text_encoder_conds[i] is not None:
-                                    text_encoder_conds[i] = encoded_text_encoder_conds[i]
+                        # if encoded_text_encoder_conds is not None, update cached text_encoder_conds
+                        for i in range(len(encoded_text_encoder_conds)):
+                            if encoded_text_encoder_conds[i] is not None:
+                                text_encoder_conds[i] = encoded_text_encoder_conds[i]
 
                     # sample noise, call unet, get target
                     noise_pred, target, timesteps, huber_c, weighting = self.get_noise_pred_and_target(
@@ -1237,7 +1213,6 @@ class NetworkTrainer:
                 if self.global_step >= break_at_steps:
                     break
                 steps_done += 1
-                self.comfy_pbar.update(1)
 
             if args.logging_dir is not None:
                 logs = {"loss/epoch": self.loss_recorder.moving_average}
@@ -1291,12 +1266,6 @@ def setup_parser() -> argparse.ArgumentParser:
 
     parser.add_argument("--unet_lr", type=float, default=None, help="learning rate for U-Net / U-Netの学習率")
     parser.add_argument("--text_encoder_lr", type=float, default=None, help="learning rate for Text Encoder / Text Encoderの学習率")
-    parser.add_argument(
-        "--fp8_base_unet",
-        action="store_true",
-        help="use fp8 for U-Net (or DiT), Text Encoder is fp16 or bf16"
-        " / U-Net（またはDiT）にfp8を使用する。Text Encoderはfp16またはbf16",
-    )
 
     parser.add_argument(
         "--network_weights", type=str, default=None, help="pretrained weights for network / 学習するネットワークの初期重み"
@@ -1393,9 +1362,10 @@ def setup_parser() -> argparse.ArgumentParser:
         + " / 初期ステップ数、全エポックを含むステップ数、0で最初のステップ（未指定時と同じ）。initial_epochを上書きする",
     )
     parser.add_argument(
-        "--cpu_offload_checkpointing",
+        "--fp8_base_unet",
         action="store_true",
-        help="[EXPERIMENTAL] enable offloading of tensors to CPU during checkpointing for U-Net or DiT, if supported",
+        help="use fp8 for U-Net (or DiT), Text Encoder is fp16 or bf16"
+        " / U-Net（またはDiT）にfp8を使用する。Text Encoderはfp16またはbf16",
     )
     # parser.add_argument("--loraplus_lr_ratio", default=None, type=float, help="LoRA+ learning rate ratio")
     # parser.add_argument("--loraplus_unet_lr_ratio", default=None, type=float, help="LoRA+ UNet learning rate ratio")