vllm-omni/vllm_omni/diffusion/lora/loader.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

import os
from collections import OrderedDict, defaultdict
from collections.abc import Callable

import torch
from diffusers.loaders.lora_conversion_utils import (
    _convert_non_diffusers_qwen_lora_to_diffusers,
    _convert_non_diffusers_wan_lora_to_diffusers,
)
from safetensors.torch import load_file
from vllm.logger import init_logger

from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.utils.tf_utils import get_transformer_from_pipeline
from vllm_omni.transformers_utils.repo_utils import hf_api

logger = init_logger(__name__)

lora_convert_mapping: dict[str, Callable] = {
    "QwenImagePipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
    "QwenImageEditPipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
    "QwenImageEditPlusPipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
    # 2.1's single-stream blocks keep the same leaf names the converter's
    # protected n-grams cover (`attn.to_q/to_k/to_v`, `img_mlp`, ...), so the
    # 2.0 converter applies. Non-diffusers LoRAs that target the fused `to_qkv`
    # projection are not supported (the converter would split it into `to.qkv`).
    "QwenImage21Pipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
    "Wan22Pipeline": _convert_non_diffusers_wan_lora_to_diffusers,
    "Wan22I2VPipeline": _convert_non_diffusers_wan_lora_to_diffusers,
}


def get_converter_by_pipeline(pipeline):
    for pipeline_cls in pipeline.__class__.__mro__:
        converter = lora_convert_mapping.get(pipeline_cls.__name__)
        if converter is not None:
            return converter
    return None


def _prepare_lora_delta(
    lora_state_dict,
    base_key: str,
    param_to_weight_names: dict[str, list[str]] | None = None,
    is_bias: bool = False,
    lora_a_suffix: str = "lora_A.weight",
    lora_b_suffix: str = "lora_B.weight",
    lora_bias_suffix: str = "bias",
    compute_device: torch.device | None = None,
):
    used_keys: set[str] = set()
    # stacked_params_mapping                       param_to_weight_names
    # [(".to_qkv", ".to_q.", "q")
    # (".to_qkv", ".to_k.", "k")  ========> {".to_qkv": [".to_q", ".to_k", ".to_v"]}
    # (".to_qkv", ".to_v.", "v")]

    # stacked_sd is to store packed parameter deltas (format: [str, list[torch.Tensor]])
    # example: {".to_qkv": [delta_q, delta_k, delta_v]}
    stacked_deltas = []
    is_stacked_param = False
    if param_to_weight_names is None:
        param_to_weight_names = defaultdict(list)
    for param_name, weight_names in param_to_weight_names.items():
        # Fused names can overlap by prefix (`.qkv_proj` is a substring of
        # `.qkv_proj_mot_gen`), and a substring match stacks both mappings into
        # one delta of twice the height.
        if not base_key.endswith(param_name):
            continue
        is_stacked_param = True
        # handle lora_a_key and lora_b_key together
        if is_bias:
            for weight_name in weight_names:
                lora_bias_key = f"{base_key.replace(param_name, weight_name)}.{lora_bias_suffix}"
                if lora_bias_key not in lora_state_dict:
                    return None, used_keys
                delta = lora_state_dict[lora_bias_key]
                if compute_device is not None and compute_device.type != "cpu":
                    delta = delta.to(compute_device)
                stacked_deltas.append(delta)
                used_keys.add(lora_bias_key)
        else:
            for weight_name in weight_names:
                lora_a_key = f"{base_key.replace(param_name, weight_name)}.{lora_a_suffix}"
                lora_b_key = lora_a_key.replace(lora_a_suffix, lora_b_suffix)
                if lora_a_key not in lora_state_dict or lora_b_key not in lora_state_dict:
                    return None, used_keys
                a = lora_state_dict[lora_a_key]
                b = lora_state_dict[lora_b_key]
                if compute_device is not None and compute_device.type != "cpu":
                    a = a.to(compute_device)
                    b = b.to(compute_device)
                delta = torch.matmul(b, a)
                stacked_deltas.append(delta)
                used_keys.add(lora_a_key)
                used_keys.add(lora_b_key)
        break

    if is_stacked_param:
        return torch.concat(stacked_deltas), used_keys

    if is_bias:
        lora_bias_key = f"{base_key}.{lora_bias_suffix}"
        if lora_bias_key not in lora_state_dict:
            return None, used_keys
        used_keys.add(lora_bias_key)
        delta = lora_state_dict[lora_bias_key]
        if compute_device is not None and compute_device.type != "cpu":
            delta = delta.to(compute_device)
        return delta, used_keys

    lora_a_key = f"{base_key}.{lora_a_suffix}"
    lora_b_key = f"{base_key}.{lora_b_suffix}"
    if lora_a_key not in lora_state_dict or lora_b_key not in lora_state_dict:
        return None, used_keys
    a = lora_state_dict[lora_a_key]
    b = lora_state_dict[lora_b_key]
    if compute_device is not None and compute_device.type != "cpu":
        a = a.to(compute_device)
        b = b.to(compute_device)
    delta = torch.matmul(b, a)
    used_keys.add(lora_a_key)
    used_keys.add(lora_b_key)
    return delta, used_keys


def _resolve_lora_compute_device() -> torch.device | None:
    """Return the local accelerator device if actually usable, else None.

    get_local_device() may return cuda:N/npu:N even when no physical device is
    present, so probe with a tiny allocation before trusting the lookup.
    """
    try:
        dev = get_local_device()
        if dev.type != "cpu":
            torch.empty(1, device=dev)
            return dev
    except Exception:
        pass
    return None


def _load_lora_state_dict(
    pretrained_model_name_or_path: str,
    weights_name: str | None = None,
    use_safetensors: bool = True,
    subfolder: str | None = None,
    cache_dir: str | None = None,
    force_download: bool = False,
    local_files_only: bool | None = None,
):
    # first we try to load it from local disk
    if use_safetensors and pretrained_model_name_or_path.endswith(".safetensors"):
        return load_file(pretrained_model_name_or_path)

    if os.path.isdir(pretrained_model_name_or_path):
        model_file = pretrained_model_name_or_path
        if subfolder:
            model_file = os.path.join(model_file, subfolder)
        if weights_name is None:
            raise ValueError("weights_name is required when loading from a directory")
        model_file = os.path.join(model_file, weights_name)
        return load_file(model_file)

    # finally, we try to load it from the internet
    try:
        model_file = hf_api().hf_hub_download(
            pretrained_model_name_or_path,
            filename=weights_name,
            subfolder=subfolder,
            cache_dir=cache_dir,
            force_download=force_download,
            local_files_only=local_files_only,
        )
        return load_file(model_file)
    except Exception as e:
        logger.error(f"Failed to download {weights_name} from {pretrained_model_name_or_path}: {e}")
        raise e


def _remap_state_dict_keys(sd, rules):
    """
    Remap keys in a state_dict by sequentially applying a list of substring substitution rules.

    This utility function transforms parameter names by applying a sequence of
    (old_substring, new_substring) pairs. For each key in the input dictionary,
    **all rules are evaluated in the given order**, and every rule whose
    `old_substring` is found within the current key triggers a replacement.
    Substitutions are cumulative: the key is updated after each match, and
    subsequent rules operate on the result of previous replacements.

    The order of rules matters; users are responsible for arranging them to
    achieve the desired final key. Keys that do not match any rule are copied
    unchanged.

    Args:
        sd (dict[str, Any]): Original state_dict mapping parameter names to tensors.
        rules (List[Tuple[str, str]]): A list of (old_substring, new_substring) pairs.
            Rules are applied sequentially, and **all matching rules are executed**.

    Returns:
        dict[str, Any]: A new state_dict with remapped keys. Values are unchanged.
    """
    new_sd = OrderedDict()
    for k, v in sd.items():
        new_key = k
        matched = False
        for p1, p2 in rules:
            if p1 not in k:
                continue
            new_key = new_key.replace(p1, p2)
            matched = True
        if matched:
            new_sd[new_key] = v
        else:
            new_sd[k] = v
    return new_sd


def _apply_diffusers_lora_alpha_scaling(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    """Fold Diffusers-format LoRA alpha values into their B matrices."""
    scaled_state_dict = dict(state_dict)
    for alpha_key in (key for key in state_dict if key.endswith(".alpha")):
        base_key = alpha_key[: -len(".alpha")]
        lora_a_key = f"{base_key}.lora_A.weight"
        lora_b_key = f"{base_key}.lora_B.weight"
        if lora_a_key not in state_dict or lora_b_key not in state_dict:
            raise ValueError(f"LoRA alpha key {alpha_key!r} does not have matching lora_A and lora_B weights.")

        rank = state_dict[lora_a_key].shape[0]
        alpha = state_dict[alpha_key]

        scaled_state_dict[lora_b_key] = state_dict[lora_b_key] * (alpha.item() / rank)
        scaled_state_dict.pop(alpha_key)

    return scaled_state_dict


class LoraLoaderMixin:
    transformer_name = "transformer"

    # Set by the pipeline this mixin is combined with; annotated, not assigned.
    transformer: torch.nn.Module
    _lora_loaded: dict[str | None, dict[str, torch.Tensor]]
    _lora_is_fused: bool

    # Lazy initialization to avoid MRO issues: __init__ may not be called
    # when mixin with nn.Module
    @property
    def lora_loaded(self):
        if not hasattr(self, "_lora_loaded"):
            self._lora_loaded = {}
        return self._lora_loaded

    @lora_loaded.setter
    def lora_loaded(self, value):
        self._lora_loaded = value

    @property
    def lora_is_fused(self) -> bool:
        """True when LoRA weights are fused into base weights.

        Mixin loads are in-place fusions, so non-empty `_lora_loaded` implies fused.
        """
        return getattr(self, "_lora_is_fused", False) or bool(getattr(self, "_lora_loaded", None))

    @lora_is_fused.setter
    def lora_is_fused(self, value: bool):
        self._lora_is_fused = value

    @classmethod
    def load_lora_into_module(
        cls,
        state_dict,
        module,
        prefix: str = "transformer",
        lora_a_suffix: str = "lora_A.weight",
        lora_b_suffix: str = "lora_B.weight",
        lora_bias_suffix: str = "bias",
    ):
        lora_loaded_keys = set()

        # prepare stacked parameters mapping for later use
        # stacked_params_mapping                       param_to_weight_names
        # [(".to_qkv", ".to_q.", "q")
        # (".to_qkv", ".to_k.", "k")  ========> {".to_qkv": [".to_q", ".to_k", ".to_v"]}
        # (".to_qkv", ".to_v.", "v")]
        param_to_weight_names = defaultdict(list)
        if hasattr(module, "stacked_params_mapping"):
            for param_name, weight_name, _ in module.stacked_params_mapping:
                param_to_weight_names[param_name].append(weight_name)

        # Resolve the accelerator fallback once, outside the parameter loop.
        accel_dev = _resolve_lora_compute_device()

        for name, params in module.named_parameters(prefix):
            is_bias = False
            if name.endswith(".bias"):
                base_key = name[: -len(".bias")]
                is_bias = True
            elif name.endswith(".weight"):
                base_key = name[: -len(".weight")]
            else:
                continue

            compute_dev = params.device if params.device.type != "cpu" else accel_dev
            delta, used_keys = _prepare_lora_delta(
                state_dict,
                base_key,
                param_to_weight_names,
                is_bias,
                lora_a_suffix,
                lora_b_suffix,
                lora_bias_suffix,
                compute_device=compute_dev,
            )
            if delta is None:
                continue
            lora_loaded_keys.update(used_keys)

            delta = delta.to(device=params.device, dtype=params.dtype)
            with torch.no_grad():
                params.add_(delta)
            del delta

        missing_keys = set(state_dict.keys()) - lora_loaded_keys
        for k in missing_keys:
            logger.warning(f"Missing loading lora key: {k}")

        logger.info(f"{len(lora_loaded_keys)} lora keys loaded into {module.__class__.__name__}.")
        return lora_loaded_keys

    @classmethod
    def unload_module_lora(
        cls,
        state_dict,
        module,
        prefix: str = "transformer",
        lora_a_suffix: str = "lora_A.weight",
        lora_b_suffix: str = "lora_B.weight",
        lora_bias_suffix: str = "bias",
    ):
        lora_unloaded_keys = set()

        param_to_weight_names = defaultdict(list)
        if hasattr(module, "stacked_params_mapping"):
            for param_name, weight_name, _ in module.stacked_params_mapping:
                param_to_weight_names[param_name].append(weight_name)

        # Resolve the accelerator fallback once, outside the parameter loop.
        accel_dev = _resolve_lora_compute_device()

        for name, param in module.named_parameters(prefix):
            is_bias = False
            if name.endswith(".bias"):
                base_key = name[: -len(".bias")]
                is_bias = True
            elif name.endswith(".weight"):
                base_key = name[: -len(".weight")]
            else:
                continue

            compute_dev = param.device if param.device.type != "cpu" else accel_dev
            delta, used_keys = _prepare_lora_delta(
                state_dict,
                base_key,
                param_to_weight_names,
                is_bias,
                lora_a_suffix,
                lora_b_suffix,
                lora_bias_suffix,
                compute_device=compute_dev,
            )
            if delta is None:
                continue
            lora_unloaded_keys.update(used_keys)

            delta = delta.to(device=param.device, dtype=param.dtype)
            with torch.no_grad():
                param.sub_(delta)
            del delta

        missing_keys = set(state_dict.keys()) - lora_unloaded_keys
        for k in missing_keys:
            logger.warning(f"Missing unloading lora key: {k}")

        logger.info(f"{len(lora_unloaded_keys)} lora keys unloaded from {module.__class__.__name__}.")
        return lora_unloaded_keys


class QwenImageLoraLoaderMixin(LoraLoaderMixin):
    def load_lora_weights(
        self,
        pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
        adapter_name: str | None = None,
    ):
        if adapter_name in self.lora_loaded:
            return

        if isinstance(pretrained_model_name_or_path_or_dict, dict):
            state_dict = pretrained_model_name_or_path_or_dict
        else:
            state_dict = _load_lora_state_dict(pretrained_model_name_or_path_or_dict)

        has_alpha = any(k.endswith(".alpha") for k in state_dict)
        is_non_diffusers_format = any(k.startswith("diffusion_model.") for k in state_dict)
        is_diffusers_lora_format = any(k.endswith(".lora_A.weight") for k in state_dict)
        if has_alpha and is_diffusers_lora_format:
            state_dict = _apply_diffusers_lora_alpha_scaling(state_dict)

        if has_alpha or is_non_diffusers_format:
            converter = get_converter_by_pipeline(self)
            if converter is None:
                raise ValueError(f"Converter for Lora weights not found for {self.__class__.__name__}")

            state_dict = converter(state_dict)

        state_dict = _remap_state_dict_keys(state_dict, [(".to_out.0.", ".to_out.")])

        self.load_lora_into_module(
            state_dict,
            self.transformer,
            prefix=self.transformer_name,
        )

        self.lora_loaded[adapter_name] = state_dict

    def unload_lora_weights(self, adapter_name: str):
        if adapter_name not in self.lora_loaded:
            return

        state_dict = self.lora_loaded[adapter_name]

        transformer = get_transformer_from_pipeline(self)
        self.unload_module_lora(state_dict, transformer, prefix=self.transformer_name)

        del self.lora_loaded[adapter_name]


class WanLoraLoaderMixin(LoraLoaderMixin):
    transformer_2_name = "transformer_2"
    has_transformer_2: bool

    def load_lora_weights(
        self,
        pretrained_model_name_or_path: str | list[str],
        adapter_name: str | None = None,
    ):
        """Load LoRA weights into the pipeline's transformer(s).

        Args:
            pretrained_model_name_or_path: Path(s) to concrete .safetensors LoRA files.
                - str: a single .safetensors file. Only valid for Wan2.1
                  (single-transformer pipelines).
                - list[str]: one file per transformer, matched **by position** to
                  [transformer, transformer_2]. Required for Wan2.2 MoE.
            adapter_name: Name to register the adapter under.
        """
        if adapter_name in self.lora_loaded:
            return

        # Normalize to list[(lora_path, target_module_name)].
        lora_paths = None
        if isinstance(pretrained_model_name_or_path, str):
            lora_paths = [pretrained_model_name_or_path]
        else:
            lora_paths = list(pretrained_model_name_or_path)

        target_modules = (
            [self.transformer_name, self.transformer_2_name] if self.has_transformer_2 else [self.transformer_name]
        )
        if len(lora_paths) != len(target_modules):
            raise ValueError(
                f"{self.__class__.__name__} expects {len(target_modules)} LoRA file(s) "
                f"(target_modules={target_modules}, paths={lora_paths}), got {len(lora_paths)}. Pass a list of concrete "  # noqa
                f".safetensors files, one per transformer."
            )

        module_to_lora_sd = dict()
        for lora_path, target_module_name in zip(lora_paths, target_modules):
            state_dict = _load_lora_state_dict(lora_path)
            is_non_diffusers_format = any(k.startswith("diffusion_model.") for k in state_dict)
            if is_non_diffusers_format:
                converter = get_converter_by_pipeline(self)
                if converter is None:
                    raise ValueError(f"Converter for Lora weights not found for {self.__class__.__name__}")

                state_dict = converter(state_dict)

            state_dict = _remap_state_dict_keys(
                state_dict,
                [
                    (".ffn.net.0.", ".ffn.net_0."),
                    (".ffn.net.2.", ".ffn.net_2."),
                    (
                        ".to_out.0.",
                        ".to_out.",
                    ),
                    # '.bias_B.bias' is a horrible name in diffuser's.
                    # So we remap it to '.bias'
                    (".lora_B.bias", ".bias"),
                ],
            )

            module = get_transformer_from_pipeline(self, transformer_name=target_module_name)
            if module is None:
                logger.warning(
                    f"Skip LoRA {lora_path}: target module '{target_module_name}' "
                    f"is not loaded in pipeline {self.__class__.__name__} "
                    f"(e.g. disabled by boundary_ratio)."
                )
                continue

            self.load_lora_into_module(state_dict, module, prefix=self.transformer_name, lora_bias_suffix="bias")
            module_to_lora_sd[target_module_name] = state_dict

        self.lora_loaded[adapter_name] = module_to_lora_sd

    def unload_lora_weights(self, adapter_name: str):
        if adapter_name not in self.lora_loaded:
            return

        module_to_lora_sd = self.lora_loaded[adapter_name]
        for module_name, lora_sd in module_to_lora_sd.items():
            module = get_transformer_from_pipeline(self, transformer_name=module_name)
            self.unload_module_lora(lora_sd, module, prefix=self.transformer_name)

        del self.lora_loaded[adapter_name]