From cb5c232af61dcbefb21398af21fc4984fffa9942 Mon Sep 17 00:00:00 2001 From: nonioke Date: Wed, 29 Apr 2026 21:18:05 +0900 Subject: [PATCH] Initial release: RogoAI Prompt Relay for Wan2.2 I2V --- .gitignore | 1 + README.md | 220 +++++++++++++++++++++++++++++++++++++++++++++++ __init__.py | 2 + nodes.py | 203 +++++++++++++++++++++++++++++++++++++++++++ prompt_relay.py | 130 ++++++++++++++++++++++++++++ requirements.txt | 5 ++ 6 files changed, 561 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 __init__.py create mode 100644 nodes.py create mode 100644 prompt_relay.py create mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c18dd8d --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..ecfdea7 --- /dev/null +++ b/README.md @@ -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 \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) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..58fa8fe --- /dev/null +++ b/__init__.py @@ -0,0 +1,2 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..29f0dc2 --- /dev/null +++ b/nodes.py @@ -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", +} diff --git a/prompt_relay.py b/prompt_relay.py new file mode 100644 index 0000000..f078272 --- /dev/null +++ b/prompt_relay.py @@ -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 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..47230ca --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +# ComfyUI-WanVideoWrapper がインストール済みであれば追加インストール不要です +# 以下は参考として記載しています +torch +accelerate +tqdm