Layout dm

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 LayoutDMCondition:
    """Strong and weak token constraints for conditional LayoutDM sampling."""

    input_ids: Int[torch.Tensor, "batch tokens"]
    mask: Bool[torch.Tensor, "batch tokens"]
    type: Literal["c", "cwh", "partial", "refinement"]
    num_element: Int[torch.Tensor, "batch"] | None = None
    original_input_ids: Int[torch.Tensor, "batch tokens"] | None = None
    weak_mask: Bool[torch.Tensor, "batch tokens"] | None = None
    weak_logits: Float[torch.Tensor, "batch tokens vocab"] | None = None

class LayoutDMConfig(ConfigMixin):
    """Serializable LayoutDM architecture and tokenizer configuration.

    Args:
        dataset_name: Dataset name or alias used to initialize labels.
        id2label: Optional persisted label-id mapping.
        max_seq_length: Maximum number of layout elements.
        num_bin_bboxes: Number of bins per bounding-box attribute.
        var_order: Per-element token order.
        shared_bbox_vocab: Bounding-box vocabulary sharing mode.
        bbox_quantization: Bounding-box quantization mode.
        special_tokens: Special token names. ``mask`` must be last for LayoutDM.
        cluster_centers: Optional bbox cluster centers stored with tokenizer files.
        cluster_centers_path: Optional local path to released cluster centers.
        hidden_size: Transformer hidden size.
        num_attention_heads: Number of attention heads.
        num_hidden_layers: Number of transformer layers.
        intermediate_size: Feed-forward hidden size.
        dropout: Transformer dropout probability.
        timestep_type: Timestep-conditioning type.
        num_timesteps: Number of diffusion timesteps.
        q_type: Diffusion transition type.
        att_1: Initial keep probability schedule value.
        att_T: Final keep probability schedule value.
        ctt_1: Initial mask probability schedule value.
        ctt_T: Final mask probability schedule value.

    Examples:
        >>> cfg = LayoutDMConfig(dataset_name="publaynet")
        >>> cfg.vocab_size > cfg.num_categories
        True
    """

    config_name = "layout_dm_config.json"

    @register_to_config
    def __init__(
        self,
        *,
        dataset_name: DatasetName | str,
        id2label: dict[int | str, str] | None = None,
        max_seq_length: int = 25,
        num_bin_bboxes: int = 32,
        var_order: str = "c-x-y-w-h",
        shared_bbox_vocab: str = "x-y-w-h",
        bbox_quantization: str = "kmeans",
        special_tokens: tuple[str, ...] = ("pad", "mask"),
        cluster_centers: dict[str, list[float]] | None = None,
        cluster_centers_path: str | None = None,
        hidden_size: int = 464,
        num_attention_heads: int = 8,
        num_hidden_layers: int = 4,
        intermediate_size: int = 1856,
        dropout: float = 0.0,
        timestep_type: str | None = "adalayernorm",
        num_timesteps: int = 100,
        q_type: str = "constrained",
        att_1: float = 0.99999,
        att_T: float = 0.000009,
        ctt_1: float = 0.000009,
        ctt_T: float = 0.99999,
    ) -> None:
        """Initialize a serializable LayoutDM configuration."""
        self.dataset_name = str(normalize_dataset_name(dataset_name))
        raw_id2label = id2label or id2label_for_dataset(self.dataset_name)
        self.id2label = {int(k): v for k, v in raw_id2label.items()}

        self.max_seq_length = max_seq_length
        self.num_bin_bboxes = num_bin_bboxes
        self.var_order = var_order
        self.shared_bbox_vocab = shared_bbox_vocab
        self.bbox_quantization = bbox_quantization
        self.special_tokens = tuple(special_tokens)
        self.cluster_centers = cluster_centers
        self.cluster_centers_path = cluster_centers_path

        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.num_hidden_layers = num_hidden_layers
        self.intermediate_size = intermediate_size
        self.dropout = dropout
        self.timestep_type = timestep_type
        self.num_timesteps = num_timesteps
        self.q_type = q_type

        self.att_1 = att_1
        self.att_T = att_T
        self.ctt_1 = ctt_1
        self.ctt_T = ctt_T

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

    @property
    def num_categories(self) -> int:
        """Return the number of dataset categories."""
        return len(self.id2label)

    @property
    def num_bbox_tokens(self) -> int:
        """Return the number of bounding-box vocabulary tokens."""
        return self.num_bin_bboxes * len(self.shared_bbox_vocab.split("-"))

    @property
    def num_special_tokens(self) -> int:
        """Return the number of special tokens."""
        return len(self.special_tokens)

    @property
    def vocab_size(self) -> int:
        """Return the full tokenizer vocabulary size."""
        return self.num_categories + self.num_bbox_tokens + self.num_special_tokens

    @property
    def pad_token_id(self) -> int:
        """Return the full vocabulary id of the padding token."""
        return (
            self.num_categories
            + self.num_bbox_tokens
            + self.special_tokens.index("pad")
        )

    @property
    def mask_token_id(self) -> int:
        """Return the full vocabulary id of the mask token."""
        return (
            self.num_categories
            + self.num_bbox_tokens
            + self.special_tokens.index("mask")
        )

    @property
    def num_attributes_per_element(self) -> int:
        """Return the number of tokens used for each layout element."""
        return len(self.var_order.split("-"))

    @property
    def max_token_length(self) -> int:
        """Return the flattened token sequence length."""
        return self.max_seq_length * self.num_attributes_per_element

    @property
    def bbox_slices(self) -> dict[str, tuple[int, int]]:
        """Return full-vocabulary slices for bbox attributes."""
        slices: dict[str, tuple[int, int]] = {}
        for i, key in enumerate(("x", "y", "w", "h")):
            start = self.num_categories + i * self.num_bin_bboxes
            slices[key] = (start, start + self.num_bin_bboxes)
        return slices

class LayoutDMDenoiser(ModelMixin, ConfigMixin):
    """Diffusers-compatible LayoutDM denoiser.

    Args:
        vocab_size: Size of the LayoutDM tokenizer vocabulary.
        max_token_length: Flattened token sequence length.
        hidden_size: Transformer hidden size.
        num_attention_heads: Number of attention heads.
        num_hidden_layers: Number of transformer layers.
        intermediate_size: Feed-forward hidden size.
        dropout: Dropout probability.
        timestep_type: Timestep-conditioning type.

    Examples:
        >>> model = LayoutDMDenoiser(vocab_size=10, max_token_length=5, hidden_size=8,
        ...     num_attention_heads=2, num_hidden_layers=1, intermediate_size=16)
        >>> model.config.vocab_size
        10
    """

    config_name = "denoiser_config.json"

    @register_to_config
    def __init__(
        self,
        *,
        vocab_size: int,
        max_token_length: int,
        hidden_size: int = 464,
        num_attention_heads: int = 8,
        num_hidden_layers: int = 4,
        intermediate_size: int = 1856,
        dropout: float = 0.0,
        timestep_type: Literal["adalayernorm", "adalayernorm_abs"]
        | None = "adalayernorm",
    ) -> None:
        """Initialize the categorical transformer denoiser."""
        super().__init__()
        self.transformer = CategoricalTransformer(
            vocab_size=vocab_size,
            max_token_length=max_token_length,
            hidden_size=hidden_size,
            num_attention_heads=num_attention_heads,
            num_hidden_layers=num_hidden_layers,
            intermediate_size=intermediate_size,
            dropout=dropout,
            timestep_type=timestep_type,
        )
        self.apply(_init_layoutdm_weights)

    def forward(
        self,
        input_ids: Int[torch.Tensor, "batch tokens"],
        timesteps: Int[torch.Tensor, "batch"],
    ) -> LayoutDMDenoiserOutput:
        """Predict token logits for noised LayoutDM sequences."""
        return LayoutDMDenoiserOutput(
            logits=self.transformer(input_ids, timestep=timesteps)["logits"]
        )

    def predict_start_log_probs(
        self,
        input_ids: Int[torch.Tensor, "batch tokens"],
        timesteps: Int[torch.Tensor, "batch"],
    ) -> Float[torch.Tensor, "batch tokens vocab"]:
        """Predict log probabilities for the denoised start sequence."""
        logits = self(input_ids=input_ids, timesteps=timesteps).logits[:, :, :-1]
        log_pred = F.log_softmax(logits.double(), dim=-1).float()
        zero_mask = torch.full(
            (*log_pred.shape[:2], 1),
            -70.0,
            device=log_pred.device,
            dtype=log_pred.dtype,
        )
        return torch.cat((log_pred, zero_mask), dim=-1).clamp(-70.0, 0.0)

@dataclass
class LayoutDMDenoiserOutput(BaseOutput):
    """Denoiser output containing token logits."""

    logits: Float[torch.Tensor, "batch tokens vocab"]

class LayoutDMPipeline(DiffusionPipeline):
    """Generate layouts with a converted LayoutDM denoiser and scheduler.

    Args:
        denoiser: LayoutDM denoiser model.
        scheduler: Discrete diffusion scheduler.
        tokenizer: Structured layout tokenizer.
        processor: Optional input processor. A default processor is created
            when omitted.

    Examples:
        >>> from collections.abc import Mapping, Sequence
    from pathlib import Path
        >>> path = Path(".cache/layout-dm/converted/layoutdm-rico25")
        >>> path.exists()  # doctest: +SKIP
        True
        >>> pipe = LayoutDMPipeline.from_pretrained(path)  # doctest: +SKIP
        >>> out = pipe(batch_size=1, seed=0, num_inference_steps=1)  # doctest: +SKIP
        >>> out.bbox.shape[-1]  # doctest: +SKIP
        4
    """

    model_cpu_offload_seq = "denoiser"

    def __init__(
        self,
        denoiser: LayoutDMDenoiser,
        scheduler: LayoutDMScheduler,
        tokenizer: LayoutDMTokenizer,
        processor: LayoutDMProcessor | None = None,
    ) -> None:
        """Initialize and register LayoutDM pipeline modules."""
        super().__init__()
        self.register_modules(
            denoiser=denoiser, scheduler=scheduler, tokenizer=tokenizer
        )
        self.tokenizer = tokenizer
        self.processor = processor or LayoutDMProcessor(tokenizer)
        self.denoiser.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,
        sampling: SamplingMode | str = SamplingMode.random,
        temperature: float = 1.0,
        top_k: int = 5,
        top_p: float = 0.9,
        output_type: Literal["dataclass", "dict"] = "dataclass",
        return_intermediates: bool = False,
        **model_kwargs: str | int | float | bool | None,
    ) -> LayoutGenerationOutput | dict[str, Shaped[torch.Tensor, "..."]]:
        """Run unconditional or conditional layout generation.

        Args:
            batch_size: Number of layouts generated for unconditional sampling.
            seed: Optional seed used only when ``generator`` is omitted.
            generator: Optional torch generator. Takes precedence over ``seed``.
            condition_type: Canonical condition type or supported source alias.
            labels: Optional labels used by conditional modes.
            bbox: Optional boxes used by conditional modes.
            mask: Optional valid-element mask for conditional inputs.
            num_elements: Reserved compatibility argument.
            box_format: Format of conditional input boxes.
            normalized: Whether conditional boxes are already normalized.
            canvas_size: Pixel canvas size used when ``normalized=False``.
            num_inference_steps: Optional shortened diffusion step count.
            sampling: Sampling strategy.
            temperature: Random sampling temperature.
            top_k: Top-k value for top-k modes.
            top_p: Top-p value for top-p modes.
            output_type: ``"dataclass"`` or ``"dict"``.
            return_intermediates: Whether to return sampling trajectory data.
            **model_kwargs: Reserved compatibility keyword arguments.

        Returns:
            ``LayoutGenerationOutput`` by default, or a dictionary when
            ``output_type="dict"``.

        Raises:
            ValueError: If a conditional mode is missing ``bbox`` or ``labels``,
                or if ``output_type`` is unsupported.
        """
        _ = (num_elements, 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)
        condition = None
        if canonical is not ConditionType.unconditional:
            missing_inputs = [
                name
                for name, value in (("bbox", bbox), ("labels", labels))
                if value is None
            ]
            if missing_inputs:
                message = (
                    f"bbox and labels are required for condition_type={condition_type}"
                )
                raise ValueError(message)

            processor_inputs = {
                "bbox": bbox,
                "labels": labels,
                "mask": mask,
                "box_format": box_format,
                "normalized": normalized,
                "canvas_size": canvas_size,
            }
            processed = self.processor(**processor_inputs)
            decoded_input = self.tokenizer.decode_layout(processed["input_ids"])
            condition = build_condition(
                self.tokenizer,
                cond_type=canonical,
                bbox=decoded_input["bbox"],
                labels=decoded_input["labels"],
                mask=decoded_input["mask"],
            )
            batch_size = condition.input_ids.shape[0]
        sampling_config = LayoutDMSamplingConfig(
            name=sampling,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            num_inference_steps=num_inference_steps,
        )
        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
        sample = self.scheduler.initial_sample(
            batch_size,
            self.tokenizer.config.max_token_length,
            device=self.device,
            condition=condition,
        )
        trajectory = [] if return_intermediates else None
        previous_timestep = self.scheduler.config.num_timesteps
        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.denoiser(input_ids=input_ids, timesteps=timestep_batch).logits
            out = self.scheduler.step(
                logits,
                timestep_batch,
                sample,
                previous_timestep=previous_timestep,
                sampling=sampling_config,
                condition=condition,
                generator=generator,
            )
            sample = out.prev_sample
            previous_timestep = int(timestep.item())
            if trajectory is not None:
                trajectory.append(log_onehot_to_index(sample).detach().cpu())
        sequences = log_onehot_to_index(sample).detach().cpu()
        decoded = self.tokenizer.decode_layout(sequences)
        output = LayoutGenerationOutput(
            bbox=decoded["bbox"],
            labels=decoded["labels"],
            mask=decoded["mask"],
            id2label=self.tokenizer.config.id2label,
            sequences=sequences,
            trajectory=trajectory,
            intermediates={"condition_type": str(canonical)}
            if return_intermediates
            else None,
        )
        if output_type == "dict":
            return 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: LayoutDMPipelineKwarg
    ) -> None:
        """Save the pipeline and tokenizer to a Diffusers directory."""
        super().save_pretrained(save_directory, **kwargs)

    @classmethod
    def from_pretrained(
        cls, pretrained_model_name_or_path: str | Path, **kwargs: LayoutDMPipelineKwarg
    ) -> "LayoutDMPipeline":
        """Load a LayoutDM pipeline from a local directory or Hub repo.

        Args:
            pretrained_model_name_or_path: Diffusers pipeline directory or Hub id.
            **kwargs: Additional arguments forwarded to Diffusers.

        Returns:
            Loaded pipeline with a matching ``LayoutDMProcessor``.
        """
        tokenizer = LayoutDMTokenizer.from_pretrained(pretrained_model_name_or_path)
        kwargs.setdefault("tokenizer", tokenizer)
        pipe = super().from_pretrained(pretrained_model_name_or_path, **kwargs)
        pipe.processor = LayoutDMProcessor(pipe.tokenizer)
        return pipe

class LayoutDMProcessor(ProcessorMixin):
    """Normalize layout arrays and encode them with ``LayoutDMTokenizer``.

    Args:
        tokenizer: Tokenizer used to encode processed layouts.

    Examples:
        >>> from layout_dm.configuration_layout_dm import LayoutDMConfig
        >>> from layout_dm.tokenization_layout_dm import LayoutDMTokenizer
        >>> processor = LayoutDMProcessor(LayoutDMTokenizer(LayoutDMConfig(dataset_name="publaynet")))
        >>> sorted(processor(bbox=[[[0.5, 0.5, 0.2, 0.2]]], labels=[[0]]))
        ['attention_mask', 'input_ids', 'mask']
    """

    config_name = "processor_config.json"
    tokenizer_class = "LayoutDMTokenizer"

    def __init__(self, tokenizer: LayoutDMTokenizer) -> None:
        """Initialize the processor with a tokenizer."""
        super().__init__(tokenizer=tokenizer)

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str | PathLike[str],
        cache_dir: str | PathLike[str] | None = None,
        force_download: bool = False,
        local_files_only: bool = False,
        token: str | bool | None = None,
        revision: str = "main",
        **kwargs: str | int | float | bool | None,
    ) -> "LayoutDMProcessor":
        """Load a processor with the LayoutDM tokenizer implementation."""
        tokenizer = LayoutDMTokenizer.from_pretrained(
            pretrained_model_name_or_path,
            cache_dir=cache_dir,
            force_download=force_download,
            local_files_only=local_files_only,
            token=token,
            revision=revision,
            **kwargs,
        )
        return cls(tokenizer=tokenizer)

    def __call__(
        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,
        return_tensors: Literal["pt"] = "pt",
    ) -> dict[str, Shaped[torch.Tensor, "..."]]:
        """Process a layout batch into model input tensors.

        Args:
            bbox: Layout boxes in ``box_format``.
            labels: Integer labels matching the layout boxes.
            mask: Optional valid-element mask. All elements are valid when omitted.
            box_format: Input box format.
            normalized: Whether boxes are already normalized to ``[0, 1]``.
            canvas_size: Pixel canvas size required when ``normalized=False``.
            return_tensors: Tensor backend. Only ``"pt"`` is supported.

        Returns:
            Tokenizer output containing ``input_ids``, ``attention_mask``, and
            ``mask`` tensors.

        Raises:
            ValueError: If ``return_tensors`` is not ``"pt"`` or if
                ``canvas_size`` is missing for pixel-space boxes.
        """
        if return_tensors != "pt":
            raise ValueError("LayoutDMProcessor only supports return_tensors='pt'")

        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,
        )
        return self.tokenizer.encode_layout(bbox=bbox_t, labels=labels_t, mask=mask_t)

@dataclass
class LayoutDMSamplingConfig:
    """Sampling parameters passed from the pipeline to the scheduler."""

    name: SamplingMode | str = SamplingMode.random
    temperature: float = 1.0
    top_k: int = 5
    top_p: float = 0.9
    num_inference_steps: int | None = None
    time_difference: float = 0.0
    refine_lambda: float = 3.0
    refine_mode: Literal["uniform", "gaussian", "negative"] = "uniform"
    refine_offset_ratio: float = 0.1

    def __post_init__(self) -> None:
        """Normalize public string sampling values to ``SamplingMode``."""
        self.name = normalize_sampling_mode(self.name)