Update app.py

app.py

@@ -4,15 +4,7 @@ import os
 import shutil
 import tempfile
 
-"""
-# Set the PATH and LD_LIBRARY_PATH for CUDA 12.3
-cuda_bin_path = "/usr/local/cuda/bin"
-cuda_lib_path = "/usr/local/cuda/lib64"
-
-# Update the environment variables
-os.environ['PATH'] = f"{cuda_bin_path}:{os.environ.get('PATH', '')}"
-os.environ['LD_LIBRARY_PATH'] = f"{cuda_lib_path}:{os.environ.get('LD_LIBRARY_PATH', '')}"
-"""
+is_shared_ui = True if "fffiloni/YuE" in os.environ['SPACE_ID'] else False
 
 # Install required package
 def install_flash_attn():
@@ -136,7 +128,7 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
         "--output_dir", f"{output_dir}",
         "--cuda_idx", "0",
         "--max_new_tokens", f"{max_new_tokens}",
-        "--disable_offload_model"
+        #"--disable_offload_model"
     ]
 
     # Set up environment variables for CUDA with optimized settings
@@ -187,7 +179,20 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
 
 with gr.Blocks() as demo:
     with gr.Column():
-        gr.Markdown("# YuE")
+        gr.Markdown("# YuE: Open Music Foundation Models for Full-Song Generation")
+        gr.HTML("""
+        <div style="display:flex;column-gap:4px;">
+            <a href="https://github.com/multimodal-art-projection/YuE">
+                <img src='https://img.shields.io/badge/GitHub-Repo-blue'>
+            </a> 
+            <a href="https://map-yue.github.io">
+                <img src='https://img.shields.io/badge/Project-Page-green'>
+            </a>
+            <a href="https://huggingface.co/spaces/fffiloni/YuE?duplicate=true">
+                <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-sm.svg" alt="Duplicate this Space">
+            </a>
+        </div>
+        """)
         with gr.Row():
             with gr.Column():
                 genre_txt = gr.Textbox(label="Genre")
@@ -207,20 +212,18 @@ Don't let this moment fade, hold me close tonight
 With you here beside me, everything's alright
 Can't imagine life alone, don't want to let you go
 Stay with me forever, let our love just flow
-
-[verse]
-Moonlight paints a picture upon your lovely face
-Every glance between us fills the empty space
-Time stands still around us when you're in my arms
-Nothing else can matter, safe from any harm
                             """
                         ]
                     ], 
                     inputs = [genre_txt, lyrics_txt]
                 )
             with gr.Column():
-                num_segments = gr.Number(label="Number of Song Segments", info="number of paragraphs", value=2, interactive=True)
-                max_new_tokens = gr.Slider(label="Max New Tokens / Duration", info="1000 token = 10 seconds", minimum=500, maximum="24000", step=500, value=1500, interactive=True)
+                if is_shared_ui:
+                    num_segments = gr.Number(label="Number of Song Segments", value=2, interactive=False)
+                    max_new_tokens = gr.Slider(label="Max New Tokens", minimum=500, maximum="24000", step=500, value=3000, interactive=False)
+                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)
                 submit_btn = gr.Button("Submit")
                 music_out = gr.Audio(label="Audio Result")