setup/setup_windows.py

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()