Update app.py

app.py

@@ -7,50 +7,14 @@ import spaces
 import torch
 import os
 import sys
-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'))
-import argparse
-import numpy as np
-import json
-from omegaconf import OmegaConf
-import torchaudio
-from torchaudio.transforms import Resample
-import soundfile as sf
-
-import uuid
-from tqdm import tqdm
-from einops import rearrange
-from codecmanipulator import CodecManipulator
-from mmtokenizer import _MMSentencePieceTokenizer
-from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
-import glob
-import time
-import copy
-from collections import Counter
-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
-
-is_shared_ui = True if "innova-ai/YuE-music-generator-demo" in os.environ['SPACE_ID'] else False
-
-# Install required package
-def install_flash_attn():
-    try:
-        print("Installing flash-attn...")
-        # Install flash attention
-        subprocess.run(
-            "pip install flash-attn --no-build-isolation",
-            env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
-            shell=True,
-        )
-        print("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()
+print("Installing flash-attn...")
+# Install flash attention
+subprocess.run(
+    "pip install flash-attn --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
 
 from huggingface_hub import snapshot_download 
 
@@ -78,6 +42,33 @@ except FileNotFoundError:
     print(f"Directory not found: {inference_dir}")
     exit(1)
 
+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'))
+import argparse
+import numpy as np
+import json
+from omegaconf import OmegaConf
+import torchaudio
+from torchaudio.transforms import Resample
+import soundfile as sf
+
+import uuid
+from tqdm import tqdm
+from einops import rearrange
+from codecmanipulator import CodecManipulator
+from mmtokenizer import _MMSentencePieceTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
+import glob
+import time
+import copy
+from collections import Counter
+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
+
+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)