141 lines
7.5 KiB
Python
141 lines
7.5 KiB
Python
# 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)
|