From 93270ec2d7e3e201d169cee33c863a74e4efb850 Mon Sep 17 00:00:00 2001
From: Pierre Chapuis <git@catwell.info>
Date: Tue, 30 Jan 2024 14:01:09 +0100
Subject: [PATCH] add inject / eject test for t2i adapter

---
 tests/adapters/test_t2i_adapter.py | 52 ++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)
 create mode 100644 tests/adapters/test_t2i_adapter.py

diff --git a/tests/adapters/test_t2i_adapter.py b/tests/adapters/test_t2i_adapter.py
new file mode 100644
index 0000000..3e97e6d
--- /dev/null
+++ b/tests/adapters/test_t2i_adapter.py
@@ -0,0 +1,52 @@
+from typing import overload
+
+import pytest
+import torch
+
+from refiners.fluxion.utils import no_grad
+from refiners.foundationals.latent_diffusion import SD1T2IAdapter, SD1UNet, SDXLT2IAdapter, SDXLUNet
+from refiners.foundationals.latent_diffusion.t2i_adapter import T2IFeatures
+
+
+@overload
+def new_adapter(target: SD1UNet, name: str) -> SD1T2IAdapter:
+    ...
+
+
+@overload
+def new_adapter(target: SDXLUNet, name: str) -> SDXLT2IAdapter:
+    ...
+
+
+def new_adapter(target: SD1UNet | SDXLUNet, name: str) -> SD1T2IAdapter | SDXLT2IAdapter:
+    if isinstance(target, SD1UNet):
+        return SD1T2IAdapter(target=target, name=name)
+    else:
+        return SDXLT2IAdapter(target=target, name=name)
+
+
+@no_grad()
+@pytest.mark.parametrize("k_unet", [SD1UNet, SDXLUNet])
+def test_inject_eject(k_unet: type[SD1UNet] | type[SDXLUNet], test_device: torch.device):
+    unet = k_unet(in_channels=4, device=test_device)
+    adapter_1 = new_adapter(unet, "adapter 1").inject()
+
+    with pytest.raises(AssertionError) as already_injected_error:
+        new_adapter(unet, "adapter 1").inject()
+
+    assert str(already_injected_error.value) == "T2I-Adapter named adapter 1 is already injected"
+
+    adapter_2 = new_adapter(unet, "adapter 2").inject()
+
+    adapter_1.eject()
+
+    new_adapter_1 = new_adapter(unet, "adapter 1").inject()
+    new_adapter_1.eject()
+
+    assert unet.parent == adapter_2
+    assert unet.find(T2IFeatures) is not None
+
+    adapter_2.eject()
+
+    assert unet.parent is None
+    assert unet.find(T2IFeatures) is None