77 lines
3.1 KiB
Python
77 lines
3.1 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
正式训练脚本：脸部 LoRA（Qwen-Image Edit-2511 / 16G 显存 / 32G 内存优化）
|
||
用法: python 训练脚本.py
|
||
前置: output/train_dataset/ 里有 img_001.jpg + img_001.txt 打标
|
||
"""
|
||
import json
|
||
import os
|
||
import subprocess
|
||
import sys
|
||
|
||
os.environ["PYTHONUTF8"] = "1"
|
||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||
|
||
VENV = r"M:\AI\sd\musubi-tuner\.venv\Scripts"
|
||
PROJ = r"D:\F\NewI\opencode\daily-workspace\projects\脸部LoRA训练-Qwen-Image"
|
||
MODELS = r"M:\AI\sd\models\qwen-edit-2511"
|
||
|
||
# 动态目录配置（GUI 训练 Tab 写入 config/train_config.json）
|
||
_cfg = {}
|
||
_cfg_path = os.path.join(PROJ, "config", "train_config.json")
|
||
if os.path.exists(_cfg_path):
|
||
try:
|
||
_cfg = json.load(open(_cfg_path, encoding="utf-8-sig")) # utf-8-sig 处理 BOM
|
||
except Exception as _e:
|
||
print(f"[WARN] train_config.json 读取失败: {_e}")
|
||
_cfg = {}
|
||
TRAIN_DATASET = _cfg.get("train_dataset", os.path.join(PROJ, "output", "train_dataset"))
|
||
CKPT_OUT = _cfg.get("output_dir", os.path.join(PROJ, "output", "checkpoints"))
|
||
|
||
# 动态生成数据集配置（image_directory 指向所选训练集）
|
||
_dataset_toml = os.path.join(PROJ, "config", "dataset_active.toml")
|
||
_toml_body = f"""# 由 GUI 训练 Tab 动态生成（勿手改）
|
||
[general]
|
||
resolution = 1024
|
||
caption_extension = ".txt"
|
||
batch_size = 1
|
||
enable_bucket = true
|
||
bucket_no_upscale = false
|
||
|
||
[[datasets]]
|
||
image_directory = "{TRAIN_DATASET.replace(chr(92), '/')}"
|
||
num_repeats = 1
|
||
"""
|
||
with open(_dataset_toml, "w", encoding="utf-8") as _f:
|
||
_f.write(_toml_body)
|
||
|
||
cmd = [
|
||
os.path.join(VENV, "accelerate.exe"), "launch",
|
||
"--num_cpu_threads_per_process", "1", "--mixed_precision", "bf16",
|
||
r"D:\AI\sd\musubi-tuner\src\musubi_tuner\qwen_image_train_network.py",
|
||
"--dit", os.path.join(MODELS, "transformer", "diffusion_pytorch_model-00001-of-00005.safetensors"),
|
||
"--vae", os.path.join(MODELS, "diffusion_pytorch_model.safetensors"),
|
||
"--text_encoder", os.path.join(MODELS, "text_encoder", "model-00001-of-00004.safetensors"),
|
||
"--model_version", "edit-2511",
|
||
"--dataset_config", _dataset_toml,
|
||
"--sdpa", "--mixed_precision", "bf16",
|
||
"--timestep_sampling", "shift", "--weighting_scheme", "none", "--discrete_flow_shift", "2.2",
|
||
"--optimizer_type", "adamw8bit", "--learning_rate", "5e-5",
|
||
"--gradient_checkpointing",
|
||
"--network_module", "networks.lora_qwen_image", "--network_dim", "16",
|
||
# 续炼：加载旧 LoRA 000060（脸已成形、细节未过度固化），用重调后的 v7 数据增量修正标签
|
||
"--network_weights", r"M:\AI\sd\novelai-webui-aki-v3-r\models\Lora\v6\myface_lora-000060.safetensors",
|
||
"--fp8_base", "--fp8_scaled", "--fp8_vl", "--blocks_to_swap", "24",
|
||
"--max_train_epochs", "50", "--save_every_n_epochs", "10", "--seed", "42",
|
||
"--sample_prompts", os.path.join(PROJ, "config", "sample_prompts.txt"),
|
||
"--sample_every_n_epochs", "10",
|
||
"--output_dir", CKPT_OUT,
|
||
"--output_name", "myface_v7",
|
||
]
|
||
print("训练集:", TRAIN_DATASET)
|
||
print("输出目录:", CKPT_OUT)
|
||
print("CMD:", " ".join(cmd))
|
||
r = subprocess.run(cmd)
|
||
print("EXIT CODE:", r.returncode)
|
||
sys.exit(r.returncode)
|