Upload 2 files

packages/ltx-core/src/ltx_core/guidance/__init__.py

@@ -0,0 +1,15 @@
+"""Guidance and perturbation utilities for attention manipulation."""
+
+from ltx_core.guidance.perturbations import (
+    BatchedPerturbationConfig,
+    Perturbation,
+    PerturbationConfig,
+    PerturbationType,
+)
+
+__all__ = [
+    "BatchedPerturbationConfig",
+    "Perturbation",
+    "PerturbationConfig",
+    "PerturbationType",
+]

packages/ltx-core/src/ltx_core/guidance/perturbations.py

@@ -1,6 +1,3 @@
-# Copyright (c) 2025 Lightricks. All rights reserved.
-# Created by Andrew Kvochko
-
 from dataclasses import dataclass
 from enum import Enum
 
@@ -9,6 +6,8 @@ from torch._prims_common import DeviceLikeType
 
 
 class PerturbationType(Enum):
+    """Types of attention perturbations for STG (Spatio-Temporal Guidance)."""
+
     SKIP_A2V_CROSS_ATTN = "skip_a2v_cross_attn"
     SKIP_V2A_CROSS_ATTN = "skip_v2a_cross_attn"
     SKIP_VIDEO_SELF_ATTN = "skip_video_self_attn"
@@ -17,6 +16,8 @@ class PerturbationType(Enum):
 
 @dataclass(frozen=True)
 class Perturbation:
+    """A single perturbation specifying which attention type to skip and in which blocks."""
+
     type: PerturbationType
     blocks: list[int] | None  # None means all blocks
 
@@ -32,6 +33,8 @@ class Perturbation:
 
 @dataclass(frozen=True)
 class PerturbationConfig:
+    """Configuration holding a list of perturbations for a single sample."""
+
     perturbations: list[Perturbation] | None
 
     def is_perturbed(self, perturbation_type: PerturbationType, block: int) -> bool:
@@ -47,6 +50,8 @@ class PerturbationConfig:
 
 @dataclass(frozen=True)
 class BatchedPerturbationConfig:
+    """Perturbation configurations for a batch, with utilities for generating attention masks."""
+
     perturbations: list[PerturbationConfig]
 
     def mask(