diff --git a/lollms/services/tti/diffusers/lollms_diffusers.py b/lollms/services/tti/diffusers/lollms_diffusers.py index 5ef1b1d..71994ad 100644 --- a/lollms/services/tti/diffusers/lollms_diffusers.py +++ b/lollms/services/tti/diffusers/lollms_diffusers.py @@ -6,7 +6,7 @@ from pathlib import Path import sys from lollms.app import LollmsApplication -from lollms.utilities import PackageManager, check_and_install_torch, find_next_available_filename, install_cuda, check_torch_version +from lollms.utilities import PackageManager, check_and_install_torch, find_next_available_filename import sys import requests diff --git a/lollms/services/tti/sd/lollms_sd.py b/lollms/services/tti/sd/lollms_sd.py index f9ce66a..7e2adb0 100644 --- a/lollms/services/tti/sd/lollms_sd.py +++ b/lollms/services/tti/sd/lollms_sd.py @@ -29,12 +29,13 @@ from typing import List, Dict, Any from ascii_colors import ASCIIColors, trace_exception from lollms.paths import LollmsPaths from lollms.tti import LollmsTTI -from lollms.utilities import git_pull, show_yes_no_dialog, run_script_in_env, create_env +from lollms.utilities import git_pull, show_yes_no_dialog, EnvironmentManager import subprocess import shutil from tqdm import tqdm import threading +em = EnvironmentManager() def download_file(url, folder_path, local_filename): @@ -64,7 +65,7 @@ def install_sd(lollms_app:LollmsApplication): shared_folder = root_dir/"shared" sd_folder = shared_folder / "auto_sd" ASCIIColors.cyan("Installing autosd conda environment with python 3.10") - create_env("autosd","3.10") + em.create_env("autosd","3.10") ASCIIColors.cyan("Done") if os.path.exists(str(sd_folder)): print("Repository already exists. Pulling latest changes...") diff --git a/lollms/utilities.py b/lollms/utilities.py index b477e66..5de74ea 100644 --- a/lollms/utilities.py +++ b/lollms/utilities.py @@ -35,106 +35,253 @@ import sys import git import mimetypes +import sys +import platform import subprocess - +import pkg_resources from functools import partial import pipmaster as pm if not pm.is_installed("Pillow"): pm.install("Pillow") from PIL import Image +import os +import subprocess +import sys +import platform +from enum import Enum +from pathlib import Path +import shutil -def create_env(env_name, python_version): - pass +class EnvManager(Enum): + CONDA = 'conda' + VENV = 'venv' + PYENV = 'pyenv' + PIP = 'pip' -def run_pip_in_env(env_name, pip_args, cwd=None): - import platform - # Set the current working directory if provided, otherwise use the current directory - if cwd is None: - cwd = os.getcwd() - - # Activate the Conda environment - python_path = Path(sys.executable).parent.parent/"miniconda3"/"envs"/env_name/"python" - ASCIIColors.yellow(f"Executing: {python_path} -m pip {pip_args}") - process = subprocess.Popen(f'{python_path} -m pip {pip_args}', shell=True) - - # Wait for the process to finish - process.wait() +class EnvironmentManager: + def __init__(self, preferred_manager=None): + """ + Initialize environment manager with optional preferred manager. + Args: + preferred_manager (str, optional): 'conda', 'venv', 'pyenv', or 'pip' + """ + self.preferred_manager = EnvManager(preferred_manager) if preferred_manager else None + self.manager = self._detect_env_manager() + self.manager_path = self._get_env_manager_path() + def _get_env_manager_path(self): + """Get the path of the environment manager executable""" + if platform.system() == 'Windows': + ext = '.exe' + else: + ext = '' + + if self.manager == EnvManager.CONDA: + # Check for portable conda + portable_conda = os.getenv('PORTABLE_CONDA_PATH') + if portable_conda: + return os.path.join(portable_conda, 'condabin', f'conda{ext}') + + # Check standard conda locations + possible_paths = [ + os.path.join(sys.prefix, 'condabin', f'conda{ext}'), + os.path.join(os.path.expanduser('~'), 'miniconda3', 'condabin', f'conda{ext}'), + os.path.join(os.path.expanduser('~'), 'anaconda3', 'condabin', f'conda{ext}') + ] + + for path in possible_paths: + if os.path.exists(path): + return path + + # Use which/where to find the executable + try: + if platform.system() == 'Windows': + result = subprocess.run(['where', self.manager.value], capture_output=True, text=True) + else: + result = subprocess.run(['which', self.manager.value], capture_output=True, text=True) + if result.returncode == 0: + return result.stdout.strip().split('\n')[0] + except: + pass + + return None -def run_python_script_in_env(env_name, script_path, cwd=None, wait=True): - import platform - # Set the current working directory if provided, otherwise use the current directory - if cwd is None: - cwd = os.getcwd() - - # Activate the Conda environment - python_path = Path(sys.executable).parent.parent/"miniconda3"/"envs"/env_name/"python" - ASCIIColors.yellow(f"Executing: {python_path} {script_path}") - process = subprocess.Popen(f'{python_path} {script_path}', shell=True) - - # Wait for the process to finish - if wait: - process.wait() - return process + def _detect_env_manager(self): + """Detect which environment manager to use based on context and preference""" + # If preferred manager is specified and available, use it + if self.preferred_manager: + if self._check_manager_available(self.preferred_manager): + return self.preferred_manager + + # Check if we're in a venv + if sys.prefix != sys.base_prefix: + return EnvManager.VENV + + # Check for conda (both portable and installed) + if self._check_manager_available(EnvManager.CONDA): + return EnvManager.CONDA + + # Check for pyenv + if self._check_manager_available(EnvManager.PYENV): + return EnvManager.PYENV + + # Default to pip + return EnvManager.PIP -def run_script_in_env(env_name, script_path, cwd=None): - import platform - # Set the current working directory if provided, otherwise use the current directory - if cwd is None: - cwd = os.path.dirname(script_path) - - # Activate the Conda environment - if platform.system()=="Windows": - python_path = Path(sys.executable).parent.parent/"miniconda3"/"condabin"/"conda" - subprocess.Popen(f'{python_path} activate {env_name} && {script_path}', shell=True, cwd=cwd) - else: - python_path = Path(sys.executable).parent.parent.parent/"miniconda3"/"bin"/"conda" - subprocess.Popen(f'source {python_path} activate {env_name} && {script_path}', shell=True, cwd=cwd) - # Activate the Conda environment - ASCIIColors.red("Python path:") - ASCIIColors.yellow(python_path) - #run_command(Commands.RUN, "-n", env_name, str(script_path), cwd=cwd) + def _check_manager_available(self, manager): + """Check if a specific environment manager is available""" + try: + if platform.system() == 'Windows': + result = subprocess.run(['where', manager.value], capture_output=True) + else: + result = subprocess.run(['which', manager.value], capture_output=True) + return result.returncode == 0 + except: + return False -def get_conda_path(): - import platform - if platform.system() == "Windows": - return Path(sys.executable).parent.parent / "miniconda3" / "condabin" / "conda" - else: - return Path(sys.executable).parent.parent.parent / "miniconda3" / "bin" / "conda" + def _get_env_python_path(self, env_name): + """Get the Python executable path for the environment""" + if platform.system() == 'Windows': + return os.path.join(env_name, 'Scripts', 'python.exe') + return os.path.join(env_name, 'bin', 'python') -def environment_exists(env_name): - from lollms.security import sanitize_shell_code - env_name = sanitize_shell_code(env_name) - env_name = Path(sys.executable).parent.parent / "miniconda3" / "env" / env_name - conda_path = get_conda_path() - ASCIIColors.yellow(f"Using conda from : {conda_path}") + def create_env(self, env_name, python_version): + """ + Create a new environment with specified Python version + Args: + env_name (str): Name of the environment + python_version (str): Python version to install (e.g., '3.8') + """ + try: + if self.manager == EnvManager.CONDA: + subprocess.run([self.manager_path, 'create', '-n', env_name, + f'python={python_version}', '-y'], check=True) + elif self.manager == EnvManager.VENV: + subprocess.run([f'python{python_version}', '-m', 'venv', env_name], check=True) + elif self.manager == EnvManager.PYENV: + subprocess.run([self.manager_path, 'install', python_version], check=True) + subprocess.run([self.manager_path, 'virtualenv', python_version, env_name], check=True) + else: + subprocess.run([sys.executable, '-m', 'venv', env_name], check=True) + except subprocess.CalledProcessError as e: + raise RuntimeError(f"Failed to create environment: {str(e)}") - result = subprocess.run(f'{conda_path} env list --json', shell=True, capture_output=True, text=True) - envs_info = json.loads(result.stdout) - env_names = [Path(env).name for env in envs_info['envs']] - return env_name in env_names + def run_pip_in_env(self, env_name, pip_args, cwd=None): + """ + Run pip commands in the specified environment + Args: + env_name (str): Name of the environment + pip_args (str): Arguments to pass to pip + cwd (str, optional): Working directory + """ + try: + if self.manager == EnvManager.CONDA: + cmd = f'"{self.manager_path}" run -n {env_name} pip {pip_args}' + else: + if platform.system() == 'Windows': + pip_path = os.path.join(env_name, 'Scripts', 'pip') + else: + pip_path = os.path.join(env_name, 'bin', 'pip') + cmd = f'"{pip_path}" {pip_args}' + + subprocess.run(cmd, shell=True, cwd=cwd, check=True) + except subprocess.CalledProcessError as e: + raise RuntimeError(f"Failed to run pip: {str(e)}") -def get_python_version(env_name): - from lollms.security import sanitize_shell_code - env_name = sanitize_shell_code(env_name) - conda_path = get_conda_path() - if environment_exists(env_name): - result = subprocess.run(f'{conda_path} run -n {env_name} python --version', shell=True, capture_output=True, text=True) - return result.stdout.strip() - else: - return "Environment does not exist." + def run_python_script_in_env(self, env_name, script_path, cwd=None, wait=True): + """ + Run a Python script in the specified environment + Args: + env_name (str): Name of the environment + script_path (str): Path to the Python script + cwd (str, optional): Working directory + wait (bool): Whether to wait for the script to complete + """ + try: + if self.manager == EnvManager.CONDA: + cmd = f'"{self.manager_path}" run -n {env_name} python "{script_path}"' + else: + python_path = self._get_env_python_path(env_name) + cmd = f'"{python_path}" "{script_path}"' + + if wait: + subprocess.run(cmd, shell=True, cwd=cwd, check=True) + else: + subprocess.Popen(cmd, shell=True, cwd=cwd) + except subprocess.CalledProcessError as e: + raise RuntimeError(f"Failed to run script: {str(e)}") -def remove_environment(env_name): - from lollms.security import sanitize_shell_code - env_name = sanitize_shell_code(env_name) - conda_path = get_conda_path() - if environment_exists(env_name): - process = subprocess.Popen(f'{conda_path} env remove --name {env_name} -y', shell=True) - process.wait() - return f"Environment '{env_name}' has been removed." - else: - return "Environment does not exist." + def run_script_in_env(self, env_name, script_path, cwd=None): + """ + Run any script in the specified environment + Args: + env_name (str): Name of the environment + script_path (str): Path to the script + cwd (str, optional): Working directory + """ + try: + if self.manager == EnvManager.CONDA: + cmd = f'"{self.manager_path}" run -n {env_name} "{script_path}"' + else: + cmd = os.path.join(env_name, 'bin', script_path) + + subprocess.run(cmd, shell=True, cwd=cwd, check=True) + except subprocess.CalledProcessError as e: + raise RuntimeError(f"Failed to run script: {str(e)}") + + def environment_exists(self, env_name): + """ + Check if the specified environment exists + Args: + env_name (str): Name of the environment + Returns: + bool: True if environment exists, False otherwise + """ + if self.manager == EnvManager.CONDA: + result = subprocess.run([self.manager_path, 'env', 'list'], + capture_output=True, text=True) + return env_name in result.stdout + else: + return os.path.exists(env_name) and os.path.isdir(env_name) + + def get_python_version(self, env_name): + """ + Get Python version of the specified environment + Args: + env_name (str): Name of the environment + Returns: + str: Python version + """ + try: + if self.manager == EnvManager.CONDA: + cmd = f'"{self.manager_path}" run -n {env_name} python --version' + else: + python_path = self._get_env_python_path(env_name) + cmd = f'"{python_path}" --version' + + result = subprocess.run(cmd, shell=True, capture_output=True, text=True, check=True) + return result.stdout.strip() + except subprocess.CalledProcessError as e: + raise RuntimeError(f"Failed to get Python version: {str(e)}") + + def remove_environment(self, env_name): + """ + Remove the specified environment + Args: + env_name (str): Name of the environment + """ + try: + if self.manager == EnvManager.CONDA: + subprocess.run([self.manager_path, 'env', 'remove', '-n', env_name, '-y'], + check=True) + elif self.manager == EnvManager.PYENV: + subprocess.run([self.manager_path, 'virtualenv-delete', env_name], check=True) + else: + if os.path.exists(env_name): + shutil.rmtree(env_name) + except (subprocess.CalledProcessError, OSError) as e: + raise RuntimeError(f"Failed to remove environment: {str(e)}") def process_ai_output(output, images, output_folder): if not PackageManager.check_package_installed("cv2"): @@ -673,131 +820,145 @@ def remove_text_from_string(string, text_to_find): return string -# Pytorch and cuda tools -def check_torch_version(min_version, min_cuda_versio=12): - import torch +import sys +import platform +import subprocess - if "+" in torch.__version__ and int(torch.__version__.split("+")[-1][2:4])= min_version - -def install_ninja(): - import conda.cli - try: - ASCIIColors.info("Installing ninja") # -c nvidia/label/cuda-12.3.2 -c nvidia -c conda-forge - result = conda.cli.main("install", "-c", "nvidia/label/cuda-12.3.2", "-c", "nvidia", "-c", "conda-forge", "ninja", "-y","--force-reinstall") - except Exception as ex: - ASCIIColors.error(ex) - -def install_cuda(): - import conda.cli - try: - ASCIIColors.info("Installing cuda 12.3.2") # -c nvidia/label/cuda-12.3.2 -c nvidia -c conda-forge - result = conda.cli.main("install", "-c", "nvidia/label/cuda-12.3.2", "-c", "nvidia", "-c", "conda-forge", "cuda-toolkit","-y","--force-reinstall") - except Exception as ex: - ASCIIColors.error(ex) - try: - ASCIIColors.info("Installing cuda compiler") # -c nvidia/label/cuda-12.3.2 -c nvidia -c conda-forge - result = conda.cli.main("install", "-c", "nvidia/label/cuda-12.3.2", "-c", "nvidia", "-c", "conda-forge", "cuda-compiler", "-y","--force-reinstall") - except Exception as ex: - ASCIIColors.error(ex) - -def install_cmake(): - import conda.cli - try: - ASCIIColors.info("Installing cmake") # -c nvidia/label/cuda-12.3.2 -c nvidia -c conda-forge - result = conda.cli.main("install", "-c", "nvidia/label/cuda-12.3.2", "-c", "nvidia", "-c", "conda-forge", "cmake","-y","--force-reinstall") - except Exception as ex: - ASCIIColors.error(ex) - def reinstall_pytorch_with_cuda(): + """ + Reinstall PyTorch with CUDA support using pip + Platform-aware: Windows and Linux will use CUDA, Mac will use default + """ try: - import conda.cli - ASCIIColors.info("Installing cuda 12.3.2") # -c nvidia/label/cuda-12.3.2 -c nvidia -c conda-forge - result = conda.cli.main("install", "-c", "nvidia/label/cuda-12.3.2", "-c", "nvidia", "-c", "conda-forge", "cuda-toolkit","-y","--force-reinstall") - except Exception as ex: - ASCIIColors.error(ex) - try: - ASCIIColors.info("Installing ninja") # -c nvidia/label/cuda-12.3.2 -c nvidia -c conda-forge - result = conda.cli.main("install", "-c", "nvidia/label/cuda-12.3.2", "-c", "nvidia", "-c", "conda-forge", "ninja", "-y","--force-reinstall") - except Exception as ex: - ASCIIColors.error(ex) - try: - ASCIIColors.info("Installing cuda compiler") # -c nvidia/label/cuda-12.3.2 -c nvidia -c conda-forge - result = conda.cli.main("install", "-c", "nvidia/label/cuda-12.3.2", "-c", "nvidia", "-c", "conda-forge", "cuda-compiler", "-y","--force-reinstall") - except Exception as ex: - ASCIIColors.error(ex) - try: - ASCIIColors.info("Installing pytorch 2.2.1") - result = subprocess.run([sys.executable, "-m", "pip", "install", "--upgrade", "torch", "torchvision", "torchaudio", "--no-cache-dir", "--index-url", "https://download.pytorch.org/whl/cu121"]) - except Exception as ex: - ASCIIColors.error(ex) - if result.returncode != 0: - ASCIIColors.warning("Couldn't find Cuda build tools on your PC. Reverting to CPU.") - result = subprocess.run([sys.executable, "-m", "pip", "install", "--upgrade", "torch", "torchvision", "torchaudio", "--no-cache-dir"]) - if result.returncode != 0: - ASCIIColors.error("Couldn't install pytorch !!") - else: - ASCIIColors.error("Pytorch installed successfully!!") - + system = platform.system().lower() + subprocess.check_call([sys.executable, "-m", "pip", "uninstall", "torch", "torchvision", "torchaudio", "-y"]) + + if system in ['windows', 'linux']: + subprocess.check_call([ + sys.executable, "-m", "pip", "install", + "torch", "torchvision", "torchaudio", + "--index-url", "https://download.pytorch.org/whl/cu121" + ]) + elif system == 'darwin': + print("Note: Installing default MacOS version as CUDA is not supported on MacOS") + subprocess.check_call([ + sys.executable, "-m", "pip", "install", + "torch", "torchvision", "torchaudio" + ]) + print("PyTorch reinstalled with CUDA support (where applicable)") + except subprocess.CalledProcessError as e: + print(f"Error reinstalling PyTorch: {e}") def reinstall_pytorch_with_rocm(): - result = subprocess.run([sys.executable, "-m", "pip", "install", "--upgrade", "torch", "torchvision", "torchaudio", "--no-cache-dir", "--index-url", "https://download.pytorch.org/whl/rocm5.6"]) - if result.returncode != 0: - ASCIIColors.warning("Couldn't find Cuda build tools on your PC. Reverting to CPU.") - result = subprocess.run([sys.executable, "-m", "pip", "install", "--upgrade", "torch", "torchvision", "torchaudio", "--no-cache-dir"]) - if result.returncode != 0: - ASCIIColors.error("Couldn't install pytorch !!") - else: - ASCIIColors.error("Pytorch installed successfully!!") - - + """ + Reinstall PyTorch with ROCm support using pip + """ + try: + subprocess.check_call([sys.executable, "-m", "pip", "uninstall", "torch", "torchvision", "torchaudio", "-y"]) + subprocess.check_call([ + sys.executable, "-m", "pip", "install", + "torch", "torchvision", "torchaudio", + "--index-url", "https://download.pytorch.org/whl/rocm5.6" + ]) + print("PyTorch reinstalled with ROCm support") + except subprocess.CalledProcessError as e: + print(f"Error reinstalling PyTorch: {e}") def reinstall_pytorch_with_cpu(): - result = subprocess.run([sys.executable, "-m", "pip", "install", "--upgrade", "torch", "torchvision", "torchaudio", "--no-cache-dir"]) - if result.returncode != 0: - ASCIIColors.warning("Couldn't find Cuda build tools on your PC. Reverting to CPU.") - result = subprocess.run([sys.executable, "-m", "pip", "install", "--upgrade", "torch", "torchvision", "torchaudio", "--no-cache-dir"]) - if result.returncode != 0: - ASCIIColors.error("Couldn't install pytorch !!") - else: - ASCIIColors.error("Pytorch installed successfully!!") + """ + Reinstall PyTorch CPU-only version using pip + """ + try: + subprocess.check_call([sys.executable, "-m", "pip", "uninstall", "torch", "torchvision", "torchaudio", "-y"]) + subprocess.check_call([ + sys.executable, "-m", "pip", "install", + "torch", "torchvision", "torchaudio" + ]) + print("PyTorch reinstalled (CPU-only version)") + except subprocess.CalledProcessError as e: + print(f"Error reinstalling PyTorch: {e}") -def check_and_install_torch(enable_gpu:bool, version:float=2.2): - if enable_gpu: - ASCIIColors.yellow("This installation has enabled GPU support. Trying to install with GPU support") - ASCIIColors.info("Checking pytorch") - try: - import torch - import torchvision - if torch.cuda.is_available(): - ASCIIColors.success(f"CUDA is supported.\nCurrent version is {torch.__version__}.") - if not check_torch_version(version): - ASCIIColors.yellow("Torch version is old. Installing new version") - reinstall_pytorch_with_cuda() + +def check_and_install_torch(enable_gpu: bool, version: float = 2.2): + """ + Check and install PyTorch with specified configuration + Args: + enable_gpu (bool): Whether to install GPU version + version (float): Minimum required PyTorch version + """ + try: + import torch + current_version = torch.__version__ + system = platform.system().lower() + + # Check if current installation meets requirements + if pkg_resources.parse_version(current_version) >= pkg_resources.parse_version(str(version)): + if enable_gpu: + if system == 'darwin': + # For Mac, MPS is the GPU solution + print(f"Current PyTorch installation ({current_version}) is compatible with Mac GPU (MPS)") + return True + elif torch.cuda.is_available(): + print(f"Current PyTorch installation ({current_version}) has CUDA support") + return True else: - ASCIIColors.yellow("Torch OK") + print("GPU version requested but CUDA not available. Reinstalling...") else: - ASCIIColors.warning("CUDA is not supported. Trying to reinstall PyTorch with CUDA support.") - reinstall_pytorch_with_cuda() - except Exception as ex: - ASCIIColors.info("Pytorch not installed. Reinstalling ...") - reinstall_pytorch_with_cuda() + print(f"Current CPU PyTorch installation ({current_version}) meets version requirement") + return True + except ImportError: + print("PyTorch not found. Installing...") + + # Perform installation based on requirements + if enable_gpu: + if system == 'darwin': + reinstall_pytorch_with_cpu() # Mac uses default installation for MPS + else: + reinstall_pytorch_with_cuda() else: - try: - import torch - import torchvision - if check_torch_version(version): - ASCIIColors.warning("Torch version is too old. Trying to reinstall PyTorch with CUDA support.") - reinstall_pytorch_with_cpu() - except Exception as ex: - ASCIIColors.info("Pytorch not installed. Reinstalling ...") - reinstall_pytorch_with_cpu() + reinstall_pytorch_with_cpu() + + # Verify installation + try: + import torch + print(f"PyTorch {torch.__version__} installed successfully") + if enable_gpu: + if system == 'darwin': + print("MPS (Mac GPU) support available if hardware supports it") + elif torch.cuda.is_available(): + print(f"CUDA version: {torch.version.cuda}") + else: + print("Warning: GPU version requested but CUDA not available") + except ImportError: + print("Installation failed") + return False + + return True class NumpyEncoderDecoder(json.JSONEncoder):