#!/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()