diff --git a/__init__.py b/__init__.py index 10130ef..51b5973 100644 --- a/__init__.py +++ b/__init__.py @@ -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'] \ No newline at end of file diff --git a/nodes_resize.py b/nodes_resize.py new file mode 100644 index 0000000..ac6a1b8 --- /dev/null +++ b/nodes_resize.py @@ -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)", +} \ No newline at end of file