Update app.py

app.py

@@ -46,7 +46,6 @@ except FileNotFoundError:
 sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer'))
 sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer', 'descriptaudiocodec'))
 
-
 # don't change above code
 
 import argparse
@@ -106,7 +105,6 @@ codec_model.eval()
 #vocal_decoder.eval()
 #inst_decoder.eval()
 
-
 @spaces.GPU(duration=120)
 def generate_music(
         max_new_tokens=5,
@@ -119,6 +117,7 @@ def generate_music(
         prompt_end_time=30.0,
         cuda_idx=0,
         rescale=False,
+        batch_size=1
 ):
     if use_audio_prompt and not audio_prompt_path:
         raise FileNotFoundError("Please offer audio prompt filepath using '--audio_prompt_path', when you enable 'use_audio_prompt'!")
@@ -155,7 +154,8 @@ def generate_music(
 
         # Call the function and print the result
         stage1_output_set = []
-
+        vocals_list = []
+        instrumentals_list = []
         genres = genre_txt.strip()
         lyrics = split_lyrics(lyrics_txt + "\n")
         # intruction
@@ -163,139 +163,188 @@ def generate_music(
         prompt_texts = [f"Generate music from the given lyrics segment by segment.\n[Genre] {genres}\n{full_lyrics}"]
         prompt_texts += lyrics
 
-        random_id = uuid.uuid4()
-        output_seq = None
-        # Here is suggested decoding config
-        top_p = 0.93
-        temperature = 1.0
-        repetition_penalty = 1.2
         # special tokens
         start_of_segment = mmtokenizer.tokenize('[start_of_segment]')
         end_of_segment = mmtokenizer.tokenize('[end_of_segment]')
 
-        raw_output = None
-
         # Format text prompt
         run_n_segments = min(run_n_segments + 1, len(lyrics))
-
-        print(list(enumerate(tqdm(prompt_texts[:run_n_segments]))))
-
-        for i, p in enumerate(tqdm(prompt_texts[:run_n_segments])):
-            section_text = p.replace('[start_of_segment]', '').replace('[end_of_segment]', '')
-            guidance_scale = 1.5 if i <= 1 else 1.2
-            if i == 0:
-                continue
-            if i == 1:
-                if use_audio_prompt:
-                    audio_prompt = load_audio_mono(audio_prompt_path)
-                    audio_prompt.unsqueeze_(0)
-                    with torch.no_grad():
-                        raw_codes = codec_model.encode(audio_prompt.to(device), target_bw=0.5)
-                    raw_codes = raw_codes.transpose(0, 1)
-                    raw_codes = raw_codes.cpu().numpy().astype(np.int16)
-                    # Format audio prompt
-                    code_ids = codectool.npy2ids(raw_codes[0])
-                    audio_prompt_codec = code_ids[int(prompt_start_time * 50): int(prompt_end_time * 50)]  # 50 is tps of xcodec
-                    audio_prompt_codec_ids = [mmtokenizer.soa] + codectool.sep_ids + audio_prompt_codec + [
-                        mmtokenizer.eoa]
-                    sentence_ids = mmtokenizer.tokenize("[start_of_reference]") + audio_prompt_codec_ids + mmtokenizer.tokenize(
-                        "[end_of_reference]")
-                    head_id = mmtokenizer.tokenize(prompt_texts[0]) + sentence_ids
+        
+        batches = [prompt_texts[i:i + batch_size] for i in range(0, run_n_segments, batch_size)]
+
+        print(batches)
+
+        for batch_idx, batch in enumerate(tqdm(batches)):
+            random_ids = [uuid.uuid4() for _ in range(len(batch))]
+            raw_outputs = [None] * len(batch)
+
+            # Here is suggested decoding config
+            top_p = 0.93
+            temperature = 1.0
+            repetition_penalty = 1.2
+
+            for i, p in enumerate(batch):
+                section_text = p.replace('[start_of_segment]', '').replace('[end_of_segment]', '')
+                # Adjust guidance scale for the first two sections to be lower
+                guidance_scale = 1.5 if (batch_idx*batch_size + i) <= 1 else 1.2
+
+                if (batch_idx*batch_size + i) == 0:
+                    continue  # Skip the first instruction
+
+                if (batch_idx * batch_size + i) == 1:
+                    if use_audio_prompt:
+                        audio_prompt = load_audio_mono(audio_prompt_path)
+                        audio_prompt.unsqueeze_(0)
+                        with torch.no_grad():
+                            raw_codes = codec_model.encode(audio_prompt.to(device), target_bw=0.5)
+                        raw_codes = raw_codes.transpose(0, 1)
+                        raw_codes = raw_codes.cpu().numpy().astype(np.int16)
+                        # Format audio prompt
+                        code_ids = codectool.npy2ids(raw_codes[0])
+                        audio_prompt_codec = code_ids[int(prompt_start_time * 50): int(prompt_end_time * 50)]  # 50 is tps of xcodec
+                        audio_prompt_codec_ids = [mmtokenizer.soa] + codectool.sep_ids + audio_prompt_codec + [
+                            mmtokenizer.eoa]
+                        sentence_ids = mmtokenizer.tokenize("[start_of_reference]") + audio_prompt_codec_ids + mmtokenizer.tokenize(
+                            "[end_of_reference]")
+                        head_id = mmtokenizer.tokenize(prompt_texts[0]) + sentence_ids
+                    else:
+                        head_id = mmtokenizer.tokenize(prompt_texts[0])
+                    prompt_ids = head_id + start_of_segment + mmtokenizer.tokenize(section_text) + [
+                        mmtokenizer.soa] + codectool.sep_ids
+                else:
+                    prompt_ids = end_of_segment + start_of_segment + mmtokenizer.tokenize(section_text) + [
+                        mmtokenizer.soa] + codectool.sep_ids
+                
+                prompt_ids = torch.as_tensor(prompt_ids).unsqueeze(0).to(device)
+                input_ids = torch.cat([raw_outputs[i], prompt_ids], dim=1) if (batch_idx * batch_size + i) > 1 else prompt_ids
+
+                # Use window slicing in case output sequence exceeds the context of model
+                max_context = 16384 - max_new_tokens - 1
+                if input_ids.shape[-1] > max_context:
+                    print(
+                        f'Section {(batch_idx * batch_size + i)}: output length {input_ids.shape[-1]} exceeding context length {max_context}, now using the last {max_context} tokens.')
+                    input_ids = input_ids[:, -(max_context):]
+
+                with torch.inference_mode(), torch.autocast(device_type='cuda', dtype=torch.float16):
+                    output_seq = model.generate(
+                        input_ids=input_ids,
+                        max_new_tokens=max_new_tokens,
+                        min_new_tokens=100,
+                        do_sample=True,
+                        top_p=top_p,
+                        temperature=temperature,
+                        repetition_penalty=repetition_penalty,
+                        eos_token_id=mmtokenizer.eoa,
+                        pad_token_id=mmtokenizer.eoa,
+                        logits_processor=LogitsProcessorList(
+                            [BlockTokenRangeProcessor(0, 32002), BlockTokenRangeProcessor(32016, 32016)]),
+                        guidance_scale=guidance_scale,
+                        use_cache=True
+                    )
+                    if output_seq[0][-1].item() != mmtokenizer.eoa:
+                        tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(model.device)
+                        output_seq = torch.cat((output_seq, tensor_eoa), dim=1)
+
+                if (batch_idx * batch_size + i) > 1:
+                    raw_outputs[i] = torch.cat([raw_outputs[i], prompt_ids, output_seq[:, input_ids.shape[-1]:]], dim=1)
                 else:
-                    head_id = mmtokenizer.tokenize(prompt_texts[0])
-                prompt_ids = head_id + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
-            else:
-                prompt_ids = end_of_segment + start_of_segment + mmtokenizer.tokenize(section_text) + [mmtokenizer.soa] + codectool.sep_ids
-
-            prompt_ids = torch.as_tensor(prompt_ids).unsqueeze(0).to(device)
-            input_ids = torch.cat([raw_output, prompt_ids], dim=1) if i > 1 else prompt_ids
-            # Use window slicing in case output sequence exceeds the context of model
-            max_context = 16384 - max_new_tokens - 1
-            if input_ids.shape[-1] > max_context:
-                print(
-                    f'Section {i}: output length {input_ids.shape[-1]} exceeding context length {max_context}, now using the last {max_context} tokens.')
-                input_ids = input_ids[:, -(max_context):]
-            with torch.inference_mode(), torch.autocast(device_type='cuda', dtype=torch.float16):
-                output_seq = model.generate(
-                    input_ids=input_ids,
-                    max_new_tokens=max_new_tokens,
-                    min_new_tokens=100,
-                    do_sample=True,
-                    top_p=top_p,
-                    temperature=temperature,
-                    repetition_penalty=repetition_penalty,
-                    eos_token_id=mmtokenizer.eoa,
-                    pad_token_id=mmtokenizer.eoa,
-                    logits_processor=LogitsProcessorList([BlockTokenRangeProcessor(0, 32002), BlockTokenRangeProcessor(32016, 32016)]),
-                    guidance_scale=guidance_scale,
-                    use_cache=True
-                )
-                if output_seq[0][-1].item() != mmtokenizer.eoa:
-                    tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(model.device)
-                    output_seq = torch.cat((output_seq, tensor_eoa), dim=1)
-            if i > 1:
-                raw_output = torch.cat([raw_output, prompt_ids, output_seq[:, input_ids.shape[-1]:]], dim=1)
-            else:
-                raw_output = output_seq
-            print(len(raw_output))
-
-        # save raw output and check sanity
-        ids = raw_output[0].cpu().numpy()
-        soa_idx = np.where(ids == mmtokenizer.soa)[0].tolist()
-        eoa_idx = np.where(ids == mmtokenizer.eoa)[0].tolist()
-        if len(soa_idx) != len(eoa_idx):
-            raise ValueError(f'invalid pairs of soa and eoa, Num of soa: {len(soa_idx)}, Num of eoa: {len(eoa_idx)}')
-
-        vocals = []
-        instrumentals = []
-        range_begin = 1 if use_audio_prompt else 0
-        for i in range(range_begin, len(soa_idx)):
-            codec_ids = ids[soa_idx[i] + 1:eoa_idx[i]]
-            if codec_ids[0] == 32016:
-                codec_ids = codec_ids[1:]
-            codec_ids = codec_ids[:2 * (codec_ids.shape[0] // 2)]
-            vocals_ids = codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[0])
-            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)
+                    raw_outputs[i] = output_seq
+
+            for i, raw_output in enumerate(raw_outputs):
+                # save raw output and check sanity
+                ids = raw_output[0].cpu().numpy()
+                soa_idx = np.where(ids == mmtokenizer.soa)[0].tolist()
+                eoa_idx = np.where(ids == mmtokenizer.eoa)[0].tolist()
+                if len(soa_idx) != len(eoa_idx):
+                    raise ValueError(f'invalid pairs of soa and eoa, Num of soa: {len(soa_idx)}, Num of eoa: {len(eoa_idx)}')
+                
+                range_begin = 1 if use_audio_prompt and batch_idx == 0 else 0
+
+                vocals_batch = []
+                instrumentals_batch = []
+                for j in range(range_begin, len(soa_idx)):
+                    codec_ids = ids[soa_idx[j] + 1:eoa_idx[j]]
+                    if codec_ids[0] == 32016:
+                        codec_ids = codec_ids[1:]
+                    codec_ids = codec_ids[:2 * (codec_ids.shape[0] // 2)]
+                    vocals_ids = codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[0])
+                    vocals_batch.append(vocals_ids)
+                    instrumentals_ids = codectool.ids2npy(rearrange(codec_ids, "(n b) -> b n", b=2)[1])
+                    instrumentals_batch.append(instrumentals_ids)
+
+                vocals_batch = np.concatenate(vocals_batch, axis=1)
+                instrumentals_batch = np.concatenate(instrumentals_batch, axis=1)
+                
+                vocals_list.append(vocals_batch)
+                instrumentals_list.append(instrumentals_batch)
+
+        vocals = np.concatenate(vocals_list, axis=1)
+        instrumentals = np.concatenate(instrumentals_list, axis=1)
+
+        vocal_save_path = os.path.join(stage1_output_dir, f"vocal_{random_ids[0]}".replace('.', '@') + '.npy')
+        inst_save_path = os.path.join(stage1_output_dir, f"instrumental_{random_ids[0]}".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...")
 
-        # 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
-
+        # 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)
 
         # reconstruct tracks
-        vocal_waveform = decode_audio_batch(vocals)
-        instrumental_waveform = decode_audio_batch(instrumentals)
-
+        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 = []
+
+        vocal_stem = None
+        instrumental_stem = None
+        sr = None
+
+        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)
+            if 'vocal' in npy:
+                vocal_stem = decodec_rlt.numpy()
+            elif 'instrumental' in npy:
+                instrumental_stem = decodec_rlt.numpy()
+            sr = 16000
+        
         # mix tracks
-        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
-
+        if vocal_stem is not None and instrumental_stem is not None:
+            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))
 
+        else:
+            print("Missing Vocal or Instrumental Stem")
+            return None, None, None
 
 def infer(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=15):
     # Execute the command
     try:
         mixed_audio_data, vocal_audio_data, instrumental_audio_data = generate_music(genre_txt=genre_txt_content, lyrics_txt=lyrics_txt_content, run_n_segments=num_segments,
-                               cuda_idx=0, max_new_tokens=max_new_tokens)
+                               cuda_idx=0, max_new_tokens=max_new_tokens, batch_size=4)
         return mixed_audio_data, vocal_audio_data, instrumental_audio_data
     except Exception as e:
         gr.Warning("An Error Occured: " + str(e))
@@ -303,7 +352,6 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=
     finally:
         print("Temporary files deleted.")
 
-
 # Gradio
 with gr.Blocks() as demo:
     with gr.Column():