v1.2 edit

diffrhythm/infer/infer_utils.py

@@ -250,6 +250,7 @@ def get_reference_latent(device, max_frames, edit, pred_segments, ref_song, vae_
             mean, scale = latent.chunk(2, dim=1)
             prompt, _ = vae_sample(mean, scale)
             prompt = prompt.transpose(1, 2) # [b t d]
+            prompt = prompt[:,:max_frames,:] if prompt.shape[1] >= max_frames else torch.nn.functional.pad(prompt, (0, 0, 0, max_frames - prompt.shape[1]), mode="constant", value=0)
         
         pred_segments = json.loads(pred_segments)
         # import pdb; pdb.set_trace()