106 lines
3.4 KiB
Python
106 lines
3.4 KiB
Python
"""
|
|
Create GPU-enabled OpenCLIP environment with complete dependency isolation
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import subprocess
|
|
import platform
|
|
|
|
def create_gpu_environment():
|
|
"""Create isolated GPU environment for OpenCLIP"""
|
|
print("Creating GPU-enabled OpenCLIP environment...")
|
|
|
|
venv_path = "openclip_gpu_env"
|
|
|
|
# Remove existing environment
|
|
if os.path.exists(venv_path):
|
|
print(f"Removing existing environment: {venv_path}")
|
|
import shutil
|
|
shutil.rmtree(venv_path)
|
|
|
|
# Create virtual environment
|
|
print(f"Creating virtual environment: {venv_path}")
|
|
result = subprocess.run([
|
|
sys.executable, '-m', 'venv', venv_path
|
|
], capture_output=True, text=True)
|
|
|
|
if result.returncode != 0:
|
|
print(f"Virtual environment creation failed: {result.stderr}")
|
|
return False
|
|
|
|
# Get virtual environment Python
|
|
if platform.system() == "Windows":
|
|
venv_python = os.path.join(venv_path, "Scripts", "python.exe")
|
|
else:
|
|
venv_python = os.path.join(venv_path, "bin", "python")
|
|
|
|
# Install PyTorch with CUDA 11.8
|
|
print("Installing PyTorch with CUDA 11.8...")
|
|
torch_cmd = [
|
|
venv_python, '-m', 'pip', 'install',
|
|
'torch==2.0.1+cu118', 'torchvision==0.15.2+cu118',
|
|
'--index-url', 'https://download.pytorch.org/whl/cu118'
|
|
]
|
|
|
|
result = subprocess.run(torch_cmd, capture_output=True, text=True)
|
|
if result.returncode != 0:
|
|
print(f"PyTorch installation failed: {result.stderr}")
|
|
return False
|
|
|
|
# Install OpenCLIP
|
|
print("Installing OpenCLIP...")
|
|
openclip_cmd = [venv_python, '-m', 'pip', 'install', 'open-clip-torch']
|
|
result = subprocess.run(openclip_cmd, capture_output=True, text=True)
|
|
if result.returncode != 0:
|
|
print(f"OpenCLIP installation failed: {result.stderr}")
|
|
return False
|
|
|
|
# Install Pillow
|
|
print("Installing Pillow...")
|
|
pillow_cmd = [venv_python, '-m', 'pip', 'install', 'Pillow']
|
|
result = subprocess.run(pillow_cmd, capture_output=True, text=True)
|
|
if result.returncode != 0:
|
|
print(f"Pillow installation failed: {result.stderr}")
|
|
return False
|
|
|
|
# Test the environment
|
|
print("Testing GPU environment...")
|
|
test_script = """
|
|
import torch
|
|
import open_clip
|
|
|
|
print("=== GPU ENVIRONMENT TEST ===")
|
|
print(f"PyTorch version: {torch.__version__}")
|
|
print(f"CUDA available: {torch.cuda.is_available()}")
|
|
if torch.cuda.is_available():
|
|
print(f"CUDA version: {torch.version.cuda}")
|
|
print(f"GPU count: {torch.cuda.device_count()}")
|
|
for i in range(torch.cuda.device_count()):
|
|
print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
|
|
|
|
# Test OpenCLIP
|
|
model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
|
|
print("OpenCLIP model loaded successfully")
|
|
|
|
if torch.cuda.is_available():
|
|
model = model.cuda()
|
|
print("Model moved to GPU - SUCCESS")
|
|
else:
|
|
print("Running on CPU")
|
|
"""
|
|
|
|
result = subprocess.run([venv_python, '-c', test_script], capture_output=True, text=True)
|
|
print(result.stdout)
|
|
|
|
if result.returncode == 0:
|
|
print("GPU environment setup successful!")
|
|
return True
|
|
else:
|
|
print("GPU environment test failed")
|
|
print(result.stderr)
|
|
return False
|
|
|
|
if __name__ == "__main__":
|
|
success = create_gpu_environment()
|
|
sys.exit(0 if success else 1) |