# Krea 2 Eric Insert: photo + request -> the same photo with er1cshen added.
# The pipeline is workflow_api.json (open it in ComfyUI to see it), run by
# ComfyUI's own executor. This file only feeds it: the photo, the request, the
# seed, the checked placement phrase; plus resizing to 2 MP and an NSFW check
# (Krea 2 license §4.2). Models load once at startup, as ZeroGPU recommends:
# https://huggingface.co/blog/run-comfyui-workflows-on-spaces
import copy
import json
import os
import re
import sys
import tempfile
import uuid

import gradio as gr
import spaces
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from transformers import pipeline

HERE = os.path.dirname(os.path.abspath(__file__))
COMFYUI = os.path.join(HERE, "ComfyUI")
MODELS = os.path.join(COMFYUI, "models")
PIXELS = 2 * 1024 * 1024  # every photo is resized to this; smaller faces come out with artifacts
DEFAULT_PLACE = "standing next to the person"
EMPTY_PLACE = "standing in the scene"  # used when the photo has no people, so the model doesn't invent one to stand next to
PERSON_WORDS = re.compile(r"\b(him|her|them|his|their|person|people|next to|beside)\b")
STARTS = {"on", "in", "at", "behind", "beside", "next", "near", "between", "under", "above", "by", "with", "from", "inside", "outside", "across", "against", "along", "around", "among", "to"}
CHANGE_WORDS = re.compile(r"\b(turn|turning|transform|transforming|into|remove|removing|replace|replacing|change|changing|make|making)\b")

for f in (
    "diffusion_models/krea2_turbo_bf16.safetensors",
    "text_encoders/qwen3vl_4b_bf16.safetensors",
    "vae/qwen_image_vae.safetensors",
):
    hf_hub_download("Comfy-Org/Krea-2", f, local_dir=MODELS)
hf_hub_download(
    "yuncuns/krea2-eric-insert",
    "krea2-eric-insert.safetensors",
    local_dir=os.path.join(MODELS, "loras"),
)

sys.path.insert(0, COMFYUI)
import asyncio  # noqa: E402

import execution  # noqa: E402
import folder_paths  # noqa: E402
import nodes  # noqa: E402
import server  # noqa: E402
from comfy import model_management  # noqa: E402

loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
server_instance = server.PromptServer(loop)
execution.PromptQueue(server_instance)
asyncio.run(nodes.init_extra_nodes())
# Same arguments ComfyUI's main.py passes for --cache-classic.
executor = execution.PromptExecutor(
    server_instance, cache_type=execution.CacheType.CLASSIC, cache_args={"lru": 0, "ram": 0, "ram_inactive": 0}
)

WORKFLOW = json.load(open(os.path.join(HERE, "workflow_api.json")))


def node(title: str) -> str:
    return next(k for k, v in WORKFLOW.items() if v["_meta"]["title"] == title)


PHOTO, REQUEST, REWRITE_OUT, PLACEMENT, PROMPT_START, PROMPT_END, SAMPLER, OUTPUT = map(
    node, ("photo", "request", "rewrite output", "placement", "prompt start", "prompt end", "sampler", "output")
)


def run(prompt: dict, outputs: list[str]) -> dict:
    executor.execute(prompt, str(uuid.uuid4()), {}, outputs)
    if not executor.success:
        raise gr.Error(f"Generation failed: {executor.status_messages[-1:]}")
    return executor.history_result["outputs"]


class GPUPreload:
    """Startup-only output node: loads the graph's models onto the GPU once."""

    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"model": ("MODEL",), "clip": ("CLIP",), "vae": ("VAE",)}}

    RETURN_TYPES = ()
    FUNCTION = "load"
    OUTPUT_NODE = True

    def load(self, model, clip, vae):
        model_management.load_models_gpu([model, clip.patcher, vae.patcher])
        return {}


nodes.NODE_CLASS_MAPPINGS["GPUPreload"] = GPUPreload
patch_inputs = WORKFLOW[node("Krea2 Edit (source patch)")]["inputs"]
encode_inputs = WORKFLOW[node("Krea2 Edit (grounded encode)")]["inputs"]
warmup = copy.deepcopy(WORKFLOW)
warmup["preload"] = {
    "class_type": "GPUPreload",
    "inputs": {"model": patch_inputs["model"], "clip": encode_inputs["clip"], "vae": patch_inputs["vae"]},
    "_meta": {"title": "preload"},
}
# The executor always runs under torch.inference_mode(); weights loaded that way
# become inference tensors that ZeroGPU cannot hand to its GPU worker
# ("Inference tensors do not track version counter"). Load them with it off.
_inference_mode = torch.inference_mode
torch.inference_mode = lambda *args, **kwargs: _inference_mode(False)
try:
    run(warmup, ["preload"])
finally:
    torch.inference_mode = _inference_mode

nsfw = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
people = pipeline("object-detection", model="hustvl/yolos-tiny", device="cpu")


def check_placement(text: str) -> str:
    phrase = text.split("</think>")[-1].strip().strip(".\"'").lower()
    first = phrase.split(" ", 1)[0]
    if (
        not re.fullmatch(r"[a-z][a-z ,'-]{2,80}", phrase)
        or not (first.endswith("ing") or first in STARTS)
        or CHANGE_WORDS.search(phrase)
    ):
        return DEFAULT_PLACE
    return phrase


@spaces.GPU(duration=75)
def generate(image: Image.Image, request: str, seed: int) -> tuple[Image.Image, str]:
    if image is None:
        raise gr.Error("Add a photo.")
    image = image.convert("RGB")
    # sides must be multiples of 16, or the model leaves a strip of noise along the edge
    scale = (PIXELS / (image.width * image.height)) ** 0.5
    image = image.resize((int(image.width * scale) // 16 * 16, int(image.height * scale) // 16 * 16), Image.LANCZOS)
    with tempfile.NamedTemporaryFile(
        suffix=".png", dir=folder_paths.get_input_directory(), delete=False
    ) as f:
        image.save(f.name)
        path = f.name

    prompt = copy.deepcopy(WORKFLOW)
    prompt[REQUEST]["inputs"]["value"] = request.strip()[:300]
    text = run(prompt, [REWRITE_OUT])[REWRITE_OUT]["text"][0]
    placement = check_placement(text)
    if PERSON_WORDS.search(placement) and not any(d["label"] == "person" for d in people(image, threshold=0.7)):
        placement = EMPTY_PLACE
    prompt[PLACEMENT]["inputs"]["value"] = placement
    # the prompt the edit model gets, shown under the result and in the logs (the workflow joins the same three parts)
    final = prompt[PROMPT_START]["inputs"]["value"] + placement + prompt[PROMPT_END]["inputs"]["value"]
    print(f"seed {seed} · request: {request.strip()[:300]}\nprompt: {final}", flush=True)
    prompt[PHOTO]["inputs"]["image"] = os.path.basename(path)
    prompt[SAMPLER]["inputs"]["seed"] = int(seed)
    result = run(prompt, [OUTPUT])[OUTPUT]["images"][0]
    os.unlink(path)

    out_path = os.path.join(folder_paths.get_temp_directory(), result["subfolder"], result["filename"])
    out = Image.open(out_path).convert("RGB")
    os.unlink(out_path)
    if any(r["label"] == "nsfw" and r["score"] > 0.5 for r in nsfw(out)):
        raise gr.Error("Blocked by the content filter.")
    return out, final


REQUEST_HINT = WORKFLOW[REQUEST]["inputs"]["value"]
SEED = WORKFLOW[SAMPLER]["inputs"]["seed"]

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Image(type="pil", label="Photo"),
        gr.Textbox(value=REQUEST_HINT, label="Prompt", lines=2),
        gr.Number(value=SEED, label="Seed", precision=0),
    ],
    outputs=[gr.Image(label="Result", format="png"), gr.Textbox(label="Final prompt", interactive=False)],
    examples=[
        ["reagan.png", REQUEST_HINT, SEED],
        ["einstein.jpg", REQUEST_HINT, SEED],
        ["obama-biden.jpg", "insert a ridiculously handsome man walking with them", SEED],
        ["apollo.jpg", "insert a ridiculously handsome man", SEED],
    ],
    cache_examples=False,
    flagging_mode="never",
)

if __name__ == "__main__":
    demo.launch()