Update app.py

app.py

@@ -2,7 +2,6 @@ import os
 import subprocess
 import torch
 import gradio as gr
-from transformers import AutoModelForTextToSpeech, AutoTokenizer, AutoConfig
 import scipy.io.wavfile
 import numpy as np
 from datetime import datetime
@@ -14,14 +13,21 @@ logger = logging.getLogger(__name__)
 
 # Function to install dependencies
 def install_dependencies():
-    """Install required dependencies for VibeVoice."""
+    """Install required dependencies, including Microsoft's custom transformers fork."""
     try:
-        # Install required packages
+        # Install core dependencies from requirements.txt
         subprocess.run(["pip", "install", "-r", "requirements.txt"], check=True)
+        
+        # Install Microsoft's custom transformers fork for VibeVoice
+        subprocess.run([
+            "pip", "install", "git+https://github.com/microsoft/VibeVoice.git#subdirectory=transformers"
+        ], check=True)
+        
         # Install flash-attn only if GPU is available
         if torch.cuda.is_available():
             subprocess.run(["pip", "install", "flash-attn==2.6.3", "--no-build-isolation"], check=True)
-        # Install ffmpeg for audio processing (Linux-specific; adjust for other OS)
+        
+        # Install ffmpeg for audio processing
         if os.name == "posix":
             subprocess.run(["apt", "update"], check=True)
             subprocess.run(["apt", "install", "ffmpeg", "-y"], check=True)
@@ -31,6 +37,20 @@ def install_dependencies():
         logger.error(f"Failed to install dependencies: {e}")
         raise
 
+# Verify transformers fork installation
+def verify_transformers_fork():
+    """Verify that the custom transformers fork is installed correctly."""
+    try:
+        from transformers import AutoModelForTextToSpeech
+        logger.info("Custom transformers fork with AutoModelForTextToSpeech loaded successfully")
+    except ImportError as e:
+        logger.error(
+            f"Failed to import AutoModelForTextToSpeech. Ensure the custom transformers fork is installed.\n"
+            f"Run: pip install git+https://github.com/microsoft/VibeVoice.git#subdirectory=transformers\n"
+            f"Error: {e}"
+        )
+        raise
+
 # Determine device and attention implementation
 def get_device():
     """Determine the device (CPU/GPU) and attention implementation."""
@@ -43,11 +63,13 @@ def get_device():
 # Load model and tokenizer
 def load_model(model_name="microsoft/VibeVoice-1.5B"):
     """Load VibeVoice model and tokenizer with memory optimization."""
+    from transformers import AutoModelForTextToSpeech, AutoTokenizer, AutoConfig
+    
     device, attn_impl = get_device()
     try:
         # Load config with specified attention implementation
         config = AutoConfig.from_pretrained(model_name, attn_implementation=attn_impl)
-        # Enable memory-efficient loading with torch_dtype
+        # Enable memory-efficient loading
         model = AutoModelForTextToSpeech.from_pretrained(
             model_name, config=config, torch_dtype=torch.float16 if device == "cuda" else torch.float32
         )
@@ -82,14 +104,14 @@ def generate_audio(text_input, speaker_names, model, tokenizer, device, output_p
         # Generate audio with memory management
         with torch.no_grad():
             if device == "cuda":
-                with torch.cuda.amp.autocast():  # Mixed precision for GPU
+                with torch.cuda.amp.autocast():
                     audio_output = model.generate(**inputs)
             else:
                 audio_output = model.generate(**inputs)
 
         # Convert to numpy array
         audio_np = audio_output.cpu().numpy().squeeze()
-        if audio_np.ndim > 1:  # Handle potential multi-channel output
+        if audio_np.ndim > 1:
             audio_np = audio_np[0]
 
         # Normalize audio to [-1, 1]
@@ -121,8 +143,9 @@ def gradio_interface(text_input, speakers_str):
     except ValueError as e:
         return str(e)
 
-# Global model and tokenizer (loaded once)
+# Global model and tokenizer
 try:
+    verify_transformers_fork()
     model, tokenizer, device = load_model()
 except Exception as e:
     logger.error(f"Failed to initialize: {e}")