Method apply_lora

lightx2v/utils/lora_loader.py:341–452

Apply LoRA weights to model weights. Args: weight_dict: The model weights dictionary (will be modified in place) lora_weights: The LoRA weights dictionary alpha: Global alpha scaling factor strength: Additional strength factor for LoR

(
        self,
        weight_dict: Dict[str, torch.Tensor],
        lora_weights: Dict[str, torch.Tensor],
        alpha: float = None,
        strength: float = 1.0,
    )

        return lora_diffs

    def apply_lora(
        self,
        weight_dict: Dict[str, torch.Tensor],
        lora_weights: Dict[str, torch.Tensor],
        alpha: float = None,
        strength: float = 1.0,
    ) -> int:
        """
        Apply LoRA weights to model weights.

        Args:
            weight_dict: The model weights dictionary (will be modified in place)
            lora_weights: The LoRA weights dictionary
            alpha: Global alpha scaling factor
            strength: Additional strength factor for LoRA deltas

        Returns:
            Number of LoRA weights successfully applied
        """
        # Extract LoRA pairs, diffs, and alphas
        lora_pairs = self.extract_lora_pairs(lora_weights)
        lora_diffs = self.extract_lora_diffs(lora_weights)

        applied_count = 0
        used_lora_keys = set()

        # Apply LoRA pairs (matrix multiplication)
        for model_key, pair_info in lora_pairs.items():
            if model_key not in weight_dict:
                logger.debug(f"Model key not found: {model_key}")
                continue

            param = weight_dict[model_key]
            up_key = pair_info["up_key"]
            down_key = pair_info["down_key"]

            # Track used keys
            used_lora_keys.add(up_key)
            used_lora_keys.add(down_key)
            if pair_info["mid_key"]:
                used_lora_keys.add(pair_info["mid_key"])

            try:
                lora_up = lora_weights[up_key].to(param.device, param.dtype)
                lora_down = lora_weights[down_key].to(param.device, param.dtype)

                # Get LoRA-specific alpha if available, otherwise use global alpha
                # Apply LoRA: W' = W + (alpha/rank) * B @ A
                # where B = up (out_features, rank), A = down (rank, in_features)
                if pair_info["alpha"]:
                    lora_scale = pair_info["alpha"] / lora_down.shape[0]
                elif alpha is not None:
                    lora_scale = alpha / lora_down.shape[0]
                else:
                    lora_scale = 1

                if len(lora_down.shape) == 2 and len(lora_up.shape) == 2:
                    lora_delta = torch.mm(lora_up, lora_down) * lora_scale