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,
)
| 47 | |
| 48 | ## Model loading |
| 49 | def load_model( |
| 50 | model_config: Optional[Dict[str, Any]] = None, |
| 51 | model_ckpt_path: Optional[str] = None, |
| 52 | pretrained_name: Optional[str] = None, |
| 53 | pretransform_ckpt_path: Optional[str] = None, |
| 54 | device: torch.device = DEVICE, |
| 55 | ) -> Tuple[torch.nn.Module, Dict[str, Any]]: |
| 56 | """Load the AudioGen model and its configuration. |
| 57 | |
| 58 | Either a pretrained model (via `pretrained_name`) or a freshly constructed one |
| 59 | (via `model_config` + `model_ckpt_path`) will be loaded. |
| 60 | |
| 61 | Args: |
| 62 | model_config: Configuration dict for creating the model. |
| 63 | model_ckpt_path: Path to a model checkpoint file. |
| 64 | pretrained_name: Name of a model to load from the repo. |
| 65 | pretransform_ckpt_path: Optional path to a pretransform checkpoint. |
| 66 | device: Torch device to map the model to. |
| 67 | |
| 68 | Returns: |
| 69 | A tuple of (model, model_config), where `model` is in eval mode |
| 70 | and cast to float, and `model_config` contains sample_rate/size, etc. |
| 71 | """ |
| 72 | |
| 73 | if pretrained_name is not None: |
| 74 | logging.info("Loading pretrained model: %s", pretrained_name) |
| 75 | model, model_config = get_pretrained_model(pretrained_name) |
| 76 | |
| 77 | elif model_config is not None: |
| 78 | if model_ckpt_path is None: |
| 79 | raise ValueError( |
| 80 | "model_ckpt_path must be provided when specifying model_config" |
| 81 | ) |
| 82 | logging.info("Creating model from config") |
| 83 | model = create_model_from_config(model_config) |
| 84 | |
| 85 | logging.info("Loading model checkpoint from: %s", model_ckpt_path) |
| 86 | |
| 87 | # Load checkpoint |
| 88 | copy_state_dict(model, load_ckpt_state_dict(model_ckpt_path)) |
| 89 | logging.info("Done loading model checkpoint") |
| 90 | |
| 91 | if pretransform_ckpt_path is not None: |
| 92 | logging.info("Loading pretransform checkpoint from %r", pretransform_ckpt_path) |
| 93 | model.pretransform.load_state_dict( |
| 94 | load_ckpt_state_dict(pretransform_ckpt_path), strict=False |
| 95 | ) |
| 96 | logging.info("Done loading pretransform.") |
| 97 | |
| 98 | model.to(device).eval().requires_grad_(False) |
| 99 | model.pretransform.model_half = False |
| 100 | model = model.to(torch.float) |
| 101 | |
| 102 | return model, model_config |
| 103 | |
| 104 | |
| 105 | ## Utility functions for conditioners |