# -*- coding: utf-8 -*-
# 2026-10-01 PD 지시 — 체크포인트 gameIconInstitute_v4XL(SDXL)로 사이트 분위기 그대로 아이콘을 만든다(FR 톤으로 밀지 않는다).
#   python comfy_gii_xl.py <출력폴더> <round 이름> <스타일: 3dicon|xieshi|2dguofeng|2dXIANTIAO>
# 모델 설명(civitai 47800 · V4_XL): 스타일 태그 gmic_\(3dicon\) · gmic_\(2dguofeng\) · gmic_\(2dXIANTIAO\) · gmic icon_\(xieshi\). 예시 그림의 생성 값은 공개돼 있지 않아
# 같은 제작자의 XL 모델(426087) 권장값(DPM++ 2M · 20 steps · CFG 7 · clip skip 2)을 따른다. SD1.5 용 LoRA(anfg v3)는 XL 과 맞지 않아 쓰지 않는다.
import io, json, os, sys, time, uuid, urllib.request, urllib.parse
from PIL import Image, ImageDraw

HOST = "http://127.0.0.1:8000"
CKPT = "gameIconInstitute_v4XL.safetensors"
STYLE = {"3dicon": "gmic_\\(3dicon\\)", "xieshi": "gmic icon_\\(xieshi\\)", "2dguofeng": "gmic_\\(2dguofeng\\)", "2dXIANTIAO": "gmic_\\(2dXIANTIAO\\)"}
NEG = "worst quality, low quality, normal quality, lowres, blurry, cropped, text, watermark, signature, jpeg artifacts, person, hands"

# 1차(권장 설정 그대로) 때와 같은 담백한 문구 — 톤을 밀지 않는다
SUBJECTS = [
    ("ring", "a ring with a gemstone"),
    ("shield", "a shield with an emblem"),
    ("quiver", "a quiver full of arrows"),
    ("spellbook", "a closed spellbook, magic book"),
    ("dagger", "a dagger"),
    ("throwing", "three throwing knives"),
    ("orb", "a magic crystal orb on a stand"),
    ("offhand_sword", "a short sword"),
]


def post(path, data):
    req = urllib.request.Request(HOST + path, data=json.dumps(data).encode("utf-8"), headers={"Content-Type": "application/json"})
    return json.loads(urllib.request.urlopen(req, timeout=60).read())


def get(path):
    return urllib.request.urlopen(HOST + path, timeout=180).read()


def workflow(prompt, seed, size=1024, steps=20, cfg=7.0, batch=1):
    return {
        "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}},
        "6": {"class_type": "EmptyLatentImage", "inputs": {"width": size, "height": size, "batch_size": batch}},
        "7": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["6", 0], "seed": seed, "steps": steps, "cfg": cfg, "sampler_name": "dpmpp_2m", "scheduler": "karras", "denoise": 1.0}},
        "8": {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0], "vae": ["1", 2]}},
        "9": {"class_type": "SaveImage", "inputs": {"images": ["8", 0], "filename_prefix": "gii_xl"}},
    }


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


def main():
    out, tag, style = sys.argv[1], sys.argv[2], sys.argv[3]
    seeds = int(sys.argv[4]) if len(sys.argv) > 4 else 2
    os.makedirs(out, exist_ok=True)
    cell, pad, label = 256, 6, 18
    sheet = Image.new("RGB", (len(SUBJECTS) * (cell + pad) + pad, seeds * (cell + pad) + pad + label), (40, 40, 40))
    draw = ImageDraw.Draw(sheet)
    for i, (name, body) in enumerate(SUBJECTS):
        prompt = "%s, game icon, %s, simple background, centered, single object" % (STYLE[style], body)
        for j in range(seeds):
            t0 = time.time()
            im = run(prompt, 2000 + i * 11 + j * 101)[0]
            print("%s %s #%d %.1f초" % (tag, name, j, time.time() - t0)); sys.stdout.flush()
            im.save(os.path.join(out, "%s_%s_%d.png" % (tag, name, j)))
            sheet.paste(im.resize((cell, cell), Image.LANCZOS), (pad + i * (cell + pad), label + pad + j * (cell + pad)))
        draw.text((pad + i * (cell + pad) + 4, 3), name, fill=(240, 220, 160))
    path = os.path.join(out, "%s_sheet.png" % tag)
    sheet.save(path)
    print("sheet", path)


if __name__ == "__main__":
    main()