144 lines
8.5 KiB
Python
144 lines
8.5 KiB
Python
|
# -*- coding: utf-8 -*-
|
||
|
|
# 2026-10-02 PD 지시(자율 반복) — 합격본(dataset.json)으로 FR 전용 아이콘 LoRA 를 학습한다. 학습을 거듭할수록 합격본이 늘고, 글자만으로 원하는 풍이 나오게 하는 것이 목표.
|
||
|
|
# python comfy_anfg_train.py <버전 이름 예: v1> [steps=800] [PD 합격본 반복 배수=2]
|
||
|
|
# 구성: 기반 mobaGameEquipment_v10 위에 ANFG LoRA(0.8)를 얹은 모델을 학습 대상으로 넣는다 → 새 LoRA 는 "ANFG 풍에서 FR 합격본 쪽으로의 차이"만 배운다.
|
||
|
|
# (쓸 때도 ANFG LoRA 와 함께 건다.) 캡션 = "fricon, game icon institute, game icon, ahfg, <묘사>, black background".
|
||
|
|
# 학습 노드 구성 · 입력 이름은 comfy_moba_lora_train.py(2026-10-02) · comfy_lora_train.py(2026-09-25 실측)와 같다.
|
||
|
|
import io, json, os, shutil, sys, time, uuid, glob, urllib.request
|
||
|
|
from PIL import Image, ImageOps
|
||
|
|
import icon_dataset
|
||
|
|
|
||
|
|
HOST = "http://127.0.0.1:8000"
|
||
|
|
VER = sys.argv[1]
|
||
|
|
STEPS = int(sys.argv[2]) if len(sys.argv) > 2 else 800
|
||
|
|
PD_REPEAT = int(sys.argv[3]) if len(sys.argv) > 3 else 2
|
||
|
|
NAME = "FR_anfg_%s" % VER
|
||
|
|
DS_NAME = "fr_anfg_%s" % VER
|
||
|
|
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", NAME + ".safetensors")
|
||
|
|
CKPT = "mobaGameEquipment_v10.safetensors"
|
||
|
|
BASE_LORA, BASE_LORA_STRENGTH = "giI_anfg_v3_000018.safetensors", 0.8
|
||
|
|
TRIGGER = "fricon, game icon institute, game icon, ahfg"
|
||
|
|
RES = 512
|
||
|
|
LR, RANK, BATCH, SEED = 2e-4, 16, 2, 11
|
||
|
|
|
||
|
|
|
||
|
|
# a11 검토(v3): 글자만 생성이 "화면 가득한 금속 원반"으로 무너졌다. 원인으로 본 것 — ① 학습 그림이 캔버스의 82% 를 채워 여백을 못 배웠다
|
||
|
|
# ② 갑옷 · 목걸이 · 장화가 많아 둥근 덩어리로 쏠렸다 ③ PD 합격본 윤곽에서 나온 복제(genpd)가 섞였다 ④ 캡션에 구도 낱말이 없었다.
|
||
|
|
# v4 부터: 그림을 줄여 여백을 넣고(MARGIN) · 종류당 장수를 제한하고(CAP) · genpd 유래를 빼고 · 캡션 앞에 생성 때와 같은 구도 낱말을 붙인다.
|
||
|
|
MARGIN = float(os.environ.get("ANFG_TRAIN_MARGIN", "0.84")) # 82% 로 맞춘 그림을 다시 이 비율로 줄인다 → 물체가 캔버스의 약 69%
|
||
|
|
CAP = int(os.environ.get("ANFG_TRAIN_CAP", "8")) # 종류당 최대 장수(PD 합격본은 제한 밖)
|
||
|
|
COMPOSE = "single item, centered, full object in frame, wide margin"
|
||
|
|
|
||
|
|
|
||
|
|
def kind_of(src):
|
||
|
|
"""원본 파일 이름(<회차>_<대상>_s..)에서 종류를 읽는다 — 대상이 '종류-길잡이' 꼴이면 앞부분"""
|
||
|
|
import re
|
||
|
|
m = re.match(r"^[a-z0-9]+_(.+)_s\d+_\d+\.png$", src)
|
||
|
|
return m.group(1).split("-")[0] if m else "?"
|
||
|
|
|
||
|
|
|
||
|
|
def select(data):
|
||
|
|
"""학습에 쓸 항목: PD 합격본 전부 + 나머지는 genpd 유래를 빼고 종류당 CAP 장(뒤 회차일수록 설정이 다듬어졌으므로 뒤에서부터 고른다)"""
|
||
|
|
# 2026-10-03 PD 지적("원색 계열만 … 저급해 보인다") 뒤: 예전 합격본(금 + 적/청)으로 학습한 v4 가 원색으로 끌어당긴다.
|
||
|
|
# ANFG_TRAIN_TAGS="a26,a27" 처럼 회차를 주면 그 회차(재질색 + 작은 강조색)의 합격본만으로 학습한다 — PD 가 예전에 고른 14장(원색 세트)도 뺀다.
|
||
|
|
tags = [x for x in os.environ.get("ANFG_TRAIN_TAGS", "").split(",") if x]
|
||
|
|
if tags:
|
||
|
|
data = [e for e in data if e["src"].split("_")[0] in tags]
|
||
|
|
picked = [e for e in data if e["source"] == "pd"]
|
||
|
|
count = {}
|
||
|
|
for e in reversed([e for e in data if e["source"] != "pd"]):
|
||
|
|
if "-genpd" in e["src"]:
|
||
|
|
continue
|
||
|
|
k = kind_of(e["src"])
|
||
|
|
if count.get(k, 0) >= CAP:
|
||
|
|
continue
|
||
|
|
count[k] = count.get(k, 0) + 1
|
||
|
|
picked.append(e)
|
||
|
|
print("학습 선택 %d / %d · 종류별 %s" % (len(picked), len(data), sorted(count.items(), key=lambda x: -x[1])), flush=True)
|
||
|
|
return picked
|
||
|
|
|
||
|
|
|
||
|
|
def build_dataset():
|
||
|
|
data = select(icon_dataset.load())
|
||
|
|
os.makedirs(DS, exist_ok=True)
|
||
|
|
n = 0
|
||
|
|
for e in data:
|
||
|
|
im = Image.open(os.path.join(icon_dataset.DS_DIR, e["file"])).convert("RGBA")
|
||
|
|
side = max(1, int(RES * MARGIN))
|
||
|
|
small = im.resize((side, side), Image.LANCZOS)
|
||
|
|
im = Image.new("RGBA", (RES, RES), (0, 0, 0, 0)); im.paste(small, ((RES - side) // 2, (RES - side) // 2), small)
|
||
|
|
base = os.path.splitext(e["file"])[0]
|
||
|
|
# PD 가 직접 고른 그림은 여러 번 넣어 비중을 높인다
|
||
|
|
# 검토에서 "최상급"으로 뽑힌 그림(best_*.json → dataset.json 의 best 표시)은 두 번 넣는다(장식 밀도 · 붓결이 좋은 쪽으로 끌기 위해)
|
||
|
|
reps = PD_REPEAT if e["source"] == "pd" else (2 if e.get("best") else 1)
|
||
|
|
for r in range(reps):
|
||
|
|
for tag, v in (("a", im), ("b", ImageOps.mirror(im))):
|
||
|
|
bg = Image.new("RGB", (RES, RES), (0, 0, 0)); bg.paste(v, (0, 0), v)
|
||
|
|
stem = "%s_%d%s" % (base, r, tag)
|
||
|
|
bg.save(os.path.join(DS, stem + ".png"))
|
||
|
|
with io.open(os.path.join(DS, stem + ".txt"), "w", encoding="utf-8") as f:
|
||
|
|
f.write("%s, %s, %s, black background" % (TRIGGER, COMPOSE, e["caption"]))
|
||
|
|
n += 1
|
||
|
|
print("dataset %s: %d images (원본 %d)" % (DS, n, len(data)), flush=True)
|
||
|
|
return n
|
||
|
|
|
||
|
|
|
||
|
|
def get(path):
|
||
|
|
return json.load(urllib.request.urlopen(HOST + path, timeout=120))
|
||
|
|
|
||
|
|
|
||
|
|
def train():
|
||
|
|
wf = {
|
||
|
|
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": CKPT}},
|
||
|
|
"10": {"class_type": "LoraLoader", "inputs": {"model": ["1", 0], "clip": ["1", 1], "lora_name": BASE_LORA, "strength_model": BASE_LORA_STRENGTH, "strength_clip": BASE_LORA_STRENGTH}},
|
||
|
|
"2": {"class_type": "LoadImageTextDataSetFromFolder", "inputs": {"folder": DS_NAME}},
|
||
|
|
"3": {"class_type": "MakeTrainingDataset", "inputs": {"images": ["2", 0], "vae": ["1", 2], "clip": ["10", 1], "texts": ["2", 1]}},
|
||
|
|
"4": {"class_type": "TrainLoraNode", "inputs": {
|
||
|
|
"model": ["10", 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/" + NAME, "steps": ["4", 2]}},
|
||
|
|
"6": {"class_type": "LossGraphNode", "inputs": {"loss": ["4", 1], "filename_prefix": NAME + "_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")[:2000]); return None
|
||
|
|
if r.get("node_errors"):
|
||
|
|
print("NODE ERRORS:", json.dumps(r["node_errors"], ensure_ascii=False)[:2000]); return None
|
||
|
|
pid = r["prompt_id"]
|
||
|
|
print("queued %s · steps %d · batch %d · lr %s · rank %d" % (pid, STEPS, BATCH, LR, RANK), flush=True)
|
||
|
|
t0 = time.time()
|
||
|
|
while time.time() - t0 < 7200:
|
||
|
|
time.sleep(20)
|
||
|
|
h = get("/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("done in %.0fs" % (time.time() - t0), flush=True)
|
||
|
|
return h[pid]["outputs"]
|
||
|
|
print("TIMEOUT"); return None
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
if os.path.isdir(DS) and os.listdir(DS):
|
||
|
|
print("데이터셋 폴더가 이미 있다(버전 이름을 바꿔 다시 실행):", DS); sys.exit(1)
|
||
|
|
build_dataset()
|
||
|
|
if train() is None:
|
||
|
|
sys.exit(1)
|
||
|
|
files = sorted(glob.glob(os.path.join(LORA_OUT, NAME + "*.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: %s -> %s (%.1f MB)" % (files[-1], LORA_DST, os.path.getsize(LORA_DST) / 2**20))
|