coreml-tools/validate_stateful_multi_lora.py

#!/usr/bin/env python3
"""Validate base and exact multi-LoRA parity for Clover's Core ML U-Net."""

from __future__ import annotations

import argparse
import json
import math
import re
from pathlib import Path

import coremltools as ct
import numpy as np
import torch
from diffusers import DiffusionPipeline
from safetensors.torch import load_file

from python_coreml_stable_diffusion import unet as coreml_unet


class MultiLoRAConv2d(torch.nn.Module):
    def __init__(self, base, components):
        super().__init__()
        self.base = base
        self.downs = torch.nn.ModuleList()
        self.ups = torch.nn.ModuleList()
        self.scales = []
        for down_weight, up_weight, scale in components:
            rank, input_channels = down_weight.shape
            output_channels, up_rank = up_weight.shape
            if rank != up_rank:
                raise ValueError("LoRA ranks do not match")
            down = torch.nn.Conv2d(input_channels, rank, 1, bias=False)
            up = torch.nn.Conv2d(rank, output_channels, 1, bias=False)
            down.weight.data.copy_(down_weight[:, :, None, None])
            up.weight.data.copy_(up_weight[:, :, None, None])
            self.downs.append(down)
            self.ups.append(up)
            self.scales.append(scale)

    def forward(self, hidden_states):
        residual = sum(
            scale * up(down(hidden_states))
            for scale, down, up in zip(self.scales, self.downs, self.ups)
        )
        return self.base(hidden_states) + residual


def inject_loras(model, adapters, scales):
    states = [load_file(str(path)) for path in adapters]
    target_names = sorted(
        {
            key.removeprefix("unet.").split(".lora.")[0]
            for key in states[0]
        }
    )
    for target_name in target_names:
        prefix = f"unet.{target_name}.lora"
        down_shape = states[0][f"{prefix}.down.weight"].shape
        up_shape = states[0][f"{prefix}.up.weight"].shape
        candidates = [target_name]
        up_match = re.match(r"up_blocks\.(\d+)\.(.+)", target_name)
        if up_match:
            candidates.append(
                f"up_blocks.{int(up_match.group(1)) + 1}.{up_match.group(2)}"
            )
        resolved_name = ""
        base = None
        for candidate in candidates:
            try:
                module = model.get_submodule(candidate)
            except AttributeError:
                continue
            if (
                isinstance(module, torch.nn.Conv2d)
                and down_shape[1] == module.in_channels
                and up_shape[0] == module.out_channels
            ):
                resolved_name = candidate
                base = module
                break
        if base is None:
            raise AttributeError(f"No compatible projection for {target_name}")
        parent_name, child_name = resolved_name.rsplit(".", 1)
        parent = model.get_submodule(parent_name)
        components = [
            (
                state[f"{prefix}.down.weight"],
                state[f"{prefix}.up.weight"],
                scale,
            )
            for state, scale in zip(states, scales)
        ]
        setattr(parent, child_name, MultiLoRAConv2d(base, components))
    if len(target_names) != 72:
        raise RuntimeError(f"Expected 72 Clover LoRA targets, found {len(target_names)}")


def psnr(reference, actual):
    error = np.asarray(reference, dtype=np.float64) - np.asarray(actual, dtype=np.float64)
    rmse = math.sqrt(float(np.mean(error * error)))
    if rmse == 0:
        return math.inf
    dynamic_range = float(reference.max() - reference.min())
    return 20 * math.log10(dynamic_range / rmse)


def source_to_state(value, source_key, state_shape, slot, scale):
    output = np.zeros(state_shape, dtype=np.float32)
    value = value.numpy().astype(np.float32)
    if ".lora.down." in source_key:
        rank = value.shape[0]
        output[slot * rank : (slot + 1) * rank, :, 0, 0] = value
    elif ".lora.up." in source_key:
        rank = value.shape[1]
        output[:, slot * rank : (slot + 1) * rank, 0, 0] = value * scale
    else:
        raise ValueError(f"Unknown LoRA direction: {source_key}")
    return output


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-version", required=True)
    parser.add_argument("--coreml-model", type=Path, required=True)
    parser.add_argument("--adapter-schema", type=Path, required=True)
    parser.add_argument("--lora-weights", type=Path, action="append", required=True)
    parser.add_argument("--lora-scale", type=float, action="append")
    parser.add_argument("--minimum-psnr", type=float, default=35.0)
    args = parser.parse_args()
    scales = args.lora_scale or [1.0] * len(args.lora_weights)
    if len(scales) != len(args.lora_weights):
        raise ValueError("Provide one --lora-scale for each --lora-weights")

    schema = json.loads(args.adapter_schema.read_text())
    maximum = schema["max_adapter_count"]
    if not 1 <= len(args.lora_weights) <= maximum:
        raise ValueError(f"The model supports one to {maximum} adapters")

    torch.manual_seed(20260813)
    pipeline = DiffusionPipeline.from_pretrained(
        args.model_version,
        local_files_only=True,
        torch_dtype=torch.float16,
        variant="fp16",
        use_safetensors=True,
    )
    reference = coreml_unet.UNet2DConditionModel(**pipeline.unet.config).eval()
    reference.load_state_dict(pipeline.unet.state_dict())
    del pipeline

    sample = torch.rand(1, reference.config.in_channels, 64, 64)
    timestep = torch.tensor([981.0], dtype=torch.float32)
    hidden_states = torch.rand(1, 768, 1, 77)
    inputs = {
        "sample": sample.numpy().astype(np.float16),
        "timestep": timestep.numpy().astype(np.float16),
        "encoder_hidden_states": hidden_states.numpy().astype(np.float16),
    }
    with torch.no_grad():
        base_reference = reference(sample, timestep, hidden_states)[0].numpy()

    model = ct.models.MLModel(
        str(args.coreml_model.resolve()),
        compute_units=ct.ComputeUnit.CPU_AND_GPU,
    )
    state = model.make_state()
    base_actual = model.predict(inputs, state=state)["noise_pred"]
    base_score = psnr(base_reference, base_actual)

    adapter_states = [load_file(str(path.resolve())) for path in args.lora_weights]
    for record in schema["states"]:
        state_value = np.zeros(record["state_shape"], dtype=np.float32)
        for slot, (adapter, scale) in enumerate(zip(adapter_states, scales)):
            state_value += source_to_state(
                adapter[record["source_key"]],
                record["source_key"],
                record["state_shape"],
                slot,
                scale,
            )
        state.write_state(name=record["state_name"], value=state_value)

    inject_loras(reference, args.lora_weights, scales)
    with torch.no_grad():
        styled_reference = reference(sample, timestep, hidden_states)[0].numpy()
    styled_actual = model.predict(inputs, state=state)["noise_pred"]
    styled_score = psnr(styled_reference, styled_actual)

    result = {
        "base_psnr_db": base_score,
        "multi_lora_psnr_db": styled_score,
        "adapter_count": len(args.lora_weights),
        "maximum_adapter_count": maximum,
        "state_count": schema["state_count"],
    }
    print(json.dumps(result, indent=2))
    if min(base_score, styled_score) < args.minimum_psnr:
        raise RuntimeError(
            f"Core ML parity fell below {args.minimum_psnr:.1f} dB: {result}"
        )


if __name__ == "__main__":
    main()