Add files via upload

https://note.com/tori29umai/n/n40965720363e
This commit is contained in:
ethanfel
2025-06-25 17:27:46 +02:00
committed by GitHub
parent dd9dfd08a4
commit dfe4747036
2 changed files with 872 additions and 9 deletions
+11 -9
View File
@@ -1,9 +1,11 @@
from .nodes import NODE_CLASS_MAPPINGS as NODES_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as NODES_DISPLAY_MAPPINGS
from .nodes_F1 import NODE_CLASS_MAPPINGS as F1_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as F1_DISPLAY_MAPPINGS
from .nodes_one import NODE_CLASS_MAPPINGS as ONE_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as ONE_DISPLAY_MAPPINGS
# Combine the mappings
NODE_CLASS_MAPPINGS = {**NODES_MAPPINGS, **F1_MAPPINGS, **ONE_MAPPINGS}
NODE_DISPLAY_NAME_MAPPINGS = {**NODES_DISPLAY_MAPPINGS, **F1_DISPLAY_MAPPINGS, **ONE_DISPLAY_MAPPINGS}
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
from .nodes import NODE_CLASS_MAPPINGS as NODES_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as NODES_DISPLAY_MAPPINGS
from .nodes_F1 import NODE_CLASS_MAPPINGS as F1_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as F1_DISPLAY_MAPPINGS
from .nodes_one import NODE_CLASS_MAPPINGS as ONE_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as ONE_DISPLAY_MAPPINGS
from .nodes_resize import NODE_CLASS_MAPPINGS as RESIZE_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as RESIZE_DISPLAY_MAPPINGS
# Combine the mappings
NODE_CLASS_MAPPINGS = {**NODES_MAPPINGS, **F1_MAPPINGS, **ONE_MAPPINGS, **RESIZE_MAPPINGS}
NODE_DISPLAY_NAME_MAPPINGS = {**NODES_DISPLAY_MAPPINGS, **F1_DISPLAY_MAPPINGS, **ONE_DISPLAY_MAPPINGS, **RESIZE_DISPLAY_MAPPINGS}
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
+861
View File
@@ -0,0 +1,861 @@
import numpy as np
from PIL import Image
import logging
import math
from tqdm import tqdm
from latent_preview import prepare_callback
import torch
logger = logging.getLogger(__name__)
import comfy.model_base
import comfy.model_management as mm
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
# Constants
VAE_SCALING_FACTOR = 0.476986
DEFAULT_BLEND_ALPHA = 0.5
OVERLAP_RATIO = 0.5 # 50% オーバーラップ
class TileProcessor:
"""50%オーバーラップを持つ2次元画像タイル処理を担当するクラス"""
def __init__(self, tile_size: int, tile_resized: int):
self.tile_size = tile_size
self.tile_resized = tile_resized
self.upscale_factor = float(tile_resized) / float(tile_size)
self.step_size = int(tile_size * (1 - OVERLAP_RATIO)) # 50%オーバーラップのステップサイズ
def split_image_to_tiles(self, image):
"""画像を50%オーバーラップでタイルに分割し、白パディングを追加して指定サイズにリサイズ"""
pil_image = self._convert_to_pil(image)
original_width, original_height = pil_image.size
# オーバーラップを考慮したタイル数を計算
tiles_x = self._calculate_tile_count(original_width)
tiles_y = self._calculate_tile_count(original_height)
self._log_split_info(original_width, original_height, tiles_x, tiles_y)
metadata = self._create_metadata(original_width, original_height, tiles_x, tiles_y)
tiles = self._extract_overlapping_tiles(pil_image, tiles_x, tiles_y, metadata)
return tiles, metadata
def _calculate_tile_count(self, dimension: int) -> int:
"""オーバーラップを考慮したタイル数を計算"""
if dimension <= self.tile_size:
return 1
# 最後のタイルが完全に収まるように計算
return math.ceil((dimension - self.tile_size) / self.step_size) + 1
def _convert_to_pil(self, image):
"""ComfyUIのIMAGE形式からPIL形式に変換"""
if isinstance(image, torch.Tensor):
image_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
return Image.fromarray(image_np)
return image
def _log_split_info(self, width: int, height: int, tiles_x: int, tiles_y: int):
"""分割情報をログ出力"""
print(f"=== 50%オーバーラップ画像分割情報 ===")
print(f"元画像サイズ: {width} × {height}")
print(f"分割タイルサイズ: {self.tile_size} × {self.tile_size}")
print(f"リサイズ後タイルサイズ: {self.tile_resized} × {self.tile_resized}")
print(f"アップスケール倍率: {self.upscale_factor:.2f}x")
print(f"ステップサイズ: {self.step_size} (50%オーバーラップ)")
print(f"タイル数: {tiles_x} × {tiles_y} = {tiles_x * tiles_y}個")
def _create_metadata(self, width: int, height: int, tiles_x: int, tiles_y: int) -> dict:
"""メタデータを作成"""
return {
"original_size": [width, height],
"tile_size": self.tile_size,
"tile_resized": self.tile_resized,
"upscale_factor": self.upscale_factor,
"step_size": self.step_size,
"overlap_ratio": OVERLAP_RATIO,
"tiles_x": tiles_x,
"tiles_y": tiles_y,
"tiles": []
}
def _extract_overlapping_tiles(self, pil_image: Image.Image, tiles_x: int, tiles_y: int, metadata: dict) -> list:
"""50%オーバーラップでタイルを切り出す"""
tiles = []
tile_index = 0
for y in range(tiles_y):
for x in range(tiles_x):
tile_tensor, tile_info = self._process_overlapping_tile(
pil_image, x, y, tile_index
)
tiles.append(tile_tensor)
metadata["tiles"].append(tile_info)
tile_index += 1
return tiles
def _process_overlapping_tile(self, pil_image: Image.Image, x: int, y: int, tile_index: int):
"""50%オーバーラップを考慮した単一タイルを処理"""
# オーバーラップを考慮した切り出し位置を計算
left = x * self.step_size
top = y * self.step_size
right = min(left + self.tile_size, pil_image.width)
bottom = min(top + self.tile_size, pil_image.height)
# 画像端で調整(右端・下端のタイルが小さくなる場合)
if right - left < self.tile_size:
left = max(0, right - self.tile_size)
if bottom - top < self.tile_size:
top = max(0, bottom - self.tile_size)
# タイルを切り出し
tile = pil_image.crop((left, top, right, bottom))
tile_width, tile_height = tile.size
# パディング処理
padding_needed = (tile_width < self.tile_size or tile_height < self.tile_size)
if padding_needed:
tile = self._add_white_padding(tile)
# アップスケール
upscaled_tile = tile.resize((self.tile_resized, self.tile_resized), Image.LANCZOS)
# テンソル変換
tile_tensor = torch.from_numpy(np.array(upscaled_tile)).float() / 255.0
tile_tensor = tile_tensor.unsqueeze(0) # バッチ次元を追加
# オーバーラップ領域の情報を計算
overlap_info = self._calculate_overlap_info(x, y, left, top, right, bottom, pil_image.size)
# メタデータ作成
tile_info = {
"index": tile_index,
"grid_position": [x, y],
"crop_box": [left, top, right, bottom],
"original_tile_size": [tile_width, tile_height],
"padding_needed": padding_needed,
"padding_right": self.tile_size - tile_width if padding_needed else 0,
"padding_bottom": self.tile_size - tile_height if padding_needed else 0,
"overlap_info": overlap_info
}
print(f"タイル {tile_index}: グリッド({x},{y}) 切り出し({left},{top},{right},{bottom}) "
f"オーバーラップ:{overlap_info}")
return tile_tensor, tile_info
def _calculate_overlap_info(self, grid_x: int, grid_y: int, left: int, top: int,
right: int, bottom: int, image_size: tuple) -> dict:
"""オーバーラップ領域の情報を計算"""
image_width, image_height = image_size
overlap_info = {
"has_left_overlap": grid_x > 0,
"has_right_overlap": right < image_width and grid_x < self._calculate_tile_count(image_width) - 1,
"has_top_overlap": grid_y > 0,
"has_bottom_overlap": bottom < image_height and grid_y < self._calculate_tile_count(image_height) - 1,
}
# オーバーラップピクセル数を計算
overlap_pixels = int(self.tile_size * OVERLAP_RATIO)
overlap_info.update({
"left_overlap_pixels": overlap_pixels if overlap_info["has_left_overlap"] else 0,
"right_overlap_pixels": overlap_pixels if overlap_info["has_right_overlap"] else 0,
"top_overlap_pixels": overlap_pixels if overlap_info["has_top_overlap"] else 0,
"bottom_overlap_pixels": overlap_pixels if overlap_info["has_bottom_overlap"] else 0,
})
return overlap_info
def _add_white_padding(self, tile: Image.Image) -> Image.Image:
"""白パディングを追加"""
padded_tile = Image.new('RGB', (self.tile_size, self.tile_size), color='white')
padded_tile.paste(tile, (0, 0))
return padded_tile
class ImageReconstructor:
"""50%オーバーラップしたタイルから画像を再構成するクラス"""
def __init__(self, metadata: dict):
self.metadata = metadata
self.upscale_factor = metadata["upscale_factor"]
self.tile_size = metadata["tile_size"]
self.step_size = metadata["step_size"]
self.overlap_ratio = metadata["overlap_ratio"]
def reconstruct_image(self, processed_tiles: list):
"""50%オーバーラップしたタイルから画像を再構成"""
final_width, final_height = self._calculate_final_size()
self._log_reconstruction_info(final_width, final_height, len(processed_tiles))
# 最終画像とブレンドウェイトマップを初期化
final_image = np.zeros((final_height, final_width, 3), dtype=np.float32)
weight_map = np.zeros((final_height, final_width), dtype=np.float32)
# 処理されたタイル画像を準備
processed_tile_images = self._preprocess_tiles(processed_tiles)
# オーバーラップブレンドでタイルを配置
self._place_tiles_with_overlap_blend(final_image, weight_map, processed_tile_images)
# 正規化して最終画像を作成
final_image = self._normalize_blended_image(final_image, weight_map)
return self._convert_to_tensor(Image.fromarray(final_image.astype(np.uint8)))
def _calculate_final_size(self) -> tuple:
"""最終画像サイズを計算"""
original_width, original_height = self.metadata["original_size"]
final_width = int(original_width * self.upscale_factor)
final_height = int(original_height * self.upscale_factor)
return final_width, final_height
def _log_reconstruction_info(self, width: int, height: int, tile_count: int):
"""再構成情報をログ出力"""
print(f"=== 50%オーバーラップ画像再構成情報 ===")
print(f"最終画像サイズ: {width} × {height} ({self.upscale_factor:.2f}x)")
print(f"処理タイル数: {tile_count}")
print(f"オーバーラップ比率: {self.overlap_ratio*100:.0f}%")
def _preprocess_tiles(self, processed_tiles: list) -> list:
"""タイル画像を前処理"""
processed_tile_images = []
for i, (tile_tensor, tile_info) in enumerate(zip(processed_tiles, self.metadata["tiles"])):
tile_image = self._tensor_to_pil(tile_tensor)
# パディング部分を除去
if tile_info["padding_needed"]:
tile_image = self._remove_padding(tile_image, tile_info)
processed_tile_images.append((tile_image, tile_info))
return processed_tile_images
def _tensor_to_pil(self, tile_tensor) -> Image.Image:
"""テンソルをPIL画像に変換"""
if isinstance(tile_tensor, torch.Tensor):
tile_np = (tile_tensor[0].cpu().numpy() * 255).astype(np.uint8)
return Image.fromarray(tile_np)
return tile_tensor
def _remove_padding(self, tile_image: Image.Image, tile_info: dict) -> Image.Image:
"""パディング部分を除去"""
original_tile_width, original_tile_height = tile_info["original_tile_size"]
content_width = int(original_tile_width * self.upscale_factor)
content_height = int(original_tile_height * self.upscale_factor)
return tile_image.crop((0, 0, content_width, content_height))
def _place_tiles_with_overlap_blend(self, final_image: np.ndarray, weight_map: np.ndarray,
processed_tile_images: list):
"""オーバーラップ領域でブレンドしながらタイルを配置"""
for i, (tile_image, tile_info) in enumerate(processed_tile_images):
grid_x, grid_y = tile_info["grid_position"]
crop_box = tile_info["crop_box"]
overlap_info = tile_info["overlap_info"]
# 元の切り出し位置をアップスケール倍率で拡大
left, top, right, bottom = crop_box
paste_left = int(left * self.upscale_factor)
paste_top = int(top * self.upscale_factor)
paste_right = int(right * self.upscale_factor)
paste_bottom = int(bottom * self.upscale_factor)
# タイル画像を numpy 配列に変換
tile_array = np.array(tile_image, dtype=np.float32)
tile_height, tile_width = tile_array.shape[:2]
# 最終画像の境界内に収まるように調整
paste_right = min(paste_right, final_image.shape[1])
paste_bottom = min(paste_bottom, final_image.shape[0])
actual_width = paste_right - paste_left
actual_height = paste_bottom - paste_top
# タイルサイズが配置領域と異なる場合は調整
if actual_width != tile_width or actual_height != tile_height:
if actual_width > 0 and actual_height > 0:
tile_array = tile_array[:actual_height, :actual_width]
else:
print(f"警告: タイル {i} の配置領域が無効です")
continue
# ブレンド用のウェイトマスクを作成
tile_weight = self._create_blend_weight_mask(
tile_array.shape[:2], overlap_info, crop_box, self.upscale_factor
)
# 加算ブレンド方式で配置
final_image[paste_top:paste_bottom, paste_left:paste_right] += \
tile_array * tile_weight[..., np.newaxis]
weight_map[paste_top:paste_bottom, paste_left:paste_right] += tile_weight
print(f"タイル {i} 配置: 元座標({left},{top},{right},{bottom}) -> "
f"配置座標({paste_left},{paste_top},{paste_right},{paste_bottom}) "
f"サイズ({actual_width}×{actual_height})")
def _create_blend_weight_mask(self, tile_shape: tuple, overlap_info: dict,
crop_box: list, upscale_factor: float) -> np.ndarray:
"""オーバーラップ領域用のウェイトマスクを作成"""
height, width = tile_shape
weight_mask = np.ones((height, width), dtype=np.float32)
# オーバーラップピクセル数を計算(アップスケール後)
overlap_pixels = int(self.tile_size * self.overlap_ratio * upscale_factor)
# 実際の切り出しボックスを使ってオーバーラップ領域を正確に計算
left, top, right, bottom = crop_box
# 左端のフェード(左にオーバーラップがある場合)
if overlap_info["has_left_overlap"]:
# 実際のオーバーラップ開始位置を計算
overlap_start_x = left % self.step_size
if overlap_start_x > 0:
fade_width = min(int(overlap_start_x * upscale_factor), width)
if fade_width > 0:
fade_weights = np.linspace(0, 1, fade_width)
weight_mask[:, :fade_width] *= fade_weights[np.newaxis, :]
# 右端のフェード(右にオーバーラップがある場合)
if overlap_info["has_right_overlap"]:
# 右端のオーバーラップ幅を計算
next_tile_start = (left // self.step_size + 1) * self.step_size
if right > next_tile_start:
overlap_width = right - next_tile_start
fade_width = min(int(overlap_width * upscale_factor), width)
if fade_width > 0:
fade_weights = np.linspace(1, 0, fade_width)
weight_mask[:, -fade_width:] *= fade_weights[np.newaxis, :]
# 上端のフェード(上にオーバーラップがある場合)
if overlap_info["has_top_overlap"]:
# 実際のオーバーラップ開始位置を計算
overlap_start_y = top % self.step_size
if overlap_start_y > 0:
fade_height = min(int(overlap_start_y * upscale_factor), height)
if fade_height > 0:
fade_weights = np.linspace(0, 1, fade_height)
weight_mask[:fade_height, :] *= fade_weights[:, np.newaxis]
# 下端のフェード(下にオーバーラップがある場合)
if overlap_info["has_bottom_overlap"]:
# 下端のオーバーラップ高さを計算
next_tile_start = (top // self.step_size + 1) * self.step_size
if bottom > next_tile_start:
overlap_height = bottom - next_tile_start
fade_height = min(int(overlap_height * upscale_factor), height)
if fade_height > 0:
fade_weights = np.linspace(1, 0, fade_height)
weight_mask[-fade_height:, :] *= fade_weights[:, np.newaxis]
return weight_mask
def _normalize_blended_image(self, final_image: np.ndarray, weight_map: np.ndarray) -> np.ndarray:
"""ブレンドされた画像を正規化"""
# ゼロ除算を避けるため、重みが0の部分には小さな値を設定
weight_map = np.maximum(weight_map, 1e-8)
# 正規化
final_image = final_image / weight_map[..., np.newaxis]
# 値の範囲を0-255にクランプ
final_image = np.clip(final_image, 0, 255)
return final_image
def _convert_to_tensor(self, final_image: Image.Image) -> torch.Tensor:
"""PIL画像をComfyUI形式のテンソルに変換"""
final_tensor = torch.from_numpy(np.array(final_image)).float() / 255.0
return final_tensor.unsqueeze(0) # バッチ次元を追加
class VAEProcessor:
"""VAE処理を担当するクラス"""
@staticmethod
def tiled_decode(vae, latents: torch.Tensor, tile_size: int = 256,
overlap: int = 64, temporal_size: int = 64,
temporal_overlap: int = 8) -> torch.Tensor:
"""タイル化VAEデコード"""
B, C, T, H, W = latents.shape
device = latents.device
dtype = latents.dtype
# 出力テンソルを初期化 (VAEは8倍アップスケール)
output_h, output_w = H * 8, W * 8
output = torch.zeros(B, T, output_h, output_w, 3, device=device, dtype=dtype)
# タイル分割数を計算
h_tiles, w_tiles, t_tiles = VAEProcessor._calculate_tile_counts(
H, W, T, tile_size, overlap, temporal_size, temporal_overlap
)
print(f"タイル化デコード: H={h_tiles}, W={w_tiles}, T={t_tiles}")
# タイル処理実行
VAEProcessor._process_tiles(
vae, latents, output, h_tiles, w_tiles, t_tiles,
tile_size, overlap, temporal_size, temporal_overlap, device, dtype
)
return output
@staticmethod
def _calculate_tile_counts(H: int, W: int, T: int, tile_size: int,
overlap: int, temporal_size: int, temporal_overlap: int) -> tuple:
"""タイル分割数を計算"""
h_tiles = max(1, (H - overlap) // (tile_size - overlap) +
(1 if (H - overlap) % (tile_size - overlap) > 0 else 0))
w_tiles = max(1, (W - overlap) // (tile_size - overlap) +
(1 if (W - overlap) % (tile_size - overlap) > 0 else 0))
t_tiles = max(1, (T - temporal_overlap) // (temporal_size - temporal_overlap) +
(1 if (T - temporal_overlap) % (temporal_size - temporal_overlap) > 0 else 0))
return h_tiles, w_tiles, t_tiles
@staticmethod
def _process_tiles(vae, latents: torch.Tensor, output: torch.Tensor,
h_tiles: int, w_tiles: int, t_tiles: int,
tile_size: int, overlap: int, temporal_size: int,
temporal_overlap: int, device, dtype):
"""実際のタイル処理を実行"""
B, C, T, H, W = latents.shape
for t_idx in range(t_tiles):
t_start, t_end = VAEProcessor._calculate_range(
t_idx, temporal_size, temporal_overlap, T
)
for h_idx in range(h_tiles):
h_start, h_end = VAEProcessor._calculate_range(
h_idx, tile_size, overlap, H
)
for w_idx in range(w_tiles):
w_start, w_end = VAEProcessor._calculate_range(
w_idx, tile_size, overlap, W
)
# タイルを切り出してデコード
tile_latent = latents[:, :, t_start:t_end, h_start:h_end, w_start:w_end]
try:
with torch.autocast(device_type=mm.get_autocast_device(device),
dtype=dtype, enabled=True):
decoded_tile = vae.decode(tile_latent)
# 出力に配置(オーバーラップ処理含む)
VAEProcessor._place_decoded_tile(
output, decoded_tile, t_start, t_end,
h_start, h_end, w_start, w_end, h_idx, w_idx
)
except Exception as e:
print(f"タイル({t_idx},{h_idx},{w_idx})のデコードに失敗: {e}")
continue
@staticmethod
def _calculate_range(idx: int, size: int, overlap: int, max_size: int) -> tuple:
"""範囲を計算"""
start = idx * (size - overlap)
end = min(start + size, max_size)
start = max(0, end - size) # 境界調整
return start, end
@staticmethod
def _place_decoded_tile(output: torch.Tensor, decoded_tile: torch.Tensor,
t_start: int, t_end: int, h_start: int, h_end: int,
w_start: int, w_end: int, h_idx: int, w_idx: int):
"""デコードされたタイルを出力に配置"""
out_t_start, out_t_end = t_start, t_end
out_h_start, out_h_end = h_start * 8, h_end * 8
out_w_start, out_w_end = w_start * 8, w_end * 8
# オーバーラップ処理
if h_idx > 0 or w_idx > 0:
existing = output[:, out_t_start:out_t_end, out_h_start:out_h_end, out_w_start:out_w_end]
if existing.sum() > 0: # 既存データがある場合のみブレンド
output[:, out_t_start:out_t_end, out_h_start:out_h_end, out_w_start:out_w_end] = \
DEFAULT_BLEND_ALPHA * existing + (1 - DEFAULT_BLEND_ALPHA) * decoded_tile
else:
output[:, out_t_start:out_t_end, out_h_start:out_h_end, out_w_start:out_w_end] = decoded_tile
else:
# 最初のタイルはそのまま配置
output[:, out_t_start:out_t_end, out_h_start:out_h_end, out_w_start:out_w_end] = decoded_tile
class FramePackSingleFrameResizeSampler:
"""メインのサンプラークラス"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("FramePackMODEL",),
"vae": ("VAE",),
"positive": ("CONDITIONING",),
"negative": ("CONDITIONING",),
"input_image": ("IMAGE", {"tooltip": "Input image to be upscaled"}),
"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'}),
"tile_size": ("INT", {"default": 256, "min": 256, "max": 1024, "step": 64, "tooltip": "Size of each tile for initial split"}),
"tile_resized": ("INT", {"default": 1024, "min": 512, "max": 2048, "step": 64, "tooltip": "Size to resize each tile to before processing"}),
"vae_tile_size": ("INT", {"default": 256, "min": 64, "max": 512, "step": 64, "tooltip": "Size of each VAE tile for decode"}),
"vae_overlap": ("INT", {"default": 64, "min": 0, "max": 256, "step": 8, "tooltip": "Overlap between VAE tiles"}),
"vae_temporal_size": ("INT", {"default": 64, "min": 1, "max": 128, "step": 1, "tooltip": "Temporal size for VAE tiled decode"}),
"vae_temporal_overlap": ("INT", {"default": 8, "min": 0, "max": 32, "step": 1, "tooltip": "Temporal overlap for VAE tiled decode"}),
},
"optional": {
"image_embeds": ("CLIP_VISION_OUTPUT", ),
"denoise_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
}
}
RETURN_TYPES = ("IMAGE", )
RETURN_NAMES = ("upscaled_image",)
FUNCTION = "process"
CATEGORY = "FramePackWrapper"
DESCRIPTION = "Variable ratio image upscaling using 50% overlapping tile-based FramePack single frame sampling"
def process(self, model, vae, shift, positive, negative, input_image, latent_window_size,
use_teacache, teacache_rel_l1_thresh, steps, cfg, guidance_scale, seed,
sampler, gpu_memory_preservation, tile_size=512, tile_resized=1024,
vae_tile_size=256, vae_overlap=64, vae_temporal_size=64,
vae_temporal_overlap=8, image_embeds=None, denoise_strength=1.0):
upscale_factor = float(tile_resized) / float(tile_size)
self._log_process_start(upscale_factor, vae_tile_size, vae_overlap,
vae_temporal_size, vae_temporal_overlap, seed)
# デバイス設定とモデル準備
transformer, base_dtype, device, offload_device = self._setup_model(model)
# 1. 画像を50%オーバーラップでタイルに分割
tile_processor = TileProcessor(tile_size, tile_resized)
tiles, metadata = tile_processor.split_image_to_tiles(input_image)
# 2. テキストエンベッディング準備
text_embeddings = self._prepare_text_embeddings(
positive, negative, cfg, device, base_dtype
)
# 3. 画像エンベッディング準備
image_encoder_embeddings = self._prepare_image_embeddings(
image_embeds, device, base_dtype
)
# 4. ComfyUIセットアップ
callback = self._setup_comfy_ui(transformer, device, steps, base_dtype)
# 5. モデル準備
self._prepare_model(transformer, device, gpu_memory_preservation,
use_teacache, steps, teacache_rel_l1_thresh)
# 6. 各タイルを処理
processed_tiles = self._process_all_tiles(
tiles, vae, transformer, sampler, denoise_strength, latent_window_size,
cfg, guidance_scale, shift, steps, seed, device, base_dtype,
text_embeddings, image_encoder_embeddings, callback,
vae_tile_size, vae_overlap, vae_temporal_size, vae_temporal_overlap
)
# 7. 50%オーバーラップブレンドで最終画像を再構成
reconstructor = ImageReconstructor(metadata)
final_image = reconstructor.reconstruct_image(processed_tiles)
# クリーンアップ
self._cleanup(transformer, offload_device)
print(f"=== 50%オーバーラップ {upscale_factor:.2f}xアップスケール完了: {final_image.shape} ===")
return (final_image,)
def _log_process_start(self, upscale_factor: float, vae_tile_size: int,
vae_overlap: int, vae_temporal_size: int,
vae_temporal_overlap: int, seed: int):
"""処理開始のログ出力"""
print(f"=== FramePack 50%オーバーラップ Resize Sampler 開始 ({upscale_factor:.2f}x) ===")
print(f"VAE Tile設定: size={vae_tile_size}, overlap={vae_overlap}, "
f"temporal_size={vae_temporal_size}, temporal_overlap={vae_temporal_overlap}")
print(f"使用シード: {seed} (全タイル共通)")
print(f"オーバーラップ比率: {OVERLAP_RATIO*100:.0f}%")
def _setup_model(self, model: dict) -> tuple:
"""モデルとデバイスのセットアップ"""
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()
return transformer, base_dtype, device, offload_device
def _prepare_text_embeddings(self, positive, negative, cfg: float,
device, base_dtype) -> dict:
"""テキストエンベッディングを準備"""
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)
return {
"llama_vec": llama_vec,
"llama_attention_mask": llama_attention_mask,
"clip_l_pooler": clip_l_pooler,
"llama_vec_n": llama_vec_n,
"llama_attention_mask_n": llama_attention_mask_n,
"clip_l_pooler_n": clip_l_pooler_n
}
def _prepare_image_embeddings(self, image_embeds, device, base_dtype):
"""画像エンベッディングを準備"""
if image_embeds is not None:
return image_embeds["last_hidden_state"].to(device, base_dtype)
return None
def _setup_comfy_ui(self, transformer, device, steps, base_dtype):
"""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"))
return prepare_callback(patcher, steps)
def _prepare_model(self, transformer, device, gpu_memory_preservation: float,
use_teacache: bool, steps: int, teacache_rel_l1_thresh: float):
"""モデルの準備(GPU移動とTeaCache設定)"""
# モデルを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)
def _process_all_tiles(self, tiles: list, vae, transformer, sampler: str,
denoise_strength: float, latent_window_size: int,
cfg: float, guidance_scale: float, shift: float,
steps: int, seed: int, device, base_dtype,
text_embeddings: dict, image_encoder_embeddings,
callback, vae_tile_size: int, vae_overlap: int,
vae_temporal_size: int, vae_temporal_overlap: int) -> list:
"""全てのタイルを処理"""
processed_tiles = []
print(f"=== {len(tiles)}個のタイルを順次処理(50%オーバーラップ、各タイルで同一シード使用) ===")
for i, tile_tensor in enumerate(tqdm(tiles, desc="オーバーラップタイル処理")):
print(f"\n--- タイル {i+1}/{len(tiles)} 処理中(シード: {seed}) ---")
# 各タイルで同じシードを使用するためにジェネレータをリセット
rnd = torch.Generator("cpu").manual_seed(seed)
processed_tile = self._process_single_tile(
tile_tensor, vae, transformer, sampler, denoise_strength,
latent_window_size, cfg, guidance_scale, shift, steps, rnd,
device, base_dtype, text_embeddings, image_encoder_embeddings,
callback, vae_tile_size, vae_overlap, vae_temporal_size, vae_temporal_overlap
)
processed_tiles.append(processed_tile)
print(f"オーバーラップタイル {i+1} 完了: {processed_tile.shape}")
return processed_tiles
def _process_single_tile(self, tile_tensor: torch.Tensor, vae, transformer,
sampler: str, denoise_strength: float, latent_window_size: int,
cfg: float, guidance_scale: float, shift: float,
steps: int, rnd: torch.Generator, device, base_dtype,
text_embeddings: dict, image_encoder_embeddings,
callback, vae_tile_size: int, vae_overlap: int,
vae_temporal_size: int, vae_temporal_overlap: int) -> torch.Tensor:
"""単一タイルを処理"""
tile_tensor = tile_tensor.to(device)
with torch.autocast(device_type=mm.get_autocast_device(device),
dtype=base_dtype, enabled=True):
# VAEエンコード
start_latent = self._encode_tile(vae, tile_tensor)
# FramePack推論のためのインデックス設定
latent_indices, clean_latents, clean_latent_indices = self._setup_latent_indices(
start_latent, latent_window_size
)
# FramePack推論実行
generated_latents = self._run_framepack_inference(
transformer, sampler, denoise_strength, start_latent.shape,
cfg, guidance_scale, shift, steps, rnd, device, base_dtype,
text_embeddings, image_encoder_embeddings, latent_indices,
clean_latents, clean_latent_indices, callback
)
# VAEタイル化デコード
decoded_tile = self._decode_tile(
vae, generated_latents, vae_tile_size, vae_overlap,
vae_temporal_size, vae_temporal_overlap
)
# 時間次元を除去して画像形式に変換 [B, T, H, W, C] -> [B, H, W, C]
if len(decoded_tile.shape) == 5:
decoded_tile = decoded_tile.squeeze(1)
return decoded_tile
def _encode_tile(self, vae, tile_tensor: torch.Tensor) -> torch.Tensor:
"""タイルをVAEエンコード"""
start_latent = vae.encode(tile_tensor)
start_latent = start_latent * VAE_SCALING_FACTOR
print(f"エンコード後latent: {start_latent.shape}")
# latentを5次元に変換 [B, C, T, H, W]
if len(start_latent.shape) == 4:
start_latent = start_latent.unsqueeze(2) # T次元を追加
return start_latent
def _setup_latent_indices(self, start_latent: torch.Tensor,
latent_window_size: int) -> tuple:
"""FramePack推論のためのインデックス設定"""
B, C, T, H, W = start_latent.shape
# 1フレーム推論のためのインデックス設定
sample_num_frames = 1
latent_padding = 0
latent_padding_size = latent_padding * latent_window_size
indices = torch.arange(0, sum([1, latent_padding_size, latent_window_size, 1, 2, 16])).unsqueeze(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:
split_sizes = [1, latent_padding_size, latent_window_size, 1, 2, 16]
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フレーム推論設定
all_indices = torch.arange(0, latent_window_size).unsqueeze(0)
latent_indices = all_indices[:, -1:]
clean_latents_pre = start_latent.to(torch.float32).cpu()
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.tensor([[0]], dtype=clean_latent_indices_pre.dtype, device=clean_latent_indices_pre.device)
clean_latents = clean_latents[:, :, :1, :, :]
return latent_indices, clean_latents, clean_latent_indices
def _run_framepack_inference(self, transformer, sampler: str, denoise_strength: float,
latent_shape: tuple, cfg: float, guidance_scale: float,
shift: float, steps: int, rnd: torch.Generator,
device, base_dtype, text_embeddings: dict,
image_encoder_embeddings, latent_indices: torch.Tensor,
clean_latents: torch.Tensor, clean_latent_indices: torch.Tensor,
callback) -> torch.Tensor:
"""FramePack推論を実行"""
B, C, T, H, W = latent_shape
sample_num_frames = 1
generated_latents = sample_hunyuan(
transformer=transformer,
sampler=sampler,
initial_latent=None,
strength=denoise_strength,
width=W * 8,
height=H * 8,
frames=sample_num_frames,
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=text_embeddings["llama_vec"],
prompt_embeds_mask=text_embeddings["llama_attention_mask"],
prompt_poolers=text_embeddings["clip_l_pooler"],
negative_prompt_embeds=text_embeddings["llama_vec_n"],
negative_prompt_embeds_mask=text_embeddings["llama_attention_mask_n"],
negative_prompt_poolers=text_embeddings["clip_l_pooler_n"],
device=device,
dtype=base_dtype,
image_embeddings=image_encoder_embeddings,
latent_indices=latent_indices,
clean_latents=clean_latents,
clean_latent_indices=clean_latent_indices,
clean_latents_2x=None, # 2x, 4x latentsを無効化
clean_latent_2x_indices=None,
clean_latents_4x=None,
clean_latent_4x_indices=None,
callback=callback,
)
return generated_latents
def _decode_tile(self, vae, generated_latents: torch.Tensor, vae_tile_size: int,
vae_overlap: int, vae_temporal_size: int, vae_temporal_overlap: int) -> torch.Tensor:
"""タイルをVAEデコード"""
generated_latents = generated_latents / VAE_SCALING_FACTOR
print(f"デコード前latent: {generated_latents.shape}")
decoded_tile = VAEProcessor.tiled_decode(
vae, generated_latents,
tile_size=vae_tile_size,
overlap=vae_overlap,
temporal_size=vae_temporal_size,
temporal_overlap=vae_temporal_overlap
)
return decoded_tile
def _cleanup(self, transformer, offload_device):
"""リソースのクリーンアップ"""
transformer.to(offload_device)
mm.soft_empty_cache()
# ノードクラスマッピング
NODE_CLASS_MAPPINGS = {
"FramePackSingleFrameResizeSampler": FramePackSingleFrameResizeSampler,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"FramePackSingleFrameResizeSampler": "FramePack Single Frame Resize Sampler (50% Overlap)",
}