418 lines
17 KiB
Python
418 lines
17 KiB
Python
from collections import OrderedDict
|
|
import os
|
|
import sqlite3
|
|
import asyncio
|
|
import concurrent.futures
|
|
from extensions_built_in.sd_trainer.SDTrainer import SDTrainer
|
|
from typing import Literal, Optional
|
|
import threading
|
|
import time
|
|
import signal
|
|
from toolkit.basic import flush
|
|
from toolkit.print import print_acc
|
|
|
|
AITK_Status = Literal["running", "stopped", "error", "completed"]
|
|
|
|
|
|
class DiffusionTrainer(SDTrainer):
|
|
def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
|
|
super(DiffusionTrainer, self).__init__(process_id, job, config, **kwargs)
|
|
self.sqlite_db_path = self.config.get("sqlite_db_path", "./aitk_db.db")
|
|
self.job_id = os.environ.get("AITK_JOB_ID", None)
|
|
self.job_id = self.job_id.strip() if self.job_id is not None else None
|
|
self.is_ui_trainer = True
|
|
if not os.path.exists(self.sqlite_db_path):
|
|
self.is_ui_trainer = False
|
|
else:
|
|
print(f"Using SQLite database at {self.sqlite_db_path}")
|
|
if self.job_id is None:
|
|
self.is_ui_trainer = False
|
|
else:
|
|
print(f"Job ID: \"{self.job_id}\"")
|
|
|
|
if self.is_ui_trainer:
|
|
self.is_stopping = False
|
|
# Create a thread pool for database operations
|
|
self.thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=1)
|
|
# Track all async tasks
|
|
self._async_tasks = []
|
|
# Initialize the status
|
|
self._run_async_operation(self._update_status("running", "Starting"))
|
|
self._stop_watcher_started = False
|
|
if os.name == "nt":
|
|
# On Windows the stop route cannot send us SIGINT from outside
|
|
# (no console to deliver a Ctrl+C to), so watch the stop flag
|
|
# and raise the interrupt from inside. On Linux the route
|
|
# sends a real SIGINT to the pid and this is unnecessary.
|
|
self.start_stop_watcher(interval_sec=2.0)
|
|
|
|
def start_stop_watcher(self, interval_sec: float = 5.0):
|
|
"""
|
|
Start a daemon thread that periodically checks should_stop()
|
|
and terminates the process immediately when triggered.
|
|
"""
|
|
if not self.is_ui_trainer:
|
|
return
|
|
if getattr(self, "_stop_watcher_started", False):
|
|
return
|
|
self._stop_watcher_started = True
|
|
t = threading.Thread(
|
|
target=self._stop_watcher_thread, args=(interval_sec,), daemon=True
|
|
)
|
|
t.start()
|
|
|
|
def _stop_watcher_thread(self, interval_sec: float):
|
|
while True:
|
|
try:
|
|
if self.should_stop():
|
|
if self.is_stopping:
|
|
# maybe_stop() already started the graceful shutdown;
|
|
# a second interrupt would only break its cleanup.
|
|
return
|
|
print("")
|
|
print("****************************************************")
|
|
print(" Stop signal received; terminating process. ")
|
|
print("****************************************************")
|
|
# Deliver a real KeyboardInterrupt to the main thread so
|
|
# on_error runs the normal shutdown (final DB write, last
|
|
# log). os.kill(pid, SIGINT) must not be used here: on
|
|
# Windows it is TerminateProcess and kills us instantly.
|
|
# Leave the thread pool alone -- on_error still needs it.
|
|
signal.raise_signal(signal.SIGINT)
|
|
return
|
|
time.sleep(interval_sec)
|
|
except Exception:
|
|
time.sleep(interval_sec)
|
|
|
|
def _run_async_operation(self, coro):
|
|
"""Helper method to run an async coroutine and track the task."""
|
|
try:
|
|
loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
# No event loop exists, create a new one
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
|
|
# Create a task and track it
|
|
if loop.is_running():
|
|
task = asyncio.run_coroutine_threadsafe(coro, loop)
|
|
self._async_tasks.append(asyncio.wrap_future(task))
|
|
else:
|
|
task = loop.create_task(coro)
|
|
self._async_tasks.append(task)
|
|
loop.run_until_complete(task)
|
|
|
|
async def _execute_db_operation(self, operation_func):
|
|
"""Execute a database operation in a separate thread with retry on lock."""
|
|
loop = asyncio.get_event_loop()
|
|
return await loop.run_in_executor(
|
|
self.thread_pool, lambda: self._retry_db_operation(operation_func)
|
|
)
|
|
|
|
def _db_connect(self):
|
|
"""Create a new connection for each operation to avoid locking."""
|
|
conn = sqlite3.connect(self.sqlite_db_path, timeout=30.0)
|
|
conn.isolation_level = None # Enable autocommit mode
|
|
return conn
|
|
|
|
def _retry_db_operation(self, operation_func, max_retries=3, base_delay=2.0):
|
|
"""Retry a database operation with exponential backoff on lock errors."""
|
|
last_error = None
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
return operation_func()
|
|
except sqlite3.OperationalError as e:
|
|
if "database is locked" in str(e):
|
|
last_error = e
|
|
if attempt < max_retries:
|
|
delay = base_delay * (2 ** attempt) # 2s, 4s, 8s
|
|
print(f"[AITK] Database locked (attempt {attempt + 1}/{max_retries + 1}), retrying in {delay:.1f}s...")
|
|
time.sleep(delay)
|
|
else:
|
|
print(f"[AITK] Database locked after {max_retries + 1} attempts, giving up.")
|
|
else:
|
|
raise
|
|
raise last_error
|
|
|
|
def should_stop(self):
|
|
if not self.is_ui_trainer:
|
|
return False
|
|
def _check_stop():
|
|
with self._db_connect() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"SELECT stop FROM Job WHERE id = ?", (self.job_id,))
|
|
stop = cursor.fetchone()
|
|
return False if stop is None else stop[0] == 1
|
|
|
|
return self._retry_db_operation(_check_stop)
|
|
|
|
def should_return_to_queue(self):
|
|
if not self.is_ui_trainer:
|
|
return False
|
|
def _check_return_to_queue():
|
|
with self._db_connect() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"SELECT return_to_queue FROM Job WHERE id = ?", (self.job_id,))
|
|
return_to_queue = cursor.fetchone()
|
|
return False if return_to_queue is None else return_to_queue[0] == 1
|
|
|
|
return self._retry_db_operation(_check_return_to_queue)
|
|
|
|
def maybe_stop(self):
|
|
if not self.is_ui_trainer:
|
|
return
|
|
if self.should_stop():
|
|
self._run_async_operation(
|
|
self._update_status("stopped", "Job stopped"))
|
|
self.is_stopping = True
|
|
raise Exception("Job stopped")
|
|
if self.should_return_to_queue():
|
|
self._run_async_operation(
|
|
self._update_status("queued", "Job queued"))
|
|
self.is_stopping = True
|
|
raise Exception("Job returning to queue")
|
|
|
|
def should_save(self):
|
|
if not self.is_ui_trainer:
|
|
return False
|
|
def _check_save():
|
|
with self._db_connect() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"SELECT save_now FROM Job WHERE id = ?", (self.job_id,))
|
|
save_now = cursor.fetchone()
|
|
return False if save_now is None else save_now[0] == 1
|
|
|
|
return self._retry_db_operation(_check_save)
|
|
|
|
def maybe_save(self):
|
|
if not self.is_ui_trainer:
|
|
return
|
|
if self.should_save():
|
|
self.update_db_key("save_now", 0)
|
|
if self.progress_bar is not None:
|
|
self.progress_bar.pause()
|
|
print_acc(f"\nSaving at step {self.step_num}")
|
|
# clear any grads
|
|
self.optimizer.zero_grad()
|
|
self.save(self.step_num)
|
|
self.ensure_params_requires_grad()
|
|
flush()
|
|
if self.progress_bar is not None:
|
|
self.progress_bar.unpause()
|
|
|
|
def should_sample(self):
|
|
if not self.is_ui_trainer:
|
|
return False
|
|
def _check_sample():
|
|
with self._db_connect() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute(
|
|
"SELECT sample_now FROM Job WHERE id = ?", (self.job_id,))
|
|
sample_now = cursor.fetchone()
|
|
return False if sample_now is None else sample_now[0] == 1
|
|
|
|
return self._retry_db_operation(_check_sample)
|
|
|
|
def maybe_sample(self):
|
|
if not self.is_ui_trainer:
|
|
return
|
|
if self.should_sample():
|
|
self.update_db_key("sample_now", 0)
|
|
if self.progress_bar is not None:
|
|
self.progress_bar.pause()
|
|
print_acc(f"\nSampling at step {self.step_num}")
|
|
# clear any grads
|
|
self.optimizer.zero_grad()
|
|
if self.train_config.free_u:
|
|
self.sd.pipeline.disable_freeu()
|
|
self.sample(self.step_num)
|
|
if self.train_config.unload_text_encoder:
|
|
# make sure the text encoder is unloaded
|
|
self.sd.text_encoder_to('cpu')
|
|
self.ensure_params_requires_grad()
|
|
flush()
|
|
if self.progress_bar is not None:
|
|
self.progress_bar.unpause()
|
|
|
|
async def _update_key(self, key, value):
|
|
if not self.accelerator.is_main_process:
|
|
return
|
|
|
|
def _do_update():
|
|
with self._db_connect() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
# Convert the value to string if it's not already
|
|
if isinstance(value, str):
|
|
value_to_insert = value
|
|
else:
|
|
value_to_insert = str(value)
|
|
|
|
# Use parameterized query for both the column name and value
|
|
update_query = f"UPDATE Job SET {key} = ? WHERE id = ?"
|
|
cursor.execute(
|
|
update_query, (value_to_insert, self.job_id))
|
|
finally:
|
|
cursor.execute("COMMIT")
|
|
|
|
await self._execute_db_operation(_do_update)
|
|
|
|
def update_step(self):
|
|
"""Non-blocking update of the step count."""
|
|
if self.accelerator.is_main_process and self.is_ui_trainer:
|
|
self._run_async_operation(self._update_key("step", self.step_num))
|
|
|
|
def update_db_key(self, key, value):
|
|
"""Non-blocking update a key in the database."""
|
|
if self.accelerator.is_main_process and self.is_ui_trainer:
|
|
self._run_async_operation(self._update_key(key, value))
|
|
|
|
async def _update_status(self, status: AITK_Status, info: Optional[str] = None):
|
|
if not self.accelerator.is_main_process or not self.is_ui_trainer:
|
|
return
|
|
|
|
def _do_update():
|
|
with self._db_connect() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
if info is not None:
|
|
cursor.execute(
|
|
"UPDATE Job SET status = ?, info = ? WHERE id = ?",
|
|
(status, info, self.job_id)
|
|
)
|
|
else:
|
|
cursor.execute(
|
|
"UPDATE Job SET status = ? WHERE id = ?",
|
|
(status, self.job_id)
|
|
)
|
|
finally:
|
|
cursor.execute("COMMIT")
|
|
|
|
await self._execute_db_operation(_do_update)
|
|
|
|
def update_status(self, status: AITK_Status, info: Optional[str] = None):
|
|
"""Non-blocking update of status."""
|
|
if self.accelerator.is_main_process and self.is_ui_trainer:
|
|
self._run_async_operation(self._update_status(status, info))
|
|
|
|
async def wait_for_all_async(self):
|
|
"""Wait for all tracked async operations to complete."""
|
|
if not self._async_tasks:
|
|
return
|
|
|
|
try:
|
|
await asyncio.gather(*self._async_tasks)
|
|
except Exception as e:
|
|
pass
|
|
finally:
|
|
# Clear the task list after completion
|
|
self._async_tasks.clear()
|
|
|
|
def on_error(self, e: Exception):
|
|
super(DiffusionTrainer, self).on_error(e)
|
|
if self.is_ui_trainer:
|
|
try:
|
|
if isinstance(e, KeyboardInterrupt):
|
|
# SIGINT (UI stop button or ctrl+c) is a stop, not an error
|
|
self.is_stopping = True
|
|
progress_bar = getattr(self, "progress_bar", None)
|
|
if progress_bar is not None:
|
|
# silence the bar so tqdm doesn't repaint it at interpreter exit
|
|
progress_bar.disable = True
|
|
progress_bar.close()
|
|
if self.accelerator.is_main_process:
|
|
self.update_status("stopped", "Job stopped")
|
|
elif self.accelerator.is_main_process and not self.is_stopping:
|
|
self.update_status("error", str(e))
|
|
self.update_db_key("step", self.last_save_step)
|
|
asyncio.run(self.wait_for_all_async())
|
|
except Exception as db_err:
|
|
print(f"[AITK] Warning: failed to update DB during error handling: {db_err}")
|
|
finally:
|
|
self.thread_pool.shutdown(wait=True)
|
|
|
|
def handle_timing_print_hook(self, timing_dict):
|
|
if "train_loop" not in timing_dict:
|
|
print("train_loop not found in timing_dict", timing_dict)
|
|
return
|
|
seconds_per_iter = timing_dict["train_loop"]
|
|
# determine iter/sec or sec/iter
|
|
if seconds_per_iter < 1:
|
|
iters_per_sec = 1 / seconds_per_iter
|
|
self.update_db_key("speed_string", f"{iters_per_sec:.2f} iter/sec")
|
|
else:
|
|
self.update_db_key(
|
|
"speed_string", f"{seconds_per_iter:.2f} sec/iter")
|
|
|
|
def done_hook(self):
|
|
super(DiffusionTrainer, self).done_hook()
|
|
if self.is_ui_trainer:
|
|
self.update_status("completed", "Training completed")
|
|
# Wait for all async operations to finish before shutting down
|
|
asyncio.run(self.wait_for_all_async())
|
|
self.thread_pool.shutdown(wait=True)
|
|
|
|
def end_step_hook(self):
|
|
super(DiffusionTrainer, self).end_step_hook()
|
|
if self.is_ui_trainer:
|
|
self.update_step()
|
|
self.maybe_stop()
|
|
self.maybe_save()
|
|
self.maybe_sample()
|
|
|
|
def hook_before_model_load(self):
|
|
super().hook_before_model_load()
|
|
if self.is_ui_trainer:
|
|
self.maybe_stop()
|
|
self.update_status("running", "Loading model")
|
|
|
|
def before_dataset_load(self):
|
|
super().before_dataset_load()
|
|
if self.is_ui_trainer:
|
|
self.maybe_stop()
|
|
self.update_status("running", "Loading dataset")
|
|
|
|
def hook_before_train_loop(self):
|
|
super().hook_before_train_loop()
|
|
if self.is_ui_trainer:
|
|
self.maybe_stop()
|
|
self.update_step()
|
|
self.update_status("running", "Training")
|
|
self.timer.add_after_print_hook(self.handle_timing_print_hook)
|
|
|
|
def status_update_hook_func(self, string):
|
|
self.update_status("running", string)
|
|
|
|
def hook_after_sd_init_before_load(self):
|
|
super().hook_after_sd_init_before_load()
|
|
if self.is_ui_trainer:
|
|
self.maybe_stop()
|
|
self.sd.add_status_update_hook(self.status_update_hook_func)
|
|
|
|
def sample_step_hook(self, img_num, total_imgs):
|
|
super().sample_step_hook(img_num, total_imgs)
|
|
if self.is_ui_trainer:
|
|
self.maybe_stop()
|
|
self.update_status(
|
|
"running", f"Generating images - {img_num + 1}/{total_imgs}")
|
|
|
|
def sample(self, step=None, is_first=False):
|
|
self.maybe_stop()
|
|
total_imgs = len(self.sample_config.prompts)
|
|
self.update_status("running", f"Generating images - 0/{total_imgs}")
|
|
super().sample(step, is_first)
|
|
self.maybe_stop()
|
|
self.update_status("running", "Training")
|
|
|
|
def save(self, step=None):
|
|
self.maybe_stop()
|
|
self.update_status("running", "Saving model")
|
|
super().save(step)
|
|
self.maybe_stop()
|
|
self.update_status("running", "Training")
|