mirror of https://github.com/bmaltais/kohya_ss
266 lines
9.5 KiB
Python
266 lines
9.5 KiB
Python
import subprocess
|
|
import os
|
|
import filecmp
|
|
import logging
|
|
import shutil
|
|
import sysconfig
|
|
import setup_common
|
|
import argparse
|
|
|
|
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 cudnn_install():
|
|
log.info("Installing nvidia-cudnn-cu11 8.9.5.29...")
|
|
setup_common.install(
|
|
"--upgrade nvidia-cudnn-cu11==8.9.5.29",
|
|
"nvidia-cudnn-cu11 8.9.5.29",
|
|
reinstall=True,
|
|
)
|
|
|
|
# Original path with "..\\venv"
|
|
original_path = os.path.join(
|
|
os.path.dirname(os.path.realpath(__file__)),
|
|
"..\\venv\\Lib\\site-packages\\nvidia\\cudnn\\bin",
|
|
)
|
|
# Normalize the path to resolve "..\\venv"
|
|
cudnn_src = os.path.abspath(original_path)
|
|
cudnn_dest = os.path.join(sysconfig.get_paths()["purelib"], "torch", "lib")
|
|
|
|
log.info(f"Copying CUDNN files from {cudnn_src} to {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):
|
|
if file.lower().endswith(".dll"): # Check if the file is a .dll file
|
|
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 .dll 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_torch2(headless: bool = False):
|
|
setup_common.check_repo_version()
|
|
if not setup_common.check_python_version():
|
|
exit(1)
|
|
|
|
setup_common.update_submodule()
|
|
|
|
setup_common.install("pip")
|
|
|
|
setup_common.install_requirements(
|
|
"requirements_windows_torch2.txt", check_no_verify_flag=False
|
|
)
|
|
|
|
setup_common.configure_accelerate(
|
|
run_accelerate=not headless
|
|
) # False if headless is True and vice versa
|
|
|
|
|
|
def install_bitsandbytes_0_35_0():
|
|
log.info("Installing bitsandbytes 0.35.0...")
|
|
setup_common.install(
|
|
"--upgrade bitsandbytes==0.35.0", "bitsandbytes 0.35.0", reinstall=True
|
|
)
|
|
sync_bits_and_bytes_files()
|
|
|
|
|
|
def install_bitsandbytes_0_40_1():
|
|
log.info("Installing bitsandbytes 0.40.1...")
|
|
setup_common.install(
|
|
"--upgrade https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.40.1.post1-py3-none-win_amd64.whl",
|
|
"bitsandbytes 0.40.1",
|
|
reinstall=True,
|
|
)
|
|
|
|
|
|
def install_bitsandbytes_0_41_1():
|
|
log.info("Installing bitsandbytes 0.41.1...")
|
|
setup_common.install(
|
|
"--upgrade https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.41.1-py3-none-win_amd64.whl",
|
|
"bitsandbytes 0.41.1",
|
|
reinstall=True,
|
|
)
|
|
|
|
|
|
def install_bitsandbytes_0_41_2():
|
|
log.info("Installing bitsandbytes 0.41.2...")
|
|
setup_common.install(
|
|
"--upgrade https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.41.2.post2-py3-none-win_amd64.whl",
|
|
"bitsandbytes 0.41.2",
|
|
reinstall=True,
|
|
)
|
|
|
|
|
|
def install_triton_2_1_0():
|
|
log.info("Installing triton 2.1.0...")
|
|
setup_common.install(
|
|
"--upgrade https://huggingface.co/Rodeszones/CogVLM-grounding-generalist-hf-quant4/resolve/main/triton-2.1.0-cp310-cp310-win_amd64.whl?download=true",
|
|
"triton 2.1.0",
|
|
reinstall=True,
|
|
)
|
|
|
|
|
|
def main_menu(headless: bool = False):
|
|
if headless:
|
|
install_kohya_ss_torch2(headless=headless)
|
|
else:
|
|
setup_common.clear_screen()
|
|
while True:
|
|
print("\nKohya_ss setup menu:\n")
|
|
print("1. Install kohya_ss GUI")
|
|
print(
|
|
"2. (Optional) Install CuDNN files (to use the latest supported CuDNN version)"
|
|
)
|
|
print("3. (Optional) Install Triton 2.1.0 for Windows")
|
|
print("4. (Optional) Install specific version of bitsandbytes")
|
|
print("5. (Optional) Manually configure Accelerate")
|
|
print("6. (Optional) Launch Kohya_ss GUI in browser")
|
|
print("7. Exit Setup")
|
|
|
|
choice = input("\nSelect an option: ")
|
|
print("")
|
|
|
|
if choice == "1":
|
|
install_kohya_ss_torch2()
|
|
elif choice == "2":
|
|
cudnn_install()
|
|
elif choice == "3":
|
|
install_triton_2_1_0()
|
|
elif choice == "4":
|
|
while True:
|
|
print("\nBitsandBytes Installation Menu:")
|
|
print("1. Force install Bitsandbytes 0.35.0")
|
|
print(
|
|
"2. Force install Bitsandbytes 0.40.1 (supports new optimizer options, pre-bugfix results)"
|
|
)
|
|
print(
|
|
"3. Force installation Bitsandbytes 0.41.1 (supports new optimizer options)"
|
|
)
|
|
print(
|
|
"4. (Recommended) Force install Bitsandbytes 0.41.2 (supports new optimizer options)"
|
|
)
|
|
print(
|
|
"5. (Warning) Install bitsandbytes-windows (may cause issues, use with caution)"
|
|
)
|
|
print("6. Return to Previous Menu:")
|
|
choice_torch = input("\nSelect an option: ")
|
|
print("")
|
|
|
|
if choice_torch == "1":
|
|
install_bitsandbytes_0_35_0()
|
|
break
|
|
elif choice_torch == "2":
|
|
install_bitsandbytes_0_40_1()
|
|
break
|
|
elif choice_torch == "3":
|
|
install_bitsandbytes_0_41_1()
|
|
break
|
|
elif choice_torch == "4":
|
|
install_bitsandbytes_0_41_2()
|
|
break
|
|
elif choice_torch == "5":
|
|
setup_common.install(
|
|
"--upgrade bitsandbytes-windows", reinstall=True
|
|
)
|
|
break
|
|
elif choice_torch == "6":
|
|
break
|
|
else:
|
|
print("Invalid choice. Please chose an option between 1-6.")
|
|
elif choice == "5":
|
|
setup_common.run_cmd("accelerate config")
|
|
elif choice == "6":
|
|
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 == "7":
|
|
print("Exiting setup.")
|
|
break
|
|
else:
|
|
print("Invalid selection. Please choose an option between 1-7.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
setup_common.ensure_base_requirements()
|
|
setup_common.setup_logging()
|
|
|
|
# Setup argument parser
|
|
parser = argparse.ArgumentParser(description="Your Script Description")
|
|
parser.add_argument("--headless", action="store_true", help="Run in headless mode")
|
|
|
|
# Parse arguments
|
|
args = parser.parse_args()
|
|
|
|
main_menu(headless=args.headless)
|