feat(lora): applied_lora ctx manager — validate, apply, revert on exit

lora.py

@@ -3,8 +3,10 @@ from __future__ import annotations
 
 import json
 import struct
+from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Any, Iterator
 
 ZIMAGE_LORA_PREFIXES = ("transformer.", "dit.", "model.transformer.")
 
@@ -68,3 +70,50 @@ def sniff(path: Path | str) -> LoRAInfo:
         target="transformer",
         size_bytes=path.stat().st_size,
     )
+
+
+@contextmanager
+def applied_lora(pipe: Any, path: Path | str | None, strength: float) -> Iterator[None]:
+    """Apply a LoRA to the pipeline's dit for the duration of the context.
+
+    Reverts on exit (including exception path) so the cached GPU model is left clean.
+    If ``path`` is ``None``, this is a no-op.
+
+    Validates the LoRA file with :func:`sniff` before touching the pipeline so a bad
+    file is rejected before any GPU work begins.
+    """
+    if path is None:
+        yield
+        return
+
+    sniff(path)  # raises LoRAValidationError on bad input
+    _apply_lora_impl(pipe, path, strength)
+    try:
+        yield
+    finally:
+        _revert_lora_impl(pipe)
+
+
+def _apply_lora_impl(pipe: Any, path: Path | str, strength: float) -> None:
+    """Apply a LoRA to ``pipe.dit``. Imports DiffSynth lazily for testability."""
+    from diffsynth.utils.lora import merge_lora
+    merge_lora(pipe.dit, str(path), alpha=float(strength))
+
+
+def _revert_lora_impl(pipe: Any) -> None:
+    """Revert the most recent LoRA from ``pipe.dit``.
+
+    Tries DiffSynth's ``unmerge_lora`` first; falls back to re-fetching clean
+    weights from the model pool if unavailable.
+    """
+    try:
+        from diffsynth.utils.lora import unmerge_lora
+        unmerge_lora(pipe.dit)
+        return
+    except ImportError:
+        pass
+
+    if hasattr(pipe, "model_pool"):
+        variant = getattr(pipe.dit, "_zis_variant", None)
+        if variant:
+            pipe.dit = pipe.model_pool.fetch_model("z_image_dit", variant=variant)

tests/test_lora.py

@@ -43,3 +43,61 @@ def test_sniff_rejects_non_zimage_keys(tmp_path):
         lora.sniff(p)
     msg = str(exc.value).lower()
     assert "down_blocks" in msg or "unexpected" in msg
+
+
+class _FakePipe:
+    """Minimal stand-in for DiffSynth's ZImagePipeline.dit hook surface."""
+    def __init__(self):
+        self.applied = []   # list of (path, strength) tuples
+        self.reverted = []
+
+
+def test_applied_lora_calls_apply_then_revert(tmp_path, monkeypatch):
+    p = tmp_path / "ok.safetensors"
+    _write_safetensors(p, {
+        "transformer.x.lora_A.weight": {"dtype": "BF16", "shape": [32, 3840]},
+        "transformer.x.lora_B.weight": {"dtype": "BF16", "shape": [3840, 32]},
+    })
+    pipe = _FakePipe()
+
+    def fake_apply(pipe, path, strength):
+        pipe.applied.append((str(path), strength))
+    def fake_revert(pipe):
+        pipe.reverted.append(True)
+    monkeypatch.setattr(lora, "_apply_lora_impl", fake_apply)
+    monkeypatch.setattr(lora, "_revert_lora_impl", fake_revert)
+
+    with lora.applied_lora(pipe, p, strength=0.8):
+        assert pipe.applied == [(str(p), 0.8)]
+        assert pipe.reverted == []
+
+    assert pipe.reverted == [True]
+
+
+def test_applied_lora_with_none_is_a_noop(tmp_path, monkeypatch):
+    pipe = _FakePipe()
+    sentinel = []
+    monkeypatch.setattr(lora, "_apply_lora_impl", lambda *a, **k: sentinel.append("apply"))
+    monkeypatch.setattr(lora, "_revert_lora_impl", lambda *a, **k: sentinel.append("revert"))
+
+    with lora.applied_lora(pipe, None, strength=0.0):
+        pass
+
+    assert sentinel == []
+
+
+def test_applied_lora_reverts_on_exception(tmp_path, monkeypatch):
+    p = tmp_path / "ok.safetensors"
+    _write_safetensors(p, {
+        "transformer.x.lora_A.weight": {"dtype": "BF16", "shape": [16, 3840]},
+        "transformer.x.lora_B.weight": {"dtype": "BF16", "shape": [3840, 16]},
+    })
+    pipe = _FakePipe()
+    monkeypatch.setattr(lora, "_apply_lora_impl", lambda pipe, p, s: pipe.applied.append((p, s)))
+    monkeypatch.setattr(lora, "_revert_lora_impl", lambda pipe: pipe.reverted.append(True))
+
+    with pytest.raises(RuntimeError):
+        with lora.applied_lora(pipe, p, strength=1.0):
+            raise RuntimeError("inference failed mid-step")
+
+    assert pipe.reverted == [True], "must still revert on exception"