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)", }