import copy import torch import comfy.sample import comfy.utils import comfy.model_management import latent_preview from comfy.nested_tensor import NestedTensor class H3LoopingSampler: @classmethod def INPUT_TYPES(s): return { "required": { "noise": ("NOISE",), "guider": ("GUIDER",), "sampler": ("SAMPLER",), "sigmas": ("SIGMAS",), "latent_image": ("LATENT",), "temporal_tile_size": ( "INT", { "default": 81, "min": 17, "max": 257, "step": 4, "tooltip": "Tamanho do tile temporal em frames de latente (vídeo)", }, ), "temporal_overlap": ( "INT", { "default": 17, "min": 5, "max": 65, "step": 4, }, ), "temporal_overlap_strength": ( "FLOAT", { "default": 0.65, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Força da nova chunk na zona de overlap (0 = mantém só o anterior, 1 = blend total)", }, ), "horizontal_tiles": ("INT", {"default": 1, "min": 1, "max": 4}), "vertical_tiles": ("INT", {"default": 1, "min": 1, "max": 4}), "spatial_overlap": ( "INT", { "default": 8, "min": 0, "max": 32, }, ), }, "optional": { "adain_factor": ( "FLOAT", { "default": 0.15, "min": 0.0, "max": 1.0, "step": 0.01, }, ), }, } RETURN_TYPES = ("LATENT", "LATENT") RETURN_NAMES = ("output", "denoised_output") FUNCTION = "sample" CATEGORY = "sampling/H3" DESCRIPTION = "H3 Looping / Tiled Sampler com output + denoised_output (compatível com SplitSigmas) ComfyGuy9000" def _is_nested(self, samples): return isinstance(samples, NestedTensor) or getattr(samples, "is_nested", False) def _get_tensors(self, samples): if self._is_nested(samples): if hasattr(samples, "tensors"): return list(samples.tensors) return list(samples.unbind()) return [samples] def _make_nested(self, tensors): if len(tensors) == 1: return tensors[0] return NestedTensor(tensors) def _get_video(self, samples): return self._get_tensors(samples)[0] def _slice_video_temporal(self, video, start, end): return video[:, :, start:end].clone() def _slice_video_spatial(self, video, v_start, v_end, h_start, h_end): return video[:, :, :, v_start:v_end, h_start:h_end].clone() def _create_spatial_weights(self, shape, v, h, vertical_tiles, horizontal_tiles, spatial_overlap, device, dtype): weights = torch.ones(shape, device=device, dtype=dtype) if spatial_overlap > 0: if h > 0: blend = torch.linspace(0, 1, spatial_overlap, device=device, dtype=dtype) weights[..., :spatial_overlap] *= blend.view(1, 1, 1, 1, -1) if h < horizontal_tiles - 1: blend = torch.linspace(1, 0, spatial_overlap, device=device, dtype=dtype) weights[..., -spatial_overlap:] *= blend.view(1, 1, 1, 1, -1) if v > 0: blend = torch.linspace(0, 1, spatial_overlap, device=device, dtype=dtype) weights[..., :spatial_overlap, :] *= blend.view(1, 1, 1, -1, 1) if v < vertical_tiles - 1: blend = torch.linspace(1, 0, spatial_overlap, device=device, dtype=dtype) weights[..., -spatial_overlap:, :] *= blend.view(1, 1, 1, -1, 1) return weights def _adain(self, source, target, factor): if factor <= 0.0: return source src_mean = source.mean(dim=(2, 3, 4), keepdim=True) src_std = source.std(dim=(2, 3, 4), keepdim=True) + 1e-5 tgt_mean = target.mean(dim=(2, 3, 4), keepdim=True) tgt_std = target.std(dim=(2, 3, 4), keepdim=True) + 1e-5 normalized = (source - src_mean) / src_std stylized = normalized * tgt_std + tgt_mean return source * (1.0 - factor) + stylized * factor def sample( self, noise, guider, sampler, sigmas, latent_image, temporal_tile_size, temporal_overlap, temporal_overlap_strength, horizontal_tiles, vertical_tiles, spatial_overlap, adain_factor=0.15, ): original_latent = latent_image samples = latent_image["samples"] video = self._get_video(samples) if video.ndim != 5: raise ValueError(f"Expected video [B,C,T,H,W], got {tuple(video.shape)}") B, C, T, H, W = video.shape print(f"\n========== H3LoopingSampler ComfyGuy9000 ==========") print(f"Input video latent: {video.shape}") print(f"Tiles: {vertical_tiles}x{horizontal_tiles} | spatial_overlap={spatial_overlap}") print(f"Temporal tile={temporal_tile_size} | overlap={temporal_overlap} | strength={temporal_overlap_strength}") original_tensors = self._get_tensors(samples) has_audio = len(original_tensors) > 1 full_audio = original_tensors[1] if has_audio else None temporal_tile_size = min(temporal_tile_size, T) temporal_overlap = min(temporal_overlap, max(4, temporal_tile_size - 4)) if vertical_tiles > 1: base_tile_h = (H + (vertical_tiles - 1) * spatial_overlap) // vertical_tiles else: base_tile_h = H if horizontal_tiles > 1: base_tile_w = (W + (horizontal_tiles - 1) * spatial_overlap) // horizontal_tiles else: base_tile_w = W print(f"Base tile size (latent): {base_tile_h} x {base_tile_w}") final_video = None final_denoised_video = None weights = None first_seed = noise.seed disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED tile_count = 0 for v in range(vertical_tiles): for h in range(horizontal_tiles): v_start = v * (base_tile_h - spatial_overlap) h_start = h * (base_tile_w - spatial_overlap) v_end = min(v_start + base_tile_h, H) if v < vertical_tiles - 1 else H h_end = min(h_start + base_tile_w, W) if h < horizontal_tiles - 1 else W tile_count += 1 print(f"\n>>> Spatial tile {tile_count}/{vertical_tiles*horizontal_tiles} ({v},{h})") print(f" H[{v_start}:{v_end}] W[{h_start}:{h_end}]") spatial_video = self._slice_video_spatial(video, v_start, v_end, h_start, h_end) tile_out_video = None tile_denoised_video = None first_chunk_ref = None step = max(1, temporal_tile_size - temporal_overlap) starts = list(range(0, max(1, T - temporal_overlap), step)) for i, start in enumerate(starts): end = min(start + temporal_tile_size, T) print(f" Temporal chunk {i}: [{start}:{end}]") chunk_video = self._slice_video_temporal(spatial_video, start, end) if has_audio: chunk_samples = self._make_nested([chunk_video, full_audio]) else: chunk_samples = chunk_video chunk_latent = {"samples": chunk_samples} if "noise_mask" in latent_image: chunk_latent["noise_mask"] = latent_image["noise_mask"] noise.seed = first_seed + start * (vertical_tiles * horizontal_tiles) + v * horizontal_tiles + h # === Captura do x0 (denoised) === x0_output = {} callback = latent_preview.prepare_callback( guider.model_patcher, sigmas.shape[-1] - 1, x0_output ) noise_mask = chunk_latent.get("noise_mask", None) out_samples = guider.sample( noise.generate_noise(chunk_latent), chunk_samples, sampler, sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise.seed, ) out_samples = out_samples.to(comfy.model_management.intermediate_device()) chunk_out_video = self._get_video(out_samples) # Pega a versão denoised (x0) se disponível if "x0" in x0_output: x0 = x0_output["x0"] if self._is_nested(out_samples) and not self._is_nested(x0): try: latent_shapes = [t.shape for t in self._get_tensors(out_samples)] x0 = NestedTensor(comfy.utils.unpack_latents(x0, latent_shapes)) except: pass chunk_denoised_video = self._get_video(x0) try: chunk_denoised_video = guider.model_patcher.model.process_latent_out( chunk_denoised_video.cpu() ).to(chunk_out_video.device) except: chunk_denoised_video = chunk_denoised_video.to(chunk_out_video.device) else: chunk_denoised_video = chunk_out_video # AdaIN (aplica nos dois) if first_chunk_ref is None: first_chunk_ref = chunk_out_video.detach() else: ref = first_chunk_ref if ref.shape[2] != chunk_out_video.shape[2]: ref = first_chunk_ref[:, :, :1].expand_as(chunk_out_video) else: ref = first_chunk_ref[:, :, :chunk_out_video.shape[2]] chunk_out_video = self._adain(chunk_out_video, ref, adain_factor) chunk_denoised_video = self._adain(chunk_denoised_video, ref, adain_factor) # === Blend temporal (CORRIGIDO) === if tile_out_video is None: tile_out_video = chunk_out_video tile_denoised_video = chunk_denoised_video else: overlap = temporal_overlap if overlap > 0 and tile_out_video.shape[2] >= overlap: alpha = torch.linspace( 1.0, 0.0, overlap, device=tile_out_video.device, dtype=tile_out_video.dtype ).view(1, 1, -1, 1, 1) # Fórmula corrigida: # strength = 0.0 → mantém só o anterior # strength = 1.0 → blend linear normal prev = tile_out_video[:, :, -overlap:] new = chunk_out_video[:, :, :overlap] blended = prev * (1.0 - (1.0 - alpha) * temporal_overlap_strength) + \ new * (1.0 - alpha) * temporal_overlap_strength tile_out_video = torch.cat( [tile_out_video[:, :, :-overlap], blended, chunk_out_video[:, :, overlap:]], dim=2 ) # Blend denoised prev_d = tile_denoised_video[:, :, -overlap:] new_d = chunk_denoised_video[:, :, :overlap] blended_d = prev_d * (1.0 - (1.0 - alpha) * temporal_overlap_strength) + \ new_d * (1.0 - alpha) * temporal_overlap_strength tile_denoised_video = torch.cat( [tile_denoised_video[:, :, :-overlap], blended_d, chunk_denoised_video[:, :, overlap:]], dim=2 ) else: tile_out_video = torch.cat([tile_out_video, chunk_out_video], dim=2) tile_denoised_video = torch.cat([tile_denoised_video, chunk_denoised_video], dim=2) # Acumula spatial if final_video is None: out_T = tile_out_video.shape[2] final_video = torch.zeros(B, C, out_T, H, W, device=tile_out_video.device, dtype=tile_out_video.dtype) final_denoised_video = torch.zeros_like(final_video) weights = torch.zeros_like(final_video) if tile_out_video.shape[2] != final_video.shape[2]: if tile_out_video.shape[2] > final_video.shape[2]: tile_out_video = tile_out_video[:, :, :final_video.shape[2]] tile_denoised_video = tile_denoised_video[:, :, :final_video.shape[2]] else: pad = final_video.shape[2] - tile_out_video.shape[2] tile_out_video = torch.nn.functional.pad(tile_out_video, (0, 0, 0, 0, 0, pad)) tile_denoised_video = torch.nn.functional.pad(tile_denoised_video, (0, 0, 0, 0, 0, pad)) w = self._create_spatial_weights( tile_out_video.shape, v, h, vertical_tiles, horizontal_tiles, spatial_overlap, tile_out_video.device, tile_out_video.dtype ) tile_out_video = tile_out_video.to(final_video.device) tile_denoised_video = tile_denoised_video.to(final_video.device) w = w.to(final_video.device) final_video[:, :, :, v_start:v_end, h_start:h_end] += tile_out_video * w final_denoised_video[:, :, :, v_start:v_end, h_start:h_end] += tile_denoised_video * w weights[:, :, :, v_start:v_end, h_start:h_end] += w final_video = final_video / (weights + 1e-8) final_denoised_video = final_denoised_video / (weights + 1e-8) noise.seed = first_seed # Monta NestedTensor para as duas saídas def make_output_latent(video_tensor): out_tensors = [video_tensor] if has_audio: out_tensors.append(full_audio.to(video_tensor.device)) out_samples = self._make_nested(out_tensors) out_latent = copy.deepcopy(original_latent) out_latent["samples"] = out_samples return out_latent output_latent = make_output_latent(final_video) denoised_latent = make_output_latent(final_denoised_video) print(f"\n[H3LoopingSampler] Final video shape: {final_video.shape}") print(f"Total spatial tiles: {tile_count}") print("Saídas: output + denoised_output") print("========================================\n") return (output_latent, denoised_latent) NODE_CLASS_MAPPINGS = { "H3LoopingSampler": H3LoopingSampler } NODE_DISPLAY_NAME_MAPPINGS = { "H3LoopingSampler": "H3 Looping / Tiled Sampler (ComfyGuy9000)" }