By deepseek prev was stable

app.py

@@ -1,11 +1,11 @@
 import gradio as gr
 import subprocess
-import os 
+import os
 import shutil
 import tempfile
 import spaces
 import torch
-import os
+import torch.nn.functional as F
 import sys
 
 print("Installing flash-attn...")
@@ -16,7 +16,7 @@ subprocess.run(
     shell=True,
 )
 
-from huggingface_hub import snapshot_download 
+from huggingface_hub import snapshot_download
 
 # Create xcodec_mini_infer folder
 folder_path = './xcodec_mini_infer'
@@ -29,8 +29,8 @@ else:
     print(f"Folder already exists at: {folder_path}")
 
 snapshot_download(
-    repo_id = "m-a-p/xcodec_mini_infer",
-    local_dir = "./xcodec_mini_infer"
+    repo_id="m-a-p/xcodec_mini_infer",
+    local_dir="./xcodec_mini_infer"
 )
 
 # Change to the "inference" directory
@@ -69,11 +69,12 @@ from models.soundstream_hubert_new import SoundStream
 from vocoder import build_codec_model, process_audio
 from post_process_audio import replace_low_freq_with_energy_matched
 import re
+import multiprocessing
 
 def empty_output_folder(output_dir):
     # List all files in the output directory
     files = os.listdir(output_dir)
-    
+
     # Iterate over the files and remove them
     for file in files:
         file_path = os.path.join(output_dir, file)
@@ -87,37 +88,27 @@ def empty_output_folder(output_dir):
         except Exception as e:
             print(f"Error deleting file {file_path}: {e}")
 
-# Function to create a temporary file with string content
-def create_temp_file(content, prefix, suffix=".txt"):
-    temp_file = tempfile.NamedTemporaryFile(delete=False, mode="w", prefix=prefix, suffix=suffix)
-    # Ensure content ends with newline and normalize line endings
-    content = content.strip() + "\n\n"  # Add extra newline at end
-    content = content.replace("\r\n", "\n").replace("\r", "\n")
-    temp_file.write(content)
-    temp_file.close()
-    
-    # Debug: Print file contents
-    print(f"\nContent written to {prefix}{suffix}:")
-    print(content)
-    print("---")
-    
-    return temp_file.name
-
 device = "cuda:0"
 
+# --- Model Loading and Quantization ---
 model = AutoModelForCausalLM.from_pretrained(
-    "m-a-p/YuE-s1-7B-anneal-en-cot", 
+    "m-a-p/YuE-s1-7B-anneal-en-cot",
     torch_dtype=torch.float16,
-    attn_implementation="flash_attention_2", # To enable flashattn, you have to install flash-attn
-    )
+    attn_implementation="flash_attention_2",  # To enable flashattn, you have to install flash-attn
+)
 model.to(device)
 model.eval()
 
-basic_model_config='./xcodec_mini_infer/final_ckpt/config.yaml'
-resume_path='./xcodec_mini_infer/final_ckpt/ckpt_00360000.pth'
-config_path='./xcodec_mini_infer/decoders/config.yaml'
-vocal_decoder_path='./xcodec_mini_infer/decoders/decoder_131000.pth'
-inst_decoder_path='./xcodec_mini_infer/decoders/decoder_151000.pth'
+# Apply dynamic quantization
+model = torch.quantization.quantize_dynamic(
+    model, {torch.nn.Linear}, dtype=torch.qint8
+)
+
+basic_model_config = './xcodec_mini_infer/final_ckpt/config.yaml'
+resume_path = './xcodec_mini_infer/final_ckpt/ckpt_00360000.pth'
+config_path = './xcodec_mini_infer/decoders/config.yaml'
+vocal_decoder_path = './xcodec_mini_infer/decoders/decoder_131000.pth'
+inst_decoder_path = './xcodec_mini_infer/decoders/decoder_151000.pth'
 
 mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
 
@@ -129,26 +120,53 @@ codec_model.load_state_dict(parameter_dict['codec_model'])
 codec_model.to(device)
 codec_model.eval()
 
+# --- Parallel Audio Processing ---
+def process_audio_wrapper(args):
+    # Unpack arguments and call the original process_audio function
+    npy, output_path, rescale, other_args, decoder, codec_model = args
+    return process_audio(npy, output_path, rescale, other_args, decoder, codec_model)
+
+def parallel_process_audio(stage1_output_set, vocoder_stems_dir, rescale, other_args, vocal_decoder, inst_decoder,
+                           codec_model, num_processes=4):
+    with multiprocessing.Pool(processes=num_processes) as pool:
+        tasks = []
+        for npy in stage1_output_set:
+            if 'instrumental' in npy:
+                output_path = os.path.join(vocoder_stems_dir, 'instrumental.mp3')
+                decoder = inst_decoder
+            else:
+                output_path = os.path.join(vocoder_stems_dir, 'vocal.mp3')
+                decoder = vocal_decoder
+            tasks.append((npy, output_path, rescale, other_args, decoder, codec_model))
+
+        results = pool.map(process_audio_wrapper, tasks)
+
+    return results
+
+# --- Optimized Music Generation ---
 def generate_music(
-    max_new_tokens=5,
-    run_n_segments=2,
-    genre_txt=None,
-    lyrics_txt=None,
-    use_audio_prompt=False,
-    audio_prompt_path="",
-    prompt_start_time=0.0,
-    prompt_end_time=30.0,
-    output_dir="./output",
-    cuda_idx=0,
-    rescale=False,
+        max_new_tokens=5,
+        run_n_segments=2,
+        genre_txt=None,
+        lyrics_txt=None,
+        use_audio_prompt=False,
+        audio_prompt_path="",
+        prompt_start_time=0.0,
+        prompt_end_time=30.0,
+        output_dir="./output",
+        rescale=False,
+        beam_width=3,  # Add beam search
+        length_penalty=1.0,  # Add length penalty
+        repetition_penalty=1.5, # Add repetition penalty
+        batch_size=2
 ):
     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'!")
-    cuda_idx = cuda_idx
-    max_new_tokens = max_new_tokens*100
+        raise FileNotFoundError(
+            "Please offer audio prompt filepath using '--audio_prompt_path', when you enable 'use_audio_prompt'!")
+    max_new_tokens = max_new_tokens * 100
     stage1_output_dir = os.path.join(output_dir, f"stage1")
     os.makedirs(stage1_output_dir, exist_ok=True)
-    
+
     class BlockTokenRangeProcessor(LogitsProcessor):
         def __init__(self, start_id, end_id):
             self.blocked_token_ids = list(range(start_id, end_id))
@@ -177,109 +195,158 @@ def generate_music(
     stage1_output_set = []
 
     genres = genre_txt.strip()
-    lyrics = split_lyrics(lyrics_txt+"\n")
+    lyrics = split_lyrics(lyrics_txt + "\n")
     # intruction
     full_lyrics = "\n".join(lyrics)
     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
+    segment_cache = {}  # Cache for repeated segments
 
     # Format text prompt
-    run_n_segments = min(run_n_segments+1, len(lyrics))
+    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
+    # Modified loop for batching and caching
+    for i in range(1, run_n_segments, batch_size):
+        batch_segments = []
+        batch_prompts = []
+        for j in range(i, min(i + batch_size, run_n_segments)):
+            section_text = prompt_texts[j].replace('[start_of_segment]', '').replace('[end_of_segment]', '')
+
+            # Check cache
+            if section_text in segment_cache:
+                cached_output = segment_cache[section_text]
+                if j > 1:
+                    raw_output = torch.cat([raw_output, cached_output], dim=1)
+                else:
+                    raw_output = cached_output
+                continue
+
+            batch_segments.append(section_text)
+            guidance_scale = 1.5 if j <= 1 else 1.2
+
+            if j == 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:
-                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):]
+                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 j > 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 {j}: output length {input_ids.shape[-1]} exceeding context length {max_context}, now using the last {max_context} tokens.')
+                input_ids = input_ids[:, -(max_context):]
+
+            batch_prompts.append(input_ids)
+
+        if not batch_prompts:
+            continue  # All segments in the batch were cached
+
+        # Pad prompts in the batch to the same length
+        max_len = max(p.size(1) for p in batch_prompts)
+        padded_prompts = []
+        for p in batch_prompts:
+            pad_len = max_len - p.size(1)
+            padded_prompt = F.pad(p, (0, pad_len), value=mmtokenizer.eoa)
+            padded_prompts.append(padded_prompt)
+
+        batch_input_ids = torch.cat(padded_prompts, dim=0)
+
         with torch.no_grad():
-            output_seq = model.generate(
-                input_ids=input_ids, 
-                max_new_tokens=max_new_tokens, 
-                min_new_tokens=100, 
-                do_sample=True, 
+            output_seqs = model.generate(
+                input_ids=batch_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, 
+                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)]),
+                logits_processor=LogitsProcessorList(
+                    [BlockTokenRangeProcessor(0, 32002), BlockTokenRangeProcessor(32016, 32016)]),
                 guidance_scale=guidance_scale,
                 use_cache=True,
-                )
+                num_beams=beam_width,  # Use beam search
+                length_penalty=length_penalty,  # Apply length penalty
+            )
+
+        # Process each output in the batch
+        for k, output_seq in enumerate(output_seqs):
             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))
+            if i > 1:
+                raw_output = torch.cat([raw_output, batch_prompts[k][:, :batch_input_ids.shape[-1]],
+                                        output_seq[:, batch_input_ids.shape[-1]:]], dim=1)
+            else:
+                raw_output = output_seq
+
+            # Cache the generated output if not already cached
+            if batch_segments[k] not in segment_cache:
+                segment_cache[batch_segments[k]] = output_seq[:, batch_input_ids.shape[-1]:].cpu()
 
     # 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):
+    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]]
+        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_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_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"cot_{genres.replace(' ', '-')}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens}_vocal_{random_id}".replace('.', '@')+'.npy')
-    inst_save_path = os.path.join(stage1_output_dir, f"cot_{genres.replace(' ', '-')}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens}_instrumental_{random_id}".replace('.', '@')+'.npy')
+    vocal_save_path = os.path.join(stage1_output_dir,
+                                   f"cot_{genres.replace(' ', '-')}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens}_vocal_{random_id}".replace(
+                                       '.', '@') + '.npy')
+    inst_save_path = os.path.join(stage1_output_dir,
+                                  f"cot_{genres.replace(' ', '-')}_tp{top_p}_T{temperature}_rp{repetition_penalty}_maxtk{max_new_tokens}_instrumental_{random_id}".replace(
+                                      '.', '@') + '.npy')
     np.save(vocal_save_path, vocals)
     np.save(inst_save_path, instrumentals)
     stage1_output_set.append(vocal_save_path)
@@ -296,6 +363,7 @@ def generate_music(
         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
     recons_output_dir = os.path.join(output_dir, "recons")
     recons_mix_dir = os.path.join(recons_output_dir, 'mix')
@@ -303,9 +371,11 @@ def generate_music(
     tracks = []
     for npy in stage1_output_set:
         codec_result = np.load(npy)
-        decodec_rlt=[]
+        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 = 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)
@@ -316,7 +386,7 @@ def generate_music(
     for inst_path in tracks:
         try:
             if (inst_path.endswith('.wav') or inst_path.endswith('.mp3')) \
-                and 'instrumental' in inst_path:
+                    and 'instrumental' in inst_path:
                 # find pair
                 vocal_path = inst_path.replace('instrumental', 'vocal')
                 if not os.path.exists(vocal_path):
@@ -329,7 +399,6 @@ def generate_music(
                 sf.write(recons_mix, mix_stem, sr)
         except Exception as e:
             print(e)
-    
 
     # vocoder to upsample audios
     vocal_decoder, inst_decoder = build_codec_model(config_path, vocal_decoder_path, inst_decoder_path)
@@ -338,53 +407,41 @@ def generate_music(
     vocoder_mix_dir = os.path.join(vocoder_output_dir, 'mix')
     os.makedirs(vocoder_mix_dir, exist_ok=True)
     os.makedirs(vocoder_stems_dir, exist_ok=True)
-    instrumental_output = None
-    vocal_output = None
-    for npy in stage1_output_set:
-        if 'instrumental' in npy:
-            # Process instrumental
-            instrumental_output = process_audio(
-                npy,
-                os.path.join(vocoder_stems_dir, 'instrumental.mp3'),
-                rescale,
-                argparse.Namespace(**locals()), # Convert local variables to argparse.Namespace
-                inst_decoder,
-                codec_model
-            )
-        else:
-            # Process vocal
-            vocal_output = process_audio(
-                npy,
-                os.path.join(vocoder_stems_dir, 'vocal.mp3'),
-                rescale,
-                 argparse.Namespace(**locals()), # Convert local variables to argparse.Namespace
-                vocal_decoder,
-                codec_model
-            )
-    # mix tracks
-    try:
-        mix_output = instrumental_output + vocal_output
-        vocoder_mix = os.path.join(vocoder_mix_dir, os.path.basename(recons_mix))
-        save_audio(mix_output, vocoder_mix, 44100, rescale)
-        print(f"Created mix: {vocoder_mix}")
-    except RuntimeError as e:
-        print(e)
-        print(f"mix {vocoder_mix} failed! inst: {instrumental_output.shape}, vocal: {vocal_output.shape}")
+
+    # Use parallel processing for vocoding
+    parallel_process_audio(stage1_output_set, vocoder_stems_dir, rescale, argparse.Namespace(**locals()), vocal_decoder,
+                           inst_decoder, codec_model)
+
+    # mix tracks after parallel processing
+    instrumental_output_path = os.path.join(vocoder_stems_dir, 'instrumental.mp3')
+    vocal_output_path = os.path.join(vocoder_stems_dir, 'vocal.mp3')
+
+    if os.path.exists(instrumental_output_path) and os.path.exists(vocal_output_path):
+      instrumental_output, sr = torchaudio.load(instrumental_output_path)
+      vocal_output, _ = torchaudio.load(vocal_output_path)
+      try:
+          mix_output = instrumental_output + vocal_output
+          vocoder_mix = os.path.join(vocoder_mix_dir, os.path.basename(recons_mix))
+          save_audio(mix_output, vocoder_mix, 44100, rescale)
+          print(f"Created mix: {vocoder_mix}")
+      except RuntimeError as e:
+          print(e)
+          print(f"mix {vocoder_mix} failed! inst: {instrumental_output.shape}, vocal: {vocal_output.shape}")
+    else:
+      print("Skipping mix creation, instrumental or vocal output missing.")
 
     # Post process
     replace_low_freq_with_energy_matched(
-        a_file=recons_mix,     # 16kHz
-        b_file=vocoder_mix,     # 48kHz
+        a_file=recons_mix,  # 16kHz
+        b_file=vocoder_mix,  # 48kHz
         c_file=os.path.join(output_dir, os.path.basename(recons_mix)),
         cutoff_freq=5500.0
     )
     print("All process Done")
     return recons_mix
 
-
 @spaces.GPU(duration=120)
-def infer(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=10):
-
+def infer(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=5):
     # Ensure the output folder exists
     output_dir = "./output"
     os.makedirs(output_dir, exist_ok=True)
@@ -394,15 +451,16 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments=2, max_new_tokens=
 
     # Execute the command
     try:
-        music = generate_music(genre_txt=genre_txt_content, lyrics_txt=lyrics_txt_content, run_n_segments=num_segments, output_dir=output_dir, cuda_idx=0, max_new_tokens=max_new_tokens)
+        music = generate_music(genre_txt=genre_txt_content, lyrics_txt=lyrics_txt_content, run_n_segments=num_segments,
+                               output_dir=output_dir, cuda_idx=0, max_new_tokens=max_new_tokens)
         return music
     except Exception as e:
         gr.Warning("An Error Occured: " + str(e))
-        return none
+        return None
     finally:
         print("Temporary files deleted.")
 
-# Gradio 
+# Gradio
 
 with gr.Blocks() as demo:
     with gr.Column():
@@ -424,15 +482,16 @@ with gr.Blocks() as demo:
             with gr.Column():
                 genre_txt = gr.Textbox(label="Genre")
                 lyrics_txt = gr.Textbox(label="Lyrics")
-                
+
             with gr.Column():
                 num_segments = gr.Number(label="Number of Segments", value=2, interactive=True)
-                max_new_tokens = gr.Slider(label="Duration of song", minimum=1, maximum=30, step=1, value=5, interactive=True)
+                max_new_tokens = gr.Slider(label="Duration of song", minimum=1, maximum=30, step=1, value=5,
+                                           interactive=True)
                 submit_btn = gr.Button("Submit")
                 music_out = gr.Audio(label="Audio Result")
 
         gr.Examples(
-            examples = [
+            examples=[
                 [
                     "female blues airy vocal bright vocal piano sad romantic guitar jazz",
                     """[verse]
@@ -467,17 +526,17 @@ Through the highs and lows, I'mma keep it real
 Living out my dreams with this mic and a deal
                     """
                 ]
-            ], 
-             inputs = [genre_txt, lyrics_txt],
-            outputs = [music_out],
-            cache_examples = True,
+            ],
+            inputs=[genre_txt, lyrics_txt],
+            outputs=[music_out],
+            cache_examples=True,
             cache_mode="eager",
             fn=infer
         )
-    
+
     submit_btn.click(
-        fn = infer, 
-        inputs = [genre_txt, lyrics_txt, num_segments, max_new_tokens],
-        outputs = [music_out]
+        fn=infer,
+        inputs=[genre_txt, lyrics_txt, num_segments, max_new_tokens],
+        outputs=[music_out]
     )
-demo.queue().launch(show_error=True)
+demo.queue().launch(show_error=True)