add boft support in stable-diffusion

Makefile

@@ -80,7 +80,7 @@ slow_tests_custom_file_input: test_installs
 # Run single-card non-regression tests
 slow_tests_1x: test_installs
 	python -m pytest tests/test_examples.py -v -s -k "single_card"
-	python -m pip install peft==0.10.0
+	python -m pip install peft==0.12.0
 	python -m pytest tests/test_peft_inference.py
 	python -m pytest tests/test_pipeline.py
 
@@ -96,7 +96,7 @@ slow_tests_deepspeed: test_installs
 slow_tests_diffusers: test_installs
 	python -m pytest tests/test_diffusers.py -v -s -k "test_no_"
 	python -m pytest tests/test_diffusers.py -v -s -k "test_textual_inversion"
-	python -m pip install peft==0.7.0
+	python -m pip install peft==0.12.0
 	python -m pytest tests/test_diffusers.py -v -s -k "test_train_text_to_image_"
 	python -m pytest tests/test_diffusers.py -v -s -k "test_train_controlnet"
 	python -m pytest tests/test_diffusers.py -v -s -k "test_deterministic_image_generation"

examples/stable-diffusion/text_to_image_generation.py

@@ -464,19 +464,50 @@ def main():
                 )
 
                 if args.unet_adapter_name_or_path is not None:
-                    from peft import PeftModel
+                    from peft import PeftModel, tuners
+                    from peft.utils import PeftType
+
+                    from optimum.habana.peft.layer import GaudiBoftGetDeltaWeight
+
+                    tuners.boft.layer.Linear.get_delta_weight = GaudiBoftGetDeltaWeight
+                    tuners.boft.layer.Conv2d.get_delta_weight = GaudiBoftGetDeltaWeight
+                    tuners.boft.layer._FBD_CUDA = False
 
                     pipeline.unet = PeftModel.from_pretrained(pipeline.unet, args.unet_adapter_name_or_path)
-                    pipeline.unet = pipeline.unet.merge_and_unload()
+                    if pipeline.unet.peft_type in [PeftType.OFT, PeftType.BOFT]:
+                        # WA torch.inverse issue in Synapse AI 1.17 for oft and boft
+                        if args.bf16:
+                            pipeline.unet = pipeline.unet.to(torch.float32)
+                        pipeline.unet = pipeline.unet.merge_and_unload()
+                        if args.bf16:
+                            pipeline.unet = pipeline.unet.to(torch.bfloat16)
+                    else:
+                        with torch.autocast(device_type="hpu", dtype=torch.bfloat16, enabled=args.bf16):
+                            pipeline.unet = pipeline.unet.merge_and_unload()
 
                 if args.text_encoder_adapter_name_or_path is not None:
-                    from peft import PeftModel
+                    from peft import PeftModel, tuners
+                    from peft.utils import PeftType
+
+                    from optimum.habana.peft.layer import GaudiBoftGetDeltaWeight
+
+                    tuners.boft.layer.Linear.get_delta_weight = GaudiBoftGetDeltaWeight
+                    tuners.boft.layer.Conv2d.get_delta_weight = GaudiBoftGetDeltaWeight
+                    tuners.boft.layer._FBD_CUDA = False
 
                     pipeline.text_encoder = PeftModel.from_pretrained(
                         pipeline.text_encoder, args.text_encoder_adapter_name_or_path
                     )
-                    pipeline.text_encoder = pipeline.text_encoder.merge_and_unload()
-
+                    if pipeline.text_encoder.peft_type in [PeftType.OFT, PeftType.BOFT]:
+                        # WA torch.inverse issue in Synapse AI 1.17 for oft and boft
+                        if args.bf16:
+                            pipeline.text_encoder = pipeline.text_encoder.to(torch.float32)
+                        pipeline.text_encoder = pipeline.text_encoder.merge_and_unload()
+                        if args.bf16:
+                            pipeline.text_encoder = pipeline.text_encoder.to(torch.bfloat16)
+                    else:
+                        with torch.autocast(device_type="hpu", dtype=torch.bfloat16, enabled=args.bf16):
+                            pipeline.text_encoder = pipeline.text_encoder.merge_and_unload()
             else:
                 # SD LDM3D use-case
                 from optimum.habana.diffusers import GaudiStableDiffusionLDM3DPipeline as GaudiStableDiffusionPipeline

examples/stable-diffusion/training/README.md

@@ -319,10 +319,10 @@ Prior-preservation is used to avoid overfitting and language-drift. Refer to the
 According to the paper, it's recommended to generate `num_epochs * num_samples` images for prior-preservation. 200-300 works well for most cases. The `num_class_images` flag sets the number of images to generate with the class prompt. You can place existing images in `class_data_dir`, and the training script will generate any additional images so that `num_class_images` are present in `class_data_dir` during training time.
 
 ### PEFT model finetune
-We provide example for dreambooth to use lora/lokr/loha/oft to finetune unet or text encoder.
+We provide example for dreambooth to use `lora`, `lokr`, `loha`, `oft` and `boft` to finetune unet or text encoder.
 
-**___Note: When using peft method we can use a much higher learning rate compared to vanilla dreambooth. Here we
-use *1e-4* instead of the usual *5e-6*.___**
+> [!NOTE]
+> When using peft method we can use a much higher learning rate compared to vanilla dreambooth. Here we use *1e-4* instead of the usual *5e-6*.
 
 Launch the multi-card training using:
 ```bash
@@ -355,19 +355,18 @@ python ../../gaudi_spawn.py --world_size 8 --use_mpi train_dreambooth.py \
   lora --unet_r 8 --unet_alpha 8
 
 ```
-Similar command could be applied to loha, lokr, oft.
+Similar command could be applied to `loha`, `lokr`, `oft` or `boft`.
 You could check each adapter specific args by "--help", like you could use following command to check oft specific args.
 
 ```bash
-python3 train_dreambooth.py oft --help
+python train_dreambooth.py oft --help
 
 ```
 
-**___Note: oft could not work with hpu graphs mode. since "torch.inverse" need to fallback to cpu.
-there's error like "cpu fallback is not supported during hpu graph capturing"___**
-
+> [!NOTE]
+> `boft` and `oft` do not work with hpu graphs mode since `torch.inverse` `torch.linalg.solve` need to fallback to cpu. Pls. remove `--use_hpu_graphs_for_training` and `--use_hpu_graphs_for_inference` to avoid `cpu fallback is not supported during hpu graph capturing` error.
 
-You could use text_to_image_generation.py to generate picture using the peft adapter like
+You could use `text_to_image_generation.py` script to generate picture using the peft adapter:
 
 ```bash
 python ../text_to_image_generation.py \
@@ -384,8 +383,8 @@ python ../text_to_image_generation.py \
 ```
 
 ### DreamBooth training example for Stable Diffusion XL
-You could use the dog images as example as well.
-You can launch training using:
+You could use the dog images as example for training:
+
 ```bash
 export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0"
 export INSTANCE_DIR="dog"
@@ -415,7 +414,7 @@ python ../../gaudi_spawn.py --world_size 8 --use_mpi train_dreambooth_lora_sdxl.
 
 ```
 
-You could use text_to_image_generation.py to generate picture using the peft adapter like
+Then, you could use `text_to_image_generation.py` to generate picture using the peft adapter:
 
 ```bash
 python ../text_to_image_generation.py \

examples/stable-diffusion/training/requirements.txt

@@ -1,2 +1,2 @@
 imagesize
-peft == 0.10.0
+peft == 0.12.0

examples/stable-diffusion/training/train_dreambooth.py

@@ -52,7 +52,7 @@
 from diffusers.utils.torch_utils import is_compiled_module
 from habana_frameworks.torch.hpu import memory_stats
 from huggingface_hub import HfApi
-from peft import LoHaConfig, LoKrConfig, LoraConfig, OFTConfig, get_peft_model
+from peft import BOFTConfig, LoHaConfig, LoKrConfig, LoraConfig, OFTConfig, get_peft_model, tuners
 from PIL import Image
 from torch.utils.data import Dataset
 from torchvision import transforms
@@ -108,7 +108,9 @@ def import_model_class_from_model_name_or_path(pretrained_model_name_or_path: st
         raise ValueError(f"{model_class} is not supported.")
 
 
-def create_unet_adapter_config(args: argparse.Namespace) -> Union[LoraConfig, LoHaConfig, LoKrConfig, OFTConfig]:
+def create_unet_adapter_config(
+    args: argparse.Namespace,
+) -> Union[LoraConfig, LoHaConfig, LoKrConfig, OFTConfig, BOFTConfig]:
     if args.adapter == "full":
         raise ValueError("Cannot create unet adapter config for full parameter")
 
@@ -152,6 +154,21 @@ def create_unet_adapter_config(args: argparse.Namespace) -> Union[LoraConfig, Lo
             coft=args.unet_use_coft,
             eps=args.unet_eps,
         )
+    elif args.adapter == "boft":
+        config = BOFTConfig(
+            boft_block_size=args.unet_block_size,
+            boft_block_num=args.unet_block_num,
+            boft_n_butterfly_factor=args.unet_n_butterfly_factor,
+            target_modules=UNET_TARGET_MODULES,
+            boft_dropout=args.unet_dropout,
+            bias=args.unet_bias,
+        )
+        from optimum.habana.peft.layer import GaudiBoftConv2dForward, GaudiBoftLinearForward
+
+        tuners.boft.layer.Linear.forward = GaudiBoftLinearForward
+        tuners.boft.layer.Conv2d.forward = GaudiBoftConv2dForward
+        tuners.boft.layer._FBD_CUDA = False
+
     else:
         raise ValueError(f"Unknown adapter type {args.adapter}")
 
@@ -160,7 +177,7 @@ def create_unet_adapter_config(args: argparse.Namespace) -> Union[LoraConfig, Lo
 
 def create_text_encoder_adapter_config(
     args: argparse.Namespace,
-) -> Union[LoraConfig, LoHaConfig, LoKrConfig, OFTConfig]:
+) -> Union[LoraConfig, LoHaConfig, LoKrConfig, OFTConfig, BOFTConfig]:
     if args.adapter == "full":
         raise ValueError("Cannot create text_encoder adapter config for full parameter")
 
@@ -202,6 +219,20 @@ def create_text_encoder_adapter_config(
             coft=args.te_use_coft,
             eps=args.te_eps,
         )
+    elif args.adapter == "boft":
+        config = BOFTConfig(
+            boft_block_size=args.te_block_size,
+            boft_block_num=args.te_block_num,
+            boft_n_butterfly_factor=args.te_n_butterfly_factor,
+            target_modules=TEXT_ENCODER_TARGET_MODULES,
+            boft_dropout=args.te_dropout,
+            bias=args.te_bias,
+        )
+        from optimum.habana.peft.layer import GaudiBoftConv2dForward, GaudiBoftLinearForward
+
+        tuners.boft.layer.Linear.forward = GaudiBoftLinearForward
+        tuners.boft.layer.Conv2d.forward = GaudiBoftConv2dForward
+        tuners.boft.layer._FBD_CUDA = False
     else:
         raise ValueError(f"Unknown adapter type {args.adapter}")
 
@@ -632,6 +663,44 @@ def parse_args(input_args=None):
         help="The control strength of COFT for text_encoder, only used if `train_text_encoder` is True",
     )
 
+    # boft adapter
+    boft = subparsers.add_parser("boft", help="Use Boft adapter")
+    boft.add_argument("--unet_block_size", type=int, default=8, help="Boft block_size for unet")
+    boft.add_argument("--unet_block_num", type=int, default=0, help="Boft block_num for unet")
+    boft.add_argument("--unet_n_butterfly_factor", type=int, default=1, help="Boft n_butterfly_factor for unet")
+    boft.add_argument("--unet_dropout", type=float, default=0.1, help="Boft dropout for unet")
+    boft.add_argument("--unet_bias", type=str, default="boft_only", help="Boft bias for unet")
+    boft.add_argument(
+        "--te_block_size",
+        type=int,
+        default=8,
+        help="Boft block_size for text_encoder,only used if `train_text_encoder` is True",
+    )
+    boft.add_argument(
+        "--te_block_num",
+        type=int,
+        default=0,
+        help="Boft block_num for text_encoder,only used if `train_text_encoder` is True",
+    )
+    boft.add_argument(
+        "--te_n_butterfly_factor",
+        type=int,
+        default=1,
+        help="Boft n_butterfly_factor for text_encoder,only used if `train_text_encoder` is True",
+    )
+    boft.add_argument(
+        "--te_dropout",
+        type=float,
+        default=0.1,
+        help="Boft dropout for text_encoder,only used if `train_text_encoder` is True",
+    )
+    boft.add_argument(
+        "--te_bias",
+        type=str,
+        default="boft_only",
+        help="Boft bias for text_encoder, only used if `train_text_encoder` is True",
+    )
+
     if input_args is not None:
         args = parser.parse_args(input_args)
     else:

optimum/habana/peft/__init__.py

@@ -2,6 +2,9 @@
     GaudiAdaloraLayerSVDLinearForward,
     GaudiAdaptedAttention_getattr,
     GaudiAdaptedAttentionPreAttnForward,
+    GaudiBoftConv2dForward,
+    GaudiBoftGetDeltaWeight,
+    GaudiBoftLinearForward,
     GaudiPolyLayerLinearForward,
 )
 from .peft_model import gaudi_generate, gaudi_prepare_inputs_for_generation