imagepipeline/imagepipeline/modules/comfy_flux_edit.py

from __future__ import annotations

import json
import random
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path

from imagepipeline.core.context import ModuleContext
from imagepipeline.core.exceptions import DependencyError
from imagepipeline.core.params import Param
from imagepipeline.modules.ai_base import AIModule
from imagepipeline.modules.registry import register

REPO_ROOT = Path(__file__).resolve().parents[2]


@register
class ComfyFluxEditModule(AIModule):
    name = "comfy_flux_edit"
    description = "Experimental FLUX img2img editing via ComfyUI HTTP API (CPU: very slow)"

    @classmethod
    def parameters(cls) -> dict[str, Param]:
        params = dict(super().parameters())
        params.update(
            {
                "prompt": Param(
                    "string",
                    required=True,
                    help="Edit prompt for the ComfyUI workflow",
                ),
                "denoise": Param(
                    "float",
                    default=0.35,
                    help="KSampler denoise strength",
                ),
                "seed": Param(
                    "int",
                    default=-1,
                    help="Random seed (-1 = random per image)",
                ),
                "server_url": Param(
                    "string",
                    default="http://127.0.0.1:8188",
                    help="ComfyUI server base URL",
                ),
                "workflow_path": Param(
                    "path",
                    default=REPO_ROOT / "workflows" / "comfy" / "flux_klein_edit_api.json",
                    help="ComfyUI workflow exported in API format",
                ),
                "poll_interval": Param(
                    "float",
                    default=2.0,
                    help="Seconds between history polls",
                ),
            }
        )
        return params

    @classmethod
    def check_dependencies(cls) -> None:
        pass

    def run(self, ctx: ModuleContext) -> None:
        workflow_path = Path(ctx.params["workflow_path"])
        if not workflow_path.is_file():
            raise DependencyError(
                f"ComfyUI workflow not found: {workflow_path}. "
                "Export a Flux Klein img2img workflow to this path (see workflows/comfy/README.md)."
            )

        server_url = ctx.params["server_url"].rstrip("/")
        self._ensure_server(server_url)
        if ctx.logger is not None:
            ctx.logger.warning("ComfyUI on CPU: expect hours per full-resolution image")

        workflow_template = json.loads(workflow_path.read_text(encoding="utf-8"))
        denoise = ctx.params["denoise"]
        prompt = ctx.params["prompt"]
        poll_interval = ctx.params["poll_interval"]

        def process(src: Path, dst: Path, _index: int, _total: int) -> None:
            uploaded_name = self._upload_image(server_url, src)
            seed = ctx.params["seed"]
            if seed < 0:
                seed = random.randint(0, 2**32 - 1)
            workflow = self._patch_workflow(
                workflow_template,
                image_name=uploaded_name,
                prompt=prompt,
                denoise=denoise,
                seed=seed,
            )
            prompt_id = self._queue_prompt(server_url, workflow)
            output_info = self._wait_for_output(server_url, prompt_id, poll_interval=poll_interval)
            image_bytes = self._download_view(server_url, output_info)
            dst.write_bytes(image_bytes)

        self.iter_input_images(ctx, process)

    @staticmethod
    def _ensure_server(server_url: str) -> None:
        try:
            urllib.request.urlopen(f"{server_url}/system_stats", timeout=5)
        except urllib.error.URLError as exc:
            raise DependencyError(f"ComfyUI server not reachable at {server_url}: {exc}") from exc

    @staticmethod
    def _upload_image(server_url: str, src: Path) -> str:
        boundary = "----imagepipelineboundary"
        body_parts = [
            f"--{boundary}\r\n".encode(),
            (
                f'Content-Disposition: form-data; name="image"; filename="{src.name}"\r\n'
                f"Content-Type: application/octet-stream\r\n\r\n"
            ).encode(),
            src.read_bytes(),
            b"\r\n",
            f"--{boundary}\r\n".encode(),
            b'Content-Disposition: form-data; name="overwrite"\r\n\r\n',
            b"true\r\n",
            f"--{boundary}--\r\n".encode(),
        ]
        body = b"".join(body_parts)
        request = urllib.request.Request(
            f"{server_url}/upload/image",
            data=body,
            headers={"Content-Type": f"multipart/form-data; boundary={boundary}"},
            method="POST",
        )
        with urllib.request.urlopen(request, timeout=120) as response:
            data = json.loads(response.read().decode("utf-8"))
        return data["name"]

    @staticmethod
    def _patch_workflow(
        template: dict,
        *,
        image_name: str,
        prompt: str,
        denoise: float,
        seed: int,
    ) -> dict:
        workflow = json.loads(json.dumps(template))
        for node in workflow.values():
            if not isinstance(node, dict):
                continue
            class_type = node.get("class_type", "")
            inputs = node.get("inputs", {})
            if class_type == "LoadImage":
                inputs["image"] = image_name
            elif class_type in {"CLIPTextEncode", "TextEncode"} and "text" in inputs:
                inputs["text"] = prompt
            elif class_type == "KSampler":
                inputs["denoise"] = denoise
                inputs["seed"] = seed
            elif "denoise" in inputs:
                inputs["denoise"] = denoise
            if "seed" in inputs and class_type != "KSampler":
                inputs["seed"] = seed
        return workflow

    @staticmethod
    def _queue_prompt(server_url: str, workflow: dict) -> str:
        payload = json.dumps({"prompt": workflow}).encode("utf-8")
        request = urllib.request.Request(
            f"{server_url}/prompt",
            data=payload,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(request, timeout=60) as response:
            data = json.loads(response.read().decode("utf-8"))
        return data["prompt_id"]

    @staticmethod
    def _wait_for_output(
        server_url: str,
        prompt_id: str,
        *,
        poll_interval: float,
        timeout: float = 86400.0,
    ) -> dict:
        deadline = time.time() + timeout
        while time.time() < deadline:
            request = urllib.request.Request(f"{server_url}/history/{prompt_id}")
            with urllib.request.urlopen(request, timeout=30) as response:
                history = json.loads(response.read().decode("utf-8"))
            if prompt_id in history:
                outputs = history[prompt_id].get("outputs") or {}
                for node_output in outputs.values():
                    images = node_output.get("images") or []
                    if images:
                        return images[0]
                status = history[prompt_id].get("status", {})
                if status.get("status_str") == "error":
                    raise RuntimeError(f"ComfyUI workflow failed: {status}")
            time.sleep(poll_interval)
        raise TimeoutError(f"ComfyUI prompt {prompt_id} did not finish within {timeout}s")

    @staticmethod
    def _download_view(server_url: str, image_info: dict) -> bytes:
        query = urllib.parse.urlencode(
            {
                "filename": image_info["filename"],
                "subfolder": image_info.get("subfolder", ""),
                "type": image_info.get("type", "output"),
            }
        )
        with urllib.request.urlopen(f"{server_url}/view?{query}", timeout=120) as response:
            return response.read()