diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index edec6ce2..6d3c7994 100644 --- a/modelscope/metainfo.py +++ b/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' diff --git a/modelscope/metrics/builder.py b/modelscope/metrics/builder.py index 2bd64613..2bc756e6 100644 --- a/modelscope/metrics/builder.py +++ b/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], } diff --git a/modelscope/models/multi_modal/__init__.py b/modelscope/models/multi_modal/__init__.py index 8bf9f018..e85c48fb 100644 --- a/modelscope/models/multi_modal/__init__.py +++ b/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 diff --git a/modelscope/models/multi_modal/efficient_diffusion_tuning/__init__.py b/modelscope/models/multi_modal/efficient_diffusion_tuning/__init__.py new file mode 100644 index 00000000..42e2b700 --- /dev/null +++ b/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={}, + ) diff --git a/modelscope/models/multi_modal/efficient_diffusion_tuning/efficient_stable_diffusion.py b/modelscope/models/multi_modal/efficient_diffusion_tuning/efficient_stable_diffusion.py new file mode 100644 index 00000000..747aecd8 --- /dev/null +++ b/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 diff --git a/modelscope/msdatasets/ms_dataset.py b/modelscope/msdatasets/ms_dataset.py index 06f47874..492f0169 100644 --- a/modelscope/msdatasets/ms_dataset.py +++ b/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, diff --git a/modelscope/pipelines/multi_modal/efficient_diffusion_tuning_pipeline.py b/modelscope/pipelines/multi_modal/efficient_diffusion_tuning_pipeline.py new file mode 100644 index 00000000..c36c373f --- /dev/null +++ b/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} diff --git a/modelscope/preprocessors/__init__.py b/modelscope/preprocessors/__init__.py index 06833b82..a35f130a 100644 --- a/modelscope/preprocessors/__init__.py +++ b/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', diff --git a/modelscope/preprocessors/cv/__init__.py b/modelscope/preprocessors/cv/__init__.py index b832f1e6..e02869cb 100644 --- a/modelscope/preprocessors/cv/__init__.py +++ b/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 diff --git a/modelscope/preprocessors/multi_modal.py b/modelscope/preprocessors/multi_modal.py index a4f77684..bd37c620 100644 --- a/modelscope/preprocessors/multi_modal.py +++ b/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( diff --git a/modelscope/trainers/multi_modal/efficient_diffusion_tuning/__init__.py b/modelscope/trainers/multi_modal/efficient_diffusion_tuning/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/trainers/multi_modal/efficient_diffusion_tuning/efficient_diffusion_tuning_trainer.py b/modelscope/trainers/multi_modal/efficient_diffusion_tuning/efficient_diffusion_tuning_trainer.py new file mode 100644 index 00000000..6ec9a2d7 --- /dev/null +++ b/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}.') diff --git a/modelscope/trainers/trainer.py b/modelscope/trainers/trainer.py index 74281564..683ff2f5 100644 --- a/modelscope/trainers/trainer.py +++ b/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 """ diff --git a/modelscope/tuners/__init__.py b/modelscope/tuners/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/tuners/control_sd_lora.py b/modelscope/tuners/control_sd_lora.py new file mode 100644 index 00000000..2585daa1 --- /dev/null +++ b/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 diff --git a/modelscope/tuners/lora.py b/modelscope/tuners/lora.py new file mode 100644 index 00000000..ba1e92e1 --- /dev/null +++ b/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 diff --git a/modelscope/tuners/sd_lora.py b/modelscope/tuners/sd_lora.py new file mode 100644 index 00000000..d740f7c0 --- /dev/null +++ b/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)) diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index c01b6e6b..2382825a 100644 --- a/modelscope/utils/constant.py +++ b/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): diff --git a/requirements/multi-modal.txt b/requirements/multi-modal.txt index 49b79f2c..8cabb1f3 100644 --- a/requirements/multi-modal.txt +++ b/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 diff --git a/tests/pipelines/test_efficient_diffusion_tuning.py b/tests/pipelines/test_efficient_diffusion_tuning.py new file mode 100644 index 00000000..9dc5e412 --- /dev/null +++ b/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() diff --git a/tests/trainers/test_efficient_diffusion_tuning_trainer.py b/tests/trainers/test_efficient_diffusion_tuning_trainer.py new file mode 100644 index 00000000..2484e24d --- /dev/null +++ b/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() diff --git a/tests/tuners/__init__.py b/tests/tuners/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/tuners/test_lora.py b/tests/tuners/test_lora.py new file mode 100644 index 00000000..2f52a4d3 --- /dev/null +++ b/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()