import subprocess
import os
import filecmp
import logging
import shutil
import sysconfig
import setup_common
errors = 0 # Define the 'errors' variable before using it
log = logging.getLogger('sd')
# ANSI escape code for yellow color
YELLOW = '\033[93m'
RESET_COLOR = '\033[0m'
def cudann_install():
cudnn_src = os.path.join(
os.path.dirname(os.path.realpath(__file__)), '..\cudnn_windows'
)
cudnn_dest = os.path.join(sysconfig.get_paths()['purelib'], 'torch', 'lib')
log.info(f'Checking for CUDNN files in {cudnn_dest}...')
if os.path.exists(cudnn_src):
if os.path.exists(cudnn_dest):
# check for different files
filecmp.clear_cache()
for file in os.listdir(cudnn_src):
src_file = os.path.join(cudnn_src, file)
dest_file = os.path.join(cudnn_dest, file)
# if dest file exists, check if it's different
if os.path.exists(dest_file):
if not filecmp.cmp(src_file, dest_file, shallow=False):
shutil.copy2(src_file, cudnn_dest)
else:
shutil.copy2(src_file, cudnn_dest)
log.info('Copied CUDNN 8.6 files to destination')
else:
log.warning(f'Destination directory {cudnn_dest} does not exist')
else:
log.error(f'Installation Failed: "{cudnn_src}" could not be found.')
def sync_bits_and_bytes_files():
import filecmp
"""
Check for "different" bitsandbytes Files and copy only if necessary.
This function is specific for Windows OS.
"""
# Only execute on Windows
if os.name != 'nt':
print('This function is only applicable to Windows OS.')
return
try:
log.info(f'Copying bitsandbytes files...')
# Define source and destination directories
source_dir = os.path.join(os.getcwd(), 'bitsandbytes_windows')
dest_dir_base = os.path.join(
sysconfig.get_paths()['purelib'], 'bitsandbytes'
)
# Clear file comparison cache
filecmp.clear_cache()
# Iterate over each file in source directory
for file in os.listdir(source_dir):
source_file_path = os.path.join(source_dir, file)
# Decide the destination directory based on file name
if file in ('main.py', 'paths.py'):
dest_dir = os.path.join(dest_dir_base, 'cuda_setup')
else:
dest_dir = dest_dir_base
dest_file_path = os.path.join(dest_dir, file)
# Compare the source file with the destination file
if os.path.exists(dest_file_path) and filecmp.cmp(
source_file_path, dest_file_path
):
log.debug(
f'Skipping {source_file_path} as it already exists in {dest_dir}'
)
else:
# Copy file from source to destination, maintaining original file's metadata
log.debug(f'Copy {source_file_path} to {dest_dir}')
shutil.copy2(source_file_path, dest_dir)
except FileNotFoundError as fnf_error:
log.error(f'File not found error: {fnf_error}')
except PermissionError as perm_error:
log.error(f'Permission error: {perm_error}')
except Exception as e:
log.error(f'An unexpected error occurred: {e}')
def install_kohya_ss_torch1():
setup_common.check_repo_version()
setup_common.check_python()
# Upgrade pip if needed
setup_common.install('--upgrade pip')
if setup_common.check_torch() == 2:
input(
f'{YELLOW}\nTorch 2 is already installed in the venv. To install Torch 1 delete the venv and re-run setup.bat\n\nHit enter to continue...{RESET_COLOR}'
)
return
# setup_common.install(
# 'torch==1.12.1+cu116 torchvision==0.13.1+cu116 --index-url https://download.pytorch.org/whl/cu116',
# 'torch torchvision'
# )
# setup_common.install(
# 'https://github.com/C43H66N12O12S2/stable-diffusion-webui/releases/download/f/xformers-0.0.14.dev0-cp310-cp310-win_amd64.whl -U -I --no-deps',
# 'xformers-0.0.14'
# )
setup_common.install_requirements('requirements_windows_torch1.txt', check_no_verify_flag=False)
sync_bits_and_bytes_files()
setup_common.configure_accelerate(run_accelerate=True)
# run_cmd(f'accelerate config')
def install_kohya_ss_torch2():
setup_common.check_repo_version()
setup_common.check_python()
# Upgrade pip if needed
setup_common.install('--upgrade pip')
if setup_common.check_torch() == 1:
input(
f'{YELLOW}\nTorch 1 is already installed in the venv. To install Torch 2 delete the venv and re-run setup.bat\n\nHit any key to acknowledge.{RESET_COLOR}'
)
return
# setup_common.install(
# 'torch==2.0.1+cu118 torchvision==0.15.2+cu118 --index-url https://download.pytorch.org/whl/cu118',
# 'torch torchvision'
# )
setup_common.install_requirements('requirements_windows_torch2.txt', check_no_verify_flag=False)
# install('https://huggingface.co/r4ziel/xformers_pre_built/resolve/main/triton-2.0.0-cp310-cp310-win_amd64.whl', 'triton', reinstall=reinstall)
sync_bits_and_bytes_files()
setup_common.configure_accelerate(run_accelerate=True)
# run_cmd(f'accelerate config')
def main_menu():
setup_common.clear_screen()
while True:
print('\nKohya_ss GUI setup menu:\n')
print('1. Install kohya_ss gui')
print('2. (Optional) Install cudann files')
print('3. (Optional) Install bitsandbytes-windows')
print('4. (Optional) Manually configure accelerate')
print('5. (Optional) Start Kohya_ss GUI in browser')
print('6. Quit')
choice = input('\nEnter your choice: ')
print('')
if choice == '1':
while True:
print('1. Torch 1 (legacy)')
print('2. Torch 2 (recommended)')
print('3. Cancel')
choice_torch = input('\nEnter your choice: ')
print('')
if choice_torch == '1':
install_kohya_ss_torch1()
break
elif choice_torch == '2':
install_kohya_ss_torch2()
break
elif choice_torch == '3':
break
else:
print('Invalid choice. Please enter a number between 1-3.')
elif choice == '2':
cudann_install()
elif choice == '3':
setup_common.install('--upgrade bitsandbytes-windows', reinstall=True)
elif choice == '4':
setup_common.run_cmd('accelerate config')
elif choice == '5':
subprocess.Popen('start cmd /k .\gui.bat --inbrowser', shell=True) # /k keep the terminal open on quit. /c would close the terminal instead
elif choice == '6':
print('Quitting the program.')
break
else:
print('Invalid choice. Please enter a number between 1-5.')
if __name__ == '__main__':
setup_common.ensure_base_requirements()
setup_common.setup_logging()
main_menu()