189 lines
3.5 KiB
Markdown
189 lines
3.5 KiB
Markdown
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.
|
||
|
||
```python
|
||
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
|
||
|