101 lines
2.8 KiB
Plaintext
101 lines
2.8 KiB
Plaintext
import os, json, time, uuid, requests
|
|
from pathlib import Path
|
|
from urllib.parse import quote
|
|
from dotenv import load_dotenv
|
|
from mcp.server.fastmcp import FastMCP
|
|
|
|
load_dotenv()
|
|
|
|
mcp = FastMCP("Alfred ComfyUI")
|
|
BASE = os.getenv("COMFYUI_URL", "http://127.0.0.1:8188").rstrip("/")
|
|
WF = Path(os.getenv("WORKFLOWS_DIR", Path(__file__).parent / "workflows"))
|
|
|
|
|
|
def submit(workflow):
|
|
r = requests.post(
|
|
BASE + "/prompt",
|
|
json={
|
|
"prompt": workflow,
|
|
"client_id": str(uuid.uuid4())
|
|
},
|
|
timeout=30
|
|
)
|
|
r.raise_for_status()
|
|
return r.json()["prompt_id"]
|
|
|
|
|
|
def wait(pid):
|
|
deadline = time.time() + int(os.getenv("COMFYUI_TIMEOUT", "300"))
|
|
while time.time() < deadline:
|
|
h = requests.get(f"{BASE}/history/{pid}", timeout=15).json()
|
|
if pid in h:
|
|
return h[pid]
|
|
time.sleep(1)
|
|
raise TimeoutError("ComfyUI timeout")
|
|
|
|
|
|
def build_image_url(filename, subfolder="", img_type="output"):
|
|
return (
|
|
f"{BASE}/view?"
|
|
f"filename={quote(str(filename))}"
|
|
f"&subfolder={quote(str(subfolder))}"
|
|
f"&type={quote(str(img_type))}"
|
|
)
|
|
|
|
|
|
def extract_images(history):
|
|
images = []
|
|
|
|
outputs = history.get("outputs", {})
|
|
for node_id, node_output in outputs.items():
|
|
for img in node_output.get("images", []):
|
|
filename = img.get("filename")
|
|
subfolder = img.get("subfolder", "")
|
|
img_type = img.get("type", "output")
|
|
|
|
if not filename:
|
|
continue
|
|
|
|
images.append({
|
|
"node_id": node_id,
|
|
"filename": filename,
|
|
"subfolder": subfolder,
|
|
"type": img_type,
|
|
"url": build_image_url(filename, subfolder, img_type)
|
|
})
|
|
|
|
return images
|
|
|
|
|
|
@mcp.tool()
|
|
def generate_image(prompt: str):
|
|
"""Generate an image with the locked text_to_image workflow."""
|
|
w = json.loads((WF / "text_to_image.json").read_text())
|
|
node = os.getenv("COMFY_PROMPT_NODE", "6")
|
|
w[node]["inputs"]["text"] = prompt
|
|
|
|
history = wait(submit(w))
|
|
images = extract_images(history)
|
|
|
|
return {
|
|
"ok": True,
|
|
"message": f"Image générée pour : {prompt}",
|
|
"prompt": prompt,
|
|
"images": images,
|
|
"first_image_url": images[0]["url"] if images else None,
|
|
"first_image_filename": images[0]["filename"] if images else None
|
|
}
|
|
|
|
|
|
@mcp.tool()
|
|
def analyze_image(image_path: str, question: str):
|
|
"""Analyze a camera/local image using the locked VLM workflow. Workflow node IDs are configured in .env."""
|
|
w = json.loads((WF / "vision_analyze.json").read_text())
|
|
w[os.getenv("COMFY_VISION_IMAGE_NODE", "1")]["inputs"]["image"] = image_path
|
|
w[os.getenv("COMFY_VISION_PROMPT_NODE", "2")]["inputs"]["text"] = question
|
|
return wait(submit(w))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
mcp.run()
|