RepConTUR / ct_volume_encoder.py

Add RepConTUR-1B weights, remote code, tokenizer and model card

c6be546 verified 6 days ago

6.5 kB

"""
CT Volume Encoder Wrapper for RepConTUR.
Wraps CTViT or NNUNet3DEncoder to encode full 3D CT volumes into
a flattened token sequence of shape (B, N, dim).
The encoder is loaded from a CT-CLIP checkpoint and fully frozen.
"""

from pathlib import Path
from typing import Optional, List

import torch
import torch.nn as nn
from ._compat import print_log


class CTVolumeEncoder(nn.Module):
    """
    Wraps CTViT or NNUNet3DEncoder for 3D CT volume encoding.
    Input:  (B, 1, D, H, W) — full 3D CT volume
    Output: (B, T*H'*W', dim) — flattened spatial/temporal tokens
    The encoder is always frozen. Only downstream modules (attentional
    pooler, MLP projection) are trainable.
    """

    def __init__(
        self,
        encoder_type: str = 'ctvit',
        checkpoint_path: Optional[str] = None,
        ct_clip_source_path: Optional[str] = None,
        # CTViT params
        dim: int = 512,
        codebook_size: int = 8192,
        image_size: int = 360,
        patch_size: int = 20,
        temporal_patch_size: int = 10,
        spatial_depth: int = 4,
        temporal_depth: int = 4,
        dim_head: int = 32,
        heads: int = 8,
        # NNUNet params
        nnunet_depth: int = 5,
        nnunet_channels: Optional[List[int]] = None,
        nnunet_stride_list: Optional[List[List[int]]] = None,
    ):
        super().__init__()
        self.encoder_type = encoder_type
        self.dim = dim

        if nnunet_channels is None:
            nnunet_channels = [32, 64, 128, 256, 320, 512]
        if nnunet_stride_list is None:
            nnunet_stride_list = [[1, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]]

        # transformer_maskgit is vendored alongside this module (ct_clip_source_path ignored).
        if encoder_type == 'ctvit':
            from .ctvit import CTViT
            self.encoder = CTViT(
                dim=dim,
                codebook_size=codebook_size,
                image_size=image_size,
                patch_size=patch_size,
                temporal_patch_size=temporal_patch_size,
                spatial_depth=spatial_depth,
                temporal_depth=temporal_depth,
                dim_head=dim_head,
                heads=heads,
            )
        elif encoder_type == 'nnunet':
            from .nnunet import NNUNet3DEncoder
            self.encoder = NNUNet3DEncoder(
                depth=nnunet_depth,
                channels=nnunet_channels,
                stride_list=nnunet_stride_list,
            )
            self.dim = nnunet_channels[nnunet_depth]  # output channel dim
        else:
            raise ValueError(f"Unknown encoder_type: {encoder_type}. Use 'ctvit' or 'nnunet'.")

        # Load checkpoint
        if checkpoint_path is not None:
            self._load_checkpoint(checkpoint_path)

        # Freeze the entire encoder
        self.encoder.requires_grad_(False)
        self.encoder.eval()

        total_params = sum(p.numel() for p in self.encoder.parameters())
        print_log(
            f"CTVolumeEncoder ({encoder_type}): loaded with {total_params:,} params (all frozen)",
            logger='current'
        )

    def _load_checkpoint(self, checkpoint_path: str):
        """Load encoder weights from a CT-CLIP checkpoint."""
        path = Path(checkpoint_path)
        if not path.exists():
            raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")

        if path.suffix == '.safetensors':
            from safetensors import safe_open
            f = safe_open(str(path), framework='pt', device='cpu')
            all_keys = list(f.keys())

            encoder_state = {}
            for prefix in ['ctclip.visual_transformer.', 'visual_transformer.']:
                for k in all_keys:
                    if k.startswith(prefix):
                        encoder_state[k[len(prefix):]] = f.get_tensor(k)
                if encoder_state:
                    break

            if not encoder_state:
                for k in all_keys:
                    encoder_state[k] = f.get_tensor(k)
        else:
            # Standard PyTorch checkpoint
            ckpt = torch.load(str(path), map_location='cpu')
            if isinstance(ckpt, dict) and 'state_dict' in ckpt:
                ckpt = ckpt['state_dict']

            encoder_state = {}
            for prefix in ['ctclip.visual_transformer.', 'visual_transformer.']:
                for k, v in ckpt.items():
                    if k.startswith(prefix):
                        encoder_state[k[len(prefix):]] = v
                if encoder_state:
                    break

            if not encoder_state:
                encoder_state = ckpt

        # Load with strict=False to skip decoder keys (to_pixels*, dec_*)
        missing, unexpected = self.encoder.load_state_dict(encoder_state, strict=False)
        print_log(
            f"CTVolumeEncoder: loaded checkpoint from {checkpoint_path}\n"
            f"  Matched keys: {len(encoder_state) - len(unexpected)}\n"
            f"  Missing keys (expected for encoder-only): {len(missing)}\n"
            f"  Unexpected keys (decoder, skipped): {len(unexpected)}",
            logger='current'
        )

    def train(self, mode=True):
        """Override train to keep encoder always in eval mode."""
        super().train(mode)
        self.encoder.eval()
        return self

    def forward(self, volume: torch.Tensor) -> torch.Tensor:
        """
        Args:
            volume: (B, 1, D, H, W) — preprocessed 3D CT volume
        Returns:
            tokens: (B, N, dim) — flattened tokens where N = T*H'*W'
                    For CTViT with (180, 360, 360): N = 18*18*18 = 5832, dim = 512
                    For NNUNet with (180, 360, 360): N varies, dim = channels[-1]
        """

        # Query the encoder's actual parameter dtype at runtime.
        # .to(bf16) in __init__ may not persist (DeepSpeed can reset
        # frozen-parameter dtypes), so we always match dynamically.
        encoder_dtype = next(self.encoder.parameters()).dtype
        volume = volume.to(dtype=encoder_dtype)

        with torch.no_grad(), torch.cuda.amp.autocast(enabled=False):
            tokens = self.encoder(volume, return_encoded_tokens=True)
            # CTViT: (B, T, H', W', dim) e.g. (B, 18, 18, 18, 512)
            # NNUNet: (B, D', H', W', dim) e.g. (B, 11, 11, 11, 512)

        B = tokens.shape[0]
        tokens = tokens.reshape(B, -1, self.dim)  # (B, N, dim)
        return tokens

	"""
	CT Volume Encoder Wrapper for RepConTUR.

	Wraps CTViT or NNUNet3DEncoder to encode full 3D CT volumes into
	a flattened token sequence of shape (B, N, dim).

	The encoder is loaded from a CT-CLIP checkpoint and fully frozen.
	"""

	from pathlib import Path
	from typing import Optional, List

	import torch
	import torch.nn as nn
	from ._compat import print_log


	class CTVolumeEncoder(nn.Module):
	"""
	Wraps CTViT or NNUNet3DEncoder for 3D CT volume encoding.

	Input: (B, 1, D, H, W) — full 3D CT volume
	Output: (B, TH'W', dim) — flattened spatial/temporal tokens

	The encoder is always frozen. Only downstream modules (attentional
	pooler, MLP projection) are trainable.
	"""

	def __init__(
	self,
	encoder_type: str = 'ctvit',
	checkpoint_path: Optional[str] = None,
	ct_clip_source_path: Optional[str] = None,
	# CTViT params
	dim: int = 512,
	codebook_size: int = 8192,
	image_size: int = 360,
	patch_size: int = 20,
	temporal_patch_size: int = 10,
	spatial_depth: int = 4,
	temporal_depth: int = 4,
	dim_head: int = 32,
	heads: int = 8,
	# NNUNet params
	nnunet_depth: int = 5,
	nnunet_channels: Optional[List[int]] = None,
	nnunet_stride_list: Optional[List[List[int]]] = None,
	):
	super().__init__()
	self.encoder_type = encoder_type
	self.dim = dim

	if nnunet_channels is None:
	nnunet_channels = [32, 64, 128, 256, 320, 512]
	if nnunet_stride_list is None:
	nnunet_stride_list = [[1, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]]

	# transformer_maskgit is vendored alongside this module (ct_clip_source_path ignored).
	if encoder_type == 'ctvit':
	from .ctvit import CTViT
	self.encoder = CTViT(
	dim=dim,
	codebook_size=codebook_size,
	image_size=image_size,
	patch_size=patch_size,
	temporal_patch_size=temporal_patch_size,
	spatial_depth=spatial_depth,
	temporal_depth=temporal_depth,
	dim_head=dim_head,
	heads=heads,
	)
	elif encoder_type == 'nnunet':
	from .nnunet import NNUNet3DEncoder
	self.encoder = NNUNet3DEncoder(
	depth=nnunet_depth,
	channels=nnunet_channels,
	stride_list=nnunet_stride_list,
	)
	self.dim = nnunet_channels[nnunet_depth] # output channel dim
	else:
	raise ValueError(f"Unknown encoder_type: {encoder_type}. Use 'ctvit' or 'nnunet'.")

	# Load checkpoint
	if checkpoint_path is not None:
	self._load_checkpoint(checkpoint_path)

	# Freeze the entire encoder
	self.encoder.requires_grad_(False)
	self.encoder.eval()

	total_params = sum(p.numel() for p in self.encoder.parameters())
	print_log(
	f"CTVolumeEncoder ({encoder_type}): loaded with {total_params:,} params (all frozen)",
	logger='current'
	)

	def _load_checkpoint(self, checkpoint_path: str):
	"""Load encoder weights from a CT-CLIP checkpoint."""
	path = Path(checkpoint_path)
	if not path.exists():
	raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")

	if path.suffix == '.safetensors':
	from safetensors import safe_open
	f = safe_open(str(path), framework='pt', device='cpu')
	all_keys = list(f.keys())

	encoder_state = {}
	for prefix in ['ctclip.visual_transformer.', 'visual_transformer.']:
	for k in all_keys:
	if k.startswith(prefix):
	encoder_state[k[len(prefix):]] = f.get_tensor(k)
	if encoder_state:
	break

	if not encoder_state:
	for k in all_keys:
	encoder_state[k] = f.get_tensor(k)
	else:
	# Standard PyTorch checkpoint
	ckpt = torch.load(str(path), map_location='cpu')
	if isinstance(ckpt, dict) and 'state_dict' in ckpt:
	ckpt = ckpt['state_dict']

	encoder_state = {}
	for prefix in ['ctclip.visual_transformer.', 'visual_transformer.']:
	for k, v in ckpt.items():
	if k.startswith(prefix):
	encoder_state[k[len(prefix):]] = v
	if encoder_state:
	break

	if not encoder_state:
	encoder_state = ckpt

	# Load with strict=False to skip decoder keys (to_pixels, dec_)
	missing, unexpected = self.encoder.load_state_dict(encoder_state, strict=False)
	print_log(
	f"CTVolumeEncoder: loaded checkpoint from {checkpoint_path}\n"
	f" Matched keys: {len(encoder_state) - len(unexpected)}\n"
	f" Missing keys (expected for encoder-only): {len(missing)}\n"
	f" Unexpected keys (decoder, skipped): {len(unexpected)}",
	logger='current'
	)

	def train(self, mode=True):
	"""Override train to keep encoder always in eval mode."""
	super().train(mode)
	self.encoder.eval()
	return self

	def forward(self, volume: torch.Tensor) -> torch.Tensor:
	"""
	Args:
	volume: (B, 1, D, H, W) — preprocessed 3D CT volume

	Returns:
	tokens: (B, N, dim) — flattened tokens where N = TH'W'
	For CTViT with (180, 360, 360): N = 181818 = 5832, dim = 512
	For NNUNet with (180, 360, 360): N varies, dim = channels[-1]
	"""

	# Query the encoder's actual parameter dtype at runtime.
	# .to(bf16) in __init__ may not persist (DeepSpeed can reset
	# frozen-parameter dtypes), so we always match dynamically.
	encoder_dtype = next(self.encoder.parameters()).dtype
	volume = volume.to(dtype=encoder_dtype)

	with torch.no_grad(), torch.cuda.amp.autocast(enabled=False):
	tokens = self.encoder(volume, return_encoded_tokens=True)
	# CTViT: (B, T, H', W', dim) e.g. (B, 18, 18, 18, 512)
	# NNUNet: (B, D', H', W', dim) e.g. (B, 11, 11, 11, 512)

	B = tokens.shape[0]
	tokens = tokens.reshape(B, -1, self.dim) # (B, N, dim)
	return tokens