Update app.py

app.py

@@ -258,64 +258,37 @@ def generate_music(
             vocals.append(vocals_ids)
             instrumentals_ids = codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[1])
             instrumentals.append(instrumentals_ids)
+
         vocals = np.concatenate(vocals, axis=1)
         instrumentals = np.concatenate(instrumentals, axis=1)
 
-        vocal_save_path = os.path.join(stage1_output_dir, f"vocal_{random_id}".replace('.', '@') + '.npy')
-        inst_save_path = os.path.join(stage1_output_dir, f"instrumental_{random_id}".replace('.', '@') + '.npy')
-        np.save(vocal_save_path, vocals)
-        np.save(inst_save_path, instrumentals)
-        stage1_output_set.append(vocal_save_path)
-        stage1_output_set.append(inst_save_path)
-
         print("Converting to Audio...")
 
-        # convert audio tokens to audio
-        def save_audio(wav: torch.Tensor, path, sample_rate: int, rescale: bool = False):
-            folder_path = os.path.dirname(path)
-            if not os.path.exists(folder_path):
-                os.makedirs(folder_path)
-            limit = 0.99
-            max_val = wav.abs().max()
-            wav = wav * min(limit / max_val, 1) if rescale else wav.clamp(-limit, limit)
-            torchaudio.save(str(path), wav, sample_rate=sample_rate, encoding='PCM_S', bits_per_sample=16)
+        # batching audio
+        def decode_audio_batch(codec_result, batch_size=4):
+            decoded_waveforms = []
+            with torch.no_grad():
+                for i in range(0, codec_result.shape[-1], batch_size):
+                    batch = codec_result[:,i:i+batch_size]
+                    batch_tensor = torch.as_tensor(batch.astype(np.int16), dtype=torch.long).unsqueeze(0).permute(1, 0, 2).to(device)
+                    decoded_waveform = codec_model.decode(batch_tensor)
+                    decoded_waveforms.append(decoded_waveform)
+                decoded_waveforms = torch.cat(decoded_waveforms, dim=-1).squeeze(0).cpu()
+            return decoded_waveforms
+
 
         # reconstruct tracks
-        recons_output_dir = os.path.join(output_dir, "recons")
-        recons_mix_dir = os.path.join(recons_output_dir, 'mix')
-        os.makedirs(recons_mix_dir, exist_ok=True)
-        tracks = []
-        for npy in stage1_output_set:
-            codec_result = np.load(npy)
-            decodec_rlt = []
-            with torch.no_grad():
-                decoded_waveform = codec_model.decode(
-                    torch.as_tensor(codec_result.astype(np.int16), dtype=torch.long).unsqueeze(0).permute(1, 0, 2).to(
-                        device))
-            decoded_waveform = decoded_waveform.cpu().squeeze(0)
-            decodec_rlt.append(torch.as_tensor(decoded_waveform))
-            decodec_rlt = torch.cat(decodec_rlt, dim=-1)
-            save_path = os.path.join(recons_output_dir, os.path.splitext(os.path.basename(npy))[0] + ".mp3")
-            tracks.append(save_path)
-            save_audio(decodec_rlt, save_path, 16000)
+        vocal_waveform = decode_audio_batch(vocals)
+        instrumental_waveform = decode_audio_batch(instrumentals)
+
         # mix tracks
-        for inst_path in tracks:
-            try:
-                if (inst_path.endswith('.wav') or inst_path.endswith('.mp3')) \
-                        and 'instrumental' in inst_path:
-                    # find pair
-                    vocal_path = inst_path.replace('instrumental', 'vocal')
-                    if not os.path.exists(vocal_path):
-                        continue
-                    # mix
-                    recons_mix = os.path.join(recons_mix_dir, os.path.basename(inst_path).replace('instrumental', 'mixed'))
-                    vocal_stem, sr = sf.read(inst_path)
-                    instrumental_stem, _ = sf.read(vocal_path)
-                    mix_stem = (vocal_stem + instrumental_stem) / 1
-                    return (sr, (mix_stem * 32767).astype(np.int16)), (sr, (vocal_stem * 32767).astype(np.int16)), (sr, (instrumental_stem * 32767).astype(np.int16))
-            except Exception as e:
-                print(e)
-                return None, None, None
+        try:
+            mix_waveform = (vocal_waveform + instrumental_waveform) / 1
+            return (16000, (mix_waveform * 32767).numpy().astype(np.int16)), (16000, (vocal_waveform * 32767).numpy().astype(np.int16)), (16000, (instrumental_waveform * 32767).numpy().astype(np.int16))
+        except Exception as e:
+            print(e)
+            return None, None, None
+
 
 
 def infer(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=15):