import gradio as gr
from pathlib import Path
import subprocess, json, tempfile, sys

COMFY_DIR = Path("ComfyUI")
WORKFLOW_JSON = Path("workflow_api.json")   # we’ll create it later

def patch_workflow(wf, pic_path, face_path, prompt):
    for node in wf.values():
        if node.get("_meta", {}).get("title") == "Picture to swap":
            node["inputs"]["image"] = str(pic_path)
        if node.get("_meta", {}).get("title") == "Input face":
            node["inputs"]["image"] = str(face_path)
        if node["class_type"] == "easy positive":
            node["inputs"]["text"] = prompt
    return wf

def run(picture, face, positive):
    tmpdir = Path(tempfile.mkdtemp())
    pic_path  = tmpdir / "pic.png"
    face_path = tmpdir / "face.png"
    picture.save(pic_path)
    face.save(face_path)

    with open(WORKFLOW_JSON) as f:
        wf = json.load(f)
    wf = patch_workflow(wf, pic_path, face_path, positive)
    tmp_wf = tmpdir / "wf.json"
    with open(tmp_wf, "w") as f:
        json.dump(wf, f)

    subprocess.run([
        sys.executable, "ComfyUI/main.py",
        "--disable-auto-launch",
        "--input", str(tmpdir),
        "--output", str(tmpdir),
        "--workflow", str(tmp_wf)
    ], cwd=COMFY_DIR, check=True)

    out = list(tmpdir.glob("*.png"))
    if not out:
        raise RuntimeError("No output image")
    return out[0]

demo = gr.Interface(
    fn=run,
    inputs=[
        gr.Image(type="pil", label="Picture to swap"),
        gr.Image(type="pil", label="Face input"),
        gr.Textbox(label="Positive prompt", lines=3,
                   placeholder="portrait of a woman, detailed face...")
    ],
    outputs=gr.Image(type="filepath", label="Swapped"),
    title="ComfyUI Face Swap"
)
demo.queue().launch()