Dlt

class DLTConfig(ConfigMixin):
    """Pipeline-level DLT configuration persisted with converted checkpoints.

    Args:
        dataset_name: Canonical dataset name.
        id2label: Optional public label mapping. When omitted, shared dataset
            labels are used.
        max_num_comp: Maximum number of layout elements.
        categories_num: Internal category count including pad and mask/drop ids.
        latent_dim: Transformer latent dimension.
        num_layers: Number of transformer encoder layers.
        num_heads: Number of attention heads.
        dropout_r: Dropout probability.
        activation: Transformer activation.
        cond_emb_size: Box-condition embedding size.
        cat_emb_size: Category embedding size.
        num_cont_timesteps: Continuous DDPM training timesteps.
        num_discrete_steps: Discrete category diffusion steps.
        beta_schedule: DDPM beta schedule.
        coordinate_range: Public coordinate range.
    """

    config_name = "dlt_config.json"

    @register_to_config
    def __init__(
        self,
        *,
        dataset_name: str = "publaynet",
        id2label: dict[int | str, str] | None = None,
        max_num_comp: int | None = None,
        categories_num: int | None = None,
        latent_dim: int = 512,
        num_layers: int = 4,
        num_heads: int = 8,
        dropout_r: float = 0.0,
        activation: str = "gelu",
        cond_emb_size: int = 224,
        cat_emb_size: int = 64,
        num_cont_timesteps: int = 100,
        num_discrete_steps: int = 10,
        beta_schedule: str = "squaredcos_cap_v2",
        coordinate_range: DLTCoordinateRange | str = DLTCoordinateRange.normalized_0_1,
    ) -> None:
        """Initialize DLT configuration."""
        dataset = normalize_dataset(dataset_name)
        labels = default_id2label(dataset) if id2label is None else id2label

        self.dataset_name = str(dataset)
        self.id2label = {int(key): value for key, value in labels.items()}
        self.max_num_comp = max_num_comp or max_elements_for_dataset(dataset)
        self.categories_num = categories_num or len(self.id2label) + 2

        self.latent_dim = latent_dim
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.dropout_r = dropout_r
        self.activation = activation
        self.cond_emb_size = cond_emb_size
        self.cat_emb_size = cat_emb_size
        self.num_cont_timesteps = num_cont_timesteps
        self.num_discrete_steps = num_discrete_steps
        self.beta_schedule = beta_schedule
        self.coordinate_range = str(DLTCoordinateRange(coordinate_range))

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

class DLT(ModelMixin, ConfigMixin):
    """Joint continuous/discrete DLT denoiser.

    The module names intentionally match released checkpoint keys so
    ``model.save_pretrained`` directories can load without key rewriting.

    Args:
        categories_num: Internal category count including pad and mask/drop ids.
        latent_dim: Transformer latent dimension.
        num_layers: Number of transformer encoder layers.
        num_heads: Number of attention heads.
        dropout_r: Dropout probability.
        activation: Transformer activation.
        cond_emb_size: Box-condition embedding size.
        cat_emb_size: Category embedding size.
    """

    config_name = "model_config.json"

    @register_to_config
    def __init__(
        self,
        categories_num: int,
        latent_dim: int = 256,
        num_layers: int = 4,
        num_heads: int = 4,
        dropout_r: float = 0.0,
        activation: str = "gelu",
        cond_emb_size: int = 224,
        cat_emb_size: int = 64,
    ) -> None:
        """Initialize the DLT denoiser."""
        super().__init__()
        self.latent_dim = latent_dim
        self.dropout_r = dropout_r
        self.categories_num = categories_num
        self.seq_pos_enc = PositionalEncoding(self.latent_dim, self.dropout_r)
        self.cat_emb = nn.Parameter(torch.randn(self.categories_num, cat_emb_size))
        self.cond_mask_box_emb = nn.Parameter(torch.randn(2, cond_emb_size))
        self.cond_mask_cat_emb = nn.Parameter(torch.randn(2, cat_emb_size))

        seq_trans_encoder_layer = nn.TransformerEncoderLayer(
            d_model=self.latent_dim,
            nhead=num_heads,
            dim_feedforward=self.latent_dim * 2,
            dropout=dropout_r,
            activation=activation,
        )
        self.seqTransEncoder = nn.TransformerEncoder(
            seq_trans_encoder_layer, num_layers=num_layers
        )
        self.embed_timestep = TimestepEmbedder(self.latent_dim, self.seq_pos_enc)
        self.output_process = nn.Sequential(nn.Linear(self.latent_dim, 4))
        self.output_cls = nn.Sequential(nn.Linear(self.latent_dim, categories_num))
        self.size_emb = nn.Sequential(nn.Linear(2, cond_emb_size))
        self.loc_emb = nn.Sequential(nn.Linear(2, cond_emb_size))

    def forward(
        self,
        sample: dict[
            str,
            Float[torch.Tensor, "batch elements channels"]
            | Int[torch.Tensor, "batch elements"],
        ],
        noisy_sample: dict[
            str,
            Float[torch.Tensor, "batch elements channels"]
            | Int[torch.Tensor, "batch elements"],
        ],
        timesteps: Int[torch.Tensor, "batch"],
        return_dict: bool = False,
    ) -> (
        DLTModelOutput
        | tuple[
            Float[torch.Tensor, "batch elements 4"],
            Float[torch.Tensor, "batch elements categories"],
        ]
    ):
        """Predict clean boxes and category logits for a noisy layout.

        Args:
            sample: DLT-format conditioning batch with ``box_cond``,
                ``cat``, ``mask_box``, and ``mask_cat``.
            noisy_sample: Current noisy ``box`` and ``cat`` tensors.
            timesteps: Continuous diffusion timestep per batch item.
            return_dict: Whether to return ``DLTModelOutput``.

        Returns:
            Either a two-tuple ``(box, logits)`` for checkpoint-compatible
            callers or a dataclass output.
        """
        cat_input = (
            noisy_sample["cat"] * sample["mask_cat"]
            + (1 - sample["mask_cat"]) * sample["cat"]
        )
        cat_input_flat = rearrange(cat_input, "b c -> (b c)")
        sample_tensor = (
            sample["mask_box"] * noisy_sample["box"]
            + (1 - sample["mask_box"]) * sample["box_cond"]
        )

        xy = sample_tensor[:, :, :2]
        wh = sample_tensor[:, :, 2:]

        elem_cat_emb = self.cat_emb[cat_input_flat, :]
        elem_cat_emb = rearrange(
            elem_cat_emb, "(b c) d -> b c d", b=noisy_sample["box"].shape[0]
        )

        def mask_to_emb(
            mask: Int[torch.Tensor, "batch elements"],
            cond_mask_emb: Float[torch.Tensor, "mask channels"],
        ) -> Float[torch.Tensor, "batch elements channels"]:
            mask_flat = rearrange(mask, "b c -> (b c)").long()
            mask_all_emb = cond_mask_emb[mask_flat, :]
            return rearrange(mask_all_emb, "(b c) d -> b c d", b=mask.shape[0])

        emb_mask_wh = mask_to_emb(sample["mask_box"][:, :, 2], self.cond_mask_box_emb)
        emb_mask_xy = mask_to_emb(sample["mask_box"][:, :, 0], self.cond_mask_box_emb)
        emb_mask_cl = mask_to_emb(sample["mask_cat"], self.cond_mask_cat_emb)
        t_emb = self.embed_timestep(timesteps)

        size_emb = self.size_emb(wh) + emb_mask_wh
        loc_emb = self.loc_emb(xy) + emb_mask_xy
        elem_cat_emb = elem_cat_emb + emb_mask_cl

        tokens_emb = torch.cat([size_emb, loc_emb, elem_cat_emb], dim=-1)
        tokens_emb = rearrange(tokens_emb, "b c d -> c b d")
        xseq = torch.cat((t_emb, tokens_emb), dim=0)
        xseq = self.seq_pos_enc(xseq)

        output = self.seqTransEncoder(xseq)[1:]
        output = rearrange(output, "c b d -> b c d")
        output_box = self.output_process(output)
        output_cls = self.output_cls(output)
        if not return_dict:
            return output_box, output_cls
        return DLTModelOutput(box=output_box, logits=output_cls)

    def save_pretrained(
        self,
        save_directory: str | os.PathLike[str],
        is_main_process: bool = True,
        save_function: Callable[..., None] | None = None,
        safe_serialization: bool = False,
        variant: str | None = None,
        max_shard_size: int | str = "10GB",
        push_to_hub: bool = False,
        use_flashpack: bool = False,
        **kwargs: str | int | bool | float | None,
    ) -> None:
        """Save the model with PyTorch serialization by default.

        DLT keeps shared positional-encoding buffers that safetensors refuses
        to flatten.
        """
        super().save_pretrained(
            save_directory,
            is_main_process=is_main_process,
            save_function=save_function,
            safe_serialization=safe_serialization,
            variant=variant,
            max_shard_size=max_shard_size,
            push_to_hub=push_to_hub,
            use_flashpack=use_flashpack,
            **kwargs,
        )

@dataclass
class DLTModelOutput(BaseOutput):
    """Output returned by the DLT denoiser.

    Attributes:
        box: Predicted clean internal-range boxes.
        logits: Category logits.
    """

    box: Float[torch.Tensor, "batch elements 4"]
    logits: Float[torch.Tensor, "batch elements categories"]

class DLTConditionAlias(StrEnum):
    """DLT checkpoint condition aliases."""

    all = auto()
    whole_box = auto()
    loc = auto()

class DLTPipeline(DiffusionPipeline):
    """Generate layouts with a converted DLT checkpoint.

    Args:
        model: DLT denoiser.
        scheduler: Joint box/category scheduler.
        config: Pipeline configuration.
        processor: Layout processor.
    """

    model_cpu_offload_seq = "model"
    _optional_components = ["processor"]

    def __init__(
        self,
        model: DLT,
        scheduler: DLTJointDiffusionScheduler,
        config: DLTConfig,
        processor: DLTProcessor | None = None,
    ) -> None:
        """Initialize a DLT pipeline."""
        super().__init__()
        self.register_modules(model=model, scheduler=scheduler)
        self.dlt_config = config
        self.processor = processor or DLTProcessor(
            dataset=self.dlt_config.dataset_name,
            labels=tuple(self.dlt_config.id2label.values()),
            max_num_comp=self.dlt_config.max_num_comp,
        )
        self.model.eval()

    @property
    def components(self) -> PipelineComponents:
        """Return serializable pipeline components."""
        return {
            "model": self.model,
            "scheduler": self.scheduler,
            "processor": self.processor,
        }

    @torch.no_grad()
    def __call__(
        self,
        *,
        batch_size: int = 1,
        seed: int | None = None,
        generator: torch.Generator | None = None,
        condition_type: ConditionType | str | None = ConditionType.unconditional,
        labels: Int[torch.Tensor, "batch elements"] | None = None,
        bbox: Float[torch.Tensor, "batch elements 4"] | None = None,
        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,
        num_inference_steps: int | None = None,
        temperature: float | None = None,
        output_type: OutputType | str = OutputType.dataclass,
        return_intermediates: bool = False,
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Float[torch.Tensor, "..."]
            | Int[torch.Tensor, "..."]
            | Bool[torch.Tensor, "..."]
            | dict[int, str]
            | list[Float[torch.Tensor, "..."]]
            | dict[str, str]
            | None,
        ]
    ):
        """Run DLT joint denoising and return generated layouts.

        Args:
            batch_size: Number of layouts to generate.
            seed: Optional seed used when ``generator`` is absent.
            generator: Optional torch generator. Takes precedence over ``seed``.
            condition_type: Canonical condition or DLT checkpoint alias.
            labels: Optional public labels for conditioned modes.
            bbox: Optional public boxes for conditioned modes.
            mask: Optional valid-element mask.
            num_elements: Optional valid element count for unconditional calls.
            box_format: Input box format.
            normalized: Whether input boxes are normalized.
            canvas_size: Pixel canvas size for non-normalized boxes.
            num_inference_steps: Number of reverse diffusion steps.
            temperature: Optional category sampling temperature override.
            output_type: ``"dataclass"`` or ``"dict"``.
            return_intermediates: Whether to include denoising trajectory.

        Returns:
            Layout generation output dataclass or dictionary.

        Raises:
            ValueError: If the condition or output type is unsupported.
        """
        canonical = normalize_condition_type(condition_type)
        output_kind = OutputType(output_type)
        if generator is None and seed is not None:
            generator = torch.Generator(device=self.device).manual_seed(seed)
        if canonical is ConditionType.unconditional:
            processed = self.processor.empty_condition(
                batch_size=batch_size, device=self.device
            )
            if num_elements is not None:
                lengths = torch.as_tensor(
                    num_elements, dtype=torch.long, device=self.device
                )
                if lengths.ndim == 0:
                    lengths = lengths.repeat(batch_size)
                processed["mask"] = (
                    torch.arange(self.processor.max_num_comp, device=self.device)[
                        None, :
                    ]
                    < lengths[:, None]
                )
        else:
            _require_condition_inputs(
                condition_type=condition_type, bbox=bbox, labels=labels
            )
            processed = self.processor(
                bbox=bbox,
                labels=labels,
                mask=mask,
                box_format=box_format,
                normalized=normalized,
                canvas_size=canvas_size,
                device=self.device,
            )
            batch_size = processed["box"].shape[0]
        mask_box, mask_cat = self.processor.condition_masks(
            str(canonical), mask=processed["mask"]
        )
        sample = {
            "box": processed["box"],
            "box_cond": processed["box_cond"],
            "cat": processed["cat"],
            "mask_box": mask_box,
            "mask_cat": mask_cat,
        }
        noisy_batch = {
            "box": torch.randn(
                processed["box"].shape,
                dtype=processed["box"].dtype,
                device=self.device,
                generator=generator,
            ),
            "cat": torch.full(
                processed["cat"].shape,
                self.processor.mask_category_id,
                dtype=torch.long,
                device=self.device,
            ),
        }
        original_temperature = self.scheduler.temperature
        if temperature is not None:
            self.scheduler.temperature = temperature
        steps = max(1, num_inference_steps or self.scheduler.num_train_timesteps)
        trajectory: list[Float[torch.Tensor, "batch elements 4"]] | None = (
            [] if return_intermediates else None
        )
        bbox_step: DLTJointSchedulerOutput | None = None
        try:
            for i in range(steps - 1, -1, -1):
                t_value = min(i, self.scheduler.num_train_timesteps - 1)
                t = torch.tensor([t_value] * batch_size, device=self.device)
                bbox_pred, cat_pred = self.model(sample, noisy_batch, timesteps=t)
                bbox_step, cat_step = self.scheduler.step_jointly(
                    bbox_pred,
                    {"cat": cat_pred},
                    timestep=t,
                    sample=noisy_batch["box"],
                    generator=generator,
                )
                noisy_batch["box"] = bbox_step.prev_sample
                noisy_batch["cat"] = cat_step["cat"]
                if trajectory is not None:
                    trajectory.append(noisy_batch["box"].detach().cpu())
        finally:
            self.scheduler.temperature = original_temperature
        if bbox_step is None:
            raise RuntimeError("DLT denoising did not run any scheduler steps")

        final_box = (
            sample["mask_box"] * bbox_step.pred_original_sample
            + (1 - sample["mask_box"]) * sample["box_cond"]
        )
        final_cat = (
            sample["mask_cat"] * noisy_batch["cat"]
            + (1 - sample["mask_cat"]) * sample["cat"]
        )
        valid_mask = processed["mask"].detach().cpu()
        output = LayoutGenerationOutput(
            bbox=self.processor.internal_to_public_boxes(final_box).detach().cpu()
            * valid_mask.unsqueeze(-1),
            labels=self.processor.internal_to_public_labels(
                final_cat, processed["mask"]
            )
            .detach()
            .cpu(),
            mask=valid_mask,
            id2label=self.processor.id2label,
            trajectory=trajectory,
            intermediates={"condition_type": str(canonical)}
            if return_intermediates
            else None,
        )
        return _format_pipeline_output(output, output_kind)

    generate = __call__

    def save_pretrained(self, save_directory: str | Path) -> None:
        """Persist DLT model, scheduler, and pipeline metadata."""
        super().save_pretrained(save_directory, safe_serialization=False)
        self.dlt_config.save_config(save_directory)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str | Path) -> Self:
        """Load a saved DLT pipeline."""
        config_dict, _ = DLTConfig.load_config(
            pretrained_model_name_or_path, return_unused_kwargs=True
        )
        config = cast(DLTConfig, DLTConfig.from_config(config_dict))
        pipe = super().from_pretrained(pretrained_model_name_or_path, config=config)
        pipe.dlt_config = config
        pipe.processor = DLTProcessor(
            dataset=config.dataset_name,
            labels=tuple(config.id2label.values()),
            max_num_comp=config.max_num_comp,
        )
        return pipe

class OutputType(StrEnum):
    """DLT pipeline output containers."""

    dataclass = auto()
    dict = auto()

class DLTProcessor(ProcessorMixin):
    """Encode DLT public inputs into the package tensor format.

    Args:
        dataset: Canonical dataset name.
        labels: Ordered public labels without internal pad/drop ids.
        max_num_comp: Maximum number of layout elements.
    """

    config_name = "processor_config.json"

    def __init__(
        self,
        dataset: DatasetName | str,
        labels: Sequence[str],
        max_num_comp: int,
    ) -> None:
        """Initialize processor metadata."""
        super().__init__()
        self.dataset = str(normalize_dataset(dataset))
        self.labels = tuple(str(label) for label in labels)
        self.max_num_comp = max_num_comp

    @classmethod
    def from_dataset(cls, dataset: DatasetName | str) -> "DLTProcessor":
        """Create a processor from shared dataset metadata."""
        canonical = normalize_dataset(dataset)
        return cls(
            dataset=canonical,
            labels=tuple(default_id2label(canonical).values()),
            max_num_comp=9
            if canonical in {DatasetName.publaynet, DatasetName.rico13}
            else 33,
        )

    @property
    def id2label(self) -> dict[int, str]:
        """Return public label names keyed by dataset-local ids."""
        return dict(enumerate(self.labels))

    @property
    def categories_num(self) -> int:
        """Return internal category count including pad and mask/drop ids."""
        return len(self.labels) + 2

    @property
    def pad_category_id(self) -> int:
        """Return the internal padding category id."""
        return 0

    @property
    def mask_category_id(self) -> int:
        """Return the internal mask/drop category id."""
        return self.categories_num - 1

    def __call__(
        self,
        *,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Float[np.ndarray, "batch elements 4"]
        | LayoutInput,
        labels: Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | LayoutInput,
        mask: Bool[torch.Tensor, "batch elements"]
        | Bool[np.ndarray, "batch elements"]
        | LayoutInput
        | None = None,
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        device: torch.device | str | None = None,
    ) -> DLTProcessedBatch:
        """Convert public layout tensors into padded internal tensors."""
        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,
        )
        if device is not None:
            bbox_t = bbox_t.to(device)
            labels_t = labels_t.to(device)
            mask_t = mask_t.to(device)
        bbox_t, labels_t, mask_t = self.pad(bbox_t, labels_t, mask_t)
        return cast(
            DLTProcessedBatch,
            {
                "box": self.public_to_internal_boxes(bbox_t) * mask_t.unsqueeze(-1),
                "box_cond": self.public_to_internal_boxes(bbox_t)
                * mask_t.unsqueeze(-1),
                "cat": self.public_to_internal_labels(labels_t, mask_t),
                "mask": mask_t,
            },
        )

    def empty_condition(
        self,
        *,
        batch_size: int,
        device: torch.device | str,
        dtype: torch.dtype = torch.float32,
    ) -> DLTProcessedBatch:
        """Return an empty unconditional internal batch."""
        device = torch.device(device)
        bbox = torch.zeros(batch_size, self.max_num_comp, 4, dtype=dtype, device=device)
        labels = torch.zeros(
            batch_size, self.max_num_comp, dtype=torch.long, device=device
        )
        mask = torch.ones(
            batch_size, self.max_num_comp, dtype=torch.bool, device=device
        )
        return cast(
            DLTProcessedBatch,
            {
                "box": self.public_to_internal_boxes(bbox),
                "box_cond": self.public_to_internal_boxes(bbox),
                "cat": self.public_to_internal_labels(labels, mask),
                "mask": mask,
            },
        )

    def pad(
        self,
        bbox: Float[torch.Tensor, "batch elements 4"],
        labels: Int[torch.Tensor, "batch elements"],
        mask: Bool[torch.Tensor, "batch elements"] | None = None,
    ) -> tuple[
        Float[torch.Tensor, "batch max_elements 4"],
        Int[torch.Tensor, "batch max_elements"],
        Bool[torch.Tensor, "batch max_elements"],
    ]:
        """Pad a layout batch to ``max_num_comp``."""
        if bbox.shape[1] > self.max_num_comp:
            raise ValueError(f"DLT supports at most {self.max_num_comp} elements")

        if mask is None:
            mask = torch.ones(labels.shape, dtype=torch.bool, device=labels.device)
        pad_count = self.max_num_comp - bbox.shape[1]
        if pad_count:
            bbox = torch.cat(
                [
                    bbox,
                    torch.zeros(
                        bbox.shape[0],
                        pad_count,
                        4,
                        dtype=bbox.dtype,
                        device=bbox.device,
                    ),
                ],
                dim=1,
            )
            labels = torch.cat(
                [
                    labels,
                    torch.zeros(
                        labels.shape[0],
                        pad_count,
                        dtype=labels.dtype,
                        device=labels.device,
                    ),
                ],
                dim=1,
            )
            mask = torch.cat(
                [
                    mask,
                    torch.zeros(
                        mask.shape[0], pad_count, dtype=torch.bool, device=mask.device
                    ),
                ],
                dim=1,
            )
        return bbox, labels, mask

    def public_to_internal_boxes(
        self, bbox: Float[torch.Tensor, "batch elements 4"]
    ) -> Float[torch.Tensor, "batch elements 4"]:
        """Map public normalized ``xywh`` boxes to DLT's internal range."""
        return bbox.clamp(0.0, 1.0) * 4.0 - 2.0

    def internal_to_public_boxes(
        self, bbox: Float[torch.Tensor, "batch elements 4"]
    ) -> Float[torch.Tensor, "batch elements 4"]:
        """Map DLT internal-range boxes to public normalized ``xywh``."""
        return (bbox / 2.0 + 1.0).div(2.0).clamp(0.0, 1.0)

    def public_to_internal_labels(
        self,
        labels: Int[torch.Tensor, "batch elements"],
        mask: Bool[torch.Tensor, "batch elements"],
    ) -> Int[torch.Tensor, "batch elements"]:
        """Shift public labels into DLT's internal category ids."""
        shifted = labels.long() + 1
        return torch.where(mask.bool(), shifted, torch.zeros_like(shifted))

    def internal_to_public_labels(
        self,
        labels: Int[torch.Tensor, "batch elements"],
        mask: Bool[torch.Tensor, "batch elements"],
    ) -> Int[torch.Tensor, "batch elements"]:
        """Shift DLT internal category ids back to public dataset-local ids."""
        public = (labels.long() - 1).clamp(0, len(self.labels) - 1)
        return public * mask.long()

    def condition_masks(
        self, condition_type: str, *, mask: Bool[torch.Tensor, "batch elements"]
    ) -> tuple[
        Int[torch.Tensor, "batch elements 4"],
        Int[torch.Tensor, "batch elements"],
    ]:
        """Return DLT ``mask_box`` and ``mask_cat`` tensors.

        ``1`` means generated/noised and ``0`` means conditioned.
        """
        mask_box = torch.ones(
            mask.shape[0], mask.shape[1], 4, dtype=torch.long, device=mask.device
        )
        mask_cat = torch.ones(mask.shape, dtype=torch.long, device=mask.device)
        if condition_type == "label":
            mask_cat.zero_()
        elif condition_type == "label_size":
            mask_box[:, :, 2:] = 0
            mask_cat.zero_()
        elif condition_type == "unconditional":
            pass
        else:
            raise ValueError(f"Unsupported DLT condition_type: {condition_type}")

        mask_box = mask_box * mask.unsqueeze(-1).long()
        mask_cat = mask_cat * mask.long()
        return mask_box, mask_cat