204 lines
8.1 KiB
Python
204 lines
8.1 KiB
Python
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",
|
|
}
|