# -*- coding: utf-8 -*-
# 2026-10-02 PD 지적 — 1차 다시 칠하기(comfy_moba_repaint.py · 512) 결과에 "퀄리티 좀 아쉬운데?".
#   원인(직접 본 것): 512 로 줄여 칠해 1024 원본의 세부가 뭉개졌고, 세기를 낮춰(0.5~0.65) MOBA 풍의 질감이 덜 올라왔다.
# 고품질 방식(이 스크립트):
#   1) 원본에서 윤곽선을 뽑아 ControlNet(선화 또는 캐니)으로 형태를 붙잡는다 → 세기를 높여(0.75~0.9) 질감 · 빛을 MOBA 풍으로 새로 칠해도 형태가 남는다.
#   2) 768(모델 권장 최대)에서 칠한 뒤 4x-UltraSharp 로 키우고 1024 로 맞춰, 약한 세기(0.3)로 한 번 더 다듬는다(세부 선명도).
#   3) 배경 제거(Inspyrenet) → 1024 투명 PNG. 채점은 1차와 같다(원본 실루엣 겹침 · 한 덩어리 · 가장자리).
#   python comfy_moba_repaint_hq.py <출력폴더> <round 이름> <번호들 "1015,1253" 또는 all> [윤곽 "lineart,canny"] [세기 "0.75,0.9"] [시드 수=1] [윤곽 강도=0.8]
import io, json, os, sys, time, uuid, urllib.parse
from PIL import Image, ImageDraw
import comfy_moba_repaint as base          # 올리기 · 채점 · 캡션 표는 1차 스크립트의 것을 그대로 쓴다

CKPT = base.CKPT
PREFIX = "(masterpiece:1.2), (best quality:1.2), moba game icon"
STYLE = "highly detailed, intricate details, sharp focus, polished metal, strong specular highlights, rim light, rich colors, hand painted game equipment icon"
# 시험(h1)에서 갑옷에 소매 · 팔이 덧붙어 몸 관련 낱말을 더했다
NEG = base.NEG + ", flat color, sketch, monochrome, washed out, sleeves, arms, torso, body, mannequin"
RES1, RES2 = 768, 1024      # 1차로 칠하는 크기 · 최종 크기
CN = {"lineart": "control_v11p_sd15_lineart_fp16.safetensors", "canny": "control_v11p_sd15_canny_fp16.safetensors"}
UPSCALER = "4x-UltraSharp.pth"
REFINE = 0.30               # 2차 다듬기 세기


def prepare(path, res):
    """원본 아이콘을 res 캔버스 가운데에 여백을 두고 놓는다 — (어두운 바탕 RGB, 투명 RGBA)"""
    im = Image.open(path).convert("RGBA")
    box = im.getchannel("A").getbbox()
    if box:
        im = im.crop(box)
    scale = res * base.FIT / float(max(im.size))
    im = im.resize((max(1, int(im.width * scale)), max(1, int(im.height * scale))), Image.LANCZOS)
    rgba = Image.new("RGBA", (res, res), (0, 0, 0, 0))
    rgba.paste(im, ((res - im.width) // 2, (res - im.height) // 2), im)
    rgb = Image.new("RGB", (res, res), base.BG)
    rgb.paste(rgba, (0, 0), rgba)
    return rgb, rgba


def workflow(src768, src1024, prompt, seed, denoise, kind, cn_strength):
    """768 칠하기(ControlNet) → 4배 확대 → 1024 → 약하게 다듬기(ControlNet) → 배경 제거"""
    wf = {
        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": CKPT}},
        "3": {"class_type": "CLIPSetLastLayer", "inputs": {"clip": ["1", 1], "stop_at_clip_layer": -2}},
        "4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["3", 0], "text": prompt}},
        "5": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["3", 0], "text": NEG}},
        "20": {"class_type": "LoadImage", "inputs": {"image": src768}},
        "22": {"class_type": "LoadImage", "inputs": {"image": src1024}},
        "30": {"class_type": "ControlNetLoader", "inputs": {"control_net_name": CN[kind]}},
        # 윤곽은 1024 원본에서 뽑는다(가는 선까지 남긴다)
        "31": ({"class_type": "LineArtPreprocessor", "inputs": {"image": ["22", 0], "coarse": "disable", "resolution": 1024}} if kind == "lineart"
               else {"class_type": "Canny", "inputs": {"image": ["22", 0], "low_threshold": 0.15, "high_threshold": 0.45}}),
        "32": {"class_type": "ControlNetApplyAdvanced", "inputs": {"positive": ["4", 0], "negative": ["5", 0], "control_net": ["30", 0], "image": ["31", 0],
                                                                  "strength": cn_strength, "start_percent": 0.0, "end_percent": 0.85, "vae": ["1", 2]}},
        "21": {"class_type": "VAEEncode", "inputs": {"pixels": ["20", 0], "vae": ["1", 2]}},
        "7": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["32", 0], "negative": ["32", 1], "latent_image": ["21", 0], "seed": seed,
                                                   "steps": 30, "cfg": 7.0, "sampler_name": "dpmpp_2m", "scheduler": "karras", "denoise": denoise}},
        "8": {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0], "vae": ["1", 2]}},
        "40": {"class_type": "UpscaleModelLoader", "inputs": {"model_name": UPSCALER}},
        "41": {"class_type": "ImageUpscaleWithModel", "inputs": {"upscale_model": ["40", 0], "image": ["8", 0]}},
        "42": {"class_type": "ImageScale", "inputs": {"image": ["41", 0], "upscale_method": "lanczos", "width": RES2, "height": RES2, "crop": "disabled"}},
        "43": {"class_type": "VAEEncode", "inputs": {"pixels": ["42", 0], "vae": ["1", 2]}},
        "44": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["32", 0], "negative": ["32", 1], "latent_image": ["43", 0], "seed": seed + 1,
                                                    "steps": 24, "cfg": 7.0, "sampler_name": "dpmpp_2m", "scheduler": "karras", "denoise": REFINE}},
        "45": {"class_type": "VAEDecode", "inputs": {"samples": ["44", 0], "vae": ["1", 2]}},
        "12": {"class_type": "InspyrenetRembg", "inputs": {"image": ["45", 0], "torchscript_jit": "default"}},
        "9": {"class_type": "SaveImage", "inputs": {"images": ["12", 0], "filename_prefix": "moba_hq"}},
    }
    return wf


def run(wf):
    pid = base.post("/prompt", {"prompt": wf, "client_id": str(uuid.uuid4())})["prompt_id"]
    t0 = time.time()
    while time.time() - t0 < 900:
        hist = json.loads(base.get("/history/" + pid))
        if pid in hist:
            st = hist[pid].get("status", {})
            if st.get("status_str") == "error":
                raise RuntimeError(json.dumps(st)[:1500])
            for node in hist[pid].get("outputs", {}).values():
                for im in node.get("images", []):
                    q = urllib.parse.urlencode({"filename": im["filename"], "subfolder": im.get("subfolder", ""), "type": im.get("type", "output")})
                    return Image.open(io.BytesIO(base.get("/view?" + q))).convert("RGBA")
        time.sleep(1.0)
    raise TimeoutError("comfy timeout")


def main():
    out, tag = sys.argv[1], sys.argv[2]
    cap = base.load_captions()
    ids = sorted(cap.keys()) if sys.argv[3] == "all" else [int(v) for v in sys.argv[3].split(",")]
    kinds = (sys.argv[4] if len(sys.argv) > 4 else "lineart,canny").split(",")
    denoises = [float(v) for v in (sys.argv[5] if len(sys.argv) > 5 else "0.75,0.9").split(",")]
    seeds = int(sys.argv[6]) if len(sys.argv) > 6 else 1
    cn_strength = float(sys.argv[7]) if len(sys.argv) > 7 else 0.8
    old_dir = sys.argv[8] if len(sys.argv) > 8 else ""          # 1차 결과(고른 것) 폴더 — 있으면 비교 칸으로 넣는다
    os.makedirs(out, exist_ok=True)
    show = len(kinds) * len(denoises) * seeds
    extra = 2 if old_dir else 1
    cell, pad = 256, 4
    sheet = Image.new("RGB", (pad + (show + extra) * (cell + pad), len(ids) * (cell + pad) + pad), (40, 40, 40))
    draw = ImageDraw.Draw(sheet)
    scores = {}

    def put(im, col, y, text):
        bg = base.checker(cell); small = im.resize((cell, cell), Image.LANCZOS); bg.paste(small, (0, 0), small)
        sheet.paste(bg, (pad + col * (cell + pad), y))
        draw.text((pad + col * (cell + pad) + 3, y + 2), text, fill=(255, 255, 255))

    for i, gid in enumerate(ids):
        path = os.path.join(base.SRC, "%d.png" % gid)
        rgb1, _ = prepare(path, RES1)
        rgb2, rgba2 = prepare(path, RES2)
        n1 = base.upload(rgb1, "bt_moba_hq768_%d.png" % gid)
        n2 = base.upload(rgb2, "bt_moba_hq1024_%d.png" % gid)
        src_mask = base.mask(rgba2)
        prompt = "%s, (%s:1.2), %s, no humans, single item, isolated, dark background" % (PREFIX, cap[gid], STYLE)
        y = pad + i * (cell + pad)
        put(rgba2, 0, y, "%d 원본" % gid)
        col = 1
        if old_dir:
            old = os.path.join(old_dir, "%d_0.png" % gid)
            if os.path.exists(old):
                put(Image.open(old).convert("RGBA"), 1, y, "1차(512)")
            col = 2
        results = []
        t0 = time.time()
        for kind in kinds:
            for d in denoises:
                for j in range(seeds):
                    fn = "%s_%d_%s_d%02d_%d.png" % (tag, gid, kind, int(d * 100), j)
                    # 이미 만든 그림은 다시 만들지 않는다(긴 작업이 끊겨도 이어서 돌린다)
                    if os.path.exists(os.path.join(out, fn)):
                        im = Image.open(os.path.join(out, fn)).convert("RGBA")
                    else:
                        im = run(workflow(n1, n2, prompt, 40000 + gid * 7 + j * 911 + int(d * 100), d, kind, cn_strength))
                        im.save(os.path.join(out, fn))
                    s, iou, mainc, edge = base.score(im, src_mask)
                    results.append(dict(file=fn, score=round(s, 3), iou=round(iou, 3), main=round(mainc, 3), edge=round(edge, 3), denoise=d, kind=kind))
                    put(im, col, y, "%s d%.2f #%d IoU%.2f" % (kind, d, j, iou)); col += 1
        scores[gid] = results
        print("%s %d · %d장 · %.0f초" % (tag, gid, len(results), time.time() - t0)); sys.stdout.flush()
        sheet.save(os.path.join(out, "%s_sheet.png" % tag))
    json.dump(scores, io.open(os.path.join(out, "%s_scores.json" % tag), "w", encoding="utf-8"), ensure_ascii=False, indent=1)
    print("sheet", os.path.join(out, "%s_sheet.png" % tag))


if __name__ == "__main__":
    main()