chatterbox/_setup_helper.py at master · TheCommandCraft/chatterbox · GitHub

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
import subprocess
import sys
import re


def detect_cuda_version():
    try:
        result = subprocess.run(
            ["nvcc", "--version"],
            capture_output=True, text=True, timeout=10,
        )
        m = re.search(r"release (\d+)\.(\d+)", result.stdout)
        if m:
            return int(m.group(1)), int(m.group(2))
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass
    return None, None


def wheel_suffix(major, minor):
    if major is None:
        return None
    if major >= 13:
        return "cu130"   # CUDA 13.x — use cu130 for native Blackwell/sm_12x support
    if major == 12:
        if minor >= 6:
            return "cu126"
        elif minor >= 4:
            return "cu124"
        elif minor >= 1:
            return "cu121"
        else:
            return "cu118"
    if major == 11:
        return "cu118"
    return None


def run(cmd):
    print(f"\n>>> {' '.join(cmd)}\n")
    subprocess.run(cmd, check=True)


major, minor = detect_cuda_version()
suffix = wheel_suffix(major, minor)

py_ver = sys.version_info[:2]
is_py314_plus = py_ver >= (3, 14)

if is_py314_plus:
    # pyproject.toml requires torch>=2.9.0 for Python 3.14+; don't pin so pip picks the latest.
    torch_reqs = ["torch", "torchaudio"]
else:
    torch_reqs = ["torch==2.6.0", "torchaudio==2.6.0"]

if suffix:
    print(f"Detected CUDA {major}.{minor} -> using {suffix}")
    print(f"Python {py_ver[0]}.{py_ver[1]} -> torch requirements: {torch_reqs}")
    index_url = f"https://download.pytorch.org/whl/{suffix}"
    run([sys.executable, "-m", "pip", "install"] + torch_reqs + ["--index-url", index_url])
else:
    print("WARNING: nvcc not found -- could not detect CUDA version.")
    print("         Installing CPU-only PyTorch. GPU will NOT be used.")
    print("         Ensure nvcc.exe is on your PATH and re-run setup.bat")
    run([sys.executable, "-m", "pip", "install"] + torch_reqs)

print("\nInstalling ChatterBox without spacy-pkuseg (optional Chinese-only dep that has Windows build issues)...")
run([sys.executable, "-m", "pip", "install", "-e", ".", "--no-deps"])

print("\nInstalling remaining dependencies (excluding spacy-pkuseg)...")
other_deps = [
    "librosa==0.11.0", "s3tokenizer", "transformers==5.2.0", "diffusers==0.29.0",
    "conformer==0.3.2", "safetensors==0.5.3", "pykakasi", "gradio==6.8.0",
    "pyloudnorm", "omegaconf", "git+https://github.com/resemble-ai/Perth.git@master",
]
run([sys.executable, "-m", "pip", "install"] + other_deps)

print("\n=== Setup complete! Run run.bat to launch the GUI. ===")