# -*- coding: utf-8 -*-
# 2026-10-01 PD 지시 — "game Icon Institute_anfg_v3-000018 로라를 적용해서 다시" · 트리거 "game icon institute, game icon, ahfg".
# 사이트 예시(직접 확인): 굵고 양식화된 손그림풍 MMO 아이콘 · 진한 채도 · 두꺼운 형태 · 발광 · 검은 배경.
# 앞선 결과가 얇고 사실적으로 나온 원인 = 기반 모델(제작자 권장 "游戏图标研究所 3.0" SD1.5 가 이 PC 에 없어 사실풍 difconsistency 로 대체).
# 여기서는 같은 기반에서 LoRA 를 더 세게 걸고(변형별 강도), 사실풍을 부정어로 누르고, 2배 다시 그리기를 빼서 LoRA 의 그림체를 최대한 살린다.
#   python comfy_gii_style.py <출력폴더> <round 이름> <LoRA 강도> <clip skip(음수)> [rembg 0|1] [시드 수]
import io, json, os, sys, time, uuid, urllib.request, urllib.parse
from PIL import Image, ImageDraw

HOST = "http://127.0.0.1:8000"
CKPT = "difconsistencyCheckpoint_v10.safetensors"
LORA = "game Icon Institute_anfg_v3-000018.safetensors"          # PD 가 loras 폴더에 넣어 둔 원래 이름의 파일
TRIGGER = "game icon institute, game icon, ahfg"
# 페이지의 부정어 "(easynegative:0.8), black, dark, Low resolution" + 사실풍 누르기
NEG = "black, dark, low resolution, worst quality, low quality, blurry, photo, photorealistic, realistic, 3d render, thin, text, watermark, signature, cropped, person, hands"

SUBJECTS = [
    ("sword", "sword"),
    ("blunt", "hammer, mace"),
    ("bow", "bow"),
    ("dagger", "dagger"),
    ("staff", "staff, gem"),
    ("magic_dagger", "dagger, glowing, gem"),
    ("armor", "armor"),
    ("belt", "belt"),
    ("gloves", "gloves, gauntlets"),
    ("boots", "boots"),
    ("ring", "ring, gem"),
    ("necklace", "necklace, gem"),
    ("shield", "shield"),
    ("quiver", "quiver, arrows"),
    ("spellbook", "book"),
    ("throwing", "shuriken"),
    ("orb", "orb, crystal, glowing"),
    ("offhand_sword", "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, strength, clip_skip, rembg):
    wf = {
        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": CKPT}},
        "2": {"class_type": "LoraLoader", "inputs": {"model": ["1", 0], "clip": ["1", 1], "lora_name": LORA, "strength_model": strength, "strength_clip": strength}},
        "3": {"class_type": "CLIPSetLastLayer", "inputs": {"clip": ["2", 1], "stop_at_clip_layer": clip_skip}},
        "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": 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": 28, "cfg": 7.0, "sampler_name": "euler_ancestral", "scheduler": "normal", "denoise": 1.0}},
        "8": {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0], "vae": ["1", 2]}},
    }
    if rembg:
        wf["12"] = {"class_type": "InspyrenetRembg", "inputs": {"image": ["8", 0], "torchscript_jit": "default"}}
        wf["9"] = {"class_type": "SaveImage", "inputs": {"images": ["12", 0], "filename_prefix": "gii_style"}}
    else:
        wf["9"] = {"class_type": "SaveImage", "inputs": {"images": ["8", 0], "filename_prefix": "gii_style"}}
    return wf


def run(prompt, seed, strength, clip_skip, rembg):
    cid = str(uuid.uuid4())
    pid = post("/prompt", {"prompt": workflow(prompt, seed, strength, clip_skip, rembg), "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:
                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")})
                        return Image.open(io.BytesIO(get("/view?" + q))).convert("RGBA")
        time.sleep(1.2)
    raise TimeoutError("comfy timeout")


def checker(size, n=16):
    im = Image.new("RGB", (size, size), (58, 58, 58))
    d = ImageDraw.Draw(im)
    for y in range(0, size, n):
        for x in range(0, size, n):
            if (x // n + y // n) % 2:
                d.rectangle([x, y, x + n - 1, y + n - 1], fill=(74, 74, 74))
    return im


def main():
    out, tag = sys.argv[1], sys.argv[2]
    strength, clip_skip = float(sys.argv[3]), int(sys.argv[4])
    rembg = int(sys.argv[5]) if len(sys.argv) > 5 else 0
    seeds = int(sys.argv[6]) if len(sys.argv) > 6 else 2
    lo = int(sys.argv[7]) if len(sys.argv) > 7 else 0
    hi = int(sys.argv[8]) if len(sys.argv) > 8 else len(SUBJECTS)
    os.makedirs(out, exist_ok=True)
    subs = SUBJECTS[lo:hi]
    cell, pad, label = 200, 4, 16
    cols = 9
    rows = (len(subs) + cols - 1) // cols
    sheet = Image.new("RGB", (cols * (cell + pad) + pad, rows * (seeds * (cell + pad) + label) + pad), (40, 40, 40))
    draw = ImageDraw.Draw(sheet)
    for i, (name, body) in enumerate(subs):
        prompt = "%s, %s, still life, black background" % (TRIGGER, body)
        cx = pad + (i % cols) * (cell + pad)
        cy = pad + (i // cols) * (seeds * (cell + pad) + label)
        draw.text((cx + 3, cy + 2), name, fill=(240, 220, 160))
        for j in range(seeds):
            t0 = time.time()
            im = run(prompt, 5000 + (lo + i) * 17 + j * 307, strength, clip_skip, rembg)
            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)))
            bg = checker(cell) if rembg else Image.new("RGB", (cell, cell), (0, 0, 0))
            small = im.resize((cell, cell), Image.LANCZOS)
            bg.paste(small, (0, 0), small)
            sheet.paste(bg, (cx, cy + label + j * (cell + pad)))
    path = os.path.join(out, "%s_sheet.png" % tag)
    sheet.save(path)
    print("sheet", path)


if __name__ == "__main__":
    main()