1580158
/
kohya_ss
mirror of https://github.com/bmaltais/kohya_ss
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Feature/update gui uv for colab parity (#3350) 

* feat: Update gui-uv.sh for Colab parity and add local scripts

This commit introduces changes to align the local setup via `gui-uv.sh`
more closely with the procedures in the reference Colab notebook,
particularly concerning `bitsandbytes` compilation.

Changes:
- Added `kohya_ss_colab.ipynb` to the repository for reference.
- Created `local_train.py`, a Python script adapted from the Colab
  notebook, providing an alternative setup and launch method.
- Modified `gui-uv.sh` to include a documented, optional section
  for users to build `bitsandbytes` from source if the `uv`-installed
  version is insufficient. This enhances flexibility and aids in
  troubleshooting CUDA/GPU-specific issues.

The primary goal is to make the local experience with `gui-uv.sh`
more robust and provide users with options similar to the Colab
environment's setup for critical dependencies like `bitsandbytes`.

* Update collab

* Update colab script

---------

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
pull/3354/head
bmaltais 2025-07-14 20:25:28 -04:00 committed by GitHub
parent 307e197a28
commit f0d3e3f5b3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 220 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
gui-uv.sh
View File

@ -52,4 +52,40 @@ if [[ "$uv_quiet" == "--quiet" ]]; then
fi
git submodule update --init --recursive
# --- bitsandbytes compilation (optional, similar to Colab setup) ---
# The pyproject.toml specifies bitsandbytes, which uv will install.
# However, if you encounter issues with the pre-compiled version,
# you might need to compile it from source, similar to the Colab notebook.
# Uncomment and adapt the following lines if needed.
# Ensure you have the necessary build tools (like CUDA toolkit for cuda11x).
# BUILD_BITSANDBYTES_FROM_SOURCE=false # Set to true to enable
# if [ "$BUILD_BITSANDBYTES_FROM_SOURCE" = true ]; then
# echo "Attempting to build bitsandbytes from source..."
# if [ ! -d "bitsandbytes" ]; then
# git clone -b 0.41.0 https://github.com/TimDettmers/bitsandbytes
# fi
# cd bitsandbytes || exit 1
# # IMPORTANT: Adjust CUDA_VERSION if necessary for your setup (e.g., 118 for CUDA 11.8, 12x for CUDA 12.x)
# # The Colab notebook used CUDA_VERSION=118
# # Common options: cuda11x (for 11.0-11.8), cuda12x (for 12.0-12.x)
# # Ensure the corresponding CUDA toolkit is installed and in PATH.
# # For ROCm, the build process is different.
# # Check bitsandbytes documentation for the correct make target for your GPU/driver.
# # Example for CUDA 11.8:
# # CUDA_VERSION=118 make cuda11x
# # Example for CUDA 12.1:
# # make cuda12x
# # Then, install using uv within the environment (or pip if uv is not active yet for this part)
# # Assuming uv environment is active or will be activated by the main command:
# uv pip install .
# # Or, if building before activating the main uv environment:
# # python setup.py install
# cd .. || exit 1
# echo "bitsandbytes build attempt finished."
# fi
# --- end of bitsandbytes compilation ---
echo "Launching Kohya GUI via uv..."
uv run $uv_quiet kohya_gui.py --noverify "${args[@]}"

103
kohya_ss_colab.ipynb Normal file
View File

@ -0,0 +1,103 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github"
},
"source": [
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/bmaltais/kohya_ss/blob/feature/update-gui-uv-for-colab-parity/kohya_ss_colab.ipynb)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#@title Train with Kohya's Stable Diffusion Trainers\n",
"%cd /content\n",
"\n",
"import os\n",
"from google.colab import drive\n",
"drive.mount('/content/drive')\n",
"\n",
"%cd /content\n",
"!git clone -b dev https://github.com/bmaltais/kohya_ss\n",
"%cd /content/kohya_ss\n",
"!git checkout dev\n",
"\n",
"os.environ[\"TERM\"] = \"dumb\"\n",
"!./gui-uv.sh --quiet --share --headless"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#@title Convert Safetensors to Diffusers\n",
"from_safetensors_url = '' #@param {type:\"string\"}\n",
"!wget -q https://raw.githubusercontent.com/huggingface/diffusers/v0.17.1/scripts/convert_original_stable_diffusion_to_diffusers.py\n",
"!wget {from_safetensors_url} -O /content/model.safetensors\n",
"!python3 convert_original_stable_diffusion_to_diffusers.py --half --from_safetensors --checkpoint_path model.safetensors --dump_path /content/model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#@title Push to HF.co\n",
"\n",
"import os\n",
"from huggingface_hub import create_repo, upload_folder\n",
"\n",
"hf_token = 'HF_WRITE_TOKEN' #@param {type:\"string\"}\n",
"repo_id = 'username/reponame' #@param {type:\"string\"}\n",
"commit_message = '\\u2764' #@param {type:\"string\"}\n",
"create_repo(repo_id, private=True, token=hf_token)\n",
"model_path = '/content/train/model' #@param {type:\"string\"}\n",
"upload_folder(folder_path=f'{model_path}', path_in_repo='', repo_id=repo_id, commit_message=commit_message, token=hf_token)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#@title Push to DagsHub.com\n",
"\n",
"!pip -q install dagshub\n",
"from dagshub.upload import Repo, create_repo\n",
"\n",
"repo_id = 'reponame' #@param {type:\"string\"}\n",
"org_name = 'orgname' #@param {type:\"string\"}\n",
"commit_message = '\\u2764' #@param {type:\"string\"}\n",
"create_repo(f\"{repo_id}\", org_name=f\"{org_name}\")\n",
"repo = Repo(f\"{org_name}\", f\"{repo_id}\")\n",
"model_path = '/content/train/model' #@param {type:\"string\"}\n",
"repo.upload(\"/content/models\", remote_path=\"data\", commit_message=f\"{commit_message}\", force=True)"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuType": "T4",
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
},
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 0
}

81
local_train.py Normal file
View File

@ -0,0 +1,81 @@
import os
import subprocess
def run_command(command):
"""Runs a shell command and prints its output."""
print(f"Running command: {command}")
process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
for line in process.stdout:
print(line, end="")
process.wait()
if process.returncode != 0:
raise Exception(f"Command failed with return code {process.returncode}: {command}")
def main():
print("Starting local training setup...")
# Target directory for cloning and operations, can be made configurable
base_dir = "kohya_local_setup"
os.makedirs(base_dir, exist_ok=True)
original_dir = os.getcwd()
os.chdir(base_dir)
print(f"Working directory: {os.getcwd()}")
# 1. Install dependencies (This should ideally be handled by gui-uv.sh using requirements.txt)
# For now, we'll list them. These were in the notebook:
# dadaptation==3.1 diffusers[torch]==0.17.1 easygui==0.98.3 einops==0.6.0
# fairscale==0.4.13 ftfy==6.1.1 gradio==3.36.1 huggingface-hub==0.14.1
# lion-pytorch==0.0.6 lycoris_lora==1.8.0.dev6 open-clip-torch==2.20.0
# prodigyopt==1.0 pytorch-lightning==1.9.0 safetensors==0.3.1 timm==0.6.12
# tk==0.1.0 transformers==4.30.2 voluptuous==0.13.1 wandb==0.15.0
# xformers==0.0.20 omegaconf
print("Ensure all dependencies are installed via requirements.txt by gui-uv.sh")
# 2. Clone bitsandbytes
if not os.path.exists("bitsandbytes"):
run_command("git clone -b 0.41.0 https://github.com/TimDettmers/bitsandbytes")
else:
print("bitsandbytes directory already exists. Skipping clone.")
# 3. Build and install bitsandbytes
# Note: This might need specific CUDA versions. The notebook used CUDA_VERSION=118
# This part is tricky for a generic local setup and might need user intervention or
# pre-compiled versions if not handled by the main gui-uv.sh environment.
# For now, we are replicating the notebook's steps.
# It's better if bitsandbytes is installed as a pip package if possible.
print("Building and installing bitsandbytes...")
os.chdir("bitsandbytes")
run_command("make cuda11x") # This assumes CUDA 11.8 development toolkit is available
run_command("python setup.py install")
os.chdir("..") # Back to base_dir
# 4. Clone kohya_ss
if not os.path.exists("kohya_ss"):
run_command("git clone -b v1.0 https://github.com/camenduru/kohya_ss")
else:
print("kohya_ss directory already exists. Skipping clone.")
# Optionally, add a git pull here to update
# os.chdir("kohya_ss")
# run_command("git pull")
# os.chdir("..")
# 5. Launch Kohya GUI
print("Launching Kohya GUI...")
os.chdir("kohya_ss")
# The notebook uses --share --headless.
# --share might not be needed or desired for local use.
# --headless might be for Colab's environment. For local GUI, we might not want headless.
# We need to check how gui-uv.sh expects to launch this.
# For now, using the notebook's command but this will likely need adjustment.
try:
run_command("python kohya_gui.py --headless")
except Exception as e:
print(f"Error launching Kohya GUI: {e}")
print("Please ensure all dependencies, including CUDA and xformers, are correctly installed.")
finally:
os.chdir(original_dir) # Change back to the original directory
if __name__ == "__main__":
main()