Spaces:
Paused
Paused
File size: 670 Bytes
9d7cf7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | from .skin_cvae_model import SkinCVAEModel
from .skin_fsq_cvae_model import SkinFSQCVAEModel
def get_model_cvae(
pretrained_path: str=None,
**kwargs
) -> SkinCVAEModel:
model = SkinCVAEModel(**kwargs)
if pretrained_path is not None:
state_dict = torch.load(pretrained_path, weights_only=True)
model.load_state_dict(state_dict)
return model
def get_model_fsq_cvae(
pretrained_path: str=None,
**kwargs
) -> SkinFSQCVAEModel:
model = SkinFSQCVAEModel(**kwargs)
if pretrained_path is not None:
state_dict = torch.load(pretrained_path, weights_only=True)
model.load_state_dict(state_dict)
return model |