Update app.py

app.py

@@ -3,176 +3,181 @@ import subprocess
 import os 
 import shutil
 import tempfile
+import torch
+import logging
+import numpy as np
+from concurrent.futures import ThreadPoolExecutor
+from functools import lru_cache
 
-# Install required package
+# 로깅 설정
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler('yue_generation.log'),
+        logging.StreamHandler()
+    ]
+)
+
+# GPU 설정 최적화
+def optimize_gpu_settings():
+    if torch.cuda.is_available():
+        # L40S에 최적화된 설정
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cudnn.deterministic = False
+        torch.backends.cudnn.enabled = True
+        
+        # GPU 메모리 설정
+        torch.cuda.empty_cache()
+        torch.cuda.set_device(0)
+        
+        logging.info(f"Using GPU: {torch.cuda.get_device_name(0)}")
+        logging.info(f"Available GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
+    else:
+        logging.warning("GPU not available!")
+
+# flash-attn 설치 함수 개선
 def install_flash_attn():
     try:
-        print("Installing flash-attn...")
+        logging.info("Installing flash-attn...")
         subprocess.run(
             ["pip", "install", "flash-attn", "--no-build-isolation"], 
-            check=True
+            check=True,
+            capture_output=True
         )
-        print("flash-attn installed successfully!")
+        logging.info("flash-attn installed successfully!")
     except subprocess.CalledProcessError as e:
-        print(f"Failed to install flash-attn: {e}")
-        exit(1)
-
-# Install flash-attn
-install_flash_attn()
+        logging.error(f"Failed to install flash-attn: {e}")
+        raise
 
-from huggingface_hub import snapshot_download 
-
-# Create xcodec_mini_infer folder
-folder_path = './inference/xcodec_mini_infer'
+# 초기화 함수
+def initialize_system():
+    optimize_gpu_settings()
+    install_flash_attn()
+    
+    from huggingface_hub import snapshot_download
+    
+    # xcodec_mini_infer 폴더 생성
+    folder_path = './inference/xcodec_mini_infer'
+    os.makedirs(folder_path, exist_ok=True)
+    logging.info(f"Created folder at: {folder_path}")
 
-# Create the folder if it doesn't exist
-if not os.path.exists(folder_path):
-    os.mkdir(folder_path)
-    print(f"Folder created at: {folder_path}")
-else:
-    print(f"Folder already exists at: {folder_path}")
+    # 모델 다운로드
+    snapshot_download(
+        repo_id="m-a-p/xcodec_mini_infer",
+        local_dir="./inference/xcodec_mini_infer",
+        resume_download=True
+    )
 
-snapshot_download(
-    repo_id = "m-a-p/xcodec_mini_infer",
-    local_dir = "./inference/xcodec_mini_infer"
-)
+    # inference 디렉토리로 이동
+    try:
+        os.chdir("./inference")
+        logging.info(f"Working directory changed to: {os.getcwd()}")
+    except FileNotFoundError as e:
+        logging.error(f"Directory error: {e}")
+        raise
 
-# Change to the "inference" directory
-inference_dir = "./inference"
-try:
-    os.chdir(inference_dir)
-    print(f"Changed working directory to: {os.getcwd()}")
-except FileNotFoundError:
-    print(f"Directory not found: {inference_dir}")
-    exit(1)
+# 캐시를 활용한 파일 관리
+@lru_cache(maxsize=100)
+def get_cached_file_path(content_hash, prefix):
+    return create_temp_file(content_hash, prefix)
 
 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)
-        try:
-            if os.path.isdir(file_path):
-                # If it's a directory, remove it recursively
-                shutil.rmtree(file_path)
-            else:
-                # If it's a file, delete it
-                os.remove(file_path)
-        except Exception as e:
-            print(f"Error deleting file {file_path}: {e}")
-
-# Function to create a temporary file with string content
+    try:
+        shutil.rmtree(output_dir)
+        os.makedirs(output_dir)
+        logging.info(f"Output folder cleaned: {output_dir}")
+    except Exception as e:
+        logging.error(f"Error cleaning output folder: {e}")
+        raise
+
 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.strip() + "\n\n"
     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("---")
-    
+    logging.debug(f"Temporary file created: {temp_file.name}")
     return temp_file.name
 
 def get_last_mp3_file(output_dir):
-    # List all files in the output directory
-    files = os.listdir(output_dir)
-    
-    # Filter only .mp3 files
-    mp3_files = [file for file in files if file.endswith('.mp3')]
-    
+    mp3_files = [f for f in os.listdir(output_dir) if f.endswith('.mp3')]
     if not mp3_files:
-        print("No .mp3 files found in the output folder.")
+        logging.warning("No MP3 files found")
         return None
     
-    # Get the full path for the mp3 files
-    mp3_files_with_path = [os.path.join(output_dir, file) for file in mp3_files]
-    
-    # Sort the files based on the modification time (most recent first)
-    mp3_files_with_path.sort(key=lambda x: os.path.getmtime(x), reverse=True)
-    
-    # Return the most recent .mp3 file
+    mp3_files_with_path = [os.path.join(output_dir, f) for f in mp3_files]
+    mp3_files_with_path.sort(key=os.path.getmtime, reverse=True)
     return mp3_files_with_path[0]
 
+# L40S에 최적화된 추론 함수
 def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
-    # Create temporary files
-    genre_txt_path = create_temp_file(genre_txt_content, prefix="genre_")
-    lyrics_txt_path = create_temp_file(lyrics_txt_content, prefix="lyrics_")
-
-    print(f"Genre TXT path: {genre_txt_path}")
-    print(f"Lyrics TXT path: {lyrics_txt_path}")
-
-    # Ensure the output folder exists
-    output_dir = "./output"
-    os.makedirs(output_dir, exist_ok=True)
-    print(f"Output folder ensured at: {output_dir}")
-
-    empty_output_folder(output_dir)
- 
-    # Command and arguments with optimized settings
-    command = [
-        "python", "infer.py",
-        "--stage1_model", "m-a-p/YuE-s1-7B-anneal-en-cot",
-        "--stage2_model", "m-a-p/YuE-s2-1B-general",
-        "--genre_txt", f"{genre_txt_path}",
-        "--lyrics_txt", f"{lyrics_txt_path}",
-        "--run_n_segments", f"{num_segments}",
-        "--stage2_batch_size", "4",
-        "--output_dir", f"{output_dir}",
-        "--cuda_idx", "0",
-        "--max_new_tokens", f"{max_new_tokens}",
-        "--disable_offload_model"
-    ]
-
-    # Set up environment variables for CUDA with optimized settings
-    env = os.environ.copy()
-    env.update({
-        "CUDA_VISIBLE_DEVICES": "0",
-        "CUDA_HOME": "/usr/local/cuda",
-        "PATH": f"/usr/local/cuda/bin:{env.get('PATH', '')}",
-        "LD_LIBRARY_PATH": f"/usr/local/cuda/lib64:{env.get('LD_LIBRARY_PATH', '')}"
-    })
-    
-    # Execute the command
     try:
-        subprocess.run(command, check=True, env=env)
-        print("Command executed successfully!")
+        # 임시 파일 생성
+        genre_txt_path = create_temp_file(genre_txt_content, prefix="genre_")
+        lyrics_txt_path = create_temp_file(lyrics_txt_content, prefix="lyrics_")
         
-        # Check and print the contents of the output folder
-        output_files = os.listdir(output_dir)
-        if output_files:
-            print("Output folder contents:")
-            for file in output_files:
-                print(f"- {file}")
-
-            last_mp3 = get_last_mp3_file(output_dir)
-
-            if last_mp3:
-                print("Last .mp3 file:", last_mp3)
-                return last_mp3
-            else:
-                return None
+        output_dir = "./output"
+        os.makedirs(output_dir, exist_ok=True)
+        empty_output_folder(output_dir)
+
+        # L40S에 최적화된 명령어
+        command = [
+            "python", "infer.py",
+            "--stage1_model", "m-a-p/YuE-s1-7B-anneal-en-cot",
+            "--stage2_model", "m-a-p/YuE-s2-1B-general",
+            "--genre_txt", genre_txt_path,
+            "--lyrics_txt", lyrics_txt_path,
+            "--run_n_segments", str(num_segments),
+            "--stage2_batch_size", "8",  # L40S에 맞게 증가
+            "--output_dir", output_dir,
+            "--cuda_idx", "0",
+            "--max_new_tokens", str(max_new_tokens),
+            "--disable_offload_model",
+            "--use_flash_attention_2",  # Flash Attention 2 활성화
+            "--bf16"  # BF16 정밀도 사용
+        ]
+
+        # CUDA 환경 변수 설정
+        env = os.environ.copy()
+        env.update({
+            "CUDA_VISIBLE_DEVICES": "0",
+            "CUDA_HOME": "/usr/local/cuda",
+            "PATH": f"/usr/local/cuda/bin:{env.get('PATH', '')}",
+            "LD_LIBRARY_PATH": f"/usr/local/cuda/lib64:{env.get('LD_LIBRARY_PATH', '')}",
+            "PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:512"
+        })
+
+        # 명령 실행
+        process = subprocess.run(command, env=env, check=True, capture_output=True)
+        logging.info("Inference completed successfully")
+
+        # 결과 처리
+        last_mp3 = get_last_mp3_file(output_dir)
+        if last_mp3:
+            logging.info(f"Generated audio file: {last_mp3}")
+            return last_mp3
         else:
-            print("Output folder is empty.")
+            logging.warning("No output audio file generated")
             return None
-    except subprocess.CalledProcessError as e:
-        print(f"Error occurred: {e}")
-        return None
+
+    except Exception as e:
+        logging.error(f"Inference error: {e}")
+        raise
     finally:
-        # Clean up temporary files
-        os.remove(genre_txt_path)
-        os.remove(lyrics_txt_path)
-        print("Temporary files deleted.")
+        # 임시 파일 정리
+        for file in [genre_txt_path, lyrics_txt_path]:
+            try:
+                os.remove(file)
+                logging.debug(f"Removed temporary file: {file}")
+            except Exception as e:
+                logging.warning(f"Failed to remove temporary file {file}: {e}")
 
-# Gradio 
+# Gradio 인터페이스
 with gr.Blocks() as demo:
     with gr.Column():
-        gr.Markdown("# YuE: Open Music Foundation Models for Full-Song Generation")
+        gr.Markdown("# YuE: Open Music Foundation Models for Full-Song Generation (L40S Optimized)")
         gr.HTML("""
         <div style="display:flex;column-gap:4px;">
             <a href="https://github.com/multimodal-art-projection/YuE">
@@ -181,24 +186,43 @@ with gr.Blocks() as demo:
             <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")
-                lyrics_txt = gr.Textbox(label="Lyrics")
+                genre_txt = gr.Textbox(
+                    label="Genre",
+                    placeholder="Enter music genre and style descriptions..."
+                )
+                lyrics_txt = gr.Textbox(
+                    label="Lyrics",
+                    placeholder="Enter song lyrics...",
+                    lines=10
+                )
                 
             with gr.Column():
-                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")
+                num_segments = gr.Number(
+                    label="Number of Song Segments",
+                    value=2,
+                    minimum=1,
+                    maximum=4,
+                    step=1,
+                    interactive=True
+                )
+                max_new_tokens = gr.Slider(
+                    label="Max New Tokens",
+                    minimum=500,
+                    maximum=32000,  # L40S의 큰 메모리를 활용
+                    step=500,
+                    value=4000,
+                    interactive=True
+                )
+                submit_btn = gr.Button("Generate Music", variant="primary")
+                music_out = gr.Audio(label="Generated Audio")
 
         gr.Examples(
-            examples = [
+            examples=[
                 [
                     "female blues airy vocal bright vocal piano sad romantic guitar jazz",
                     """[verse]
@@ -233,13 +257,26 @@ 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]
+            ],
+            inputs=[genre_txt, lyrics_txt]
         )
+
+    # 시스템 초기화
+    initialize_system()
     
+    # 이벤트 핸들러
     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_api=True, show_error=True)
+
+# 서버 설정으로 실행
+demo.queue(concurrency_count=2).launch(
+    server_name="0.0.0.0",
+    server_port=7860,
+    share=True,
+    enable_queue=True,
+    show_api=True,
+    show_error=True
+)