File size: 1,332 Bytes
03de09d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | from .noise_encoders import get_noise_encoder
from .text_encoders import get_text_encoder
from .model import ScorePredictor
NOISE_ENCODERS = ['residualconv']
TEXT_ENCODERS = ['attnpool', 'lightattnpool', 'pertokenscalar']
def get_model(
noise_enc: str = 'residualconv',
text_enc: str = 'attnpool',
dropout: float = 0.1,
num_heads: int = 1,
spatial_size: int = 128,
in_channels: int = 4,
embed_dim: int = 2048,
seq_len: int = 77,
pos_encoding: str = 'none',
) -> ScorePredictor:
if noise_enc not in NOISE_ENCODERS:
raise ValueError(f"Unknown noise encoder: {noise_enc}. Available: {NOISE_ENCODERS}")
if text_enc not in TEXT_ENCODERS:
raise ValueError(f"Unknown text encoder: {text_enc}. Available: {TEXT_ENCODERS}")
text_encoder = get_text_encoder(text_enc, embed_dim=embed_dim, seq_len=seq_len, pos_encoding=pos_encoding)
noise_encoder = get_noise_encoder(spatial_size=spatial_size, in_channels=in_channels)
return ScorePredictor(
noise_encoder=noise_encoder,
text_encoder=text_encoder,
dropout=dropout,
num_heads=num_heads,
)
__all__ = [
'get_model',
'ScorePredictor',
'get_text_encoder',
'get_noise_encoder',
'NOISE_ENCODERS',
'TEXT_ENCODERS',
]
|