feat: enhance Python environment setup with improved detection and verification methods

backend/setup.py

@@ -1,12 +1,13 @@
 import os
 import sys
+import argparse
 import subprocess
 import venv
 
 # Force line buffering
 sys.stdout.reconfigure(line_buffering=True)
 
-def run_command(command, description):
+def run_command(command, description, required=True):
     print(f"[STATUS] {description}...")
     try:
         process = subprocess.Popen(
@@ -17,30 +18,42 @@ def run_command(command, description):
             shell=True,
             bufsize=1
         )
-        
-        # Read output in real-time
+
         for line in process.stdout:
             line_str = line.strip()
-            # If the output has pip download status, keep it readable but don't spam too much
             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] Command failed with code {process.returncode}", file=sys.stderr)
-            return False
+            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 occurred: {str(e)}", file=sys.stderr)
-        return False
+        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():
-    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    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")
-    
-    # 1. Create venv if not exists
+    requirements_path = os.path.join(backend_dir, "requirements.txt")
+
     if not os.path.exists(venv_dir):
         print("[STATUS] Creating Python virtual environment...")
         try:
@@ -52,33 +65,42 @@ def main():
     else:
         print("[STATUS] Virtual environment already exists.")
 
-    # 2. Path to pip inside the venv
     if os.name == 'nt':
-        pip_path = os.path.join(venv_dir, "Scripts", "pip.exe")
         python_path = os.path.join(venv_dir, "Scripts", "python.exe")
     else:
-        pip_path = os.path.join(venv_dir, "bin", "pip")
         python_path = os.path.join(venv_dir, "bin", "python")
 
-    if not os.path.exists(pip_path):
-        print("[ERROR] pip executable not found inside virtual environment!", file=sys.stderr)
+    if not os.path.exists(python_path):
+        print("[ERROR] Python executable not found inside virtual environment!", file=sys.stderr)
         sys.exit(1)
 
-    # 3. Upgrade pip
-    run_command(f'"{pip_path}" install --upgrade pip', "Upgrading pip to latest version")
+    run_command(
+        f'"{python_path}" -m pip install --upgrade pip',
+        "Upgrading pip to latest version",
+        required=False
+    )
 
-    # 4. Install dependencies
-    requirements_path = os.path.join(base_dir, "backend", "requirements.txt")
-    
-    # Use extra index URL for CPU-only torch to download faster (150MB instead of 2.5GB+ GPU version)
-    # This is perfect for single images on desktop!
-    pip_install_cmd = f'"{pip_path}" install -r "{requirements_path}" --extra-index-url https://download.pytorch.org/whl/cpu'
-    
-    success = run_command(pip_install_cmd, "Installing PyTorch, Transformers, and PIL dependencies")
-    if not success:
+    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")