Layoutdiffusion

class LayoutDiffusionConfig(ConfigMixin):
    """Serializable LayoutDiffusion model, tokenizer, and scheduler settings.

    Args:
        dataset_name: Dataset name or alias.
        id2label: Optional persisted dataset-local label mapping.
        vocab: Optional token-to-id vocabulary loaded from ``vocab.json``.
        seq_length: Full internal token sequence length.
        max_num_elements: Maximum number of layout elements.
        num_coordinate_bins: Number of coordinate tokens.
        diffusion_steps: Number of training diffusion timesteps.
        noise_schedule: Reference diffusion schedule name.
        num_channels: OpenAI timestep embedding dimension.
        bert_config_name: Name of the BERT config used by the checkpoint.
        max_position_embeddings: BERT position embedding count.
        hidden_size: Transformer hidden size.
        num_hidden_layers: BERT encoder layer count.
        num_attention_heads: Attention head count.
        intermediate_size: Feed-forward hidden size.
        dropout: Dropout probability.
        training_mode: Reference training mode.
        vocab_size: Full vocabulary size including mask.
        refine_start_step: Dataset-specific refinement start step.
        type_start_step: Reference type-conditioned start step.
        element_count_prior: Optional 20-entry unconditional element count prior.
        pow_num: Gaussian transition exponent.
        mul_num: Gaussian transition multiplier.

    Examples:
        >>> cfg = LayoutDiffusionConfig(dataset_name="publaynet")
        >>> cfg.mask_token_id == cfg.vocab_size - 1
        True
    """

    config_name = "layoutdiffusion_config.json"

    @register_to_config
    def __init__(
        self,
        *,
        dataset_name: DatasetName | str = DatasetName.rico25,
        id2label: dict[int | str, str] | None = None,
        vocab: dict[str, int] | None = None,
        seq_length: int = 121,
        max_num_elements: int = 20,
        num_coordinate_bins: int = 128,
        diffusion_steps: int = 200,
        noise_schedule: str = "gaussian_refine_pow2.5",
        num_channels: int = 128,
        bert_config_name: str = "bert-base-uncased",
        max_position_embeddings: int = 512,
        hidden_size: int = 768,
        num_hidden_layers: int = 12,
        num_attention_heads: int = 12,
        intermediate_size: int = 3072,
        dropout: float = 0.1,
        training_mode: str = "discrete",
        vocab_size: int | None = None,
        refine_start_step: int | None = None,
        type_start_step: int = 160,
        element_count_prior: list[float] | None = None,
        pow_num: float = 2.5,
        mul_num: float = 12.4,
    ) -> None:
        """Initialize LayoutDiffusion configuration."""
        self.dataset_name = str(normalize_dataset_name(dataset_name))
        raw_id2label = id2label or default_id2label(self.dataset_name)
        self.id2label = {int(k): v for k, v in raw_id2label.items()}

        self.vocab = vocab or self.default_vocab()
        self.seq_length = seq_length
        self.max_num_elements = max_num_elements
        self.num_coordinate_bins = num_coordinate_bins

        self.diffusion_steps = diffusion_steps
        self.noise_schedule = noise_schedule
        self.num_channels = num_channels
        self.bert_config_name = bert_config_name

        self.max_position_embeddings = max_position_embeddings
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.dropout = dropout

        self.training_mode = training_mode
        self.vocab_size = vocab_size or len(self.vocab)
        self.refine_start_step = refine_start_step or (
            60 if self.dataset_name == str(DatasetName.publaynet) else 50
        )
        self.type_start_step = type_start_step
        self.element_count_prior = element_count_prior or self.default_element_prior()
        self.pow_num = pow_num
        self.mul_num = mul_num

    @property
    def special_token_ids(self) -> dict[str, int]:
        """Return LayoutDiffusion special-token ids."""
        return {
            "START": self.vocab["START"],
            "END": self.vocab["END"],
            "UNK": self.vocab["UNK"],
            "PAD": self.vocab["PAD"],
            "|": self.vocab["|"],
        }

    @property
    def pad_token_id(self) -> int:
        """Return the padding token id."""
        return self.vocab["PAD"]

    @property
    def mask_token_id(self) -> int:
        """Return the mask token id."""
        return self.vocab_size - 1

    @property
    def label_token_offset(self) -> int:
        """Return the first label-token id."""
        return 5

    @property
    def num_labels(self) -> int:
        """Return the dataset label count."""
        return len(self.id2label)

    @property
    def coordinate_token_offset(self) -> int:
        """Return the first coordinate-token id."""
        return self.label_token_offset + self.num_labels

    @property
    def max_token_length(self) -> int:
        """Return the full LayoutDiffusion token sequence length."""
        return self.seq_length

    @property
    def label2id(self) -> dict[str, int]:
        """Return inverse public label mapping."""
        return {v: k for k, v in self.id2label.items()}

    @property
    def type_classes(self) -> int:
        """Return the number of reference type classes."""
        return self.vocab_size - 1 - self.num_coordinate_bins - 5

    def default_vocab(self) -> dict[str, int]:
        """Build the default LayoutDiffusion vocabulary for the dataset."""
        vocab = {"START": 0, "END": 1, "UNK": 2, "PAD": 3, "|": 4}
        for label in default_id2label(self.dataset_name).values():
            vocab[label] = len(vocab)
        for coord in range(self.num_coordinate_bins):
            vocab[str(coord)] = len(vocab)
        vocab["MASK"] = len(vocab)
        return vocab

    def default_element_prior(self) -> list[float]:
        """Return the reference unconditional element-count prior."""
        if self.dataset_name == str(DatasetName.publaynet):
            return [
                0.00321776,
                0.03342678,
                0.04233181,
                0.04218409,
                0.05404355,
                0.07231605,
                0.08247029,
                0.0905211,
                0.0949399,
                0.0959322,
                0.08953522,
                0.07810608,
                0.0619627,
                0.04775897,
                0.03585776,
                0.0261788,
                0.018812,
                0.01404317,
                0.00972071,
                0.00664104,
            ]
        return [
            0.04849498,
            0.03704171,
            0.0534486,
            0.06045308,
            0.06354515,
            0.07585032,
            0.08045687,
            0.0644917,
            0.05676153,
            0.05742412,
            0.05471067,
            0.04944153,
            0.04552912,
            0.04190068,
            0.04426705,
            0.0387455,
            0.03533792,
            0.03167792,
            0.02997413,
            0.0304474,
        ]

class LayoutDiffusionTransformer(ModelMixin, ConfigMixin):
    """BERT-encoder denoiser for LayoutDiffusion token sequences.

    Args:
        vocab_size: Full tokenizer vocabulary size including mask.
        num_channels: OpenAI timestep embedding dimension.
        hidden_size: BERT hidden size.
        num_hidden_layers: Number of BERT encoder layers.
        num_attention_heads: Number of attention heads.
        intermediate_size: BERT feed-forward size.
        dropout: Hidden dropout probability.
        max_position_embeddings: Position embedding count.
        constrained: Optional reference constraint mode.

    Examples:
        >>> model = LayoutDiffusionTransformer(
        ...     vocab_size=16, hidden_size=32, num_channels=8,
        ...     num_hidden_layers=1, num_attention_heads=4, intermediate_size=64,
        ... )
        >>> out = model(torch.zeros(2, 5, dtype=torch.long), torch.zeros(2, dtype=torch.long))
        >>> out.logits.shape
        torch.Size([2, 15, 5])
    """

    config_name = "transformer_config.json"

    position_ids: Int[torch.Tensor, "1 max_positions"]

    @register_to_config
    def __init__(
        self,
        *,
        vocab_size: int,
        num_channels: int = 128,
        hidden_size: int = 768,
        num_hidden_layers: int = 12,
        num_attention_heads: int = 12,
        intermediate_size: int = 3072,
        dropout: float = 0.1,
        max_position_embeddings: int = 512,
        constrained: str | None = None,
    ) -> None:
        """Initialize the transformer."""
        super().__init__()
        config = BertConfig(
            hidden_size=hidden_size,
            num_hidden_layers=num_hidden_layers,
            num_attention_heads=num_attention_heads,
            intermediate_size=intermediate_size,
            hidden_dropout_prob=dropout,
            attention_probs_dropout_prob=dropout,
            max_position_embeddings=max_position_embeddings,
        )
        self.constrained = constrained
        self.in_channels = 768
        self.model_channels = num_channels
        self.out_channels = vocab_size - 1
        self.word_embedding = nn.Embedding(vocab_size, self.in_channels)
        time_embed_dim = num_channels * 4
        self.time_embed = nn.Sequential(
            nn.Linear(num_channels, time_embed_dim),
            nn.SiLU(),
            nn.Linear(time_embed_dim, hidden_size),
        )
        self.input_up_proj = nn.Sequential(
            nn.Linear(self.in_channels, hidden_size),
            nn.Tanh(),
            nn.Linear(hidden_size, hidden_size),
        )
        self.input_transformers = BertEncoder(config)
        self.register_buffer(
            "position_ids",
            torch.arange(max_position_embeddings).expand((1, -1)),
            persistent=False,
        )
        self.position_embeddings = nn.Embedding(max_position_embeddings, hidden_size)
        self.LayerNorm = nn.LayerNorm(hidden_size, eps=config.layer_norm_eps)
        self.dropout_layer = nn.Dropout(dropout)
        self.output_down_proj = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.Tanh(),
            nn.Linear(hidden_size, self.out_channels),
        )

    def get_embeds(
        self, input_ids: Int[torch.Tensor, "batch tokens"]
    ) -> Float[torch.Tensor, "batch tokens channels"]:
        """Return token embeddings for parity diagnostics."""
        return self.word_embedding(input_ids)

    def forward(
        self,
        input_ids: Int[torch.Tensor, "batch tokens"],
        timesteps: Int[torch.Tensor, "batch"],
        condition_ids: Int[torch.Tensor, "batch tokens"] | None = None,
        condition_type: str | None = None,
        return_dict: bool = True,
    ) -> (
        LayoutDiffusionTransformerOutput
        | tuple[Float[torch.Tensor, "batch vocab tokens"]]
    ):
        """Predict start-token logits for a reverse diffusion step.

        Args:
            input_ids: Current token ids shaped ``(B, L)``.
            timesteps: Diffusion timestep per batch item.
            condition_ids: Optional internal condition token ids.
            condition_type: Optional condition mode.
            return_dict: Whether to return a dataclass output.

        Returns:
            Logits shaped ``(B, vocab_size - 1, L)``.
        """
        x = input_ids
        if condition_ids is not None and condition_type == "label":
            mask = condition_ids.le(self.out_channels - 129).unsqueeze(-1)
            hidden = self.word_embedding(condition_ids) * mask + self.word_embedding(
                x
            ) * (~mask)
        elif condition_ids is not None and condition_type == "completion":
            keep = torch.tensor(
                [1] * 6 + [0] * (condition_ids.shape[1] - 6),
                device=x.device,
                dtype=torch.bool,
            ).expand(condition_ids.shape[0], -1)
            hidden = self.word_embedding(condition_ids) * keep.unsqueeze(
                -1
            ) + self.word_embedding(x) * (~keep).unsqueeze(-1)
        else:
            hidden = self.word_embedding(x)
        emb = self.time_embed(
            get_timestep_embedding(
                timesteps,
                self.model_channels,
                flip_sin_to_cos=True,
                downscale_freq_shift=0,
            ).to(hidden)
        )
        seq_length = hidden.size(1)
        position_ids = self.position_ids[:, :seq_length]
        inputs = self.input_up_proj(hidden)
        inputs = inputs + self.position_embeddings(position_ids) + emb.unsqueeze(1)
        inputs = self.dropout_layer(self.LayerNorm(inputs))
        encoded = self.input_transformers(inputs).last_hidden_state
        logits = rearrange(self.output_down_proj(encoded), "b l c -> b c l").type(
            hidden.dtype
        )
        if not return_dict:
            return (logits,)
        return LayoutDiffusionTransformerOutput(logits=logits)

class LayoutDiffusionPipeline(DiffusionPipeline):
    """Generate layouts with a converted LayoutDiffusion pipeline.

    Args:
        transformer: LayoutDiffusion transformer denoiser.
        scheduler: Categorical diffusion scheduler.
        tokenizer: LayoutDiffusion layout tokenizer.
        processor: Optional processor.

    Examples:
        >>> from layoutdiffusion import LayoutDiffusionConfig, LayoutDiffusionTokenizer
        >>> from layoutdiffusion import LayoutDiffusionScheduler, LayoutDiffusionTransformer
        >>> from layoutdiffusion.sampling import LayoutDiffusionSamplingConfig
        >>> cfg = LayoutDiffusionConfig(dataset_name="publaynet", hidden_size=32, num_hidden_layers=1, num_attention_heads=4, intermediate_size=64, num_channels=8)
        >>> tok = LayoutDiffusionTokenizer(cfg)
        >>> pipe = LayoutDiffusionPipeline(
        ...     LayoutDiffusionTransformer(vocab_size=cfg.vocab_size, hidden_size=32, num_hidden_layers=1, num_attention_heads=4, intermediate_size=64, num_channels=8),
        ...     LayoutDiffusionScheduler(vocab_size=cfg.vocab_size, mask_token_id=cfg.mask_token_id, type_classes=cfg.type_classes, num_train_timesteps=2),
        ...     tok,
        ... )
        >>> pipe(batch_size=1, seed=0, sampling=LayoutDiffusionSamplingConfig(num_inference_steps=1)).bbox.shape[-1]
        4
    """

    model_cpu_offload_seq = "transformer"

    def __init__(
        self,
        transformer: LayoutDiffusionTransformer,
        scheduler: LayoutDiffusionScheduler,
        tokenizer: LayoutDiffusionTokenizer,
        processor: LayoutDiffusionProcessor | None = None,
    ) -> None:
        """Initialize and register pipeline modules."""
        super().__init__()
        self.register_modules(
            transformer=transformer,
            scheduler=scheduler,
            tokenizer=tokenizer,
        )
        self.tokenizer = tokenizer
        self.processor = processor or LayoutDiffusionProcessor(tokenizer)
        self.transformer.eval()

    @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: Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Float[np.ndarray, "batch elements 4"]
        | Sequence[ArrayLikeInput]
        | None = None,
        mask: Bool[torch.Tensor, "batch elements"]
        | Bool[np.ndarray, "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: Literal["dataclass", "dict"] = "dataclass",
        return_intermediates: bool = False,
        sampling: LayoutDiffusionSamplingConfig,
        **model_kwargs: str | int | float | bool | None,
    ) -> LayoutGenerationOutput | LayoutDiffusionOutputDict:
        """Run LayoutDiffusion generation.

        Args:
            batch_size: Number of layouts for unconditional generation.
            seed: Seed used only when ``generator`` is omitted.
            generator: Optional torch generator. Takes precedence over ``seed``.
            condition_type: Canonical condition type or supported alias.
            labels: Optional conditional labels.
            bbox: Optional conditional boxes.
            mask: Optional conditional valid mask.
            num_elements: Optional element counts.
            box_format: Input box format.
            normalized: Whether conditional boxes are normalized.
            canvas_size: Pixel canvas size for unnormalized inputs.
            num_inference_steps: Optional shortened inference steps.
            output_type: ``"dataclass"`` or ``"dict"``.
            return_intermediates: Whether to include trajectories.
            sampling: Sampling config.
            **model_kwargs: Reserved compatibility kwargs.

        Returns:
            Layout output dataclass or dictionary.

        Raises:
            ValueError: If ``output_type`` is unsupported.
        """
        _ = model_kwargs
        if generator is None and seed is not None:
            generator = torch.Generator(device=self.device).manual_seed(seed)
        canonical = normalize_condition_type(condition_type)
        processed = self.processor(
            bbox=bbox,
            labels=labels,
            mask=mask,
            num_elements=num_elements,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
        )
        condition_input = processed.get("input_ids")
        processed_labels = None if labels is None else torch.as_tensor(labels)
        counts = processed.get("num_elements")
        condition = build_condition(
            self.tokenizer,
            condition_type=canonical,
            input_ids=condition_input,
            labels=processed_labels,
            num_elements=counts,
        )
        if condition is not None and condition.input_ids is not None:
            batch_size = condition.input_ids.shape[0]
        if condition is not None and canonical is ConditionType.refinement:
            if condition.input_ids is None:
                raise ValueError("refinement condition is missing input_ids")

            start_ids = condition.input_ids.to(self.device)
        else:
            start_ids = self.tokenizer.build_initial_tokens(
                batch_size=batch_size,
                num_elements=counts,
                labels=processed_labels,
                condition_type=str(canonical),
                generator=generator,
                device=self.device,
            )
        sample = index_to_log_onehot(start_ids, self.scheduler.config.vocab_size)
        start_step = None if condition is None else condition.start_step
        self.scheduler.set_timesteps(
            sampling.num_inference_steps or num_inference_steps,
            start_step=start_step,
            device=self.device,
        )
        trajectory = [] if return_intermediates else None
        for timestep in self.scheduler.timesteps:
            timestep_batch = torch.full(
                (batch_size,),
                int(timestep.item()),
                device=self.device,
                dtype=torch.long,
            )
            input_ids = log_onehot_to_index(sample)
            logits = self.transformer(
                input_ids=input_ids,
                timesteps=timestep_batch,
                condition_ids=None
                if condition is None or condition.input_ids is None
                else condition.input_ids.to(self.device),
                condition_type=None if condition is None else str(condition.type),
            ).logits
            out = self.scheduler.step(
                logits,
                timestep_batch,
                sample,
                sampling=sampling,
                condition=condition,
                generator=generator,
            )
            sample = out.prev_sample
            if trajectory is not None:
                trajectory.append(log_onehot_to_index(sample).detach().cpu())
        output = self._decode_final_sample(
            sample=sample,
            trajectory=trajectory,
            condition_type=str(canonical),
            return_intermediates=return_intermediates,
        )
        return self._coerce_output(output=output, output_type=output_type)

    def _decode_final_sample(
        self,
        *,
        sample: Float[torch.Tensor, "batch vocab tokens"],
        trajectory: list[Int[torch.Tensor, "batch tokens"]] | None,
        condition_type: str,
        return_intermediates: bool,
    ) -> LayoutGenerationOutput:
        """Decode final token logits into the public layout output schema."""
        sequences = log_onehot_to_index(sample).detach().cpu()
        layout = self.tokenizer.decode_layout(sequences)
        metadata = {"condition_type": condition_type} if return_intermediates else None
        return LayoutGenerationOutput(
            bbox=layout["bbox"],
            labels=layout["labels"],
            mask=layout["mask"],
            id2label=self.tokenizer.config.id2label,
            sequences=sequences if return_intermediates else None,
            trajectory=trajectory,
            intermediates=metadata,
        )

    @staticmethod
    def _coerce_output(
        *,
        output: LayoutGenerationOutput,
        output_type: Literal["dataclass", "dict"],
    ) -> LayoutGenerationOutput | LayoutDiffusionOutputDict:
        """Return the requested output container."""
        if output_type == "dict":
            return cast(LayoutDiffusionOutputDict, dict(output))
        if output_type != "dataclass":
            raise ValueError(f"Unsupported output_type: {output_type}")

        return output

    generate = __call__

    def save_pretrained(
        self, save_directory: str | Path, **kwargs: LayoutDiffusionPipelineKwarg
    ) -> None:
        """Save a Diffusers pipeline directory."""
        super().save_pretrained(save_directory, **kwargs)

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str | Path,
        **kwargs: LayoutDiffusionPipelineKwarg,
    ) -> "LayoutDiffusionPipeline":
        """Load a LayoutDiffusion pipeline and rebuild its processor."""
        tokenizer = kwargs.pop("tokenizer", None)
        if tokenizer is None:
            tokenizer = LayoutDiffusionTokenizer.from_pretrained(
                pretrained_model_name_or_path
            )
        pipe = super().from_pretrained(
            pretrained_model_name_or_path,
            tokenizer=tokenizer,
            **kwargs,
        )
        pipe.processor = LayoutDiffusionProcessor(pipe.tokenizer)
        return pipe