125 lines
5.2 KiB
Python
125 lines
5.2 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""直接调 kohya train_network.py 真正并行训练 LoRA
|
||
|
||
绕过 photo_service 的串行阻塞，用 subprocess 直接起 train_network.py 进程，
|
||
通过 ThreadPoolExecutor 控制并发数。
|
||
|
||
用法:
|
||
python /home/ubuntu/change_hair/train_lora_parallel.py
|
||
|
||
前提：step1 训练数据已准备好（data/train_material/<hid>/images/1_hairstyle/*.png）
|
||
"""
|
||
import os
|
||
import sys
|
||
import time
|
||
import subprocess
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
|
||
TRAIN_DIR = "/home/ubuntu/change_hair/data/train_material"
|
||
KOHYA_WORKDIR = "/home/ubuntu/change_hair/project/kohya_ss_home/kohya_ss"
|
||
BASE_MODEL = "/home/ubuntu/change_hair/project/onediff/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors"
|
||
KOHYA_ACCEL = "/home/ubuntu/miniconda3/envs/kohya/bin/accelerate"
|
||
LOG_DIR = "/home/ubuntu/change_hair/project/logs"
|
||
|
||
# 本次要训练的发型
|
||
HAIRSTYLES = ["chang_tuoyuan", "chang_zhixian", "chang_huaban", "chang_xinxing"]
|
||
# chang_bolang 已完成，跳过
|
||
PARALLEL = 2
|
||
|
||
|
||
def log(msg):
|
||
print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)
|
||
|
||
|
||
def train_one(hair_id):
|
||
"""训练单个发型的 LoRA，返回 (hair_id, success, log_path)"""
|
||
images_dir = os.path.join(TRAIN_DIR, hair_id, "images")
|
||
model_dir = os.path.join(TRAIN_DIR, hair_id, "model")
|
||
os.makedirs(model_dir, exist_ok=True)
|
||
|
||
# 准备 sample prompt 文件
|
||
sample_dir = os.path.join(model_dir, "sample")
|
||
os.makedirs(sample_dir, exist_ok=True)
|
||
sample_txt = os.path.join(sample_dir, "prompt.txt")
|
||
with open(sample_txt, "w") as f:
|
||
f.write('titor hairstyle, easyphoto, faceless, no human, white background, simple background, '
|
||
' --n low quality, worst quality, bad anatomy, bad composition, poor, low effort --h 768 '
|
||
'--w 768 --s 30 --l 7')
|
||
|
||
log_path = os.path.join(LOG_DIR, f"train_{hair_id}.log")
|
||
|
||
cmd = [
|
||
"bash", "-c",
|
||
f'cd {KOHYA_WORKDIR} && '
|
||
f'HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 CUDA_VISIBLE_DEVICES=0 '
|
||
f'{KOHYA_ACCEL} launch --num_cpu_threads_per_process=2 "./train_network.py" --enable_bucket '
|
||
f'--min_bucket_reso=256 --max_bucket_reso=2048 --pretrained_model_name_or_path="{BASE_MODEL}" '
|
||
f'--train_data_dir={images_dir} --resolution="2000,2000" '
|
||
f'--output_dir={model_dir} '
|
||
'--network_alpha="64" --save_model_as=safetensors --network_module=networks.lora --text_encoder_lr=5e-05 '
|
||
'--unet_lr=0.0001 --network_dim=128 --output_name="hairstyle_hd_lora" --lr_scheduler_num_cycles="20" '
|
||
'--no_half_vae --learning_rate="0.0001" --lr_scheduler="cosine" --lr_warmup_steps="650" --train_batch_size="1" '
|
||
'--max_train_steps="1500" --save_every_n_epochs="100" --mixed_precision="fp16" --save_precision="fp16" '
|
||
'--caption_extension=".txt" --sample_sampler=ddim '
|
||
f'--sample_prompts={sample_txt} --sample_every_n_epochs="1000" '
|
||
'--seed="1234" --cache_latents --optimizer_type="AdamW" --max_data_loader_n_workers="0" --bucket_reso_steps=64 '
|
||
'--sdpa --bucket_no_upscale --noise_offset=0.0'
|
||
]
|
||
|
||
log(f" ▶ {hair_id} 开始训练，日志: {log_path}")
|
||
t0 = time.time()
|
||
try:
|
||
with open(log_path, "w") as lf:
|
||
proc = subprocess.run(cmd, stdout=lf, stderr=subprocess.STDOUT, timeout=1800)
|
||
ok = proc.returncode == 0
|
||
lora = os.path.join(model_dir, "hairstyle_hd_lora.safetensors")
|
||
ok = ok and os.path.exists(lora) and os.path.getsize(lora) > 100000000
|
||
elapsed = time.time() - t0
|
||
sz = os.path.getsize(lora) // 1024 // 1024 if os.path.exists(lora) else 0
|
||
log(f" {'✓' if ok else '✗'} {hair_id} {'训练完成' if ok else '训练失败'} ({sz}MB, {elapsed:.0f}s)")
|
||
return (hair_id, ok, log_path)
|
||
except subprocess.TimeoutExpired:
|
||
log(f" ✗ {hair_id} 训练超时")
|
||
return (hair_id, False, log_path)
|
||
except Exception as e:
|
||
log(f" ✗ {hair_id} 训练异常: {e}")
|
||
return (hair_id, False, log_path)
|
||
|
||
|
||
def main():
|
||
t0 = time.time()
|
||
log(f"并行 LoRA 训练: {HAIRSTYLES} (并发={PARALLEL})")
|
||
|
||
results = {}
|
||
with ThreadPoolExecutor(max_workers=PARALLEL) as ex:
|
||
futures = {ex.submit(train_one, hid): hid for hid in HAIRSTYLES}
|
||
for fut in as_completed(futures):
|
||
hid = futures[fut]
|
||
try:
|
||
hair_id, ok, log_path = fut.result()
|
||
results[hair_id] = ok
|
||
except Exception as e:
|
||
log(f" ✗ {hid} 异常: {e}")
|
||
results[hid] = False
|
||
|
||
log(f"\n{'='*60}")
|
||
log(f"全部完成，耗时 {time.time()-t0:.0f}s")
|
||
for hid, ok in results.items():
|
||
lora = os.path.join(TRAIN_DIR, hid, "model", "hairstyle_hd_lora.safetensors")
|
||
sz = os.path.getsize(lora) // 1024 // 1024 if os.path.exists(lora) else 0
|
||
log(f" {'✓' if ok else '✗'} {hid} ({sz}MB)")
|
||
success = [h for h, ok in results.items() if ok]
|
||
log(f"成功: {success}")
|
||
|
||
# 写完成清单
|
||
done_file = os.path.join(LOG_DIR, "train_lora_done.json")
|
||
import json
|
||
with open(done_file, "w") as f:
|
||
json.dump({"completed": success, "all": HAIRSTYLES}, f)
|
||
log(f"完成清单: {done_file}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|