BurningTimesAi/공유/기획/FirstRising/아이템/comfy_icons4.py

# 4차 시험: FR 장비 아이콘으로 학습한 스타일 LoRA(FR_icons_v1) 로 생성 → Inspyrenet 배경 제거 → 여백 보정 → 비교표 (BT25 추가 지시 2026-09-25)
#   A열 = 기존 부위 재현 8(화풍 학습 여부 판정용) · B열 = 신규 부위 8(투구2·반지·방패·오브·마도서·화살통·목걸이)
# 사용: python comfy_icons4.py [lora_file=FR_icons_v1.safetensors] [strength=1.0] [tag=4차]
import json, os, sys, time, uuid, glob, shutil, urllib.request, urllib.parse
from PIL import Image, ImageDraw, ImageFont

HOST = "http://127.0.0.1:8000"
S = os.path.dirname(os.path.abspath(__file__))
LORA = sys.argv[1] if len(sys.argv) > 1 else "FR_icons_v1.safetensors"
STRENGTH = float(sys.argv[2]) if len(sys.argv) > 2 else 1.0
TAG = sys.argv[3] if len(sys.argv) > 3 else "4차"
EXTRA = (", " + sys.argv[4]) if len(sys.argv) > 4 else ""  # 4차b: 잘림·광택 보정용 추가 문구
OUT = os.path.join(os.environ.get("BT_SCRATCH", S), f"comfy_out_{TAG}"); os.makedirs(OUT, exist_ok=True)
CKPT = "difconsistencyCheckpoint_v10.safetensors"
TRIGGER = "frdicon"
# 학습 캡션과 같은 형식 유지: "<trigger>, game equipment item icon, <내용>, black background"
POS = TRIGGER + ", game equipment item icon, {subject}, black background"
NEG = "text, letters, watermark, signature, cropped, cut off, out of frame, multiple objects, blurry, lowres, frame, border, badge, character, hands, cartoon, anime"
REPRO = [("sword", "a long dark steel sword with a gold crossguard and an orange glowing rune"),
         ("mace", "a heavy iron war hammer with orange glowing cracks and a dark handle"),
         ("necklace", "a dark ornate round amulet with an orange glowing gem"),
         ("glove", "a pair of dark leather gauntlets with gold trim and orange glowing runes"),
         ("armor", "a black plate armor with gold trim and an orange glowing emblem"),
         ("belt", "a dark leather belt with gold buckles and orange gems"),
         ("boots", "a pair of black high boots with gold trim and orange glowing accents"),
         ("dagger", "a dark steel dagger with a cyan glowing gem")]
NEW = [("helmet_knight", "a closed dark steel knight helmet with gold trim"),
       ("helmet_horned", "a horned dark iron great helm with orange glowing eyes"),
       ("ring_ruby", "a thick black iron ring with gold inlay and an orange glowing gem"),
       ("shield_kite", "a dark steel kite shield with gold rivets and a glowing rune"),
       ("orb_arcane", "a dark iron claw holding a glowing orange crystal orb"),
       ("grimoire", "a black leather grimoire with gold clasps and a glowing rune"),
       ("quiver", "a black leather arrow quiver with gold buckles and arrows"),
       ("amulet", "a black iron gothic amulet with gold filigree and a cyan gem")]
KO = {"sword": "검", "mace": "둔기", "necklace": "목걸이", "glove": "장갑", "armor": "갑옷", "belt": "벨트", "boots": "신발", "dagger": "단검",
      "helmet_knight": "투구(기사)", "helmet_horned": "투구(뿔)", "ring_ruby": "반지", "shield_kite": "방패", "orb_arcane": "오브", "grimoire": "마도서", "quiver": "화살통", "amulet": "목걸이(신)"}


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


def post(wf):
    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("REJECTED:", e.read().decode(errors="replace")[:600]); return None
    pid = r["prompt_id"]
    for _ in range(240):
        time.sleep(2)
        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("ERROR:", m[1].get("node_type"), m[1].get("exception_message", "")[:600])
                return None
            for nid, o in h[pid]["outputs"].items():
                for im in o.get("images", []):
                    q = urllib.parse.urlencode({"filename": im["filename"], "subfolder": im.get("subfolder", ""), "type": im.get("type", "output")})
                    return urllib.request.urlopen(HOST + "/view?" + q, timeout=120).read()
    return None


def upload(path):
    boundary = "----bt25" + uuid.uuid4().hex
    with open(path, "rb") as f:
        data = f.read()
    body = (f"--{boundary}\r\nContent-Disposition: form-data; name=\"image\"; filename=\"{os.path.basename(path)}\"\r\nContent-Type: image/png\r\n\r\n").encode() + data + (f"\r\n--{boundary}\r\nContent-Disposition: form-data; name=\"overwrite\"\r\n\r\ntrue\r\n--{boundary}--\r\n").encode()
    req = urllib.request.Request(HOST + "/upload/image", data=body, headers={"Content-Type": f"multipart/form-data; boundary={boundary}"})
    return json.load(urllib.request.urlopen(req, timeout=120))["name"]


def gen(subject, seed):
    return post({"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": CKPT}},
                 "2": {"class_type": "LoraLoaderModelOnly", "inputs": {"model": ["1", 0], "lora_name": LORA, "strength_model": STRENGTH}},
                 "4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["1", 1], "text": POS.format(subject=subject + EXTRA)}},
                 "5": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["1", 1], "text": NEG}},
                 "6": {"class_type": "EmptyLatentImage", "inputs": {"width": 512, "height": 512, "batch_size": 1}},
                 "7": {"class_type": "KSampler", "inputs": {"model": ["2", 0], "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["6", 0], "seed": seed, "steps": 30, "cfg": 7.0, "sampler_name": "dpmpp_2m", "scheduler": "karras", "denoise": 1.0}},
                 "8": {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0], "vae": ["1", 2]}},
                 "9": {"class_type": "UpscaleModelLoader", "inputs": {"model_name": "4x-UltraSharp.pth"}},
                 "10": {"class_type": "ImageUpscaleWithModel", "inputs": {"upscale_model": ["9", 0], "image": ["8", 0]}},
                 "11": {"class_type": "ImageScale", "inputs": {"image": ["10", 0], "upscale_method": "lanczos", "width": 1024, "height": 1024, "crop": "disabled"}},
                 "13": {"class_type": "SaveImage", "inputs": {"images": ["11", 0], "filename_prefix": "bt25_icon4"}}})


def rembg(path):
    return post({"1": {"class_type": "LoadImage", "inputs": {"image": upload(path)}},
                 "2": {"class_type": "InspyrenetRembg", "inputs": {"image": ["1", 0], "torchscript_jit": "default"}},
                 "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0], "filename_prefix": "bt25_inspy4"}}})


def fit_margin(path, margin=0.12):
    """알파 경계로 잘라 정사각 1024 에 여백을 두고 배치 — 테두리에 닿은(잘린) 오브젝트는 True"""
    im = Image.open(path).convert("RGBA"); a = im.getchannel("A").point(lambda v: 255 if v > 24 else 0)
    box = a.getbbox()
    if not box:
        return False
    touching = box[0] <= 2 or box[1] <= 2 or box[2] >= im.width - 2 or box[3] >= im.height - 2
    crop = im.crop(box); side = int(max(crop.size) / (1 - 2 * margin))
    canvas = Image.new("RGBA", (side, side), (0, 0, 0, 0)); canvas.paste(crop, ((side - crop.width) // 2, (side - crop.height) // 2), crop)
    canvas.resize((1024, 1024), Image.LANCZOS).save(path)
    return touching


t0 = time.time(); cut = []; made = {}
for i, (key, subject) in enumerate(REPRO + NEW):
    data = gen(subject, 44100 + i)
    if data is None:
        print(key, "FAILED"); continue
    p = os.path.join(OUT, key + ".png")
    with open(p, "wb") as f:
        f.write(data)
    d2 = rembg(p)
    if d2 is None:
        print(key, "rembg failed"); continue
    p2 = p[:-4] + "_inspy.png"
    with open(p2, "wb") as f:
        f.write(d2)
    if fit_margin(p2):
        cut.append(key)
    made[key] = p2
    print(key, "ok", f"{time.time() - t0:.0f}s", flush=True)
print(f"total {time.time() - t0:.0f}s · lora {LORA} x{STRENGTH} · 테두리에 잘린 것: {cut}")

u = r"C:\BurningTimes\FirstRising\Assets\GameAssets\UI\EquipmentItemIcon"
ref = [(os.path.join(u, n + ".png"), t) for n, t in [("1001", "FR 검"), ("1015", "FR 둔기"), ("1085", "FR 목걸이"), ("1155", "FR 장갑"), ("1225", "FR 갑옷"), ("1295", "FR 벨트"), ("1365", "FR 신발"), ("1043", "FR 단검")]]
rowA = [(made[k], f"{TAG} 재현 {KO[k]}" + (" (잘림)" if k in cut else "")) for k, _ in REPRO if k in made]
rowB = [(made[k], f"{TAG} 신규 {KO[k]}" + (" (잘림)" if k in cut else "")) for k, _ in NEW if k in made]
cell, cols = 190, 8
sheet = Image.new("RGB", (cols * cell, 3 * (cell + 26)), (58, 58, 58)); d = ImageDraw.Draw(sheet)
font = ImageFont.truetype("C:/Windows/Fonts/malgunbd.ttf", 13)
for row, items in enumerate((ref, rowA, rowB)):
    for i, (p, t) in enumerate(items[:cols]):
        im = Image.open(p).convert("RGBA"); im.thumbnail((cell - 10, cell - 10))
        bg = Image.new("RGBA", (cell, cell), (58, 58, 58, 255)); bg.paste(im, ((cell - im.width) // 2, (cell - im.height) // 2), im)
        x, y = i * cell, row * (cell + 26); sheet.paste(bg.convert("RGB"), (x, y))
        d.text((x + 4, y + cell + 5), t, fill=(240, 220, 160) if row else (220, 220, 220), font=font)
sheet_path = os.path.join(OUT, f"비교_{TAG}.png"); sheet.save(sheet_path); print("sheet", sheet_path, len(rowA), len(rowB))