refiners/tests/adapters/test_ella_adapter.py

42 lines
1.3 KiB
Python

import torch
import refiners.fluxion.layers as fl
from refiners.fluxion.utils import no_grad
from refiners.foundationals.latent_diffusion import SD1ELLAAdapter, SD1UNet
from refiners.foundationals.latent_diffusion.ella_adapter import ELLACrossAttentionAdapter
def new_adapter(target: SD1UNet) -> SD1ELLAAdapter:
return SD1ELLAAdapter(target=target)
@no_grad()
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
adapter.inject()
assert repr(unet) != initial_repr
adapter.eject()
assert repr(unet) == initial_repr
adapter.inject()
assert repr(unet) != initial_repr
adapter.eject()
assert repr(unet) == initial_repr
@no_grad()
def test_ella_cross_attention(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:
return isinstance(p, ELLACrossAttentionAdapter) and isinstance(m, fl.UseContext)
for m, _ in unet.walk(predicate):
assert isinstance(m, fl.UseContext)
assert m.context == "ella"
assert m.key == "latents"
assert len(adapter.sub_adapters) == 32