Model Card for LayoutGAN++¶

class LayoutGANPPConfig(PretrainedConfig):
    """Configuration for the LayoutGAN++ generator.

    Args:
        dataset_name: Dataset key or alias used to resolve labels and sequence length.
        latent_size: Size of each per-element latent vector.
        num_labels: Optional label vocabulary size override.
        id2label: Optional mapping from label IDs to display labels.
        label2id: Optional mapping from display labels to label IDs.
        d_model: Transformer hidden size used by the generator.
        nhead: Number of transformer attention heads.
        num_layers: Number of transformer encoder layers.
        bbox_format: Bounding-box format produced by the model.
        bbox_normalized: Whether generated boxes are normalized to the canvas.
        max_position_embeddings: Maximum element count for generated layouts.
        **kwargs: Extra `PretrainedConfig` keyword arguments.

    Examples:
        >>> config = LayoutGANPPConfig(dataset_name="rico")
        >>> config.model_type
        'layoutganpp'
    """

    model_type = "layoutganpp"

    def __init__(
        self,
        dataset_name: DatasetName | str = DatasetName.rico13,
        latent_size: int = 4,
        num_labels: int | None = None,
        id2label: Id2LabelMapping | None = None,
        label2id: dict[str, int] | None = None,
        d_model: int = 512,
        nhead: int = 8,
        num_layers: int = 4,
        bbox_format: BoxFormat | str = BoxFormat.xywh,
        bbox_normalized: bool = True,
        max_position_embeddings: int | None = None,
        **kwargs: LayoutGANPPConfigValue,
    ) -> None:
        """Initialize a LayoutGAN++ config.

        Args:
            dataset_name: Dataset key or alias used to resolve labels and metadata.
            latent_size: Size of each latent vector passed to the generator.
            num_labels: Optional explicit label vocabulary size.
            id2label: Optional label ID to text mapping.
            label2id: Optional label text to ID mapping.
            d_model: Transformer hidden size.
            nhead: Number of attention heads.
            num_layers: Number of transformer encoder layers.
            bbox_format: Format of generated bounding boxes.
            bbox_normalized: Whether generated boxes are normalized.
            max_position_embeddings: Optional maximum layout length override.
            **kwargs: Extra `PretrainedConfig` keyword arguments.

        Raises:
            ValueError: If `dataset_name` is not a supported LayoutGAN++ dataset.

        Examples:
            >>> LayoutGANPPConfig(dataset_name="publaynet").num_labels
            5
        """
        metadata = dataset_metadata(dataset_name)
        raw_id2label = id2label or id2label_for_dataset(dataset_name)
        normalized_id2label = {int(k): v for k, v in raw_id2label.items()}
        normalized_label2id = label2id or {
            label: i for i, label in normalized_id2label.items()
        }
        resolved_num_labels = num_labels or len(normalized_id2label)
        super().__init__(
            id2label=normalized_id2label,
            label2id=normalized_label2id,
        )
        for key, value in kwargs.items():
            setattr(self, key, value)
        self.dataset_name = str(metadata["name"])
        self.latent_size = latent_size
        self.num_labels = resolved_num_labels
        self.d_model = d_model
        self.nhead = nhead
        self.num_layers = num_layers
        self.bbox_format = str(normalize_box_format(bbox_format))
        self.bbox_normalized = bbox_normalized
        self.max_position_embeddings = (
            max_position_embeddings or max_elements_for_dataset(metadata["name"])
        )
        self.architectures = ["LayoutGANPPModel"]

class LayoutGANPPModel(PreTrainedModel):
    """Transformers-compatible LayoutGAN++ generator.

    Args:
        config: LayoutGAN++ model configuration.

    Examples:
        >>> config = LayoutGANPPConfig(num_labels=2, id2label={0: "a", 1: "b"})
        >>> model = LayoutGANPPModel(config)
        >>> model.config.model_type
        'layoutganpp'
    """

    config_class = LayoutGANPPConfig
    base_model_prefix = "layoutganpp"
    supports_gradient_checkpointing = False

    def __init__(self, config: LayoutGANPPConfig) -> None:
        """Initialize the LayoutGAN++ generator layers.

        Args:
            config: LayoutGAN++ model configuration.

        Examples:
            >>> model = LayoutGANPPModel(LayoutGANPPConfig())
            >>> model.base_model_prefix
            'layoutganpp'
        """
        super().__init__(config)
        self.fc_z = nn.Linear(config.latent_size, config.d_model // 2)
        self.emb_label = nn.Embedding(config.num_labels, config.d_model // 2)
        self.fc_in = nn.Linear(config.d_model, config.d_model)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=config.d_model,
            nhead=config.nhead,
            dim_feedforward=config.d_model // 2,
            batch_first=False,
        )
        self.transformer = nn.TransformerEncoder(
            encoder_layer, num_layers=config.num_layers
        )
        self.fc_out = nn.Linear(config.d_model, 4)
        self.post_init()

    def _init_weights(self, module: nn.Module) -> None:
        """Keep PyTorch constructor initialization during Transformers setup."""
        del module

    def forward(
        self,
        latents: Float[torch.Tensor, "batch elements latent"],
        labels: Int[torch.Tensor, "batch elements"],
        attention_mask: Bool[torch.Tensor, "batch elements"] | None = None,
        padding_mask: Bool[torch.Tensor, "batch elements"] | None = None,
        return_dict: bool = True,
    ) -> (
        LayoutGANPPModelOutput
        | tuple[
            Float[torch.Tensor, "batch elements 4"],
            Int[torch.Tensor, "batch elements"],
            Bool[torch.Tensor, "batch elements"],
        ]
    ):
        """Run a forward pass from latents and label IDs.

        Args:
            latents: Per-element latent vectors shaped `(batch, sequence, latent_size)`.
            labels: Label IDs shaped `(batch, sequence)`.
            attention_mask: Optional mask where true values mark valid labels.
            padding_mask: Optional mask where true values mark padded labels.
            return_dict: Whether to return a `LayoutGANPPModelOutput`.

        Returns:
            Model output dataclass or tuple containing boxes, labels, and mask.

        Raises:
            ValueError: If labels or latents have invalid shape or label IDs.

        Examples:
            >>> model = LayoutGANPPModel(LayoutGANPPConfig(num_labels=2))
            >>> labels = torch.tensor([[0, 1]])
            >>> latents = torch.zeros(1, 2, model.config.latent_size)
            >>> tuple(model(latents=latents, labels=labels).bbox.shape)
            (1, 2, 4)
        """
        labels = labels.to(dtype=torch.long)
        if labels.ndim != 2:
            raise ValueError("labels must have shape (batch, sequence)")

        if latents.shape[:2] != labels.shape:
            raise ValueError("latents must have shape (batch, sequence, latent_size)")

        if latents.shape[-1] != self.config.latent_size:
            raise ValueError(
                f"latents last dimension must be {self.config.latent_size}"
            )

        if labels.numel() and (
            int(labels.min().item()) < 0
            or int(labels.max().item()) >= self.config.num_labels
        ):
            raise ValueError("labels contain ids outside config.num_labels")

        if padding_mask is None:
            if attention_mask is None:
                padding_mask = torch.zeros(
                    labels.shape, dtype=torch.bool, device=labels.device
                )
            else:
                padding_mask = ~attention_mask.to(
                    device=labels.device, dtype=torch.bool
                )
        else:
            padding_mask = padding_mask.to(device=labels.device, dtype=torch.bool)
        latents = latents.to(device=labels.device, dtype=self.dtype)
        z = self.fc_z(latents)
        label_emb = self.emb_label(labels)
        hidden = torch.cat([z, label_emb], dim=-1)
        hidden = torch.relu(self.fc_in(hidden)).permute(1, 0, 2)
        hidden = self.transformer(hidden, src_key_padding_mask=padding_mask)
        bbox = torch.sigmoid(self.fc_out(hidden.permute(1, 0, 2)))
        mask = ~padding_mask
        if not return_dict:
            return bbox, labels, mask
        return LayoutGANPPModelOutput(
            bbox=bbox, labels=labels, mask=mask, latents=latents
        )

    @torch.no_grad()
    def generate(
        self,
        *,
        batch_size: int = 1,
        condition_type: ConditionType | str = ConditionType.label,
        bbox: Float[torch.Tensor, "batch elements 4"] | None = None,
        labels: Int[torch.Tensor, "batch elements"] | None = None,
        mask: Bool[torch.Tensor, "batch elements"] | None = None,
        attention_mask: Bool[torch.Tensor, "batch elements"] | 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,
        seed: int | None = None,
        generator: torch.Generator | None = None,
        num_inference_steps: int | None = None,
        output_type: OutputType | str = OutputType.dataclass,
        return_intermediates: bool = False,
        latents: Float[torch.Tensor, "batch elements latent"] | None = None,
    ) -> LayoutGenerationOutput | LayoutGANPPOutputDict:
        """Generate layouts from label conditions.

        Args:
            batch_size: Requested batch size; label shape determines the final value.
            condition_type: Condition type or alias. LayoutGAN++ supports label conditions.
            bbox: Reserved compatibility argument.
            labels: Required label IDs for generation.
            mask: Optional valid-element mask.
            attention_mask: Optional valid-element mask.
            num_elements: Reserved compatibility argument.
            box_format: Reserved compatibility argument.
            normalized: Reserved compatibility argument.
            canvas_size: Reserved compatibility argument.
            seed: Optional random seed for latent sampling.
            generator: Optional PyTorch random generator.
            num_inference_steps: Reserved compatibility argument.
            output_type: Return format, either `dataclass` or `dict`.
            return_intermediates: Whether to include generation intermediates.
            latents: Optional fixed latent vectors.

        Returns:
            A layout generation dataclass or dictionary.

        Raises:
            ValueError: If labels are missing, generation options are unsupported,
                or output type is invalid.

        Examples:
            >>> model = LayoutGANPPModel(LayoutGANPPConfig(num_labels=2))
            >>> out = model.generate(labels=torch.tensor([[0, 1]]), seed=0)
            >>> tuple(out.bbox.shape)
            (1, 2, 4)
        """
        del bbox, num_elements, normalized, canvas_size, num_inference_steps
        normalize_box_format(box_format)
        canonical = normalize_condition_type(condition_type)
        if canonical is ConditionType.unconditional:
            raise ValueError(
                "layoutganpp v1 requires labels; unconditional is unsupported"
            )

        if canonical is not ConditionType.label:
            raise ValueError(f"Unsupported condition_type for layoutganpp: {canonical}")

        if labels is None:
            raise ValueError("labels are required for layoutganpp generation")

        device = next(self.parameters()).device
        labels = torch.as_tensor(labels, dtype=torch.long, device=device)
        if labels.ndim == 1:
            labels = labels.unsqueeze(0)
        batch_size = labels.shape[0]
        if mask is not None:
            attention_mask = mask
        if attention_mask is None:
            attention_mask = torch.ones(labels.shape, dtype=torch.bool, device=device)
        else:
            attention_mask = torch.as_tensor(
                attention_mask, dtype=torch.bool, device=device
            )
            if attention_mask.ndim == 1:
                attention_mask = attention_mask.unsqueeze(0)
        if latents is None:
            latents = self._sample_latents(
                (batch_size, labels.shape[1], self.config.latent_size),
                seed=seed,
                generator=generator,
                device=device,
                dtype=self.dtype,
            )
        else:
            latents = torch.as_tensor(latents, dtype=self.dtype, device=device)
        out = self.forward(
            latents=latents,
            labels=labels,
            attention_mask=attention_mask,
            return_dict=True,
        )
        assert isinstance(out, LayoutGANPPModelOutput)
        assert out.labels is not None
        assert out.mask is not None
        assert out.latents is not None
        layout = LayoutGenerationOutput(
            bbox=out.bbox.detach().cpu(),
            labels=out.labels.detach().cpu(),
            mask=out.mask.detach().cpu(),
            id2label={int(k): v for k, v in self.config.id2label.items()},
            intermediates={
                "condition_type": canonical,
                "latents": out.latents.detach().cpu() if return_intermediates else None,
            }
            if return_intermediates
            else None,
        )
        resolved_output_type = normalize_output_type(output_type)
        if resolved_output_type is OutputType.dict:
            return cast(LayoutGANPPOutputDict, dict(layout))
        if resolved_output_type is OutputType.dataclass:
            return layout
        assert_never(resolved_output_type)

    def _sample_latents(
        self,
        shape: tuple[int, int, int],
        *,
        seed: int | None,
        generator: torch.Generator | None,
        device: torch.device,
        dtype: torch.dtype,
    ) -> Float[torch.Tensor, "batch elements latent"]:
        generator = resolve_torch_generator(generator=generator, seed=seed)
        return randn(shape, generator=generator, device=device, dtype=dtype)

@dataclass
class LayoutGANPPModelOutput(ModelOutput):
    """Raw LayoutGAN++ model output.

    Args:
        bbox: Generated normalized `xywh` boxes with shape `(batch, sequence, 4)`.
        labels: Optional label IDs used for generation.
        mask: Optional valid-element mask.
        latents: Optional latent vectors used by the generator.

    Examples:
        >>> out = LayoutGANPPModelOutput(bbox=torch.zeros(1, 1, 4))
        >>> tuple(out.bbox.shape)
        (1, 1, 4)
    """

    bbox: Float[torch.Tensor, "batch elements 4"]
    labels: Int[torch.Tensor, "batch elements"] | None = None
    mask: Bool[torch.Tensor, "batch elements"] | None = None
    latents: Float[torch.Tensor, "batch elements latent"] | None = None

class OutputType(StrEnum):
    """Supported LayoutGAN++ generation output formats."""

    dataclass = auto()
    dict = auto()

class LayoutGANPPPipeline(LayoutGenerationPipeline):
    """Transformers pipeline for LayoutGAN++ label-conditioned generation.

    Args:
        model: LayoutGAN++ model instance.
        processor: Optional processor for label encoding and decoding.
        config: Optional root pipeline config. Defaults to `model.config`.
        device: Optional torch device passed to the base pipeline.
        binary_output: Whether the base pipeline should produce binary output.

    Examples:
        >>> model = LayoutGANPPModel(LayoutGANPPConfig(num_labels=2))
        >>> pipe = LayoutGANPPPipeline(model=model)
        >>> pipe.model.config.model_type
        'layoutganpp'
    """

    config_class: ClassVar[type[PretrainedConfig]] = LayoutGANPPConfig
    component_specs: ClassVar[dict[str, PipelineComponentSpec]] = (
        model_processor_component_specs(
            model_loader=_load_model_component,
            processor_loader=_load_processor_component,
        )
    )

    config: LayoutGANPPConfig
    model: LayoutGANPPModel
    processor: LayoutGANPPProcessor

    def __init__(
        self,
        model: LayoutGANPPModel,
        processor: LayoutGANPPProcessor | None = None,
        config: LayoutGANPPConfig | None = None,
        device: int | torch.device | None = None,
        binary_output: bool = False,
    ) -> None:
        """Initialize a LayoutGAN++ pipeline.

        Args:
            model: LayoutGAN++ model instance.
            processor: Optional processor for label encoding and decoding.
            config: Optional root pipeline config.
            device: Optional torch device passed to the base pipeline.
            binary_output: Whether the base pipeline should produce binary output.

        Examples:
            >>> pipe = LayoutGANPPPipeline(LayoutGANPPModel(LayoutGANPPConfig()))
            >>> isinstance(pipe.processor, LayoutGANPPProcessor)
            True
        """
        _ = binary_output
        super().__init__(config or model.config)
        self.config = config or model.config
        self.model = model
        self.processor = processor or LayoutGANPPProcessor(
            dataset_name=model.config.dataset_name,
            id2label=model.config.id2label,
        )
        if device is not None:
            resolved_device = (
                torch.device("cpu")
                if isinstance(device, int) and device < 0
                else torch.device(f"cuda:{device}")
                if isinstance(device, int)
                else device
            )
            self.to(resolved_device)

    @classmethod
    def _from_pretrained_components(
        cls,
        *,
        config: PretrainedConfig,
        components: Mapping[str, PipelineComponent | None],
    ) -> LayoutGANPPPipeline:
        """Build a pipeline from loaded root components."""
        return cls(
            config=cast(LayoutGANPPConfig, config),
            model=cast(LayoutGANPPModel, components["model"]),
            processor=cast(LayoutGANPPProcessor, components["processor"]),
        )

    def _sanitize_parameters(
        self, **kwargs: LayoutGANPPPipelineKwarg
    ) -> tuple[
        dict[str, LayoutGANPPPipelineKwarg],
        dict[str, LayoutGANPPPipelineKwarg],
        dict[str, LayoutGANPPPipelineKwarg],
    ]:
        sanitized = cast(
            dict[str, LayoutGANPPPipelineKwarg],
            kwargs,
        )
        return {}, sanitized, {}

    def preprocess(
        self,
        input_: list[list[str | int]]
        | list[str | int]
        | Int[torch.Tensor, "batch elements"]
        | None = None,
        **preprocess_parameters: LayoutGANPPPipelineKwarg,
    ) -> BatchEncoding:
        """Encode pipeline inputs into model inputs.

        Args:
            input_: Labels supplied as the positional pipeline input.
            **preprocess_parameters: Keyword labels and generation arguments.

        Returns:
            Batch encoding containing label IDs, attention mask, and generation kwargs.

        Raises:
            ValueError: If labels are not supplied or cannot be encoded.

        Examples:
            >>> pipe = LayoutGANPPPipeline(LayoutGANPPModel(LayoutGANPPConfig()))
            >>> "labels" in pipe.preprocess(["Toolbar"])
            True
        """
        labels = preprocess_parameters.pop("labels", input_)
        if labels is None:
            raise ValueError("labels are required for LayoutGANPPPipeline")

        encoded = self._layoutganpp_processor()(
            cast(list[list[str | int]] | list[str | int] | torch.Tensor, labels)
        )
        encoded.update(preprocess_parameters)
        return encoded

    def _forward(
        self,
        model_inputs: dict[str, LayoutGANPPPipelineKwarg],
        **forward_params: LayoutGANPPPipelineKwarg,
    ) -> LayoutGenerationOutput | LayoutGANPPOutputDict:
        del forward_params
        labels = torch.as_tensor(model_inputs.pop("labels"), dtype=torch.long)
        attention_mask = torch.as_tensor(
            model_inputs.pop("attention_mask"), dtype=torch.bool
        )
        condition_type = cast(
            ConditionType | str, model_inputs.pop("condition_type", ConditionType.label)
        )
        bbox = cast(torch.Tensor | None, model_inputs.pop("bbox", None))
        mask = cast(torch.Tensor | None, model_inputs.pop("mask", None))
        num_elements = cast(
            int | list[int] | torch.Tensor | None,
            model_inputs.pop("num_elements", None),
        )
        box_format = cast(
            BoxFormat | str, model_inputs.pop("box_format", BoxFormat.xywh)
        )
        normalized = cast(bool, model_inputs.pop("normalized", True))
        canvas_size = cast(
            tuple[int, int] | None, model_inputs.pop("canvas_size", None)
        )
        seed = cast(int | None, model_inputs.pop("seed", None))
        generator = cast(torch.Generator | None, model_inputs.pop("generator", None))
        num_inference_steps = cast(
            int | None, model_inputs.pop("num_inference_steps", None)
        )
        output_type = cast(
            OutputType | str, model_inputs.pop("output_type", OutputType.dataclass)
        )
        return_intermediates = cast(
            bool, model_inputs.pop("return_intermediates", False)
        )
        latents = cast(torch.Tensor | None, model_inputs.pop("latents", None))
        if model_inputs:
            unknown = ", ".join(sorted(model_inputs))
            raise ValueError(f"Unsupported generation kwargs: {unknown}")

        return self._layoutganpp_model().generate(
            condition_type=condition_type,
            bbox=bbox,
            labels=labels,
            mask=mask,
            attention_mask=attention_mask,
            num_elements=num_elements,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
            seed=seed,
            generator=generator,
            num_inference_steps=num_inference_steps,
            output_type=output_type,
            return_intermediates=return_intermediates,
            latents=latents,
        )

    def postprocess(
        self,
        model_outputs: LayoutGenerationOutput | LayoutGANPPOutputDict,
        **kwargs: LayoutGANPPPipelineKwarg,
    ) -> LayoutGenerationOutput | LayoutGANPPOutputDict:
        """Return generated layouts from the pipeline output.

        Args:
            model_outputs: Output produced by `LayoutGANPPModel.generate`.
            **kwargs: Reserved post-processing keyword arguments.

        Returns:
            The generated layout output unchanged.

        Examples:
            >>> output = LayoutGenerationOutput(bbox=torch.zeros(1, 1, 4))
            >>> pipe = LayoutGANPPPipeline(LayoutGANPPModel(LayoutGANPPConfig()))
            >>> pipe.postprocess(output) is output
            True
        """
        del kwargs
        return model_outputs

    @torch.no_grad()
    def __call__(
        self,
        labels: list[list[str | int]]
        | list[str | int]
        | Int[torch.Tensor, "batch elements"]
        | None = None,
        *,
        batch_size: int = 1,
        condition_type: ConditionType | str = ConditionType.label,
        bbox: Float[torch.Tensor, "batch elements 4"] | None = None,
        mask: Bool[torch.Tensor, "batch elements"] | None = None,
        attention_mask: Bool[torch.Tensor, "batch elements"] | 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,
        seed: int | None = None,
        generator: torch.Generator | None = None,
        num_inference_steps: int | None = None,
        output_type: OutputType | str = OutputType.dataclass,
        return_intermediates: bool = False,
        latents: Float[torch.Tensor, "batch elements latent"] | None = None,
    ) -> LayoutGenerationOutput | LayoutGANPPOutputDict:  # ty: ignore[invalid-method-override]
        """Generate LayoutGAN++ boxes from labels.

        Args:
            labels: Label strings or label IDs to condition on.
            batch_size: Reserved compatibility argument.
            condition_type: Condition type or alias.
            bbox: Reserved compatibility argument.
            mask: Optional valid-element mask.
            attention_mask: Optional valid-element mask.
            num_elements: Reserved compatibility argument.
            box_format: Reserved compatibility argument.
            normalized: Reserved compatibility argument.
            canvas_size: Reserved compatibility argument.
            seed: Optional random seed for latent sampling.
            generator: Optional PyTorch random generator.
            num_inference_steps: Reserved compatibility argument.
            output_type: Return format, either `dataclass` or `dict`.
            return_intermediates: Whether to include generation intermediates.
            latents: Optional fixed latent vectors.

        Returns:
            A layout generation dataclass or dictionary.

        Raises:
            ValueError: If labels are missing or generation options are invalid.

        Examples:
            >>> pipe = LayoutGANPPPipeline(LayoutGANPPModel(LayoutGANPPConfig()))
            >>> out = pipe(labels=["Toolbar"], seed=0)
            >>> tuple(out.bbox.shape)
            (1, 1, 4)
        """
        del batch_size
        if labels is None:
            raise ValueError("labels are required for layoutganpp v1")

        if isinstance(labels, torch.Tensor):
            encoded_labels = labels
            resolved_mask = attention_mask if attention_mask is not None else mask
        else:
            encoded = self._layoutganpp_processor()(labels)
            encoded_labels = encoded["labels"]
            if attention_mask is not None:
                resolved_mask = attention_mask
            else:
                resolved_mask = encoded["attention_mask"] if mask is None else mask
        return self._layoutganpp_model().generate(
            condition_type=condition_type,
            bbox=bbox,
            labels=cast(torch.Tensor, encoded_labels),
            mask=cast(torch.Tensor | None, resolved_mask),
            num_elements=num_elements,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
            seed=seed,
            generator=generator,
            num_inference_steps=num_inference_steps,
            output_type=output_type,
            return_intermediates=return_intermediates,
            latents=latents,
        )

    def _layoutganpp_model(self) -> LayoutGANPPModel:
        return self.model

    def _layoutganpp_processor(self) -> LayoutGANPPProcessor:
        return self.processor