Update app.py

app.py

@@ -1,21 +1,9 @@
 import os
 import subprocess
 import sys
-import fileinput
-import gradio as gr
+from pathlib import Path
 
-# --- 1. Setup Environment ---
-# Install ffmpeg if it's not already installed by packages.txt
-print("Updating apt and installing ffmpeg...")
-try:
-    subprocess.run(["apt-get", "update", "-y"], check=True)
-    subprocess.run(["apt-get", "install", "ffmpeg", "-y"], check=True)
-    print("ffmpeg installed successfully.")
-except subprocess.CalledProcessError as e:
-    print(f"Error installing ffmpeg: {e.stderr}")
-    # Continue anyway, as it might be pre-installed
-
-# --- 2. Clone the VibeVoice Repository ---
+# --- 1. Clone the VibeVoice Repository ---
 repo_dir = "VibeVoice"
 if not os.path.exists(repo_dir):
     print("Cloning the VibeVoice repository...")
@@ -33,9 +21,10 @@ if not os.path.exists(repo_dir):
 else:
     print("Repository already exists. Skipping clone.")
 
-# --- 3. Install the Package ---
+# --- 2. Install the Package ---
 os.chdir(repo_dir)
 print(f"Changed directory to: {os.getcwd()}")
+
 print("Installing the VibeVoice package...")
 try:
     subprocess.run(
@@ -49,38 +38,51 @@ except subprocess.CalledProcessError as e:
     print(f"Error installing package: {e.stderr}")
     sys.exit(1)
 
-
-# --- 4. Modify the demo script for CPU execution ---
-demo_script_path = "demo/gradio_demo.py"
+# --- 3. Modify the demo script for CPU execution (Robust Method) ---
+demo_script_path = Path("demo/gradio_demo.py")
 print(f"Modifying {demo_script_path} for CPU execution...")
 
-# Use fileinput to perform in-place replacement
-with fileinput.FileInput(demo_script_path, inplace=True) as file:
-    for line in file:
-        # Change the model loading line to be CPU-compatible
-        if 'VibeVoiceForConditionalGenerationInference.from_pretrained(' in line:
-            print('        self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(')
-            print('            self.model_path,')
-            print('            torch_dtype=torch.float32,  # Use float32 for CPU')
-            print('            device_map="cpu",')
-            print('        )')
-            # Skip the next few lines of the original arguments
-            next(file, None)
-            next(file, None)
-            next(file, None)
-        else:
-            print(line, end='')
+try:
+    # Read the entire file content
+    file_content = demo_script_path.read_text()
+
+    # Define the original GPU-specific model loading block
+    original_block = """        self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(
+            self.model_path,
+            torch_dtype=torch.bfloat16,
+            device_map='cuda',
+            attn_implementation="flash_attention_2",
+        )"""
+
+    # Define the new CPU-compatible block
+    replacement_block = """        self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(
+            self.model_path,
+            torch_dtype=torch.float32,  # Use float32 for CPU
+            device_map="cpu",
+        )"""
 
-print("Script modified successfully.")
+    # Replace the entire block
+    if original_block in file_content:
+        modified_content = file_content.replace(original_block, replacement_block)
+        
+        # Write the modified content back to the file
+        demo_script_path.write_text(modified_content)
+        print("Script modified successfully.")
+    else:
+        print("Warning: GPU-specific model loading block not found. The script might have been updated. Proceeding without modification.")
+
+except Exception as e:
+    print(f"An error occurred while modifying the script: {e}")
+    sys.exit(1)
 
 
-# --- 5. Launch the Gradio Demo ---
+# --- 4. Launch the Gradio Demo ---
 model_id = "microsoft/VibeVoice-1.5B"
 
 # Construct the command as specified in the README
 command = [
     "python",
-    demo_script_path,
+    str(demo_script_path),
     "--model_path",
     model_id,
     "--share"