Model Card for RADM¶

class RADMConfig(ConfigMixin):
    """Configuration saved with RADM pipeline artifacts.

    Args:
        dataset_name: Poster dataset variant.
        id2label: Public label mapping. Defaults to the shared CGL labels.
        original_id2label: Label spellings from the checked RADM source.
        num_proposals: Number of proposal boxes sampled during inference.
        num_classes: Number of semantic classes predicted by the denoiser.
        hidden_dim: Proposal and condition hidden dimension.
        text_feature_dim: Last dimension of RADM text-feature tensors.
        max_text_num: Maximum text-feature rows per example.
        num_heads: Number of repeated proposal refinement blocks.
        num_attention_heads: Proposal self-attention head count.
        dim_feedforward: Proposal feed-forward hidden dimension.
        num_dynamic: Number of dynamic-convolution projections.
        dim_dynamic: Intermediate dynamic-convolution dimension.
        num_cls: Number of repeated classification projection layers.
        num_reg: Number of repeated box-regression projection layers.
        roi_resolution: ROIAlign output resolution.
        roi_sampling_ratio: ROIAlign sampling ratio.
        with_vtram: Whether visual-textual relation features are enabled.
        with_gram: Whether geometric relation features are enabled.
        deep_supervision: Whether intermediate proposal heads are returned.
        backbone_depth: ResNet backbone depth.
        backbone_freeze_at: Number of initial ResNet stages to freeze.
        image_size: Short-side image preprocessing target.
        size_divisibility: Padding multiple for image tensors.
        pixel_mean: RGB mean in 0-255 space.
        pixel_std: RGB standard deviation in 0-255 space.
        inference_steps: Default number of reverse-diffusion steps.
        num_train_timesteps: Number of training diffusion timesteps.
        snr_scale: RADM signal-to-noise scaling metadata.
        sample_step: Original sampling stride metadata.
        class_threshold: Default confidence threshold.
        nms_threshold: Default class-wise NMS threshold.
        denoiser_subfolder: Pipeline denoiser component subfolder.
        scheduler_subfolder: Pipeline scheduler component subfolder.
        processor_subfolder: Pipeline processor component subfolder.
        conversion_report: Conversion metadata persisted with local artifacts.
        checkpoint_status: Human-readable checkpoint status.
        license_status: Human-readable source/weight license status.

    Examples:
        >>> config = RADMConfig(num_proposals=2, hidden_dim=8)
        >>> config.num_labels
        5
    """

    config_name: str = "radm_config.json"

    @register_to_config
    def __init__(
        self,
        *,
        dataset_name: DatasetName | str = DatasetName.cgl,
        id2label: Mapping[int | str, str] | None = None,
        original_id2label: Mapping[int | str, str] | None = None,
        num_proposals: int = 100,
        num_classes: int = 5,
        hidden_dim: int = 256,
        text_feature_dim: int = 768,
        max_text_num: int = 20,
        num_heads: int = 6,
        num_attention_heads: int = 8,
        dim_feedforward: int = 2048,
        num_dynamic: int = 2,
        dim_dynamic: int = 64,
        num_cls: int = 1,
        num_reg: int = 3,
        roi_resolution: int = 7,
        roi_sampling_ratio: int = 2,
        with_vtram: bool = True,
        with_gram: bool = True,
        deep_supervision: bool = True,
        backbone_depth: int = 50,
        backbone_freeze_at: int = 2,
        image_size: int = 800,
        size_divisibility: int = 32,
        pixel_mean: Sequence[float] = (123.675, 116.280, 103.530),
        pixel_std: Sequence[float] = (58.395, 57.120, 57.375),
        inference_steps: int = 50,
        num_train_timesteps: int = 1000,
        snr_scale: float = 2.0,
        sample_step: int = 1,
        class_threshold: float = 0.30,
        nms_threshold: float = 0.50,
        denoiser_subfolder: str = "denoiser",
        scheduler_subfolder: str = "scheduler",
        processor_subfolder: str = "processor",
        conversion_report: Mapping[str, YamlValue] | None = None,
        checkpoint_status: str = "RADM README confirms datasets only; no released checkpoint is published",
        license_status: str = "checked original source has no license file",
    ) -> None:
        """Initialize RADM settings."""
        dataset = normalize_dataset_name(dataset_name)
        if dataset not in {DatasetName.cgl, DatasetName.cgl_v2}:
            raise ValueError(f"Unsupported RADM dataset_name: {dataset_name}")

        raw_id2label = id2label or default_id2label(dataset)
        raw_original = original_id2label or default_id2label(
            dataset, label_mode=RADMLabelMode.original
        )
        self.dataset_name = str(dataset)
        self.id2label = {int(key): value for key, value in raw_id2label.items()}
        self.original_id2label = {
            int(key): value for key, value in raw_original.items()
        }
        self.num_proposals = int(num_proposals)
        self.num_classes = int(num_classes)
        self.hidden_dim = int(hidden_dim)
        self.text_feature_dim = int(text_feature_dim)
        self.max_text_num = int(max_text_num)
        self.num_heads = int(num_heads)
        self.num_attention_heads = int(num_attention_heads)
        self.dim_feedforward = int(dim_feedforward)
        self.num_dynamic = int(num_dynamic)
        self.dim_dynamic = int(dim_dynamic)
        self.num_cls = int(num_cls)
        self.num_reg = int(num_reg)
        self.roi_resolution = int(roi_resolution)
        self.roi_sampling_ratio = int(roi_sampling_ratio)
        self.with_vtram = bool(with_vtram)

        self.with_gram = bool(with_gram)
        self.deep_supervision = bool(deep_supervision)
        self.backbone_depth = int(backbone_depth)
        self.backbone_freeze_at = int(backbone_freeze_at)
        self.image_size = int(image_size)
        self.size_divisibility = int(size_divisibility)
        self.pixel_mean = tuple(float(value) for value in pixel_mean)
        self.pixel_std = tuple(float(value) for value in pixel_std)
        self.inference_steps = int(inference_steps)
        self.num_train_timesteps = int(num_train_timesteps)
        self.snr_scale = float(snr_scale)
        self.sample_step = int(sample_step)
        self.class_threshold = float(class_threshold)
        self.nms_threshold = float(nms_threshold)
        self.denoiser_subfolder = denoiser_subfolder

        self.scheduler_subfolder = scheduler_subfolder
        self.processor_subfolder = processor_subfolder
        self.conversion_report = dict(conversion_report or {})
        self.checkpoint_status = checkpoint_status
        self.license_status = license_status

    @property
    def label2id(self) -> dict[str, int]:
        """Return the public label-name to integer-id mapping."""
        return {value: key for key, value in self.id2label.items()}

    @property
    def num_labels(self) -> int:
        """Return the configured public class count."""
        return len(self.id2label)

class RADMLabelMode(StrEnum):
    """Closed set of public RADM label vocabularies."""

    english = auto()
    original = auto()

class RADMImageProcessor(BaseImageProcessor):
    """Prepare RGB images for RADM inference.

    Args:
        image_size: Short-side resize target.
        size_divisibility: Padding multiple used by FPN-style backbones.
        pixel_mean: RGB mean in 0-255 space.
        pixel_std: RGB standard deviation in 0-255 space.

    Examples:
        >>> processor = RADMImageProcessor(image_size=32)
        >>> batch = processor.preprocess(Image.new("RGB", (16, 24)))
        >>> tuple(batch["pixel_values"].shape[:2])
        (1, 3)
    """

    model_input_names = ["pixel_values"]

    def __init__(
        self,
        image_size: int = 800,
        size_divisibility: int = 32,
        pixel_mean: Sequence[float] = (123.675, 116.280, 103.530),
        pixel_std: Sequence[float] = (58.395, 57.120, 57.375),
        **kwargs: Unpack[ImagesKwargs],
    ) -> None:
        """Initialize image normalization settings."""
        super().__init__(**kwargs)
        self.image_size = int(image_size)
        self.size_divisibility = int(size_divisibility)
        self.pixel_mean = tuple(float(value) for value in pixel_mean)
        self.pixel_std = tuple(float(value) for value in pixel_std)

    @classmethod
    def from_config(cls, config: RADMConfig) -> "RADMImageProcessor":
        """Build an image processor from RADM configuration.

        Args:
            config: RADM pipeline configuration.

        Returns:
            Image processor with matching resize and normalization metadata.
        """
        return cls(
            image_size=config.image_size,
            size_divisibility=config.size_divisibility,
            pixel_mean=config.pixel_mean,
            pixel_std=config.pixel_std,
        )

    def preprocess(
        self,
        images: ImageInput | Sequence[ImageInput],
        **kwargs: Unpack[ImagesKwargs],
    ) -> BatchFeature:
        """Preprocess one image or an image batch.

        Args:
            images: PIL, numpy, or torch image input.
            return_tensors: Tensor framework. Only ``pt`` is supported.
            kwargs: Ignored compatibility kwargs.

        Returns:
            Batch feature with padded ``pixel_values`` and original sizes.

        Raises:
            ValueError: If ``return_tensors`` is not ``pt``.
        """
        return_tensors = kwargs.pop("return_tensors", "pt")
        if return_tensors != "pt":
            raise ValueError("RADMImageProcessor only supports return_tensors='pt'")

        tensors: list[Float[torch.Tensor, "channels height width"]] = []
        original_sizes: list[tuple[int, int]] = []
        resized_sizes: list[tuple[int, int]] = []
        for image in _ensure_pil_batch(images):
            width, height = image.size
            original_sizes.append((height, width))
            resized = _resize_short_side(image.convert("RGB"), self.image_size)
            resized_width, resized_height = resized.size
            resized_sizes.append((resized_height, resized_width))
            array = np.asarray(resized, dtype=np.float32)
            mean = np.asarray(self.pixel_mean, dtype=np.float32)
            std = np.asarray(self.pixel_std, dtype=np.float32)

            tensors.append(torch.from_numpy(((array - mean) / std).transpose(2, 0, 1)))

        padded, padding_mask = _pad_tensors(tensors, self.size_divisibility)
        return BatchFeature(
            {
                "pixel_values": padded.float(),
                "pixel_mask": padding_mask,
                "original_sizes": torch.tensor(original_sizes, dtype=torch.long),
                "resized_sizes": torch.tensor(resized_sizes, dtype=torch.long),
            }
        )

class RADMDenoiser(ModelMixin, ConfigMixin):
    """Genuine RADM proposal denoiser shared by inference and training."""

    config_name: str = "denoiser_config.json"

    def __init__(
        self,
        *,
        config: RADMConfig | None = None,
        num_classes: int | None = None,
        num_proposals: int | None = None,
        hidden_dim: int | None = None,
        text_feature_dim: int | None = None,
        max_text_num: int | None = None,
        num_heads: int | None = None,
        num_attention_heads: int | None = None,
        dim_feedforward: int | None = None,
        num_dynamic: int | None = None,
        dim_dynamic: int | None = None,
        num_cls: int | None = None,
        num_reg: int | None = None,
        roi_resolution: int | None = None,
        roi_sampling_ratio: int | None = None,
        with_vtram: bool | None = None,
        with_gram: bool | None = None,
        deep_supervision: bool | None = None,
        backbone_depth: int | None = None,
        backbone_freeze_at: int | None = None,
        num_train_timesteps: int | None = None,
        snr_scale: float | None = None,
    ) -> None:
        """Initialize from an explicit package config or serialized model values."""
        super().__init__()
        if config is None:
            values = {
                "num_classes": num_classes,
                "num_proposals": num_proposals,
                "hidden_dim": hidden_dim,
                "text_feature_dim": text_feature_dim,
                "max_text_num": max_text_num,
                "num_heads": num_heads,
                "num_attention_heads": num_attention_heads,
                "dim_feedforward": dim_feedforward,
                "num_dynamic": num_dynamic,
                "dim_dynamic": dim_dynamic,
                "num_cls": num_cls,
                "num_reg": num_reg,
                "roi_resolution": roi_resolution,
                "roi_sampling_ratio": roi_sampling_ratio,
                "with_vtram": with_vtram,
                "with_gram": with_gram,
                "deep_supervision": deep_supervision,
                "backbone_depth": backbone_depth,
                "backbone_freeze_at": backbone_freeze_at,
                "num_train_timesteps": num_train_timesteps,
                "snr_scale": snr_scale,
            }
            if any(value is None for value in values.values()):
                raise TypeError("RADMDenoiser requires an explicit RADMConfig")

            config = RADMConfig(**values)  # type: ignore[arg-type]

        self.register_to_config(
            **{
                "dataset_name": config.dataset_name,
                "id2label": config.id2label,
                "original_id2label": config.original_id2label,
                "num_classes": config.num_classes,
                "num_proposals": config.num_proposals,
                "hidden_dim": config.hidden_dim,
                "text_feature_dim": config.text_feature_dim,
                "max_text_num": config.max_text_num,
                "num_heads": config.num_heads,
                "num_attention_heads": config.num_attention_heads,
                "dim_feedforward": config.dim_feedforward,
                "num_dynamic": config.num_dynamic,
                "dim_dynamic": config.dim_dynamic,
                "num_cls": config.num_cls,
                "num_reg": config.num_reg,
                "roi_resolution": config.roi_resolution,
                "roi_sampling_ratio": config.roi_sampling_ratio,
                "with_vtram": config.with_vtram,
                "with_gram": config.with_gram,
                "deep_supervision": config.deep_supervision,
                "backbone_depth": config.backbone_depth,
                "backbone_freeze_at": config.backbone_freeze_at,
                "num_train_timesteps": config.num_train_timesteps,
                "snr_scale": config.snr_scale,
                "sample_step": config.sample_step,
            }
        )
        if config.backbone_depth not in (18, 50):
            raise ValueError("backbone_depth must be 18 or 50")

        self.num_classes = config.num_classes
        self.num_proposals = config.num_proposals
        self.hidden_dim = config.hidden_dim
        self.backbone = RADMBackbone(
            depth=cast(Literal[18, 50], config.backbone_depth),
            freeze_at=config.backbone_freeze_at,
        )
        self.head = RADMProposalHead(config)
        self._register_diffusion_buffers(config)

    def _register_diffusion_buffers(self, config: RADMConfig) -> None:
        betas = cosine_beta_schedule(config.num_train_timesteps)
        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)
        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
        self.register_buffer("betas", betas)
        self.register_buffer("alphas_cumprod", alphas_cumprod)
        self.register_buffer("alphas_cumprod_prev", alphas_cumprod_prev)
        self.register_buffer("sqrt_alphas_cumprod", alphas_cumprod.sqrt())
        self.register_buffer(
            "sqrt_one_minus_alphas_cumprod", (1.0 - alphas_cumprod).sqrt()
        )
        self.register_buffer(
            "log_one_minus_alphas_cumprod", torch.log(1.0 - alphas_cumprod)
        )
        self.register_buffer("sqrt_recip_alphas_cumprod", (1.0 / alphas_cumprod).sqrt())
        self.register_buffer(
            "sqrt_recipm1_alphas_cumprod", (1.0 / alphas_cumprod - 1.0).sqrt()
        )
        posterior_variance = (
            betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
        )
        self.register_buffer("posterior_variance", posterior_variance)
        self.register_buffer(
            "posterior_log_variance_clipped",
            torch.log(posterior_variance.clamp(min=1e-20)),
        )
        self.register_buffer(
            "posterior_mean_coef1",
            betas * alphas_cumprod_prev.sqrt() / (1.0 - alphas_cumprod),
        )
        self.register_buffer(
            "posterior_mean_coef2",
            (1.0 - alphas_cumprod_prev) * alphas.sqrt() / (1.0 - alphas_cumprod),
        )

    def q_sample(
        self,
        x_start: Float[torch.Tensor, "batch proposals 4"],
        timesteps: Int[torch.Tensor, "batch"],
        noise: Float[torch.Tensor, "batch proposals 4"] | None = None,
    ) -> Float[torch.Tensor, "batch proposals 4"]:
        """Sample the configured cosine forward diffusion process."""
        sampled_noise = torch.randn_like(x_start) if noise is None else noise
        return (
            _extract(self.sqrt_alphas_cumprod, timesteps, x_start.shape) * x_start
            + _extract(self.sqrt_one_minus_alphas_cumprod, timesteps, x_start.shape)
            * sampled_noise
        )

    def predict_noise_from_start(
        self,
        x_t: Float[torch.Tensor, "batch proposals 4"],
        timesteps: Int[torch.Tensor, "batch"],
        x_start: Float[torch.Tensor, "batch proposals 4"],
    ) -> Float[torch.Tensor, "batch proposals 4"]:
        """Recover epsilon from a predicted clean diffusion sample."""
        return (
            _extract(self.sqrt_recip_alphas_cumprod, timesteps, x_t.shape) * x_t
            - x_start
        ) / _extract(self.sqrt_recipm1_alphas_cumprod, timesteps, x_t.shape)

    def prepare_diffusion_concat(
        self,
        boxes_cxcywh: Float[torch.Tensor, "elements 4"],
        *,
        generator: torch.Generator | None = None,
    ) -> tuple[
        Float[torch.Tensor, "proposals 4"],
        Float[torch.Tensor, "proposals 4"],
        Int[torch.Tensor, "1"],
    ]:
        """Prepare one image's padded proposals using the configured sampler."""
        device = boxes_cxcywh.device
        timestep = randint(
            0,
            self.config.num_train_timesteps,
            (1,),
            generator=generator,
            device=device,
        ).long()
        noise = randn(
            self.num_proposals,
            4,
            generator=generator,
            device=device,
        )
        count = boxes_cxcywh.shape[0]
        if count == 0:
            boxes_cxcywh = boxes_cxcywh.new_tensor([[0.5, 0.5, 1.0, 1.0]])
            count = 1

        if count < self.num_proposals:
            placeholders = (
                randn(
                    self.num_proposals - count,
                    4,
                    generator=generator,
                    device=device,
                )
                / 6.0
                + 0.5
            )
            placeholders[:, 2:] = placeholders[:, 2:].clamp_min(1e-4)
            clean = torch.cat((boxes_cxcywh, placeholders), dim=0)
        elif count > self.num_proposals:
            clean = boxes_cxcywh[: self.num_proposals]
        else:
            clean = boxes_cxcywh

        clean = (clean * 2.0 - 1.0) * self.config.snr_scale
        diffused = self.q_sample(clean, timestep, noise=noise)
        diffused = diffused.clamp(-self.config.snr_scale, self.config.snr_scale)
        diffused = (diffused / self.config.snr_scale + 1.0) / 2.0
        return _cxcywh_to_xyxy(diffused), noise, timestep

    def forward(
        self,
        boxes_xyxy: Float[torch.Tensor, "batch proposals 4"],
        timesteps: Int[torch.Tensor, "batch"]
        | Int[torch.Tensor, ""]
        | Float[torch.Tensor, "batch"]
        | Float[torch.Tensor, ""],
        text_features: Float[torch.Tensor, "batch text text_dim"],
        text_mask: Bool[torch.Tensor, "batch text 1"] | None = None,
        images: Float[torch.Tensor, "batch channels height width"] | None = None,
        image_scales: Float[torch.Tensor, "batch 4"] | None = None,
        absolute_boxes_xyxy: Float[torch.Tensor, "batch proposals 4"] | None = None,
    ) -> RADMDenoiserOutput:
        """Predict proposal classes and denoised boxes."""
        if text_features.ndim != 3:
            raise ValueError("text_features must have shape (batch, text, dim)")

        if boxes_xyxy.shape[1] != self.num_proposals:
            raise ValueError(f"expected {self.num_proposals} proposals")

        if timesteps.ndim == 0:
            timesteps = timesteps.repeat(boxes_xyxy.shape[0])

        timestep_batch = timesteps.to(
            device=boxes_xyxy.device, dtype=torch.long
        ).reshape(-1)
        if images is None:
            features = self._fallback_features(boxes_xyxy)
            default_scale = boxes_xyxy.new_ones((boxes_xyxy.shape[0], 4))
        else:
            features = self.backbone(images)
            default_scale = boxes_xyxy.new_tensor(
                (images.shape[-1], images.shape[-2], images.shape[-1], images.shape[-2])
            ).expand(boxes_xyxy.shape[0], -1)

        resolved_scales = default_scale if image_scales is None else image_scales
        auxiliary_logits, auxiliary_boxes_absolute = self.head(
            features,
            boxes_xyxy,
            text_features,
            text_mask,
            timestep_batch,
            resolved_scales,
            absolute_boxes_xyxy,
        )
        image_scale = resolved_scales[:, None, :].to(device=boxes_xyxy.device)
        auxiliary_boxes = auxiliary_boxes_absolute / image_scale
        logits = auxiliary_logits[-1]
        pred_original = auxiliary_boxes[-1]
        return RADMDenoiserOutput(
            logits=logits,
            boxes_xyxy=pred_original,
            pred_original_sample=pred_original,
            pred_noise=boxes_xyxy - pred_original,
            auxiliary_logits=auxiliary_logits,
            auxiliary_boxes_xyxy=auxiliary_boxes,
            auxiliary_boxes_absolute_xyxy=auxiliary_boxes_absolute,
        )

    def _fallback_features(
        self, boxes_xyxy: Float[torch.Tensor, "batch proposals 4"]
    ) -> dict[str, Float[torch.Tensor, "batch channels height width"]]:
        """Create a deterministic proposal-only feature pyramid for serialization tests."""
        zero = boxes_xyxy.new_zeros(boxes_xyxy.shape[0], 256, 1, 1)
        return {name: zero for name in ("p2", "p3", "p4", "p5")}