Stable Diffusion动态加载Lora过程中的实验、原理与说明

对比实验

显存占用情况

使用StableDiffusionXLPipeline.from_pretrained() 方法SDXL半精度加载显存占用约7G左右。

使用load_lora_weights()加载了5个Lora模型后显存提升到8G，平均每个Lora的大小在200M左右。
使用unload_lora_weights()后显存没有发生变化，还是8G，说明该方法不会清空已经加载到显存的Lora模型，但这时候再调用模型生成图片已经丢失Lora的效果了。

推理耗时

这里使用的Lora平均每个的大小在200M左右，从上表不难发现单个Lora耗时约增加4秒左右。

代码分析与原理说明

1）加载Lora

通过调用load_lora_weights()来加载不同的Lora权重，这些权重的张量都会加载到显存中，但注意只有第一次调用该方法的Lora才会生效，可通过get_active_adapters()查看。

def load_lora_weights(
        self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], adapter_name=None, **kwargs
    ):
    
        ...
        # lora_state_dict 实际执行把tensor加载到显存中，同时返回2个字典记录所添加的lora的名称和配置信息
        state_dict, network_alphas = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
        
       
        # state_dict 和 network_alphas 是上面返回的2个参数
        # 加载后默认调用的lora是第一次load进来的lora
        self.load_lora_into_unet(
            state_dict,
            network_alphas=network_alphas,
            unet=getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet,
            low_cpu_mem_usage=low_cpu_mem_usage,
            adapter_name=adapter_name,
            _pipeline=self,
        )
        self.load_lora_into_text_encoder(
            state_dict,
            network_alphas=network_alphas,
            text_encoder=getattr(self, self.text_encoder_name)
            if not hasattr(self, "text_encoder")
            else self.text_encoder,
            lora_scale=self.lora_scale,
            low_cpu_mem_usage=low_cpu_mem_usage,
            adapter_name=adapter_name,
            _pipeline=self,
        )

用一张图来表示load_lora_weights()加载过程，蓝色表示生效的张量。

因此如果不特别指定调用哪个Lora的话，默认输入的张量需要跟Base和第一个Lora的权重分别相乘再叠加

2）卸载Lora

卸载Lora时需要调用unload_lora_weights()方法，关键是把config删除，并不清显存，如代码中的del self.unet.peft_config操作。如果后续需要再调用Lora的话，则需要重头开始加载Lora权重。

  def unload_lora_weights(self):
        
        if not USE_PEFT_BACKEND:
            if version.parse(__version__) > version.parse("0.23"):
                logger.warn(
                    "You are using `unload_lora_weights` to disable and unload lora weights. If you want to iteratively enable and disable adapter weights,"
                    "you can use `pipe.enable_lora()` or `pipe.disable_lora()`. After installing the latest version of PEFT."
                )

            for _, module in self.unet.named_modules():
                if hasattr(module, "set_lora_layer"):
                    module.set_lora_layer(None)
        else:
            recurse_remove_peft_layers(self.unet)
            if hasattr(self.unet, "peft_config"):
            # 关键是把config删除，并不清显存
            # 因此执行unload_lora_weights 会丢失所有记载的lora的config信息，但不会释放显存
                del self.unet.peft_config

3）指定Lora

在拥有多个Lora加载后，如果需要指定某个或者某几个Lora，则需要用set_adapters()方法。该方法对Unet和Text Encoder都执行set_adapter()操作，间接调用了peft_utils模块中的set_weights_and_activate_adapters() 方法为Unet或者Text Encoder的某些层指定Lora类型和权重。

 def set_adapters(
        self,
        adapter_names: Union[List[str], str],
        adapter_weights: Optional[List[float]] = None,
    ):
        # 对unet和text encoder都执行set adapter操作。间接调用 set_weights_and_activate_adapters 方法
        self.unet.set_adapters(adapter_names, adapter_weights)

        # Handle the Text Encoder
        if hasattr(self, "text_encoder"):
            self.set_adapters_for_text_encoder(adapter_names, self.text_encoder, adapter_weights)
        if hasattr(self, "text_encoder_2"):
            self.set_adapters_for_text_encoder(adapter_names, self.text_encoder_2, adapter_weights)

# src/diffusers/utils/peft_utils.py
def set_weights_and_activate_adapters(model, adapter_names, weights):
    from peft.tuners.tuners_utils import BaseTunerLayer

    # iterate over each adapter, make it active and set the corresponding scaling weight
    for adapter_name, weight in zip(adapter_names, weights):
        for module in model.modules():
            if isinstance(module, BaseTunerLayer):
                # For backward compatbility with previous PEFT versions
                if hasattr(module, "set_adapter"):
                    module.set_adapter(adapter_name)
                else:
                    module.active_adapter = adapter_name
                module.set_scale(adapter_name, weight)

如果已经执行过set_adapter()，而下次又需要使用不同的Lora，则重新执行set_adapter()选择需要的Lora即可。

4）解除Lora

如果不需要Lora，想用Base模型直接生成图片，这时候可以通过disable_lora()方法来解除已经指定好的Lora。这时候再生成图片，输入的张量只会跟Base的矩阵张量相乘。

注意：后续如果需要重新使用Lora，必须先执行 enable_lora() 方法!

# src/diffusers/utils/peft_utils.py#L227
# disable_lora底层调用了该方法
# 会把Unet和Text Encoder的某些层注释掉Lora，因此达到解除Lora的效果
# 可以看到效果和指定Lora时刚好是相反的

def set_adapter_layers(model, enabled=True):
    from peft.tuners.tuners_utils import BaseTunerLayer

    for module in model.modules():
        if isinstance(module, BaseTunerLayer):
            # The recent version of PEFT needs to call `enable_adapters` instead
            if hasattr(module, "enable_adapters"):
            # false的时候注释掉 lora适配器
                module.enable_adapters(enabled=enabled)
            else:
                module.disable_adapters = not enabled

5）查看Lora

通过get_active_adapters()方法即可查看指定的Lora，即前面示例图中蓝色的Lora模块。

    def get_active_adapters(self) -> List[str]:
     
        if not USE_PEFT_BACKEND:
            raise ValueError(
                "PEFT backend is required for this method. Please install the latest version of PEFT `pip install -U peft`"
            )

        from peft.tuners.tuners_utils import BaseTunerLayer

        active_adapters = []

        for module in self.unet.modules():
            if isinstance(module, BaseTunerLayer):
                active_adapters = module.active_adapters
                break

        return active_adapters

6）融合Lora

前面已经实验过，当有多个Lora调用时，推理时间几乎会线性增加，这是因为输入的张量除了与Base的模型相乘外，还需要与每个Lora相乘再叠加，因此推理时间会变长。一种解决方案是将需要的Lora与Base进行合并，即通过fuse_lora()来实现，这样推理时间可与Base单独使用时一致。下图中深蓝色表示Base的权重已经叠加了需要的Lora的权重，因此输入张量无需再经过Lora。但是缺点是多次融合Lora后，Base将无法恢复，需要后续需要单独使用Base，只能重新加载Base！