BurningTimesAi/공유/기획/FirstRising/아이템/comfy_gii_site.py

# -*- coding: utf-8 -*-
# 2026-10-01 PD 지시 — civitai 143301 "Game Icon Institute anfg v3" 페이지의 아이콘 풍을 최대한 따라 만들고, 배경은 반드시 투명.
#   python comfy_gii_site.py <출력폴더> <round 이름> <시작 번호> <끝 번호(미포함)> [LoRA 강도=0.6] [시드 수=2]
# 페이지에 적힌 사용법 그대로: 프롬프트 "game icon institute, game icon, ahfg, <대상>, still life, black background" · LoRA 0.6 ·
#   부정 "(easynegative:0.8), black, dark, Low resolution"(easynegative 임베딩은 이 PC 에 없어 같은 뜻의 낱말로 대신한다).
# 기반 모델: 페이지 권장은 gameIconInstitute v3.0(SD1.5)인데 이 PC 에 없어 보유한 SD1.5 체크포인트를 쓴다(바뀌면 CKPT 만 고친다).
# 512 로 그린 뒤 2배로 다시 그려(1024) 선명도를 올리고, InspyrenetRembg 로 배경을 빼 RGBA 로 저장한다.
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 = "GameIconInstitute_anfg_v3.safetensors"
TRIGGER = "game icon institute, game icon, ahfg"
NEG = "black, dark, low resolution, worst quality, low quality, blurry, text, watermark, signature, jpeg artifacts, cropped, out of frame, person, hands, multiple views"

# FR 장비 전 부위 — 페이지 예시처럼 대상 낱말만 담백하게 쓴다(학습 태그에 있던 낱말을 우선: ring · shield · gem · armor · helmet · gauntlet · boots · sword · belt · necklace)
SUBJECTS = [
    ("sword", "sword, longsword"),
    ("blunt", "mace, war hammer, blunt weapon"),
    ("bow", "bow, longbow"),
    ("dagger", "dagger"),
    ("staff", "magic staff, wizard staff with crystal"),
    ("magic_dagger", "magic dagger, enchanted dagger, glowing"),
    ("armor", "armor, breastplate"),
    ("belt", "belt, leather belt with metal buckle"),
    ("gloves", "gloves, gauntlets"),
    ("boots", "boots, armored boots"),
    ("ring", "ring, gem"),
    ("necklace", "necklace, pendant, gem"),
    ("shield", "shield"),
    ("quiver", "quiver, leather quiver with arrows inside"),
    ("spellbook", "book, magic book, closed spellbook"),
    ("throwing", "throwing knife, kunai"),
    ("orb", "orb, crystal ball, glowing"),
    ("offhand_sword", "short sword"),
]

# 2차(다시 돌리기) — 1차에서 형태가 약했던 것만. 직접 보고 고친 점:
#  검: 두 자루가 교차 → "single" 과 부정어 · 둔기: 창처럼 나옴 → 머리 모양을 적음 · 활: 활이 아님 → 시위 · 굽은 몸체를 적음
#  지팡이: 조각남 → 한 자루 · 화살통: 가방처럼 나옴 → 원통 · 깃 · 마법책: 펼쳐짐/상자 → 닫힌 책 한 권 · 투척 무기: 검처럼 나옴 → 작은 표창
RETRY = [
    ("sword", "sword, single one-handed sword, straight blade, crossguard, diagonal"),
    ("blunt", "mace, single flanged mace, heavy round metal head, short handle"),
    ("bow", "bow \\(weapon\\), single recurve bow, curved wooden limbs, bowstring"),
    ("staff", "staff, single long wooden staff, large crystal on top, diagonal"),
    ("magic_dagger", "dagger, single curved dagger, glowing blue runes on blade"),
    ("quiver", "quiver, cylindrical leather arrow quiver, many arrows with feather fletching sticking out of the top, strap"),
    ("spellbook", "book, single closed thick book, leather cover, metal corners, gem on cover"),
    ("throwing", "shuriken, single four-pointed steel throwing star"),
]
RETRY_NEG = ", two weapons, crossed weapons, multiple objects, duplicate, open book, pages, chest, box"

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):
    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": strength, "strength_clip": strength}},
        "3": {"class_type": "CLIPSetLastLayer", "inputs": {"clip": ["2", 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": 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": "dpmpp_2m", "scheduler": "karras", "denoise": 1.0}},
        # 2배 다시 그리기 — 형태는 지키고 질감만 살리도록 약하게
        "10": {"class_type": "LatentUpscaleBy", "inputs": {"samples": ["7", 0], "upscale_method": "bislerp", "scale_by": 2.0}},
        "11": {"class_type": "KSampler", "inputs": {"model": ["2", 0], "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["10", 0], "seed": seed, "steps": 20, "cfg": 7.0, "sampler_name": "dpmpp_2m", "scheduler": "karras", "denoise": 0.42}},
        "8": {"class_type": "VAEDecode", "inputs": {"samples": ["11", 0], "vae": ["1", 2]}},
        "12": {"class_type": "InspyrenetRembg", "inputs": {"image": ["8", 0], "torchscript_jit": "default"}},
        "9": {"class_type": "SaveImage", "inputs": {"images": ["12", 0], "filename_prefix": "gii_site"}},
    }


def run(prompt, seed, strength):
    cid = str(uuid.uuid4())
    pid = post("/prompt", {"prompt": workflow(prompt, seed, strength), "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.5)
    raise TimeoutError("comfy timeout")


def checker(size, a=(58, 58, 58), b=(74, 74, 74), n=16):
    """투명이 보이도록 깔아 두는 체크무늬 바탕"""
    im = Image.new("RGB", (size, size), a)
    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=b)
    return im


def main():
    out, tag = sys.argv[1], sys.argv[2]
    lo, hi = int(sys.argv[3]), int(sys.argv[4])
    strength = float(sys.argv[5]) if len(sys.argv) > 5 else 0.6
    seeds = int(sys.argv[6]) if len(sys.argv) > 6 else 2
    os.makedirs(out, exist_ok=True)
    # round 이름이 w 로 시작하면 다시 돌리기 목록을 쓴다
    global NEG
    retry = tag.startswith("w")
    if retry:
        NEG = NEG + RETRY_NEG
    subs = (RETRY if retry else SUBJECTS)[lo:hi]
    cell, pad, label = 256, 6, 18
    sheet = Image.new("RGB", (len(subs) * (cell + pad) + pad, seeds * (cell + pad) + pad + label), (40, 40, 40))
    draw = ImageDraw.Draw(sheet)
    for i, (name, body) in enumerate(subs):
        prompt = "%s, %s, still life, black background" % (TRIGGER, body)
        for j in range(seeds):
            t0 = time.time()
            im = run(prompt, 3000 + (lo + i) * 13 + j * 211, strength)
            alpha = im.getchannel("A")
            opaque = sum(1 for v in alpha.getdata() if v > 16) / float(alpha.width * alpha.height)
            print("%s %s #%d %.1f초 불투명 %.0f%%" % (tag, name, j, time.time() - t0, opaque * 100)); sys.stdout.flush()
            im.save(os.path.join(out, "%s_%s_%d.png" % (tag, name, j)))
            bg = checker(cell)
            small = im.resize((cell, cell), Image.LANCZOS)
            bg.paste(small, (0, 0), small)
            sheet.paste(bg, (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_%d_%d.png" % (tag, lo, hi))
    sheet.save(path)
    print("sheet", path)


if __name__ == "__main__":
    main()