572 lines
23 KiB
Python
572 lines
23 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": 30, "min": 1}),
|
||
"use_teacache": (
|
||
"BOOLEAN",
|
||
{"default": True, "tooltip": "Use teacache for faster sampling."},
|
||
),
|
||
"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"}),
|
||
"use_kisekaeichi": (
|
||
"BOOLEAN",
|
||
{
|
||
"default": False,
|
||
"tooltip": "Enable Kisekaeichi mode for style transfer",
|
||
},
|
||
),
|
||
},
|
||
"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_latent": (
|
||
"LATENT",
|
||
{"tooltip": "Reference image latent for kisekaeichi mode"},
|
||
),
|
||
"reference_image_embeds": (
|
||
"CLIP_VISION_OUTPUT",
|
||
{"tooltip": "Reference image CLIP embeds for kisekaeichi mode"},
|
||
),
|
||
"target_index": (
|
||
"INT",
|
||
{
|
||
"default": 1,
|
||
"min": 0,
|
||
"max": 8,
|
||
"step": 1,
|
||
"tooltip": "Target index for kisekaeichi mode (recommended: 1)",
|
||
},
|
||
),
|
||
"history_index": (
|
||
"INT",
|
||
{
|
||
"default": 13,
|
||
"min": 0,
|
||
"max": 16,
|
||
"step": 1,
|
||
"tooltip": "History index for kisekaeichi mode (recommended: 13)",
|
||
},
|
||
),
|
||
"input_mask": (
|
||
"MASK",
|
||
{"tooltip": "Input mask for selective application"},
|
||
),
|
||
"reference_mask": (
|
||
"MASK",
|
||
{"tooltip": "Reference mask for selective features"},
|
||
),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("LATENT",)
|
||
RETURN_NAMES = ("samples",)
|
||
FUNCTION = "process"
|
||
CATEGORY = "FramePackWrapper"
|
||
DESCRIPTION = "Single frame sampler with Kisekaeichi (style transfer) support"
|
||
|
||
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,
|
||
use_kisekaeichi=False,
|
||
reference_latent=None,
|
||
reference_image_embeds=None,
|
||
target_index=1,
|
||
history_index=13,
|
||
input_mask=None,
|
||
reference_mask=None,
|
||
):
|
||
print("=== 1フレーム推論モード ===")
|
||
if use_kisekaeichi:
|
||
print("Kisekaeichi（着せ替え）モード有効")
|
||
print(f"target_index: {target_index}, history_index: {history_index}")
|
||
|
||
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
|
||
if use_kisekaeichi and reference_latent is not None:
|
||
reference_latent = reference_latent["samples"] * vae_scaling_factor
|
||
print(f"参照画像latent: {reference_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
|
||
|
||
if use_kisekaeichi and reference_image_embeds is not None:
|
||
reference_image_encoder_last_hidden_state = reference_image_embeds[
|
||
"last_hidden_state"
|
||
].to(device, base_dtype)
|
||
print("参照画像のCLIP embeddingを設定しました")
|
||
else:
|
||
reference_image_encoder_last_hidden_state = None
|
||
|
||
# テキストエンベッディング処理
|
||
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 reference_latent is not None:
|
||
print("=== Kisekaeichi モード設定（完全版） ===")
|
||
|
||
# 一つ目のコードのone_frame_inference処理を完全再現
|
||
one_frame_inference = set()
|
||
one_frame_inference.add(f"target_index={target_index}")
|
||
one_frame_inference.add(f"history_index={history_index}")
|
||
|
||
# 公式実装に従った処理
|
||
latent_indices = indices[:, -1:] # デフォルトは最後のフレーム
|
||
|
||
# パラメータ解析と処理（一つ目のコードと同じ）
|
||
for one_frame_param in one_frame_inference:
|
||
if one_frame_param.startswith("target_index="):
|
||
target_idx = int(one_frame_param.split("=")[1])
|
||
latent_indices[:, 0] = target_idx
|
||
print(f"latent_indices設定: target_index={target_idx}")
|
||
|
||
elif one_frame_param.startswith("history_index="):
|
||
history_idx = int(one_frame_param.split("=")[1])
|
||
clean_latent_indices_post[:, 0] = history_idx
|
||
print(
|
||
f"clean_latent_indices_post設定: history_index={history_idx}"
|
||
)
|
||
|
||
# history_latentsのダミーを作成（一つ目のコードと同じ構造）
|
||
history_latents = torch.zeros(
|
||
size=(1, 16, 1 + 2 + 16, H, W), dtype=torch.float32, device="cpu"
|
||
)
|
||
|
||
# clean_latents_preの設定（入力画像）
|
||
clean_latents_pre = start_latent.to(history_latents.dtype).to(
|
||
history_latents.device
|
||
)
|
||
if len(clean_latents_pre.shape) < 5:
|
||
clean_latents_pre = clean_latents_pre.unsqueeze(2)
|
||
|
||
# マスクの適用（入力画像）
|
||
if input_mask is not None:
|
||
print("入力画像マスクを適用中...")
|
||
try:
|
||
height_latent, width_latent = clean_latents_pre.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)
|
||
)
|
||
input_mask_resized = input_mask_resized.to(
|
||
clean_latents_pre.device
|
||
)[None, None, None, :, :]
|
||
clean_latents_pre = clean_latents_pre * input_mask_resized
|
||
print("入力画像マスクを適用しました")
|
||
except Exception as e:
|
||
print(f"入力マスク適用エラー: {e}")
|
||
|
||
# clean_latents_postの設定（参照画像）
|
||
clean_latents_post = (
|
||
reference_latent[:, :, 0:1, :, :]
|
||
.to(history_latents.dtype)
|
||
.to(history_latents.device)
|
||
)
|
||
|
||
# マスクの適用（参照画像）
|
||
if reference_mask is not None:
|
||
print("参照画像マスクを適用中...")
|
||
try:
|
||
height_latent, width_latent = clean_latents_post.shape[-2:]
|
||
|
||
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)
|
||
)
|
||
reference_mask_resized = reference_mask_resized.to(
|
||
clean_latents_post.device
|
||
)[None, None, None, :, :]
|
||
clean_latents_post = clean_latents_post * reference_mask_resized
|
||
print("参照画像マスクを適用しました")
|
||
except Exception as e:
|
||
print(f"参照マスク適用エラー: {e}")
|
||
|
||
# clean_latentsを結合
|
||
clean_latents = torch.cat(
|
||
[clean_latents_pre, clean_latents_post], dim=2
|
||
)
|
||
|
||
# clean_latent_indicesを結合
|
||
clean_latent_indices = torch.cat(
|
||
[clean_latent_indices_pre, clean_latent_indices_post], dim=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インデックスを無効化しました")
|
||
|
||
# 画像エンベッディングの処理（両方を活用）
|
||
if (
|
||
reference_image_encoder_last_hidden_state is not None
|
||
and start_image_encoder_last_hidden_state is not None
|
||
):
|
||
# 重み付き平均で統合
|
||
ref_weight = 0.3 # 参照画像の重み
|
||
input_weight = 1.0 - ref_weight
|
||
image_encoder_last_hidden_state = (
|
||
start_image_encoder_last_hidden_state * input_weight
|
||
+ reference_image_encoder_last_hidden_state * ref_weight
|
||
)
|
||
print(
|
||
f"画像エンベッディングを統合 (入力:{input_weight:.2f}, 参照:{ref_weight:.2f})"
|
||
)
|
||
elif reference_image_encoder_last_hidden_state is not None:
|
||
image_encoder_last_hidden_state = (
|
||
reference_image_encoder_last_hidden_state
|
||
)
|
||
else:
|
||
image_encoder_last_hidden_state = (
|
||
start_image_encoder_last_hidden_state
|
||
)
|
||
|
||
print(f"Kisekaeichi設定完了:")
|
||
print(f" - clean_latents.shape: {clean_latents.shape} (入力+参照)")
|
||
print(f" - latent_indices: {latent_indices}")
|
||
print(f" - clean_latent_indices: {clean_latent_indices}")
|
||
print(f" - sample_num_frames: {sample_num_frames}")
|
||
print(f" - 2x/4x無効化: True")
|
||
|
||
else:
|
||
# 通常モード（参照画像なし）
|
||
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()
|
||
|
||
# 出力処理
|
||
mode_info = (
|
||
"Kisekaeichi（latent+エンベッディング）" if use_kisekaeichi else "通常"
|
||
)
|
||
print(
|
||
f"=== 1フレーム生成完了 ({mode_info}モード): {generated_latents.shape} ==="
|
||
)
|
||
|
||
return ({"samples": generated_latents / vae_scaling_factor},)
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"FramePackSingleFrameSampler": FramePackSingleFrameSampler,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"FramePackSingleFrameSampler": "FramePack Single Frame Sampler",
|
||
}
|