Load the AudioGen model and its configuration. Either a pretrained model (via `pretrained_name`) or a freshly constructed one (via `model_config` + `model_ckpt_path`) will be loaded. Args: model_config: Configuration dict for creating the model. model_ckpt_path: Path to
(
model_config: Optional[Dict[str, Any]] = None,
model_ckpt_path: Optional[str] = None,
pretrained_name: Optional[str] = None,
pretransform_ckpt_path: Optional[str] = None,
device: torch.device = DEVICE,
)
| 21 | |
| 22 | ## Model loading |
| 23 | def load_model( |
| 24 | model_config: Optional[Dict[str, Any]] = None, |
| 25 | model_ckpt_path: Optional[str] = None, |
| 26 | pretrained_name: Optional[str] = None, |
| 27 | pretransform_ckpt_path: Optional[str] = None, |
| 28 | device: torch.device = DEVICE, |
| 29 | ) -> Tuple[torch.nn.Module, Dict[str, Any]]: |
| 30 | """Load the AudioGen model and its configuration. |
| 31 | |
| 32 | Either a pretrained model (via `pretrained_name`) or a freshly constructed one |
| 33 | (via `model_config` + `model_ckpt_path`) will be loaded. |
| 34 | |
| 35 | Args: |
| 36 | model_config: Configuration dict for creating the model. |
| 37 | model_ckpt_path: Path to a model checkpoint file. |
| 38 | pretrained_name: Name of a model to load from the repo. |
| 39 | pretransform_ckpt_path: Optional path to a pretransform checkpoint. |
| 40 | device: Torch device to map the model to. |
| 41 | |
| 42 | Returns: |
| 43 | A tuple of (model, model_config), where `model` is in eval mode |
| 44 | and cast to float, and `model_config` contains sample_rate/size, etc. |
| 45 | """ |
| 46 | |
| 47 | if pretrained_name is not None: |
| 48 | logging.info("Loading pretrained model: %s", pretrained_name) |
| 49 | model, model_config = get_pretrained_model(pretrained_name) |
| 50 | |
| 51 | elif model_config is not None: |
| 52 | if model_ckpt_path is None: |
| 53 | raise ValueError( |
| 54 | "model_ckpt_path must be provided when specifying model_config" |
| 55 | ) |
| 56 | logging.info("Creating model from config") |
| 57 | model = create_model_from_config(model_config) |
| 58 | |
| 59 | logging.info("Loading model checkpoint from: %s", model_ckpt_path) |
| 60 | # Load checkpoint |
| 61 | copy_state_dict(model, load_ckpt_state_dict(model_ckpt_path)) |
| 62 | logging.info("Done loading model checkpoint") |
| 63 | |
| 64 | SAMPLE_RATE = model_config["sample_rate"] |
| 65 | SAMPLE_SIZE = model_config["sample_size"] |
| 66 | |
| 67 | if pretransform_ckpt_path is not None: |
| 68 | logging.info("Loading pretransform checkpoint from %r", pretransform_ckpt_path) |
| 69 | model.pretransform.load_state_dict( |
| 70 | load_ckpt_state_dict(pretransform_ckpt_path), strict=False |
| 71 | ) |
| 72 | logging.info("Done loading pretransform.") |
| 73 | |
| 74 | model.to(device).eval().requires_grad_(False) |
| 75 | model = model.to(torch.float) |
| 76 | |
| 77 | return model, model_config |
| 78 | |
| 79 | |
| 80 | ## ----------------- Conditioners Utility Functions ------------------- |