mirror of
https://github.com/finegrain-ai/refiners.git
synced 2024-09-19 19:05:28 +00:00
42 lines
1.3 KiB
Python
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
|