add efficient tunner modules

modelscope/metainfo.py

@@ -203,6 +203,7 @@ class Models(object):
     vldoc = 'vldoc'
     hitea = 'hitea'
     soonet = 'soonet'
+    efficient_diffusion_tuning = 'efficient-diffusion-tuning'
 
     # science models
     unifold = 'unifold'
@@ -510,6 +511,7 @@ class Pipelines(object):
     gridvlp_multi_modal_classification = 'gridvlp-multi-modal-classification'
     gridvlp_multi_modal_embedding = 'gridvlp-multi-modal-embedding'
     soonet_video_temporal_grounding = 'soonet-video-temporal-grounding'
+    efficient_diffusion_tuning = 'efficient-diffusion-tuning'
 
     # science tasks
     protein_structure = 'unifold-protein-structure'
@@ -884,6 +886,7 @@ class MultiModalTrainers(object):
     ofa = 'ofa'
     mplug = 'mplug'
     mgeo_ranking_trainer = 'mgeo-ranking-trainer'
+    efficient_diffusion_tuning = 'efficient-diffusion-tuning'
 
 
 class AudioTrainers(object):
@@ -1028,6 +1031,7 @@ class Preprocessors(object):
     mgeo_ranking = 'mgeo-ranking'
     vldoc_preprocessor = 'vldoc-preprocessor'
     hitea_tasks_preprocessor = 'hitea-tasks-preprocessor'
+    diffusion_image_generation_preprocessor = 'diffusion-image-generation-preprocessor'
 
     # science preprocessor
     unifold_preprocessor = 'unifold-preprocessor'

modelscope/metrics/builder.py

@@ -75,6 +75,7 @@ task_default_metrics = {
     [Metrics.image_quality_assessment_mos_metric],
     Tasks.bad_image_detecting: [Metrics.accuracy],
     Tasks.ocr_recognition: [Metrics.ocr_recognition_metric],
+    Tasks.efficient_diffusion_tuning: [Metrics.loss_metric],
 }
 
 

modelscope/models/multi_modal/__init__.py

@@ -19,6 +19,7 @@ if TYPE_CHECKING:
         MultiStageDiffusionForTextToImageSynthesis
     from .vldoc import VLDocForDocVLEmbedding
     from .video_synthesis import TextToVideoSynthesis
+    from .efficient_diffusion_tuning import EfficientStableDiffusion
 
 else:
     _import_structure = {
@@ -36,6 +37,7 @@ else:
         ['MultiStageDiffusionForTextToImageSynthesis'],
         'vldoc': ['VLDocForDocVLEmbedding'],
         'video_synthesis': ['TextToVideoSynthesis'],
+        'efficient_diffusion_tuning': ['EfficientStableDiffusion']
     }
 
     import sys

modelscope/models/multi_modal/efficient_diffusion_tuning/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import TYPE_CHECKING
+
+from modelscope.utils.import_utils import LazyImportModule
+
+if TYPE_CHECKING:
+
+    from .efficient_stable_diffusion import EfficientStableDiffusion
+
+else:
+    _import_structure = {
+        'efficient_stable_diffusion': ['EfficientStableDiffusion'],
+    }
+
+    import sys
+
+    sys.modules[__name__] = LazyImportModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

modelscope/models/multi_modal/efficient_diffusion_tuning/efficient_stable_diffusion.py

@@ -0,0 +1,247 @@
+# Copyright 2023-2024 The Alibaba Fundamental Vision Team Authors. All rights reserved.
+# The implementation is adopted from HighCWu,
+# made pubicly available under the Apache License 2.0 License at https://github.com/HighCWu/ControlLoRA
+import os
+import os.path as osp
+from functools import partial
+from typing import Any, Callable, List, Mapping, Optional, Union
+
+import torch
+import torch.nn.functional as F
+from diffusers import (AutoencoderKL, DDPMScheduler, DiffusionPipeline,
+                       DPMSolverMultistepScheduler, UNet2DConditionModel,
+                       utils)
+from diffusers.models import cross_attention
+from diffusers.utils import deprecation_utils
+from transformers import CLIPTextModel, CLIPTokenizer
+
+from modelscope.metainfo import Models
+from modelscope.models import TorchModel
+from modelscope.models.builder import MODELS
+from modelscope.outputs import OutputKeys
+from modelscope.tuners.control_sd_lora import ControlLoRATuner
+from modelscope.tuners.sd_lora import LoRATuner
+from modelscope.utils.checkpoint import save_checkpoint, save_configuration
+from modelscope.utils.config import Config
+from modelscope.utils.constant import ModelFile, Tasks
+
+utils.deprecate = lambda *arg, **kwargs: None
+deprecation_utils.deprecate = lambda *arg, **kwargs: None
+cross_attention.deprecate = lambda *arg, **kwargs: None
+
+__tuner_MAP__ = {'lora': LoRATuner, 'control_lora': ControlLoRATuner}
+
+
+@MODELS.register_module(
+    Tasks.efficient_diffusion_tuning,
+    module_name=Models.efficient_diffusion_tuning)
+class EfficientStableDiffusion(TorchModel):
+    """ The implementation of efficient diffusion tuning model based on TorchModel.
+
+    This model is constructed with the implementation of stable diffusion model. If you want to
+    finetune lightweight parameters on your own dataset, you can define you own tuner module
+    and load in this cls.
+    """
+
+    def __init__(self, model_dir, *args, **kwargs):
+        """ Initialize a vision efficient diffusion tuning model.
+
+        Args:
+          model_dir: model id or path, where model_dir/pytorch_model.bin
+        """
+        super().__init__(model_dir, *args, **kwargs)
+        tuner_name = kwargs.pop('tuner_name', 'lora')
+        pretrained_model_name_or_path = kwargs.pop(
+            'pretrained_model_name_or_path', 'runwayml/stable-diffusion-v1-5')
+        tuner_config = kwargs.pop('tuner_config', None)
+        pretrained_tuner = kwargs.pop('pretrained_tuner', None)
+        revision = kwargs.pop('revision', None)
+        inference = kwargs.pop('inference', True)
+
+        if pretrained_tuner is not None:
+            pretrained_tuner = osp.join(model_dir, pretrained_tuner)
+
+        self.weight_dtype = torch.float32
+        self.inference = inference
+
+        self.device = torch.device(
+            'cuda' if torch.cuda.is_available() else 'cpu')
+
+        if self.inference:
+            self.pipe = DiffusionPipeline.from_pretrained(
+                pretrained_model_name_or_path,
+                revision=revision,
+                torch_dtype=self.weight_dtype,
+                safety_checker=None)
+            self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
+                self.pipe.scheduler.config)
+            self.pipe = self.pipe.to(self.device)
+            self.unet = self.pipe.unet
+        else:
+            # Load scheduler, tokenizer and models.
+            self.noise_scheduler = DDPMScheduler.from_pretrained(
+                pretrained_model_name_or_path, subfolder='scheduler')
+            self.tokenizer = CLIPTokenizer.from_pretrained(
+                pretrained_model_name_or_path,
+                subfolder='tokenizer',
+                revision=revision)
+            self.text_encoder = CLIPTextModel.from_pretrained(
+                pretrained_model_name_or_path,
+                subfolder='text_encoder',
+                revision=revision)
+            self.vae = AutoencoderKL.from_pretrained(
+                pretrained_model_name_or_path,
+                subfolder='vae',
+                revision=revision)
+            self.unet = UNet2DConditionModel.from_pretrained(
+                pretrained_model_name_or_path,
+                subfolder='unet',
+                revision=revision)
+            self.unet.requires_grad_(False)
+            self.vae.requires_grad_(False)
+            self.text_encoder.requires_grad_(False)
+        self.is_control = tuner_name.startswith('control_')
+        self.tuner_name = tuner_name
+        if tuner_name in ('lora', 'control_lora'):
+            # if not set the config of control-tuner, we add the lora tuner directly to the original framework,
+            # otherwise the control side network is also added.
+            tuner_cls = __tuner_MAP__[tuner_name]
+            tuner = tuner_cls.tune(
+                self,
+                tuner_config=osp.join(model_dir, tuner_config),
+                pretrained_tuner=pretrained_tuner)
+            self.tuner = tuner
+
+    def train(self, mode: bool = True):
+        self.training = mode
+        if hasattr(self, 'tuner'):
+            self.tuner.train(mode=mode)
+
+    def load_state_dict(self,
+                        state_dict: Mapping[str, Any],
+                        strict: bool = True):
+        if hasattr(self, 'tuner'):
+            self.tuner.load_state_dict(state_dict=state_dict, strict=strict)
+        else:
+            return super().load_state_dict(
+                state_dict=state_dict, strict=strict)
+
+    def state_dict(self):
+        if hasattr(self, 'tuner'):
+            return self.tuner.state_dict()
+        else:
+            return super().state_dict()
+
+    def tokenize_caption(self, captions):
+        """ Convert caption text to token data.
+
+        Args:
+          captions: a batch of texts.
+        Returns: token's data as tensor.
+        """
+        inputs = self.tokenizer(
+            captions,
+            max_length=self.tokenizer.model_max_length,
+            padding='max_length',
+            truncation=True,
+            return_tensors='pt')
+        return inputs.input_ids
+
+    def forward(self, prompt='', cond=None, target=None):
+        if self.inference:
+            generator = torch.Generator(device=self.device).manual_seed(0)
+            if self.is_control:
+                _ = self.tuner(cond.to(self.device)).control_states
+            images = self.pipe(
+                prompt, num_inference_steps=30, generator=generator).images
+            return images
+        else:
+            with torch.no_grad():
+                latents = self.vae.encode(
+                    target.to(dtype=self.weight_dtype)).latent_dist.sample()
+            latents = latents * self.vae.config.scaling_factor
+            # Sample noise that we'll add to the latents
+            noise = torch.randn_like(latents)
+            bsz = latents.shape[0]
+            # Sample a random timestep for each image
+            timesteps = torch.randint(
+                0,
+                self.noise_scheduler.num_train_timesteps, (bsz, ),
+                device=latents.device)
+            timesteps = timesteps.long()
+
+            # Add noise to the latents according to the noise magnitude at each timestep
+            # (this is the forward diffusion process)
+            noisy_latents = self.noise_scheduler.add_noise(
+                latents, noise, timesteps)
+
+            input_ids = self.tokenize_caption(prompt).to(self.device)
+
+            # Get the text embedding for conditioning
+            with torch.no_grad():
+                encoder_hidden_states = self.text_encoder(input_ids)[0]
+
+            # Inject control states to unet
+            if self.is_control:
+                _ = self.tuner(cond.to(dtype=self.weight_dtype)).control_states
+            # else:
+            #     tune_weights_list = self.tuner()
+
+            # Get the target for loss depending on the prediction type
+            if self.noise_scheduler.config.prediction_type == 'epsilon':
+                target = noise
+            elif self.noise_scheduler.config.prediction_type == 'v_prediction':
+                target = self.noise_scheduler.get_velocity(
+                    latents, noise, timesteps)
+            else:
+                raise ValueError(
+                    f'Unknown prediction type {self.noise_scheduler.config.prediction_type}'
+                )
+
+            # Predict the noise residual and compute loss
+            model_pred = self.unet(noisy_latents, timesteps,
+                                   encoder_hidden_states).sample
+            loss = F.mse_loss(
+                model_pred.float(), target.float(), reduction='mean')
+            output = {OutputKeys.LOSS: loss}
+            return output
+
+    def parameters(self, recurse: bool = True):
+        if hasattr(self, 'tuner'):
+            return self.tuner.parameters(recurse=recurse)
+        else:
+            return super().parameters(recurse=recurse)
+
+    def save_pretrained(self,
+                        target_folder: Union[str, os.PathLike],
+                        save_checkpoint_names: Union[str, List[str]] = None,
+                        save_function: Callable = partial(
+                            save_checkpoint, with_meta=False),
+                        config: Optional[dict] = None,
+                        save_config_function: Callable = save_configuration,
+                        **kwargs):
+
+        if config is None and hasattr(self, 'cfg'):
+            config = self.cfg
+
+        config['model']['inference'] = True
+        super().save_pretrained(target_folder, save_checkpoint_names,
+                                save_function, config, save_config_function,
+                                **kwargs)
+
+    @classmethod
+    def _instantiate(cls, model_dir, **kwargs):
+        config = Config.from_file(osp.join(model_dir, ModelFile.CONFIGURATION))
+        for k, v in kwargs.items():
+            config.model[k] = v
+
+        model = EfficientStableDiffusion(
+            model_dir,
+            pretrained_model_name_or_path=config.model.
+            pretrained_model_name_or_path,
+            tuner_name=config.model.tuner_name,
+            tuner_config=config.model.tuner_config,
+            pretrained_tuner=config.model.get('pretrained_tuner', None),
+            inference=config.model.get('inference', False))
+        model.config = config
+        return model

modelscope/msdatasets/ms_dataset.py

@@ -594,7 +594,8 @@ class MsDataset:
         columns = [
             key for key in self._hf_ds.features.keys() if key in columns
         ]
-        retained_columns = []
+        retained_numeric_columns = []
+        retained_unumeric_columns = []
         if to_tensor:
             sample = next(iter(self._hf_ds))
 
@@ -612,20 +613,23 @@ class MsDataset:
                 if not is_numpy_number(sample_res[k]):
                     logger.warning(
                         f'Data of column {k} is non-numeric, will be removed')
+                    retained_unumeric_columns.append(k)
                     continue
-                retained_columns.append(k)
+                retained_numeric_columns.append(k)
 
         import torch
 
         class MsMapDataset(torch.utils.data.Dataset):
 
             def __init__(self, dataset: Iterable, preprocessor_list,
-                         retained_columns, columns, to_tensor):
+                         retained_numeric_columns, retained_unumeric_columns,
+                         columns, to_tensor):
                 super(MsDataset).__init__()
                 self.dataset = dataset
                 self.preprocessor_list = preprocessor_list
                 self.to_tensor = to_tensor
-                self.retained_columns = retained_columns
+                self.retained_numeric_columns = retained_numeric_columns
+                self.retained_unumeric_columns = retained_unumeric_columns
                 self.columns = columns
 
             def __len__(self):
@@ -641,19 +645,21 @@ class MsDataset:
                 item_dict = self.dataset[index]
                 res = {
                     k: self.type_converter(item_dict[k])
-                    for k in self.columns
-                    if (not self.to_tensor) or k in self.retained_columns
+                    for k in self.columns if (not self.to_tensor)
+                    or k in self.retained_numeric_columns
                 }
                 for preprocessor in self.preprocessor_list:
-                    res.update({
-                        k: self.type_converter(v)
-                        for k, v in preprocessor(item_dict).items()
-                        if (not self.to_tensor) or k in self.retained_columns
-                    })
+                    for k, v in preprocessor(item_dict).items():
+                        if (not self.to_tensor) or \
+                                k in self.retained_numeric_columns:
+                            res[k] = self.type_converter(v)
+                        elif k in self.retained_unumeric_columns:
+                            res[k] = v
                 return res
 
-        return MsMapDataset(self._hf_ds, preprocessor_list, retained_columns,
-                            columns, to_tensor)
+        return MsMapDataset(self._hf_ds, preprocessor_list,
+                            retained_numeric_columns,
+                            retained_unumeric_columns, columns, to_tensor)
 
     def _to_tf_dataset_with_processors(
         self,

modelscope/pipelines/multi_modal/efficient_diffusion_tuning_pipeline.py

@@ -0,0 +1,77 @@
+# Copyright 2022-2023 The Alibaba Fundamental Vision Team Authors. All rights reserved.
+from typing import Any, Dict
+
+import cv2
+import numpy as np
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+
+from modelscope.metainfo import Pipelines
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines.base import Input, Pipeline
+from modelscope.pipelines.builder import PIPELINES
+from modelscope.preprocessors import LoadImage
+from modelscope.utils.constant import Tasks
+from modelscope.utils.logger import get_logger
+
+logger = get_logger()
+
+
+@PIPELINES.register_module(
+    Tasks.efficient_diffusion_tuning,
+    module_name=Pipelines.efficient_diffusion_tuning)
+class EfficientDiffusionTuningPipeline(Pipeline):
+
+    def __init__(self, model: str, **kwargs):
+        """
+        use `model` to create a diffusion efficient tuning pipeline for prediction
+        Args:
+            model: model id on modelscope hub.
+        Example:
+            >>> from modelscope.pipelines import pipeline
+            >>> petl_pipeline = pipeline('efficient-diffusion-tuning',
+                'damo/cv_vitb16_classification_vision-efficient-tuning-adapter')
+            >>> result = petl_pipeline(
+                'data/test/images/vision_efficient_tuning_test_1.png')
+            >>> print(f'Output: {result}.')
+        """
+        super().__init__(model=model, **kwargs)
+
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.model = self.model.to(self.device)
+        self.model.eval()
+        self.preprocessor = transforms.Compose([
+            transforms.Resize(
+                512, interpolation=transforms.InterpolationMode.BILINEAR),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5]),
+        ])
+
+    def preprocess(self, inputs: Input, **preprocess_params) -> Dict[str, Any]:
+        """ Preprocess method build from transforms or Preprocessor """
+        assert isinstance(inputs, dict)
+        result = {}
+        if 'cond' in inputs:
+            img = LoadImage.convert_to_img(inputs['cond'])
+            data = self.preprocessor(img)
+            result['cond'] = data.unsqueeze(0).to(self.device)
+        if 'prompt' in inputs:
+            result['prompt'] = inputs['prompt']
+        return result
+
+    def forward(self, inputs: Dict[str, Any],
+                **forward_params) -> Dict[str, Any]:
+        with torch.no_grad():
+            results = self.model(**inputs)
+            return results
+
+    def postprocess(self, inputs: Dict[str, Any],
+                    **post_params) -> Dict[str, Any]:
+        images = []
+        for idx, img in enumerate(inputs):
+            if isinstance(img, Image.Image):
+                img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+                images.append(img)
+                cv2.imwrite(f'{self.model.tuner_name}_{idx}.jpg', img)
+        return {OutputKeys.OUTPUT_IMGS: images}

modelscope/preprocessors/__init__.py

@@ -18,7 +18,8 @@ if TYPE_CHECKING:
                      ControllableImageGenerationPreprocessor)
     from .kws import WavToLists
     from .tts import KanttsDataPreprocessor
-    from .multi_modal import (OfaPreprocessor, MPlugPreprocessor,
+    from .multi_modal import (DiffusionImageGenerationPreprocessor,
+                              OfaPreprocessor, MPlugPreprocessor,
                               HiTeAPreprocessor)
     from .nlp import (
         DocumentSegmentationTransformersPreprocessor,
@@ -67,8 +68,10 @@ else:
         ],
         'kws': ['WavToLists'],
         'tts': ['KanttsDataPreprocessor'],
-        'multi_modal':
-        ['OfaPreprocessor', 'MPlugPreprocessor', 'HiTeAPreprocessor'],
+        'multi_modal': [
+            'DiffusionImageGenerationPreprocessor', 'OfaPreprocessor',
+            'MPlugPreprocessor', 'HiTeAPreprocessor'
+        ],
         'nlp': [
             'DocumentSegmentationTransformersPreprocessor',
             'FaqQuestionAnsweringTransformersPreprocessor',

modelscope/preprocessors/cv/__init__.py

@@ -29,7 +29,9 @@ else:
         'controllable_image_generation':
         ['ControllableImageGenerationPreprocessor'],
         'image_classification_preprocessor':
-        ['ImageClassificationPreprocessor']
+        ['ImageClassificationPreprocessor'],
+        'diffusion_image_generation_preprocessor':
+        ['DiffusionImageGenerationPreprocessor']
     }
 
     import sys

modelscope/preprocessors/multi_modal.py

@@ -9,6 +9,7 @@ import numpy as np
 import torch
 from PIL import Image
 from timm.data import create_transform
+from torchvision import transforms
 from torchvision.transforms import Compose, Normalize, Resize, ToTensor
 
 from modelscope.hub.snapshot_download import snapshot_download
@@ -26,7 +27,48 @@ from .ofa import *  # noqa
 from .ofa.utils.collate import collate_fn
 from .ofa.utils.constant import OFA_TASK_KEY_MAPPING
 
-__all__ = ['OfaPreprocessor', 'MPlugPreprocessor', 'HiTeAPreprocessor']
+__all__ = [
+    'DiffusionImageGenerationPreprocessor', 'OfaPreprocessor',
+    'MPlugPreprocessor', 'HiTeAPreprocessor'
+]
+
+
+@PREPROCESSORS.register_module(
+    Fields.multi_modal,
+    module_name=Preprocessors.diffusion_image_generation_preprocessor)
+class DiffusionImageGenerationPreprocessor(Preprocessor):
+    """ Preprocessor the data with the combination of image and text.
+        Args:
+            data: process the value as an image for keys ending with 'FILE'
+                or existing in preprocessor_image_keys and pass-through the values of other keys.
+
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.preprocessor_resolution = kwargs.pop('resolution', 512)
+        self.preprocessor_mean = kwargs.pop('mean', [0.5, 0.5, 0.5])
+        self.preprocessor_std = kwargs.pop('std', [0.5, 0.5, 0.5])
+        self.preprocessor_image_keys = set(kwargs.pop('image_keys', []))
+        self.transform_input = transforms.Compose([
+            transforms.Resize(
+                self.preprocessor_resolution,
+                interpolation=transforms.InterpolationMode.BILINEAR),
+            transforms.ToTensor(),
+            transforms.Normalize(self.preprocessor_mean,
+                                 self.preprocessor_std),
+        ])
+
+    def __call__(self, data) -> Dict[str, Any]:
+        results = {}
+        for key, value in data.items():
+            if key.endswith(':FILE') or key in self.preprocessor_image_keys:
+                image = load_image(value)
+                img = self.transform_input(image)
+                results[key.replace(':FILE', '').lower()] = img
+            else:
+                results[key.lower()] = value
+        return results
 
 
 @PREPROCESSORS.register_module(

modelscope/trainers/multi_modal/efficient_diffusion_tuning/efficient_diffusion_tuning_trainer.py

@@ -0,0 +1,73 @@
+# Copyright 2022-2023 The Alibaba Fundamental Vision Team Authors. All rights reserved.
+from typing import Union
+
+import torch
+from torch import nn
+
+from modelscope.metainfo import Trainers
+from modelscope.models.base import Model, TorchModel
+from modelscope.trainers.builder import TRAINERS
+from modelscope.trainers.optimizer.builder import build_optimizer
+from modelscope.trainers.trainer import EpochBasedTrainer
+from modelscope.utils.config import Config, ConfigDict
+
+
+@TRAINERS.register_module(module_name=Trainers.efficient_diffusion_tuning)
+class EfficientDiffusionTuningTrainer(EpochBasedTrainer):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def build_model(self) -> Union[nn.Module, TorchModel]:
+        """ Instantiate a pytorch model and return.
+
+        By default, we will create a model using config from configuration file. You can
+        override this method in a subclass.
+
+        """
+        model = Model.from_pretrained(self.model_dir, cfg_dict=self.cfg)
+        if not isinstance(model, nn.Module) and hasattr(model, 'model'):
+            return model.model
+        elif isinstance(model, nn.Module):
+            return model
+
+    def build_optimizer(self, cfg: ConfigDict, default_args: dict = None):
+        try:
+            return build_optimizer(
+                self.model.tuner, cfg=cfg, default_args=default_args)
+        except KeyError as e:
+            self.logger.error(
+                f'Build optimizer error, the optimizer {cfg} is a torch native component, '
+                f'please check if your torch with version: {torch.__version__} matches the config.'
+            )
+            raise e
+
+    def train(self, *args, **kwargs):
+        self.print_model_params_status()
+        super().train(*args, **kwargs)
+
+    def evaluate(self, *args, **kwargs):
+        eval_res = super().evaluate(*args, **kwargs)
+        return eval_res
+
+    def print_model_params_status(self, model=None, logger=None):
+        """Print the status and parameters of the model"""
+        if model is None:
+            model = self.model
+        if logger is None:
+            logger = self.logger
+        train_param_dict = {}
+        all_param_numel = 0
+        for key, val in model.named_parameters():
+            if val.requires_grad:
+                sub_key = '.'.join(key.split('.', 1)[-1].split('.', 2)[:2])
+                if sub_key in train_param_dict:
+                    train_param_dict[sub_key] += val.numel()
+                else:
+                    train_param_dict[sub_key] = val.numel()
+            all_param_numel += val.numel()
+        train_param_numel = sum(train_param_dict.values())
+        logger.info(
+            f'Load trainable params {train_param_numel} / {all_param_numel} = '
+            f'{train_param_numel/all_param_numel:.2%}, '
+            f'train part: {train_param_dict}.')

modelscope/trainers/trainer.py

@@ -87,6 +87,7 @@ class EpochBasedTrainer(BaseTrainer):
         compile (bool, optional): Compile the model with torch 2.0, default False
         compile_options (dict, optional): The compile options if compile=True,
             default None to use the default params of 'TorchModel.compile'.
+        efficient_tuners (dict, optional): The tuners to use to train the model
 
         Examples of cfg_modify_fn:
             >>> def cfg_modify_fn(cfg):
@@ -113,6 +114,7 @@ class EpochBasedTrainer(BaseTrainer):
             model_revision: Optional[str] = DEFAULT_MODEL_REVISION,
             seed: int = 42,
             callbacks: Optional[List[Hook]] = None,
+            efficient_tuners: List[Dict] = None,
             **kwargs):
 
         self._seed = seed
@@ -216,6 +218,7 @@ class EpochBasedTrainer(BaseTrainer):
         self.use_fp16 = kwargs.get('use_fp16', False)
         self.launcher = kwargs.get('launcher')
         self.device = kwargs.get('device')
+        self.tune_module(efficient_tuners)
 
         # The parallel_groups field will be initialized in the hooks' after_init stage.
         # Please check the DDPHook and MegatronHook for details.
@@ -259,6 +262,14 @@ class EpochBasedTrainer(BaseTrainer):
 
         self.print_cfg()
 
+    def tune_module(self, efficient_tuners):
+        if efficient_tuners is not None:
+            for tuner in efficient_tuners:
+                type = tuner.pop('type')
+                if type == 'lora':
+                    from modelscope.tuners.lora import LoRATuner
+                    LoRATuner.tune(self.model, **tuner)
+
     def place_model(self):
         """Place model to device, or to DDP
         """

modelscope/tuners/control_sd_lora.py

@@ -0,0 +1,912 @@
+# Copyright 2023-2024 The Alibaba Fundamental Vision Team Authors. All rights reserved.
+# The implementation is adopted from HighCWu,
+# made pubicly available under the Apache License 2.0 License at https://github.com/HighCWu/ControlLoRA
+
+import os
+from dataclasses import dataclass
+from typing import List, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.cross_attention import CrossAttention, LoRALinearLayer
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.resnet import (Downsample2D, Mish, Upsample2D,
+                                     downsample_2d, partial, upsample_2d)
+from diffusers.models.unet_2d_blocks import \
+    get_down_block as get_down_block_default
+from diffusers.utils.outputs import BaseOutput
+
+from .sd_lora import LoRACrossAttnProcessor
+
+
+@dataclass
+class ControlLoRAOutput(BaseOutput):
+    control_states: Tuple[torch.FloatTensor]
+
+
+class ControlLoRATuner(ModelMixin, ConfigMixin):
+    """ The implementation of control lora module.
+    This module conduct encoding operation for control-condition and use lora to perform efficient tuning.
+    """
+
+    @staticmethod
+    def tune(
+        model: nn.Module,
+        tuner_config=None,
+        pretrained_tuner=None,
+    ):
+        tuner = ControlLoRATuner.from_config(tuner_config)
+        if pretrained_tuner is not None and os.path.exists(pretrained_tuner):
+            tuner.load_state_dict(
+                torch.load(pretrained_tuner, map_location='cpu'), strict=True)
+
+        tune_layers_list = list(
+            [list(layer_list) for layer_list in tuner.lora_layers])
+
+        assert hasattr(model, 'unet')
+        unet = model.unet
+        tuner.to(unet.device)
+        tune_attn_procs = tuner.set_tune_layers(unet, tune_layers_list)
+        unet.set_attn_processor(tune_attn_procs)
+        return tuner
+
+    def set_tune_layers(self, unet, tune_layers_list):
+        n_ch = len(unet.config.block_out_channels)
+        control_ids = [i for i in range(n_ch)]
+        tune_attn_procs = {}
+
+        for name in unet.attn_processors.keys():
+            if name.startswith('mid_block'):
+                control_id = control_ids[-1]
+            elif name.startswith('up_blocks'):
+                block_id = int(name[len('up_blocks.')])
+                control_id = list(reversed(control_ids))[block_id]
+            elif name.startswith('down_blocks'):
+                block_id = int(name[len('down_blocks.')])
+                control_id = control_ids[block_id]
+
+            tune_layers = tune_layers_list[control_id]
+            if len(tune_layers) != 0:
+                tune_layer = tune_layers.pop(0)
+                tune_attn_procs[name] = tune_layer
+        return tune_attn_procs
+
+    @register_to_config
+    def __init__(self,
+                 in_channels: int = 3,
+                 down_block_types: Tuple[str] = (
+                     'SimpleDownEncoderBlock2D',
+                     'SimpleDownEncoderBlock2D',
+                     'SimpleDownEncoderBlock2D',
+                     'SimpleDownEncoderBlock2D',
+                 ),
+                 block_out_channels: Tuple[int] = (32, 64, 128, 256),
+                 layers_per_block: int = 1,
+                 act_fn: str = 'silu',
+                 norm_num_groups: int = 32,
+                 lora_pre_down_block_types: Tuple[str] = (
+                     None,
+                     'SimpleDownEncoderBlock2D',
+                     'SimpleDownEncoderBlock2D',
+                     'SimpleDownEncoderBlock2D',
+                 ),
+                 lora_pre_down_layers_per_block: int = 1,
+                 lora_pre_conv_skipped: bool = False,
+                 lora_pre_conv_types: Tuple[str] = (
+                     'SimpleDownEncoderBlock2D',
+                     'SimpleDownEncoderBlock2D',
+                     'SimpleDownEncoderBlock2D',
+                     'SimpleDownEncoderBlock2D',
+                 ),
+                 lora_pre_conv_layers_per_block: int = 1,
+                 lora_pre_conv_layers_kernel_size: int = 1,
+                 lora_block_in_channels: Tuple[int] = (256, 256, 256, 256),
+                 lora_block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
+                 lora_cross_attention_dims: Tuple[List[int]] = ([
+                     None, 768, None, 768, None, 768, None, 768, None, 768
+                 ], [None, 768, None, 768, None, 768, None, 768, None, 768], [
+                     None, 768, None, 768, None, 768, None, 768, None, 768
+                 ], [None, 768]),
+                 lora_rank: int = 4,
+                 lora_control_rank: int = None,
+                 lora_post_add: bool = False,
+                 lora_concat_hidden: bool = False,
+                 lora_control_channels: Tuple[int] = (None, None, None, None),
+                 lora_control_self_add: bool = True,
+                 lora_key_states_skipped: bool = False,
+                 lora_value_states_skipped: bool = False,
+                 lora_output_states_skipped: bool = False,
+                 lora_control_version: int = 1):
+        """ Initialize a control lora module instance.
+               Args:
+                   in_channels (`int`): The number of channels for input conditional data.
+                   down_block_types (Tuple[str], *optional*):
+                       The down block types for conditional data's downsample operation.
+                   block_out_channels (Tuple[int],  *optional*, defaults to (32, 64, 128, 256)):
+                       The number of channels for every down-block.
+                   layers_per_block (`int`,  *optional*, defaults to 1):
+                       The number of layers of every block.
+                   act_fn (`str`, *optional*, defaults to silu):
+                       The activation function.
+                   norm_num_groups (`int`, *optional*, defaults to 32):
+                       The number of groups for norm operation.
+                   lora_pre_down_block_types (Tuple[str], *optional*):
+                       The block'types for pre down-block.
+                   lora_pre_down_layers_per_block (`int`, *optional*, defaults to 1)
+                       The number of layers of every pre down-block block.
+                   lora_pre_conv_skipped ('bool', *optional*, defaults to False )
+                       Set to True to skip conv in pre downsample.
+                   lora_pre_conv_types (Tuple[str], *optional*):
+                       The block'types for pre conv.
+                   lora_pre_conv_layers_per_block (`int`, *optional*, defaults to 1)
+                       The number of layers of every pre conv block.
+                   lora_pre_conv_layers_kernel_size (`int`, *optional*, defaults to 1)
+                       The conv kernel size of pre conv block.
+                   lora_block_in_channels (Tuple[int],  *optional*, defaults to (256, 256, 256, 256)):
+                       The number of input channels for lora block.
+                   lora_block_out_channels (Tuple[int],  *optional*, defaults to (256, 256, 256, 256)):
+                       The number of output channels for lora block.
+                   lora_rank (int,  *optional*, defaults to 4):
+                       The rank of lora block.
+                   lora_control_rank (int,  *optional*, defaults to 4):
+                       The rank of lora block.
+                   lora_post_add (`bool`,  *optional*, defaults to False):
+                        Set to `True`, conduct weighted adding operation after lora.
+                   lora_concat_hidden (`bool`,  *optional*, defaults to False):
+                        Set to `True`, conduct concat operation for hidden embedding.
+                   lora_control_channels  (Tuple[int],  *optional*, defaults to (None, None, None, None)):
+                        The number of control channels.
+                   lora_control_self_add (`bool`,  *optional*, defaults to True):
+                        Set to `True` to perform self attn add.
+                   lora_key_states_skipped (`bool`, *optional*, defaults to False):
+                        Set to `True` for skip to perform lora on key value.
+                    value_states_skipped (`bool`, *optional*, defaults to False):
+                        Set to `True` for skip to perform lora on value.
+                    output_states_skipped (`bool`, *optional*, defaults to False):
+                        Set to `True` for skip to perform lora on output value.
+                    lora_control_version (int,  *optional*, defaults to 1):
+                        Use lora attn version: ControlLoRACrossAttnProcessor vs ControlLoRACrossAttnProcessorV2.
+               """
+
+        super().__init__()
+        lora_control_cls = ControlLoRACrossAttnProcessor
+        if lora_control_version == 2:
+            lora_control_cls = ControlLoRACrossAttnProcessorV2
+
+        assert lora_block_in_channels[0] == block_out_channels[-1]
+
+        if lora_pre_conv_skipped:
+            lora_control_channels = lora_block_in_channels
+            lora_control_self_add = False
+
+        self.layers_per_block = layers_per_block
+        self.lora_pre_down_layers_per_block = lora_pre_down_layers_per_block
+        self.lora_pre_conv_layers_per_block = lora_pre_conv_layers_per_block
+
+        self.conv_in = torch.nn.Conv2d(
+            in_channels,
+            block_out_channels[0],
+            kernel_size=3,
+            stride=1,
+            padding=1)
+
+        self.down_blocks = nn.ModuleList([])
+        self.pre_lora_layers = nn.ModuleList([])
+        self.lora_layers = nn.ModuleList([])
+
+        # pre_down
+        pre_down_blocks = []
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            pre_down_block = get_down_block(
+                down_block_type,
+                num_layers=self.layers_per_block,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                add_downsample=not is_final_block,
+                resnet_eps=1e-6,
+                downsample_padding=0,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                attn_num_head_channels=None,
+                temb_channels=None,
+            )
+            pre_down_blocks.append(pre_down_block)
+        self.down_blocks.append(nn.Sequential(*pre_down_blocks))
+        self.pre_lora_layers.append(
+            get_down_block(
+                lora_pre_conv_types[0],
+                num_layers=self.lora_pre_conv_layers_per_block,
+                in_channels=lora_block_in_channels[0],
+                out_channels=(
+                    lora_block_out_channels[0] if lora_control_channels[0] is
+                    None else lora_control_channels[0]),
+                add_downsample=False,
+                resnet_eps=1e-6,
+                downsample_padding=0,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                attn_num_head_channels=None,
+                temb_channels=None,
+                resnet_kernel_size=lora_pre_conv_layers_kernel_size,
+            ) if not lora_pre_conv_skipped else nn.Identity())
+        self.lora_layers.append(
+            nn.ModuleList([
+                lora_control_cls(
+                    lora_block_out_channels[0],
+                    cross_attention_dim=cross_attention_dim,
+                    rank=lora_rank,
+                    control_rank=lora_control_rank,
+                    post_add=lora_post_add,
+                    concat_hidden=lora_concat_hidden,
+                    control_channels=lora_control_channels[0],
+                    control_self_add=lora_control_self_add,
+                    key_states_skipped=lora_key_states_skipped,
+                    value_states_skipped=lora_value_states_skipped,
+                    output_states_skipped=lora_output_states_skipped)
+                for cross_attention_dim in lora_cross_attention_dims[0]
+            ]))
+
+        # down
+        output_channel = lora_block_in_channels[0]
+        for i, down_block_type in enumerate(lora_pre_down_block_types):
+            if i == 0:
+                continue
+            input_channel = output_channel
+            output_channel = lora_block_in_channels[i]
+
+            down_block = get_down_block(
+                down_block_type,
+                num_layers=self.lora_pre_down_layers_per_block,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                add_downsample=True,
+                resnet_eps=1e-6,
+                downsample_padding=0,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                attn_num_head_channels=None,
+                temb_channels=None,
+            )
+            self.down_blocks.append(down_block)
+
+            self.pre_lora_layers.append(
+                get_down_block(
+                    lora_pre_conv_types[i],
+                    num_layers=self.lora_pre_conv_layers_per_block,
+                    in_channels=output_channel,
+                    out_channels=(
+                        lora_block_out_channels[i] if lora_control_channels[i]
+                        is None else lora_control_channels[i]),
+                    add_downsample=False,
+                    resnet_eps=1e-6,
+                    downsample_padding=0,
+                    resnet_act_fn=act_fn,
+                    resnet_groups=norm_num_groups,
+                    attn_num_head_channels=None,
+                    temb_channels=None,
+                    resnet_kernel_size=lora_pre_conv_layers_kernel_size,
+                ) if not lora_pre_conv_skipped else nn.Identity())
+            self.lora_layers.append(
+                nn.ModuleList([
+                    lora_control_cls(
+                        lora_block_out_channels[i],
+                        cross_attention_dim=cross_attention_dim,
+                        rank=lora_rank,
+                        control_rank=lora_control_rank,
+                        post_add=lora_post_add,
+                        concat_hidden=lora_concat_hidden,
+                        control_channels=lora_control_channels[i],
+                        control_self_add=lora_control_self_add,
+                        key_states_skipped=lora_key_states_skipped,
+                        value_states_skipped=lora_value_states_skipped,
+                        output_states_skipped=lora_output_states_skipped)
+                    for cross_attention_dim in lora_cross_attention_dims[i]
+                ]))
+
+    def forward(self,
+                x: torch.FloatTensor,
+                return_dict: bool = True) -> Union[ControlLoRAOutput, Tuple]:
+        lora_layer: ControlLoRACrossAttnProcessor
+
+        orig_dtype = x.dtype
+        dtype = self.conv_in.weight.dtype
+
+        h = x.to(dtype)
+        h = self.conv_in(h)
+        control_states_list = []
+
+        # down
+        for down_block, pre_lora_layer, lora_layer_list in zip(
+                self.down_blocks, self.pre_lora_layers, self.lora_layers):
+            h = down_block(h)
+            control_states = pre_lora_layer(h)
+            if isinstance(control_states, tuple):
+                control_states = control_states[0]
+            control_states = control_states.to(orig_dtype)
+            for lora_layer in lora_layer_list:
+                lora_layer.inject_control_states(control_states)
+            control_states_list.append(control_states)
+
+        if not return_dict:
+            return tuple(control_states_list)
+
+        return ControlLoRAOutput(control_states=tuple(control_states_list))
+
+
+def get_down_block(
+    down_block_type,
+    num_layers,
+    in_channels,
+    out_channels,
+    temb_channels,
+    add_downsample,
+    resnet_eps,
+    resnet_act_fn,
+    attn_num_head_channels,
+    resnet_groups=None,
+    cross_attention_dim=None,
+    downsample_padding=None,
+    dual_cross_attention=False,
+    use_linear_projection=False,
+    only_cross_attention=False,
+    upcast_attention=False,
+    resnet_time_scale_shift='default',
+    resnet_kernel_size=3,
+):
+    down_block_type = down_block_type[7:] if down_block_type.startswith(
+        'UNetRes') else down_block_type
+    if down_block_type == 'SimpleDownEncoderBlock2D':
+        return SimpleDownEncoderBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            add_downsample=add_downsample,
+            convnet_eps=resnet_eps,
+            convnet_act_fn=resnet_act_fn,
+            convnet_groups=resnet_groups,
+            downsample_padding=downsample_padding,
+            convnet_time_scale_shift=resnet_time_scale_shift,
+            convnet_kernel_size=resnet_kernel_size)
+    else:
+        return get_down_block_default(
+            down_block_type,
+            num_layers,
+            in_channels,
+            out_channels,
+            temb_channels,
+            add_downsample,
+            resnet_eps,
+            resnet_act_fn,
+            attn_num_head_channels,
+            resnet_groups=resnet_groups,
+            cross_attention_dim=cross_attention_dim,
+            downsample_padding=downsample_padding,
+            dual_cross_attention=dual_cross_attention,
+            use_linear_projection=use_linear_projection,
+            only_cross_attention=only_cross_attention,
+            upcast_attention=upcast_attention,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            # resnet_kernel_size=resnet_kernel_size
+        )
+
+
+class ControlLoRACrossAttnProcessor(LoRACrossAttnProcessor):
+
+    def __init__(self,
+                 hidden_size,
+                 cross_attention_dim=None,
+                 rank=4,
+                 control_rank=None,
+                 post_add=False,
+                 concat_hidden=False,
+                 control_channels=None,
+                 control_self_add=True,
+                 key_states_skipped=False,
+                 value_states_skipped=False,
+                 output_states_skipped=False,
+                 **kwargs):
+        super().__init__(
+            hidden_size,
+            cross_attention_dim,
+            rank,
+            post_add=post_add,
+            key_states_skipped=key_states_skipped,
+            value_states_skipped=value_states_skipped,
+            output_states_skipped=output_states_skipped)
+
+        control_rank = rank if control_rank is None else control_rank
+        control_channels = hidden_size if control_channels is None else control_channels
+        self.concat_hidden = concat_hidden
+        self.control_self_add = control_self_add if control_channels is None else False
+        self.control_states: torch.Tensor = None
+
+        self.to_control = LoRALinearLayer(
+            control_channels + (hidden_size if concat_hidden else 0),
+            hidden_size, control_rank)
+        self.pre_loras: List[LoRACrossAttnProcessor] = []
+        self.post_loras: List[LoRACrossAttnProcessor] = []
+
+    def inject_pre_lora(self, lora_layer):
+        self.pre_loras.append(lora_layer)
+
+    def inject_post_lora(self, lora_layer):
+        self.post_loras.append(lora_layer)
+
+    def inject_control_states(self, control_states):
+        self.control_states = control_states
+
+    def process_control_states(self, hidden_states, scale=1.0):
+        control_states = self.control_states.to(hidden_states.dtype)
+        if hidden_states.ndim == 3 and control_states.ndim == 4:
+            batch, _, height, width = control_states.shape
+            control_states = control_states.permute(0, 2, 3, 1).reshape(
+                batch, height * width, -1)
+            self.control_states = control_states
+        _control_states = control_states
+        if self.concat_hidden:
+            b1, b2 = control_states.shape[0], hidden_states.shape[0]
+            if b1 != b2:
+                control_states = control_states[:, None].repeat(
+                    1, b2 // b1, *([1] * (len(control_states.shape) - 1)))
+                control_states = control_states.view(-1,
+                                                     *control_states.shape[2:])
+            _control_states = torch.cat([hidden_states, control_states], -1)
+        _control_states = scale * self.to_control(_control_states)
+        if self.control_self_add:
+            control_states = control_states + _control_states
+        else:
+            control_states = _control_states
+
+        return control_states
+
+    def __call__(self,
+                 attn: CrossAttention,
+                 hidden_states,
+                 encoder_hidden_states=None,
+                 attention_mask=None,
+                 scale=1.0):
+        pre_lora: LoRACrossAttnProcessor
+        post_lora: LoRACrossAttnProcessor
+        assert self.control_states is not None
+
+        batch_size, sequence_length, _ = hidden_states.shape
+        attention_mask = attn.prepare_attention_mask(attention_mask,
+                                                     sequence_length)
+        query = attn.to_q(hidden_states)
+        for pre_lora in self.pre_loras:
+            lora_in = query if pre_lora.post_add else hidden_states
+            if isinstance(pre_lora, ControlLoRACrossAttnProcessor):
+                lora_in = lora_in + pre_lora.process_control_states(
+                    hidden_states, scale)
+            query = query + scale * pre_lora.to_q_lora(lora_in)
+        query = query + scale * self.to_q_lora(
+            (query if self.post_add else hidden_states)
+            + self.process_control_states(hidden_states, scale))
+        for post_lora in self.post_loras:
+            lora_in = query if post_lora.post_add else hidden_states
+            if isinstance(post_lora, ControlLoRACrossAttnProcessor):
+                lora_in = lora_in + post_lora.process_control_states(
+                    hidden_states, scale)
+            query = query + scale * post_lora.to_q_lora(lora_in)
+        query = attn.head_to_batch_dim(query)
+
+        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
+
+        key = attn.to_k(encoder_hidden_states)
+        for pre_lora in self.pre_loras:
+            if not pre_lora.key_states_skipped:
+                key = key + scale * pre_lora.to_k_lora(
+                    key if pre_lora.post_add else encoder_hidden_states)
+        if not self.key_states_skipped:
+            key = key + scale * self.to_k_lora(
+                key if self.post_add else encoder_hidden_states)
+        for post_lora in self.post_loras:
+            if not post_lora.key_states_skipped:
+                key = key + scale * post_lora.to_k_lora(
+                    key if post_lora.post_add else encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+        for pre_lora in self.pre_loras:
+            if not pre_lora.value_states_skipped:
+                value = value + pre_lora.to_v_lora(
+                    value if pre_lora.post_add else encoder_hidden_states)
+        if not self.value_states_skipped:
+            value = value + scale * self.to_v_lora(
+                value if self.post_add else encoder_hidden_states)
+        for post_lora in self.post_loras:
+            if not post_lora.value_states_skipped:
+                value = value + post_lora.to_v_lora(
+                    value if post_lora.post_add else encoder_hidden_states)
+
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        out = attn.to_out[0](hidden_states)
+        for pre_lora in self.pre_loras:
+            if not pre_lora.output_states_skipped:
+                out = out + scale * pre_lora.to_out_lora(
+                    out if pre_lora.post_add else hidden_states)
+        out = out + scale * self.to_out_lora(
+            out if self.post_add else hidden_states)
+        for post_lora in self.post_loras:
+            if not post_lora.output_states_skipped:
+                out = out + scale * post_lora.to_out_lora(
+                    out if post_lora.post_add else hidden_states)
+        hidden_states = out
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        return hidden_states
+
+
+class ControlLoRACrossAttnProcessorV2(LoRACrossAttnProcessor):
+
+    def __init__(self,
+                 hidden_size,
+                 cross_attention_dim=None,
+                 rank=4,
+                 control_rank=None,
+                 control_channels=None,
+                 **kwargs):
+        super().__init__(
+            hidden_size,
+            cross_attention_dim,
+            rank,
+            post_add=False,
+            key_states_skipped=True,
+            value_states_skipped=True,
+            output_states_skipped=False)
+
+        control_rank = rank if control_rank is None else control_rank
+        control_channels = hidden_size if control_channels is None else control_channels
+        self.concat_hidden = True
+        self.control_self_add = False
+        self.control_states: torch.Tensor = None
+
+        self.to_control = LoRALinearLayer(hidden_size + control_channels,
+                                          hidden_size, control_rank)
+        self.to_control_out = LoRALinearLayer(hidden_size + control_channels,
+                                              hidden_size, control_rank)
+        self.pre_loras: List[LoRACrossAttnProcessor] = []
+        self.post_loras: List[LoRACrossAttnProcessor] = []
+
+    def inject_pre_lora(self, lora_layer):
+        self.pre_loras.append(lora_layer)
+
+    def inject_post_lora(self, lora_layer):
+        self.post_loras.append(lora_layer)
+
+    def inject_control_states(self, control_states):
+        self.control_states = control_states
+
+    def process_control_states(self, hidden_states, scale=1.0, is_out=False):
+        control_states = self.control_states.to(hidden_states.dtype)
+        if hidden_states.ndim == 3 and control_states.ndim == 4:
+            batch, _, height, width = control_states.shape
+            control_states = control_states.permute(0, 2, 3, 1).reshape(
+                batch, height * width, -1)
+            self.control_states = control_states
+        _control_states = control_states
+        if self.concat_hidden:
+            b1, b2 = control_states.shape[0], hidden_states.shape[0]
+            if b1 != b2:
+                control_states = control_states[:, None].repeat(
+                    1, b2 // b1, *([1] * (len(control_states.shape) - 1)))
+                control_states = control_states.view(-1,
+                                                     *control_states.shape[2:])
+            _control_states = torch.cat([hidden_states, control_states], -1)
+        _control_states = scale * (self.to_control_out
+                                   if is_out else self.to_control)(
+                                       _control_states)
+        if self.control_self_add:
+            control_states = control_states + _control_states
+        else:
+            control_states = _control_states
+
+        return control_states
+
+    def __call__(self,
+                 attn: CrossAttention,
+                 hidden_states,
+                 encoder_hidden_states=None,
+                 attention_mask=None,
+                 scale=1.0):
+        pre_lora: LoRACrossAttnProcessor
+        post_lora: LoRACrossAttnProcessor
+        assert self.control_states is not None
+
+        batch_size, sequence_length, _ = hidden_states.shape
+        attention_mask = attn.prepare_attention_mask(attention_mask,
+                                                     sequence_length)
+        for pre_lora in self.pre_loras:
+            if isinstance(pre_lora, ControlLoRACrossAttnProcessorV2):
+                hidden_states = hidden_states + pre_lora.process_control_states(
+                    hidden_states, scale)
+        hidden_states = hidden_states + self.process_control_states(
+            hidden_states, scale)
+        for post_lora in self.post_loras:
+            if isinstance(post_lora, ControlLoRACrossAttnProcessorV2):
+                hidden_states = hidden_states + post_lora.process_control_states(
+                    hidden_states, scale)
+        query = attn.to_q(hidden_states)
+        for pre_lora in self.pre_loras:
+            lora_in = query if pre_lora.post_add else hidden_states
+            query = query + scale * pre_lora.to_q_lora(lora_in)
+        query = query + scale * self.to_q_lora(
+            query if self.post_add else hidden_states)
+        for post_lora in self.post_loras:
+            lora_in = query if post_lora.post_add else hidden_states
+            query = query + scale * post_lora.to_q_lora(lora_in)
+        query = attn.head_to_batch_dim(query)
+
+        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
+
+        key = attn.to_k(encoder_hidden_states)
+        for pre_lora in self.pre_loras:
+            if not pre_lora.key_states_skipped:
+                key = key + scale * pre_lora.to_k_lora(
+                    key if pre_lora.post_add else encoder_hidden_states)
+        if not self.key_states_skipped:
+            key = key + scale * self.to_k_lora(
+                key if self.post_add else encoder_hidden_states)
+        for post_lora in self.post_loras:
+            if not post_lora.key_states_skipped:
+                key = key + scale * post_lora.to_k_lora(
+                    key if post_lora.post_add else encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+        for pre_lora in self.pre_loras:
+            if not pre_lora.value_states_skipped:
+                value = value + pre_lora.to_v_lora(
+                    value if pre_lora.post_add else encoder_hidden_states)
+        if not self.value_states_skipped:
+            value = value + scale * self.to_v_lora(
+                value if self.post_add else encoder_hidden_states)
+        for post_lora in self.post_loras:
+            if not post_lora.value_states_skipped:
+                value = value + post_lora.to_v_lora(
+                    value if post_lora.post_add else encoder_hidden_states)
+
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        for pre_lora in self.pre_loras:
+            if isinstance(pre_lora, ControlLoRACrossAttnProcessorV2):
+                hidden_states = hidden_states + pre_lora.process_control_states(
+                    hidden_states, scale, is_out=True)
+        hidden_states = hidden_states + self.process_control_states(
+            hidden_states, scale, is_out=True)
+        for post_lora in self.post_loras:
+            if isinstance(post_lora, ControlLoRACrossAttnProcessorV2):
+                hidden_states = hidden_states + post_lora.process_control_states(
+                    hidden_states, scale, is_out=True)
+        out = attn.to_out[0](hidden_states)
+        for pre_lora in self.pre_loras:
+            if not pre_lora.output_states_skipped:
+                out = out + scale * pre_lora.to_out_lora(
+                    out if pre_lora.post_add else hidden_states)
+        out = out + scale * self.to_out_lora(
+            out if self.post_add else hidden_states)
+        for post_lora in self.post_loras:
+            if not post_lora.output_states_skipped:
+                out = out + scale * post_lora.to_out_lora(
+                    out if post_lora.post_add else hidden_states)
+        hidden_states = out
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        return hidden_states
+
+
+class ConvBlock2D(nn.Module):
+
+    def __init__(
+        self,
+        *,
+        in_channels,
+        out_channels=None,
+        conv_kernel_size=3,
+        dropout=0.0,
+        temb_channels=512,
+        groups=32,
+        groups_out=None,
+        pre_norm=True,
+        eps=1e-6,
+        non_linearity='swish',
+        time_embedding_norm='default',
+        kernel=None,
+        output_scale_factor=1.0,
+        up=False,
+        down=False,
+    ):
+        super().__init__()
+        self.pre_norm = pre_norm
+        self.pre_norm = True
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.time_embedding_norm = time_embedding_norm
+        self.up = up
+        self.down = down
+        self.output_scale_factor = output_scale_factor
+
+        if groups_out is None:
+            groups_out = groups
+
+        self.norm1 = torch.nn.GroupNorm(
+            num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+
+        self.conv1 = torch.nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=conv_kernel_size,
+            stride=1,
+            padding=conv_kernel_size // 2)
+
+        if temb_channels is not None:
+            if self.time_embedding_norm == 'default':
+                time_emb_proj_out_channels = out_channels
+            elif self.time_embedding_norm == 'scale_shift':
+                time_emb_proj_out_channels = out_channels * 2
+            else:
+                raise ValueError(
+                    f'unknown time_embedding_norm : {self.time_embedding_norm} '
+                )
+
+            self.time_emb_proj = torch.nn.Linear(temb_channels,
+                                                 time_emb_proj_out_channels)
+        else:
+            self.time_emb_proj = None
+
+        self.norm2 = torch.nn.GroupNorm(
+            num_groups=groups_out,
+            num_channels=out_channels,
+            eps=eps,
+            affine=True)
+        self.dropout = torch.nn.Dropout(dropout)
+
+        if non_linearity == 'swish':
+            self.nonlinearity = lambda x: F.silu(x)
+        elif non_linearity == 'mish':
+            self.nonlinearity = Mish()
+        elif non_linearity == 'silu':
+            self.nonlinearity = nn.SiLU()
+
+        self.upsample = self.downsample = None
+        if self.up:
+            if kernel == 'fir':
+                fir_kernel = (1, 3, 3, 1)
+                self.upsample = lambda x: upsample_2d(x, kernel=fir_kernel)
+            elif kernel == 'sde_vp':
+                self.upsample = partial(
+                    F.interpolate, scale_factor=2.0, mode='nearest')
+            else:
+                self.upsample = Upsample2D(in_channels, use_conv=False)
+        elif self.down:
+            if kernel == 'fir':
+                fir_kernel = (1, 3, 3, 1)
+                self.downsample = lambda x: downsample_2d(x, kernel=fir_kernel)
+            elif kernel == 'sde_vp':
+                self.downsample = partial(
+                    F.avg_pool2d, kernel_size=2, stride=2)
+            else:
+                self.downsample = Downsample2D(
+                    in_channels, use_conv=False, padding=1, name='op')
+
+    def forward(self, input_tensor, temb):
+        hidden_states = input_tensor
+
+        hidden_states = self.norm1(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+
+        if self.upsample is not None:
+            # upsample_nearest_nhwc fails with large batch sizes.
+            # see https://github.com/huggingface/diffusers/issues/984
+            if hidden_states.shape[0] >= 64:
+                input_tensor = input_tensor.contiguous()
+                hidden_states = hidden_states.contiguous()
+            _ = self.upsample(input_tensor)
+            hidden_states = self.upsample(hidden_states)
+        elif self.downsample is not None:
+            _ = self.downsample(input_tensor)
+            hidden_states = self.downsample(hidden_states)
+
+        hidden_states = self.conv1(hidden_states)
+
+        if temb is not None:
+            temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None,
+                                                               None]
+
+        if temb is not None and self.time_embedding_norm == 'default':
+            hidden_states = hidden_states + temb
+
+        hidden_states = self.norm2(hidden_states)
+
+        if temb is not None and self.time_embedding_norm == 'scale_shift':
+            scale, shift = torch.chunk(temb, 2, dim=1)
+            hidden_states = hidden_states * (1 + scale) + shift
+
+        hidden_states = self.nonlinearity(hidden_states)
+
+        output_tensor = self.dropout(hidden_states)
+
+        return output_tensor
+
+
+class SimpleDownEncoderBlock2D(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        convnet_eps: float = 1e-6,
+        convnet_time_scale_shift: str = 'default',
+        convnet_act_fn: str = 'swish',
+        convnet_groups: int = 32,
+        convnet_pre_norm: bool = True,
+        convnet_kernel_size: int = 3,
+        output_scale_factor=1.0,
+        add_downsample=True,
+        downsample_padding=1,
+    ):
+        super().__init__()
+        convnets = []
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            convnets.append(
+                ConvBlock2D(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    temb_channels=None,
+                    eps=convnet_eps,
+                    groups=convnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=convnet_time_scale_shift,
+                    non_linearity=convnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=convnet_pre_norm,
+                    conv_kernel_size=convnet_kernel_size,
+                ))
+        in_channels = in_channels if num_layers == 0 else out_channels
+
+        self.convnets = nn.ModuleList(convnets)
+
+        if add_downsample:
+            self.downsamplers = nn.ModuleList([
+                Downsample2D(
+                    in_channels,
+                    use_conv=True,
+                    out_channels=out_channels,
+                    padding=downsample_padding,
+                    name='op')
+            ])
+        else:
+            self.downsamplers = None
+
+    def forward(self, hidden_states):
+        for convnet in self.convnets:
+            hidden_states = convnet(hidden_states, temb=None)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+        return hidden_states

modelscope/tuners/lora.py

@@ -0,0 +1,624 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
+import logging
+import math
+import os.path
+import types
+from typing import Dict, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+
+class LoRATuner:
+
+    @staticmethod
+    def tune(model: nn.Module,
+             rank=6,
+             replace_modules=None,
+             lora_alpha=1.,
+             lora_dropout=0.,
+             merge_weights=True,
+             fan_in_fan_out=False,
+             bias='none',
+             pretrained_tuner=None):
+        """Tune a model with lora.
+
+        Args:
+            model: The torch.nn.Module containing the target module to be patched.
+            rank: The lora rank.
+            replace_modules: The module names to be replaced, the replacing strategy is `end with`.
+            lora_alpha: The alpha value for lora module.
+            lora_dropout: The dropout value for lora module.
+            merge_weights: If merge_weights set to True, when the module turns to `eval`, the lora weights
+                will be added into the origin weight to reduce calculation.
+            fan_in_fan_out: Set this to True if the layer to replace stores weight like (fan_in, fan_out).
+            bias: The grad strategy for bias, can be `none`, 'all' or 'lora_only'.
+            pretrained_tuner: The pretrained file of lora.
+
+        Returns:
+            The lora modules
+        """
+        modules = LoRATuner._dynamic_patch_lora(
+            model,
+            replace_modules=replace_modules,
+            r=rank,
+            lora_alpha=lora_alpha,
+            lora_dropout=lora_dropout,
+            merge_weights=merge_weights,
+            fan_in_fan_out=fan_in_fan_out)
+
+        mark_only_lora_as_trainable(model, bias)
+
+        def state_dict_hook(module, destination, prefix, local_metadata):
+            return lora_state_dict(destination, bias)
+
+        model.state_dict_hook_handle = model._register_state_dict_hook(
+            state_dict_hook)
+
+        def warning_hook(module, incompatible_keys):
+            logger.info(
+                f'The {module.__class__.__name__} module has unmatched keys: {incompatible_keys},'
+                f'this is converted to a notice with respect to LoRA')
+            for ik in incompatible_keys:
+                ik.clear()
+
+        if hasattr(model, 'register_load_state_dict_post_hook'):
+            model.load_state_dict_hook_handle = model.register_load_state_dict_post_hook(
+                warning_hook)
+        else:
+
+            def load_state_dict(self, state_dict, strict=True):
+                return self.load_state_dict_origin(state_dict, False)
+
+            model.load_state_dict_origin = model.load_state_dict
+            model.load_state_dict = types.MethodType(load_state_dict, model)
+
+        if pretrained_tuner is not None and os.path.isfile(pretrained_tuner):
+            logger.info(f'Loading LoRA weights from file: {pretrained_tuner}')
+            model.load_state_dict(torch.load(pretrained_tuner))
+
+        return modules
+
+    @staticmethod
+    def _dynamic_patch_lora(model, replace_modules, **kwargs):
+        """Dynamic patch lora to model
+
+        Args:
+            model: The torch.nn.Module containing the target module to be patched.
+            replace_modules: The module names to be replaced, the replacing strategy is `end with`.
+            **kwargs: The arguments passed from `tune` which are needed by lora.
+
+        Returns:
+            The lora modules
+        """
+        modules = []
+        module_keys = [key for key, _ in model.named_modules()]
+        assert isinstance(replace_modules, (str, list))
+        if isinstance(replace_modules, str):
+            replace_modules = [replace_modules]
+
+        for module_key in module_keys:
+            if any([module_key.endswith(name)
+                    for name in replace_modules]):  # noqa
+                parts = module_key.split('.')
+                module = model.get_submodule('.'.join(parts[:-1]))
+                sub_module = model.get_submodule(module_key)
+                _key = parts[-1]
+
+                lora_module = None
+                if isinstance(sub_module, torch.nn.Linear):
+                    lora_module = Linear(
+                        sub_module.in_features,
+                        sub_module.out_features,
+                        bias=sub_module.bias is not None,
+                        **kwargs)
+                elif isinstance(sub_module, torch.nn.Conv2d):
+                    kwargs.pop('fan_in_fan_out', None)
+                    lora_module = Conv2d(
+                        sub_module.in_channels,
+                        sub_module.out_channels,
+                        kernel_size=sub_module.kernel_size,
+                        stride=sub_module.stride,
+                        padding=sub_module.padding,
+                        dilation=sub_module.dilation,
+                        groups=sub_module.groups,
+                        **kwargs)
+
+                if lora_module is not None:
+                    lora_module.weight = sub_module.weight
+                    if sub_module.bias is not None:
+                        lora_module.bias = sub_module.bias
+                    lora_module.to(sub_module.weight.device).to(
+                        sub_module.weight.dtype)
+                    setattr(module, _key, lora_module)
+                    modules.append(lora_module)
+        return modules
+
+    @staticmethod
+    def unpatch_lora(model, replace_modules):
+        """Unpatch lora modules and merge the weights to original modules.
+
+        Args:
+            model: The model called with `tune` function.
+            replace_modules: The module names to be replaced, the replacing strategy is `end with`.
+
+        Returns:
+            The lora modules.
+        """
+        modules = []
+        module_keys = [key for key, _ in model.named_modules()]
+        assert isinstance(replace_modules, (str, list))
+        if isinstance(replace_modules, str):
+            replace_modules = [replace_modules]
+
+        for module_key in module_keys:
+            if any([module_key.endswith(name)
+                    for name in replace_modules]):  # noqa
+                parts = module_key.split('.')
+                module = model.get_submodule('.'.join(parts[:-1]))
+                sub_module = model.get_submodule(module_key)
+                _key = parts[-1]
+
+                origin_module = None
+                if isinstance(sub_module, Linear):
+                    origin_module = torch.nn.Linear(
+                        sub_module.in_features,
+                        sub_module.out_features,
+                        bias=sub_module.bias is not None)
+                elif isinstance(sub_module, Conv2d):
+                    origin_module = torch.nn.Conv2d(
+                        sub_module.in_channels,
+                        sub_module.out_channels,
+                        kernel_size=sub_module.kernel_size,
+                        stride=sub_module.stride,
+                        padding=sub_module.padding,
+                        dilation=sub_module.dilation,
+                        groups=sub_module.groups)
+
+                if origin_module is not None:
+                    sub_module.merge_weights = True
+                    sub_module.eval()
+                    origin_module.weight = sub_module.weight
+                    if sub_module.bias is not None:
+                        origin_module.bias = sub_module.bias
+                    origin_module.to(sub_module.weight.device).to(
+                        sub_module.weight.dtype)
+                    setattr(module, _key, origin_module)
+                    modules.append(sub_module)
+
+        model.state_dict_hook_handle.remove()
+        if hasattr(model, 'load_state_dict_hook_handle'):
+            model.load_state_dict_hook_handle.remove()
+        else:
+            model.load_state_dict = model.load_state_dict_origin
+        return modules
+
+
+class LoRALayer:
+
+    def __init__(
+        self,
+        r: int,
+        lora_alpha: int,
+        lora_dropout: float,
+        merge_weights: bool,
+    ):
+        self.r = r
+        self.lora_alpha = lora_alpha
+        # Optional dropout
+        if lora_dropout > 0.:
+            self.lora_dropout = nn.Dropout(p=lora_dropout)
+        else:
+            self.lora_dropout = lambda x: x
+        # Mark the weight as unmerged
+        self.merged = False
+        self.merge_weights = merge_weights
+
+
+class Embedding(nn.Embedding, LoRALayer):
+    # LoRA implemented in a dense layer
+    def __init__(self,
+                 num_embeddings: int,
+                 embedding_dim: int,
+                 r: int = 0,
+                 lora_alpha: int = 1,
+                 merge_weights: bool = True,
+                 **kwargs):
+        nn.Embedding.__init__(self, num_embeddings, embedding_dim, **kwargs)
+        LoRALayer.__init__(
+            self,
+            r=r,
+            lora_alpha=lora_alpha,
+            lora_dropout=0,
+            merge_weights=merge_weights)
+        # Actual trainable parameters
+        if r > 0:
+            self.lora_A = nn.Parameter(
+                self.weight.new_zeros((r, num_embeddings)))
+            self.lora_B = nn.Parameter(
+                self.weight.new_zeros((embedding_dim, r)))
+            self.scaling = self.lora_alpha / self.r
+            # Freezing the pre-trained weight matrix
+            self.weight.requires_grad = False
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.Embedding.reset_parameters(self)
+        if hasattr(self, 'lora_A'):
+            # initialize A the same way as the default for nn.Linear and B to zero
+            nn.init.zeros_(self.lora_A)
+            nn.init.normal_(self.lora_B)
+
+    def train(self, mode: bool = True):
+        nn.Embedding.train(self, mode)
+        self.lora_A.requires_grad = mode
+        self.lora_B.requires_grad = mode
+        if mode and self.merge_weights and self.merged:
+            # Make sure that the weights are not merged
+            if self.r > 0:
+                self.weight.data -= (self.lora_B
+                                     @ self.lora_A).T * self.scaling
+            self.merged = False
+        if not mode and self.merge_weights and not self.merged:
+            # Merge the weights and mark it
+            if self.r > 0:
+                self.weight.data += (self.lora_B
+                                     @ self.lora_A).T * self.scaling
+            self.merged = True
+
+    def eval(self):
+        nn.Embedding.eval(self)
+        self.lora_A.requires_grad = False
+        self.lora_B.requires_grad = False
+        if self.merge_weights and not self.merged:
+            # Merge the weights and mark it
+            if self.r > 0:
+                self.weight.data += (self.lora_B @ self.lora_A) * self.scaling
+            self.merged = True
+
+    def forward(self, x: torch.Tensor):
+        if self.r > 0 and not self.merged:
+            result = nn.Embedding.forward(self, x)
+            if self.r > 0:
+                after_A = F.embedding(x, self.lora_A.T, self.padding_idx,
+                                      self.max_norm, self.norm_type,
+                                      self.scale_grad_by_freq, self.sparse)
+                result += (after_A @ self.lora_B.T) * self.scaling
+            return result
+        else:
+            return nn.Embedding.forward(self, x)
+
+
+class Linear(nn.Linear, LoRALayer):
+    # LoRA implemented in a dense layer
+    def __init__(
+            self,
+            in_features: int,
+            out_features: int,
+            r: int = 0,
+            lora_alpha: int = 1,
+            lora_dropout: float = 0.,
+            fan_in_fan_out: bool = False,
+            # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
+            merge_weights: bool = True,
+            **kwargs):
+        nn.Linear.__init__(self, in_features, out_features, **kwargs)
+        LoRALayer.__init__(
+            self,
+            r=r,
+            lora_alpha=lora_alpha,
+            lora_dropout=lora_dropout,
+            merge_weights=merge_weights)
+
+        self.fan_in_fan_out = fan_in_fan_out
+        # Actual trainable parameters
+        if r > 0:
+            self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features)))
+            self.lora_B = nn.Parameter(
+                self.weight.new_zeros((out_features, r)))
+            self.scaling = self.lora_alpha / self.r
+            # Freezing the pre-trained weight matrix
+            self.weight.requires_grad = False
+        self.reset_parameters()
+        if fan_in_fan_out:
+            self.weight.data = self.weight.data.T
+
+    def reset_parameters(self):
+        nn.Linear.reset_parameters(self)
+        if hasattr(self, 'lora_A'):
+            # initialize A the same way as the default for nn.Linear and B to zero
+            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+            nn.init.zeros_(self.lora_B)
+
+    def train(self, mode: bool = True):
+
+        def T(w):
+            return w.T if self.fan_in_fan_out else w
+
+        nn.Linear.train(self, mode)
+        self.lora_A.requires_grad = mode
+        self.lora_B.requires_grad = mode
+        if mode and self.merge_weights and self.merged:
+            # Make sure that the weights are not merged
+            if self.r > 0:
+                self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling
+            self.merged = False
+        if not mode and self.merge_weights and not self.merged:
+            # Merge the weights and mark it
+            if self.r > 0:
+                self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling
+            self.merged = True
+
+    def eval(self):
+
+        def T(w):
+            return w.T if self.fan_in_fan_out else w
+
+        nn.Linear.eval(self)
+        self.lora_A.requires_grad = False
+        self.lora_B.requires_grad = False
+        if self.merge_weights and not self.merged:
+            # Merge the weights and mark it
+            if self.r > 0:
+                self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling
+            self.merged = True
+
+    def forward(self, x: torch.Tensor):
+
+        def T(w):
+            return w.T if self.fan_in_fan_out else w
+
+        if self.r > 0 and not self.merged:
+            result = F.linear(x, T(self.weight), bias=self.bias)
+            if self.r > 0:
+                result += (self.lora_dropout(x) @ self.lora_A.T
+                           @ self.lora_B.T) * self.scaling
+            return result
+        else:
+            return F.linear(x, T(self.weight), bias=self.bias)
+
+
+class MergedLinear(nn.Linear, LoRALayer):
+    # LoRA implemented in a dense layer
+    def __init__(self,
+                 in_features: int,
+                 out_features: int,
+                 r: int = 0,
+                 lora_alpha: int = 1,
+                 lora_dropout: float = 0.,
+                 enable_lora: List[bool] = [False],
+                 fan_in_fan_out: bool = False,
+                 merge_weights: bool = True,
+                 **kwargs):
+        nn.Linear.__init__(self, in_features, out_features, **kwargs)
+        LoRALayer.__init__(
+            self,
+            r=r,
+            lora_alpha=lora_alpha,
+            lora_dropout=lora_dropout,
+            merge_weights=merge_weights)
+        assert out_features % len(enable_lora) == 0, \
+            'The length of enable_lora must divide out_features'
+        self.enable_lora = enable_lora
+        self.fan_in_fan_out = fan_in_fan_out
+        # Actual trainable parameters
+        if r > 0 and any(enable_lora):
+            self.lora_A = nn.Parameter(
+                self.weight.new_zeros((r * sum(enable_lora), in_features)))
+            self.lora_B = nn.Parameter(
+                self.weight.new_zeros(
+                    (out_features // len(enable_lora) * sum(enable_lora),
+                     r)))  # weights for Conv1D with groups=sum(enable_lora)
+            self.scaling = self.lora_alpha / self.r
+            # Freezing the pre-trained weight matrix
+            self.weight.requires_grad = False
+            # Compute the indices
+            self.lora_ind = self.weight.new_zeros(
+                (out_features, ), dtype=torch.bool).view(len(enable_lora), -1)
+            self.lora_ind[enable_lora, :] = True
+            self.lora_ind = self.lora_ind.view(-1)
+        self.reset_parameters()
+        if fan_in_fan_out:
+            self.weight.data = self.weight.data.T
+
+    def reset_parameters(self):
+        nn.Linear.reset_parameters(self)
+        if hasattr(self, 'lora_A'):
+            # initialize A the same way as the default for nn.Linear and B to zero
+            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+            nn.init.zeros_(self.lora_B)
+
+    def zero_pad(self, x):
+        result = x.new_zeros((*x.shape[:-1], self.out_features))
+        result = result.view(-1, self.out_features)
+        result[:, self.lora_ind] = x.reshape(
+            -1,
+            self.out_features // len(self.enable_lora) * sum(self.enable_lora))
+        return result.view((*x.shape[:-1], self.out_features))
+
+    def train(self, mode: bool = True):
+
+        def T(w):
+            return w.T if self.fan_in_fan_out else w
+
+        nn.Linear.train(self, mode)
+        self.lora_A.requires_grad = mode
+        self.lora_B.requires_grad = mode
+        if mode and self.merge_weights and self.merged:
+            # Make sure that the weights are not merged
+            if self.r > 0 and any(self.enable_lora):
+                delta_w = F.conv1d(
+                    self.lora_A.data.unsqueeze(0),
+                    self.lora_B.data.unsqueeze(-1),
+                    groups=sum(self.enable_lora)).squeeze(0)
+                self.weight.data -= self.zero_pad(T(delta_w * self.scaling))
+            self.merged = False
+        if not mode and self.merge_weights and not self.merged:
+            if self.r > 0 and any(self.enable_lora):
+                delta_w = F.conv1d(
+                    self.lora_A.data.unsqueeze(0),
+                    self.lora_B.data.unsqueeze(-1),
+                    groups=sum(self.enable_lora)).squeeze(0)
+                self.weight.data += self.zero_pad(T(delta_w * self.scaling))
+            self.merged = True
+
+    def eval(self):
+
+        def T(w):
+            return w.T if self.fan_in_fan_out else w
+
+        nn.Linear.eval(self)
+        self.lora_A.requires_grad = False
+        self.lora_B.requires_grad = False
+        if self.merge_weights and not self.merged:
+            # Merge the weights and mark it
+            if self.r > 0 and any(self.enable_lora):
+                delta_w = F.conv1d(
+                    self.lora_A.data.unsqueeze(0),
+                    self.lora_B.data.unsqueeze(-1),
+                    groups=sum(self.enable_lora)).squeeze(0)
+                self.weight.data += self.zero_pad(T(delta_w * self.scaling))
+            self.merged = True
+
+    def forward(self, x: torch.Tensor):
+
+        def T(w):
+            return w.T if self.fan_in_fan_out else w
+
+        if self.merged:
+            return F.linear(x, T(self.weight), bias=self.bias)
+        else:
+            result = F.linear(x, T(self.weight), bias=self.bias)
+            if self.r > 0:
+                after_A = F.linear(self.lora_dropout(x), self.lora_A)
+                after_B = F.conv1d(
+                    after_A.transpose(-2, -1),
+                    self.lora_B.unsqueeze(-1),
+                    groups=sum(self.enable_lora)).transpose(-2, -1)
+                result += self.zero_pad(after_B) * self.scaling
+            return result
+
+
+class Conv2d(nn.Conv2d, LoRALayer):
+    # LoRA implemented in a dense layer
+    def __init__(self,
+                 in_channels: int,
+                 out_channels: int,
+                 kernel_size: int,
+                 r: int = 0,
+                 lora_alpha: int = 1,
+                 lora_dropout: float = 0.,
+                 merge_weights: bool = True,
+                 **kwargs):
+        nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size,
+                           **kwargs)
+        LoRALayer.__init__(
+            self,
+            r=r,
+            lora_alpha=lora_alpha,
+            lora_dropout=lora_dropout,
+            merge_weights=merge_weights)
+        assert type(kernel_size) is int
+        # Actual trainable parameters
+        if r > 0:
+            self.lora_A = nn.Parameter(
+                self.weight.new_zeros(
+                    (r * kernel_size, in_channels * kernel_size)))
+            self.lora_B = nn.Parameter(
+                self.weight.new_zeros(
+                    (out_channels * kernel_size, r * kernel_size)))
+            self.scaling = self.lora_alpha / self.r
+            # Freezing the pre-trained weight matrix
+            self.weight.requires_grad = False
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.Conv2d.reset_parameters(self)
+        if hasattr(self, 'lora_A'):
+            # initialize A the same way as the default for nn.Linear and B to zero
+            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+            nn.init.zeros_(self.lora_B)
+
+    def train(self, mode: bool = True):
+        nn.Conv2d.train(self, mode)
+        self.lora_A.requires_grad = mode
+        self.lora_B.requires_grad = mode
+        if mode and self.merge_weights and self.merged:
+            # Make sure that the weights are not merged
+            self.weight.data -= (self.lora_B @ self.lora_A).view(
+                self.weight.shape) * self.scaling
+            self.merged = False
+        if not mode and self.merge_weights and not self.merged:
+            self.weight.data += (self.lora_B @ self.lora_A).view(
+                self.weight.shape) * self.scaling
+            self.merged = True
+
+    def eval(self):
+        nn.Conv2d.eval(self)
+        self.lora_A.requires_grad = False
+        self.lora_B.requires_grad = False
+        if self.merge_weights and not self.merged:
+            # Merge the weights and mark it
+            self.weight.data += (self.lora_B @ self.lora_A).view(
+                self.weight.shape) * self.scaling
+            self.merged = True
+
+    def forward(self, x: torch.Tensor):
+        if self.r > 0 and not self.merged:
+            return F.conv2d(
+                x,
+                self.weight +  # noqa
+                (self.lora_B @ self.lora_A).view(self.weight.shape)  # noqa
+                * self.scaling,
+                self.bias,
+                self.stride,
+                self.padding,
+                self.dilation,
+                self.groups)
+        return nn.Conv2d.forward(self, x)
+
+
+def mark_only_lora_as_trainable(model: nn.Module, bias: str = 'none') -> None:
+    for n, p in model.named_parameters():
+        if 'lora_' not in n:
+            p.requires_grad = False
+    if bias == 'none':
+        return
+    elif bias == 'all':
+        for n, p in model.named_parameters():
+            if 'bias' in n:
+                p.requires_grad = True
+    elif bias == 'lora_only':
+        for m in model.modules():
+            if isinstance(m, LoRALayer) and \
+                    hasattr(m, 'bias') and \
+                    m.bias is not None:
+                m.bias.requires_grad = True
+    else:
+        raise NotImplementedError
+
+
+def lora_state_dict(state_dict, bias: str = 'none') -> Dict[str, torch.Tensor]:
+    if bias == 'none':
+        return {k: state_dict[k] for k in state_dict if 'lora_' in k}
+    elif bias == 'all':
+        return {
+            k: state_dict[k]
+            for k in state_dict if 'lora_' in k or 'bias' in k
+        }
+    elif bias == 'lora_only':
+        to_return = {}
+        for k in state_dict:
+            if 'lora_' in k:
+                to_return[k] = state_dict[k]
+                bias_name = k.split('lora_')[0] + 'bias'
+                if bias_name in state_dict:
+                    to_return[bias_name] = state_dict[bias_name]
+        return to_return
+    else:
+        raise NotImplementedError

modelscope/tuners/sd_lora.py

@@ -0,0 +1,217 @@
+# Copyright 2023-2024 The Alibaba Fundamental Vision Team Authors. All rights reserved.
+# The implementation is adopted from HighCWu,
+# made pubicly available under the Apache License 2.0 License at https://github.com/HighCWu/ControlLoRA
+import os
+from dataclasses import dataclass
+from typing import List, Tuple, Union
+
+import torch
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.cross_attention import CrossAttention, LoRALinearLayer
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils.outputs import BaseOutput
+
+
+@dataclass
+class TunerOutput(BaseOutput):
+    lora_states: Tuple[torch.FloatTensor]
+
+
+class LoRACrossAttnProcessor(nn.Module):
+    """ The implementation of lora attention module.
+    """
+
+    def __init__(self,
+                 hidden_size,
+                 cross_attention_dim=None,
+                 rank=4,
+                 post_add=False,
+                 key_states_skipped=False,
+                 value_states_skipped=False,
+                 output_states_skipped=False):
+        """ Initialize a lora attn instance.
+        Args:
+            hidden_size (`int`): The number of channels in embedding.
+            cross_attention_dim (`int`, *optional*):
+                The number of channels in the hidden_states. If not given, defaults to `hidden_size`.
+            rank (`int`,  *optional*, defaults to 4): The number of rank of lora.
+            post_add (`bool`,  *optional*, defaults to False): Set to `True`, conduct weighted
+            adding operation after lora.
+            key_states_skipped (`bool`, *optional*, defaults to False):
+                Set to `True` for skip to perform lora on key value.
+            value_states_skipped (`bool`, *optional*, defaults to False):
+                Set to `True` for skip to perform lora on value.
+            output_states_skipped (`bool`, *optional*, defaults to False):
+                Set to `True` for skip to perform lora on output value.
+        """
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.cross_attention_dim = cross_attention_dim
+        self.rank = rank
+        self.post_add = post_add
+
+        self.to_q_lora = LoRALinearLayer(hidden_size, hidden_size, rank)
+        if not key_states_skipped:
+            self.to_k_lora = LoRALinearLayer(
+                hidden_size if post_add else
+                (cross_attention_dim or hidden_size), hidden_size, rank)
+        if not value_states_skipped:
+            self.to_v_lora = LoRALinearLayer(
+                hidden_size if post_add else
+                (cross_attention_dim or hidden_size), hidden_size, rank)
+        if not output_states_skipped:
+            self.to_out_lora = LoRALinearLayer(hidden_size, hidden_size, rank)
+
+        self.key_states_skipped: bool = key_states_skipped
+        self.value_states_skipped: bool = value_states_skipped
+        self.output_states_skipped: bool = output_states_skipped
+
+    def skip_key_states(self, is_skipped: bool = True):
+        if not is_skipped:
+            assert hasattr(self, 'to_k_lora')
+        self.key_states_skipped = is_skipped
+
+    def skip_value_states(self, is_skipped: bool = True):
+        if not is_skipped:
+            assert hasattr(self, 'to_q_lora')
+        self.value_states_skipped = is_skipped
+
+    def skip_output_states(self, is_skipped: bool = True):
+        if not is_skipped:
+            assert hasattr(self, 'to_out_lora')
+        self.output_states_skipped = is_skipped
+
+    def __call__(self,
+                 attn: CrossAttention,
+                 hidden_states,
+                 encoder_hidden_states=None,
+                 attention_mask=None,
+                 scale=1.0):
+        batch_size, sequence_length, _ = hidden_states.shape
+        attention_mask = attn.prepare_attention_mask(attention_mask,
+                                                     sequence_length,
+                                                     batch_size)
+
+        query = attn.to_q(hidden_states)
+        query = query + scale * self.to_q_lora(
+            query if self.post_add else hidden_states)
+        query = attn.head_to_batch_dim(query)
+
+        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
+
+        key = attn.to_k(encoder_hidden_states)
+        if not self.key_states_skipped:
+            key = key + scale * self.to_k_lora(
+                key if self.post_add else encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+        if not self.value_states_skipped:
+            value = value + scale * self.to_v_lora(
+                value if self.post_add else encoder_hidden_states)
+
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        out = attn.to_out[0](hidden_states)
+        if not self.output_states_skipped:
+            out = out + scale * self.to_out_lora(
+                out if self.post_add else hidden_states)
+        hidden_states = out
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        return hidden_states
+
+
+class LoRATuner(ModelMixin, ConfigMixin):
+
+    @staticmethod
+    def tune(
+        model: nn.Module,
+        tuner_config=None,
+        pretrained_tuner=None,
+    ):
+        tuner = LoRATuner.from_config(tuner_config)
+        if pretrained_tuner is not None and os.path.exists(pretrained_tuner):
+            tuner.load_state_dict(
+                torch.load(pretrained_tuner, map_location='cpu'), strict=True)
+        tune_layers_list = list(
+            [list(layer_list) for layer_list in tuner.lora_layers])
+        assert hasattr(model, 'unet')
+        unet = model.unet
+        tuner.to(unet.device)
+        tune_attn_procs = tuner.set_tune_layers(unet, tune_layers_list)
+        unet.set_attn_processor(tune_attn_procs)
+        return tuner
+
+    def set_tune_layers(self, unet, tune_layers_list):
+        n_ch = len(unet.config.block_out_channels)
+        control_ids = [i for i in range(n_ch)]
+        tune_attn_procs = {}
+
+        for name in unet.attn_processors.keys():
+            if name.startswith('mid_block'):
+                control_id = control_ids[-1]
+            elif name.startswith('up_blocks'):
+                block_id = int(name[len('up_blocks.')])
+                control_id = list(reversed(control_ids))[block_id]
+            elif name.startswith('down_blocks'):
+                block_id = int(name[len('down_blocks.')])
+                control_id = control_ids[block_id]
+
+            tune_layers = tune_layers_list[control_id]
+            if len(tune_layers) != 0:
+                tune_layer = tune_layers.pop(0)
+                tune_attn_procs[name] = tune_layer
+        return tune_attn_procs
+
+    @register_to_config
+    def __init__(
+        self,
+        lora_block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
+        lora_cross_attention_dims: Tuple[List[int]] = ([
+            None, 768, None, 768, None, 768, None, 768, None, 768
+        ], [None, 768, None, 768, None, 768, None, 768, None,
+            768], [None, 768, None, 768, None, 768, None, 768, None,
+                   768], [None, 768]),
+        lora_rank: int = 4,
+        lora_post_add: bool = False,
+        lora_key_states_skipped: bool = False,
+        lora_value_states_skipped: bool = False,
+        lora_output_states_skipped: bool = False,
+    ):
+        super().__init__()
+
+        lora_cls = LoRACrossAttnProcessor
+
+        self.lora_layers = nn.ModuleList([])
+
+        for i, lora_cross_attention_dim in enumerate(
+                lora_cross_attention_dims):
+            self.lora_layers.append(
+                nn.ModuleList([
+                    lora_cls(
+                        lora_block_out_channels[i],
+                        cross_attention_dim=cross_attention_dim,
+                        rank=lora_rank,
+                        post_add=lora_post_add,
+                        key_states_skipped=lora_key_states_skipped,
+                        value_states_skipped=lora_value_states_skipped,
+                        output_states_skipped=lora_output_states_skipped)
+                    for cross_attention_dim in lora_cross_attention_dim
+                ]))
+
+    def forward(self) -> Union[TunerOutput, Tuple]:
+        lora_states_list = []
+        tune_layers_list = list(
+            [list(layer_list) for layer_list in self.lora_layers])
+        for tune_list in tune_layers_list:
+            for tune_layer in tune_list:
+                lora_states_list.append(tune_layer.to_q_lora.down.weight)
+        return TunerOutput(lora_states=tuple(lora_states_list))

modelscope/utils/constant.py

@@ -246,6 +246,7 @@ class MultiModalTasks(object):
     video_question_answering = 'video-question-answering'
     video_temporal_grounding = 'video-temporal-grounding'
     text_to_video_synthesis = 'text-to-video-synthesis'
+    efficient_diffusion_tuning = 'efficient-diffusion-tuning'
 
 
 class ScienceTasks(object):
@@ -266,6 +267,7 @@ class TasksIODescriptions(object):
     visual_question_answering = 'visual_question_answering',
     visual_entailment = 'visual_entailment',
     generative_multi_modal_embedding = 'generative_multi_modal_embedding'
+    efficient_diffusion_tuning = 'efficient_diffusion_tuning'
 
 
 class Tasks(CVTasks, NLPTasks, AudioTasks, MultiModalTasks, ScienceTasks):

requirements/multi-modal.txt

@@ -1,5 +1,5 @@
 accelerate
-diffusers>=0.11.1
+diffusers>=0.13.1
 ftfy>=6.0.3
 librosa<=0.9.2
 opencv-python

tests/pipelines/test_efficient_diffusion_tuning.py

@@ -0,0 +1,63 @@
+# Copyright 2022-2023 The Alibaba Fundamental Vision Team Authors. All rights reserved.
+import unittest
+
+from modelscope.models import Model
+from modelscope.models.multi_modal import EfficientStableDiffusion
+from modelscope.pipelines import pipeline
+from modelscope.utils.constant import Tasks
+from modelscope.utils.demo_utils import DemoCompatibilityCheck
+from modelscope.utils.test_utils import test_level
+
+
+class EfficientDiffusionTuningTest(unittest.TestCase, DemoCompatibilityCheck):
+
+    def setUp(self) -> None:
+        self.task = Tasks.efficient_diffusion_tuning
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_efficient_diffusion_tuning_lora_run_pipeline(self):
+        model_id = 'damo/multi-modal_efficient-diffusion-tuning-lora'
+        inputs = {'prompt': 'pale golden rod circle with old lace background'}
+        edt_pipeline = pipeline(self.task, model_id)
+        result = edt_pipeline(inputs)
+        print(f'Efficient-diffusion-tuning-lora output: {result}.')
+
+    @unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
+    def test_efficient_diffusion_tuning_lora_load_model_from_pretrained(self):
+        model_id = 'damo/multi-modal_efficient-diffusion-tuning-lora'
+        model = Model.from_pretrained(model_id)
+        self.assertTrue(model.__class__ == EfficientStableDiffusion)
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_efficient_diffusion_tuning_lora_demo_compatibility(self):
+        self.model_id = 'damo/multi-modal_efficient-diffusion-tuning-lora'
+        self.compatibility_check()
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_efficient_diffusion_tuning_control_lora_run_pipeline(self):
+        model_id = 'damo/multi-modal_efficient-diffusion-tuning-control-lora'
+        inputs = {
+            'prompt':
+            'pale golden rod circle with old lace background',
+            'cond':
+            'data/test/images/efficient_diffusion_tuning_sd_control_lora_source.png'
+        }
+        edt_pipeline = pipeline(self.task, model_id)
+        result = edt_pipeline(inputs)
+        print(f'Efficient-diffusion-tuning-control-lora output: {result}.')
+
+    @unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
+    def test_efficient_diffusion_tuning_control_lora_load_model_from_pretrained(
+            self):
+        model_id = 'damo/multi-modal_efficient-diffusion-tuning-control-lora'
+        model = Model.from_pretrained(model_id)
+        self.assertTrue(model.__class__ == EfficientStableDiffusion)
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_efficient_diffusion_tuning_control_lora_demo_compatibility(self):
+        self.model_id = 'damo/multi-modal_efficient-diffusion-tuning-control-lora'
+        self.compatibility_check()
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/trainers/test_efficient_diffusion_tuning_trainer.py

@@ -0,0 +1,142 @@
+# Copyright 2022-2023 The Alibaba Fundamental Vision Team Authors. All rights reserved.
+import os
+import shutil
+import tempfile
+import unittest
+
+from modelscope.metainfo import Trainers
+from modelscope.msdatasets import MsDataset
+from modelscope.trainers import build_trainer
+from modelscope.utils.constant import DownloadMode
+from modelscope.utils.test_utils import test_level
+
+
+class TestEfficientDiffusionTuningTrainer(unittest.TestCase):
+
+    def setUp(self):
+        print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
+
+        self.train_dataset = MsDataset.load(
+            'controlnet_dataset_condition_fill50k',
+            namespace='damo',
+            split='train',
+            download_mode=DownloadMode.FORCE_REDOWNLOAD).to_hf_dataset(  # noqa
+            ).select(range(100))  # noqa
+        self.eval_dataset = MsDataset.load(
+            'controlnet_dataset_condition_fill50k',
+            namespace='damo',
+            split='validation',
+            download_mode=DownloadMode.FORCE_REDOWNLOAD).to_hf_dataset(  # noqa
+            ).select(range(20))  # noqa
+
+        self.max_epochs = 1
+
+        self.tmp_dir = tempfile.TemporaryDirectory().name
+        if not os.path.exists(self.tmp_dir):
+            os.makedirs(self.tmp_dir)
+
+    def tearDown(self):
+        shutil.rmtree(self.tmp_dir)
+        super().tearDown()
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_efficient_diffusion_tuning_lora_train(self):
+        model_id = 'damo/multi-modal_efficient-diffusion-tuning-lora'
+
+        def cfg_modify_fn(cfg):
+            cfg.train.max_epochs = self.max_epochs
+            cfg.train.lr_scheduler.T_max = self.max_epochs
+            cfg.model.inference = False
+            return cfg
+
+        kwargs = dict(
+            model=model_id,
+            work_dir=self.tmp_dir,
+            train_dataset=self.train_dataset,
+            eval_dataset=self.eval_dataset,
+            cfg_modify_fn=cfg_modify_fn)
+
+        trainer = build_trainer(
+            name=Trainers.efficient_diffusion_tuning, default_args=kwargs)
+        trainer.train()
+        result = trainer.evaluate()
+        print(f'Efficient-diffusion-tuning-lora train output: {result}.')
+
+        results_files = os.listdir(self.tmp_dir)
+        self.assertIn(f'{trainer.timestamp}.log.json', results_files)
+        for i in range(self.max_epochs):
+            self.assertIn(f'epoch_{i+1}.pth', results_files)
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_efficient_diffusion_tuning_lora_eval(self):
+        model_id = 'damo/multi-modal_efficient-diffusion-tuning-lora'
+
+        def cfg_modify_fn(cfg):
+            cfg.model.inference = False
+            return cfg
+
+        kwargs = dict(
+            model=model_id,
+            work_dir=self.tmp_dir,
+            train_dataset=None,
+            eval_dataset=self.eval_dataset,
+            cfg_modify_fn=cfg_modify_fn)
+
+        trainer = build_trainer(
+            name=Trainers.efficient_diffusion_tuning, default_args=kwargs)
+        result = trainer.evaluate()
+        print(f'Efficient-diffusion-tuning-lora eval output: {result}.')
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_efficient_diffusion_tuning_control_lora_train(self):
+        model_id = 'damo/multi-modal_efficient-diffusion-tuning-control-lora'
+
+        def cfg_modify_fn(cfg):
+            cfg.train.max_epochs = self.max_epochs
+            cfg.train.lr_scheduler.T_max = self.max_epochs
+            cfg.model.inference = False
+            return cfg
+
+        kwargs = dict(
+            model=model_id,
+            work_dir=self.tmp_dir,
+            train_dataset=self.train_dataset,
+            eval_dataset=self.eval_dataset,
+            cfg_modify_fn=cfg_modify_fn)
+
+        trainer = build_trainer(
+            name=Trainers.efficient_diffusion_tuning, default_args=kwargs)
+        trainer.train()
+        result = trainer.evaluate()
+        print(
+            f'Efficient-diffusion-tuning-control-lora train output: {result}.')
+
+        results_files = os.listdir(self.tmp_dir)
+        self.assertIn(f'{trainer.timestamp}.log.json', results_files)
+        for i in range(self.max_epochs):
+            self.assertIn(f'epoch_{i+1}.pth', results_files)
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_efficient_diffusion_tuning_control_lora_eval(self):
+        model_id = 'damo/multi-modal_efficient-diffusion-tuning-control-lora'
+
+        def cfg_modify_fn(cfg):
+            cfg.model.inference = False
+            return cfg
+
+        kwargs = dict(
+            model=model_id,
+            work_dir=self.tmp_dir,
+            train_dataset=None,
+            eval_dataset=self.eval_dataset,
+            cfg_modify_fn=cfg_modify_fn)
+
+        trainer = build_trainer(
+            name=Trainers.efficient_diffusion_tuning, default_args=kwargs)
+        result = trainer.evaluate()
+        print(
+            f'Efficient-diffusion-tuning-control-lora eval output: {result}.')
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/tuners/test_lora.py

@@ -0,0 +1,120 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import os
+import shutil
+import tempfile
+import unittest
+
+import numpy as np
+import torch
+
+from modelscope.hub.snapshot_download import snapshot_download
+from modelscope.models.base import Model
+from modelscope.msdatasets import MsDataset
+from modelscope.pipelines import pipeline
+from modelscope.trainers import build_trainer
+from modelscope.tuners.lora import (Linear, LoRATuner,
+                                    mark_only_lora_as_trainable)
+from modelscope.utils.constant import ModelFile, Tasks
+from modelscope.utils.test_utils import test_level
+
+
+class TestLora(unittest.TestCase):
+
+    def setUp(self):
+        print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
+        self.tmp_dir = tempfile.TemporaryDirectory().name
+        if not os.path.exists(self.tmp_dir):
+            os.makedirs(self.tmp_dir)
+
+    def tearDown(self):
+        shutil.rmtree(self.tmp_dir)
+        super().tearDown()
+
+    @unittest.skipUnless(test_level() >= 0, 'skip in this level')
+    def test_lora_base(self):
+
+        class TestModel(torch.nn.Module):
+
+            def __init__(self):
+                super().__init__()
+                self.lora = Linear(16, 16, r=4)
+
+        model = TestModel()
+        mark_only_lora_as_trainable(model)
+        model.train()
+        loss = model.lora(torch.ones(16, 16))
+        loss = loss.sum()
+        loss.backward()
+
+        model = TestModel()
+        mark_only_lora_as_trainable(model)
+        model.eval()
+        loss = model.lora(torch.ones(16, 16))
+        loss = loss.sum()
+        try:
+            loss.backward()
+        except Exception:
+            pass
+        else:
+            raise Exception('No tensor needs grad, should throw en error here')
+
+    @unittest.skipUnless(test_level() >= 0, 'skip in this level')
+    def test_lora_smoke_test(self):
+        dataset = MsDataset.load(
+            'clue', subset_name='afqmc',
+            split='train').to_hf_dataset().select(range(2))
+
+        model_dir = snapshot_download(
+            'damo/nlp_structbert_sentence-similarity_chinese-tiny')
+        model = Model.from_pretrained(
+            'damo/nlp_structbert_sentence-similarity_chinese-tiny',
+            adv_grad_factor=None)
+
+        cfg_file = os.path.join(model_dir, 'configuration.json')
+
+        kwargs = dict(
+            model=model,
+            cfg_file=cfg_file,
+            train_dataset=dataset,
+            eval_dataset=dataset,
+            work_dir=self.tmp_dir,
+            efficient_tuners=[{
+                'type': 'lora',
+                'replace_modules': ['query', 'key', 'value']
+            }])
+
+        trainer = build_trainer(default_args=kwargs)
+        trainer.train()
+        output_dir = os.path.join(self.tmp_dir, ModelFile.TRAIN_OUTPUT_DIR)
+
+        def pipeline_sentence_similarity(model_dir):
+            model = Model.from_pretrained(model_dir)
+            LoRATuner.tune(model, replace_modules=['query', 'key', 'value'])
+            model.load_state_dict(
+                torch.load(os.path.join(output_dir, 'pytorch_model.bin')))
+            model.eval()
+            pipeline_ins = pipeline(
+                task=Tasks.sentence_similarity, model=model)
+            return pipeline_ins(input=('test', 'this is a test'))
+
+        output1 = pipeline_sentence_similarity(
+            'damo/nlp_structbert_sentence-similarity_chinese-tiny')
+
+        LoRATuner.unpatch_lora(model, ['query', 'key', 'value'])
+        model.save_pretrained(
+            output_dir, save_checkpoint_names='pytorch_model.bin')
+
+        def pipeline_sentence_similarity_origin():
+            model = Model.from_pretrained(output_dir)
+            model.eval()
+            pipeline_ins = pipeline(
+                task=Tasks.sentence_similarity, model=model)
+            return pipeline_ins(input=('test', 'this is a test'))
+
+        output2 = pipeline_sentence_similarity_origin()
+        print(output1, output2)
+        self.assertTrue(all(np.isclose(output1['scores'], output2['scores'])))
+
+
+if __name__ == '__main__':
+    unittest.main()