711 lines
31 KiB
Python
711 lines
31 KiB
Python
import logging
|
||
import math
|
||
|
||
import torch
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
import comfy.latent_formats
|
||
import comfy.model_base
|
||
import comfy.model_management as mm
|
||
from comfy.utils import common_upscale
|
||
|
||
from .diffusers_helper.memory import move_model_to_device_with_memory_preservation
|
||
from .diffusers_helper.pipelines.k_diffusion_hunyuan import sample_hunyuan
|
||
from .diffusers_helper.utils import crop_or_pad_yield_mask
|
||
from .nodes import HyVideoModel, HyVideoModelConfig
|
||
|
||
vae_scaling_factor = 0.476986
|
||
|
||
|
||
class FramePackSingleFrameSampler:
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"model": ("FramePackMODEL",),
|
||
"positive": ("CONDITIONING",),
|
||
"negative": ("CONDITIONING",),
|
||
"start_latent": (
|
||
"LATENT",
|
||
{"tooltip": "init Latents to use for image2image"},
|
||
),
|
||
"steps": ("INT", {"default": 20, "min": 1, "tooltip": "Sampling steps (10=fast preview, 20=balanced, 30=high quality)"}),
|
||
"use_teacache": (
|
||
"BOOLEAN",
|
||
{"default": False, "tooltip": "Use TeaCache for faster sampling (recommended for 15+ steps)."},
|
||
),
|
||
"teacache_rel_l1_thresh": (
|
||
"FLOAT",
|
||
{
|
||
"default": 0.15,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01,
|
||
"tooltip": "The threshold for the relative L1 loss.",
|
||
},
|
||
),
|
||
"cfg": (
|
||
"FLOAT",
|
||
{"default": 1.0, "min": 0.0, "max": 30.0, "step": 0.01},
|
||
),
|
||
"guidance_scale": (
|
||
"FLOAT",
|
||
{"default": 10.0, "min": 0.0, "max": 32.0, "step": 0.01},
|
||
),
|
||
"shift": (
|
||
"FLOAT",
|
||
{"default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.01},
|
||
),
|
||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
||
"latent_window_size": (
|
||
"INT",
|
||
{
|
||
"default": 9,
|
||
"min": 1,
|
||
"max": 33,
|
||
"step": 1,
|
||
"tooltip": "The size of the latent window to use for sampling.",
|
||
},
|
||
),
|
||
"gpu_memory_preservation": (
|
||
"FLOAT",
|
||
{
|
||
"default": 6.0,
|
||
"min": 0.0,
|
||
"max": 128.0,
|
||
"step": 0.1,
|
||
"tooltip": "The amount of GPU memory to preserve.",
|
||
},
|
||
),
|
||
"sampler": (["unipc_bh1", "unipc_bh2"], {"default": "unipc_bh1"}),
|
||
},
|
||
"optional": {
|
||
"image_embeds": ("CLIP_VISION_OUTPUT",),
|
||
"initial_samples": (
|
||
"LATENT",
|
||
{"tooltip": "init Latents to use for image2image variation"},
|
||
),
|
||
"denoise_strength": (
|
||
"FLOAT",
|
||
{"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01},
|
||
),
|
||
"reference_latents": (
|
||
"REFERENCE_LATENT_LIST",
|
||
{"tooltip": "List of reference image latents for kisekaeichi mode (up to 8)"},
|
||
),
|
||
"reference_image_embeds": (
|
||
"REFERENCE_EMBEDS_LIST",
|
||
{"tooltip": "List of reference image CLIP embeds for kisekaeichi mode (up to 8)"},
|
||
),
|
||
"reference_masks": (
|
||
"REFERENCE_MASK_LIST",
|
||
{"tooltip": "List of reference masks for selective features (up to 8, optional)"},
|
||
),
|
||
"target_index": (
|
||
"INT",
|
||
{
|
||
"default": 5,
|
||
"min": 0,
|
||
"max": 16,
|
||
"step": 1,
|
||
"tooltip": "Target index for kisekaeichi mode (recommended: 1)",
|
||
},
|
||
),
|
||
"control_index": (
|
||
"STRING",
|
||
{
|
||
"default": "0;10",
|
||
"tooltip": "Control indices separated by semicolon (preferred) or comma: 0;7;8;9;10 or 0,7,8,9,10 (default: 0;10)",
|
||
},
|
||
),
|
||
"input_mask": (
|
||
"MASK",
|
||
{"tooltip": "Input mask for selective application"},
|
||
),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("LATENT",)
|
||
RETURN_NAMES = ("samples",)
|
||
FUNCTION = "process"
|
||
CATEGORY = "FramePackWrapper"
|
||
DESCRIPTION = "Single frame sampler with auto-enabled Kisekaeichi (style transfer) support for multiple references"
|
||
|
||
def process(
|
||
self,
|
||
model,
|
||
shift,
|
||
positive,
|
||
negative,
|
||
latent_window_size,
|
||
use_teacache,
|
||
teacache_rel_l1_thresh,
|
||
steps,
|
||
cfg,
|
||
guidance_scale,
|
||
seed,
|
||
sampler,
|
||
gpu_memory_preservation,
|
||
start_latent=None,
|
||
image_embeds=None,
|
||
initial_samples=None,
|
||
denoise_strength=1.0,
|
||
reference_latents=None,
|
||
reference_image_embeds=None,
|
||
reference_masks=None,
|
||
target_index=1,
|
||
control_index="0;10",
|
||
input_mask=None,
|
||
):
|
||
print("=== 1フレーム推論モード（musubi-tuner完全互換） ===")
|
||
|
||
# Use control_index parameter directly (no more one_frame_inference)
|
||
|
||
# Auto-enable kisekaeichi when reference inputs are provided
|
||
use_kisekaeichi = reference_latents is not None and len(reference_latents) > 0
|
||
if use_kisekaeichi:
|
||
print("Kisekaeichi（着せ替え）モード有効 (自動検出)")
|
||
print(f"target_index: {target_index}")
|
||
print(f"control_index raw value: '{control_index}'")
|
||
print(f"control_index type: {type(control_index)}")
|
||
|
||
# Parse control indices from parameter (support both semicolon and comma for backward compatibility)
|
||
if ';' in control_index:
|
||
# New format (musubi-tuner compatible): semicolon-separated
|
||
control_indices_list = [int(x.strip()) for x in control_index.split(';') if x.strip()]
|
||
else:
|
||
# Old format: comma-separated for backward compatibility
|
||
control_indices_list = [int(x.strip()) for x in control_index.split(',') if x.strip()]
|
||
|
||
if len(control_indices_list) == 0:
|
||
control_indices_list = [0, 10] # default
|
||
print(f"parsed control_index: {control_indices_list}")
|
||
|
||
transformer = model["transformer"]
|
||
base_dtype = model["dtype"]
|
||
device = mm.get_torch_device()
|
||
offload_device = mm.unet_offload_device()
|
||
|
||
mm.unload_all_models()
|
||
mm.cleanup_models()
|
||
mm.soft_empty_cache()
|
||
|
||
# Latent処理
|
||
if start_latent is not None:
|
||
start_latent = start_latent["samples"] * vae_scaling_factor
|
||
if initial_samples is not None:
|
||
initial_samples = initial_samples["samples"] * vae_scaling_factor
|
||
# Process reference latents if provided
|
||
processed_reference_latents = []
|
||
if use_kisekaeichi and reference_latents is not None:
|
||
for i, ref_latent in enumerate(reference_latents):
|
||
processed_latent = ref_latent["samples"] * vae_scaling_factor
|
||
processed_reference_latents.append(processed_latent)
|
||
print(f"参照画像latent {i+1}: {processed_latent.shape}")
|
||
|
||
print("start_latent", start_latent.shape)
|
||
B, C, T, H, W = start_latent.shape
|
||
|
||
# 画像エンベッディング処理
|
||
if image_embeds is not None:
|
||
start_image_encoder_last_hidden_state = image_embeds[
|
||
"last_hidden_state"
|
||
].to(device, base_dtype)
|
||
else:
|
||
start_image_encoder_last_hidden_state = None
|
||
|
||
# Process reference image embeds if provided
|
||
processed_reference_embeds = []
|
||
if use_kisekaeichi and reference_image_embeds is not None:
|
||
for i, ref_embeds in enumerate(reference_image_embeds):
|
||
embed_state = ref_embeds["last_hidden_state"].to(device, base_dtype)
|
||
processed_reference_embeds.append(embed_state)
|
||
print(f"参照画像 {i+1} のCLIP embeddingを設定しました")
|
||
|
||
# Validate that latents and embeds have same count
|
||
if use_kisekaeichi:
|
||
if len(processed_reference_latents) != len(processed_reference_embeds):
|
||
raise ValueError(f"Reference latents count ({len(processed_reference_latents)}) must match embeds count ({len(processed_reference_embeds)})")
|
||
if len(processed_reference_latents) == 0:
|
||
raise ValueError("No reference latents provided for kisekaeichi mode")
|
||
|
||
# テキストエンベッディング処理
|
||
llama_vec = positive[0][0].to(device, base_dtype)
|
||
clip_l_pooler = positive[0][1]["pooled_output"].to(device, base_dtype)
|
||
|
||
if not math.isclose(cfg, 1.0):
|
||
llama_vec_n = negative[0][0].to(device, base_dtype)
|
||
clip_l_pooler_n = negative[0][1]["pooled_output"].to(device, base_dtype)
|
||
else:
|
||
llama_vec_n = torch.zeros_like(llama_vec, device=device)
|
||
clip_l_pooler_n = torch.zeros_like(clip_l_pooler, device=device)
|
||
|
||
llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512)
|
||
llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask(
|
||
llama_vec_n, length=512
|
||
)
|
||
|
||
# シード設定
|
||
rnd = torch.Generator("cpu").manual_seed(seed)
|
||
|
||
# === 一つ目のコードに完全準拠した設定 ===
|
||
|
||
# 1フレームモード固定設定
|
||
sample_num_frames = 1
|
||
total_latent_sections = 1
|
||
latent_padding = 0
|
||
latent_padding_size = latent_padding * latent_window_size # 0
|
||
|
||
# 一つ目のコードと同じインデックス構造
|
||
indices = torch.arange(
|
||
0, sum([1, latent_padding_size, latent_window_size, 1, 2, 16])
|
||
).unsqueeze(0)
|
||
split_sizes = [1, latent_padding_size, latent_window_size, 1, 2, 16]
|
||
|
||
# latent_padding_sizeが0の場合の分割（一つ目のコードと完全同一）
|
||
if latent_padding_size == 0:
|
||
clean_latent_indices_pre = indices[:, 0:1]
|
||
latent_indices = indices[:, 1 : 1 + latent_window_size]
|
||
clean_latent_indices_post = indices[
|
||
:, 1 + latent_window_size : 2 + latent_window_size
|
||
]
|
||
clean_latent_2x_indices = indices[
|
||
:, 2 + latent_window_size : 4 + latent_window_size
|
||
]
|
||
clean_latent_4x_indices = indices[
|
||
:, 4 + latent_window_size : 20 + latent_window_size
|
||
]
|
||
blank_indices = torch.empty((1, 0), dtype=torch.long)
|
||
else:
|
||
(
|
||
clean_latent_indices_pre,
|
||
blank_indices,
|
||
latent_indices,
|
||
clean_latent_indices_post,
|
||
clean_latent_2x_indices,
|
||
clean_latent_4x_indices,
|
||
) = indices.split(split_sizes, dim=1)
|
||
|
||
# 一つ目のコードの重要な処理：1フレームモード時のone_frame_inference処理
|
||
if sample_num_frames == 1:
|
||
if use_kisekaeichi and len(processed_reference_latents) > 0:
|
||
print("=== Kisekaeichi モード設定（完全版） ===")
|
||
|
||
# 一つ目のコードのone_frame_inference処理を完全再現
|
||
one_frame_inference = set()
|
||
one_frame_inference.add(f"target_index={target_index}")
|
||
# control_indexは後でcontrol_index_listから設定されるため、ここでは追加しない
|
||
|
||
# musubi-tuner仕様：latent_indicesを初期化
|
||
latent_indices = torch.zeros((1, 1), dtype=torch.int64) # 1x1 latent index for target image
|
||
latent_indices[:, 0] = latent_window_size # musubi-tunerのデフォルト値
|
||
|
||
# Apply target_index parameter (musubi-tuner仕様)
|
||
latent_indices[:, 0] = target_index
|
||
print(f"target_index設定 (musubi-tuner仕様): {target_index}")
|
||
|
||
# この部分は削除（後でまとめて処理するため）
|
||
|
||
# control_latentsのダミーを作成（一つ目のコードと同じ構造）
|
||
control_latents = torch.zeros(
|
||
size=(1, 16, 1 + 2 + 16, H, W), dtype=torch.float32, device="cpu"
|
||
)
|
||
|
||
# musubi-tuner仕様：control_latentsを直接結合（マスク適用は後で実行）
|
||
control_latents = []
|
||
|
||
# 入力画像をcontrol_latentsに追加
|
||
clean_latents_pre = start_latent.to(torch.float32).cpu()
|
||
if len(clean_latents_pre.shape) < 5:
|
||
clean_latents_pre = clean_latents_pre.unsqueeze(2)
|
||
control_latents.append(clean_latents_pre)
|
||
|
||
# 参照画像をcontrol_latentsに追加
|
||
for i, ref_latent in enumerate(processed_reference_latents):
|
||
clean_latent_post = ref_latent[:, :, 0:1, :, :].to(torch.float32).cpu()
|
||
control_latents.append(clean_latent_post)
|
||
|
||
# musubi-tuner仕様：ゼロlatentをclean_latents_postとして追加
|
||
control_latents.append(torch.zeros((1, 16, 1, H, W), dtype=torch.float32))
|
||
print(f"Add zero latents as clean latents post for one frame inference.")
|
||
|
||
# musubi-tuner仕様：control_latentsを直接結合
|
||
clean_latents = torch.cat(control_latents, dim=2) # (1, 16, num_control_images, H//8, W//8)
|
||
|
||
# musubi-tuner仕様：clean_latent_indicesを設定
|
||
clean_latent_indices = torch.zeros((1, len(control_latents)), dtype=torch.int64)
|
||
if "no_post" not in []: # one_frame_inferenceは空なので常にtrue
|
||
clean_latent_indices[:, -1] = 1 + latent_window_size # default index for clean latents post
|
||
|
||
# musubi-tuner仕様：control_indices_listをclean_latent_indicesに動的適用
|
||
if control_indices_list:
|
||
print(f"control_indices動的適用開始: {control_indices_list}")
|
||
i = 0
|
||
while i < len(control_indices_list) and i < clean_latent_indices.shape[1]:
|
||
control_index = int(control_indices_list[i])
|
||
clean_latent_indices[:, i] = control_index
|
||
print(f"clean_latent_indices[:, {i}] = {control_index}")
|
||
i += 1
|
||
print(f"最終的なclean_latent_indices: {clean_latent_indices}")
|
||
else:
|
||
print("control_indices_listが空のため、デフォルト設定を使用")
|
||
|
||
# musubi-tuner仕様：マスク適用（clean_latents生成後）
|
||
for i in range(len(control_latents)):
|
||
mask_image = None
|
||
|
||
# 入力画像のマスク適用
|
||
if i == 0 and input_mask is not None:
|
||
print("入力画像マスクを適用中（musubi-tuner仕様）...")
|
||
try:
|
||
height_latent, width_latent = clean_latents.shape[-2:]
|
||
if isinstance(input_mask, torch.Tensor):
|
||
input_mask_tensor = input_mask
|
||
else:
|
||
input_mask_tensor = torch.from_numpy(input_mask)
|
||
|
||
input_mask_resized = (
|
||
common_upscale(
|
||
input_mask_tensor.unsqueeze(0).unsqueeze(0),
|
||
width_latent,
|
||
height_latent,
|
||
"bilinear",
|
||
"center",
|
||
)
|
||
.squeeze(0)
|
||
.squeeze(0)
|
||
)
|
||
mask_image = input_mask_resized.to(clean_latents.device)[None, None, :, :]
|
||
print("入力画像マスクを適用しました")
|
||
except Exception as e:
|
||
print(f"入力マスク適用エラー: {e}")
|
||
|
||
# 参照画像のマスク適用
|
||
elif i > 0 and i < len(control_latents) - 1: # -1はゼロlatentなので除外
|
||
ref_index = i - 1 # 最初が入力画像なので-1
|
||
if reference_masks is not None and ref_index < len(reference_masks):
|
||
print(f"参照画像 {ref_index+1} マスクを適用中（musubi-tuner仕様）...")
|
||
try:
|
||
height_latent, width_latent = clean_latents.shape[-2:]
|
||
reference_mask = reference_masks[ref_index]
|
||
|
||
if isinstance(reference_mask, torch.Tensor):
|
||
reference_mask_tensor = reference_mask
|
||
else:
|
||
reference_mask_tensor = torch.from_numpy(reference_mask)
|
||
|
||
reference_mask_resized = (
|
||
common_upscale(
|
||
reference_mask_tensor.unsqueeze(0).unsqueeze(0),
|
||
width_latent,
|
||
height_latent,
|
||
"bilinear",
|
||
"center",
|
||
)
|
||
.squeeze(0)
|
||
.squeeze(0)
|
||
)
|
||
mask_image = reference_mask_resized.to(clean_latents.device)[None, None, :, :]
|
||
print(f"参照画像 {ref_index+1} マスクを適用しました")
|
||
except Exception as e:
|
||
print(f"参照マスク {ref_index+1} 適用エラー: {e}")
|
||
|
||
# マスクを適用
|
||
if mask_image is not None:
|
||
clean_latents[:, :, i : i + 1, :, :] = clean_latents[:, :, i : i + 1, :, :] * mask_image
|
||
|
||
# Always disable 2x/4x processing for kisekaeichi mode
|
||
clean_latents_2x_param = None
|
||
clean_latent_2x_indices = None
|
||
clean_latents_2x = None
|
||
|
||
clean_latents_4x_param = None
|
||
clean_latent_4x_indices = None
|
||
clean_latents_4x = None
|
||
|
||
print("Kisekaeichi: 2x/4xインデックスを無効化しました")
|
||
|
||
# 画像エンベッディングの処理（musubi-tuner仕様：セクション毎の個別処理）
|
||
# musubi-tunerではセクション毎に個別のimage_encoder_last_hidden_stateを使用
|
||
# 複数リファレンス時は最初の参照画像のエンベッディングを使用（musubi-tuner互換）
|
||
if len(processed_reference_embeds) > 0:
|
||
# musubi-tuner互換：最初の参照画像エンベッディングを使用
|
||
image_encoder_last_hidden_state = processed_reference_embeds[0]
|
||
print(f"musubi-tuner互換：最初の参照画像エンベッディングを使用 (参照画像数: {len(processed_reference_embeds)})")
|
||
else:
|
||
image_encoder_last_hidden_state = start_image_encoder_last_hidden_state
|
||
|
||
print(f"Kisekaeichi設定完了（musubi-tuner完全互換）:")
|
||
print(f" - clean_latents.shape: {clean_latents.shape} (入力+参照{len(processed_reference_latents)}個+ゼロlatent)")
|
||
print(f" - latent_indices: {latent_indices} (初期値: {latent_window_size} -> target: {target_index})")
|
||
print(f" - clean_latent_indices: {clean_latent_indices} (control_indices適用済み)")
|
||
print(f" - sample_num_frames: {sample_num_frames}")
|
||
print(f" - control_indices適用: {control_indices_list}")
|
||
print(f" - control_latents数: {len(control_latents)}")
|
||
print(f" - マスク適用: clean_latents生成後")
|
||
print(f" - 2x/4x無効化: True")
|
||
|
||
else:
|
||
# 通常モード（参照画像なし）
|
||
# control_indicesは通常モードでは使用しない
|
||
all_indices = torch.arange(0, latent_window_size).unsqueeze(0)
|
||
latent_indices = all_indices[:, -1:]
|
||
|
||
clean_latents_pre = start_latent.to(torch.float32).cpu()
|
||
if len(clean_latents_pre.shape) < 5:
|
||
clean_latents_pre = clean_latents_pre.unsqueeze(2)
|
||
|
||
clean_latents_post = torch.zeros_like(clean_latents_pre)
|
||
clean_latents = torch.cat(
|
||
[clean_latents_pre, clean_latents_post], dim=2
|
||
)
|
||
clean_latent_indices = torch.cat(
|
||
[clean_latent_indices_pre, clean_latent_indices_post], dim=1
|
||
)
|
||
|
||
# 通常モードでのインデックス調整
|
||
clean_latent_indices = torch.tensor(
|
||
[[0]],
|
||
dtype=clean_latent_indices.dtype,
|
||
device=clean_latent_indices.device,
|
||
)
|
||
clean_latents = clean_latents[:, :, :1, :, :]
|
||
|
||
clean_latents_2x_param = None
|
||
clean_latents_4x_param = None
|
||
clean_latent_2x_indices = None
|
||
clean_latent_4x_indices = None
|
||
|
||
# 2x, 4x latentsの設定も無効化
|
||
clean_latents_2x = None
|
||
clean_latents_4x = None
|
||
|
||
print("Kisekaeichi: 2x/4xインデックスを無効化しました")
|
||
|
||
image_encoder_last_hidden_state = start_image_encoder_last_hidden_state
|
||
|
||
print("通常モード設定:")
|
||
print(f" - clean_latents.shape: {clean_latents.shape}")
|
||
print(f" - latent_indices: {latent_indices}")
|
||
print(f" - clean_latent_indices: {clean_latent_indices}")
|
||
|
||
# 初期サンプルの処理
|
||
input_init_latents = None
|
||
if initial_samples is not None:
|
||
input_init_latents = initial_samples[:, :, 0:1, :, :].to(device)
|
||
print("初期サンプルを設定しました")
|
||
|
||
# ComfyUI用のセットアップ
|
||
comfy_model = HyVideoModel(
|
||
HyVideoModelConfig(base_dtype),
|
||
model_type=comfy.model_base.ModelType.FLOW,
|
||
device=device,
|
||
)
|
||
patcher = comfy.model_patcher.ModelPatcher(
|
||
comfy_model, device, torch.device("cpu")
|
||
)
|
||
from latent_preview import prepare_callback
|
||
|
||
callback = prepare_callback(patcher, steps)
|
||
|
||
# モデルをGPUに移動
|
||
move_model_to_device_with_memory_preservation(
|
||
transformer,
|
||
target_device=device,
|
||
preserved_memory_gb=gpu_memory_preservation,
|
||
)
|
||
|
||
# TeaCacheの設定
|
||
if use_teacache:
|
||
transformer.initialize_teacache(
|
||
enable_teacache=True,
|
||
num_steps=steps,
|
||
rel_l1_thresh=teacache_rel_l1_thresh,
|
||
)
|
||
else:
|
||
transformer.initialize_teacache(enable_teacache=False)
|
||
|
||
print("=== サンプリング開始 ===")
|
||
print(f"sample_num_frames: {sample_num_frames}")
|
||
print(f"clean_latents使用フレーム数: {clean_latents.shape[2]}")
|
||
print(f"clean_latent_2x_indices: {clean_latent_2x_indices}")
|
||
print(f"clean_latent_4x_indices: {clean_latent_4x_indices}")
|
||
|
||
with torch.autocast(
|
||
device_type=mm.get_autocast_device(device), dtype=base_dtype, enabled=True
|
||
):
|
||
generated_latents = sample_hunyuan(
|
||
transformer=transformer,
|
||
sampler=sampler,
|
||
initial_latent=input_init_latents,
|
||
strength=denoise_strength,
|
||
width=W * 8,
|
||
height=H * 8,
|
||
frames=sample_num_frames, # 1フレーム固定
|
||
real_guidance_scale=cfg,
|
||
distilled_guidance_scale=guidance_scale,
|
||
guidance_rescale=0,
|
||
shift=shift if shift != 0 else None,
|
||
num_inference_steps=steps,
|
||
generator=rnd,
|
||
prompt_embeds=llama_vec,
|
||
prompt_embeds_mask=llama_attention_mask,
|
||
prompt_poolers=clip_l_pooler,
|
||
negative_prompt_embeds=llama_vec_n,
|
||
negative_prompt_embeds_mask=llama_attention_mask_n,
|
||
negative_prompt_poolers=clip_l_pooler_n,
|
||
device=device,
|
||
dtype=base_dtype,
|
||
image_embeddings=image_encoder_last_hidden_state,
|
||
latent_indices=latent_indices,
|
||
clean_latents=clean_latents,
|
||
clean_latent_indices=clean_latent_indices,
|
||
clean_latents_2x=clean_latents_2x,
|
||
clean_latent_2x_indices=clean_latent_2x_indices,
|
||
clean_latents_4x=clean_latents_4x,
|
||
clean_latent_4x_indices=clean_latent_4x_indices,
|
||
callback=callback,
|
||
)
|
||
|
||
# クリーンアップ
|
||
transformer.to(offload_device)
|
||
mm.soft_empty_cache()
|
||
|
||
# 出力処理
|
||
if use_kisekaeichi:
|
||
mode_info = f"Kisekaeichi（musubi-tuner完全互換、参照{len(processed_reference_latents)}個）"
|
||
else:
|
||
mode_info = "通常"
|
||
print(
|
||
f"=== 1フレーム生成完了 ({mode_info}モード): {generated_latents.shape} ==="
|
||
)
|
||
|
||
return ({"samples": generated_latents / vae_scaling_factor},)
|
||
|
||
|
||
class ReferenceLatentList:
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {},
|
||
"optional": {
|
||
"reference_1": ("LATENT", {"tooltip": "Reference latent 1"}),
|
||
"reference_2": ("LATENT", {"tooltip": "Reference latent 2"}),
|
||
"reference_3": ("LATENT", {"tooltip": "Reference latent 3"}),
|
||
"reference_4": ("LATENT", {"tooltip": "Reference latent 4"}),
|
||
"reference_5": ("LATENT", {"tooltip": "Reference latent 5"}),
|
||
"reference_6": ("LATENT", {"tooltip": "Reference latent 6"}),
|
||
"reference_7": ("LATENT", {"tooltip": "Reference latent 7"}),
|
||
"reference_8": ("LATENT", {"tooltip": "Reference latent 8"}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("REFERENCE_LATENT_LIST",)
|
||
RETURN_NAMES = ("reference_latents",)
|
||
FUNCTION = "create_list"
|
||
CATEGORY = "FramePackWrapper/References"
|
||
DESCRIPTION = "Combine up to 8 reference latents into a list for multi-reference Kisekaeichi mode"
|
||
|
||
def create_list(self, **kwargs):
|
||
reference_list = []
|
||
for i in range(1, 9):
|
||
ref_key = f"reference_{i}"
|
||
if ref_key in kwargs and kwargs[ref_key] is not None:
|
||
reference_list.append(kwargs[ref_key])
|
||
|
||
if len(reference_list) == 0:
|
||
return (None,)
|
||
|
||
print(f"Created reference latent list with {len(reference_list)} items")
|
||
return (reference_list,)
|
||
|
||
|
||
class ReferenceEmbedsList:
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {},
|
||
"optional": {
|
||
"embeds_1": ("CLIP_VISION_OUTPUT", {"tooltip": "Reference embeds 1"}),
|
||
"embeds_2": ("CLIP_VISION_OUTPUT", {"tooltip": "Reference embeds 2"}),
|
||
"embeds_3": ("CLIP_VISION_OUTPUT", {"tooltip": "Reference embeds 3"}),
|
||
"embeds_4": ("CLIP_VISION_OUTPUT", {"tooltip": "Reference embeds 4"}),
|
||
"embeds_5": ("CLIP_VISION_OUTPUT", {"tooltip": "Reference embeds 5"}),
|
||
"embeds_6": ("CLIP_VISION_OUTPUT", {"tooltip": "Reference embeds 6"}),
|
||
"embeds_7": ("CLIP_VISION_OUTPUT", {"tooltip": "Reference embeds 7"}),
|
||
"embeds_8": ("CLIP_VISION_OUTPUT", {"tooltip": "Reference embeds 8"}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("REFERENCE_EMBEDS_LIST",)
|
||
RETURN_NAMES = ("reference_embeds",)
|
||
FUNCTION = "create_list"
|
||
CATEGORY = "FramePackWrapper/References"
|
||
DESCRIPTION = "Combine up to 8 reference CLIP vision embeddings into a list for multi-reference Kisekaeichi mode"
|
||
|
||
def create_list(self, **kwargs):
|
||
embeds_list = []
|
||
for i in range(1, 9):
|
||
embeds_key = f"embeds_{i}"
|
||
if embeds_key in kwargs and kwargs[embeds_key] is not None:
|
||
embeds_list.append(kwargs[embeds_key])
|
||
|
||
if len(embeds_list) == 0:
|
||
return (None,)
|
||
|
||
print(f"Created reference embeds list with {len(embeds_list)} items")
|
||
return (embeds_list,)
|
||
|
||
|
||
class ReferenceMaskList:
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {},
|
||
"optional": {
|
||
"mask_1": ("MASK", {"tooltip": "Reference mask 1"}),
|
||
"mask_2": ("MASK", {"tooltip": "Reference mask 2"}),
|
||
"mask_3": ("MASK", {"tooltip": "Reference mask 3"}),
|
||
"mask_4": ("MASK", {"tooltip": "Reference mask 4"}),
|
||
"mask_5": ("MASK", {"tooltip": "Reference mask 5"}),
|
||
"mask_6": ("MASK", {"tooltip": "Reference mask 6"}),
|
||
"mask_7": ("MASK", {"tooltip": "Reference mask 7"}),
|
||
"mask_8": ("MASK", {"tooltip": "Reference mask 8"}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("REFERENCE_MASK_LIST",)
|
||
RETURN_NAMES = ("reference_masks",)
|
||
FUNCTION = "create_list"
|
||
CATEGORY = "FramePackWrapper/References"
|
||
DESCRIPTION = "Combine up to 8 reference masks into a list for multi-reference Kisekaeichi mode (optional)"
|
||
|
||
def create_list(self, **kwargs):
|
||
mask_list = []
|
||
for i in range(1, 9):
|
||
mask_key = f"mask_{i}"
|
||
if mask_key in kwargs and kwargs[mask_key] is not None:
|
||
mask_list.append(kwargs[mask_key])
|
||
|
||
if len(mask_list) == 0:
|
||
return (None,)
|
||
|
||
print(f"Created reference mask list with {len(mask_list)} items")
|
||
return (mask_list,)
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"FramePackSingleFrameSampler": FramePackSingleFrameSampler,
|
||
"ReferenceLatentList": ReferenceLatentList,
|
||
"ReferenceEmbedsList": ReferenceEmbedsList,
|
||
"ReferenceMaskList": ReferenceMaskList,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"FramePackSingleFrameSampler": "FramePack Single Frame Sampler",
|
||
"ReferenceLatentList": "Reference Latent List",
|
||
"ReferenceEmbedsList": "Reference Embeds List",
|
||
"ReferenceMaskList": "Reference Mask List",
|
||
}
|