Finegrain/refiners
1
0
Fork 0
mirror of https://github.com/finegrain-ai/refiners.git synced 2024-11-09 15:02:01 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

add context_key getter and setter to RangeAdapter2d

Browse source
This commit is contained in:
Laurent 2024-02-14 15:21:49 +00:00 committed by Laureηt
parent 35b6e2f7c5
commit a54808e757
1 changed files with 19 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
src/refiners/foundationals/latent_diffusion/range_adapter.py
View file

@ -55,14 +55,31 @@ class RangeAdapter2d(fl.Sum, Adapter[fl.Conv2d]):
) -> None:
self.channels = channels
self.embedding_dim = embedding_dim
self.context_key = context_key
with self.setup_adapter(target):
super().__init__(
target,
fl.Chain(
fl.UseContext("range_adapter", context_key),
fl.SiLU(),
fl.Linear(in_features=embedding_dim, out_features=channels, device=device, dtype=dtype),
fl.Linear(
in_features=embedding_dim,
out_features=channels,
device=device,
dtype=dtype,
),
fl.Reshape(channels, 1, 1),
),
)
@property
def context_key(self) -> str:
use_context_module = self.ensure_find(fl.UseContext)
assert use_context_module.context == "range_adapter"
return use_context_module.key
@context_key.setter
def context_key(self, value: str) -> None:
use_context_module = self.ensure_find(fl.UseContext)
assert use_context_module.context == "range_adapter"
use_context_module.key = value