# FR 장비 아이콘 스타일 LoRA 학습 — ComfyUI 코어 TrainLoraNode (BT25 추가 지시 2026-09-25 "FR에서 사용한 아이템을 학습시킨 로라로 테스트")
# 1) 데이터셋: 기존 장비 아이콘 32장 중 금색 'S' 오버레이(ItemClass 1) 6장 제외 → 26장을 검정 배경 512² 로 합성, 원본·좌우반전·축소(0.85) 3벌 = 78장 + 캡션 .txt
#    → C:\Comfyui\input\fr_icons_v1  (LoadImageTextDataSetFromFolder 가 input 하위 폴더만 읽음)
# 2) 학습: difconsistency(SD1.5) · LoRA rank 16 · lr 4e-4 · batch 2 · N step (인수, 기본 500) → output/loras → models/loras/FR_icons_v1.safetensors 복사
# 사용: python comfy_lora_train.py [steps]
import json, os, shutil, sys, time, uuid, glob, urllib.request, urllib.parse
from PIL import Image, ImageOps

HOST = "http://127.0.0.1:8000"
SRC = r"C:\BurningTimes\FirstRising\Assets\GameAssets\UI\EquipmentItemIcon"
DS_NAME = "fr_icons_v1"
DS = os.path.join(r"C:\Comfyui\input", DS_NAME)
LORA_OUT = r"C:\Comfyui\output\loras"
LORA_DST = os.path.join(r"C:\Comfyui\models\loras", "FR_icons_v1.safetensors")
CKPT = "difconsistencyCheckpoint_v10.safetensors"
TRIGGER = "frdicon"
RES = 512
STEPS = int(sys.argv[1]) if len(sys.argv) > 1 else 500
LR, RANK, BATCH, SEED = 4e-4, 16, 2, 7

# 캡션은 '내용(형태·색)'만 적고 화풍은 트리거 토큰에 묶는다. 실제 아이콘을 보고 작성(fr_icons_sheet32.png).
CAP = {
    1015: "a war hammer with a dark wooden handle and a cracked iron head with orange sparks",
    1029: "a simple curved wooden bow with bone colored limbs",
    1043: "a dark steel dagger with an orange glowing gem in the hilt",
    1057: "a dark iron flanged mace with orange glowing spikes",
    1071: "a long thin dark rod dagger with an orange glowing tip",
    1435: "a plain iron short sword with a brown leather grip",
    1099: "a round dark amulet with a cyan glowing gem",
    1113: "a silver shield shaped pendant with an icy blue gem",
    1127: "a dark spiked diamond pendant with a gold and orange star gem",
    1141: "a round red medallion with a skull",
    1169: "a pair of dark leather gloves with cyan glowing rune lines",
    1183: "a pair of brown leather gloves with engraved patterns",
    1197: "a pair of dark leather gloves with orange glowing accents",
    1211: "a pair of dark iron gauntlets with brown leather",
    1239: "a dark plate armor with a fur collar and a cyan gem",
    1253: "an ornate silver grey plate armor",
    1267: "a grey steel plate armor with engraved patterns",
    1281: "a chainmail armor with spiked shoulder pauldrons",
    1309: "a brown leather belt with cyan gems and hanging chains",
    1323: "a dark cloth waist wrap with a gold plate and a dagger",
    1337: "a dark leather belt with a red rope and iron buckles",
    1351: "a dark leather belt with a skull buckle and small cyan gems",
    1379: "a pair of brown leather lace up boots with cyan gems",
    1393: "a pair of brown leather boots with buckles",
    1407: "a pair of blue and brown leather high boots with straps",
    1421: "a pair of black plate high boots with straps",
}
EXCLUDED_S = [1001, 1085, 1155, 1225, 1295, 1365]  # 금색 S 오버레이 — 화풍에 섞이지 않게 제외


def build_dataset():
    if os.path.isdir(DS):
        shutil.rmtree(DS)
    os.makedirs(DS)
    n = 0
    for gid, desc in CAP.items():
        im = Image.open(os.path.join(SRC, f"{gid}.png")).convert("RGBA").resize((RES, RES), Image.LANCZOS)
        variants = {"a": im, "b": ImageOps.mirror(im)}
        small = im.resize((int(RES * 0.85), int(RES * 0.85)), Image.LANCZOS)
        c = Image.new("RGBA", (RES, RES), (0, 0, 0, 0)); c.paste(small, ((RES - small.width) // 2, (RES - small.height) // 2), small)
        variants["c"] = c
        for tag, v in variants.items():
            bg = Image.new("RGB", (RES, RES), (0, 0, 0)); bg.paste(v, (0, 0), v)
            bg.save(os.path.join(DS, f"{gid}_{tag}.png"))
            with open(os.path.join(DS, f"{gid}_{tag}.txt"), "w", encoding="utf-8") as f:
                f.write(f"{TRIGGER}, game equipment item icon, {desc}, black background")
            n += 1
    print(f"dataset {DS}: {n} images ({len(CAP)} icons x 3)", flush=True)
    return n


def get(path):
    return json.load(urllib.request.urlopen(HOST + path, timeout=120))


def vram():
    try:
        d = get("/system_stats")["devices"][0]
        return f"VRAM free {d['vram_free'] / 2**30:.1f}/{d['vram_total'] / 2**30:.1f} GB"
    except Exception as e:
        return f"vram? {e}"


def train():
    wf = {
        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": CKPT}},
        "2": {"class_type": "LoadImageTextDataSetFromFolder", "inputs": {"folder": DS_NAME}},
        "3": {"class_type": "MakeTrainingDataset", "inputs": {"images": ["2", 0], "vae": ["1", 2], "clip": ["1", 1], "texts": ["2", 1]}},
        "4": {"class_type": "TrainLoraNode", "inputs": {
            "model": ["1", 0], "latents": ["3", 0], "positive": ["3", 1],
            "batch_size": BATCH, "grad_accumulation_steps": 1, "steps": STEPS, "learning_rate": LR, "rank": RANK,
            "optimizer": "AdamW", "loss_function": "MSE", "seed": SEED, "training_dtype": "bf16", "lora_dtype": "bf16",
            "quantized_backward": False, "algorithm": "LoRA", "gradient_checkpointing": True, "checkpoint_depth": 1,
            "offloading": False, "existing_lora": "[None]", "bucket_mode": False, "bypass_mode": False}},
        "5": {"class_type": "SaveLoRA", "inputs": {"lora": ["4", 0], "prefix": "loras/FR_icons_v1", "steps": ["4", 2]}},
        "6": {"class_type": "LossGraphNode", "inputs": {"loss": ["4", 1], "filename_prefix": "FR_icons_v1_loss"}},
    }
    req = urllib.request.Request(HOST + "/prompt", data=json.dumps({"prompt": wf, "client_id": str(uuid.uuid4())}).encode(), headers={"Content-Type": "application/json"})
    try:
        r = json.load(urllib.request.urlopen(req, timeout=60))
    except urllib.error.HTTPError as e:
        print("PROMPT REJECTED:", e.read().decode(errors="replace")[:1500]); return None
    if r.get("node_errors"):
        print("NODE ERRORS:", json.dumps(r["node_errors"], ensure_ascii=False)[:1500]); return None
    pid = r["prompt_id"]
    print(f"queued {pid} · steps {STEPS} · batch {BATCH} · lr {LR} · rank {RANK} · {vram()}", flush=True)
    t0 = time.time()
    while time.time() - t0 < 5400:
        time.sleep(15)
        h = get(f"/history/{pid}")
        if pid in h:
            st = h[pid].get("status", {})
            if st.get("status_str") == "error":
                for m in st.get("messages", []):
                    if m[0] == "execution_error":
                        print("TRAIN ERROR:", m[1].get("node_type"), m[1].get("exception_message", "")[:1500])
                        print("\n".join(m[1].get("traceback", [])[-6:])[:2000])
                return None
            print(f"done in {time.time() - t0:.0f}s · {vram()}", flush=True)
            return h[pid]["outputs"]
        print(f"  training… {time.time() - t0:.0f}s elapsed · {vram()}", flush=True)
    print("TIMEOUT"); return None


if __name__ == "__main__":
    build_dataset()
    print(vram(), flush=True)
    out = train()
    if out is None:
        sys.exit(1)
    files = sorted(glob.glob(os.path.join(LORA_OUT, "FR_icons_v1*.safetensors")), key=os.path.getmtime)
    if not files:
        print("no lora file in", LORA_OUT); sys.exit(1)
    shutil.copy2(files[-1], LORA_DST)
    print("LORA:", files[-1], "->", LORA_DST, f"{os.path.getsize(LORA_DST) / 2**20:.1f} MB")
    for nid, o in out.items():
        for im in o.get("images", []):
            print("output image:", im)