handler.py

import base64
import json
import os
import subprocess
import time
import uuid
from io import BytesIO
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

import requests
import runpod
from PIL import Image


COMFY_DIR = Path(os.environ.get("COMFY_DIR", "/workspace/ComfyUI"))
COMFY_HOST = os.environ.get("COMFY_HOST", "127.0.0.1")
COMFY_PORT = int(os.environ.get("COMFY_PORT", "8188"))
COMFY_BASE_URL = f"http://{COMFY_HOST}:{COMFY_PORT}"

WORKFLOW_T2V_PATH = Path(os.environ.get("WORKFLOW_T2V_PATH", "/app/workflows/ltx2_t2v_full.json"))
WORKFLOW_I2V_PATH = Path(os.environ.get("WORKFLOW_I2V_PATH", "/app/workflows/ltx2_i2v_full.json"))

COMFY_OUTPUT_DIR = Path(os.environ.get("COMFY_OUTPUT_DIR", str(COMFY_DIR / "output")))
COMFY_INPUT_DIR = Path(os.environ.get("COMFY_INPUT_DIR", str(COMFY_DIR / "input")))


_comfy_proc: Optional[subprocess.Popen] = None


class UserError(Exception):
    pass


def _strip_data_url_prefix(b64_or_data_url: str) -> str:
    if "," in b64_or_data_url and b64_or_data_url.strip().lower().startswith("data:"):
        return b64_or_data_url.split(",", 1)[1]
    return b64_or_data_url


def _ensure_comfy_running() -> None:
    global _comfy_proc

    try:
        r = requests.get(f"{COMFY_BASE_URL}/system_stats", timeout=2)
        if r.ok:
            return
    except Exception:
        pass

    if _comfy_proc is not None and _comfy_proc.poll() is None:
        # Process exists but server might still be starting.
        _wait_for_comfy_ready()
        return

    cmd = [
        "python",
        str(COMFY_DIR / "main.py"),
        "--listen",
        COMFY_HOST,
        "--port",
        str(COMFY_PORT),
    ]
    env = dict(os.environ)
    env.setdefault("PYTHONUNBUFFERED", "1")
    env.setdefault("COMFYUI_PATH", str(COMFY_DIR))

    _comfy_proc = subprocess.Popen(
        cmd,
        cwd=str(COMFY_DIR),
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )

    _wait_for_comfy_ready()


def _wait_for_comfy_ready(timeout_s: int = 120) -> None:
    start = time.time()
    last_err = None
    while time.time() - start < timeout_s:
        try:
            r = requests.get(f"{COMFY_BASE_URL}/system_stats", timeout=2)
            if r.ok:
                return
            last_err = f"status={r.status_code}"
        except Exception as e:
            last_err = str(e)
        time.sleep(1)
    raise RuntimeError(f"ComfyUI did not become ready in {timeout_s}s (last_err={last_err})")


def _load_workflow(path: Path) -> Dict[str, Any]:
    """
    Load a ComfyUI workflow JSON.

    This function is aware of two shapes:
    - Raw /prompt payloads (mapping node_id -> node dict)
    - Full exported workflow templates from ComfyUI which store the
      backend prompt under extra.prompt.
    """
    if not path.exists():
        raise UserError(f"Workflow file not found: {path}")

    data = json.loads(path.read_text(encoding="utf-8"))

    # If this looks like a full workflow template, extract the prompt section.
    extra = data.get("extra")
    if isinstance(extra, dict) and isinstance(extra.get("prompt"), dict):
        return extra["prompt"]

    # Otherwise assume it's already a /prompt payload.
    return data


def _pick_workflow(mode: str) -> Path:
    if mode == "t2v":
        return WORKFLOW_T2V_PATH
    if mode == "i2v":
        return WORKFLOW_I2V_PATH
    raise UserError('`mode` must be "t2v" or "i2v"')


def _write_input_image(input_obj: Dict[str, Any]) -> str:
    """
    Writes an input image into ComfyUI input dir and returns the filename.
    This function does NOT know how the workflow references the image; handler must patch workflow accordingly.
    """
    image_path = input_obj.get("image_path")
    image_url = input_obj.get("image_url")
    image_base64 = input_obj.get("image_base64")

    provided = [v for v in [image_path, image_url, image_base64] if v]
    if len(provided) != 1:
        raise UserError("For mode=i2v, provide exactly one of image_path, image_url, image_base64")

    COMFY_INPUT_DIR.mkdir(parents=True, exist_ok=True)
    out_name = f"input_{uuid.uuid4().hex}.png"
    out_path = COMFY_INPUT_DIR / out_name

    if image_path:
        src = Path(image_path)
        if not src.exists():
            raise UserError(f"image_path does not exist: {image_path}")
        img = Image.open(src).convert("RGB")
        img.save(out_path, format="PNG")
        return out_name

    if image_url:
        r = requests.get(image_url, timeout=30)
        r.raise_for_status()
        img = Image.open(BytesIO(r.content)).convert("RGB")  # type: ignore[name-defined]
        img.save(out_path, format="PNG")
        return out_name

    # image_base64
    raw = base64.b64decode(_strip_data_url_prefix(str(image_base64)))
    img = Image.open(BytesIO(raw)).convert("RGB")  # type: ignore[name-defined]
    img.save(out_path, format="PNG")
    return out_name


def _submit_prompt(prompt: Dict[str, Any]) -> str:
    r = requests.post(f"{COMFY_BASE_URL}/prompt", json={"prompt": prompt}, timeout=30)
    r.raise_for_status()
    data = r.json()
    prompt_id = data.get("prompt_id")
    if not prompt_id:
        raise RuntimeError(f"ComfyUI /prompt response missing prompt_id: {data}")
    return str(prompt_id)


def _wait_for_completion(prompt_id: str, timeout_s: int = 3600) -> Dict[str, Any]:
    start = time.time()
    while time.time() - start < timeout_s:
        r = requests.get(f"{COMFY_BASE_URL}/history/{prompt_id}", timeout=30)
        if r.status_code == 404:
            time.sleep(1)
            continue
        r.raise_for_status()
        history = r.json()
        if prompt_id in history:
            return history[prompt_id]
        time.sleep(1)
    raise RuntimeError(f"Timed out waiting for ComfyUI prompt_id={prompt_id}")


def _extract_first_mp4(history_item: Dict[str, Any]) -> Tuple[Path, str]:
    """
    Extract the first mp4 output file from ComfyUI history.
    Returns (path, filename).
    """
    outputs = history_item.get("outputs") or {}
    for _node_id, out in outputs.items():
        for key in ("gifs", "videos", "files", "images"):
            items = out.get(key)
            if not items:
                continue
            for it in items:
                filename = it.get("filename")
                if not filename:
                    continue
                if filename.lower().endswith(".mp4"):
                    p = COMFY_OUTPUT_DIR / filename
                    return p, filename
    raise RuntimeError("No mp4 output found in ComfyUI history outputs")


def _base64_mp4(path: Path) -> str:
    if not path.exists():
        raise RuntimeError(f"Expected output file not found: {path}")
    raw = path.read_bytes()
    return "data:video/mp4;base64," + base64.b64encode(raw).decode("utf-8")


def _patch_workflow_inputs(workflow: Dict[str, Any], input_obj: Dict[str, Any]) -> Dict[str, Any]:
    """
    Patch the LTX-2 workflows' logical input nodes so they can be driven
    entirely from the serverless input payload.

    The IDs used here are taken from the official ComfyUI workflow templates:
    - Text-to-video:  video_ltx2_t2v.json
    - Image-to-video: video_ltx2_i2v.json
    """

    mode = str(input_obj.get("mode") or "").strip()
    prompt_text = input_obj.get("prompt")
    if not prompt_text:
        raise UserError("`prompt` is required")

    negative = (
        input_obj.get("negative_prompt")
        or input_obj.get("negative")
        or "blurry, low quality, still frame, frames, watermark, overlay, titles, has blurbox, has subtitles"
    )

    seed = input_obj.get("seed")
    width = input_obj.get("width")
    height = input_obj.get("height")
    fps = input_obj.get("fps")
    num_frames = input_obj.get("num_frames")
    ckpt_name = input_obj.get("ckpt_name") or input_obj.get("checkpoint")

    wf = workflow  # mutate in-place; caller passes fresh dict from _load_workflow

    def _ensure_node(node_id: str, expected_type: Optional[str] = None) -> Dict[str, Any]:
        node = wf.get(node_id)
        if not isinstance(node, dict):
            raise UserError(f"Workflow missing expected node {node_id}")
        if expected_type and node.get("class_type") != expected_type:
            raise UserError(
                f"Workflow node {node_id} is {node.get('class_type')}, "
                f"expected {expected_type}. Make sure you are using the bundled LTX-2 templates."
            )
        if "inputs" not in node or not isinstance(node["inputs"], dict):
            raise UserError(f"Workflow node {node_id} has no editable inputs")
        return node

    # Shared patching (both modes use these same logical nodes in the templates)
    # Main text prompt (node "3": CLIPTextEncode)
    n3 = _ensure_node("3", "CLIPTextEncode")
    n3["inputs"]["text"] = str(prompt_text)

    # Negative prompt (node "4": CLIPTextEncode)
    n4 = _ensure_node("4", "CLIPTextEncode")
    n4["inputs"]["text"] = str(negative)

    # Seed (node "11": RandomNoise)
    if seed is not None:
        n11 = _ensure_node("11", "RandomNoise")
        n11["inputs"]["noise_seed"] = int(seed)

    # FPS (node "23": FloatConstant)
    if fps is not None:
        n23 = _ensure_node("23")
        n23_inputs = n23.setdefault("inputs", {})
        n23_inputs["value"] = float(fps)

        # LTXVConditioning + CreateVideo also take frame_rate/fps but those are
        # wired from node 23 in the templates, so we don't need to touch them here.

    # Number of frames (node "27": INTConstant)
    if num_frames is not None:
        n27 = _ensure_node("27")
        n27_inputs = n27.setdefault("inputs", {})
        n27_inputs["value"] = int(num_frames)

    # Resolution (node "43": EmptyLTXVLatentVideo)
    if width is not None or height is not None:
        n43 = _ensure_node("43", "EmptyLTXVLatentVideo")
        if width is not None:
            n43["inputs"]["width"] = int(width)
        if height is not None:
            n43["inputs"]["height"] = int(height)

    # Checkpoint (node "1": CheckpointLoaderSimple)
    if ckpt_name:
        n1 = _ensure_node("1", "CheckpointLoaderSimple")
        n1["inputs"]["ckpt_name"] = str(ckpt_name)

    # Image-to-video specific patching: wire the uploaded image into LoadImage.
    if mode == "i2v":
        # _write_input_image() will have written the image into COMFY_INPUT_DIR
        # using a random name; we just need to patch the loader to use it.
        image_filename = input_obj.get("_comfy_image_filename")
        if image_filename:
            # Top-level UI workflow uses a LoadImage node with widgets_values[0] = filename.
            # The backend prompt for i2v uses an image input that ultimately comes from this node.
            # Since /prompt only cares about node.inputs, we update both defensively.
            # Try prompt-style node first (if present).
            load_node = wf.get("98")
            if isinstance(load_node, Dict):
                # Some forks represent the path in inputs["image"], some only via widgets_values.
                inputs = load_node.setdefault("inputs", {})
                if "image" in inputs:
                    inputs["image"] = image_filename

        # If, for some reason, the workflow did not include the LoadImage node in the
        # prompt graph, we still validated that the image exists on disk; the template's
        # default example image will be used instead.

    return wf


def handler(event: Dict[str, Any]) -> Dict[str, Any]:
    try:
        input_obj = event.get("input") or {}
        mode = str(input_obj.get("mode") or "").strip()
        wf_path = _pick_workflow(mode)
        workflow = _load_workflow(wf_path)

        _ensure_comfy_running()

        # If I2V, store image in ComfyUI input dir so the workflow can reference it.
        if mode == "i2v":
            filename = _write_input_image(input_obj)
            # Make the filename available to the patcher without mutating the original event.
            input_obj = {**input_obj, "_comfy_image_filename": filename}

        workflow = _patch_workflow_inputs(workflow, input_obj)

        prompt_id = _submit_prompt(workflow)
        history_item = _wait_for_completion(prompt_id)
        out_path, _filename = _extract_first_mp4(history_item)
        return {"video": _base64_mp4(out_path)}
    except UserError as e:
        return {"error": str(e)}
    except Exception as e:
        return {"error": f"{type(e).__name__}: {e}"}


runpod.serverless.start({"handler": handler})