# 3차 시험: 팔레트 고정(검은 강철·금 장식·주황 발광)을 프롬프트 가중치로만 얼마나 맞출 수 있는지 · Diablo LoRA 0.8 · 여백 확보 후처리 · Inspyrenet · 비교표
import json, os, time, uuid, glob, 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__))
OUT = os.path.join(S, "comfy_out3"); os.makedirs(OUT, exist_ok=True)
CKPT = "difconsistencyCheckpoint_v10.safetensors"; LORA = "Diablo4_sd15.safetensors"
STYLE = ("(single {subject}:1.3), game inventory item icon, (blackened dark steel:1.4), (gold filigree ornaments:1.3), (warm orange glowing runes and gem:1.3), "
         "(no silver, no blue, no white metal:1.2), matte black background, isolated object with wide empty margin, entire object visible, centered, "
         "painted realistic dark fantasy diablo style, dramatic warm rim light, ultra detailed, sharp")
NEG = ("silver, chrome, white, blue, cyan, pastel, bright background, emblem, badge, circular frame, ornate border, background pattern, flames in background, "
       "cropped, cut off, out of frame, text, watermark, blurry, lowres, multiple objects, character, person, hands, ui panel, cartoon, anime")
SUBJECTS = [("helmet_knight", "closed knight helmet with gold trim"), ("helmet_horned", "horned dark iron great helm"), ("ring_ruby", "thick black iron ring with gold inlay and a glowing orange gem"),
            ("shield_kite", "kite shield with gold rivets and a glowing rune"), ("orb_arcane", "dark iron claw holding a glowing orange crystal orb"), ("grimoire", "black leather grimoire with gold clasps and glowing rune"),
            ("quiver", "black leather arrow quiver with gold buckles"), ("amulet", "black iron gothic amulet with gold filigree and orange gem")]


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"})
    pid = json.load(urllib.request.urlopen(req, timeout=60))["prompt_id"]
    for _ in range(240):
        time.sleep(2)
        h = get(f"/history/{pid}")
        if pid in h:
            if h[pid].get("status", {}).get("status_str") == "error":
                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": "LoraLoader", "inputs": {"model": ["1", 0], "clip": ["1", 1], "lora_name": LORA, "strength_model": 0.8, "strength_clip": 0.8}},
                 "4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 1], "text": STYLE.format(subject=subject)}},
                 "5": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 1], "text": NEG}},
                 "6": {"class_type": "EmptyLatentImage", "inputs": {"width": 640, "height": 640, "batch_size": 1}},
                 "7": {"class_type": "KSampler", "inputs": {"model": ["2", 0], "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["6", 0], "seed": seed, "steps": 32, "cfg": 7.5, "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_icon3"}}})


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_inspy3"}}})


def fit_margin(path, margin=0.12):
    """알파 경계로 잘라 정사각 1024 에 여백을 두고 배치 — 테두리에 닿은(잘린) 오브젝트는 표시"""
    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 = []
for i, (key, subject) in enumerate(SUBJECTS):
    data = gen(subject, 33100 + 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)
    print(key, "ok")
print(f"total {time.time() - t0:.0f}s · 테두리에 잘린 것: {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 단검")]]
ko = {"helmet_knight": "투구(기사)", "helmet_horned": "투구(뿔)", "ring_ruby": "반지", "shield_kite": "방패", "orb_arcane": "오브", "grimoire": "마도서", "quiver": "화살통", "amulet": "목걸이"}
gen3 = [(os.path.join(OUT, k + "_inspy.png"), "3차 " + ko[k] + (" (잘림)" if k in cut else "")) for k, _ in SUBJECTS if os.path.exists(os.path.join(OUT, k + "_inspy.png"))]
cell, cols = 190, 8
sheet = Image.new("RGB", (cols * cell, 2 * (cell + 26)), (58, 58, 58)); d = ImageDraw.Draw(sheet)
font = ImageFont.truetype("C:/Windows/Fonts/malgunbd.ttf", 13)
for row, items in enumerate((ref, gen3)):
    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.save(os.path.join(S, "icon_compare4.png")); print("sheet ok", len(gen3))