bump requirements, add optims, ramtorch

Author: 67372a

flux_train_network.py

@@ -142,13 +142,24 @@ def load_target_model(self, args, weight_dtype, accelerator):
 
         ae = flux_utils.load_ae(args.ae, weight_dtype, "cpu", disable_mmap=args.disable_mmap_load_safetensors)
 
- 
         if args.use_ramtorch:
             logger.info("Applying RamTorch to FLUX models (DiT, T5-XXL, CLIP-L, AE).")
-            model = replace_linear_with_ramtorch(model, accelerator.device)
-            clip_l = replace_linear_with_ramtorch(clip_l, accelerator.device)
-            t5xxl = replace_linear_with_ramtorch(t5xxl, accelerator.device)
-            ae = replace_linear_with_ramtorch(ae, accelerator.device)
+            if isinstance(model, torch.nn.Module):
+                model = replace_linear_with_ramtorch(model, accelerator.device)
+                logger.info("RamTorch applied to FLUX DiT.")
+
+            if isinstance(clip_l, torch.nn.Module):
+                clip_l = replace_linear_with_ramtorch(clip_l, accelerator.device)
+                logger.info("RamTorch applied to FLUX Clip-L.")
+
+            if isinstance(t5xxl, torch.nn.Module):
+                t5xxl = replace_linear_with_ramtorch(t5xxl, accelerator.device)
+                logger.info("RamTorch applied to FLUX t5xxl.")
+
+            if isinstance(ae, torch.nn.Module):
+                ae = replace_linear_with_ramtorch(ae, accelerator.device)
+                logger.info("RamTorch applied to FLUX AE/autoencoder.")
+
 
         return flux_utils.MODEL_VERSION_FLUX_V1, [clip_l, t5xxl], ae, model
 

requirements.txt

@@ -1,38 +1,42 @@
-accelerate~=1.10.1
+accelerate~=1.11.0
 transformers~=4.50.3
-diffusers~=0.35.1
+diffusers~=0.35.2
+ftfy~=6.3.1
 opencv-python~=4.12.0.88
 einops~=0.8.1
 pytorch-lightning~=2.5.5
-bitsandbytes~=0.47.0
+bitsandbytes~=0.48.2
 prodigyopt~=1.1.2
 lion-pytorch~=0.2.3
 schedulefree~=1.4.1
 tensorboard~=2.20.0
 safetensors~=0.6.2
 toml~=0.10.2
 voluptuous~=0.15.2
-huggingface-hub[hf_xet]~=0.35.1
-matplotlib~=3.10.6
+huggingface-hub[hf_xet]~=0.36.0
+matplotlib~=3.10.7
 scipy~=1.16.2
 imagesize~=1.4.1
-rich~=14.1.0
+numpy~=2.2.6
+rich~=14.2.0
 sentencepiece~=0.2.1
 dadaptation~=3.2
-wandb~=0.22.0
+wandb~=0.22.3
 came-pytorch~=0.1.3
+torchastic~=0.1.2
+torch-optimi~=0.3.2
 pytorch_optimizer==3.8.2 # Has applied breaking changes in patches, thus pinned
 wheel~=0.45.1
 ninja~=1.13.0
 cmake~=4.1.0
 setuptools~=80.9.0
 triton~=3.3.1; platform_system == "Linux"
-triton-windows~=3.3.0.post21; platform_system == "Windows"
+triton-windows~=3.3.1.post21; platform_system == "Windows"
 prodigy-plus-schedule-free~=2.0.1
-kornia~=0.8.1
+kornia~=0.8.2
 PyWavelets~=1.9.0
-hf-xet~=1.1.10
-RamTorch~=0.2.1
-adv_optm~=1.1.4
+hf-xet~=1.2.0
+RamTorch~=0.2.2
+adv_optm~=1.2.3
 # for kohya_ss library
 -e .

sd3_train_network.py

@@ -8,6 +8,7 @@
 from accelerate import Accelerator
 from library import sd3_models, strategy_sd3, utils
 from library.device_utils import init_ipex, clean_memory_on_device
+from ramtorch.helpers import replace_linear_with_ramtorch
 
 init_ipex()
 
@@ -130,6 +131,24 @@ def load_target_model(self, args, weight_dtype, accelerator):
             args.vae, weight_dtype, "cpu", disable_mmap=args.disable_mmap_load_safetensors, state_dict=state_dict
         )
 
+        if args.use_ramtorch:
+            logger.info("Applying RamTorch to SD3 mmdit, VAE, and Text Encoders (clip l, clip g, t5xxl).")
+            if isinstance(mmdit, torch.nn.Module):
+                mmdit = replace_linear_with_ramtorch(mmdit, accelerator.device)
+                logger.info("RamTorch applied to SD3 mmdit.")
+
+            if isinstance(clip_l, torch.nn.Module):
+                clip_l = replace_linear_with_ramtorch(clip_l, accelerator.device)
+                logger.info("RamTorch applied to SD3 Clip-L.")
+
+            if isinstance(clip_g, torch.nn.Module):
+                clip_g = replace_linear_with_ramtorch(clip_g, accelerator.device)
+                logger.info("RamTorch applied to SD3 Clip-G.")
+
+            if isinstance(t5xxl, torch.nn.Module):
+                t5xxl = replace_linear_with_ramtorch(t5xxl, accelerator.device)
+                logger.info("RamTorch applied to SD3 t5xxl.")
+
         return mmdit.model_type, [clip_l, clip_g, t5xxl], vae, mmdit
 
     def get_tokenize_strategy(self, args):

sdxl_train_network.py

@@ -57,10 +57,21 @@ def load_target_model(self, args, weight_dtype, accelerator):
 
         if args.use_ramtorch:
             logger.info("Applying RamTorch to SDXL UNet, VAE, and Text Encoders.")
-            unet = replace_linear_with_ramtorch(unet, accelerator.device)
-            vae = replace_linear_with_ramtorch(vae, accelerator.device)
-            text_encoder1 = replace_linear_with_ramtorch(text_encoder1, accelerator.device)
-            text_encoder2 = replace_linear_with_ramtorch(text_encoder2, accelerator.device)
+            if isinstance(unet, torch.nn.Module):
+                unet = replace_linear_with_ramtorch(unet, accelerator.device)
+                logger.info("RamTorch applied to SDXL unet.")
+
+            if isinstance(vae, torch.nn.Module):
+                vae = replace_linear_with_ramtorch(vae, accelerator.device)
+                logger.info("RamTorch applied to SDXL vae.")
+
+            if isinstance(text_encoder1, torch.nn.Module):
+                text_encoder1 = replace_linear_with_ramtorch(text_encoder1, accelerator.device)
+                logger.info("RamTorch applied to SDXL Clip-L.")
+
+            if isinstance(text_encoder2, torch.nn.Module):
+                text_encoder2 = replace_linear_with_ramtorch(text_encoder2, accelerator.device)
+                logger.info("RamTorch applied to SDXL Clip-G.")
 
         # モデルに xformers とか memory efficient attention を組み込む
         train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers, args.sdpa)

train_network.py

@@ -31,6 +31,7 @@
 from accelerate import Accelerator
 from diffusers import DDPMScheduler
 from library import deepspeed_utils, model_util, strategy_base, strategy_sd
+from ramtorch.helpers import replace_linear_with_ramtorch
 
 import library.train_util as train_util
 from library.train_util import DreamBoothDataset
@@ -311,6 +312,20 @@ def assert_extra_args(self, args, train_dataset_group):
     def load_target_model(self, args, weight_dtype, accelerator):
         text_encoder, vae, unet, _ = train_util.load_target_model(args, weight_dtype, accelerator)
 
+        if args.use_ramtorch:
+            logger.info("Applying RamTorch to SD UNet, VAE, and Clip-L.")
+            if isinstance(unet, torch.nn.Module):
+                unet = replace_linear_with_ramtorch(unet, accelerator.device)
+                logger.info("RamTorch applied to SD unet.")
+
+            if isinstance(text_encoder, torch.nn.Module):
+                text_encoder = replace_linear_with_ramtorch(text_encoder, accelerator.device)
+                logger.info("RamTorch applied to SD Clip-L.")
+
+            if isinstance(vae, torch.nn.Module):
+                vae = replace_linear_with_ramtorch(vae, accelerator.device)
+                logger.info("RamTorch applied to SD VAE.")
+
         # モデルに xformers とか memory efficient attention を組み込む
         train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers, args.sdpa)
         if torch.__version__ >= "2.0.0":  # PyTorch 2.0.0 以上対応のxformersなら以下が使える