# -*- coding: utf-8 -*-
# 2026-10-01 PD 지시 — Civitai "Game Icon Institute anfg v3" LoRA(SD1.5)로 아이콘을 직접 생성해 본다.
# ComfyUI(Comfy Desktop, 127.0.0.1:8000) API 로 작업을 보내고 결과를 받아 한 장으로 모은다.
#   python comfy_gii_trial.py <출력폴더> [round 이름]
# LoRA 정보(실측): 기반 sd_v1 · 학습 기반 모델 gameIconInstitute_v30 · 트리거 "game icon institute, game icon, ahfg" · 권장 강도 0.4~0.8
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"      # 보유한 SD1.5 체크포인트(권장 기반 gameIconInstitute_v30 은 없다)
LORA = "GameIconInstitute_anfg_v3.safetensors"
TRIGGER = "game icon institute, game icon, ahfg"
# 2차: 1차에서 배경이 회색 · 그라데이션, 바닥 반사가 나와 배경 관련 부정어를 더했다
NEG = "low resolution, blurry, text, watermark, signature, logo, frame, border, hands, person, cropped, out of frame, jpeg artifacts, grey background, white background, gradient background, yellow background, floor, table, reflection, shadow on ground, bright colors, all gold, brown leather"

# (이름, 프롬프트 본문) — 이번 아이템 시스템에 필요한 반지 + 보조 장비 7종
SUBJECTS = [
    # 2차: FR 톤(검은 강철 + 금 장식 + 주황 발광)을 가중치로 밀고, 형태가 어색했던 화살통 · 투척 무기 문구를 고쳤다
    ("ring", "ring, (black iron ring:1.3), thin gold trim, (glowing orange gemstone:1.2), gem, glowing, dark fantasy"),
    ("shield", "shield, (black steel kite shield:1.3), gold trim, (glowing orange emblem:1.2), dark fantasy, armor"),
    ("quiver", "quiver, (upright cylindrical black leather quiver:1.3), arrows with feathers sticking out of the top, gold fittings, orange glow, dark fantasy"),
    ("spellbook", "book, (closed black leather grimoire:1.3), gold corners, (glowing orange rune on cover:1.2), dark fantasy"),
    ("dagger", "dagger, (black steel dagger:1.3), gold guard, (glowing orange edge:1.1), dark fantasy, single weapon"),
    ("throwing", "throwing knives, (three small black steel throwing knives:1.3), fanned out, gold accents, orange glow, dark fantasy"),
    ("orb", "orb, (glowing orange crystal orb:1.2), held by (black iron claws:1.3), crystal, glowing, dark fantasy"),
    ("offhand_sword", "short sword, (black steel short sword:1.3), gold hilt, orange gem in guard, dark fantasy, single weapon"),
]


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=120).read()


def workflow(prompt, seed, lora_strength=0.6, steps=28, cfg=7.0, batch=2, clip_skip=-2):
    return {
        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": CKPT}},
        "2": {"class_type": "LoraLoader", "inputs": {"model": ["1", 0], "clip": ["1", 1], "lora_name": LORA, "strength_model": lora_strength, "strength_clip": lora_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": batch}},
        "7": {"class_type": "KSampler", "inputs": {"model": ["2", 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_trial"}},
    }


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 < 600:
        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)[:600])
            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(1.5)
    raise TimeoutError("comfy timeout")


def main():
    out = sys.argv[1]
    tag = sys.argv[2] if len(sys.argv) > 2 else "r1"
    strength = float(sys.argv[3]) if len(sys.argv) > 3 else 0.6
    os.makedirs(out, exist_ok=True)
    cell, pad, label = 256, 6, 18
    cols = 2
    sheet = Image.new("RGB", (len(SUBJECTS) * (cell + pad) + pad, cols * (cell + pad) + pad + label), (40, 40, 40))
    draw = ImageDraw.Draw(sheet)
    for i, (name, body) in enumerate(SUBJECTS):
        prompt = "%s, %s, still life, (pure black background:1.3), simple background, centered, single object, dark, high contrast" % (TRIGGER, body)
        t0 = time.time()
        images = run(prompt, 1000 + i * 7, lora_strength=strength)
        print("%s %s %d장 %.1f초" % (tag, name, len(images), time.time() - t0)); sys.stdout.flush()
        for j, im in enumerate(images[:cols]):
            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()