501 lines
19 KiB
Python
501 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import math
|
|
import re
|
|
import shlex
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Sequence, Tuple
|
|
|
|
from .state_store import append_job_log, get_job, now_iso, update_job
|
|
|
|
try:
|
|
from comfy.utils import ProgressBar
|
|
except Exception: # pragma: no cover - available inside ComfyUI
|
|
ProgressBar = None
|
|
|
|
PLUGIN_DIR = Path(__file__).resolve().parent
|
|
LEGACY_WORKSPACE_PREFIX = "/workspace"
|
|
|
|
|
|
def _discover_comfyui_dir() -> Path:
|
|
env_value = os.environ.get("COMFYUI_DIR")
|
|
if env_value:
|
|
return Path(env_value).expanduser().resolve()
|
|
for parent in [PLUGIN_DIR, *PLUGIN_DIR.parents]:
|
|
if parent.name.lower() == "comfyui":
|
|
return parent
|
|
if (parent / "main.py").exists() and (parent / "custom_nodes").exists():
|
|
return parent
|
|
if PLUGIN_DIR.parent.name == "custom_nodes":
|
|
return PLUGIN_DIR.parent.parent
|
|
return PLUGIN_DIR.parent / "ComfyUI"
|
|
|
|
|
|
COMFYUI_DIR = _discover_comfyui_dir()
|
|
WORKSPACE_DIR = COMFYUI_DIR.parent
|
|
BACKEND_DIR = Path(os.environ.get("ANIMAFORGE_BACKEND_DIR", PLUGIN_DIR / "backend" / "sd-scripts")).expanduser().resolve()
|
|
DEFAULT_TRAIN_SCRIPT_PATH = str(PLUGIN_DIR / "scripts" / "anima_train_network_wrapper.py")
|
|
DEFAULT_BACKEND_PYTHON = str(BACKEND_DIR / ".venv" / ("Scripts/python.exe" if os.name == "nt" else "bin/python"))
|
|
DEFAULT_ANIMA_LORA_OUTPUT_DIR = str(COMFYUI_DIR / "models" / "loras" / "anima_lora")
|
|
|
|
|
|
def expand_path(value: str) -> str:
|
|
value = os.path.expandvars(str(value or "").strip())
|
|
if not value:
|
|
return ""
|
|
if value == LEGACY_WORKSPACE_PREFIX:
|
|
value = str(WORKSPACE_DIR)
|
|
elif value.startswith(f"{LEGACY_WORKSPACE_PREFIX}/"):
|
|
value = str(WORKSPACE_DIR / value.removeprefix(f"{LEGACY_WORKSPACE_PREFIX}/"))
|
|
elif value.replace("\\", "/").startswith("ComfyUI/"):
|
|
value = str(COMFYUI_DIR / value.replace("\\", "/").removeprefix("ComfyUI/"))
|
|
return str(Path(os.path.expanduser(value)).resolve())
|
|
|
|
|
|
def normalize_arg_path(value: str) -> str:
|
|
value = os.path.expandvars(str(value))
|
|
if value == LEGACY_WORKSPACE_PREFIX:
|
|
return str(WORKSPACE_DIR)
|
|
if value.startswith(f"{LEGACY_WORKSPACE_PREFIX}/"):
|
|
return str(WORKSPACE_DIR / value.removeprefix(f"{LEGACY_WORKSPACE_PREFIX}/"))
|
|
if value.replace("\\", "/").startswith("ComfyUI/"):
|
|
return str(COMFYUI_DIR / value.replace("\\", "/").removeprefix("ComfyUI/"))
|
|
return value
|
|
|
|
|
|
def resolve_python_executable(value: str = "") -> str:
|
|
if value.strip():
|
|
value = normalize_arg_path(str(value).strip())
|
|
return str(Path(os.path.expanduser(value)).absolute())
|
|
return DEFAULT_BACKEND_PYTHON if Path(DEFAULT_BACKEND_PYTHON).exists() else sys.executable
|
|
|
|
|
|
def resolve_script_path(script_path: str) -> Tuple[str, str]:
|
|
script_path = (script_path or "").strip()
|
|
if not script_path:
|
|
script_path = DEFAULT_TRAIN_SCRIPT_PATH
|
|
|
|
expanded = Path(expand_path(script_path))
|
|
if expanded.exists():
|
|
return str(expanded.resolve()), ""
|
|
|
|
found = shutil.which(script_path)
|
|
if found:
|
|
return found, ""
|
|
|
|
return str(expanded), f"training script not found: {expanded}"
|
|
|
|
|
|
def split_extra_args(extra_args: str) -> List[str]:
|
|
extra_args = (extra_args or "").strip()
|
|
if not extra_args:
|
|
return []
|
|
return shlex.split(extra_args)
|
|
|
|
PRESET_DEFAULTS = {
|
|
"smoke_test": {
|
|
"target_steps_per_lora": 500,
|
|
"learning_rate": "2e-5",
|
|
"network_dim": 16,
|
|
"network_alpha": 16,
|
|
"resolution": "768,768",
|
|
},
|
|
"standard": {
|
|
"target_steps_per_lora": 2000,
|
|
"learning_rate": "2e-5",
|
|
"network_dim": 32,
|
|
"network_alpha": 32,
|
|
"resolution": "1024,1024",
|
|
},
|
|
"strong": {
|
|
"target_steps_per_lora": 3000,
|
|
"learning_rate": "3e-5",
|
|
"network_dim": 32,
|
|
"network_alpha": 32,
|
|
"resolution": "1024,1024",
|
|
},
|
|
"style": {
|
|
"target_steps_per_lora": 4000,
|
|
"learning_rate": "1e-5",
|
|
"network_dim": 32,
|
|
"network_alpha": 32,
|
|
"resolution": "768,768",
|
|
},
|
|
}
|
|
|
|
|
|
def normalized_config(config: Dict[str, Any]) -> Dict[str, Any]:
|
|
normalized = dict(config)
|
|
preset = str(normalized.get("training_preset") or "standard")
|
|
if preset not in PRESET_DEFAULTS:
|
|
preset = "standard"
|
|
normalized["training_preset"] = preset
|
|
|
|
# Presets are defaults, not a separate training path. Users can still override
|
|
# the concrete fields in the node; the workflow's default values match standard.
|
|
preset_defaults = PRESET_DEFAULTS[preset]
|
|
for key, value in preset_defaults.items():
|
|
if normalized.get(key) in (None, ""):
|
|
normalized[key] = value
|
|
|
|
mode = str(normalized.get("training_strength_mode") or "target_steps")
|
|
if mode not in {"target_steps", "manual_epochs"}:
|
|
mode = "target_steps"
|
|
normalized["training_strength_mode"] = mode
|
|
|
|
save_strategy = str(normalized.get("save_strategy") or "final_only")
|
|
if save_strategy not in {"final_only", "every_n_epochs"}:
|
|
save_strategy = "final_only"
|
|
normalized["save_strategy"] = save_strategy
|
|
return normalized
|
|
|
|
|
|
def int_config(config: Dict[str, Any], key: str, default: int = 1) -> int:
|
|
try:
|
|
return int(config.get(key) if config.get(key) is not None else default)
|
|
except Exception:
|
|
return default
|
|
|
|
|
|
def task_image_count(task: Dict[str, Any]) -> int:
|
|
try:
|
|
return int(task.get("image_count") or len(task.get("images") or []))
|
|
except Exception:
|
|
return 0
|
|
|
|
|
|
def calculate_task_training_plan(task: Dict[str, Any], config: Dict[str, Any]) -> Dict[str, Any]:
|
|
config = normalized_config(config)
|
|
image_count = task_image_count(task)
|
|
batch_size = int_config(config, "train_batch_size", 1)
|
|
epochs = int_config(config, "max_train_epochs", 10)
|
|
target_steps = int_config(config, "target_steps_per_lora", 2000)
|
|
|
|
if image_count <= 0:
|
|
raise ValueError(f"image_count must be > 0 for task {task.get('folder_name') or task}")
|
|
if batch_size <= 0:
|
|
raise ValueError("train_batch_size must be > 0")
|
|
if epochs <= 0:
|
|
raise ValueError("max_train_epochs must be > 0")
|
|
|
|
mode = str(config.get("training_strength_mode") or "target_steps")
|
|
if mode == "target_steps":
|
|
if target_steps <= 0:
|
|
raise ValueError("target_steps_per_lora must be > 0")
|
|
repeats = math.ceil((target_steps * batch_size) / (image_count * epochs))
|
|
else:
|
|
repeats = int_config(config, "num_repeats", 1)
|
|
|
|
repeats = max(1, repeats)
|
|
actual_steps = math.ceil(image_count * repeats * epochs / batch_size)
|
|
warnings: List[str] = []
|
|
if mode == "target_steps":
|
|
if actual_steps < target_steps * 0.7:
|
|
warnings.append("actual_steps is below 70% of target_steps")
|
|
if actual_steps > target_steps * 1.8:
|
|
warnings.append("actual_steps is above 180% of target_steps; few images may overfit")
|
|
|
|
return {
|
|
"folder_name": str(task.get("folder_name") or ""),
|
|
"trigger_word": str(task.get("trigger_word") or task.get("folder_name") or ""),
|
|
"image_count": image_count,
|
|
"target_steps": target_steps if mode == "target_steps" else "",
|
|
"num_repeats": repeats,
|
|
"actual_steps": actual_steps,
|
|
"warnings": warnings,
|
|
}
|
|
|
|
|
|
def format_training_plan_summary(task_list: Sequence[Dict[str, Any]], config: Dict[str, Any]) -> str:
|
|
config = normalized_config(config)
|
|
lines = [
|
|
f"training_preset={config.get('training_preset')}",
|
|
f"training_strength_mode={config.get('training_strength_mode')}",
|
|
"actual_steps may be slightly above target_steps because num_repeats must be an integer.",
|
|
]
|
|
for task in task_list:
|
|
plan = calculate_task_training_plan(task, config)
|
|
target = plan["target_steps"] if plan["target_steps"] != "" else "manual"
|
|
line = (
|
|
f"{plan['folder_name']}: trigger={plan['trigger_word']}, images={plan['image_count']}, "
|
|
f"target_steps={target}, repeats={plan['num_repeats']}, actual_steps={plan['actual_steps']}"
|
|
)
|
|
if plan["warnings"]:
|
|
line += ", WARNING: " + "; ".join(plan["warnings"])
|
|
lines.append(line)
|
|
return "\n".join(lines)
|
|
|
|
|
|
def estimate_training_steps(task_list: Sequence[Dict[str, Any]], config: Dict[str, Any]) -> int:
|
|
total = 0
|
|
for task in task_list:
|
|
total += calculate_task_training_plan(task, config)["actual_steps"]
|
|
return max(1, total)
|
|
|
|
|
|
def parse_training_step(line: str) -> Tuple[int, int] | None:
|
|
if "steps:" not in line:
|
|
return None
|
|
matches = re.findall(r"(\d+)\s*/\s*(\d+)", line)
|
|
if not matches:
|
|
return None
|
|
current, total = matches[-1]
|
|
try:
|
|
return int(current), max(1, int(total))
|
|
except ValueError:
|
|
return None
|
|
|
|
|
|
def build_train_command(task: Dict[str, Any], config: Dict[str, Any]) -> Tuple[List[str], Path, str]:
|
|
config = normalized_config(config)
|
|
plan = calculate_task_training_plan(task, config)
|
|
python_executable = resolve_python_executable(str(config.get("python_executable") or ""))
|
|
script_path, script_warning = resolve_script_path(str(config.get("train_script_path") or ""))
|
|
trigger_word = str(task.get("trigger_word") or task.get("folder_name") or "anima_subject")
|
|
train_data_dir = expand_path(str(task.get("task_dir") or ""))
|
|
|
|
output_root = str(config.get("output_dir") or DEFAULT_ANIMA_LORA_OUTPUT_DIR).strip()
|
|
if output_root:
|
|
output_dir = Path(expand_path(output_root)) / trigger_word
|
|
else:
|
|
output_dir = Path(train_data_dir) / "anima_lora_output"
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
caption_extension = str(config.get("caption_extension") or ".txt").strip() or ".txt"
|
|
caption_separator = str(config.get("caption_separator") or ",").strip() or ","
|
|
keep_tokens = str(config.get("keep_tokens") if config.get("keep_tokens") is not None else "1").strip()
|
|
|
|
command = [
|
|
python_executable,
|
|
script_path,
|
|
"--train_data_dir",
|
|
train_data_dir,
|
|
"--anima_num_repeats",
|
|
str(plan["num_repeats"]),
|
|
"--output_dir",
|
|
str(output_dir),
|
|
"--output_name",
|
|
trigger_word,
|
|
"--caption_extension",
|
|
caption_extension,
|
|
"--caption_separator",
|
|
caption_separator,
|
|
]
|
|
if keep_tokens:
|
|
command.extend(["--keep_tokens", keep_tokens])
|
|
if bool(config.get("shuffle_caption", True)):
|
|
command.append("--shuffle_caption")
|
|
|
|
save_every_n_epochs = config.get("save_every_n_epochs")
|
|
if str(config.get("save_strategy") or "final_only") == "final_only":
|
|
save_every_n_epochs = config.get("max_train_epochs")
|
|
|
|
optional_pairs = [
|
|
("base_model_path", "--pretrained_model_name_or_path"),
|
|
("resolution", "--resolution"),
|
|
("train_batch_size", "--train_batch_size"),
|
|
("max_train_epochs", "--max_train_epochs"),
|
|
("network_dim", "--network_dim"),
|
|
("network_alpha", "--network_alpha"),
|
|
("learning_rate", "--learning_rate"),
|
|
("mixed_precision", "--mixed_precision"),
|
|
("optimizer_type", "--optimizer_type"),
|
|
]
|
|
for key, arg_name in optional_pairs:
|
|
value = config.get(key)
|
|
if value is None:
|
|
continue
|
|
value = str(value).strip()
|
|
if value:
|
|
if key in {"base_model_path"}:
|
|
value = expand_path(value)
|
|
command.extend([arg_name, value])
|
|
|
|
if save_every_n_epochs is not None and str(save_every_n_epochs).strip():
|
|
command.extend(["--save_every_n_epochs", str(save_every_n_epochs).strip()])
|
|
|
|
command.extend(normalize_arg_path(arg) for arg in split_extra_args(str(config.get("extra_args") or "")))
|
|
return command, output_dir, script_warning
|
|
|
|
|
|
def run_job(job_id: str, task_list: Sequence[Dict[str, Any]], config: Dict[str, Any], progress_bar=None) -> None:
|
|
config = normalized_config(config)
|
|
stop_on_error = bool(config.get("stop_on_error", True))
|
|
total = len(task_list)
|
|
total_training_steps = estimate_training_steps(task_list, config)
|
|
finished_training_steps = 0
|
|
if progress_bar is not None:
|
|
progress_bar.update_absolute(0, total_training_steps)
|
|
update_job(
|
|
job_id,
|
|
status="running",
|
|
total_tasks=total,
|
|
completed_tasks=0,
|
|
failed_tasks=0,
|
|
started_at=now_iso(),
|
|
config=config,
|
|
)
|
|
append_job_log(job_id, f"Batch training started with {total} task(s).")
|
|
append_job_log(job_id, format_training_plan_summary(task_list, config))
|
|
|
|
completed = 0
|
|
failed = 0
|
|
for index, task in enumerate(task_list, start=1):
|
|
trigger_word = str(task.get("trigger_word") or task.get("folder_name") or f"task_{index}")
|
|
update_job(job_id, current_task=trigger_word, current_index=index)
|
|
append_job_log(job_id, f"Starting task {index}/{total}: {trigger_word}")
|
|
|
|
try:
|
|
command, output_dir, script_warning = build_train_command(task, config)
|
|
if script_warning:
|
|
append_job_log(job_id, f"ERROR: {script_warning}")
|
|
failed += 1
|
|
if stop_on_error:
|
|
break
|
|
continue
|
|
|
|
log_path = output_dir / f"{trigger_word}.train.log"
|
|
append_job_log(job_id, "Command: " + shlex.join(command))
|
|
task_estimated_steps = max(1, estimate_training_steps([task], config))
|
|
task_reported_step = 0
|
|
task_reported_total = task_estimated_steps
|
|
with log_path.open("a", encoding="utf-8", errors="ignore") as log_file:
|
|
log_file.write(f"\n[{now_iso()}] Command: {shlex.join(command)}\n")
|
|
process = subprocess.Popen(
|
|
command,
|
|
cwd=str(Path(task.get("task_dir") or ".").resolve()),
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
bufsize=1,
|
|
)
|
|
assert process.stdout is not None
|
|
for line in process.stdout:
|
|
log_file.write(line)
|
|
log_file.flush()
|
|
step_progress = parse_training_step(line)
|
|
if progress_bar is not None and step_progress is not None:
|
|
task_reported_step, task_reported_total = step_progress
|
|
absolute_step = finished_training_steps + min(task_reported_step, task_reported_total)
|
|
absolute_total = max(total_training_steps, finished_training_steps + task_reported_total)
|
|
progress_bar.update_absolute(absolute_step, absolute_total)
|
|
return_code = process.wait()
|
|
|
|
if return_code == 0:
|
|
completed += 1
|
|
finished_training_steps += max(task_estimated_steps, task_reported_total)
|
|
if progress_bar is not None:
|
|
progress_bar.update_absolute(min(finished_training_steps, total_training_steps), total_training_steps)
|
|
append_job_log(job_id, f"DONE: {trigger_word} -> {output_dir}")
|
|
else:
|
|
failed += 1
|
|
append_job_log(job_id, f"FAILED: {trigger_word}, return code {return_code}. Log: {log_path}")
|
|
if stop_on_error:
|
|
break
|
|
except Exception as exc:
|
|
failed += 1
|
|
append_job_log(job_id, f"ERROR: {trigger_word}: {exc}")
|
|
if stop_on_error:
|
|
break
|
|
finally:
|
|
update_job(job_id, completed_tasks=completed, failed_tasks=failed)
|
|
|
|
status = "completed" if failed == 0 and completed == total else "failed"
|
|
if failed and completed:
|
|
status = "partial_failed"
|
|
update_job(
|
|
job_id,
|
|
status=status,
|
|
completed_tasks=completed,
|
|
failed_tasks=failed,
|
|
finished_at=now_iso(),
|
|
current_task="",
|
|
)
|
|
append_job_log(job_id, f"Batch training finished: status={status}, completed={completed}, failed={failed}.")
|
|
if progress_bar is not None:
|
|
progress_bar.update_absolute(total_training_steps, total_training_steps)
|
|
|
|
|
|
def start_batch_training(
|
|
task_list: Sequence[Dict[str, Any]],
|
|
config: Dict[str, Any],
|
|
*,
|
|
background: bool = False,
|
|
) -> Tuple[str, str]:
|
|
if not task_list:
|
|
return "", "ERROR: TASK_LIST is empty."
|
|
|
|
job_id = uuid.uuid4().hex[:12]
|
|
update_job(
|
|
job_id,
|
|
status="queued",
|
|
total_tasks=len(task_list),
|
|
completed_tasks=0,
|
|
failed_tasks=0,
|
|
created_at=now_iso(),
|
|
logs=[],
|
|
)
|
|
|
|
if background:
|
|
thread = threading.Thread(target=run_job, args=(job_id, list(task_list), dict(config), None), daemon=True)
|
|
thread.start()
|
|
return job_id, f"Started Anima batch training job: {job_id}"
|
|
|
|
progress_bar = ProgressBar(estimate_training_steps(task_list, config)) if ProgressBar is not None else None
|
|
run_job(job_id, list(task_list), dict(config), progress_bar)
|
|
status = get_job(job_id) or {}
|
|
status_text = str(status.get("status", "unknown"))
|
|
if status_text not in {"completed"}:
|
|
logs = status.get("logs") or []
|
|
tail = "\n".join(str(line) for line in logs[-20:])
|
|
raise RuntimeError(f"Anima batch training failed: {status_text}\n{tail}")
|
|
return job_id, f"Finished Anima batch training job: {job_id}, status={status_text}"
|
|
|
|
|
|
def format_job_status(job_id: str, tail_lines: int = 80) -> str:
|
|
if not job_id.strip():
|
|
return "ERROR: job_id is empty."
|
|
job = get_job(job_id.strip())
|
|
if not job:
|
|
return f"ERROR: job not found: {job_id}"
|
|
|
|
logs = job.get("logs") or []
|
|
tail_lines = max(1, int(tail_lines or 80))
|
|
visible_logs = logs[-tail_lines:]
|
|
header = [
|
|
f"job_id: {job_id}",
|
|
f"status: {job.get('status', 'unknown')}",
|
|
f"progress: {job.get('completed_tasks', 0)}/{job.get('total_tasks', 0)} completed, {job.get('failed_tasks', 0)} failed",
|
|
f"current_task: {job.get('current_task', '')}",
|
|
f"updated_at: {job.get('updated_at', '')}",
|
|
"",
|
|
"log_tail:",
|
|
]
|
|
return "\n".join(header + visible_logs)
|
|
|
|
|
|
def find_trained_loras(output_dir: str, name_filter: str = "") -> Tuple[str, str]:
|
|
root = Path(os.path.expanduser(output_dir)).resolve() if output_dir.strip() else Path.cwd()
|
|
if not root.exists():
|
|
return "", f"ERROR: output_dir does not exist: {root}"
|
|
|
|
suffixes = {".safetensors", ".pt", ".ckpt"}
|
|
name_filter = (name_filter or "").strip().lower()
|
|
matches: List[Path] = []
|
|
for path in root.rglob("*"):
|
|
if not path.is_file() or path.suffix.lower() not in suffixes:
|
|
continue
|
|
if name_filter and name_filter not in path.name.lower():
|
|
continue
|
|
matches.append(path)
|
|
|
|
matches.sort(key=lambda p: p.stat().st_mtime, reverse=True)
|
|
result = "\n".join(str(path) for path in matches)
|
|
return result, f"Found {len(matches)} trained LoRA file(s) under {root}."
|