Automated PR - 2025-04-17

2025-04-17 15:00:46 +00:00
parent e2a73515dd
commit bdde649e25
22 changed files with 4720 additions and 2642 deletions
@@ -4,7 +4,19 @@ ComfyUI-LTXVideo is a collection of custom nodes for ComfyUI, designed to provid
 The model itself is supported in the core ComfyUI [code](https://github.com/comfyanonymous/ComfyUI/tree/master/comfy/ldm/lightricks).
 The main LTXVideo repository can be found [here](https://github.com/Lightricks/LTX-Video).

-# 5.03.2025 :star: LTXVideo 0.9.5 Release :star:
+# 17.04.2025 ⭐ LTXVideo 0.9.6 Release ⭐
+
+### LTXVideo 0.9.6 introduces:
+
+1. LTXV 0.9.6 – higher quality, faster, great for final output. Download from [here](https://huggingface.co/Lightricks/LTX-Video/resolve/main/ltxv-2b-0.9.6-dev-04-25.safetensors).
+2. LTXV 0.9.6 Distilled – our fastest model yet (only 8 steps for generation), lighter, great for rapid iteration. Download from [here](https://huggingface.co/Lightricks/LTX-Video/resolve/main/ltxv-2b-0.9.6-distilled-04-25.safetensors).
+
+### Technical Updates
+
+We introduce the __STGGuiderAdvanced__ node, which applies different CFG and STG parameters at various diffusion steps. All flows have been updated to use this node and are designed to provide optimal parameters for the best quality.
+See the [Example Workflows](#example-workflows) section.
+
+# 5.03.2025 ⭐ LTXVideo 0.9.5 Release ⭐

 ### LTXVideo 0.9.5 introduces:

@@ -73,6 +85,10 @@ Note that to run the example workflows, you need to have some additional custom
 [Download workflow](assets/ltxvideo-i2v.json)
 ![workflow](assets/ltxvideo-i2v.png)

+### Image-to-video-distilled
+[Download workflow](assets/ltxvideo-i2v-distilled.json)
+![workflow](assets/ltxvideo-i2v-distilled.png)
+
 ### Text-to-video

 [Download workflow](assets/ltxvideo-t2v.json)
@@ -1,9 +1,14 @@
+from .latents import LTXVSelectLatents, LTXVSetVideoLatentNoiseMasks  # noqa: F401
 from .nodes_registry import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
 from .prompt_enhancer_nodes import (  # noqa: F401
    LTXVPromptEnhancer,
    LTXVPromptEnhancerLoader,
 )
-from .stg import LTXVApplySTG, STGGuiderNode  # noqa: F401
+from .recurrent_sampler import (  # noqa: F401
+    LinearOverlapLatentTransition,
+    LTXVRecurrentKSampler,
+)
+from .stg import LTXVApplySTG, STGGuiderAdvancedNode, STGGuiderNode  # noqa: F401
 from .tricks import NODE_CLASS_MAPPINGS as TRICKS_NODE_CLASS_MAPPINGS
 from .tricks import NODE_DISPLAY_NAME_MAPPINGS as TRICKS_NODE_DISPLAY_NAME_MAPPINGS

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:adab606ec23e98c5aadb7adbfa9d16cb1e9d460004bbe648cb03f5e89972229a
-size 1267760
+oid sha256:7899cb4e0e23f2439cc489b1ed6191eef91e91dbefdcb2c4c29545d5a8cd6c3a
+size 1292625
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:97ea4a9a45004726ae84971bef44578270a7be986b1fce533cc8d3ec1a213fdd
-size 1390379
+oid sha256:77d9f98afa1a5711fade39a6fa08b9c3eab8c439c3727f522bae639bee11cf0a
+size 1380444
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b981ed68d5f4a25e48daf31cf4cf0e156bee27b75730dd5fd213bfb805d1c8d
+size 1342913
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3e29e3b630ab07c3404e715dd0f5d62ebdb4ef24af4279239e546f2a401824a3
-size 1167784
+oid sha256:5cb7cbf30c7ad19dda71322c40e7a471697350b6ec6b17c8ab14654dad18b909
+size 1142742
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7ba0a0ab34867f35268823c6ef8ca4cd8067cd6484e0c6f5b1277c9778bb243e
-size 1379716
+oid sha256:006feae06d8679de9175eae516bf0919b273745859bb045b3d229329a34084fb
+size 1160850
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0823a77dec3324a8f696072a3988fe50431ce89482698e09047d32852f886cdc
-size 770420
+oid sha256:7a4f2fe897fd955633099a429aa73bccc2e6007627a3e1bf419b4b8db67bce80
+size 817453
@@ -0,0 +1,351 @@
+from typing import Optional
+
+import comfy_extras.nodes_lt as nodes_lt
+import torch
+
+from .nodes_registry import comfy_node
+
+
+@comfy_node(name="LTXVSelectLatents")
+class LTXVSelectLatents:
+    """
+    Selects a range of frames from a video latent.
+
+    Features:
+    - Supports positive and negative indexing
+    - Preserves batch processing capabilities
+    - Handles noise masks if present
+    - Maintains 5D tensor format
+    """
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "samples": ("LATENT",),
+                "start_index": (
+                    "INT",
+                    {"default": 0, "min": -9999, "max": 9999, "step": 1},
+                ),
+                "end_index": (
+                    "INT",
+                    {"default": -1, "min": -9999, "max": 9999, "step": 1},
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "select_latents"
+    CATEGORY = "latent/video"
+
+    def select_latents(self, samples: dict, start_index: int, end_index: int) -> tuple:
+        """
+        Selects a range of frames from the video latent.
+
+        Args:
+            samples (dict): Video latent dictionary
+            start_index (int): Starting frame index (supports negative indexing)
+            end_index (int): Ending frame index (supports negative indexing)
+
+        Returns:
+            tuple: Contains modified latent dictionary with selected frames
+
+        Raises:
+            ValueError: If indices are invalid
+        """
+        try:
+            s = samples.copy()
+            video_latent = s["samples"]
+            batch, channels, frames, height, width = video_latent.shape
+
+            # Handle negative indices
+            start_idx = frames + start_index if start_index < 0 else start_index
+            end_idx = frames + end_index if end_index < 0 else end_index
+
+            # Validate and clamp indices
+            start_idx = max(0, min(start_idx, frames - 1))
+            end_idx = max(0, min(end_idx, frames - 1))
+            if start_idx > end_idx:
+                start_idx = min(start_idx, end_idx)
+
+            # Select frames while maintaining 5D format
+            s["samples"] = video_latent[:, :, start_idx : end_idx + 1, :, :]
+
+            # Handle noise mask if present
+            if "noise_mask" in s:
+                s["noise_mask"] = s["noise_mask"][:, :, start_idx : end_idx + 1, :, :]
+
+            return (s,)
+
+        except Exception as e:
+            print(f"[LTXVSelectLatents] Error: {str(e)}")
+            raise
+
+
+@comfy_node(name="LTXVAddLatents")
+class LTXVAddLatents:
+    """
+    Concatenates two video latents along the frames dimension.
+
+    Features:
+    - Validates dimension compatibility
+    - Handles device placement
+    - Preserves noise masks with proper handling
+    - Supports batch processing
+    """
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "latents1": ("LATENT",),
+                "latents2": ("LATENT",),
+            }
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "add_latents"
+    CATEGORY = "latent/video"
+
+    def add_latents(
+        self, latents1: torch.Tensor, latents2: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Concatenates two video latents along the frames dimension.
+
+        Args:
+            latents1 (dict): First video latent dictionary
+            latents2 (dict): Second video latent dictionary
+
+        Returns:
+            tuple: Contains concatenated latent dictionary
+
+        Raises:
+            ValueError: If latent dimensions don't match
+            RuntimeError: If tensor operations fail
+        """
+        try:
+            s = latents1.copy()
+            video_latent1 = latents1["samples"]
+            video_latent2 = latents2["samples"]
+
+            # Ensure tensors are on the same device
+            target_device = video_latent1.device
+            video_latent2 = video_latent2.to(target_device)
+
+            # Validate dimensions
+            self._validate_dimensions(video_latent1, video_latent2)
+
+            # Concatenate along frames dimension
+            s["samples"] = torch.cat([video_latent1, video_latent2], dim=2)
+
+            # Handle noise masks
+            s["noise_mask"] = self._merge_noise_masks(
+                latents1, latents2, video_latent1.shape[2], video_latent2.shape[2]
+            )
+
+            return (s,)
+
+        except Exception as e:
+            print(f"[LTXVAddLatents] Error: {str(e)}")
+            raise
+
+    def _validate_dimensions(self, latent1: torch.Tensor, latent2: torch.Tensor):
+        """Validates that latent dimensions match except for frames."""
+        b1, c1, f1, h1, w1 = latent1.shape
+        b2, c2, f2, h2, w2 = latent2.shape
+
+        if not (b1 == b2 and c1 == c2 and h1 == h2 and w1 == w2):
+            raise ValueError(
+                f"Latent dimensions must match (except frames dimension).\n"
+                f"Got shapes {latent1.shape} and {latent2.shape}"
+            )
+
+    def _merge_noise_masks(
+        self, latents1: torch.Tensor, latents2: torch.Tensor, frames1: int, frames2: int
+    ) -> Optional[torch.Tensor]:
+        """Merges noise masks from both latents with proper handling."""
+        if "noise_mask" in latents1 and "noise_mask" in latents2:
+            return torch.cat([latents1["noise_mask"], latents2["noise_mask"]], dim=2)
+        elif "noise_mask" in latents1:
+            zeros = torch.zeros_like(latents1["noise_mask"][:, :, :frames2, :, :])
+            return torch.cat([latents1["noise_mask"], zeros], dim=2)
+        elif "noise_mask" in latents2:
+            zeros = torch.zeros_like(latents2["noise_mask"][:, :, :frames1, :, :])
+            return torch.cat([zeros, latents2["noise_mask"]], dim=2)
+        return None
+
+
+@comfy_node(name="LTXVSetVideoLatentNoiseMasks")
+class LTXVSetVideoLatentNoiseMasks:
+    """
+    Applies multiple masks to a video latent.
+
+    Features:
+    - Supports multiple input mask formats (2D, 3D, 4D)
+    - Automatically handles fewer masks than frames by reusing the last mask
+    - Resizes masks to match latent dimensions
+    - Preserves batch processing capabilities
+
+    Input Formats:
+    - 2D mask: Single mask [H, W]
+    - 3D mask: Multiple masks [M, H, W]
+    - 4D mask: Multiple masks with channels [M, C, H, W]
+    """
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "samples": ("LATENT",),
+                "masks": ("MASK",),
+            }
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "set_mask"
+    CATEGORY = "latent/video"
+
+    def set_mask(self, samples: dict, masks: torch.Tensor) -> tuple:
+        """
+        Applies masks to video latent frames.
+
+        Args:
+            samples (dict): Video latent dictionary containing 'samples' tensor
+            masks (torch.Tensor): Mask tensor in various possible formats
+                - 2D: [H, W] single mask
+                - 3D: [M, H, W] multiple masks
+                - 4D: [M, C, H, W] multiple masks with channels
+
+        Returns:
+            tuple: Contains modified latent dictionary with applied masks
+
+        Raises:
+            ValueError: If mask dimensions are unsupported
+            RuntimeError: If tensor operations fail
+        """
+        try:
+            s = samples.copy()
+            video_latent = s["samples"]
+            batch_size, channels, num_frames, height, width = video_latent.shape
+
+            # Initialize noise_mask if not present
+            if "noise_mask" not in s:
+                s["noise_mask"] = torch.zeros(
+                    (batch_size, 1, num_frames, height, width),
+                    dtype=video_latent.dtype,
+                    device=video_latent.device,
+                )
+
+            # Process masks
+            masks_reshaped = self._reshape_masks(masks)
+            M = masks_reshaped.shape[0]
+            resized_masks = self._resize_masks(masks_reshaped, height, width)
+
+            # Apply masks efficiently
+            self._apply_masks(s["noise_mask"], resized_masks, num_frames, M)
+            return (s,)
+
+        except Exception as e:
+            print(f"[LTXVSetVideoLatentNoiseMasks] Error: {str(e)}")
+            raise
+
+    def _reshape_masks(self, masks: torch.Tensor) -> torch.Tensor:
+        """Reshapes input masks to consistent 4D format."""
+        original_shape = tuple(masks.shape)
+        ndims = masks.ndim
+
+        if ndims == 2:
+            return masks.unsqueeze(0).unsqueeze(0)
+        elif ndims == 3:
+            return masks.reshape(masks.shape[0], 1, masks.shape[1], masks.shape[2])
+        elif ndims == 4:
+            return masks.reshape(masks.shape[0], 1, masks.shape[2], masks.shape[3])
+        else:
+            raise ValueError(
+                f"Unsupported 'masks' dimension: {original_shape}. "
+                "Must be 2D (H,W), 3D (M,H,W), or 4D (M,C,H,W)."
+            )
+
+    def _resize_masks(
+        self, masks: torch.Tensor, height: int, width: int
+    ) -> torch.Tensor:
+        """Resizes all masks to match latent dimensions."""
+        return torch.nn.functional.interpolate(
+            masks, size=(height, width), mode="bilinear", align_corners=False
+        )
+
+    def _apply_masks(
+        self,
+        noise_mask: torch.Tensor,
+        resized_masks: torch.Tensor,
+        num_frames: int,
+        M: int,
+    ) -> None:
+        """Applies resized masks to all frames."""
+        for f in range(num_frames):
+            mask_idx = min(f, M - 1)  # Reuse last mask if we run out
+            noise_mask[:, :, f] = resized_masks[mask_idx]
+
+
+@comfy_node(name="LTXVAddLatentGuide")
+class LTXVAddLatentGuide(nodes_lt.LTXVAddGuide):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "vae": ("VAE",),
+                "positive": ("CONDITIONING",),
+                "negative": ("CONDITIONING",),
+                "latent": ("LATENT",),
+                "guiding_latent": ("LATENT",),
+                "latent_idx": (
+                    "INT",
+                    {
+                        "default": 0,
+                        "min": -9999,
+                        "max": 9999,
+                        "step": 1,
+                        "tooltip": "Latent index to start the conditioning at. Can be negative to"
+                        "indicate that the conditioning is on the frames before the latent.",
+                    },
+                ),
+                "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}),
+            }
+        }
+
+    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
+    RETURN_NAMES = ("positive", "negative", "latent")
+
+    CATEGORY = "ltxtricks"
+    FUNCTION = "generate"
+
+    def generate(
+        self, vae, positive, negative, latent, guiding_latent, latent_idx, strength
+    ):
+        noise_mask = nodes_lt.get_noise_mask(latent)
+        latent = latent["samples"]
+        guiding_latent = guiding_latent["samples"]
+        scale_factors = vae.downscale_index_formula
+
+        if latent_idx <= 0:
+            frame_idx = latent_idx * scale_factors[0]
+        else:
+            frame_idx = 1 + (latent_idx - 1) * scale_factors[0]
+
+        positive, negative, latent, noise_mask = self.append_keyframe(
+            positive=positive,
+            negative=negative,
+            frame_idx=frame_idx,
+            latent_image=latent,
+            noise_mask=noise_mask,
+            guiding_latent=guiding_latent,
+            strength=strength,
+            scale_factors=scale_factors,
+        )
+
+        return (
+            positive,
+            negative,
+            {"samples": latent, "noise_mask": noise_mask},
+        )
@@ -16,7 +16,8 @@ dependencies = [
 [project.optional-dependencies]
 internal = [
    "ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9.1",
-    "av>=10.0.0"
+    "av>=10.0.0",
+    "q8-kernels>=0.0.2"
 ]
 dev = [
    "pre-commit>=4.0.1",
@@ -0,0 +1,254 @@
+import torch
+from comfy_extras.nodes_custom_sampler import CFGGuider, SamplerCustomAdvanced
+from comfy_extras.nodes_lt import LTXVAddGuide, LTXVCropGuides
+
+from .latents import LTXVAddLatents, LTXVSelectLatents
+from .nodes_registry import comfy_node
+from .tricks import AddLatentGuideNode
+
+
+@comfy_node(description="Linear transition with overlap")
+class LinearOverlapLatentTransition:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "samples1": ("LATENT",),
+                "samples2": ("LATENT",),
+                "overlap": ("INT", {"default": 1, "min": 1, "max": 256}),
+            },
+            "optional": {
+                "axis": ("INT", {"default": 0}),
+            },
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "process"
+
+    CATEGORY = "Lightricks/latent"
+
+    def get_subbatch(self, samples):
+        s = samples.copy()
+        samples = s["samples"]
+        return samples
+
+    def process(self, samples1, samples2, overlap, axis=0):
+        samples1 = self.get_subbatch(samples1)
+        samples2 = self.get_subbatch(samples2)
+
+        # Create transition coefficients
+        alpha = torch.linspace(1, 0, overlap + 2)[1:-1].to(samples1.device)
+
+        # Create shape for broadcasting based on the axis
+        shape = [1] * samples1.dim()
+        shape[axis] = alpha.size(0)
+        alpha = alpha.reshape(shape)
+
+        # Create slices for the overlap regions
+        slice_all = [slice(None)] * samples1.dim()
+        slice_overlap1 = slice_all.copy()
+        slice_overlap1[axis] = slice(-overlap, None)
+        slice_overlap2 = slice_all.copy()
+        slice_overlap2[axis] = slice(0, overlap)
+        slice_rest1 = slice_all.copy()
+        slice_rest1[axis] = slice(None, -overlap)
+        slice_rest2 = slice_all.copy()
+        slice_rest2[axis] = slice(overlap, None)
+
+        # Combine samples
+        parts = [
+            samples1[tuple(slice_rest1)],
+            alpha * samples1[tuple(slice_overlap1)]
+            + (1 - alpha) * samples2[tuple(slice_overlap2)],
+            samples2[tuple(slice_rest2)],
+        ]
+
+        combined_samples = torch.cat(parts, dim=axis)
+        combined_batch_index = torch.arange(0, combined_samples.shape[0])
+
+        return (
+            {
+                "samples": combined_samples,
+                "batch_index": combined_batch_index,
+            },
+        )
+
+
+@comfy_node(
+    name="LTXVRecurrentKSampler",
+)
+class LTXVRecurrentKSampler:
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "vae": ("VAE",),
+                "noise": ("NOISE",),
+                "sampler": ("SAMPLER",),
+                "sigmas": ("SIGMAS",),
+                "latents": ("LATENT",),
+                "chunk_sizes": ("STRING", {"default": "3", "multiline": False}),
+                "overlaps": ("STRING", {"default": "1", "multiline": False}),
+                "positive": ("CONDITIONING",),
+                "negative": ("CONDITIONING",),
+                "input_image": ("IMAGE",),
+                "linear_blend_latents": ("BOOLEAN", {"default": True}),
+                "conditioning_strength": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01},
+                ),
+            },
+            "optional": {
+                "guider": ("GUIDER",),
+            },
+        }
+
+    RETURN_TYPES = (
+        "LATENT",
+        "LATENT",
+    )
+    RETURN_NAMES = (
+        "output",
+        "denoised_output",
+    )
+
+    FUNCTION = "sample"
+
+    CATEGORY = "sampling"
+
+    def sample(
+        self,
+        model,
+        vae,
+        noise,
+        sampler,
+        sigmas,
+        latents,
+        chunk_sizes,
+        overlaps,
+        positive,
+        negative,
+        input_image,
+        linear_blend_latents,
+        conditioning_strength,
+        guider=None,
+    ):
+        select_latents = LTXVSelectLatents().select_latents
+        add_latent_guide = AddLatentGuideNode().generate
+        add_latents = LTXVAddLatents().add_latents
+        positive_orig = positive.copy()
+        negative_orig = negative.copy()
+
+        # Parse chunk sizes and overlaps from strings
+        chunk_sizes = [int(x) for x in chunk_sizes.split(",")]
+        overlaps = [int(x) for x in overlaps.split(",")]
+
+        # Extend lists if shorter than number of sigma steps
+        n_steps = len(sigmas) - 1
+        if len(chunk_sizes) < n_steps:
+            chunk_sizes.extend([chunk_sizes[-1]] * (n_steps - len(chunk_sizes)))
+        if len(overlaps) < n_steps:
+            overlaps.extend([overlaps[-1]] * (n_steps - len(overlaps)))
+
+        # Initialize working latents
+        current_latents = latents.copy()
+        t_latents = None
+        # Loop through sigma pairs for progressive denoising
+        for i in range(n_steps):
+            current_sigmas = sigmas[i : i + 2]
+            current_chunk_size = chunk_sizes[i]
+            current_overlap = overlaps[i]
+
+            print(f"\nProcessing sigma step {i} with sigmas {current_sigmas}")
+            print(
+                f"Using chunk size {current_chunk_size} and overlap {current_overlap}"
+            )
+
+            # Calculate valid chunk starts to ensure the last chunk isn't shorter than the overlap
+            total_frames = current_latents["samples"].shape[2]
+            chunk_stride = current_chunk_size - current_overlap
+            valid_chunk_starts = list(
+                range(0, total_frames - current_overlap, chunk_stride)
+            )
+
+            # If the last chunk would be too short, remove the last start position
+            if (
+                total_frames > chunk_stride
+                and (total_frames - valid_chunk_starts[-1]) < current_chunk_size
+            ):
+                print(
+                    "last chunk is too short, it will only be of size",
+                    total_frames - valid_chunk_starts[-1],
+                    "frames",
+                )
+
+            # Process each chunk for current sigma pair
+            for i_chunk, chunk_start in enumerate(valid_chunk_starts):
+                (latents_chunk,) = select_latents(
+                    current_latents, chunk_start, chunk_start + current_chunk_size - 1
+                )
+                print(f"Processing chunk {i_chunk} starting at frame {chunk_start}")
+
+                if i_chunk == 0:
+                    positive, negative, latents_chunk = LTXVAddGuide().generate(
+                        positive_orig,
+                        negative_orig,
+                        vae,
+                        latents_chunk,
+                        input_image,
+                        0,
+                        0.75,
+                    )
+                else:
+                    (cond_latent,) = select_latents(t_latents, -current_overlap, -1)
+                    model, latents_chunk = add_latent_guide(
+                        model, latents_chunk, cond_latent, 0, conditioning_strength
+                    )
+
+                if guider is None:
+                    (guider_obj,) = CFGGuider().get_guider(
+                        model, positive, negative, 1.0
+                    )
+                else:
+                    guider_obj = guider
+                (_, denoised_latents) = SamplerCustomAdvanced().sample(
+                    noise, guider_obj, sampler, current_sigmas, latents_chunk
+                )
+                (positive, negative, denoised_latents) = LTXVCropGuides().crop(
+                    positive, negative, denoised_latents
+                )
+
+                if i_chunk == 0:
+                    t_latents = denoised_latents
+                else:
+                    if linear_blend_latents and current_overlap > 1:
+                        # the first output latent is the result of a 1:8 latent
+                        # reinterpreted as a 1:1 latent, so we ignore it
+                        (denoised_latents_drop_first,) = select_latents(
+                            denoised_latents, 1, -1
+                        )
+                        (t_latents,) = LinearOverlapLatentTransition().process(
+                            t_latents,
+                            denoised_latents_drop_first,
+                            current_overlap - 1,
+                            axis=2,
+                        )
+                    else:
+                        (truncated_denoised_latents,) = select_latents(
+                            denoised_latents, current_overlap, -1
+                        )
+                        (t_latents,) = add_latents(
+                            t_latents, truncated_denoised_latents
+                        )
+
+                print(
+                    f"Completed chunk {i_chunk}, current output shape: {t_latents['samples'].shape}"
+                )
+
+            # Update current_latents for next sigma step
+            current_latents = t_latents.copy()
+            print(f"Completed sigma step {i}")
+
+        return t_latents, t_latents
@@ -1,10 +1,10 @@
 import contextlib
+import math
 from dataclasses import dataclass
 from typing import List

 import comfy.ldm.modules.attention
 import comfy.samplers
-import comfy.utils
 import torch
 from comfy.model_patcher import ModelPatcher

@@ -20,8 +20,8 @@ def stg(
    rescale_scale,
 ):
    noise_pred = (
-        noise_pred_neg
-        + cfg_scale * (noise_pred_pos - noise_pred_neg)
+        noise_pred_pos
+        + (cfg_scale - 1) * (noise_pred_pos - noise_pred_neg)
        + stg_scale * (noise_pred_pos - noise_pred_pertubed)
    )
    if rescale_scale != 0:
@@ -183,6 +183,18 @@ class STGGuider(comfy.samplers.CFGGuider):
        positive_cond = self.conds.get("positive", None)
        negative_cond = self.conds.get("negative", None)

+        if model_options.get("sigma_to_params_mapping", None) is not None:
+            cfg_value, stg_scale, stg_layer_skip_layer_indices, stg_rescale = (
+                model_options["sigma_to_params_mapping"](timestep)
+            )
+            self.stg_flag.skip_layers = stg_layer_skip_layer_indices
+            self.patch_model(self.model_patcher, self.stg_flag)
+
+        else:
+            cfg_value = self.cfg
+            stg_scale = self.stg_scale
+            stg_rescale = self.rescale_scale
+
        noise_pred_pos = comfy.samplers.calc_cond_batch(
            self.inner_model,
            [positive_cond],
@@ -194,7 +206,7 @@ class STGGuider(comfy.samplers.CFGGuider):
        noise_pred_neg = 0
        noise_pred_perturbed = 0

-        if self.cfg > 1:
+        if not math.isclose(cfg_value, 1.0):
            noise_pred_neg = comfy.samplers.calc_cond_batch(
                self.inner_model,
                [negative_cond],
@@ -203,7 +215,7 @@ class STGGuider(comfy.samplers.CFGGuider):
                model_options,
            )[0]

-        if self.stg_scale > 0:
+        if not math.isclose(stg_scale, 0.0):
            try:
                model_options["transformer_options"]["ptb_index"] = 0
                self.stg_flag.do_skip = True
@@ -222,9 +234,161 @@ class STGGuider(comfy.samplers.CFGGuider):
            noise_pred_pos,
            noise_pred_neg,
            noise_pred_perturbed,
-            self.cfg,
-            self.stg_scale,
-            self.rescale_scale,
+            cfg_value,
+            stg_scale,
+            stg_rescale,
+        )
+
+        # normally this would be done in cfg_function, but we skipped
+        # that for efficiency: we can compute the noise predictions in
+        # a single call to calc_cond_batch() (rather than two)
+        # so we replicate the hook here
+        for fn in model_options.get("sampler_post_cfg_function", []):
+            args = {
+                "denoised": stg_result,
+                "cond": positive_cond,
+                "uncond": negative_cond,
+                "model": self.inner_model,
+                "uncond_denoised": noise_pred_neg,
+                "cond_denoised": noise_pred_pos,
+                "sigma": timestep,
+                "model_options": model_options,
+                "input": x,
+                # not in the original call in samplers.py:cfg_function, but made available for future hooks
+                "perturbed_cond": positive_cond,
+                "perturbed_cond_denoised": noise_pred_perturbed,
+            }
+            stg_result = fn(args)
+
+        return stg_result
+
+
+class STGGuiderAdvanced(comfy.samplers.CFGGuider):
+    def __init__(
+        self,
+        model: ModelPatcher,
+        sigma_list,
+        cfg_list,
+        stg_scale_list,
+        stg_rescale_list,
+        stg_layers_indices_list,
+        skip_steps_sigma_threshold,
+        cfg_star_rescale,
+    ):
+        model = model.clone()
+        super().__init__(model)
+
+        self.stg_flag = STGFlag(
+            do_skip=False,
+            skip_layers=model.model_options["transformer_options"].get(
+                "skip_block_list"
+            ),
+        )
+
+        self.sigma_list = sigma_list
+        self.cfg_list = cfg_list
+        self.stg_scale_list = stg_scale_list
+        self.stg_rescale_list = stg_rescale_list
+        self.stg_layers_indices_list = stg_layers_indices_list
+        self.skip_steps_sigma_threshold = skip_steps_sigma_threshold
+        self.cfg_star_rescale = cfg_star_rescale
+        STGGuider.patch_model(model, self.stg_flag)
+
+    def sigma_to_params_mapping(self, sigma):
+        # Find the closest higher sigma value and return corresponding cfg
+        higher_sigmas = [s for s in self.sigma_list if s >= sigma]
+        if not higher_sigmas:
+            return self.cfg_list[-1]  # Return last cfg if no higher sigma exists
+        closest_higher = min(higher_sigmas)
+        closest_idx = self.sigma_list.index(closest_higher)
+        return (
+            self.cfg_list[closest_idx],
+            self.stg_scale_list[closest_idx],
+            self.stg_rescale_list[closest_idx],
+            self.stg_layers_indices_list[closest_idx],
+        )
+
+    def set_conds(self, positive, negative):
+        self.inner_set_conds({"positive": positive, "negative": negative})
+
+    def predict_noise(
+        self,
+        x: torch.Tensor,
+        timestep: torch.Tensor,
+        model_options: dict = {},
+        seed=None,
+    ):
+        # CFG zero init - skipping steps with timestep bigger than given threshold.
+        if timestep > self.skip_steps_sigma_threshold:
+            return torch.zeros_like(x)
+
+        # in CFGGuider.predict_noise, we call sampling_function(), which uses cfg_function() to compute pos & neg
+        # but we'd rather do a single batch of sampling pos, neg, and perturbed, so we call calc_cond_batch([perturbed,pos,neg]) directly
+        positive_cond = self.conds.get("positive", None)
+        negative_cond = self.conds.get("negative", None)
+
+        cfg_value, stg_scale, stg_rescale, stg_layer_skip_layer_indices = (
+            self.sigma_to_params_mapping(timestep)
+        )
+
+        if stg_layer_skip_layer_indices is not None:
+            self.stg_flag.skip_layers = stg_layer_skip_layer_indices
+            STGGuider.patch_model(self.model_patcher, self.stg_flag)
+
+        noise_pred_pos = comfy.samplers.calc_cond_batch(
+            self.inner_model,
+            [positive_cond],
+            x,
+            timestep,
+            model_options,
+        )[0]
+
+        noise_pred_neg = 0
+        noise_pred_perturbed = 0
+
+        if not math.isclose(cfg_value, 1.0):
+            noise_pred_neg = comfy.samplers.calc_cond_batch(
+                self.inner_model,
+                [negative_cond],
+                x,
+                timestep,
+                model_options,
+            )[0]
+
+            if self.cfg_star_rescale:
+                batch_size = noise_pred_pos.shape[0]
+
+                positive_flat = noise_pred_pos.view(batch_size, -1)
+                negative_flat = noise_pred_neg.view(batch_size, -1)
+                dot_product = torch.sum(
+                    positive_flat * negative_flat, dim=1, keepdim=True
+                )
+                squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
+                alpha = dot_product / squared_norm
+                noise_pred_neg = alpha * noise_pred_neg
+
+        if not math.isclose(stg_scale, 0.0):
+            try:
+                model_options["transformer_options"]["ptb_index"] = 0
+                self.stg_flag.do_skip = True
+                noise_pred_perturbed = comfy.samplers.calc_cond_batch(
+                    self.inner_model,
+                    [positive_cond],
+                    x,
+                    timestep,
+                    model_options,
+                )[0]
+            finally:
+                self.stg_flag.do_skip = False
+                del model_options["transformer_options"]["ptb_index"]
+
+        stg_result = stg(
+            noise_pred_pos,
+            noise_pred_neg,
+            noise_pred_perturbed,
+            cfg_value,
+            stg_scale,
+            stg_rescale,
        )

        # normally this would be done in cfg_function, but we skipped
@@ -290,3 +454,130 @@ class STGGuiderNode:
        guider = STGGuider(model, cfg, stg, rescale)
        guider.set_conds(positive, negative)
        return (guider,)
+
+
+@comfy_node(name="STGGuiderAdvanced")
+class STGGuiderAdvancedNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "positive": ("CONDITIONING",),
+                "negative": ("CONDITIONING",),
+                "skip_steps_sigma_threshold": (
+                    "FLOAT",
+                    {
+                        "default": 0.996,
+                        "min": 0.0,
+                        "max": 100.0,
+                        "step": 0.001,
+                        "tooltip": "Steps with sigma greater than this values will be skipped.",
+                    },
+                ),
+                "cfg_star_rescale": (
+                    "BOOLEAN",
+                    {
+                        "default": True,
+                        "tooltip": "If true, applies the cfg star rescale, scales the negative prediction according to dot product between positive and negative.",
+                    },
+                ),
+                "sigmas": (
+                    "STRING",
+                    {
+                        "default": "1.0, 0.9933, 0.9850, 0.9767, 0.9008, 0.6180",
+                        "tooltip": "Comma-separated list sigmas, the actual parameters will be selected according to the closest sigma from this list to current timestep sigma.",
+                    },
+                ),
+                "cfg_values": (
+                    "STRING",
+                    {
+                        "default": "4, 4, 4, 4, 1, 1",
+                        "tooltip": "Comma-separated list of cfg values. Should be same length as sigmas list.",
+                    },
+                ),
+                "stg_scale_values": (
+                    "STRING",
+                    {
+                        "default": "2, 2, 2, 2, 1, 0",
+                        "tooltip": "Comma-separated list of stg scale values. Should be same length as sigmas list.",
+                    },
+                ),
+                "stg_rescale_values": (
+                    "STRING",
+                    {
+                        "default": "1, 1, 1, 1, 1, 1",
+                        "tooltip": "Comma-separated list of stg rescale values. Should be same length as sigmas list.",
+                    },
+                ),
+                "stg_layers_indices": (
+                    "STRING",
+                    {
+                        "default": "[14], [14], [14], [14], [14], [14]",
+                        "tooltip": "Comma-separated list of list of layer indices. Should be same length as sigmas list.",
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("GUIDER",)
+
+    FUNCTION = "get_guider"
+    CATEGORY = "lightricks/LTXV"
+
+    @classmethod
+    def parse_stg_layers_indices(cls, stg_layers_indices: str) -> List[List[int]]:
+        # First split by "], " to separate the lists, but preserve the closing bracket
+        lists = [s + "]" for s in stg_layers_indices.split("],")[:-1]]
+        # Add the last list which already has its closing bracket
+        if stg_layers_indices.strip():
+            lists.append(stg_layers_indices.split("],")[-1])
+
+        result = []
+        for s in lists:
+            s = s.strip()
+            if s == "[]":  # Empty list case
+                result.append([])
+            else:
+                # Remove brackets and whitespace
+                s = s.strip("[]").strip()
+                if not s:  # Handle case like "[]" after split
+                    result.append([])
+                else:
+                    # Split by comma and convert to integers
+                    numbers = [int(n.strip()) for n in s.split(",") if n.strip()]
+                    result.append(numbers)
+
+        return result
+
+    def get_guider(
+        self,
+        model,
+        positive,
+        negative,
+        skip_steps_sigma_threshold,
+        cfg_star_rescale,
+        sigmas,
+        cfg_values,
+        stg_scale_values,
+        stg_rescale_values,
+        stg_layers_indices,
+    ):
+        sigma_list = [float(s.strip()) for s in sigmas.split(",")]
+        cfg_list = [float(c.strip()) for c in cfg_values.split(",")]
+        stg_scale_list = [float(s.strip()) for s in stg_scale_values.split(",")]
+        stg_rescale_list = [float(s.strip()) for s in stg_rescale_values.split(",")]
+        stg_layers_indices_list = self.parse_stg_layers_indices(stg_layers_indices)
+
+        guider = STGGuiderAdvanced(
+            model,
+            sigma_list,
+            cfg_list,
+            stg_scale_list,
+            stg_rescale_list,
+            stg_layers_indices_list,
+            skip_steps_sigma_threshold,
+            cfg_star_rescale,
+        )
+        guider.set_conds(positive, negative)
+        return (guider,)
@@ -1302,7 +1302,7 @@
        "Node name for S&R": "CheckpointLoaderSimple"
      },
      "widgets_values": [
-        "ltx-video-2b-v0.9.5.safetensors"
+        "ltxv-2b-0.9.6-dev-04-25.safetensors"
      ]
    },
    {
@@ -2547,7 +2547,7 @@
      },
      "44": {
        "inputs": {
-          "ckpt_name": "ltx-video-2b-v0.9.5.safetensors"
+          "ckpt_name": "ltxv-2b-0.9.6-dev-04-25.safetensors"
        },
        "class_type": "CheckpointLoaderSimple",
        "_meta": {
@@ -1574,7 +1574,7 @@
        "Node name for S&R": "CheckpointLoaderSimple"
      },
      "widgets_values": [
-        "ltx-video-2b-v0.9.5.safetensors"
+        "ltxv-2b-0.9.6-dev-04-25.safetensors"
      ]
    },
    {
@@ -3362,7 +3362,7 @@
      },
      "44": {
        "inputs": {
-          "ckpt_name": "ltx-video-2b-v0.9.5.safetensors"
+          "ckpt_name": "ltxv-2b-0.9.6-dev-04-25.safetensors"
        },
        "class_type": "CheckpointLoaderSimple",
        "_meta": {
@@ -10,6 +10,7 @@ class AddLatentGuideNode(nodes_lt.LTXVAddGuide):
                "latent": ("LATENT",),
                "image_latent": ("LATENT",),
                "index": ("INT", {"default": 0, "min": -1, "max": 9999, "step": 1}),
+                "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}),
            }
        }

@@ -19,7 +20,7 @@ class AddLatentGuideNode(nodes_lt.LTXVAddGuide):
    CATEGORY = "ltxtricks"
    FUNCTION = "generate"

-    def generate(self, model, latent, image_latent, index):
+    def generate(self, model, latent, image_latent, index, strength):
        noise_mask = nodes_lt.get_noise_mask(latent)
        latent = latent["samples"]

@@ -30,7 +31,7 @@ class AddLatentGuideNode(nodes_lt.LTXVAddGuide):
            noise_mask,
            image_latent,
            index,
-            1.0,
+            strength,
        )

        return (