multi_scale_vision_transformer_cloak

Module containing multi-scale input-dependent Stained Glass Transforms for computer vision problems.

Two architectures are provided, both estimating noise parameters at several spatial scales and summing per-level predictions into a single estimate:

MultiScaleVisionTransformerCloak: an FPN top-down decoder over a pretrained encoder. The FPN builds and refines the multi-scale pyramid, which suits CNN backbones that do not natively expose strong multi-scale features.
MultiScaleSegformerCloak: a SegFormer-style all-MLP fusion over a pretrained hierarchical encoder (e.g. a Mix Vision Transformer mit_b* backbone). The encoder's native pyramid is projected to a common width and fused, avoiding the redundant top-down decoder.

Classes:

Name	Description
`MultiScaleFPNModule`	Wrap a pretrained FPN encoder and pyramid blocks with per-level prediction heads.
`MultiScaleSegformerCloak`	Facade for the multi-scale input-dependent Stained Glass Transform using a SegFormer-style architecture.
`MultiScaleSegformerModule`	Wrap a pretrained hierarchical encoder with SegFormer-style all-MLP fusion and per-level prediction heads.
`MultiScaleVisionTransformerCloak`	Facade for the multi-scale input-dependent Stained Glass Transform using an FPN architecture.
`MultiScaleVisionTransformerEstimator`	Estimator that uses a multi-scale FPN module to compute noise parameters.
`ResidualTemporalBlock`	Mix encoder features across the temporal axis with a small-initialized residual `Conv3d`.

MultiScaleFPNModule ¶

Bases: _BaseMultiScaleModule

Wrap a pretrained FPN encoder and pyramid blocks with per-level prediction heads.

The module extracts multi-scale pyramid features (p2-p5) from a pretrained FPN model and applies lightweight prediction heads at each pyramid level. The per-level predictions are upsampled to the input resolution and summed to produce a single 3-channel output.

The module accepts either a 4D image batch (B, 3, H, W) or a 5D video batch (B, 3, T, H, W); the output rank matches the input rank. For video, frames are flattened into the batch axis so the pretrained 2D encoder runs per frame, and (when enable_temporal is set) small-initialized ResidualTemporalBlocks mix the encoder pyramid features across time before decoding. Because the temporal blocks start as a near-identity perturbation, a freshly built video module produces approximately the per-frame image result until temporal gradients accrue.

Parameters:

Name	Type	Description	Default
`encoder_name` ¶	`str`	The name of the encoder backbone from `segmentation_models_pytorch`.	required
`encoder_weights` ¶	`str \| None`	Pretrained weights to use for the encoder. `None` means random init.	`'imagenet'`
`pyramid_channels` ¶	`int`	Number of channels in the FPN pyramid features.	`_DEFAULT_PYRAMID_CHANNELS`
`enable_temporal` ¶	`bool`	When `True`, insert small-initialized residual temporal convolutions between the encoder and decoder so 5D video inputs gain temporal mixing. When `False` (default), no temporal parameters are created and the module is identical to the image-only architecture; 5D inputs are still accepted but processed frame-by-frame.	`False`

Methods:

Name	Description
`forward`	Compute the multi-scale noise estimate from an image or video input.
`forward_per_level`	Compute per-level multi-scale predictions at the input spatial resolution.

forward ¶

forward(x: Tensor, **kwargs: Any) -> torch.Tensor

Compute the multi-scale noise estimate from an image or video input.

Override the base implementation to build the accumulator with torch.zeros_like so it matches the rank of the per-level predictions: the base assumes a 4D (B, 3, H, W) output, but video predictions are 5D (B, 3, T, H, W).

Parameters:

Name	Type	Description	Default
`x` ¶	`Tensor`	Input tensor of shape `(B, 3, H, W)` (image) or `(B, 3, T, H, W)` (video).	required
`**kwargs` ¶	`Any`	Additional keyword arguments forwarded to `forward_per_level`.	required

Returns:

Type	Description
`torch.Tensor`	A tensor at the input spatial resolution and rank, representing the summed
`torch.Tensor`	multi-scale predictions.

forward_per_level ¶

forward_per_level(
    x: Tensor, **kwargs: Any
) -> list[torch.Tensor]

Compute per-level multi-scale predictions at the input spatial resolution.

Parameters:

Name	Type	Description	Default
`x` ¶	`Tensor`	Input tensor of shape `(B, 3, H, W)` (image) or `(B, 3, T, H, W)` (video).	required
`**kwargs` ¶	`Any`	Additional keyword arguments. Accepted for compatibility with `BaseNoiseLayer.forward` forwarding and ignored by this module.	required

Returns:

Type	Description
`list[torch.Tensor]`	A list of tensors at the input spatial resolution, one per pyramid level in descending
`list[torch.Tensor]`	scale order. Each tensor matches the input rank: `(B, 3, H, W)` for images and
`list[torch.Tensor]`	`(B, 3, T, H, W)` for video.

Raises:

Type	Description
`ValueError`	If `x` is neither a 4D image nor a 5D video tensor.

MultiScaleSegformerCloak ¶

Bases: BaseNoiseLayer[MultiScaleSegformerModule, CloakStandardDeviationParameterization, BatchwisePercentMasker | None]

Facade for the multi-scale input-dependent Stained Glass Transform using a SegFormer-style architecture.

Use a pretrained hierarchical encoder (e.g. a Mix Vision Transformer mit_b* backbone) to estimate noise parameters at multiple spatial scales. The encoder's native multi-scale pyramid is projected to a common width and fused with a SegFormer all-MLP head rather than an FPN top-down decoder, with per-level prediction heads. The per-level predictions are summed to produce the final noise estimate.

Parameters:

Name	Type	Description	Default
`encoder_name` ¶	`str`	The name of the encoder backbone from `segmentation_models_pytorch`. Hierarchical transformer backbones such as `mit_b0`-`mit_b5` are the intended fit.	required
`encoder_weights` ¶	`str \| None`	Pretrained weights to use for the encoder. `None` means random init.	`'imagenet'`
`segmentation_channels` ¶	`int`	Number of channels in the projected and fused features.	`256`
`scale` ¶	`tuple[float, float]`	The min and max scale for the standard deviation parameterization.	`(0.0001, 2.0)`
`mean_bound` ¶	`float \| None`	Optional magnitude bound on the mean field. When provided, the mean estimator uses `BoundedMeanParameterization` so the mean stays in `[-mean_bound, +mean_bound]`. When `None` (default), the mean estimator runs without a parameterization and produces a raw (unbounded) mean field.	`None`
`shallow` ¶	`float`	A hyperparameter smoothing the tanh parameterizing the std estimator. Only has an effect when `directly_learn_stds` is `False`.	`1.0`
`directly_learn_stds` ¶	`bool`	If `True`, use `DirectStandardDeviationParameterization` (clamp via absolute value) instead of the tanh-based `CloakStandardDeviationParameterization`. When `True`, `shallow` must be left at its default of `1.0`.	`False`
`value_range` ¶	`tuple[float \| None, float \| None] \| None`	The range of features to limit the output of the SGT to.	`None`
`soft_clamp_beta` ¶	`float \| None`	When provided (and positive), clamp the output with a smooth `SoftClamp` elbow of this width instead of the hard `Clamp`. The choice is serialized so a deployed transform clamps identically. When `None` (default), a hard `Clamp` is used.	`None`
`seed` ¶	`int \| None`	A seed for the RNG.	`None`
`noise_layer_dtype` ¶	`dtype \| None`	The `torch.dtype` to use for the noise layer. If `None`, the default dtype is used.	`None`
`remove_skip_connection` ¶	`bool`	Whether to remove the skip connection in the noise layer. This is highly experimental and should be used with caution.	`False`

Added in version v3.46.0.

Methods:

Name	Description
`__call__`	Transform the input data.
`__getstate__`	Prepare a copy of the noise layer's state for serialization.
`__init__`
`__setstate__`	Restore the state of the noise layer from a serialized copy.
`forward`	Compute the transformed output from the input and noise mask.
`get_applied_transform_components_factory`	Create a function that returns the elements of the transform components (`'mean'` and `'std'`) applied during the most recent
`get_transformed_output_factory`	Create a function that returns the transformed output from the most recent forward pass.
`initial_seed`	Return the initial seed of the CPU device's random number generator.
`manual_seed`	Seed each of the random number generators.
`reset_parameters`	Reinitialize parameters and buffers.
`seed`	Seed each of the random number generators using a non-deterministic random number.

call ¶

__call__(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> torch.Tensor

Transform the input data.

Parameters:

Name	Type	Description	Default
`input` ¶	`Tensor`	The input to transform.	required
`noise_mask` ¶	`Tensor \| None`	An optional mask that selects the elements of `input` to transform. Where the mask is `False`, the original `input` value is returned. Also used to select the elements of the sampled standard deviations to use to mask the `input`. If `None`, the entire `input` is transformed.	`None`
`**kwargs` ¶	`Any`	Additional keyword arguments to the estimator modules.	required

getstate ¶

__getstate__() -> dict[str, Any]

Prepare a copy of the noise layer's state for serialization.

The returned dictionary can be JSON-serialized if the state_dict key is removed (it contains tensors). Generally you should pop the state_dict, save it as a safetensors file, then JSON-serialize the rest of the dictionary.

Returns:

Type	Description
`dict[str, Any]`	A dictionary containing the configuration of the noise layer, including its type string, the state dict, and the generator
`dict[str, Any]`	states if they exist.

Changed in version v3.15.0: Added serialization support for all noise layers.

init ¶

__init__(
    encoder_name: str,
    encoder_weights: str | None = "imagenet",
    segmentation_channels: int = 256,
    scale: tuple[float, float] = (0.0001, 2.0),
    mean_bound: float | None = None,
    shallow: float = 1.0,
    directly_learn_stds: bool = False,
    value_range: (
        tuple[float | None, float | None] | None
    ) = None,
    soft_clamp_beta: float | None = None,
    seed: int | None = None,
    noise_layer_dtype: dtype | None = None,
    remove_skip_connection: bool = False,
) -> None

Changed in version v3.51.0: Added the `soft_clamp_beta` constructor argument for opt-in `SoftClamp` on the cloak's output. The choice round-trips through serialization. `None` (default) keeps the hard `Clamp` and a `state_dict` identical to prior releases.

setstate ¶

__setstate__(
    state: dict[str, Any],
    trust_remote_code: bool = False,
    third_party_model_path: (
        str | PathLike[str] | None
    ) = None,
) -> None

Restore the state of the noise layer from a serialized copy.

state_dict and _generators are both optional keys, and will be restored if they exist in the state.

Parameters:

Name	Type	Description	Default
`state` ¶	`dict[str, Any]`	The state to set.	required
`trust_remote_code` ¶	`bool`	Whether to trust remote code when loading from HuggingFace Hub.	`False`
`third_party_model_path` ¶	`str \| PathLike[str] \| None`	The path or huggingface reference to a third-party model to load. This is useful when loading SGTs whose internal structure depends on transformers which are not importable directly through transformers, but are present on the Hugging Face Hub.	`None`

Changed in version v3.15.0: Added serialization support for all noise layers.

forward ¶

forward(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> torch.Tensor

Compute the transformed output from the input and noise mask.

Parameters:

Name	Type	Description	Default
`input` ¶	`Tensor`	The input to transform.	required
`noise_mask` ¶	`Tensor \| None`	An optional noise mask to apply to the output.	`None`
`**kwargs` ¶	`Any`	Additional keyword arguments to pass to the estimators.	required

Returns:

Type	Description
`torch.Tensor`	The Stained Glass transformed output.

get_applied_transform_components_factory ¶

get_applied_transform_components_factory() -> (
    Callable[[], dict[str, torch.Tensor]]
)

Create a function that returns the elements of the transform components ('mean' and 'std') applied during the most recent forward pass.

Specifically, the applied elements are those selected by the noise mask (if supplied) and standard deviation mask (if std_estimator.masker is not None). If no masks are used, all elements are returned.

The applied transform components are returned flattened.

This function is intended to be used to log histograms of the transform components.

Returns:

Type	Description
`Callable[[], dict[str, torch.Tensor]]`	A function that returns the the elements of the transform components applied during the most recent forward pass.

Examples:

>>> from torch import nn
>>> from stainedglass_core import model as sg_model, noise_layer as sg_noise_layer
>>> base_model = nn.Linear(20, 2)
>>> noisy_model = sg_model.NoisyModel(
...     sg_noise_layer.CloakNoiseLayer1,
...     base_model,
...     target_parameter="input",
... )
>>> get_applied_transform_components = (
...     noisy_model.noise_layer.get_applied_transform_components_factory()
... )
>>> input = torch.ones(1, 20)
>>> noise_mask = torch.tensor(5 * [False] + 15 * [True])
>>> output = noisy_model(input, noise_mask=noise_mask)
>>> applied_transform_components = get_applied_transform_components()
>>> applied_transform_components
{'mean': tensor(...), 'std': tensor(...)}
>>> {
...     component_name: component.shape
...     for component_name, component in applied_transform_components.items()
... }
{'mean': torch.Size([15]), 'std': torch.Size([15])}

get_transformed_output_factory ¶

get_transformed_output_factory() -> (
    Callable[[], torch.Tensor]
)

Create a function that returns the transformed output from the most recent forward pass.

If super batching is active, only the transformed half of the super batch output is returned.

Returns:

Type	Description
`Callable[[], torch.Tensor]`	A function that returns the transformed output from the most recent forward pass.

Examples:

>>> from stainedglass_core import noise_layer as sg_noise_layer
>>> noise_layer = sg_noise_layer.CloakNoiseLayer1()
>>> get_transformed_output = noise_layer.get_transformed_output_factory()
>>> input = torch.ones(2, 3, 32, 32)
>>> output = noise_layer(input)
>>> transformed_output = get_transformed_output()
>>> assert output.equal(transformed_output)

initial_seed ¶

initial_seed() -> int

Return the initial seed of the CPU device's random number generator.

manual_seed ¶

manual_seed(
    seed: int | None, rank_dependent: bool = True
) -> None

Seed each of the random number generators.

Setting seed to None will destroy any existing generators.

Parameters:

Name	Type	Description	Default
`seed` ¶	`int \| None`	The seed to set.	required
`rank_dependent` ¶	`bool`	Whether to add the distributed rank to the seed to ensure that each process samples different noise.	`True`

reset_parameters ¶

reset_parameters() -> None

Reinitialize parameters and buffers.

seed ¶

seed() -> None

Seed each of the random number generators using a non-deterministic random number.

MultiScaleSegformerModule ¶

Bases: _BaseMultiScaleModule

Wrap a pretrained hierarchical encoder with SegFormer-style all-MLP fusion and per-level prediction heads.

Hierarchical encoders (e.g. Mix Vision Transformer mit_b* backbones) natively emit a multi-scale feature pyramid, so no FPN top-down decoder is required to build one. Each of the deepest _NUM_PYRAMID_LEVELS encoder stages is projected to a common channel width with a 1x1 convolution, all projected stages are upsampled to the finest stage and fused into a single multi-scale context map (the SegFormer all-MLP head), and that context is injected back into each projected stage before a lightweight per-level prediction head. Per-level predictions are upsampled to the input resolution and summed to produce a single 3-channel output.

Parameters:

Name	Type	Description	Default
`encoder_name` ¶	`str`	The name of the encoder backbone from `segmentation_models_pytorch`.	required
`encoder_weights` ¶	`str \| None`	Pretrained weights to use for the encoder. `None` means random init.	`'imagenet'`
`segmentation_channels` ¶	`int`	Number of channels in the projected and fused features.	`_DEFAULT_PYRAMID_CHANNELS`

Methods:

Name	Description
`forward`	Compute the multi-scale noise estimate as the sum of per-level predictions.
`forward_per_level`	Compute per-level multi-scale predictions at the input spatial resolution.

forward ¶

forward(x: Tensor, **kwargs: Any) -> torch.Tensor

Compute the multi-scale noise estimate as the sum of per-level predictions.

Parameters:

Name	Type	Description	Default
`x` ¶	`Tensor`	Input tensor of shape `(B, 3, H, W)`.	required
`**kwargs` ¶	`Any`	Additional keyword arguments forwarded to `forward_per_level`.	required

Returns:

Type	Description
`torch.Tensor`	A 3-channel tensor at the input spatial resolution, representing the summed multi-scale predictions.

forward_per_level ¶

forward_per_level(
    x: Tensor, **kwargs: Any
) -> list[torch.Tensor]

Compute per-level multi-scale predictions at the input spatial resolution.

Parameters:

Name	Type	Description	Default
`x` ¶	`Tensor`	Input tensor of shape `(B, 3, H, W)`.	required
`**kwargs` ¶	`Any`	Additional keyword arguments. Accepted for compatibility with `BaseNoiseLayer.forward` forwarding and ignored by this module.	required

Returns:

Type	Description
`list[torch.Tensor]`	A list of 3-channel tensors at the input spatial resolution, one per encoder stage
`list[torch.Tensor]`	in ascending stride order (finest first).

MultiScaleVisionTransformerCloak ¶

Bases: BaseNoiseLayer[MultiScaleFPNModule, CloakStandardDeviationParameterization, BatchwisePercentMasker | None]

Facade for the multi-scale input-dependent Stained Glass Transform using an FPN architecture.

Use a pretrained FPN encoder and pyramid decoder to estimate noise parameters at multiple spatial scales (p2-p5), with prediction heads at each pyramid level. The per-level predictions are summed to produce the final noise estimate.

The mean and std estimators may use different backbones and pyramid widths. Each of encoder_name, encoder_weights, and pyramid_channels accepts either a single value (shared by both estimators) or a (mean, std) pair to configure the two estimators independently.

Parameters:

Name	Type	Description	Default
`encoder_name` ¶	`str \| Sequence[str]`	The name of the encoder backbone from `segmentation_models_pytorch`. Pass a single name to share it across both estimators, or a `(mean, std)` pair to use a different backbone for each.	required
`encoder_weights` ¶	`str \| None \| Sequence[str \| None]`	Pretrained weights to use for the encoder. `None` means random init. Pass a single value to share it, or a `(mean, std)` pair. Weights are backbone-specific, so a pair is typically required when `encoder_name` is a pair.	`'imagenet'`
`pyramid_channels` ¶	`int \| Sequence[int]`	Number of channels in the FPN pyramid features. Pass a single value to share it, or a `(mean, std)` pair.	`256`
`scale` ¶	`tuple[float, float]`	The min and max scale for the standard deviation parameterization.	`(0.0001, 2.0)`
`mean_bound` ¶	`float \| None`	Optional magnitude bound on the mean field. When provided, the mean estimator uses `BoundedMeanParameterization` so the mean stays in `[-mean_bound, +mean_bound]`. When `None` (default), the mean estimator runs without a parameterization and produces a raw (unbounded) mean field.	`None`
`shallow` ¶	`float`	A hyperparameter smoothing the tanh parameterizing the std estimator. Only has an effect when `directly_learn_stds` is `False`.	`1.0`
`directly_learn_stds` ¶	`bool`	If `True`, use `DirectStandardDeviationParameterization` (clamp via absolute value) instead of the tanh-based `CloakStandardDeviationParameterization`. When `True`, `shallow` must be left at its default of `1.0`.	`False`
`value_range` ¶	`tuple[float \| None, float \| None] \| None`	The range of features to limit the output of the SGT to.	`None`
`soft_clamp_beta` ¶	`float \| None`	When provided (and positive), clamp the output with a smooth `SoftClamp` elbow of this width instead of the hard `Clamp`. The elbow keeps a non-zero gradient just past `value_range`, which is useful during training; the choice is serialized so a deployed transform clamps identically. When `None` (default), a hard `Clamp` is used.	`None`
`enable_temporal` ¶	`bool`	When `True`, insert small-initialized residual temporal convolutions into the FPN so 5D video inputs `(B, 3, T, H, W)` gain temporal mixing. The transform still accepts 4D image inputs `(B, 3, H, W)`, and the output rank always matches the input rank. When `False` (default), no temporal parameters exist and the architecture is identical to the image-only cloak.	`False`
`seed` ¶	`int \| None`	A seed for the RNG.	`None`
`noise_layer_dtype` ¶	`dtype \| None`	The `torch.dtype` to use for the noise layer. If `None`, the default dtype is used.	`None`
`remove_skip_connection` ¶	`bool`	Whether to remove the skip connection in the noise layer. This is highly experimental and should be used with caution.	`False`

Added in version v3.39.0.

Methods:

Name	Description
`__call__`	Transform the input data.
`__getstate__`	Prepare a copy of the noise layer's state for serialization.
`__init__`
`__setstate__`	Restore the state of the noise layer from a serialized copy.
`forward`	Compute the transformed output from the input and noise mask.
`get_applied_transform_components_factory`	Create a function that returns the elements of the transform components (`'mean'` and `'std'`) applied during the most recent
`get_transformed_output_factory`	Create a function that returns the transformed output from the most recent forward pass.
`initial_seed`	Return the initial seed of the CPU device's random number generator.
`manual_seed`	Seed each of the random number generators.
`reset_parameters`	Reinitialize parameters and buffers.
`seed`	Seed each of the random number generators using a non-deterministic random number.

call ¶

__call__(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> torch.Tensor

Transform the input data.

Parameters:

Name	Type	Description	Default
`input` ¶	`Tensor`	The input to transform.	required
`noise_mask` ¶	`Tensor \| None`	An optional mask that selects the elements of `input` to transform. Where the mask is `False`, the original `input` value is returned. Also used to select the elements of the sampled standard deviations to use to mask the `input`. If `None`, the entire `input` is transformed.	`None`
`**kwargs` ¶	`Any`	Additional keyword arguments to the estimator modules.	required

getstate ¶

__getstate__() -> dict[str, Any]

Prepare a copy of the noise layer's state for serialization.

The returned dictionary can be JSON-serialized if the state_dict key is removed (it contains tensors). Generally you should pop the state_dict, save it as a safetensors file, then JSON-serialize the rest of the dictionary.

Returns:

Type	Description
`dict[str, Any]`	A dictionary containing the configuration of the noise layer, including its type string, the state dict, and the generator
`dict[str, Any]`	states if they exist.

Changed in version v3.15.0: Added serialization support for all noise layers.

init ¶

__init__(
    encoder_name: str | Sequence[str],
    encoder_weights: (
        str | None | Sequence[str | None]
    ) = "imagenet",
    pyramid_channels: int | Sequence[int] = 256,
    scale: tuple[float, float] = (0.0001, 2.0),
    mean_bound: float | None = None,
    shallow: float = 1.0,
    directly_learn_stds: bool = False,
    value_range: (
        tuple[float | None, float | None] | None
    ) = None,
    soft_clamp_beta: float | None = None,
    enable_temporal: bool = False,
    seed: int | None = None,
    noise_layer_dtype: dtype | None = None,
    remove_skip_connection: bool = False,
) -> None

Changed in version v3.42.0: Added the `mean_bound` constructor argument for opt-in `BoundedMeanParameterization` on the mean estimator.

Changed in version v3.42.0: The `BIAS_OFFSET = 4.0` shift previously subtracted unconditionally in `MultiScaleVisionTransformerEstimator.forward` is now carried by the std parameterization's `init_shift` instead. The CloakStd path's end-to-end `raw -> noise` mapping is unchanged. The `directly_learn_stds=True` path now initializes std at `min_scale` (was `max_scale`); the unparameterized mean path is no longer biased by `-4` at init.

Changed in version v3.45.0: Exposed and serialized the `noise_layer_dtype` and `remove_skip_connection` arguments.

Changed in version v3.47.0: `encoder_name`, `encoder_weights`, and `pyramid_channels` now each accept a `(mean, std)` pair so the mean and std estimators can use different backbones and pyramid widths. Passing a single value is unchanged and still shared by both.

Changed in version v3.48.0: Added 5D video support `(B, 3, T, H, W)` via the opt-in `enable_temporal` flag, which inserts small-initialized residual temporal convolutions (a Pseudo-3D FPN). The transform now dispatches on input rank; 4D image behavior is unchanged and `enable_temporal=False` keeps the `state_dict` identical to prior releases.

Changed in version v3.51.0: Added the `soft_clamp_beta` constructor argument for opt-in `SoftClamp` on the cloak's output. The choice round-trips through serialization, so a deployed transform clamps with the same smooth elbow it was built with. `None` (default) keeps the hard `Clamp` and a `state_dict` identical to prior releases.

setstate ¶

__setstate__(
    state: dict[str, Any],
    trust_remote_code: bool = False,
    third_party_model_path: (
        str | PathLike[str] | None
    ) = None,
) -> None

Restore the state of the noise layer from a serialized copy.

state_dict and _generators are both optional keys, and will be restored if they exist in the state.

Parameters:

Name	Type	Description	Default
`state` ¶	`dict[str, Any]`	The state to set.	required
`trust_remote_code` ¶	`bool`	Whether to trust remote code when loading from HuggingFace Hub.	`False`
`third_party_model_path` ¶	`str \| PathLike[str] \| None`	The path or huggingface reference to a third-party model to load. This is useful when loading SGTs whose internal structure depends on transformers which are not importable directly through transformers, but are present on the Hugging Face Hub.	`None`

Changed in version v3.15.0: Added serialization support for all noise layers.

forward ¶

forward(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> torch.Tensor

Compute the transformed output from the input and noise mask.

Parameters:

Name	Type	Description	Default
`input` ¶	`Tensor`	The input to transform.	required
`noise_mask` ¶	`Tensor \| None`	An optional noise mask to apply to the output.	`None`
`**kwargs` ¶	`Any`	Additional keyword arguments to pass to the estimators.	required

Returns:

Type	Description
`torch.Tensor`	The Stained Glass transformed output.

get_applied_transform_components_factory ¶

get_applied_transform_components_factory() -> (
    Callable[[], dict[str, torch.Tensor]]
)

Create a function that returns the elements of the transform components ('mean' and 'std') applied during the most recent forward pass.

Specifically, the applied elements are those selected by the noise mask (if supplied) and standard deviation mask (if std_estimator.masker is not None). If no masks are used, all elements are returned.

The applied transform components are returned flattened.

This function is intended to be used to log histograms of the transform components.

Returns:

Type	Description
`Callable[[], dict[str, torch.Tensor]]`	A function that returns the the elements of the transform components applied during the most recent forward pass.

Examples:

>>> from torch import nn
>>> from stainedglass_core import model as sg_model, noise_layer as sg_noise_layer
>>> base_model = nn.Linear(20, 2)
>>> noisy_model = sg_model.NoisyModel(
...     sg_noise_layer.CloakNoiseLayer1,
...     base_model,
...     target_parameter="input",
... )
>>> get_applied_transform_components = (
...     noisy_model.noise_layer.get_applied_transform_components_factory()
... )
>>> input = torch.ones(1, 20)
>>> noise_mask = torch.tensor(5 * [False] + 15 * [True])
>>> output = noisy_model(input, noise_mask=noise_mask)
>>> applied_transform_components = get_applied_transform_components()
>>> applied_transform_components
{'mean': tensor(...), 'std': tensor(...)}
>>> {
...     component_name: component.shape
...     for component_name, component in applied_transform_components.items()
... }
{'mean': torch.Size([15]), 'std': torch.Size([15])}

get_transformed_output_factory ¶

get_transformed_output_factory() -> (
    Callable[[], torch.Tensor]
)

Create a function that returns the transformed output from the most recent forward pass.

If super batching is active, only the transformed half of the super batch output is returned.

Returns:

Type	Description
`Callable[[], torch.Tensor]`	A function that returns the transformed output from the most recent forward pass.

Examples:

>>> from stainedglass_core import noise_layer as sg_noise_layer
>>> noise_layer = sg_noise_layer.CloakNoiseLayer1()
>>> get_transformed_output = noise_layer.get_transformed_output_factory()
>>> input = torch.ones(2, 3, 32, 32)
>>> output = noise_layer(input)
>>> transformed_output = get_transformed_output()
>>> assert output.equal(transformed_output)

initial_seed ¶

initial_seed() -> int

Return the initial seed of the CPU device's random number generator.

manual_seed ¶

manual_seed(
    seed: int | None, rank_dependent: bool = True
) -> None

Seed each of the random number generators.

Setting seed to None will destroy any existing generators.

Parameters:

Name	Type	Description	Default
`seed` ¶	`int \| None`	The seed to set.	required
`rank_dependent` ¶	`bool`	Whether to add the distributed rank to the seed to ensure that each process samples different noise.	`True`

reset_parameters ¶

reset_parameters() -> None

Reinitialize parameters and buffers.

seed ¶

seed() -> None

Seed each of the random number generators using a non-deterministic random number.

MultiScaleVisionTransformerEstimator ¶

Bases: Estimator[_BaseMultiScaleModule, StandardDeviationParameterization | BoundedMeanParameterization | None, Masker | None]

Estimator that uses a multi-scale FPN module to compute noise parameters.

Resize the input to a fixed internal resolution, run the MultiScaleFPNModule, resize the output back, and apply optional parameterization and masking.

Methods:

Name	Description
`__call__`	Estimate noise components from input values.
`forward`	Compute the forward pass of the multi-scale estimator.
`reset_parameters`	Reinitialize prediction head parameters so noise starts minimal.

call ¶

__call__(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> tuple[torch.Tensor, torch.Tensor | None]

__call__(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> tuple[torch.Tensor, torch.Tensor | None]

__call__(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> tuple[torch.Tensor, torch.Tensor | None]

__call__(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> tuple[torch.Tensor, torch.Tensor | None]

Estimate noise components from input values.

Parameters:

Name	Type	Description	Default
`input` ¶	`Tensor`	The tensor to estimate noise components from.	required
`noise_mask` ¶	`Tensor \| None`	An optional mask tensor to use to select a subset of the elements of the estimated standard deviations for computing a mask to apply to the input. If `None`, the mask is computed over all of the estimated standard deviations.	`None`
`kwargs` ¶	`Any`	Additional keyword arguments to `module`.	required

Returns:

Type	Description
`tuple[torch.Tensor, torch.Tensor \| None]`	A tuple containing the estimated noise components and an input mask if there is a `masker`.

forward ¶

forward(
    input: Tensor,
    noise_mask: Tensor | None = None,
    **kwargs: Any
) -> tuple[torch.Tensor, torch.Tensor | None]

Compute the forward pass of the multi-scale estimator.

Parameters:

Name	Type	Description	Default
`input` ¶	`Tensor`	The input to the estimator block.	required
`noise_mask` ¶	`Tensor \| None`	The optional noise mask. Kept around for legacy support and expected to already be at the original input resolution; not meaningfully used in computer vision problems.	`None`
`**kwargs` ¶	`Any`	Additional keyword arguments passed to the underlying module.	required

Returns:

Type	Description
`tuple[torch.Tensor, torch.Tensor \| None]`	The estimated output and an optional mask.

reset_parameters ¶

reset_parameters() -> None

Reinitialize prediction head parameters so noise starts minimal.

Reset each head's weights and biases via its default initialization, then scale by HALF_PERCENT. Reinitializing before scaling keeps this method idempotent: repeated calls do not progressively shrink the parameters. When the module carries temporal mixers, re-initialize them as well so the temporal path stays a near-identity at init.

ResidualTemporalBlock ¶

Bases: Module

Mix encoder features across the temporal axis with a small-initialized residual Conv3d.

The block applies a depth-preserving 1D temporal convolution (kernel (3, 1, 1)) over the time axis of a pyramid feature map and adds the result back to its input. The convolution uses a Kaiming initialization scaled down by _TEMPORAL_INIT_SCALE with a zero bias, so the block starts as a near-identity small perturbation. This keeps the pretrained 2D spatial features approximately intact while temporal gradients ramp up, while avoiding the degenerate symmetry of a pure zero initialization.

Parameters:

Name	Type	Description	Default
`channels` ¶	`int`	Number of feature channels at the pyramid level this block operates on.	required

Methods:

Name	Description
`forward`	Apply temporal mixing to a frame-flattened feature map.
`reset_parameters`	Re-initialize the temporal convolution as a small Kaiming-scaled, zero-bias perturbation.

forward ¶

forward(
    x: Tensor, batch_size: int, num_frames: int
) -> torch.Tensor

Apply temporal mixing to a frame-flattened feature map.

Parameters:

Name	Type	Description	Default
`x` ¶	`Tensor`	Feature map of shape `(B * T, C, H, W)`, where frames are flattened into the batch axis in batch-major order (frame `t` of sample `b` at index `b * T + t`).	required
`batch_size` ¶	`int`	The number of videos `B` in the batch.	required
`num_frames` ¶	`int`	The number of frames `T` per video.	required

Returns:

Type	Description
`torch.Tensor`	A feature map of shape `(B * T, C, H, W)` with temporal information mixed across frames.

reset_parameters ¶

reset_parameters() -> None

Re-initialize the temporal convolution as a small Kaiming-scaled, zero-bias perturbation.

Reset the convolution via its default Kaiming initialization, scale the weights by _TEMPORAL_INIT_SCALE, and zero the bias. Re-initializing before scaling keeps this method idempotent. The Kaiming init breaks the symmetry that a pure zero init would impose, while the small scale and zero bias keep the residual branch a near-identity at the start of training.

multi_scale_vision_transformer_cloak

MultiScaleFPNModule ¶

encoder_name ¶

encoder_weights ¶

pyramid_channels ¶

enable_temporal ¶

forward ¶

x ¶

**kwargs ¶

forward_per_level ¶

x ¶

**kwargs ¶

MultiScaleSegformerCloak ¶

encoder_name ¶

encoder_weights ¶

segmentation_channels ¶

scale ¶

mean_bound ¶

shallow ¶

directly_learn_stds ¶

value_range ¶

soft_clamp_beta ¶

seed ¶

noise_layer_dtype ¶

remove_skip_connection ¶

__call__ ¶

input ¶

noise_mask ¶

**kwargs ¶

__getstate__ ¶

__init__ ¶

__setstate__ ¶

state ¶

trust_remote_code ¶

third_party_model_path ¶

forward ¶

input ¶

noise_mask ¶

**kwargs ¶

get_applied_transform_components_factory ¶

get_transformed_output_factory ¶

initial_seed ¶

manual_seed ¶

seed ¶

rank_dependent ¶

reset_parameters ¶

seed ¶

MultiScaleSegformerModule ¶

encoder_name ¶

encoder_weights ¶

segmentation_channels ¶

forward ¶

x ¶

**kwargs ¶

forward_per_level ¶

x ¶

**kwargs ¶

MultiScaleVisionTransformerCloak ¶

encoder_name ¶

encoder_weights ¶

pyramid_channels ¶

scale ¶

mean_bound ¶

shallow ¶

directly_learn_stds ¶

value_range ¶

soft_clamp_beta ¶

enable_temporal ¶

seed ¶

noise_layer_dtype ¶

remove_skip_connection ¶

__call__ ¶

input ¶

noise_mask ¶

**kwargs ¶

__getstate__ ¶

__init__ ¶

__setstate__ ¶

state ¶

trust_remote_code ¶

`encoder_name` ¶

`encoder_weights` ¶

`pyramid_channels` ¶

`enable_temporal` ¶

`x` ¶

`**kwargs` ¶

`x` ¶

`**kwargs` ¶

`encoder_name` ¶

`encoder_weights` ¶

`segmentation_channels` ¶

`scale` ¶

`mean_bound` ¶

`shallow` ¶

`directly_learn_stds` ¶

`value_range` ¶

`soft_clamp_beta` ¶

`seed` ¶

`noise_layer_dtype` ¶

`remove_skip_connection` ¶

call ¶

`input` ¶

`noise_mask` ¶

`**kwargs` ¶

getstate ¶

init ¶

setstate ¶

`state` ¶

`trust_remote_code` ¶

`third_party_model_path` ¶

`input` ¶

`noise_mask` ¶

`**kwargs` ¶

`seed` ¶

`rank_dependent` ¶

`encoder_name` ¶

`encoder_weights` ¶

`segmentation_channels` ¶

`x` ¶

`**kwargs` ¶

`x` ¶

`**kwargs` ¶

`encoder_name` ¶

`encoder_weights` ¶

`pyramid_channels` ¶

`scale` ¶

`mean_bound` ¶

`shallow` ¶

`directly_learn_stds` ¶

`value_range` ¶

`soft_clamp_beta` ¶

`enable_temporal` ¶

`seed` ¶

`noise_layer_dtype` ¶

`remove_skip_connection` ¶

call ¶

`input` ¶

`noise_mask` ¶

`**kwargs` ¶

getstate ¶

init ¶

setstate ¶

`state` ¶

`trust_remote_code` ¶

`third_party_model_path` ¶

`input` ¶

`noise_mask` ¶

`**kwargs` ¶

`seed` ¶

`rank_dependent` ¶

call ¶

`input` ¶

`noise_mask` ¶

`kwargs` ¶

`input` ¶

`noise_mask` ¶

`**kwargs` ¶

`channels` ¶

`x` ¶

`batch_size` ¶