Initial release: RogoAI Prompt Relay for Wan2.2 I2V

This commit is contained in:
nonioke
2026-04-29 21:18:05 +09:00
commit cb5c232af6
6 changed files with 561 additions and 0 deletions
+1
View File
@@ -0,0 +1 @@
__pycache__/
+220
View File
@@ -0,0 +1,220 @@
# ComfyUI-RogoAI-PromptRelay
**Wan2.2 I2V** と **LTX-Video 2.3** に対応した Prompt Relay ComfyUI カスタムノードです。
複数シーンを **Latent 空間で直接連結**することで、pixel連結方式と比べてキャラクターのドリフトを抑制し、長尺動画のIP一貫性を向上させます。
![RogoAI](https://img.shields.io/badge/RogoAI-Takeji-blue)
![ComfyUI](https://img.shields.io/badge/ComfyUI-Custom%20Node-orange)
![License](https://img.shields.io/badge/License-MIT-green)
---
## Prompt Relay とは
[Prompt Relay (arXiv:2604.10030)](https://arxiv.org/abs/2604.10030) は、NTU S-Lab の Gordon Chen らが提案した推論時プロンプトルーティング手法です。
動画生成モデルの **Cross-Attention** にペナルティ項 **C** を注入することで、各テキストプロンプトが担当する時間セグメントにのみ影響を与えます。
```
Attn(Q, K, V) = softmax(QKᵀ/√d - C(Q, K)) V
```
- **グローバルプロンプト** → ペナルティなし → 全フレームに影響
- **ローカルプロンプト** → ペナルティあり → 担当セグメントのみに影響
---
## Pixel連携 vs Latent連携
| | Pixel連携(従来) | Latent連携(Prompt Relay) |
|---|---|---|
| シーン接続 | VAEデコード → pixel → VAEエンコード | Latent空間で直接連結 |
| 情報ロス | VAE往復で劣化あり | 劣化なし |
| キャラドリフト | 蓄積しやすい | 抑制される |
| シーン境界 | ハードカット | スムーズなブレンド |
---
## 対応モデル
| モデル | フレーム数 | シーン数 | 備考 |
|--------|-----------|---------|------|
| **Wan2.2 I2V** | 81フレーム(約3秒) | 2〜3シーン | WanVideoWrapper経由 |
| **LTX-Video 2.3** | 最大993フレーム(約40秒) | 5〜8シーン以上 | Kijai workflow推奨 |
> **注意**: Wan2.2はWanVideoWrapperの81フレーム上限により、長尺のPrompt Relayには不向きです。長尺動画にはLTX-Video 2.3をお勧めします。
---
## ノード一覧
### `RogoAI Prompt Relay (Wan2.2 I2V)`
Wan2.2 I2V 用の Prompt Relay エンコードノードです。
**入力:**
| パラメータ | 型 | 説明 |
|-----------|-----|------|
| model | WANVIDEOMODEL | Wan2.2モデル |
| t5 | WANTEXTENCODER | T5テキストエンコーダー |
| global_prompt | STRING | 全シーン共通プロンプト(キャラクター・場所・照明) |
| negative_prompt | STRING | ネガティブプロンプト |
| local_prompt_1〜6 | STRING | 各シーンのプロンプト(1〜6シーン) |
| segment_lengths | STRING | 各シーンのlatentフレーム数(例: `7,7,7`) |
| frame_num | INT | 総フレーム数(デフォルト: 81) |
| width / height | INT | 解像度(デフォルト: 832×480) |
| epsilon | FLOAT | セグメント境界の減衰係数(デフォルト: 0.001) |
| force_offload | BOOLEAN | T5をオフロードするか |
**出力:**
| 出力 | 型 | 説明 |
|------|-----|------|
| model | WANVIDEOMODEL | Prompt Relayパッチ済みモデル |
| text_embeds | WANVIDEOTEXTEMBEDS | テキスト埋め込み |
---
### `RogoAI Prompt Relay Seg Calc`
フレーム数とシーン数からセグメント長を自動計算する補助ノードです。
**入力:**
| パラメータ | 型 | 説明 |
|-----------|-----|------|
| frame_num | INT | 総フレーム数 |
| num_segments | INT | シーン数(1〜6) |
**出力:**
| 出力 | 型 | 説明 |
|------|-----|------|
| segment_lengths | STRING | 自動計算されたセグメント長(例: `7,7,7`) |
| latent_frames | INT | latentフレーム総数 |
---
## インストール
### 方法1: ComfyUI Manager(推奨)
1. ComfyUI Manager を開く
2.「Custom Nodes」→「Install via Git URL」
3. 以下のURLを入力:
```
https://github.com/RogoAI-Takeji/ComfyUI-RogoAI-PromptRelay
```
### 方法2: 手動インストール
```powershell
cd <ComfyUIフォルダ>\custom_nodes
git clone https://github.com/RogoAI-Takeji/ComfyUI-RogoAI-PromptRelay
```
依存ライブラリは ComfyUI-WanVideoWrapper に含まれているため、追加インストールは不要です。
---
## 使い方
### 基本ワークフロー(Wan2.2)
```
[Load Wan Model] ──→ [RogoAI Prompt Relay (Wan2.2 I2V)] ──→ [WanVideo Sampler]
[Load T5] ──→ ↓
[Load Image] ──────────────────────────────────────→ [WanVideo I2V Encode]
```
### プロンプト設計のコツ
**グローバルプロンプト(短く・一貫性要素のみ):**
```
A young woman with long black wavy hair and blue headphones,
sitting at a PC desk, realistic cinematic lighting.
```
**ローカルプロンプト(動作・カメラのみ):**
```
Scene 1: She faces the camera directly with a calm expression.
Scene 2: She turns her head away, looking at the monitor.
Scene 3: She turns back to the camera with a slight smile.
```
### segment_lengths の計算
```
81フレーム → latent_frames = (81-1)÷4+1 = 21
3シーンで均等: 7,7,7
2シーンで均等: 11,10
※ RogoAI Prompt Relay Seg Calc ノードで自動計算できます
```
### 参照画像について
I2V の参照画像(first-frame)は **VAE latent として全フレームに影響**します。
- キャラクターの顔・外見一貫性のために**クローズアップ画像を推奨**
- グローバルプロンプトと参照画像の組み合わせがIP一貫性の二本柱です
---
## LTX-Video 2.3 での利用
LTX-Video 2.3 では [Kijai の ComfyUI-LTXVideo](https://github.com/kijai/ComfyUI-LTXVideo) の Prompt Relay workflow を使用することを推奨します(2026年4月19日公表)。
LTX-Video 2.3 が Wan2.2 より Prompt Relay の効果が高い理由:
| 要因 | Wan2.2 | LTX-Video 2.3 |
|------|--------|---------------|
| 空間トークン解像度 | 粗い(≈16px/token) | 細かい(≈8px/token) |
| Attention構造 | 空間・時間分離 | Full Spatiotemporal統合 |
| 蒸留 | あり | なし |
| 最大フレーム数 | 81(ComfyUI) | 993(約40秒) |
---
## 制限事項
- Wan2.2は81フレーム上限のため**2〜3シーンが現実的**
- 複数人物の同時アクション制御は困難(Cross-Attentionの空間解像度の制約)
- 激しい動きのシーンはPrompt Relay効果が出にくい場合があります
---
## 関連リポジトリ
- [ComfyUI-RogoAI-APIImageGen](https://github.com/RogoAI-Takeji/ComfyUI-RogoAI-APIImageGen) - Gemini / OpenAI gpt-image-2 画像生成ノード
- [ComfyUI-RogoAI-4DTimelapse](https://github.com/RogoAI-Takeji/ComfyUI-RogoAI-4DTimelapse) - 4Dタイムラプス生成ノード
- [ComfyUI-RogoAI-Music](https://github.com/RogoAI-Takeji/ComfyUI-RogoAI-Music) - ACE-Step音楽生成ノード
---
## 参考文献
```
@article{chen2026prompt,
title={Prompt Relay: Inference-Time Temporal Control for Multi-Event Video Generation},
author={Chen, Gordon and Huang, Ziqi and Liu, Ziwei},
journal={arXiv preprint arXiv:2604.10030},
year={2026}
}
```
---
## ライセンス
MIT License
---
## 老後AI (RogoAI) チャンネル
**「高齢者にはやり直す時間がない」**
YouTube: [老後AI](https://www.youtube.com/@RogoAI)
GitHub: [RogoAI-Takeji](https://github.com/RogoAI-Takeji)
+2
View File
@@ -0,0 +1,2 @@
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
+203
View File
@@ -0,0 +1,203 @@
import torch
import logging
from tqdm import tqdm
from .prompt_relay import (
get_raw_tokenizer,
map_token_indices,
distribute_segment_lengths,
build_segments,
create_mask_fn,
)
log = logging.getLogger(__name__)
VAE_STRIDE_T = 4
PATCH_SIZE_H = 2
PATCH_SIZE_W = 2
def calc_tokens_per_frame(height, width):
return (height // 8 // PATCH_SIZE_H) * (width // 8 // PATCH_SIZE_W)
def calc_latent_frames(frame_num):
return (frame_num - 1) // VAE_STRIDE_T + 1
class RogoAI_PromptRelayWan:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("WANVIDEOMODEL",),
"t5": ("WANTEXTENCODER",),
"global_prompt": ("STRING", {
"multiline": True,
"default": "A young woman with long black wavy hair and blue headphones, sitting at a PC desk, realistic cinematic lighting.",
}),
"negative_prompt": ("STRING", {
"multiline": True,
"default": "blurry, low quality, distorted face",
}),
"local_prompt_1": ("STRING", {"multiline": True, "default": "She faces the camera directly with a calm expression."}),
"local_prompt_2": ("STRING", {"multiline": True, "default": "She turns her head away, looking at the monitor."}),
"local_prompt_3": ("STRING", {"multiline": True, "default": "She turns back to the camera with a slight smile."}),
"segment_lengths": ("STRING", {"default": "7,7,7"}),
"frame_num": ("INT", {"default": 81, "min": 9, "max": 201, "step": 4}),
"width": ("INT", {"default": 832, "min": 64, "max": 2048, "step": 8}),
"height": ("INT", {"default": 480, "min": 64, "max": 2048, "step": 8}),
},
"optional": {
"local_prompt_4": ("STRING", {"multiline": True, "default": ""}),
"local_prompt_5": ("STRING", {"multiline": True, "default": ""}),
"local_prompt_6": ("STRING", {"multiline": True, "default": ""}),
"epsilon": ("FLOAT", {"default": 0.001, "min": 0.0001, "max": 0.1, "step": 0.0001}),
"force_offload": ("BOOLEAN", {"default": True}),
}
}
RETURN_TYPES = ("WANVIDEOMODEL", "WANVIDEOTEXTEMBEDS",)
RETURN_NAMES = ("model", "text_embeds",)
FUNCTION = "apply"
CATEGORY = "RogoAI/PromptRelay"
def apply(self, model, t5,
global_prompt, negative_prompt,
local_prompt_1, local_prompt_2, local_prompt_3,
segment_lengths, frame_num, width, height,
local_prompt_4="", local_prompt_5="", local_prompt_6="",
epsilon=0.001, force_offload=True):
from comfy import model_management as mm
from accelerate.utils import set_module_tensor_to_device
# 1. ローカルプロンプトをまとめる
local_prompts = [p for p in [
local_prompt_1, local_prompt_2, local_prompt_3,
local_prompt_4, local_prompt_5, local_prompt_6
] if p.strip()]
if not local_prompts:
raise ValueError("[RogoAI PromptRelay] local_promptが1つも入力されていません")
# 2. segment_lengthsを解析
latent_frames = calc_latent_frames(frame_num)
seg_lengths_raw = [int(x.strip()) for x in segment_lengths.split(",") if x.strip()]
effective_lengths = distribute_segment_lengths(
len(local_prompts), latent_frames, seg_lengths_raw or None)
log.info("[RogoAI PromptRelay] latent_frames=%d, segments=%s", latent_frames, effective_lengths)
# 3. T5エンコーダーからトークナイザーを取得
encoder = t5["model"]
dtype = t5["dtype"]
raw_tokenizer = None
for attr_name in dir(encoder):
if attr_name.startswith("_"):
continue
inner = getattr(encoder, attr_name, None)
if inner is not None and hasattr(inner, "tokenizer"):
raw_tokenizer = inner.tokenizer
break
if raw_tokenizer is None and hasattr(encoder, "tokenizer"):
raw_tokenizer = encoder.tokenizer
if raw_tokenizer is None:
raise RuntimeError("[RogoAI PromptRelay] T5のトークナイザーが見つかりません")
# 4. full_promptとtoken_rangesを取得
full_prompt, token_ranges = map_token_indices(raw_tokenizer, global_prompt, local_prompts)
log.info("[RogoAI PromptRelay] full_prompt=%d chars, %d segments", len(full_prompt), len(local_prompts))
# 5. セグメントメタデータを構築
q_token_idx = build_segments(token_ranges, effective_lengths, epsilon=epsilon)
# 6. mask_fnを生成してmodel.clone()でパッチ
tokens_per_frame = calc_tokens_per_frame(height, width)
mask_fn = create_mask_fn(q_token_idx, tokens_per_frame, latent_frames)
patcher = model.clone()
if "transformer_options" not in patcher.model_options:
patcher.model_options["transformer_options"] = {}
patcher.model_options["transformer_options"]["prompt_relay_mask_fn"] = mask_fn
log.info("[RogoAI PromptRelay] mask_fn patched via model.clone()")
# 7. T5でテキストエンコード(WanVideoTextEncodeと同じ方式)
device = mm.get_torch_device()
offload_device = mm.unet_offload_device()
if encoder.quantization == "fp8_e4m3fn":
cast_dtype = torch.float8_e4m3fn
else:
cast_dtype = encoder.dtype
params_to_keep = {"norm", "pos_embedding", "token_embedding"}
if hasattr(encoder, "state_dict"):
model_state_dict = encoder.state_dict
else:
model_state_dict = encoder.model.state_dict()
params_list = list(encoder.model.named_parameters())
pbar = tqdm(params_list, desc="[PromptRelay] Loading T5", leave=True)
for name, param in pbar:
dtype_to_use = dtype if any(k in name for k in params_to_keep) else cast_dtype
value = model_state_dict[name]
set_module_tensor_to_device(encoder.model, name, device=device, dtype=dtype_to_use, value=value)
del model_state_dict
if hasattr(encoder, "state_dict"):
del encoder.state_dict
mm.soft_empty_cache()
with torch.no_grad():
context = encoder([full_prompt], device)
context_null = encoder([negative_prompt], device)
if force_offload:
params_list2 = list(encoder.model.named_parameters())
for name, param in params_list2:
set_module_tensor_to_device(encoder.model, name, device=offload_device)
mm.soft_empty_cache()
text_embeds = {
"prompt_embeds": context,
"negative_prompt_embeds": context_null,
}
return (patcher, text_embeds,)
class RogoAI_PromptRelaySegCalc:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"frame_num": ("INT", {"default": 81, "min": 9, "max": 201, "step": 4}),
"num_segments": ("INT", {"default": 3, "min": 1, "max": 6}),
}
}
RETURN_TYPES = ("STRING", "INT",)
RETURN_NAMES = ("segment_lengths", "latent_frames",)
FUNCTION = "calc"
CATEGORY = "RogoAI/PromptRelay"
def calc(self, frame_num, num_segments):
latent_frames = calc_latent_frames(frame_num)
step = -(-latent_frames // num_segments)
lengths = []
cursor = 0
for _ in range(num_segments):
end = min(cursor + step, latent_frames)
lengths.append(max(end - cursor, 0))
cursor = end
result = ",".join(str(l) for l in lengths)
return (result, latent_frames,)
NODE_CLASS_MAPPINGS = {
"RogoAI_PromptRelayWan": RogoAI_PromptRelayWan,
"RogoAI_PromptRelaySegCalc": RogoAI_PromptRelaySegCalc,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"RogoAI_PromptRelayWan": "RogoAI Prompt Relay (Wan2.2 I2V)",
"RogoAI_PromptRelaySegCalc": "RogoAI Prompt Relay Seg Calc",
}
+130
View File
@@ -0,0 +1,130 @@
import logging
import math
import torch
log = logging.getLogger(__name__)
def build_temporal_cost(q_token_idx, Lq, Lk, device, dtype, tokens_per_frame):
offset = torch.zeros(Lq, Lk, device=device, dtype=dtype)
query_frames = torch.arange(Lq, device=device, dtype=torch.long) // tokens_per_frame
for seg in q_token_idx:
local = seg['local_token_idx'].to(device=device)
d = (query_frames.float()[:, None] - seg['midpoint']).abs()
cost = (torch.relu(d - seg['window']) ** 2) / (2 * seg['sigma'] ** 2)
offset[:, local] = cost.to(offset.dtype)
return offset
def build_temporal_cost_scaled(q_token_idx, Lq, Lk, device, dtype, latent_frames):
offset = torch.zeros(Lq, Lk, device=device, dtype=dtype)
query_frames = torch.arange(Lq, device=device, dtype=torch.float32) * latent_frames / Lq
for seg in q_token_idx:
local = seg['local_token_idx'].to(device=device)
d = (query_frames[:, None] - seg['midpoint']).abs()
cost = (torch.relu(d - seg['window']) ** 2) / (2 * seg['sigma'] ** 2)
offset[:, local] = cost.to(offset.dtype)
return offset
def create_mask_fn(q_token_idx, fallback_tokens_per_frame, latent_frames):
cache = {}
max_token_idx = max(int(seg['local_token_idx'].max().item()) for seg in q_token_idx) + 1
def mask_fn(q, k, transformer_options):
Lq, Lk = q.shape[1], k.shape[1]
if Lq == Lk:
return None
cond_or_uncond = transformer_options.get('cond_or_uncond', [])
if 1 in cond_or_uncond and 0 not in cond_or_uncond:
return None
grid_sizes = transformer_options.get('grid_sizes', None)
video_tpf = int(grid_sizes[1]) * int(grid_sizes[2]) if grid_sizes is not None else fallback_tokens_per_frame
video_lq = latent_frames * video_tpf
if Lk == video_lq or Lk < max_token_idx:
return None
mode = 'video' if Lq == video_lq else 'scaled'
key = (Lq, Lk, mode, q.device)
if key not in cache:
if mode == 'video':
cost = build_temporal_cost(q_token_idx, Lq, Lk, q.device, q.dtype, video_tpf)
else:
cost = build_temporal_cost_scaled(q_token_idx, Lq, Lk, q.device, q.dtype, latent_frames)
log.info('[PromptRelay] Built penalty matrix (%s): Lq=%d, Lk=%d, nonzero=%d/%d',
mode, Lq, Lk, (cost > 0).sum().item(), cost.numel())
cache[key] = -cost
return cache[key].to(q.dtype)
return mask_fn
def build_segments(token_ranges, segment_lengths, epsilon=1e-3):
sigma = 1.0 / math.log(1.0 / epsilon) if 0 < epsilon < 1 else 0.1448
q_token_idx = []
frame_cursor = 0
for (tok_start, tok_end), L in zip(token_ranges, segment_lengths):
if L <= 0:
frame_cursor += L
continue
midpoint = (2 * frame_cursor + L) // 2
window = max(L // 2 - 2, 0)
q_token_idx.append({
'local_token_idx': torch.arange(tok_start, tok_end),
'midpoint': midpoint,
'window': window,
'sigma': sigma,
})
frame_cursor += L
return q_token_idx
def get_raw_tokenizer(clip):
tokenizer_wrapper = clip.tokenizer
for attr_name in dir(tokenizer_wrapper):
if attr_name.startswith('_'):
continue
inner = getattr(tokenizer_wrapper, attr_name, None)
if inner is not None and hasattr(inner, 'tokenizer'):
return inner.tokenizer
raise RuntimeError(
f'Could not find raw tokenizer on CLIP object. '
f'Known attributes: {[a for a in dir(tokenizer_wrapper) if not a.startswith(chr(95))]}'
)
def map_token_indices(raw_tokenizer, global_prompt, local_prompts):
prefixed_locals = [' ' + lp for lp in local_prompts]
full_prompt = global_prompt + ''.join(prefixed_locals)
has_eos = getattr(raw_tokenizer, 'add_eos', False)
eos_adj = 1 if has_eos else 0
prev_len = len(raw_tokenizer(global_prompt)['input_ids']) - eos_adj
token_ranges = []
built = global_prompt
for plp in prefixed_locals:
built += plp
cur_len = len(raw_tokenizer(built)['input_ids']) - eos_adj
if cur_len <= prev_len:
raise ValueError(f'Local prompt produced no tokens: {plp.strip()!r}')
token_ranges.append((prev_len, cur_len))
prev_len = cur_len
return full_prompt, token_ranges
def distribute_segment_lengths(num_segments, latent_frames, specified_lengths=None):
if specified_lengths:
if len(specified_lengths) != num_segments:
raise ValueError(
f'Number of segment_lengths ({len(specified_lengths)}) '
f'must match number of local prompts ({num_segments})'
)
lengths = specified_lengths
else:
step = -(-latent_frames // num_segments)
lengths = [step] * num_segments
effective = []
cursor = 0
for L in lengths:
end = min(cursor + L, latent_frames)
effective.append(max(end - cursor, 0))
cursor = end
return effective
+5
View File
@@ -0,0 +1,5 @@
# ComfyUI-WanVideoWrapper がインストール済みであれば追加インストール不要です
# 以下は参考として記載しています
torch
accelerate
tqdm