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