Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,995 Bytes
b7658fb 6b64262 3f97053 4324db0 3f97053 1cb26d6 4324db0 6b64262 3f97053 6b64262 1cb26d6 6b64262 4324db0 1cb26d6 6b64262 1cb26d6 4324db0 1cb26d6 21a831f 4324db0 1cb26d6 4324db0 1cb26d6 4324db0 1cb26d6 4324db0 21a831f 1cb26d6 4324db0 21a831f 4324db0 6b64262 21a831f 6b64262 4324db0 6b64262 21a831f 6b64262 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 |
import os
import subprocess
import sys
from pathlib import Path
# --- 0. Hardcoded Toggle for Execution Environment ---
# Set this to True to use Hugging Face ZeroGPU
# Set this to False to use a pure CPU environment
USE_ZEROGPU = True
# --- 1. Clone the VibeVoice Repository ---
repo_dir = "VibeVoice"
if not os.path.exists(repo_dir):
print("Cloning the VibeVoice repository...")
try:
subprocess.run(
["git", "clone", "https://github.com/microsoft/VibeVoice.git"],
check=True,
capture_output=True,
text=True
)
print("Repository cloned successfully.")
except subprocess.CalledProcessError as e:
print(f"Error cloning repository: {e.stderr}")
sys.exit(1)
else:
print("Repository already exists. Skipping clone.")
# --- 2. Install Dependencies ---
os.chdir(repo_dir)
print(f"Changed directory to: {os.getcwd()}")
# Install the main package
print("Installing the VibeVoice package...")
try:
subprocess.run(
[sys.executable, "-m", "pip", "install", "-e", "."],
check=True,
capture_output=True,
text=True
)
print("Package installed successfully.")
except subprocess.CalledProcessError as e:
print(f"Error installing package: {e.stderr}")
sys.exit(1)
# Install 'spaces' if using ZeroGPU
if USE_ZEROGPU:
print("Installing the 'spaces' library for ZeroGPU...")
try:
subprocess.run(
[sys.executable, "-m", "pip", "install", "huggingface-hub", "gradio", "spaces"],
check=True,
capture_output=True,
text=True
)
print("'spaces' library installed successfully.")
except subprocess.CalledProcessError as e:
print(f"Error installing 'spaces' library: {e.stderr}")
sys.exit(1)
# --- 3. Modify the demo script based on the toggle ---
demo_script_path = Path("demo/gradio_demo.py")
print(f"Reading {demo_script_path}...")
try:
file_content = demo_script_path.read_text()
if USE_ZEROGPU:
print("Optimizing for ZeroGPU execution...")
# Ensure the original GPU block is present
original_block = """ self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(
self.model_path,
torch_dtype=torch.bfloat16,
device_map='cuda',
attn_implementation="flash_attention_2",
)"""
if original_block in file_content:
# Add 'import spaces' at the beginning of the file
modified_content = "import spaces\n" + file_content
# Decorate the model loading and generation functions with @spaces.GPU
# This is a robust way to ensure both setup and inference get GPU access
modified_content = modified_content.replace(
"class VibeVoiceGradioInterface:",
"@spaces.GPU\nclass VibeVoiceGradioInterface:"
)
print("Script modified for ZeroGPU successfully.")
# Write the modified content back to the file
demo_script_path.write_text(modified_content)
else:
print("Warning: Original GPU-specific model loading block not found. The script might have been updated. Proceeding with potential ZeroGPU compatibility.")
else:
print("Modifying for CPU execution...")
# 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",
)"""
# 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 for CPU 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)
# --- 4. Launch the Gradio Demo ---
model_id = "microsoft/VibeVoice-1.5B"
# Construct the command as specified in the README
command = [
"python",
str(demo_script_path),
"--model_path",
model_id,
"--share"
]
print(f"Launching Gradio demo with command: {' '.join(command)}")
# This command will start the Gradio server
subprocess.run(command) |