Spaces:
Sleeping
Sleeping
File size: 1,668 Bytes
443e99e |
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 |
#!/usr/bin/env python
"""
Startup script that ensures CUDA PyTorch is installed before running Flask app.
"""
import subprocess
import sys
from pathlib import Path
def ensure_cuda_torch():
"""Ensure CUDA-enabled PyTorch is installed."""
try:
import torch
if torch.cuda.is_available():
print(f"✓ CUDA available: {torch.cuda.get_device_name(0)}")
return True
else:
print("⚠ CUDA not available in current PyTorch installation")
print("Installing CUDA-enabled PyTorch...")
subprocess.run([
sys.executable, "-m", "pip", "install",
"torch", "torchvision",
"--index-url", "https://download.pytorch.org/whl/cu121",
"--upgrade"
], check=True)
# Re-import to check
import importlib
importlib.reload(torch)
if torch.cuda.is_available():
print(f"✓ CUDA now available: {torch.cuda.get_device_name(0)}")
return True
else:
print("⚠ Still no CUDA after installation. Using CPU mode.")
return False
except Exception as e:
print(f"Error checking CUDA: {e}")
return False
if __name__ == '__main__':
print("Checking GPU availability...")
ensure_cuda_torch()
print("\nStarting PDF Layout Extractor Flask App...")
print("Open your browser to http://localhost:5000\n")
from app import app
# Disable reloader to avoid environment discrepancies in child process
app.run(debug=False, use_reloader=False, host='0.0.0.0', port=5000)
|