refiners/tests/adapters/test_range_adapter.py
2023-09-01 17:29:48 +02:00

25 lines
808 B
Python

import torch
from refiners.fluxion.adapters.adapter import Adapter
from refiners.foundationals.latent_diffusion.range_adapter import RangeEncoder
from refiners.fluxion.layers import Chain, Linear
class DummyLinearAdapter(Chain, Adapter[Linear]):
def __init__(self, target: Linear):
with self.setup_adapter(target):
super().__init__(target)
def test_range_encoder_dtype_after_adaptation(test_device: torch.device): # FG-433
dtype = torch.float64
chain = Chain(RangeEncoder(320, 1280, device=test_device, dtype=dtype))
adaptee = chain.RangeEncoder.Linear_1
adapter = DummyLinearAdapter(adaptee).inject(chain.RangeEncoder)
assert adapter.parent == chain.RangeEncoder
x = torch.tensor([42], device=test_device)
y = chain(x)
assert y.dtype == dtype