Files
gateway-ComfyUI-Kie-API/docs/MODEL_IMPLEMENTATION_TEMPLATE.md
T

3.5 KiB
Raw Blame History

i# KIE Model Implementation Template

This document defines the REQUIRED structure and conventions for adding new KIE API models to this repository.

All new models MUST follow this template to ensure:

  • ComfyUI compatibility
  • shared polling & error handling
  • consistent logging
  • minimal duplication

File Placement

Each KIE model MUST live in its own file:

kie_api/<model_name>.py

i Examples:

  • nanobanana.py
  • seedream45.py
  • veo3.py

Required Imports

Models MUST reuse shared helpers. Do NOT duplicate logic.

from typing import Any

from .auth import _load_api_key
from .credits import _log_remaining_credits
from .jobs import create_task, poll_task_until_complete
from .images import image_bytes_to_comfy_tensor
from .http import TransientKieError
from .log import log

Required Constants

Each model MUST define:

MODEL_NAME = "<exact model name from KIE docs>"

Optional:

ASPECT_RATIO_OPTIONS = [...]
RESOLUTION_OPTIONS = [...]
OUTPUT_FORMAT_OPTIONS = [...]

Required Public Entry Function

Each model MUST expose exactly ONE public runner:

def run_<model_name>(
    *,
    prompt: str,
    image_inputs: list[bytes] | None,
    aspect_ratio: str,
    resolution: str,
    output_format: str,
    poll_interval_s: float,
    timeout_s: int,
    log: bool,
) -> torch.Tensor:

Rules:

-MUST validate required inputs
-MUST create task via create_task(...)
-MUST poll via poll_task_until_complete(...)
-MUST return ComfyUI IMAGE tensor (BHWC, float32, 0–1)
-MUST NOT print directly (use log helper)

Logging Rules

-Logging is controlled by log: bool
-Use shared logger only
-Avoid dumping full JSON payloads
-Summarize state transitions only

Good:

[KIE] Task <id> state: waiting (elapsed=30s)


Bad:

[KIE] Full recordInfo JSON: {...huge...}

Error Handling Rules
Use RuntimeError when:

-user input is invalid
-model fails permanently
-decoding output fails
Use TransientKieError when:
-API reports temporary failure
-retry is reasonable

Never swallow exceptions silently.

Image Output Requirements (ComfyUI)

All image outputs MUST:

-be torch.Tensor
-shape: [B, H, W, C]
-dtype: float32
-range: 0.0 – 1.0
-reside on CPU

Use ONLY:

image_bytes_to_comfy_tensor(...)


Forbidden Patterns ❌

-No polling loops inside nodes.py
-No direct requests calls in nodes.py
-No torch.frombuffer on raw bytes
-No duplicated polling code
-No printing directly to console

Node Layer Responsibility

Nodes MUST:

-define INPUT_TYPES
-define RETURN_TYPES
-call exactly one model runner
-never perform HTTP or polling logic

Versioning

Each model file SHOULD include a header comment:

# Model: Seedream 4.5
# API version: 2025-xx
# KIE Docs: <url>


## Video Models

Video-producing models MUST:
- return a local file path (string)
- use kie_api.video helpers
- NOT decode frames
- NOT return tensors

Public runner signature example:

```python
def run_<model>_video(...) -> str:
    """Returns a path to a local video file."""
yaml


This teaches Codex the rule **once**, forever.

---

### 🔹 Step 3 — Create your first KIE video model runner
Now you’re ready for Codex.

📍 **Use VS Code + Codex**

**Prompt (copy exactly):**
```text
Add a KIE video model runner following docs/MODEL_IMPLEMENTATION_TEMPLATE.md.

Tasks:
1. Create kie_api/<model>_video.py
2. Use shared helpers:
   - _poll_task_until_complete
   - _extract_result_urls
   - _download_video
   - _write_video_to_temp_file
3. Return a VIDEO file path (string).
4. Do NOT decode video frames.
5. Show diff before applying.

Always when done add the proper model to nodes.py and adjust the mappings to add this new model