tests/test_lora_runtime.py

from pathlib import Path

import torch
from safetensors.torch import save_file

from lora_runtime import LoRAHookSession, LoRASource, TargetIndex


class Attention(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.q_proj = torch.nn.Linear(3, 2, bias=False)

    def forward(self, x):
        return self.q_proj(x)


class Block(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.self_attn = Attention()


class Adapter(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.out_proj = torch.nn.Linear(3, 2, bias=False)


class DummyDiT(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.blocks = torch.nn.ModuleList([Block()])
        self.llm_adapter = Adapter()


class DummyText(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = torch.nn.Module()
        self.model.layers = torch.nn.ModuleList([Block()])


class DummyPipe:
    def __init__(self):
        self.dit = DummyDiT()
        self.text_encoder = DummyText()


def test_forge_unet_alias_and_reversible_hook(tmp_path: Path):
    pipe = DummyPipe()
    module = pipe.dit.blocks[0].self_attn.q_proj
    module.weight.data.zero_()

    down = torch.tensor([[1.0, 2.0, 3.0]])
    up = torch.tensor([[4.0], [5.0]])
    path = tmp_path / "test.safetensors"
    save_file(
        {
            "lora_unet_blocks_0_self_attn_q_proj.lora_down.weight": down,
            "lora_unet_blocks_0_self_attn_q_proj.lora_up.weight": up,
            "lora_unet_blocks_0_self_attn_q_proj.alpha": torch.tensor(1.0),
        },
        str(path),
    )

    x = torch.tensor([[1.0, 1.0, 1.0]])
    baseline = pipe.dit.blocks[0].self_attn(x)
    session = LoRAHookSession(pipe)
    report = session.apply([LoRASource(str(path), strength=0.5)])
    active = pipe.dit.blocks[0].self_attn(x)
    session.clear()
    restored = pipe.dit.blocks[0].self_attn(x)

    expected = torch.tensor([[12.0, 15.0]])  # (1+2+3) * [4,5] * 0.5
    assert torch.allclose(baseline, torch.zeros_like(baseline))
    assert torch.allclose(active, expected)
    assert torch.allclose(restored, baseline)
    assert report[0].matched_pairs == 1


def test_forge_adapter_move_alias_resolves():
    pipe = DummyPipe()
    index = TargetIndex(pipe)
    target = index.resolve("text_encoders.qwen3_06b.llm_adapter.out_proj")
    assert target is not None
    assert target.scope == "dit"
    assert target.name == "llm_adapter.out_proj"


def test_forge_qwen_alias_resolves():
    pipe = DummyPipe()
    index = TargetIndex(pipe)
    target = index.resolve("lora_te_layers_0_self_attn_q_proj")
    assert target is not None
    assert target.scope == "text_encoder"
    assert target.name == "model.layers.0.self_attn.q_proj"