Model Card for DS-GAN¶

class DSGANConfig(PretrainedConfig):
    """Store DS-GAN architecture and dataset metadata.

    Args:
        dataset_name: Dataset key. DS-GAN currently supports PKU PosterLayout.
        backbone: Timm ResNet backbone name used by the reference checkpoint.
        max_elem: Maximum number of generated elements.
        in_channels: CNN-LSTM input channels after flattening class and box planes.
        out_channels: CNN-LSTM convolution output channels.
        hidden_size: Bidirectional LSTM hidden size.
        num_layers: Number of LSTM layers.
        output_size: Combined class and box output width in the internal model.
        image_size: Processor/model input size as ``(height, width)``.
        reference_canvas_size: Reference normalization canvas as ``(width, height)``.
        backbone_feature_size: Flattened ResNet-FPN spatial size. The reference
            default is ``22 * 15 = 330`` for ``image_size=(350, 240)``.
        model_num_classes: Internal class channels including ``0 = no object``.
        id2label: Public zero-based semantic label mapping.
        model_subfolder: Pipeline subfolder for the model component.
        processor_subfolder: Pipeline subfolder for the processor component.

    Examples:
        >>> DSGANConfig().max_elem
        32
    """

    model_type = "ds_gan"

    def __init__(
        self,
        dataset_name: DatasetName | str = DatasetName.pku_posterlayout,
        backbone: str = "resnet50",
        max_elem: int = 32,
        in_channels: int = 8,
        out_channels: int = 32,
        hidden_size: int | None = None,
        num_layers: int = 4,
        output_size: int = 8,
        image_size: tuple[int, int] | list[int] = (350, 240),
        reference_canvas_size: tuple[int, int] | list[int] = (513, 750),
        backbone_feature_size: int = 330,
        model_num_classes: int = 4,
        id2label: Id2LabelMapping | None = None,
        label2id: dict[str, int] | None = None,
        model_subfolder: str = "model",
        processor_subfolder: str = "processor",
        condition_types: list[str] | tuple[str, ...] | None = None,
        architectures: list[str] | None = None,
        model_type: str | None = None,
        transformers_version: str | None = None,
        torch_dtype: str | None = None,
        dtype: str | None = None,
        name_or_path: str = "",
        _commit_hash: str | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize DS-GAN configuration."""
        _ = (model_type, transformers_version)
        dataset = normalize_dataset_name(dataset_name)
        if dataset is not DatasetName.pku_posterlayout:
            raise ValueError(f"Unsupported DS-GAN dataset_name: {dataset_name}")

        public_id2label = {
            int(k): v for k, v in (id2label or _semantic_pku_id2label()).items()
        }
        public_label2id = label2id or {v: k for k, v in public_id2label.items()}
        super().__init__(
            id2label=public_id2label,
            label2id=public_label2id,
            architectures=architectures or ["DSGANModel"],
            torch_dtype=torch_dtype,  # ty: ignore[unknown-argument]
            dtype=dtype,
            name_or_path=name_or_path,  # ty: ignore[unknown-argument]
            _commit_hash=_commit_hash,  # ty: ignore[unknown-argument]
            **kwargs,  # ty: ignore[invalid-argument-type]
        )
        self.dataset_name = str(dataset)
        self.backbone = backbone
        self.max_elem = max_elem
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.hidden_size = hidden_size if hidden_size is not None else max_elem * 8
        self.num_layers = num_layers
        self.output_size = output_size
        self.image_size = tuple(int(v) for v in image_size)
        self.reference_canvas_size = tuple(int(v) for v in reference_canvas_size)
        self.backbone_feature_size = backbone_feature_size
        self.model_num_classes = model_num_classes
        self.model_subfolder = model_subfolder
        self.processor_subfolder = processor_subfolder
        self.condition_types = list(condition_types or ["content_image"])

    @property
    def public_num_labels(self) -> int:
        """Return the number of public semantic labels."""
        return len(cast(dict[int, str], self.id2label))

class DSGANModel(PreTrainedModel):
    """Transformers-compatible DS-GAN generator.

    Args:
        config: DS-GAN model configuration.

    Examples:
        >>> config = DSGANConfig(backbone="resnet18", max_elem=4, hidden_size=32, num_layers=2, image_size=(64, 64), backbone_feature_size=16)
        >>> model = DSGANModel(config)
        >>> model.config.model_type
        'ds_gan'
    """

    config_class = DSGANConfig
    base_model_prefix = "ds_gan"
    supports_gradient_checkpointing = False

    def _init_weights(self, module: nn.Module) -> None:
        """Leave module-specific initialization unchanged during ``post_init``."""
        del module

    def __init__(
        self, config: DSGANConfig, *, backbone_weights: str | Path | None = None
    ) -> None:
        """Initialize DS-GAN generator layers."""
        super().__init__(config)
        self.resnet_fpn = ResnetBackbone(config, backbone_weights=backbone_weights)
        self.cnnlstm = CNNLSTM(config)
        self.fc1 = nn.Linear(2 * config.hidden_size, config.output_size // 2)
        self.fc2 = nn.Linear(2 * config.hidden_size, config.output_size // 2)
        self.post_init()

    def forward(
        self,
        pixel_values: Float[torch.Tensor, "batch 4 height width"],
        layout: Float[torch.Tensor, "batch elements 2 4"],
        return_dict: bool = True,
    ) -> (
        DSGANModelOutput
        | tuple[
            Float[torch.Tensor, "batch elements 4"],
            Float[torch.Tensor, "batch elements 4"],
        ]
    ):
        """Run a DS-GAN generator forward pass.

        Args:
            pixel_values: RGB plus saliency tensor shaped ``(B, 4, H, W)``.
            layout: Initial internal layout shaped ``(B, max_elem, 2, 4)``.
            return_dict: Whether to return a dataclass output.

        Returns:
            Raw class probabilities and normalized center ``xywh`` boxes.

        Raises:
            ValueError: If tensor shapes do not match the config.
        """
        if pixel_values.ndim != 4 or pixel_values.shape[1] != 4:
            raise ValueError("pixel_values must have shape (batch, 4, height, width)")

        expected_layout = (pixel_values.shape[0], self.config.max_elem, 2, 4)
        if tuple(layout.shape) != expected_layout:
            raise ValueError(f"layout must have shape {expected_layout}")

        pixel_values = pixel_values.to(dtype=self.dtype)
        layout = layout.to(device=pixel_values.device, dtype=self.dtype)
        h0 = self.resnet_fpn(pixel_values)
        lstm_output = self.cnnlstm(layout, h0)
        class_probs = torch.softmax(self.fc1(lstm_output), dim=-1)
        bbox = torch.sigmoid(self.fc2(lstm_output))
        if not return_dict:
            return class_probs, bbox
        return DSGANModelOutput(
            class_probs=class_probs,
            bbox=bbox,
            initial_layout=layout,
        )

@dataclass
class DSGANModelOutput(ModelOutput):
    """Raw DS-GAN generator output.

    Attributes:
        class_probs: Internal class probabilities with shape
            ``(batch, elements, 4)`` where id 0 is ``no object``.
        bbox: Normalized center ``xywh`` boxes with shape
            ``(batch, elements, 4)``.
        initial_layout: Initial class/box layout passed to the generator.
    """

    class_probs: Float[torch.Tensor, "batch elements 4"]
    bbox: Float[torch.Tensor, "batch elements 4"] | None = None
    initial_layout: Float[torch.Tensor, "batch elements 2 4"] | None = None

class DSGANPipeline(LayoutGenerationPipeline):
    """Transformers-side pipeline for content-aware PosterLayout generation.

    Args:
        model: DS-GAN generator.
        processor: Optional processor for images and output decoding.
        config: Optional root pipeline config.
        device: Optional runtime device.

    Examples:
        >>> config = DSGANConfig(backbone="resnet18", max_elem=4, hidden_size=32, num_layers=2, image_size=(64, 64), backbone_feature_size=16)
        >>> pipe = DSGANPipeline(DSGANModel(config))
        >>> pipe.config.model_type
        'ds_gan'
    """

    config_class: ClassVar[type[PretrainedConfig]] = DSGANConfig
    component_specs: ClassVar[dict[str, PipelineComponentSpec]] = {
        "model": PipelineComponentSpec(
            attribute_name="model",
            loader=_load_model_component,
            marker_file="config.json",
            config_subfolder_attribute="model_subfolder",
        ),
        "processor": PipelineComponentSpec(
            attribute_name="processor",
            loader=_load_processor_component,
            marker_file="processor_config.json",
            save_with_is_main_process=False,
            config_subfolder_attribute="processor_subfolder",
        ),
    }

    config: DSGANConfig
    model: DSGANModel
    processor: DSGANProcessor

    def __init__(
        self,
        model: DSGANModel,
        processor: DSGANProcessor | None = None,
        config: DSGANConfig | None = None,
        device: str | torch.device | None = None,
    ) -> None:
        """Initialize DS-GAN pipeline."""
        super().__init__(config or model.config)
        self.config = config or model.config
        self.model = model
        self.processor = processor or DSGANProcessor(
            dataset_name=self.config.dataset_name,
            id2label=cast(dict[int | str, str], self.config.id2label),
            image_size=cast(tuple[int, int], self.config.image_size),
        )
        self.model.eval()
        if device is not None:
            self.to(device)

    @classmethod
    def _from_pretrained_components(
        cls,
        *,
        config: PretrainedConfig,
        components: Mapping[str, DSGANPipelineComponent | None],
    ) -> "DSGANPipeline":
        """Build a pipeline from loaded root components."""
        return cls(
            config=cast(DSGANConfig, config),
            model=cast(DSGANModel, components["model"]),
            processor=cast(DSGANProcessor, components["processor"]),
        )

    @torch.no_grad()
    def __call__(  # ty: ignore[invalid-method-override]
        self,
        images: ImageInput
        | list[ImageInput]
        | Float[torch.Tensor, "..."]
        | None = None,
        *,
        batch_size: int = 1,
        seed: int | None = None,
        generator: torch.Generator | None = None,
        condition_type: ConditionType | str = ConditionType.content_image,
        labels: Int[torch.Tensor, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Sequence[ArrayLikeInput]
        | None = None,
        mask: Bool[torch.Tensor, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        num_elements: int | list[int] | Int[torch.Tensor, "batch"] | None = None,
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        num_inference_steps: int | None = None,
        output_type: OutputType | str = OutputType.dataclass,
        return_intermediates: bool = False,
        saliency: ImageInput
        | list[ImageInput]
        | Float[torch.Tensor, "..."]
        | None = None,
        saliency_pfpnet: ImageInput
        | list[ImageInput]
        | Float[torch.Tensor, "..."]
        | None = None,
        saliency_basnet: ImageInput
        | list[ImageInput]
        | Float[torch.Tensor, "..."]
        | None = None,
        pixel_values: Float[torch.Tensor, "batch 4 height width"] | None = None,
        initial_layout: Float[torch.Tensor, "batch elements 2 4"] | None = None,
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Shaped[torch.Tensor, "..."]
            | dict[int, str]
            | Mapping[str, Shaped[torch.Tensor, "..."] | ConditionType | str | bool]
            | None,
        ]
    ):
        """Generate layouts from content images and saliency maps.

        Args:
            images: RGB image or batch. Required unless ``pixel_values`` is given.
            batch_size: Batch size used when ``pixel_values`` is supplied.
            seed: Convenience seed. Ignored when ``generator`` is supplied.
            generator: Explicit torch generator.
            condition_type: Must normalize to ``content_image``.
            labels: Optional public labels used only with ``bbox`` to provide a
                fixed initial layout.
            bbox: Optional public boxes used with ``labels`` for a fixed layout.
            mask: Optional valid-element mask for fixed initial layouts.
            num_elements: Reserved compatibility argument.
            box_format: Format of optional ``bbox``.
            normalized: Whether optional ``bbox`` is normalized.
            canvas_size: Pixel canvas size for optional unnormalized ``bbox``.
            num_inference_steps: Reserved compatibility argument.
            output_type: ``dataclass`` or ``dict``.
            return_intermediates: Whether to include raw model tensors.
            saliency: Optional single merged saliency map.
            saliency_pfpnet: Optional PFPNet saliency map.
            saliency_basnet: Optional BASNet saliency map.
            pixel_values: Preprocessed ``(B, 4, H, W)`` tensor.
            initial_layout: Optional internal layout ``(B, max_elem, 2, 4)``.

        Returns:
            Shared layout-generation output.

        Raises:
            ValueError: If the condition, image inputs, or fixed layout are invalid.
        """
        del num_elements, num_inference_steps
        canonical = normalize_condition_type(condition_type)
        resolved_output_type = normalize_output_type(output_type)
        device = self.device or next(self.model.parameters()).device
        if pixel_values is None:
            if images is None:
                raise ValueError("images or pixel_values are required for DS-GAN")

            encoded = self.processor(
                images,
                saliency=saliency,
                saliency_pfpnet=saliency_pfpnet,
                saliency_basnet=saliency_basnet,
            )
            pixel_values = cast(torch.Tensor, encoded["pixel_values"])
        pixel_values = pixel_values.to(device=device, dtype=self.model.dtype)
        batch_size = pixel_values.shape[0] if pixel_values is not None else batch_size
        if initial_layout is None and bbox is not None and labels is not None:
            encoded_layout = self.processor.encode_layout(
                bbox=bbox,
                labels=labels,
                mask=mask,
                box_format=box_format,
                normalized=normalized,
                canvas_size=canvas_size,
                max_elem=self.config.max_elem,
            )
            initial_layout = encoded_layout["layout"]
        if initial_layout is None:
            prepared = self.prepare_generator(generator=generator, seed=seed)
            initial_layout = random_initial_layout(
                batch_size,
                self.config.max_elem,
                generator=prepared,
                device=device,
                dtype=self.model.dtype,
            )
        else:
            initial_layout = initial_layout.to(device=device, dtype=self.model.dtype)
        model_output = self.model(
            pixel_values=pixel_values,
            layout=initial_layout,
            return_dict=True,
        )
        assert isinstance(model_output, DSGANModelOutput)
        intermediates = None
        if return_intermediates:
            intermediates = {
                "condition_type": canonical,
                "initial_layout": initial_layout.detach().cpu(),
                "class_probs": model_output.class_probs.detach().cpu(),
            }
        return self.processor.decode(
            class_probs=model_output.class_probs,
            bbox=model_output.bbox,
            output_type=resolved_output_type.value,
            intermediates=intermediates,
        )

class OutputType(StrEnum):
    """Supported DS-GAN pipeline output containers."""

    dataclass = auto()
    dict = auto()

class DSGANProcessor(ProcessorMixin):
    """Prepare PosterLayout RGB/saliency inputs and decode DS-GAN outputs.

    Args:
        dataset_name: Dataset key. Only PKU PosterLayout is supported.
        id2label: Public semantic labels excluding model ``no object``.
        image_size: Resize target as ``(height, width)``.

    Examples:
        >>> processor = DSGANProcessor()
        >>> processor.id2label[0]
        'text'
    """

    config_name = "processor_config.json"

    def __init__(
        self,
        dataset_name: DatasetName | str = DatasetName.pku_posterlayout,
        id2label: dict[int | str, str] | None = None,
        image_size: tuple[int, int] | list[int] = (350, 240),
    ) -> None:
        """Initialize processor metadata."""
        self.chat_template = None
        dataset = normalize_dataset_name(dataset_name)
        if dataset is not DatasetName.pku_posterlayout:
            raise ValueError(f"Unsupported DS-GAN dataset_name: {dataset_name}")

        self.dataset_name = str(dataset)
        default_id2label = {0: "text", 1: "logo", 2: "underlay"}
        raw_id2label = id2label or default_id2label
        self.id2label = {int(k): v for k, v in raw_id2label.items()}
        self.label2id = {v: k for k, v in self.id2label.items()}
        height, width = image_size
        self.image_size: tuple[int, int] = (int(height), int(width))

    def __call__(
        self,
        images: DSGANImageInput | Sequence[DSGANImageInput],
        *,
        saliency: DSGANImageInput | Sequence[DSGANImageInput] | None = None,
        saliency_pfpnet: DSGANImageInput | Sequence[DSGANImageInput] | None = None,
        saliency_basnet: DSGANImageInput | Sequence[DSGANImageInput] | None = None,
        return_tensors: Literal["pt"] = "pt",
    ) -> BatchEncoding:
        """Encode content images into ``pixel_values``.

        Args:
            images: RGB image or batch of RGB images.
            saliency: Optional saliency image or batch. If omitted and both
                saliency maps are given, the maps are merged by pixelwise max.
            saliency_pfpnet: Optional PFPNet saliency map.
            saliency_basnet: Optional BASNet saliency map.
            return_tensors: Tensor framework. Only ``pt`` is supported.

        Returns:
            Batch encoding containing ``pixel_values`` shaped ``(B, 4, H, W)``.
        """
        if return_tensors != "pt":
            raise ValueError("DSGANProcessor only supports return_tensors='pt'")

        image_rows = _ensure_batch(images)
        saliency_rows = self._resolve_saliency(
            len(image_rows),
            saliency=saliency,
            saliency_pfpnet=saliency_pfpnet,
            saliency_basnet=saliency_basnet,
        )
        tensors = []
        for image, sal in zip(image_rows, saliency_rows, strict=True):
            rgb = _to_rgb_tensor(image, self.image_size)
            sal_tensor = (
                torch.zeros(1, *self.image_size, dtype=torch.float32)
                if sal is None
                else _to_l_tensor(sal, self.image_size)
            )
            tensors.append(torch.cat((rgb, sal_tensor), dim=0))
        return BatchEncoding({"pixel_values": torch.stack(tensors)})

    def decode(
        self,
        *,
        class_probs: Float[torch.Tensor, "batch elements 4"],
        bbox: Float[torch.Tensor, "batch elements 4"],
        output_type: Literal["dataclass", "dict"] = "dataclass",
        scores: Float[torch.Tensor, "batch elements"] | None = None,
        intermediates: Mapping[str, Shaped[torch.Tensor, "..."] | str | bool]
        | None = None,
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Shaped[torch.Tensor, "..."]
            | dict[int, str]
            | Mapping[str, Shaped[torch.Tensor, "..."] | str | bool]
            | None,
        ]
    ):
        """Decode raw DS-GAN class probabilities and boxes.

        Args:
            class_probs: Internal class probabilities shaped ``(B, E, 4)``.
            bbox: Normalized center ``xywh`` boxes shaped ``(B, E, 4)``.
            output_type: Return format.
            scores: Optional per-element class scores.
            intermediates: Optional model-specific intermediate tensors.

        Returns:
            Shared layout output with public labels and mask semantics.
        """
        class_ids = torch.argmax(class_probs, dim=-1)
        mask = class_ids != 0
        public_labels = (class_ids - 1).clamp_min(0).long()
        resolved_scores = scores
        if resolved_scores is None:
            resolved_scores = class_probs.max(dim=-1).values
        output = LayoutGenerationOutput(
            bbox=bbox.detach().cpu().clamp(0.0, 1.0),
            labels=public_labels.detach().cpu(),
            mask=mask.detach().cpu(),
            id2label=dict(self.id2label),
            scores=resolved_scores.detach().cpu(),
            intermediates=intermediates,
        )
        if output_type == "dict":
            return dict(output)
        if output_type == "dataclass":
            return output
        raise ValueError(f"Unsupported output_type: {output_type}")

    def encode_layout(
        self,
        *,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Float[np.ndarray, "batch elements 4"]
        | Sequence[ArrayLikeInput],
        labels: Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput],
        mask: Bool[torch.Tensor, "batch elements"]
        | Bool[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        max_elem: int = 32,
    ) -> dict[str, Shaped[torch.Tensor, "..."]]:
        """Encode public boxes/labels into the internal layout tensor.

        Args:
            bbox: Public boxes.
            labels: Public zero-based semantic labels.
            mask: Optional valid-element mask.
            box_format: Input box format.
            normalized: Whether the boxes are normalized.
            canvas_size: Pixel canvas size used when ``normalized=False``.
            max_elem: Output slot count.

        Returns:
            Dictionary with internal ``layout``, normalized ``bbox``, labels, and mask.
        """
        bbox_t, labels_t, mask_t = prepare_layout_tensors(
            bbox=bbox,
            labels=labels,
            mask=mask,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
            clamp_converted_normalized=True,
        )
        bbox_t, labels_t, mask_t = self.pad(bbox_t, labels_t, mask_t, max_elem=max_elem)
        model_labels = torch.zeros_like(labels_t)
        model_labels[mask_t] = labels_t[mask_t] + 1
        class_one_hot = torch.nn.functional.one_hot(model_labels, num_classes=4).to(
            dtype=bbox_t.dtype
        )
        layout = torch.stack((class_one_hot, bbox_t), dim=2)
        return {"layout": layout, "bbox": bbox_t, "labels": labels_t, "mask": mask_t}

    def pad(
        self,
        bbox: Float[torch.Tensor, "batch elements 4"],
        labels: Int[torch.Tensor, "batch elements"],
        mask: Bool[torch.Tensor, "batch elements"],
        *,
        max_elem: int,
    ) -> tuple[
        Float[torch.Tensor, "batch padded_elements 4"],
        Int[torch.Tensor, "batch padded_elements"],
        Bool[torch.Tensor, "batch padded_elements"],
    ]:
        """Pad layout tensors to DS-GAN ``max_elem`` slots."""
        if bbox.shape[1] > max_elem:
            raise ValueError(f"DS-GAN supports at most {max_elem} elements")

        pad_count = max_elem - bbox.shape[1]
        if pad_count:
            bbox = torch.cat((bbox, torch.zeros(bbox.shape[0], pad_count, 4)), dim=1)
            labels = torch.cat(
                (labels, torch.zeros(labels.shape[0], pad_count, dtype=torch.long)),
                dim=1,
            )
            mask = torch.cat(
                (mask, torch.zeros(mask.shape[0], pad_count, dtype=torch.bool)),
                dim=1,
            )
        labels = labels.clone()
        labels[~mask] = 0
        return bbox, labels, mask

    def _resolve_saliency(
        self,
        batch_size: int,
        *,
        saliency: DSGANImageInput | Sequence[DSGANImageInput] | None,
        saliency_pfpnet: DSGANImageInput | Sequence[DSGANImageInput] | None,
        saliency_basnet: DSGANImageInput | Sequence[DSGANImageInput] | None,
    ) -> list[DSGANImageInput | Float[torch.Tensor, "1 height width"] | None]:
        if saliency is not None:
            rows = _ensure_batch(saliency)
            if len(rows) != batch_size:
                raise ValueError("saliency batch size must match images")

            return cast(
                list[DSGANImageInput | Float[torch.Tensor, "1 height width"] | None],
                rows,
            )
        if saliency_pfpnet is None and saliency_basnet is None:
            return [None] * batch_size
        first = (
            _ensure_batch(saliency_pfpnet)
            if saliency_pfpnet is not None
            else [None] * batch_size
        )
        second = (
            _ensure_batch(saliency_basnet)
            if saliency_basnet is not None
            else [None] * batch_size
        )
        if len(first) != batch_size or len(second) != batch_size:
            raise ValueError("saliency batch size must match images")

        merged: list[Float[torch.Tensor, "1 height width"]] = []
        for left, right in zip(first, second, strict=True):
            if left is None:
                merged.append(_to_l_tensor(right, self.image_size))
            elif right is None:
                merged.append(_to_l_tensor(left, self.image_size))
            else:
                merged.append(_merge_saliency_native(left, right, self.image_size))
        return cast(
            list[DSGANImageInput | Float[torch.Tensor, "1 height width"] | None], merged
        )