(args)
| 125 | |
| 126 | |
| 127 | def main(args): |
| 128 | t5_checkpoint = checkpoints.load_t5x_checkpoint(args.checkpoint_path) |
| 129 | t5_checkpoint = jnp.tree_util.tree_map(onp.array, t5_checkpoint) |
| 130 | |
| 131 | gin_overrides = [ |
| 132 | "from __gin__ import dynamic_registration", |
| 133 | "from music_spectrogram_diffusion.models.diffusion import diffusion_utils", |
| 134 | "diffusion_utils.ClassifierFreeGuidanceConfig.eval_condition_weight = 2.0", |
| 135 | "diffusion_utils.DiffusionConfig.classifier_free_guidance = @diffusion_utils.ClassifierFreeGuidanceConfig()", |
| 136 | ] |
| 137 | |
| 138 | gin_file = os.path.join(args.checkpoint_path, "..", "config.gin") |
| 139 | gin_config = inference.parse_training_gin_file(gin_file, gin_overrides) |
| 140 | synth_model = inference.InferenceModel(args.checkpoint_path, gin_config) |
| 141 | |
| 142 | scheduler = DDPMScheduler(beta_schedule="squaredcos_cap_v2", variance_type="fixed_large") |
| 143 | |
| 144 | notes_encoder = SpectrogramNotesEncoder( |
| 145 | max_length=synth_model.sequence_length["inputs"], |
| 146 | vocab_size=synth_model.model.module.config.vocab_size, |
| 147 | d_model=synth_model.model.module.config.emb_dim, |
| 148 | dropout_rate=synth_model.model.module.config.dropout_rate, |
| 149 | num_layers=synth_model.model.module.config.num_encoder_layers, |
| 150 | num_heads=synth_model.model.module.config.num_heads, |
| 151 | d_kv=synth_model.model.module.config.head_dim, |
| 152 | d_ff=synth_model.model.module.config.mlp_dim, |
| 153 | feed_forward_proj="gated-gelu", |
| 154 | ) |
| 155 | |
| 156 | continuous_encoder = SpectrogramContEncoder( |
| 157 | input_dims=synth_model.audio_codec.n_dims, |
| 158 | targets_context_length=synth_model.sequence_length["targets_context"], |
| 159 | d_model=synth_model.model.module.config.emb_dim, |
| 160 | dropout_rate=synth_model.model.module.config.dropout_rate, |
| 161 | num_layers=synth_model.model.module.config.num_encoder_layers, |
| 162 | num_heads=synth_model.model.module.config.num_heads, |
| 163 | d_kv=synth_model.model.module.config.head_dim, |
| 164 | d_ff=synth_model.model.module.config.mlp_dim, |
| 165 | feed_forward_proj="gated-gelu", |
| 166 | ) |
| 167 | |
| 168 | decoder = T5FilmDecoder( |
| 169 | input_dims=synth_model.audio_codec.n_dims, |
| 170 | targets_length=synth_model.sequence_length["targets_context"], |
| 171 | max_decoder_noise_time=synth_model.model.module.config.max_decoder_noise_time, |
| 172 | d_model=synth_model.model.module.config.emb_dim, |
| 173 | num_layers=synth_model.model.module.config.num_decoder_layers, |
| 174 | num_heads=synth_model.model.module.config.num_heads, |
| 175 | d_kv=synth_model.model.module.config.head_dim, |
| 176 | d_ff=synth_model.model.module.config.mlp_dim, |
| 177 | dropout_rate=synth_model.model.module.config.dropout_rate, |
| 178 | ) |
| 179 | |
| 180 | notes_encoder = load_notes_encoder(t5_checkpoint["target"]["token_encoder"], notes_encoder) |
| 181 | continuous_encoder = load_continuous_encoder(t5_checkpoint["target"]["continuous_encoder"], continuous_encoder) |
| 182 | decoder = load_decoder(t5_checkpoint["target"]["decoder"], decoder) |
| 183 | |
| 184 | melgan = OnnxRuntimeModel.from_pretrained("kashif/soundstream_mel_decoder") |
no test coverage detected