Layoutformerpp

class ConditionType(StrEnum):
    """Canonical condition names used by layout generation interfaces."""

    unconditional = auto()
    label = auto()
    label_size = auto()
    completion = auto()
    refinement = auto()
    text = auto()
    content_image = auto()
    relation = auto()
    hierarchical = auto()
    retrieval = auto()

@dataclass
class LayoutGenerationOutput(ModelOutput):
    """Canonical layout-generation output for Transformers-style APIs.

    Attributes:
        bbox: Normalized center ``xywh`` boxes with shape
            ``(batch, elements, 4)``.
        labels: Dataset-local integer labels with shape ``(batch, elements)``.
        mask: Boolean valid-element mask with shape ``(batch, elements)``.
        id2label: Mapping from integer label ids to display names.
        sequences: Optional raw token sequences.
        scores: Optional per-token or per-element scores.
        trajectory: Optional sampling trajectory.
        intermediates: Optional model-specific debug or auxiliary data.

    Examples:
        >>> import numpy as np
        >>> output = LayoutGenerationOutput(
        ...     bbox=np.zeros((1, 1, 4), dtype=np.float32),
        ...     labels=np.zeros((1, 1), dtype=np.int64),
        ...     mask=np.ones((1, 1), dtype=bool),
        ...     id2label={0: "text"},
        ... )
        >>> output["bbox"].shape
        (1, 1, 4)
    """

    bbox: (
        Float[np.ndarray, "batch elements 4"] | Float[torch.Tensor, "batch elements 4"]
    )
    labels: Int[np.ndarray, "batch elements"] | Int[torch.Tensor, "batch elements"] = (
        cast(
            'Int[np.ndarray, "batch elements"] | Int[torch.Tensor, "batch elements"]',
            None,
        )
    )
    mask: Bool[np.ndarray, "batch elements"] | Bool[torch.Tensor, "batch elements"] = (
        cast(
            'Bool[np.ndarray, "batch elements"] | Bool[torch.Tensor, "batch elements"]',
            None,
        )
    )
    id2label: dict[int, str] = cast(dict[int, str], None)
    sequences: object | None = None
    scores: object | None = None
    trajectory: object | None = None
    intermediates: object | None = None

class LayoutFormerPPConfig(PretrainedConfig):
    """Stores model, tokenizer, and task defaults for converted checkpoints."""

    model_type = "layoutformerpp"

    def __init__(
        self,
        vocab_size: int = 0,
        max_position_embeddings: int | None = None,
        d_model: int = 512,
        encoder_layers: int = 8,
        decoder_layers: int | None = None,
        encoder_attention_heads: int = 8,
        decoder_attention_heads: int | None = None,
        dropout: float = 0.1,
        dim_feedforward: int | None = None,
        share_embedding: bool = True,
        dataset: DatasetName | str = DatasetName.rico25,
        task: LayoutFormerPPTask | ConditionType | str = LayoutFormerPPTask.gen_t,
        max_num_elements: int = 20,
        bbox_format: BoxFormat | str = BoxFormat.ltwh,
        default_box_format: BoxFormat | str = BoxFormat.xywh,
        discrete_x_grid: int = 128,
        discrete_y_grid: int = 128,
        add_sep_token: bool = True,
        sort_by_dict: bool = True,
        add_task_embedding: bool = False,
        add_task_prompt_token_in_model: bool = False,
        num_task_prompt_token: int = 1,
        task_id: int | None = None,
        decode_max_length: int | None = None,
        eval_seed: int | None = None,
        gen_t_add_unk_token: bool = False,
        gen_ts_add_unk_token: bool = False,
        gen_r_add_unk_token: bool = False,
        gen_r_compact: bool = False,
        bos_token_id: int = 0,
        eos_token_id: int = 1,
        pad_token_id: int = 2,
        is_encoder_decoder: bool = True,
        condition_type: ConditionType | str | None = None,
        model_type: str | None = None,
        transformers_version: str | None = None,
        architectures: list[str] | None = None,
        output_hidden_states: bool | None = False,
        output_attentions: bool | None = False,
        return_dict: bool | None = True,
        chunk_size_feed_forward: int = 0,
        problem_type: Literal[
            "regression", "single_label_classification", "multi_label_classification"
        ]
        | None = None,
        id2label: dict[int | str, str] | None = None,
        label2id: dict[str, int] | None = None,
        torch_dtype: str | None = None,
        dtype: str | None = None,
        tie_word_embeddings: bool = True,
        task_specific_params: dict[
            str, str | int | float | bool | list[str | int | float | bool]
        ]
        | None = None,
        name_or_path: str = "",
        _name_or_path: str | None = None,
        _commit_hash: str | None = None,
        attn_implementation: str | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize architecture and task-specific generation defaults."""
        _ = (condition_type, model_type, transformers_version)
        normalized_dataset = normalize_layoutformerpp_dataset(dataset)
        normalized_task = normalize_layoutformerpp_task(task)
        condition_type = TASK_TO_CONDITION[normalized_task]
        self.dataset = layoutformerpp_dataset_slug(normalized_dataset)
        self.task = str(normalized_task)
        self.condition_type = str(condition_type)
        defaults = TASK_DEFAULTS.get((normalized_dataset, condition_type), {})
        if max_position_embeddings is None:
            max_position_embeddings = int(defaults.get("max_position_embeddings", 150))
        if decode_max_length is None:
            decode_max_length = int(defaults.get("decode_max_length", 120))
        if eval_seed is None:
            eval_seed = int(defaults.get("eval_seed", 100))

        normalized_id2label = (
            {int(key): str(value) for key, value in id2label.items()}
            if id2label is not None
            else None
        )

        self.vocab_size = vocab_size
        self.max_position_embeddings = max_position_embeddings
        self.d_model = d_model
        self.encoder_layers = encoder_layers
        self.decoder_layers = (
            decoder_layers if decoder_layers is not None else encoder_layers
        )
        self.encoder_attention_heads = encoder_attention_heads
        self.decoder_attention_heads = (
            decoder_attention_heads
            if decoder_attention_heads is not None
            else encoder_attention_heads
        )
        self.dropout = dropout
        self.dim_feedforward = (
            dim_feedforward if dim_feedforward is not None else d_model * 4
        )
        self.share_embedding = share_embedding
        self.max_num_elements = max_num_elements
        self.bbox_format = str(normalize_box_format(bbox_format))
        self.default_box_format = str(normalize_box_format(default_box_format))

        self.discrete_x_grid = discrete_x_grid
        self.discrete_y_grid = discrete_y_grid
        self.add_sep_token = add_sep_token
        self.sort_by_dict = sort_by_dict
        self.add_task_embedding = add_task_embedding
        self.add_task_prompt_token_in_model = add_task_prompt_token_in_model
        self.num_task_prompt_token = num_task_prompt_token
        self.task_id = task_id

        self.decode_max_length = decode_max_length
        self.eval_seed = eval_seed
        self.gen_t_add_unk_token = gen_t_add_unk_token
        self.gen_ts_add_unk_token = gen_ts_add_unk_token
        self.gen_r_add_unk_token = gen_r_add_unk_token
        self.gen_r_compact = gen_r_compact

        super().__init__(
            transformers_version=transformers_version,
            architectures=architectures,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            chunk_size_feed_forward=chunk_size_feed_forward,
            problem_type=problem_type,
            is_encoder_decoder=is_encoder_decoder,
            id2label=normalized_id2label,
            label2id=label2id,
            dtype=torch_dtype or dtype,
        )
        # Transformers v5 keeps model-specific token/vocabulary fields on the
        # subclass; only the common fields above belong in the base call.
        self.vocab_size = vocab_size
        self.bos_token_id = bos_token_id
        self.eos_token_id = eos_token_id
        self.pad_token_id = pad_token_id
        self.tie_word_embeddings = tie_word_embeddings
        if output_attentions is not None:
            self.output_attentions = output_attentions
        if task_specific_params is not None:
            self.task_specific_params = task_specific_params
        self.name_or_path = name_or_path if _name_or_path is None else _name_or_path
        self._commit_hash = _commit_hash
        self._attn_implementation = attn_implementation
        # Transformers v5 passes legacy and model-specific fields through the
        # tolerant config-loading path; preserve them as ordinary attributes.
        for key, value in kwargs.items():
            setattr(self, key, value)

class LayoutFormerPPForConditionalGeneration(PreTrainedModel):
    """Transformers `PreTrainedModel` with checkpoint-compatible module names."""

    config_class = LayoutFormerPPConfig
    base_model_prefix = "layoutformerpp"
    main_input_name = "input_ids"
    _tied_weights_keys = {
        "dec_embedding.weight": "enc_embedding.weight",
        "out.weight": "dec_embedding.weight",
    }

    def __init__(self, config: LayoutFormerPPConfig) -> None:
        """Initialize checkpoint-compatible encoder/decoder modules."""
        super().__init__(config)
        self.d_model = config.d_model
        self.vocab_size = config.vocab_size
        self.bos_token_id = int(config.bos_token_id)
        self.pad_token_id = int(config.pad_token_id)
        self.eos_token_id = int(config.eos_token_id)

        self.enc_embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.enc_pos_embedding = PositionalEncoding(
            config.d_model, config.dropout, max_len=config.max_position_embeddings
        )
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=config.d_model,
            nhead=config.encoder_attention_heads,
            dropout=config.dropout,
            dim_feedforward=config.dim_feedforward,
        )
        self.encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=config.encoder_layers
        )

        self.dec_embedding = (
            self.enc_embedding
            if config.share_embedding
            else nn.Embedding(config.vocab_size, config.d_model)
        )
        self.dec_pos_embedding = PositionalEncoding(
            config.d_model, config.dropout, max_len=config.max_position_embeddings
        )
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=config.d_model,
            nhead=config.decoder_attention_heads,
            dropout=config.dropout,
            dim_feedforward=config.dim_feedforward,
        )
        self.decoder = nn.TransformerDecoder(
            decoder_layer, num_layers=config.decoder_layers
        )
        self.out = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.out.weight = self.dec_embedding.weight

        self.task_embedding = None
        if config.add_task_embedding:
            self.task_embedding = nn.Embedding(6, config.d_model)

        self.task_prompt_embed = None
        if config.add_task_prompt_token_in_model:
            self.num_task_prompt_token = config.num_task_prompt_token
            self.task_prompt_embed = nn.Parameter(
                torch.empty(6, config.num_task_prompt_token, config.d_model)
            )
            nn.init.normal_(self.task_prompt_embed)
        self.tie_weights()
        self.all_tied_weights_keys = dict(self._tied_weights_keys)

    def encode(
        self,
        input_ids: Int[torch.Tensor, "batch tokens"],
        padding_mask: Bool[torch.Tensor, "batch tokens"],
        task_ids: Int[torch.Tensor, "batch"] | None = None,
    ) -> tuple[
        Float[torch.Tensor, "seq batch channels"],
        Bool[torch.Tensor, "batch seq"],
    ]:
        """Encode input token ids with optional task prompt embeddings."""
        if self.task_prompt_embed is not None:
            if task_ids is None:
                raise ValueError(
                    "task_ids are required when task prompt embeddings are enabled"
                )

            x = self.enc_embedding(input_ids)
            prompts = self.task_prompt_embed[task_ids]
            x = torch.cat([prompts, x], dim=1).permute(1, 0, 2)
            bsz = input_ids.size(0)
            prompt_mask = padding_mask.new_zeros(
                (bsz, self.num_task_prompt_token)
            ).bool()
            enc_padding_mask = torch.cat([prompt_mask, padding_mask], dim=1)
        else:
            x = self.enc_embedding(input_ids).permute(1, 0, 2)
            enc_padding_mask = padding_mask
        enc_hs = self.encoder(
            self.enc_pos_embedding(x), src_key_padding_mask=enc_padding_mask
        )
        if self.task_embedding is not None:
            if task_ids is None:
                raise ValueError(
                    "task_ids are required when task embeddings are enabled"
                )

            enc_hs = enc_hs + self.task_embedding(task_ids).unsqueeze(0)
        return enc_hs, enc_padding_mask

    def prepare_decoder_input_ids_from_labels(
        self, labels: Int[torch.Tensor, "batch tokens"]
    ) -> Int[torch.Tensor, "batch tokens"]:
        """Shift labels right and prepend BOS."""
        bos = labels.new_full((labels.size(0), 1), self.bos_token_id)
        return torch.cat([bos, labels[:, :-1]], dim=1)

    def forward(
        self,
        input_ids: Int[torch.Tensor, "batch tokens"],
        attention_mask: Bool[torch.Tensor, "batch tokens"] | None = None,
        labels: Int[torch.Tensor, "batch tokens"] | None = None,
        decoder_input_ids: Int[torch.Tensor, "batch tokens"] | None = None,
        task_ids: Int[torch.Tensor, "batch"] | None = None,
        return_dict: bool | None = None,
    ) -> Seq2SeqLMOutput | tuple[Float[torch.Tensor, "..."], ...]:
        """Run teacher-forced LayoutFormer++ decoding."""
        if attention_mask is None:
            attention_mask = input_ids.ne(self.pad_token_id)
        padding_mask = ~attention_mask.bool()
        if decoder_input_ids is None:
            if labels is None:
                raise ValueError("decoder_input_ids or labels must be provided")

            decoder_input_ids = self.prepare_decoder_input_ids_from_labels(labels)
        enc_hs, enc_padding_mask = self.encode(input_ids, padding_mask, task_ids)
        dec_input = self.dec_pos_embedding(
            self.dec_embedding(decoder_input_ids).permute(1, 0, 2)
        )
        tgt_mask = generate_square_subsequent_mask(dec_input.size(0), dec_input.device)
        y = self.decoder(
            tgt=dec_input,
            memory=enc_hs,
            tgt_mask=tgt_mask,
            memory_key_padding_mask=enc_padding_mask,
        )
        logits = self.out(y.permute(1, 0, 2))
        loss = None
        if labels is not None:
            targets = labels.clone()
            targets[targets == self.pad_token_id] = -100
            loss = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                targets.reshape(-1),
                ignore_index=-100,
            )
        if return_dict is False:
            return (logits,) if loss is None else (loss, logits)
        return Seq2SeqLMOutput(loss=cast(torch.FloatTensor | None, loss), logits=logits)

    @torch.no_grad()
    def _generate_sequences(
        self,
        input_ids: Int[torch.Tensor, "batch tokens"],
        attention_mask: Bool[torch.Tensor, "batch tokens"] | None = None,
        *,
        max_length: int | None = None,
        do_sample: bool = False,
        top_k: int = 10,
        temperature: float = 0.7,
        generation_constraint_fn: Callable[
            [int, int, Int[torch.Tensor, "tokens"]], tuple[list[int], int | None]
        ]
        | None = None,
        task_ids: Int[torch.Tensor, "batch"] | None = None,
        generator: torch.Generator | None = None,
    ) -> Int[torch.Tensor, "batch tokens"]:
        """Run the reference greedy/top-k autoregressive loop."""
        if attention_mask is None:
            attention_mask = input_ids.ne(self.pad_token_id)
        padding_mask = ~attention_mask.bool()
        max_length = max_length or self.config.decode_max_length
        enc_hs, enc_padding_mask = self.encode(input_ids, padding_mask, task_ids)
        bsz = input_ids.size(0)
        stop = input_ids.new_zeros(bsz, dtype=torch.bool)
        pred_ids = input_ids.new_full((bsz, 1), self.bos_token_id)
        outs: list[Int[torch.Tensor, "batch"]] = []
        for idx in range(max_length):
            dec_input = self.dec_pos_embedding(
                self.dec_embedding(pred_ids).permute(1, 0, 2)
            )
            tgt_mask = generate_square_subsequent_mask(idx + 1, input_ids.device)
            y = self.decoder(
                tgt=dec_input,
                memory=enc_hs,
                tgt_mask=tgt_mask,
                memory_key_padding_mask=enc_padding_mask,
            )
            logits = self.out(y.permute(1, 0, 2)[:, -1, :])
            if generation_constraint_fn is not None:
                current = (
                    torch.stack(outs, dim=1) if outs else input_ids.new_empty((bsz, 0))
                )
                for batch_idx in range(bsz):
                    allowed, _ = generation_constraint_fn(
                        batch_idx, idx, current[batch_idx]
                    )
                    mask = torch.ones(
                        logits.size(-1), dtype=torch.bool, device=logits.device
                    )
                    mask[allowed] = False
                    logits[batch_idx].masked_fill_(mask, -math.inf)
            if do_sample:
                probs = F.softmax(top_k_logits(logits / temperature, top_k), dim=-1)
                curr = torch.multinomial(
                    probs,
                    num_samples=1,
                    generator=generator,
                ).squeeze(-1)
            else:
                curr = torch.argmax(logits, dim=-1)
            eos = curr.eq(self.eos_token_id)
            curr[stop] = self.pad_token_id
            outs.append(curr)
            pred_ids = torch.cat([pred_ids, curr.unsqueeze(1)], dim=1)
            stop = torch.logical_or(stop, eos)
            if bool(torch.all(stop)):
                break
        return torch.stack(outs, dim=1)

class LayoutFormerPPPipeline(LayoutGenerationPipeline):
    """Compose a LayoutFormer++ model and processor for layout generation.

    Args:
        model: Converted LayoutFormer++ model.
        processor: Matching processor/tokenizer.
        config: Optional root pipeline config. Defaults to `model.config`.

    Examples:
        >>> processor = LayoutFormerPPProcessor.from_config(dataset="rico", task="gen_t")
        >>> config = LayoutFormerPPConfig(vocab_size=processor.tokenizer.vocab_size)
        >>> pipe = LayoutFormerPPPipeline(
        ...     model=LayoutFormerPPForConditionalGeneration(config),
        ...     processor=processor,
        ... )
        >>> pipe.config.model_type
        'layoutformerpp'
    """

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

    config: LayoutFormerPPConfig
    model: LayoutFormerPPForConditionalGeneration
    processor: LayoutFormerPPProcessor

    def __init__(
        self,
        model: LayoutFormerPPForConditionalGeneration,
        processor: LayoutFormerPPProcessor,
        config: LayoutFormerPPConfig | None = None,
    ) -> None:
        """Initialize the pipeline with model and processor components."""
        super().__init__(config or model.config)
        self.config = config or model.config
        self.model = model
        self.processor = processor

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

    @torch.no_grad()
    def __call__(
        self,
        batch_size: int = 1,
        seed: int | None = None,
        generator: torch.Generator | None = None,
        condition_type: ConditionType | str = ConditionType.unconditional,
        labels: list[list[int | str]]
        | Int[torch.Tensor, "batch elements"]
        | None = None,
        bbox: LayoutFormerPPBBoxInput = None,
        mask: Bool[torch.Tensor, "batch elements"] | None = None,
        relations: list[list[tuple[int, int, int, int, int]]]
        | Int[torch.Tensor, "batch relations relation_attrs"]
        | 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,
        max_length: int | None = None,
        do_sample: bool | None = None,
        top_k: int = 10,
        temperature: float = 0.7,
    ) -> LayoutGenerationOutput | LayoutFormerPPOutputDict:  # ty: ignore[invalid-method-override]
        """Generate layouts by encoding conditions, generating ids, and decoding.

        Args:
            batch_size: Number of layouts to generate when labels are omitted.
            seed: Convenience seed used only when `generator` is absent.
            generator: Optional PyTorch generator; takes precedence over `seed`.
            condition_type: Canonical condition type or supported alias.
            labels: Optional label conditions.
            bbox: Optional layout boxes for size/completion/refinement conditions.
            mask: Reserved public validity mask input.
            relations: Optional relation tuples for relation-conditioned checkpoints.
            num_elements: Reserved v1 interface argument.
            box_format: Input and output bounding-box format.
            normalized: Whether public boxes are normalized.
            canvas_size: Reserved v1 interface argument.
            num_inference_steps: Reserved v1 interface argument.
            output_type: Return `dataclass` or `dict`.
            return_intermediates: Reserved output detail flag.
            max_length: Optional token decode length override.
            do_sample: Optional sampling override.
            top_k: Top-k value used by the reference sampling loop.
            temperature: Sampling temperature.

        Returns:
            Layout generation output dataclass or dictionary.

        Raises:
            ValueError: If processor inputs are invalid.
        """
        _ = (num_elements, num_inference_steps, return_intermediates)
        encoded = self.processor(
            condition_type=condition_type,
            batch_size=batch_size,
            return_tensors="pt",
            labels=labels,
            bbox=bbox,
            mask=mask,
            relations=relations,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
        )
        model_device = next(self.model.parameters()).device
        input_ids = encoded["input_ids"].to(model_device)
        attention_mask = encoded["attention_mask"].to(model_device)
        condition = self.processor.normalize_condition_type(condition_type)
        generation_generator = self.prepare_generator(
            generator=generator,
            seed=seed,
            device=model_device,
        )
        default_do_sample = condition in {
            ConditionType.unconditional,
            ConditionType.completion,
        }
        sequences = self.model._generate_sequences(
            input_ids,
            attention_mask,
            max_length=max_length,
            do_sample=default_do_sample if do_sample is None else do_sample,
            top_k=top_k,
            temperature=temperature,
            generator=generation_generator,
        )
        return self.processor.post_process_layouts(
            sequences.cpu(),
            box_format=box_format,
            output_type=output_type,
        )