mirror of
https://github.com/finegrain-ai/refiners.git
synced 2024-09-19 19:05:28 +00:00
62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
from typing import overload
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import refiners.fluxion.layers as fl
|
|
from refiners.fluxion.utils import no_grad
|
|
from refiners.foundationals.latent_diffusion import SD1IPAdapter, SD1UNet, SDXLIPAdapter, SDXLUNet
|
|
from refiners.foundationals.latent_diffusion.image_prompt import ImageCrossAttention
|
|
|
|
|
|
@overload
|
|
def new_adapter(target: SD1UNet) -> SD1IPAdapter: ...
|
|
|
|
|
|
@overload
|
|
def new_adapter(target: SDXLUNet) -> SDXLIPAdapter: ...
|
|
|
|
|
|
def new_adapter(target: SD1UNet | SDXLUNet) -> SD1IPAdapter | SDXLIPAdapter:
|
|
if isinstance(target, SD1UNet):
|
|
return SD1IPAdapter(target=target)
|
|
else:
|
|
return SDXLIPAdapter(target=target)
|
|
|
|
|
|
@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, 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()
|
|
@pytest.mark.parametrize("k_unet", [SD1UNet, SDXLUNet])
|
|
def test_scale(k_unet: type[SD1UNet] | type[SDXLUNet], test_device: torch.device):
|
|
unet = k_unet(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, ImageCrossAttention) and isinstance(m, fl.Multiply)
|
|
|
|
for m, _ in unet.walk(predicate):
|
|
assert isinstance(m, fl.Multiply)
|
|
assert m.scale == 1.0
|
|
|
|
adapter.scale = 0.42
|
|
assert adapter.scale == 0.42
|
|
for m, _ in unet.walk(predicate):
|
|
assert isinstance(m, fl.Multiply)
|
|
assert m.scale == 0.42
|