109 lines
3.7 KiB
Python
109 lines
3.7 KiB
Python
import os
|
||
import sys
|
||
import argparse
|
||
import subprocess
|
||
import venv
|
||
|
||
# Force line buffering
|
||
sys.stdout.reconfigure(line_buffering=True)
|
||
|
||
def run_command(command, description, required=True):
|
||
print(f"[STATUS] {description}...")
|
||
try:
|
||
process = subprocess.Popen(
|
||
command,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
text=True,
|
||
shell=True,
|
||
bufsize=1
|
||
)
|
||
|
||
for line in process.stdout:
|
||
line_str = line.strip()
|
||
if "Downloading" in line_str or "Installing" in line_str:
|
||
print(f"[STATUS] {line_str}")
|
||
elif line_str:
|
||
print(f"[INFO] {line_str}")
|
||
|
||
process.wait()
|
||
if process.returncode != 0:
|
||
print(f"[ERROR] {description} failed (exit code {process.returncode})", file=sys.stderr)
|
||
if required:
|
||
return False
|
||
return True
|
||
except Exception as e:
|
||
print(f"[ERROR] Exception during {description}: {str(e)}", file=sys.stderr)
|
||
return not required
|
||
|
||
def verify_install(python_path):
|
||
print("[STATUS] Verifying PyTorch installation...")
|
||
verify_cmd = (
|
||
f'"{python_path}" -c "import torch; import transformers; import PIL; '
|
||
f'print(f\\"torch={{torch.__version__}}\\")"'
|
||
)
|
||
return run_command(verify_cmd, "Verifying installed packages", required=True)
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument('--data-dir', help='Writable directory for .venv (used by packaged app builds)')
|
||
args = parser.parse_args()
|
||
|
||
backend_dir = os.path.dirname(os.path.abspath(__file__))
|
||
base_dir = os.path.abspath(args.data_dir) if args.data_dir else os.path.dirname(backend_dir)
|
||
venv_dir = os.path.join(base_dir, ".venv")
|
||
requirements_path = os.path.join(backend_dir, "requirements.txt")
|
||
|
||
if not os.path.exists(venv_dir):
|
||
print("[STATUS] Creating Python virtual environment...")
|
||
try:
|
||
venv.create(venv_dir, with_pip=True)
|
||
print("[STATUS] Virtual environment successfully created.")
|
||
except Exception as e:
|
||
print(f"[ERROR] Failed to create virtual environment: {str(e)}", file=sys.stderr)
|
||
sys.exit(1)
|
||
else:
|
||
print("[STATUS] Virtual environment already exists.")
|
||
|
||
if os.name == 'nt':
|
||
python_path = os.path.join(venv_dir, "Scripts", "python.exe")
|
||
else:
|
||
python_path = os.path.join(venv_dir, "bin", "python")
|
||
|
||
if not os.path.exists(python_path):
|
||
print("[ERROR] Python executable not found inside virtual environment!", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
run_command(
|
||
f'"{python_path}" -m pip install --upgrade pip',
|
||
"Upgrading pip to latest version",
|
||
required=False
|
||
)
|
||
|
||
torch_index = "https://download.pytorch.org/whl/cpu"
|
||
torch_cmd = (
|
||
f'"{python_path}" -m pip install torch torchvision '
|
||
f'--index-url {torch_index}'
|
||
)
|
||
if not run_command(torch_cmd, "Installing PyTorch AI Engine (CPU version)"):
|
||
print("[ERROR] PyTorch installation failed. Ensure Python 3.10–3.12 is installed and you have internet access.", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
other_deps_cmd = (
|
||
f'"{python_path}" -m pip install '
|
||
f'"transformers<5" timm einops kornia pillow accelerate'
|
||
)
|
||
if not run_command(other_deps_cmd, "Installing HuggingFace Transformers & PIL library"):
|
||
print("[ERROR] Dependency installation failed!", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
if not verify_install(python_path):
|
||
print("[ERROR] Package verification failed after installation!", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
print("[STATUS] Setup completed successfully!")
|
||
print("[STATUS] DONE")
|
||
|
||
if __name__ == '__main__':
|
||
main()
|