133 lines
4.4 KiB
Python
133 lines
4.4 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import runpy
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
PLUGIN_DIR = Path(__file__).resolve().parents[1]
|
|
BACKEND_DIR = Path(os.environ.get("ANIMAFORGE_BACKEND_DIR", PLUGIN_DIR / "backend" / "sd-scripts")).expanduser().resolve()
|
|
TRAIN_PYTHON = Path(
|
|
os.environ.get(
|
|
"ANIMAFORGE_TRAIN_PYTHON",
|
|
BACKEND_DIR / ".venv" / ("Scripts/python.exe" if os.name == "nt" else "bin/python"),
|
|
)
|
|
).expanduser()
|
|
TRAIN_VENV = TRAIN_PYTHON.parent.parent.resolve()
|
|
|
|
|
|
def _value_after(args: list[str], name: str) -> str | None:
|
|
if name not in args:
|
|
return None
|
|
index = args.index(name)
|
|
if index + 1 >= len(args):
|
|
return None
|
|
return args[index + 1]
|
|
|
|
|
|
def _pop_arg(args: list[str], name: str) -> str | None:
|
|
if name not in args:
|
|
return None
|
|
index = args.index(name)
|
|
if index + 1 >= len(args):
|
|
return None
|
|
value = args[index + 1]
|
|
del args[index:index + 2]
|
|
return value
|
|
|
|
|
|
def _pop_optional_int(args: list[str], name: str, default: int) -> int:
|
|
value = _pop_arg(args, name)
|
|
if value is None:
|
|
return default
|
|
try:
|
|
return max(1, int(value))
|
|
except Exception:
|
|
return default
|
|
|
|
|
|
def _ensure_backend_ready() -> None:
|
|
if not BACKEND_DIR.exists():
|
|
raise RuntimeError(
|
|
f"sd-scripts backend was not found: {BACKEND_DIR}\n"
|
|
"Run scripts/setup_sd_scripts_backend.ps1 on Windows before training."
|
|
)
|
|
if not (BACKEND_DIR / "anima_train_network.py").exists():
|
|
raise RuntimeError(f"anima_train_network.py was not found under backend: {BACKEND_DIR}")
|
|
if not TRAIN_PYTHON.exists():
|
|
raise RuntimeError(
|
|
f"training Python was not found: {TRAIN_PYTHON}\n"
|
|
"Run scripts/setup_sd_scripts_backend.ps1 to create the backend virtual environment."
|
|
)
|
|
|
|
|
|
def _ensure_dataset_config(args: list[str]) -> list[str]:
|
|
if "--dataset_config" in args:
|
|
return args
|
|
|
|
num_repeats = _pop_optional_int(args, "--anima_num_repeats", 1)
|
|
train_data_dir = _pop_arg(args, "--train_data_dir")
|
|
if not train_data_dir:
|
|
return args
|
|
|
|
image_dir = Path(train_data_dir).expanduser().resolve()
|
|
output_dir = Path(_value_after(args, "--output_dir") or image_dir / "anima_lora_output").expanduser().resolve()
|
|
output_name = _value_after(args, "--output_name") or image_dir.name
|
|
resolution = _value_after(args, "--resolution") or "1024,1024"
|
|
batch_size = _value_after(args, "--train_batch_size") or "1"
|
|
caption_extension = _value_after(args, "--caption_extension") or ".txt"
|
|
caption_separator = _value_after(args, "--caption_separator") or ","
|
|
keep_tokens = _value_after(args, "--keep_tokens") or "1"
|
|
|
|
try:
|
|
width, height = [item.strip() for item in resolution.split(",", 1)]
|
|
resolution_toml = f"[{int(width)}, {int(height)}]"
|
|
except Exception:
|
|
resolution_toml = "1024"
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
config_path = output_dir / f"{output_name}.dataset.toml"
|
|
config_path.write_text(
|
|
"\n".join(
|
|
[
|
|
"[general]",
|
|
"shuffle_caption = true" if "--shuffle_caption" in args else "shuffle_caption = false",
|
|
f"caption_extension = {caption_extension!r}",
|
|
f"caption_separator = {caption_separator!r}",
|
|
f"keep_tokens = {int(keep_tokens)}",
|
|
"",
|
|
"[[datasets]]",
|
|
f"resolution = {resolution_toml}",
|
|
f"batch_size = {int(batch_size)}",
|
|
"enable_bucket = true",
|
|
"bucket_no_upscale = false",
|
|
"",
|
|
" [[datasets.subsets]]",
|
|
f" image_dir = {str(image_dir)!r}",
|
|
f" num_repeats = {num_repeats}",
|
|
f" class_tokens = {image_dir.name!r}",
|
|
"",
|
|
]
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
args.extend(["--dataset_config", str(config_path)])
|
|
return args
|
|
|
|
|
|
def main() -> None:
|
|
_ensure_backend_ready()
|
|
if os.path.realpath(sys.prefix) != os.path.realpath(TRAIN_VENV):
|
|
os.execv(str(TRAIN_PYTHON), [str(TRAIN_PYTHON), __file__, *sys.argv[1:]])
|
|
|
|
sys.argv[1:] = _ensure_dataset_config(list(sys.argv[1:]))
|
|
os.chdir(BACKEND_DIR)
|
|
sys.path.insert(0, str(BACKEND_DIR))
|
|
runpy.run_path(str(BACKEND_DIR / "anima_train_network.py"), run_name="__main__")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|