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", }