Update app.py

app.py

@@ -67,8 +67,6 @@ from vocoder import build_codec_model, process_audio
 from post_process_audio import replace_low_freq_with_energy_matched
 import re
 
-is_shared_ui = True if "innova-ai/YuE-music-generator-demo" in os.environ['SPACE_ID'] else False
-
 def empty_output_folder(output_dir):
     # List all files in the output directory
     files = os.listdir(output_dir)
@@ -114,7 +112,7 @@ model.eval()
 
 def generate_music(
     # stage1_model="m-a-p/YuE-s1-7B-anneal-en-cot",
-    max_new_tokens=3000,
+    max_new_tokens=5,
     run_n_segments=2,
     genre_txt=None,
     lyrics_txt=None,
@@ -136,7 +134,7 @@ def generate_music(
     
     model = model
     cuda_idx = cuda_idx
-    max_new_tokens = max_new_tokens
+    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)
 
@@ -398,11 +396,10 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments=2, 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 subprocess.CalledProcessError as e:
-        print(f"Error occurred: {e}")
-        return None
+    except Exception as e:
+        gr.warning("An Error Occured: " + e)
+        return none
     finally:
-        # Clean up temporary files
         print("Temporary files deleted.")
 
 # Gradio 
@@ -429,12 +426,8 @@ with gr.Blocks() as demo:
                 lyrics_txt = gr.Textbox(label="Lyrics")
                 
             with gr.Column():
-                if is_shared_ui:
-                    num_segments = gr.Number(label="Number of Segments", value=2, interactive=True)
-                    max_new_tokens = gr.Slider(label="Max New Tokens", info="100 tokens equals 1 second long music", minimum=100, maximum="3000", step=100, value=500, interactive=True) # increase it after testing
-                else:
-                    num_segments = gr.Number(label="Number of Song Segments", value=2, interactive=True)
-                    max_new_tokens = gr.Slider(label="Max New Tokens", minimum=500, maximum="24000", step=500, value=3000, interactive=True)
+                num_segments = gr.Number(label="Number of Segments", value=2, interactive=True)
+                max_new_tokens = gr.Slider(label="Duration of song", minimum=100, maximum=3000, step=100, value=500, interactive=True)
                 submit_btn = gr.Button("Submit")
                 music_out = gr.Audio(label="Audio Result")
 
@@ -477,8 +470,8 @@ Living out my dreams with this mic and a deal
             ], 
              inputs = [genre_txt, lyrics_txt],
             outputs = [music_out],
-            cache_examples = False,
-            # cache_mode="lazy",
+            cache_examples = True,
+            cache_mode="eager",
             fn=infer
         )
     
@@ -487,4 +480,4 @@ Living out my dreams with this mic and a deal
         inputs = [genre_txt, lyrics_txt, num_segments, max_new_tokens],
         outputs = [music_out]
     )
-demo.queue().launch(show_api=False, show_error=True)
+demo.queue().launch(show_error=True)