743 lines
29 KiB
Python
743 lines
29 KiB
Python
from __future__ import annotations
|
|
|
|
import time
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Tuple
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageOps
|
|
|
|
from .dataset_utils import (
|
|
DEFAULT_IMAGE_EXTENSIONS,
|
|
iter_task_images,
|
|
normalize_trigger_word,
|
|
refresh_task_caption_counts,
|
|
scan_task_folders,
|
|
)
|
|
from .train_runner import find_trained_loras, format_job_status, start_batch_training
|
|
from .vlm_caption import (
|
|
CAPTION_MODES,
|
|
TRIGGER_MODES,
|
|
VlmCaptionConfig,
|
|
call_vlm_caption_api,
|
|
clean_caption,
|
|
save_caption_file,
|
|
)
|
|
|
|
|
|
CATEGORY = "Anima/Batch LoRA"
|
|
|
|
PLUGIN_DIR = Path(__file__).resolve().parent
|
|
|
|
|
|
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_BACKEND_PYTHON = BACKEND_DIR / ".venv" / ("Scripts/python.exe" if os.name == "nt" else "bin/python")
|
|
|
|
DEFAULT_TRAIN_SCRIPT_PATH = str(PLUGIN_DIR / "scripts" / "anima_train_network_wrapper.py")
|
|
DEFAULT_PYTHON_EXECUTABLE = str(DEFAULT_BACKEND_PYTHON)
|
|
DEFAULT_ANIMA_BASE_MODEL = str(COMFYUI_DIR / "models/diffusion_models/anima-base-v1.0.safetensors")
|
|
DEFAULT_ANIMA_LORA_OUTPUT_DIR = str(COMFYUI_DIR / "models/loras/anima_lora")
|
|
DEFAULT_ANIMA_EXTRA_ARGS = (
|
|
f"--qwen3 {COMFYUI_DIR / 'models/clip/qwen_3_06b_base.safetensors'} "
|
|
f"--vae {COMFYUI_DIR / 'models/vae/qwen_image_vae.safetensors'} "
|
|
"--save_model_as safetensors --network_module networks.lora_anima "
|
|
"--network_train_unet_only --lr_scheduler constant --timestep_sampling sigmoid "
|
|
"--discrete_flow_shift 1.0 --gradient_checkpointing --cache_latents "
|
|
"--vae_chunk_size 64 --vae_disable_cache --max_data_loader_n_workers 0"
|
|
)
|
|
|
|
|
|
def _derive_trigger_word(task: Dict[str, Any], image_path: Path, trigger_mode: str) -> str:
|
|
if trigger_mode == "root_folder":
|
|
root = task.get("dataset_root") or ""
|
|
if root:
|
|
return normalize_trigger_word(Path(str(root)).name)
|
|
if trigger_mode == "parent_folder":
|
|
return normalize_trigger_word(image_path.parent.name)
|
|
return normalize_trigger_word(str(task.get("trigger_word") or task.get("folder_name") or image_path.parent.name))
|
|
|
|
|
|
def _read_caption(path: Path) -> str:
|
|
try:
|
|
return path.read_text(encoding="utf-8", errors="ignore").strip()
|
|
except Exception:
|
|
return ""
|
|
|
|
|
|
def _caption_exists(path: Path) -> bool:
|
|
caption_path = path.with_suffix(".txt")
|
|
return caption_path.exists() and _read_caption(caption_path) != ""
|
|
|
|
|
|
def _load_square_image_tensor(image_path: Path, image_size: int) -> torch.Tensor:
|
|
image_size = max(64, int(image_size))
|
|
with Image.open(image_path) as image:
|
|
image = ImageOps.exif_transpose(image).convert("RGB")
|
|
width, height = image.size
|
|
scale = min(image_size / max(width, height), 1.0)
|
|
new_width = max(1, int(round(width * scale)))
|
|
new_height = max(1, int(round(height * scale)))
|
|
if (new_width, new_height) != image.size:
|
|
image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
|
canvas = Image.new("RGB", (image_size, image_size), (0, 0, 0))
|
|
left = (image_size - new_width) // 2
|
|
top = (image_size - new_height) // 2
|
|
canvas.paste(image, (left, top))
|
|
array = np.asarray(canvas).astype(np.float32) / 255.0
|
|
return torch.from_numpy(array)
|
|
|
|
|
|
def _normalize_caption_list(captions: Any) -> List[str]:
|
|
if captions is None:
|
|
return []
|
|
if isinstance(captions, (list, tuple)):
|
|
return [str(item).strip() for item in captions]
|
|
text = str(captions).strip()
|
|
if not text:
|
|
return []
|
|
if "====== Image" in text:
|
|
chunks = []
|
|
current: List[str] = []
|
|
for line in text.splitlines():
|
|
if line.strip().startswith("====== Image"):
|
|
if current:
|
|
chunks.append("\n".join(current).strip())
|
|
current = []
|
|
else:
|
|
current.append(line)
|
|
if current:
|
|
chunks.append("\n".join(current).strip())
|
|
return [chunk for chunk in chunks if chunk]
|
|
return [text]
|
|
|
|
|
|
def _unwrap_record_list(value: Any) -> List[Dict[str, Any]]:
|
|
if isinstance(value, list) and len(value) == 1 and isinstance(value[0], list):
|
|
value = value[0]
|
|
if isinstance(value, list):
|
|
return [item for item in value if isinstance(item, dict)]
|
|
return []
|
|
|
|
|
|
def _unwrap_task_list(value: Any) -> List[Dict[str, Any]]:
|
|
if isinstance(value, list) and value and all(isinstance(item, dict) for item in value):
|
|
return value
|
|
if isinstance(value, list) and len(value) == 1 and isinstance(value[0], list):
|
|
return _unwrap_task_list(value[0])
|
|
return []
|
|
|
|
|
|
class AnimaBatchFolderLoader:
|
|
RETURN_TYPES = ("TASK_LIST", "STRING")
|
|
RETURN_NAMES = ("TASK_LIST", "scan_log")
|
|
FUNCTION = "load"
|
|
CATEGORY = CATEGORY
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"dataset_root": ("STRING", {"default": "", "multiline": False}),
|
|
"image_extensions": ("STRING", {"default": DEFAULT_IMAGE_EXTENSIONS, "multiline": False}),
|
|
"folder_filter": ("STRING", {"default": "", "multiline": False}),
|
|
"include_root_images": ("BOOLEAN", {"default": False}),
|
|
"max_images_per_folder": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 1}),
|
|
}
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
return time.time()
|
|
|
|
def load(
|
|
self,
|
|
dataset_root: str,
|
|
image_extensions: str = DEFAULT_IMAGE_EXTENSIONS,
|
|
folder_filter: str = "",
|
|
include_root_images: bool = False,
|
|
max_images_per_folder: int = 0,
|
|
) -> Tuple[List[Dict[str, Any]], str]:
|
|
return scan_task_folders(
|
|
dataset_root=dataset_root,
|
|
image_extensions=image_extensions,
|
|
folder_filter=folder_filter,
|
|
include_root_images=include_root_images,
|
|
max_images_per_folder=max_images_per_folder,
|
|
)
|
|
|
|
|
|
class AnimaLlamaCppCaptionPrompt:
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("caption_prompt",)
|
|
FUNCTION = "build"
|
|
CATEGORY = CATEGORY
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"caption_mode": (CAPTION_MODES, {"default": "hybrid"}),
|
|
"nsfw_caption_mode": ("BOOLEAN", {"default": True}),
|
|
"extra_rules": ("STRING", {"default": "", "multiline": True}),
|
|
}
|
|
}
|
|
|
|
def build(self, caption_mode: str = "hybrid", nsfw_caption_mode: bool = True, extra_rules: str = ""):
|
|
if caption_mode == "tags_only":
|
|
mode_text = "Output anime illustration followed by concise Danbooru/Gelbooru-style tags."
|
|
elif caption_mode == "natural_language":
|
|
mode_text = "Output anime illustration followed by one short natural-language caption."
|
|
else:
|
|
mode_text = "Output anime illustration, Danbooru/Gelbooru-style tags, then one short natural-language phrase."
|
|
|
|
adult_rules = ""
|
|
if nsfw_caption_mode:
|
|
adult_rules = (
|
|
"\nIf the image contains adult-only visual content, describe it clinically and briefly as "
|
|
"training tags only. Do not write erotic prose. Do not describe minors or young-looking "
|
|
"subjects as sexualized. If the subject appears underage or ambiguous, avoid sexual tags "
|
|
"and use safe neutral visual tags only."
|
|
)
|
|
|
|
prompt = (
|
|
"You are a captioning engine for Anima-Base LoRA training.\n"
|
|
"Return exactly one line and nothing else.\n"
|
|
"Begin exactly with: FINAL: anime illustration,\n"
|
|
f"{mode_text}\n"
|
|
"Use concise comma-separated English tags for only visible content: subject, clothing, hair, eyes, pose, expression, composition, background, and style.\n"
|
|
"Use lowercase ordinary tags with spaces, not underscores.\n"
|
|
"Do not invent a trigger word; it will be inserted automatically later.\n"
|
|
"Preserve the visible differences between images; do not make every image use the same caption.\n"
|
|
"For character LoRA, keep fixed character traits. For clothing LoRA, keep color, material, pattern, and cut. For style LoRA, keep linework, coloring, lighting, composition, and background style.\n"
|
|
"Do not explain, reason, restate instructions, mention uncertainty, or output channel/control tokens.\n"
|
|
"Do not include artist names, copyright names, watermark, logo, signature, score tags, rating tags, or quality tags."
|
|
f"{adult_rules}"
|
|
)
|
|
if extra_rules.strip():
|
|
prompt += "\n" + extra_rules.strip()
|
|
return (prompt,)
|
|
|
|
|
|
class AnimaLoadCaptionImageBatch:
|
|
RETURN_TYPES = ("IMAGE", "IMAGE_RECORDS", "TASK_LIST", "STRING")
|
|
RETURN_NAMES = ("images", "IMAGE_RECORDS", "TASK_LIST", "load_log")
|
|
FUNCTION = "load"
|
|
CATEGORY = CATEGORY
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"TASK_LIST": ("TASK_LIST",),
|
|
"task_index": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 1}),
|
|
"start_index": ("INT", {"default": 0, "min": 0, "max": 1000000, "step": 1}),
|
|
"max_images": ("INT", {"default": 32, "min": 1, "max": 4096, "step": 1}),
|
|
"image_size": ("INT", {"default": 1024, "min": 256, "max": 4096, "step": 64}),
|
|
"skip_existing_caption": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
return time.time()
|
|
|
|
def load(
|
|
self,
|
|
TASK_LIST,
|
|
task_index: int = 0,
|
|
start_index: int = 0,
|
|
max_images: int = 32,
|
|
image_size: int = 1024,
|
|
skip_existing_caption: bool = True,
|
|
):
|
|
task_list = TASK_LIST or []
|
|
logs: List[str] = []
|
|
records: List[Dict[str, Any]] = []
|
|
tensors: List[torch.Tensor] = []
|
|
|
|
if not task_list:
|
|
blank = torch.zeros((1, image_size, image_size, 3), dtype=torch.float32)
|
|
return blank, records, task_list, "ERROR: TASK_LIST is empty."
|
|
|
|
selected_tasks = task_list
|
|
if task_index > 0:
|
|
if task_index > len(task_list):
|
|
blank = torch.zeros((1, image_size, image_size, 3), dtype=torch.float32)
|
|
return blank, records, task_list, f"ERROR: task_index {task_index} is out of range. TASK_LIST has {len(task_list)} task(s)."
|
|
selected_tasks = [task_list[task_index - 1]]
|
|
|
|
candidates: List[Tuple[Dict[str, Any], Path, str]] = []
|
|
for task in selected_tasks:
|
|
for _task, image_path, trigger_word in iter_task_images([task]):
|
|
if skip_existing_caption and _caption_exists(image_path):
|
|
continue
|
|
candidates.append((_task, image_path, trigger_word))
|
|
|
|
candidates = candidates[start_index:start_index + max_images]
|
|
for task, image_path, trigger_word in candidates:
|
|
try:
|
|
tensors.append(_load_square_image_tensor(image_path, image_size))
|
|
records.append(
|
|
{
|
|
"image_path": str(image_path),
|
|
"caption_path": str(image_path.with_suffix(".txt")),
|
|
"trigger_word": trigger_word,
|
|
"task_dir": str(task.get("task_dir") or image_path.parent),
|
|
"folder_name": str(task.get("folder_name") or image_path.parent.name),
|
|
}
|
|
)
|
|
logs.append(f"OK: loaded {image_path}")
|
|
except Exception as exc:
|
|
logs.append(f"ERROR: failed to load {image_path}: {exc}")
|
|
|
|
if tensors:
|
|
images = torch.stack(tensors, dim=0)
|
|
else:
|
|
images = torch.zeros((1, image_size, image_size, 3), dtype=torch.float32)
|
|
logs.insert(0, "WARNING: no images loaded. Output is a blank placeholder image.")
|
|
|
|
logs.insert(0, f"Loaded {len(records)} image(s) for llama-cpp_vlm captioning.")
|
|
return images, records, task_list, "\n".join(logs)
|
|
|
|
|
|
class AnimaSaveLlamaCppCaptions:
|
|
RETURN_TYPES = ("TASK_LIST", "STRING")
|
|
RETURN_NAMES = ("TASK_LIST", "caption_log")
|
|
FUNCTION = "save"
|
|
CATEGORY = CATEGORY
|
|
OUTPUT_NODE = True
|
|
INPUT_IS_LIST = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"TASK_LIST": ("TASK_LIST",),
|
|
"IMAGE_RECORDS": ("IMAGE_RECORDS",),
|
|
"caption_list": ("STRING", {"forceInput": True}),
|
|
"overwrite_existing_caption": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
return time.time()
|
|
|
|
def save(self, TASK_LIST, IMAGE_RECORDS, caption_list, overwrite_existing_caption: bool = True):
|
|
task_list = _unwrap_task_list(TASK_LIST)
|
|
records = _unwrap_record_list(IMAGE_RECORDS)
|
|
captions = _normalize_caption_list(caption_list)
|
|
if isinstance(overwrite_existing_caption, list):
|
|
overwrite_existing_caption = bool(overwrite_existing_caption[0]) if overwrite_existing_caption else True
|
|
logs: List[str] = []
|
|
saved = 0
|
|
skipped = 0
|
|
failed = 0
|
|
|
|
if not records:
|
|
return task_list, "ERROR: IMAGE_RECORDS is empty."
|
|
if len(captions) != len(records):
|
|
return (
|
|
task_list,
|
|
f"ERROR: caption count mismatch. records={len(records)}, captions={len(captions)}. Connect llama-cpp output_list, not the joined output.",
|
|
)
|
|
|
|
for record, raw_caption in zip(records, captions):
|
|
image_path = Path(str(record.get("image_path") or ""))
|
|
caption_path = Path(str(record.get("caption_path") or image_path.with_suffix(".txt")))
|
|
trigger_word = normalize_trigger_word(str(record.get("trigger_word") or image_path.parent.name))
|
|
try:
|
|
if caption_path.exists() and _read_caption(caption_path) and not overwrite_existing_caption:
|
|
skipped += 1
|
|
logs.append(f"SKIP: existing caption: {caption_path}")
|
|
continue
|
|
caption = clean_caption(raw_caption, trigger_word)
|
|
caption_path.write_text(f"{caption}\n", encoding="utf-8")
|
|
saved += 1
|
|
logs.append(f"OK: {caption_path.name} -> {caption}")
|
|
except Exception as exc:
|
|
failed += 1
|
|
logs.append(f"ERROR: {caption_path}: {exc}")
|
|
|
|
refreshed = refresh_task_caption_counts(task_list)
|
|
logs.insert(0, f"Saved llama-cpp captions. saved={saved}, skipped={skipped}, failed={failed}")
|
|
return refreshed, "\n".join(logs)
|
|
|
|
|
|
class AnimaVlmCaptionConfig:
|
|
RETURN_TYPES = ("VLM_CAPTION_CONFIG",)
|
|
RETURN_NAMES = ("VLM_CAPTION_CONFIG",)
|
|
FUNCTION = "build"
|
|
CATEGORY = CATEGORY
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"api_url": (
|
|
"STRING",
|
|
{"default": "http://127.0.0.1:8000/v1/chat/completions", "multiline": False},
|
|
),
|
|
"api_key": ("STRING", {"default": "", "multiline": False}),
|
|
"model_name": ("STRING", {"default": "gemma-vlm", "multiline": False}),
|
|
"caption_mode": (CAPTION_MODES, {"default": "hybrid"}),
|
|
"trigger_mode": (TRIGGER_MODES, {"default": "folder_name"}),
|
|
"max_tokens": ("INT", {"default": 180, "min": 16, "max": 2048, "step": 1}),
|
|
"temperature": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 2.0, "step": 0.05}),
|
|
"overwrite_existing_caption": ("BOOLEAN", {"default": False}),
|
|
"nsfw_caption_mode": ("BOOLEAN", {"default": True}),
|
|
"request_timeout": ("INT", {"default": 120, "min": 5, "max": 1800, "step": 1}),
|
|
}
|
|
}
|
|
|
|
def build(
|
|
self,
|
|
api_url: str,
|
|
api_key: str,
|
|
model_name: str,
|
|
caption_mode: str,
|
|
trigger_mode: str,
|
|
max_tokens: int,
|
|
temperature: float,
|
|
overwrite_existing_caption: bool,
|
|
nsfw_caption_mode: bool,
|
|
request_timeout: int,
|
|
):
|
|
config = VlmCaptionConfig(
|
|
api_url=api_url,
|
|
api_key=api_key,
|
|
model_name=model_name,
|
|
caption_mode=caption_mode,
|
|
trigger_mode=trigger_mode,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
overwrite_existing_caption=overwrite_existing_caption,
|
|
nsfw_caption_mode=nsfw_caption_mode,
|
|
request_timeout=request_timeout,
|
|
)
|
|
return (config.to_dict(),)
|
|
|
|
|
|
class AnimaBatchVlmCaption:
|
|
RETURN_TYPES = ("TASK_LIST", "STRING")
|
|
RETURN_NAMES = ("TASK_LIST", "caption_log")
|
|
FUNCTION = "caption"
|
|
CATEGORY = CATEGORY
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"TASK_LIST": ("TASK_LIST",),
|
|
"VLM_CAPTION_CONFIG": ("VLM_CAPTION_CONFIG",),
|
|
}
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
return time.time()
|
|
|
|
def caption(self, TASK_LIST, VLM_CAPTION_CONFIG):
|
|
task_list = TASK_LIST or []
|
|
config = VlmCaptionConfig.from_dict(VLM_CAPTION_CONFIG or {})
|
|
logs: List[str] = []
|
|
saved = 0
|
|
skipped = 0
|
|
failed = 0
|
|
|
|
if not task_list:
|
|
return task_list, "ERROR: TASK_LIST is empty."
|
|
if not config.api_url:
|
|
return task_list, "ERROR: api_url is empty."
|
|
|
|
for task, image_path, _trigger_word in iter_task_images(task_list):
|
|
trigger_word = _derive_trigger_word(task, image_path, config.trigger_mode)
|
|
caption_path = image_path.with_suffix(".txt")
|
|
existing = _read_caption(caption_path)
|
|
if existing and not config.overwrite_existing_caption:
|
|
skipped += 1
|
|
logs.append(f"SKIP: existing caption: {caption_path}")
|
|
continue
|
|
|
|
try:
|
|
raw_caption, error = call_vlm_caption_api(image_path, trigger_word, config)
|
|
if error:
|
|
failed += 1
|
|
logs.append(f"ERROR: {image_path}: {error}")
|
|
continue
|
|
caption = clean_caption(raw_caption, trigger_word)
|
|
save_caption_file(image_path, caption)
|
|
saved += 1
|
|
logs.append(f"OK: {image_path.name} -> {caption}")
|
|
except Exception as exc:
|
|
failed += 1
|
|
logs.append(f"ERROR: {image_path}: {exc}")
|
|
|
|
refreshed = refresh_task_caption_counts(task_list)
|
|
summary = f"VLM caption finished. saved={saved}, skipped={skipped}, failed={failed}"
|
|
logs.insert(0, summary)
|
|
return refreshed, "\n".join(logs)
|
|
|
|
|
|
class AnimaCaptionPrepare:
|
|
RETURN_TYPES = ("TASK_LIST", "STRING")
|
|
RETURN_NAMES = ("TASK_LIST", "prepare_log")
|
|
FUNCTION = "prepare"
|
|
CATEGORY = CATEGORY
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"TASK_LIST": ("TASK_LIST",),
|
|
"trigger_mode": (TRIGGER_MODES, {"default": "folder_name"}),
|
|
"normalize_existing_caption": ("BOOLEAN", {"default": True}),
|
|
"create_missing_caption": ("BOOLEAN", {"default": True}),
|
|
"default_caption_body": ("STRING", {"default": "anime illustration", "multiline": False}),
|
|
}
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
return time.time()
|
|
|
|
def prepare(
|
|
self,
|
|
TASK_LIST,
|
|
trigger_mode: str = "folder_name",
|
|
normalize_existing_caption: bool = True,
|
|
create_missing_caption: bool = True,
|
|
default_caption_body: str = "anime illustration",
|
|
):
|
|
task_list = TASK_LIST or []
|
|
logs: List[str] = []
|
|
created = 0
|
|
normalized = 0
|
|
skipped = 0
|
|
failed = 0
|
|
|
|
if not task_list:
|
|
return task_list, "ERROR: TASK_LIST is empty."
|
|
|
|
for task, image_path, _trigger_word in iter_task_images(task_list):
|
|
trigger_word = _derive_trigger_word(task, image_path, trigger_mode)
|
|
caption_path = image_path.with_suffix(".txt")
|
|
existing = _read_caption(caption_path)
|
|
|
|
try:
|
|
if existing:
|
|
if normalize_existing_caption:
|
|
cleaned = clean_caption(existing, trigger_word)
|
|
if cleaned != existing:
|
|
caption_path.write_text(f"{cleaned}\n", encoding="utf-8")
|
|
normalized += 1
|
|
logs.append(f"NORMALIZE: {caption_path.name}")
|
|
else:
|
|
skipped += 1
|
|
else:
|
|
skipped += 1
|
|
elif create_missing_caption:
|
|
fallback = clean_caption(default_caption_body, trigger_word)
|
|
caption_path.write_text(f"{fallback}\n", encoding="utf-8")
|
|
created += 1
|
|
logs.append(f"CREATE: {caption_path.name} -> {fallback}")
|
|
else:
|
|
skipped += 1
|
|
logs.append(f"SKIP: missing caption: {caption_path}")
|
|
except Exception as exc:
|
|
failed += 1
|
|
logs.append(f"ERROR: {caption_path}: {exc}")
|
|
|
|
refreshed = refresh_task_caption_counts(task_list)
|
|
logs.insert(0, f"Caption prepare finished. created={created}, normalized={normalized}, skipped={skipped}, failed={failed}")
|
|
return refreshed, "\n".join(logs)
|
|
|
|
|
|
class AnimaBatchTrainConfig:
|
|
RETURN_TYPES = ("ANIMA_TRAIN_CONFIG",)
|
|
RETURN_NAMES = ("ANIMA_TRAIN_CONFIG",)
|
|
FUNCTION = "build"
|
|
CATEGORY = CATEGORY
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"train_script_path": ("STRING", {"default": DEFAULT_TRAIN_SCRIPT_PATH, "multiline": False}),
|
|
"python_executable": ("STRING", {"default": DEFAULT_PYTHON_EXECUTABLE, "multiline": False}),
|
|
"base_model_path": ("STRING", {"default": DEFAULT_ANIMA_BASE_MODEL, "multiline": False}),
|
|
"output_dir": ("STRING", {"default": DEFAULT_ANIMA_LORA_OUTPUT_DIR, "multiline": False}),
|
|
"resolution": ("STRING", {"default": "1024,1024", "multiline": False}),
|
|
"train_batch_size": ("INT", {"default": 1, "min": 1, "max": 64, "step": 1}),
|
|
"max_train_epochs": ("INT", {"default": 10, "min": 1, "max": 10000, "step": 1}),
|
|
"save_every_n_epochs": ("INT", {"default": 1, "min": 1, "max": 10000, "step": 1}),
|
|
"network_dim": ("INT", {"default": 32, "min": 1, "max": 1024, "step": 1}),
|
|
"network_alpha": ("INT", {"default": 16, "min": 1, "max": 1024, "step": 1}),
|
|
"learning_rate": ("STRING", {"default": "1e-4", "multiline": False}),
|
|
"mixed_precision": (("bf16", "fp16", "no"), {"default": "bf16"}),
|
|
"optimizer_type": ("STRING", {"default": "adamw8bit", "multiline": False}),
|
|
"shuffle_caption": ("BOOLEAN", {"default": True}),
|
|
"keep_tokens": ("INT", {"default": 1, "min": 0, "max": 1000, "step": 1}),
|
|
"caption_separator": ("STRING", {"default": ",", "multiline": False}),
|
|
"caption_extension": ("STRING", {"default": ".txt", "multiline": False}),
|
|
"extra_args": ("STRING", {"default": DEFAULT_ANIMA_EXTRA_ARGS, "multiline": True}),
|
|
"stop_on_error": ("BOOLEAN", {"default": True}),
|
|
"training_preset": (("smoke_test", "standard", "strong", "style"), {"default": "standard"}),
|
|
"training_strength_mode": (("target_steps", "manual_epochs"), {"default": "target_steps"}),
|
|
"target_steps_per_lora": ("INT", {"default": 2000, "min": 1, "max": 1000000, "step": 50}),
|
|
"num_repeats": ("INT", {"default": 1, "min": 1, "max": 10000, "step": 1}),
|
|
"save_strategy": (("final_only", "every_n_epochs"), {"default": "final_only"}),
|
|
}
|
|
}
|
|
|
|
def build(
|
|
self,
|
|
train_script_path: str,
|
|
python_executable: str,
|
|
base_model_path: str,
|
|
output_dir: str,
|
|
resolution: str,
|
|
train_batch_size: int,
|
|
max_train_epochs: int,
|
|
save_every_n_epochs: int,
|
|
network_dim: int,
|
|
network_alpha: int,
|
|
learning_rate: str,
|
|
mixed_precision: str,
|
|
optimizer_type: str,
|
|
shuffle_caption: bool,
|
|
keep_tokens: int,
|
|
caption_separator: str,
|
|
caption_extension: str,
|
|
extra_args: str,
|
|
stop_on_error: bool,
|
|
training_preset: str = "standard",
|
|
training_strength_mode: str = "target_steps",
|
|
target_steps_per_lora: int = 2000,
|
|
num_repeats: int = 1,
|
|
save_strategy: str = "final_only",
|
|
):
|
|
return (
|
|
{
|
|
"train_script_path": train_script_path,
|
|
"python_executable": python_executable,
|
|
"base_model_path": base_model_path,
|
|
"output_dir": output_dir,
|
|
"resolution": resolution,
|
|
"train_batch_size": train_batch_size,
|
|
"max_train_epochs": max_train_epochs,
|
|
"save_every_n_epochs": save_every_n_epochs,
|
|
"network_dim": network_dim,
|
|
"network_alpha": network_alpha,
|
|
"learning_rate": learning_rate,
|
|
"mixed_precision": mixed_precision,
|
|
"optimizer_type": optimizer_type,
|
|
"shuffle_caption": shuffle_caption,
|
|
"keep_tokens": keep_tokens,
|
|
"caption_separator": caption_separator,
|
|
"caption_extension": caption_extension,
|
|
"extra_args": extra_args,
|
|
"stop_on_error": stop_on_error,
|
|
"training_preset": training_preset,
|
|
"training_strength_mode": training_strength_mode,
|
|
"target_steps_per_lora": target_steps_per_lora,
|
|
"num_repeats": num_repeats,
|
|
"save_strategy": save_strategy,
|
|
},
|
|
)
|
|
|
|
|
|
class AnimaBatchStartTrain:
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("job_id", "train_log")
|
|
FUNCTION = "start"
|
|
CATEGORY = CATEGORY
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"TASK_LIST": ("TASK_LIST",),
|
|
"ANIMA_TRAIN_CONFIG": ("ANIMA_TRAIN_CONFIG",),
|
|
"start_now": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
return time.time()
|
|
|
|
def start(self, TASK_LIST, ANIMA_TRAIN_CONFIG, start_now: bool = True):
|
|
if not start_now:
|
|
return "", "start_now=false, training was not started."
|
|
return start_batch_training(TASK_LIST or [], ANIMA_TRAIN_CONFIG or {})
|
|
|
|
|
|
class AnimaBatchCheckStatus:
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("status_log",)
|
|
FUNCTION = "check"
|
|
CATEGORY = CATEGORY
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"job_id": ("STRING", {"default": "", "multiline": False}),
|
|
"tail_lines": ("INT", {"default": 80, "min": 1, "max": 1000, "step": 1}),
|
|
}
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
return time.time()
|
|
|
|
def check(self, job_id: str, tail_lines: int = 80):
|
|
return (format_job_status(job_id, tail_lines),)
|
|
|
|
|
|
class AnimaFindTrainedLora:
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("lora_paths", "lora_log")
|
|
FUNCTION = "find"
|
|
CATEGORY = CATEGORY
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"output_dir": ("STRING", {"default": DEFAULT_ANIMA_LORA_OUTPUT_DIR, "multiline": False}),
|
|
"name_filter": ("STRING", {"default": "", "multiline": False}),
|
|
}
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
return time.time()
|
|
|
|
def find(self, output_dir: str, name_filter: str = ""):
|
|
return find_trained_loras(output_dir, name_filter)
|