# -*- 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)