models/styctrl.py

import math
import torch
import logging
import torch.nn as nn

from typing import Literal, Optional


logger = logging.getLogger(__name__)


class HadamardLayer(nn.Module):

    def __init__(
        self,
        features: int,
        device: Optional[torch.device] = None,
        dtype: Optional[torch.dtype] = None,
    ):
        super().__init__()

        self.dtype = dtype
        self.weight = nn.Parameter(torch.ones(features, device=device, dtype=dtype))
        self.bias = nn.Parameter(torch.zeros(features, device=device, dtype=dtype))

        nn.init.normal_(self.weight, mean=1.0, std=0.1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_dtype = x.dtype
        x = x.to(dtype=self.dtype)
        x = x * self.weight + self.bias
        return x.to(x_dtype)


class StyCtrlLoRA(nn.Module):

    def __init__(self, base_model: nn.Linear):
        super().__init__()

        self.base_model = base_model

        self.dtype = self.base_model.weight.dtype
        self.in_features = self.base_model.in_features
        self.out_features = self.base_model.out_features

        self.projector = nn.ModuleDict()
        self.lora_A = nn.ParameterDict()
        self.lora_B = nn.ParameterDict()
        self.bias = nn.ParameterDict()

        self.active_adapter = "styctrl"
        self.proj_types = {}

    def get_current_adapters(self) -> list[str]:
        return [k for k in self.lora_A.keys()]

    def add_adapter(
        self,
        rank: int,
        adapter_name: str = "default",
        proj_type: Literal["none", "low_rank_linear", "low_rank_scale", "in_scale"] = "none",
        bias: bool = False,
        device: torch.device = None,
        dtype: torch.dtype = None,
    ) -> None:
        if adapter_name in self.get_current_adapters():
            logger.warning(
                f"{__class__.__name__}: {adapter_name=} already exists in {__class__}, existing adapters are: {self.get_current_adapters()}."
            )
            return

        device = device or self.base_model.weight.device
        dtype = dtype or self.base_model.weight.dtype

        # Add strength projector
        if proj_type == "low_rank_linear":
            # Add linear projector in the low rank space
            self.projector[adapter_name] = nn.Linear(rank, rank, bias=True, device=device, dtype=dtype)
            nn.init.eye_(self.projector[adapter_name].weight)
            nn.init.zeros_(self.projector[adapter_name].bias)
        elif proj_type == "low_rank_scale":
            self.projector[adapter_name] = HadamardLayer(rank, device=device, dtype=dtype)
        elif proj_type == "in_scale":
            self.projector[adapter_name] = HadamardLayer(self.in_features, device=device, dtype=dtype)
        elif proj_type in ["none", "lora"]:
            self.projector[adapter_name] = nn.Identity()
        else:
            raise KeyError(f"{proj_type=} has not been implemented.")
        self.proj_types[adapter_name] = proj_type

        # Add main LoRA linear projectors
        self.lora_A[adapter_name] = nn.Parameter(torch.empty(rank, self.in_features)).to(device=device, dtype=dtype)
        self.lora_B[adapter_name] = nn.Parameter(torch.empty(self.out_features, rank)).to(device=device, dtype=dtype)

        nn.init.normal_(self.lora_A[adapter_name].data, mean=0.0, std=0.02)
        nn.init.normal_(self.lora_B[adapter_name].data, mean=0.0, std=0.02)

        # Add bias
        if bias:
            init_bias = 1 / math.sqrt(rank)
            self.bias[adapter_name] = nn.Parameter(torch.zeros((self.out_features,), dtype=self.dtype, device=device))
            nn.init.uniform_(self.bias[adapter_name], -init_bias, init_bias)

    def activate(self, adapter_name: str):
        if adapter_name not in self.get_current_adapters():
            logger.warning(
                f"{__class__.__name__}: {adapter_name=} not in current adapter list: {self.get_current_adapters()}"
            )
            return
        self.active_adapter = adapter_name

    def deactivate(self):
        self.active_adapter = None

    def forward(
        self,
        x: torch.Tensor,
        enable_lora: bool = False,
        w: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        x_dtype = x.dtype
        x = x.to(self.dtype)
        out = self.base_model(x)

        if enable_lora:
            adapter_name = self.active_adapter
            proj_type = self.proj_types[adapter_name]

            if w is not None and proj_type == "in_scale":
                w = w.to(self.dtype)
                w = w.unsqueeze(-1).repeat(1, self.in_features)
                w = self.projector[adapter_name](w)
                x = w * x

            lora_out = x @ self.lora_A[adapter_name].T

            if w is not None and proj_type == "low_rank_linear":
                w = w.to(self.dtype)
                while w.ndim < lora_out.ndim:
                    w = w.unsqueeze(-1)
                lora_out = self.projector[adapter_name](w * lora_out)

            if w is not None and proj_type == "low_rank_scale":
                w = w.to(self.dtype)
                w = w.unsqueeze(-1).repeat(1, lora_out.shape[-1])
                w = self.projector[adapter_name](w)
                lora_out = w * lora_out
            
            lora_out = lora_out @ self.lora_B[adapter_name].T

            if w is not None and proj_type == "lora":
                w = w.to(self.dtype)
                lora_out = w * lora_out

            out = out + lora_out

        out = out.to(x_dtype)
        return out


def register_styctrl(
    model: nn.Module,
    target_modules: list[str],
    device: torch.device = "cpu",
    dtype: torch.dtype = torch.bfloat16,
):
    r"""
    This method only insert LoRA adapter into model, not add or activate any adapter.
    If there are already lora_adapter in model, this method do nothing.
    """
    for name, module in model.named_modules():
        if not any(target in name for target in target_modules):
            continue
        if not isinstance(module, nn.Linear):
            continue

        splited_name = name.split(".")
        parent_node = model.get_submodule(".".join(splited_name[:-1]))
        target_name = splited_name[-1]

        lora_layer = StyCtrlLoRA(base_model=module).to(device=device, dtype=dtype)

        module.requires_grad_(False)
        setattr(parent_node, target_name, lora_layer)
        logger.debug(f"Inject LoRA into {name}.")


def add_styctrl(
    model: nn.Module,
    rank: int,
    lora_layer_indices: list[int],
    adapter_name: str = "default",
    proj_type: Literal["none", "linear", "scale"] = "none",
    bias: bool = False,
    device: Optional[torch.device] = None,
    dtype: Optional[torch.dtype] = None,
) -> None:
    for name, module in model.named_modules():
        if not any(
            f"transformer_blocks.{i}." in name or f"single_transformer_blocks.{i}." in name
            for i in lora_layer_indices
        ):
            continue
        if isinstance(module, StyCtrlLoRA):
            module.add_adapter(
                rank=rank,
                adapter_name=adapter_name,
                proj_type=proj_type,
                bias=bias,
                device=device,
                dtype=dtype,
            )


def activate(model: nn.Module, adapter_name: str) -> None:
    for n, module in model.named_modules():
        if isinstance(module, StyCtrlLoRA):
            logger.debug(f"Activate LoRA in {n}")
            module.activate(adapter_name)


def deactivate(model: nn.Module) -> None:
    for n, module in model.named_modules():
        if isinstance(module, StyCtrlLoRA):
            logger.debug(f"Deactivate LoRA in {n}")
            module.deactivate()