Update app.py

app.py

@@ -143,20 +143,38 @@ def optimize_gpu_settings():
 
 def install_flash_attn():
     try:
-        logging.info("Installing flash-attn...")
-        subprocess.run(
-            ["pip", "install", "flash-attn", "--no-build-isolation"], 
-            check=True,
-            capture_output=True
-        )
+        if not torch.cuda.is_available():
+            logging.warning("GPU not available, skipping flash-attn installation")
+            return False
+            
+        cuda_version = torch.version.cuda
+        if cuda_version is None:
+            logging.warning("CUDA not available, skipping flash-attn installation")
+            return False
+            
+        logging.info(f"Detected CUDA version: {cuda_version}")
+        
+        # CUDA 11.8 specific wheel for Python 3.10
+        if cuda_version.startswith("11.8"):
+            flash_attn_url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu11torch2.2cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"
+            subprocess.run(
+                ["pip", "install", flash_attn_url],
+                check=True,
+                capture_output=True
+            )
+        else:
+            logging.warning(f"Unsupported CUDA version: {cuda_version}, skipping flash-attn installation")
+            return False
+            
         logging.info("flash-attn installed successfully!")
-    except subprocess.CalledProcessError as e:
-        logging.error(f"Failed to install flash-attn: {e}")
-        raise
+        return True
+    except Exception as e:
+        logging.warning(f"Failed to install flash-attn: {e}")
+        return False
 
 def initialize_system():
     optimize_gpu_settings()
-    install_flash_attn()
+    has_flash_attn = install_flash_attn()
     
     from huggingface_hub import snapshot_download
     
@@ -177,42 +195,6 @@ def initialize_system():
         logging.error(f"Directory error: {e}")
         raise
 
-@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):
-    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)
-    content = content.strip() + "\n\n"
-    content = content.replace("\r\n", "\n").replace("\r", "\n")
-    temp_file.write(content)
-    temp_file.close()
-    logging.debug(f"Temporary file created: {temp_file.name}")
-    return temp_file.name
-
-def get_last_mp3_file(output_dir):
-    mp3_files = [f for f in os.listdir(output_dir) if f.endswith('.mp3')]
-    if not mp3_files:
-        logging.warning("No MP3 files found")
-        return None
-    
-    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]
-
-
-
-
-
 def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
     try:
         # 모델 선택 및 설정
@@ -234,7 +216,7 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
         os.makedirs(output_dir, exist_ok=True)
         empty_output_folder(output_dir)
 
-        # 명령어 구성
+        # 기본 명령어 구성
         command = [
             "python", "infer.py",
             "--stage1_model", model_path,
@@ -247,21 +229,31 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
             "--cuda_idx", "0",
             "--max_new_tokens", str(actual_max_tokens),
             "--temperature", str(config['temperature']),
-            "--disable_offload_model",
-            "--use_flash_attention_2",
-            "--bf16",
             "--chorus_strength", str(config['chorus_strength'])
         ]
 
+        # GPU가 있고 flash-attn이 설치된 경우에만 추가 옵션 적용
+        if torch.cuda.is_available():
+            command.extend([
+                "--disable_offload_model",
+                "--bf16"
+            ])
+            try:
+                import flash_attn
+                command.append("--use_flash_attention_2")
+            except ImportError:
+                logging.info("flash-attn not available, skipping flash attention option")
+
         # 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"
-        })
+        if torch.cuda.is_available():
+            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)
@@ -288,6 +280,42 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
             except Exception as e:
                 logging.warning(f"Failed to remove temporary file {file}: {e}")
 
+@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):
+    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)
+    content = content.strip() + "\n\n"
+    content = content.replace("\r\n", "\n").replace("\r", "\n")
+    temp_file.write(content)
+    temp_file.close()
+    logging.debug(f"Temporary file created: {temp_file.name}")
+    return temp_file.name
+
+def get_last_mp3_file(output_dir):
+    mp3_files = [f for f in os.listdir(output_dir) if f.endswith('.mp3')]
+    if not mp3_files:
+        logging.warning("No MP3 files found")
+        return None
+    
+    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]
+
+
+
+
+
 # Gradio 인터페이스
 with gr.Blocks() as demo:
     with gr.Column():