fix

Author: ily-R

tests/adapters/test_ella_adapter.py

@@ -1,4 +1,3 @@
-import pytest
 import torch
 
 import refiners.fluxion.layers as fl
@@ -12,9 +11,8 @@ def new_adapter(target: SD1UNet) -> SD1ELLAAdapter:
 
 
 @no_grad()
-@pytest.mark.parametrize("k_unet", [SD1UNet])
-def test_inject_eject(k_unet: type[SD1UNet], test_device: torch.device):
-    unet = k_unet(in_channels=4, device=test_device, dtype=torch.float16)
+def test_inject_eject(test_device: torch.device):
+    unet = SD1UNet(in_channels=4, device=test_device, dtype=torch.float16)
     initial_repr = repr(unet)
     adapter = new_adapter(unet)
     assert repr(unet) == initial_repr
@@ -29,9 +27,8 @@ def test_inject_eject(k_unet: type[SD1UNet], test_device: torch.device):
 
 
 @no_grad()
-@pytest.mark.parametrize("k_unet", [SD1UNet])
-def test_scale(k_unet: type[SD1UNet], test_device: torch.device):
-    unet = k_unet(in_channels=4, device=test_device, dtype=torch.float16)
+def test_scale(test_device: torch.device):
+    unet = SD1UNet(in_channels=4, device=test_device, dtype=torch.float16)
     adapter = new_adapter(unet).inject()
 
     def predicate(m: fl.Module, p: fl.Chain) -> bool:

tests/e2e/test_diffusion.py

@@ -513,12 +513,26 @@ def lda_ft_mse_weights(test_weights_path: Path) -> Path:
 
 
 @pytest.fixture(scope="module")
-def ella_adapter_weights(test_weights_path: Path) -> Path:
+def ella_weights(test_weights_path: Path) -> tuple[Path, Path]:
     ella_adapter_weights = test_weights_path / "ELLA-Adapter/ella-sd1.5-tsc-t5xl.safetensors"
     if not ella_adapter_weights.is_file():
         warn(f"could not find weights at {ella_adapter_weights}, skipping")
         pytest.skip(allow_module_level=True)
-    return ella_adapter_weights
+    t5xl_weights = test_weights_path / "QQGYLab/T5XLFP16"
+    t5xl_files = [
+        "config.json",
+        "model.safetensors",
+        "special_tokens_map.json",
+        "spiece.model",
+        "tokenizer_config.json",
+        "tokenizer.json",
+    ]
+    for file in t5xl_files:
+        if not (t5xl_weights / file).is_file():
+            warn(f"could not find weights at {t5xl_weights / file}, skipping")
+            pytest.skip(allow_module_level=True)
+
+    return (ella_adapter_weights, t5xl_weights)
 
 
 @pytest.fixture(scope="module")
@@ -605,10 +619,6 @@ def sd15_std_sde(
 def sd15_std_float16(
     text_encoder_weights: Path, lda_weights: Path, unet_weights_std: Path, test_device: torch.device
 ) -> StableDiffusion_1:
-    if test_device.type == "cpu":
-        warn("not running on CPU, skipping")
-        pytest.skip()
-
     sd15 = StableDiffusion_1(device=test_device, dtype=torch.float16)
 
     sd15.clip_text_encoder.load_from_safetensors(text_encoder_weights)
@@ -1817,15 +1827,13 @@ def test_diffusion_textual_inversion_random_init(
 @no_grad()
 def test_diffusion_ella_adapter(
     sd15_std_float16: StableDiffusion_1,
-    ella_adapter_weights: Path,
-    test_weights_path: Path,
+    ella_weights: tuple[Path, Path],
     expected_image_ella_adapter: Image.Image,
     test_device: torch.device,
 ):
     sd15 = sd15_std_float16
-    t5_encoder = T5TextEmbedder(pretrained_path=test_weights_path / "QQGYLab/T5XLFP16", max_length=128).to(
-        test_device, torch.float16
-    )
+    ella_adapter_weights, t5xl_weights = ella_weights
+    t5_encoder = T5TextEmbedder(pretrained_path=t5xl_weights, max_length=128).to(test_device, torch.float16)
 
     prompt = "a chinese man wearing a white shirt and a checkered headscarf, holds a large falcon near his shoulder. the falcon has dark feathers with a distinctive beak. the background consists of a clear sky and a fence, suggesting an outdoor setting, possibly a desert or arid region"
     negative_prompt = ""