refiners/tests/adapters/test_control_lora.py
Cédric Deltheil 176807740b control_lora: fix adapter set scale
The adapter set scale did not propagate the scale to the underlying
zero convolutions. The value set at CTOR time was used instead.

Follow up of #285
2024-02-22 10:01:05 +01:00

35 lines
1.2 KiB
Python

import torch
import refiners.fluxion.layers as fl
from refiners.foundationals.latent_diffusion import ControlLoraAdapter, SDXLUNet
from refiners.foundationals.latent_diffusion.stable_diffusion_xl.control_lora import ZeroConvolution
def test_inject_eject(test_device: torch.device):
unet = SDXLUNet(in_channels=4, device=test_device, dtype=torch.float16)
initial_repr = repr(unet)
adapter = ControlLoraAdapter(name="foo", target=unet)
assert repr(unet) == initial_repr
adapter.inject()
assert repr(unet) != initial_repr
adapter.eject()
assert repr(unet) == initial_repr
def test_scale(test_device: torch.device):
unet = SDXLUNet(in_channels=4, device=test_device, dtype=torch.float16)
adapter = ControlLoraAdapter(name="foo", target=unet, scale=0.75).inject()
def predicate(m: fl.Module, p: fl.Chain) -> bool:
return isinstance(p, ZeroConvolution) and isinstance(m, fl.Multiply)
for m, _ in unet.walk(predicate):
assert isinstance(m, fl.Multiply)
assert m.scale == 0.75
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