# -*- coding: utf-8 -*-
# 2026-10-01 PD 지시 — "내가 원하는 스타일이 아니야. 다른 체크포인트로 제작을 시도해보자."
# 보유한 다른 체크포인트로 같은 대상 6종을 시험 생성해 한 장으로 모은다(풍 비교용).
#   python comfy_alt_ckpt.py <출력폴더> <flux|noob|noob2|sd15rpg>
import io, json, os, sys, time, uuid, urllib.request, urllib.parse
from PIL import Image, ImageDraw

HOST = "http://127.0.0.1:8000"
SUBJECTS = [("sword", "a sword"), ("armor", "a breastplate armor"), ("gloves", "a pair of gauntlets"), ("ring", "a ring with a gemstone"), ("shield", "a shield"), ("bow", "a bow")]
# 목표 풍(PD 가 처음 고른 사이트 예시의 특징): 굵고 양식화된 손그림풍 MMO 아이콘 · 진한 채도 · 두꺼운 형태 · 발광 · 검은 배경
STYLE = "stylized fantasy MMORPG game item icon, hand-painted, bold thick chunky shapes, vibrant saturated colors, glowing magical accents, ornate, three-quarter view, centered single object, plain black background"

PRESETS = {
    "flux": dict(ckpt="flux1CheckpointEasyTo_v10.safetensors", flux=True, size=768, steps=20, cfg=1.0, sampler="euler", scheduler="simple", neg="", prefix="", lora=None),
    "noob": dict(ckpt="noobaiXLNAIXL_epsilonPred05Version.safetensors", flux=False, size=1024, steps=26, cfg=6.0, sampler="euler_ancestral", scheduler="normal",
                 neg="worst quality, low quality, blurry, text, watermark, person, 1girl, hands, photo", prefix="masterpiece, best quality, no humans, ", lora=None),
    "noob2": dict(ckpt="catTowerNoobaiXL_chenkinnoobV11.safetensors", flux=False, size=1024, steps=26, cfg=6.0, sampler="euler_ancestral", scheduler="normal",
                  neg="worst quality, low quality, blurry, text, watermark, person, 1girl, hands, photo", prefix="masterpiece, best quality, no humans, ", lora=None),
    "sd15rpg": dict(ckpt="difconsistencyCheckpoint_v10.safetensors", flux=False, size=512, steps=28, cfg=7.0, sampler="dpmpp_2m", scheduler="karras",
                    neg="worst quality, low quality, blurry, text, watermark, person, hands, photo, realistic", prefix="", lora=("GameItemUI_sd15.safetensors", 0.8)),
}


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


def workflow(p, prompt, seed):
    wf = {"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": p["ckpt"]}}}
    model, clip = ["1", 0], ["1", 1]
    if p["lora"]:
        wf["2"] = {"class_type": "LoraLoader", "inputs": {"model": model, "clip": clip, "lora_name": p["lora"][0], "strength_model": p["lora"][1], "strength_clip": p["lora"][1]}}
        model, clip = ["2", 0], ["2", 1]
    wf["4"] = {"class_type": "CLIPTextEncode", "inputs": {"clip": clip, "text": prompt}}
    wf["5"] = {"class_type": "CLIPTextEncode", "inputs": {"clip": clip, "text": p["neg"]}}
    pos = ["4", 0]
    if p["flux"]:
        wf["13"] = {"class_type": "FluxGuidance", "inputs": {"conditioning": ["4", 0], "guidance": 3.5}}
        pos = ["13", 0]
    wf["6"] = {"class_type": "EmptyLatentImage", "inputs": {"width": p["size"], "height": p["size"], "batch_size": 1}}
    wf["7"] = {"class_type": "KSampler", "inputs": {"model": model, "positive": pos, "negative": ["5", 0], "latent_image": ["6", 0], "seed": seed, "steps": p["steps"], "cfg": p["cfg"], "sampler_name": p["sampler"], "scheduler": p["scheduler"], "denoise": 1.0}}
    wf["8"] = {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0], "vae": ["1", 2]}}
    wf["9"] = {"class_type": "SaveImage", "inputs": {"images": ["8", 0], "filename_prefix": "alt_ckpt"}}
    return wf


def run(p, prompt, seed):
    pid = post("/prompt", {"prompt": workflow(p, prompt, seed), "client_id": str(uuid.uuid4())})["prompt_id"]
    t0 = time.time()
    while time.time() - t0 < 1500:
        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)[:900])
            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(get("/view?" + q))).convert("RGB")
        time.sleep(2)
    raise TimeoutError("comfy timeout")


def main():
    out, name = sys.argv[1], sys.argv[2]
    p = PRESETS[name]
    os.makedirs(out, exist_ok=True)
    cell, pad, label = 256, 6, 18
    sheet = Image.new("RGB", (len(SUBJECTS) * (cell + pad) + pad, cell + pad * 2 + label), (40, 40, 40))
    draw = ImageDraw.Draw(sheet)
    for i, (sname, body) in enumerate(SUBJECTS):
        prompt = "%s%s, %s" % (p["prefix"], body, STYLE)
        t0 = time.time()
        im = run(p, prompt, 7000 + i * 19)
        print("%s %s %.1f초" % (name, sname, time.time() - t0)); sys.stdout.flush()
        im.save(os.path.join(out, "%s_%s.png" % (name, sname)))
        sheet.paste(im.resize((cell, cell), Image.LANCZOS), (pad + i * (cell + pad), label + pad))
        draw.text((pad + i * (cell + pad) + 4, 3), "%s / %s" % (name, sname), fill=(240, 220, 160))
    path = os.path.join(out, "%s_sheet.png" % name)
    sheet.save(path)
    print("sheet", path)


if __name__ == "__main__":
    main()