aigc-home/e2e/fake_comfy.py

"""Standalone fake ComfyUI for HEL-654 D-batch / HEL-655 E-batch browser validation.

Mode file e2e/fake_mode.json:
{"mode": "ok"|"fail", "delay_sec": 0, "videos": false,
 "object_info": {class_type: {...}}, "no_outputs": false}
- ok + delay 0: history is ready immediately (job completes fast)
- ok + delay N: history appears N seconds after /prompt (progress state)
- fail: /prompt returns 500 (job fails with a Chinese error from the worker)
- object_info: when present, /object_info serves this catalogue verbatim —
  the E-batch script builds it from the workflow it imports so the L2 check
  can be driven all-green (every class present) or missing-node (catalogue
  omits classes). Absent → {} (every node "missing").
- no_outputs: history entries complete with no files (test-run no-output pass)
- video_file: absolute path of an mp4 served (and named) as the video output
  instead of tiny.mp4 (HEL-669 last-frame export runs real clips through it)
"""
from __future__ import annotations
import json, threading, uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

ROOT = Path(__file__).parent
MODE_FILE = ROOT / "fake_mode.json"
MP4 = ROOT / "tiny.mp4"

PNG_1PX = bytes.fromhex(
    "89504e470d0a1a0a0000000d49484452000000010000000108060000001f15c489"
    "0000000d4944415478da63f8cfc0f01f0005000105ab5e30760000000049454e44ae426082"
)

KNOWN = {"SaveImage", "PreviewImage", "SaveImageAdvanced", "SaveImageWebsocket",
         "SaveVideo", "VideoOutput", "VHS_VideoCombine", "AnimateDiffCombine"}


def read_mode() -> dict:
    try:
        return json.loads(MODE_FILE.read_text())
    except Exception:
        return {"mode": "ok", "delay_sec": 0}


class Handler(BaseHTTPRequestHandler):
    server_version = "FakeComfy/1.0"

    def do_GET(self):
        if self.path == "/system_stats":
            self._json(200, {"system": {"comfyui_version": "0.3.72-fake", "os": "linux"},
                             "devices": [{"name": "FakeGPU", "type": "cuda"}]})
            return
        if self.path == "/object_info":
            self._json(200, read_mode().get("object_info") or {})
            return
        if self.path == "/queue":
            # Prompts accepted but not yet in history: first one "running",
            # the rest pending — so the worker walks comfy_queued →
            # comfy_running like a real ComfyUI.
            live = [p for p in self.server.accepted if p not in self.server.history]
            running = [[1, live[0]]] if live else []
            pending = [[2, p] for p in live[1:]]
            self._json(200, {"queue_running": running, "queue_pending": pending})
            return
        if self.path.startswith("/history/"):
            pid = self.path[len("/history/"):]
            entry = self.server.history.get(pid)
            self._json(200, {} if entry is None else {pid: entry})
            return
        if self.path.startswith("/view"):
            if "fmt=mp4" in self.path or ".mp4" in self.path:
                src = read_mode().get("video_file")
                self._raw(200, Path(src).read_bytes() if src else MP4.read_bytes(), "video/mp4")
            else:
                src = read_mode().get("image_file")
                self._raw(200, Path(src).read_bytes() if src else PNG_1PX, "image/png")
            return
        self._json(404, {"error": "not found"})

    def do_POST(self):
        length = int(self.headers.get("Content-Length") or 0)
        body = self.rfile.read(length) if length else b""
        if self.path == "/prompt":
            mode = read_mode()
            if mode.get("mode") == "fail":
                self._json(500, {"error": "boom", "node_errors": []})
                return
            pid = uuid.uuid4().hex
            self.server.accepted.append(pid)
            videos = mode.get("videos") is True
            no_outputs = mode.get("no_outputs") is True
            video_file = mode.get("video_file")
            out_name = (Path(video_file).stem + ".mp4" if video_file else "out-1.mp4") if videos else "out-1.png"
            if mode.get("unique_filename"):
                out_name = pid + "-" + out_name
            key = "videos" if videos else "images"
            outputs = {} if no_outputs else {
                self._output_node(body): {key: [{"filename": out_name, "subfolder": "", "type": "output"}]}
            }
            entry = {
                "status": {"status_str": "success", "completed": True},
                "outputs": outputs,
                "prompt": [],
            }
            delay = float(mode.get("delay_sec") or 0)
            if delay <= 0:
                self.server.history[pid] = entry
            else:
                threading.Timer(delay, lambda: self.server.history.setdefault(pid, entry)).start()
            self.server.prompts.append(json.loads(body.decode("utf-8") or "{}"))
            self._json(200, {"prompt_id": pid})
            return
        if self.path == "/upload/image":
            self.server.uploads.append(body)
            self._json(200, {"name": "uploaded-input.png", "subfolder": "calliope"})
            return
        self._json(404, {"error": "not found"})

    def _output_node(self, body: bytes) -> str:
        try:
            prompt = json.loads(body.decode("utf-8") or "{}").get("prompt") or {}
        except json.JSONDecodeError:
            prompt = {}
        for nid, node in prompt.items():
            if isinstance(node, dict) and node.get("class_type") in KNOWN:
                return str(nid)
        return "9"

    def _json(self, status, payload):
        data = json.dumps(payload).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def _raw(self, status, data, ctype):
        self.send_response(status)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def log_message(self, fmt, *args):
        pass


if __name__ == "__main__":
    server = ThreadingHTTPServer(("127.0.0.1", 18188), Handler)
    server.history = {}
    server.prompts = []
    server.uploads = []
    server.accepted = []
    print("fake comfy on http://127.0.0.1:18188", flush=True)
    server.serve_forever()