3.5 KiB
3.5 KiB
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