vision

Utilities for locating vision encoders and extracting clean/noisy vision features for Stained Glass Transform training.

Functions:

Name	Description
extract_clean_noisy_vision_features	Extract clean (teacher) and noisy (student) vision features in a single super-batched forward.
locate_vision_encoder	Locate the vision encoder on a multimodal model and return it with its dotted path from base_model.

extract_clean_noisy_vision_features

extract_clean_noisy_vision_features(
    clean_pixels: Tensor,
    noisy_pixels: Tensor,
    extract: Callable[[Tensor], Tensor],
    *,
    pixel_mean: Tensor,
    pixel_std: Tensor,
    dtype: dtype | None = None
) -> tuple[torch.Tensor, torch.Tensor]

Extract clean (teacher) and noisy (student) vision features in a single super-batched forward.

Concatenates [clean_pixels | noisy_pixels] along the batch axis, normalizes by pixel_mean / pixel_std, optionally casts to dtype, runs the result through the extract callback, then splits the returned features in half along the batch axis. The clean half is returned first so it can be used as the detached teacher for vision_feature_cosine_similarity.

The extract callback carries all model-specific and stateful logic — frozen-encoder snapshot swaps, patchification, grid metadata, pooling, and any reshape needed so its output is batch-major (2 * B, ...). Keeping it a callback lets this helper stay free of any model or wrapper coupling.

noisy_pixels may already be a super-batched (2 * B, ...) tensor (the no-fine-tune clean/noisy forward path stashes the concatenated output); in that case its second (B, ...) half — the noisy branch — is used. A plain (B, ...) noisy tensor is used as-is.

No torch.no_grad is applied: gradients flow through the student features back to noisy_pixels. Whether the encoder weights update is the caller's responsibility (typically they are frozen, or extract swaps in a frozen snapshot).

Privacy note: the student (noisy) features carry gradients into the Stained Glass Transform; the clean pixels are treated as the teacher target. Callers comparing the two (e.g. for a cosine constraint) should detach the teacher.

Parameters:

Name	Type	Description	Default
clean_pixels	Tensor	Clean image tensor of shape (B, C, H, W).	required
noisy_pixels	Tensor	Noisy image tensor of shape (B, C, H, W) or a super-batched (2 * B, C, H, W) whose second half is the noisy branch.	required
extract	Callable[[Tensor], Tensor]	Callback mapping the normalized super-batch (2 * B, C, H, W) to batch-major features (2 * B, ...).	required
pixel_mean	Tensor	Per-channel mean for normalization, broadcastable to the super-batch.	required
pixel_std	Tensor	Per-channel standard deviation for normalization, broadcastable to the super-batch.	required
dtype	dtype | None	Optional dtype to cast the normalized pixels to before extract (e.g. the encoder's parameter dtype).	None

Returns:

Type	Description
tuple[torch.Tensor, torch.Tensor]	Tuple (teacher_features, student_features), the clean and noisy halves of the extractor output, each of shape (B, ...).

Example

clean = torch.zeros(2, 3, 4, 4) noisy = torch.ones(2, 3, 4, 4) pixel_mean = torch.zeros(3, 1, 1) pixel_std = torch.ones(3, 1, 1) teacher, student = extract_clean_noisy_vision_features( ... clean, noisy, lambda x: x, pixel_mean=pixel_mean, pixel_std=pixel_std ... ) teacher.shape, student.shape (torch.Size([2, 3, 4, 4]), torch.Size([2, 3, 4, 4])) bool((teacher == 0).all()), bool((student == 1).all()) (True, True)

Added in version v3.49.0. Run a super-batched [clean | noisy] pixel pair through a model-specific feature extractor in a single forward and split the result back into a (teacher, student) feature pair.

locate_vision_encoder

locate_vision_encoder(
    base_model: Module,
    attr_paths: Sequence[str] = (
        "vision_model",
        "vision_tower",
        "model.visual",
        "visual",
    ),
) -> tuple[nn.Module, str]

Locate the vision encoder on a multimodal model and return it with its dotted path from base_model.

Walks attr_paths in priority order and returns the first attribute that resolves to an nn.Module. The returned dotted path is suitable for use with submodule_swap and FrozenSnapshot. Pass a model-specific attr_paths to control the priority order (e.g. Qwen3-VL resolves to model.visual; Nemotron-Omni to vision_model).

Parameters:

Name	Type	Description	Default
base_model	Module	The loaded multimodal model.	required
attr_paths	Sequence[str]	Dotted attribute paths to probe, in priority order. Defaults to the common Hugging Face conventions.	('vision_model', 'vision_tower', 'model.visual', 'visual')

Returns:

Type	Description
tuple[nn.Module, str]	Tuple (encoder, path) where base_model.get_submodule(path) is encoder.

Raises:

Type	Description
AttributeError	If no path in attr_paths resolves to an nn.Module.

Example

model = nn.Module() model.vision_model = nn.Linear(4, 4) encoder, path = locate_vision_encoder(model) path 'vision_model' encoder is model.vision_model True

Added in version v3.49.0. Locate a multimodal model's vision encoder by dotted attribute path, returning it with a path suitable for `submodule_swap` / `FrozenSnapshot`.