Upload H3Loopsampler.py with huggingface_hub

H3Loopsampler.py

@@ -0,0 +1,374 @@
+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)"
+}