Posterllava

class PosterLlavaConfig(PretrainedConfig):
    """Configuration saved with a PosterLLaVA recipe checkpoint.

    Args:
        checkpoint_id: Upstream LLaVA-style checkpoint id used by local smoke
            scripts and documentation.
        dataset_name: Canonical poster/content dataset metadata key.
        id2label: Persisted label metadata. Open-vocabulary generation uses a
            batch-local map at runtime, but this config records known dataset
            labels for model cards and smoke checks.
        prompt_template: Prompt body template passed through the LLaVA
            conversation wrapper.
        default_conv_mode: Default LLaVA conversation template.
        image_aspect_ratio: Image preprocessing mode; ``"pad"`` matches the
            released checkpoint.
        max_new_tokens: Default token budget for generation.
        default_temperature: Default sampled-generation temperature.
        processor_subfolder: Subfolder used by pipeline component loading.
        model_subfolder: Optional model component subfolder.
        tokenizer_subfolder: Optional tokenizer component subfolder.
        image_processor_subfolder: Optional image processor component subfolder.
        kwargs: Extra Hugging Face config fields.

    Raises:
        ValueError: If numeric fields or enum-like fields are invalid.

    Examples:
        >>> cfg = PosterLlavaConfig(dataset_name="ad_banner")
        >>> cfg.checkpoint_id
        'posterllava/posterllava_v0'
    """

    model_type = "posterllava"
    id2label: dict[int, str]

    def __init__(
        self,
        *,
        checkpoint_id: str = DEFAULT_CHECKPOINT_ID,
        dataset_name: DatasetName | str = DatasetName.ad_banner,
        id2label: Mapping[int, str] | Mapping[str, str] | None = None,
        prompt_template: str = DEFAULT_PROMPT_TEMPLATE,
        default_conv_mode: ConversationMode | str = ConversationMode.llava_v0,
        image_aspect_ratio: str = "pad",
        max_new_tokens: int = 1024,
        default_temperature: float = 0.2,
        processor_subfolder: str = "processor",
        model_subfolder: str = "model",
        tokenizer_subfolder: str = "tokenizer",
        image_processor_subfolder: str = "image_processor",
        **kwargs: PosterLlavaConfigValue,
    ) -> None:
        """Initialize PosterLLaVA recipe configuration."""
        super().__init__(**kwargs)  # ty: ignore[invalid-argument-type]
        dataset = normalize_dataset_name(dataset_name)
        conv_mode = normalize_conversation_mode(default_conv_mode)
        if max_new_tokens <= 0:
            raise ValueError("max_new_tokens must be positive")

        if default_temperature < 0:
            raise ValueError("default_temperature must be non-negative")

        if image_aspect_ratio != "pad":
            raise ValueError("PosterLLaVA currently supports image_aspect_ratio='pad'")

        self.checkpoint_id = checkpoint_id
        self.dataset_name = str(dataset)
        self.id2label = {
            int(key): str(value)
            for key, value in (id2label or id2label_for_dataset(dataset)).items()
        }
        self.prompt_template = prompt_template
        self.default_conv_mode = str(conv_mode)
        self.image_aspect_ratio = image_aspect_ratio
        self.max_new_tokens = max_new_tokens
        self.default_temperature = default_temperature
        self.processor_subfolder = processor_subfolder
        self.model_subfolder = model_subfolder
        self.tokenizer_subfolder = tokenizer_subfolder
        self.image_processor_subfolder = image_processor_subfolder

class PosterLlavaImageProcessor(CLIPImageProcessor):
    """CLIP image processor with PosterLLaVA square-padding behavior.

    Args:
        kwargs: Keyword arguments forwarded to ``CLIPImageProcessor``.

    Examples:
        >>> from PIL import Image
        >>> processor = PosterLlavaImageProcessor()
        >>> image = Image.new("RGB", (8, 4))
        >>> processor.expand_to_square(image).size
        (8, 8)
    """

    model_input_names = ["pixel_values"]

    def expand_to_square(self, image: Image.Image) -> Image.Image:
        """Pad an image to a square using the configured CLIP mean color.

        Args:
            image: RGB image to pad.

        Returns:
            Square RGB image.
        """
        return _expand_to_square(image, self.image_mean)

    def preprocess(
        self,
        images: Image.Image | Sequence[Image.Image],
        *,
        image_aspect_ratio: Literal["pad"] = "pad",
        return_tensors: str | None = "pt",
        **kwargs: PosterLlavaImageProcessorKwarg,
    ) -> BatchFeature:
        """Preprocess images with PosterLLaVA's square-padding policy.

        Args:
            images: One image or a sequence of images.
            image_aspect_ratio: Only ``"pad"`` is supported.
            return_tensors: Tensor container requested from Transformers.
            kwargs: Additional ``CLIPImageProcessor.preprocess`` options.

        Returns:
            Batch feature with processed image tensors.

        Raises:
            ValueError: If ``image_aspect_ratio`` is unsupported.
        """
        if image_aspect_ratio != "pad":
            raise ValueError("PosterLLaVA image preprocessing only supports pad")

        image_list = [images] if isinstance(images, Image.Image) else list(images)
        padded = [self.expand_to_square(image.convert("RGB")) for image in image_list]
        return cast(
            BatchFeature,
            super().preprocess(padded, return_tensors=return_tensors, **kwargs),
        )

class PosterLlavaPipeline(LayoutGenerationPipeline):
    """Generate poster layouts with a LLaVA-style causal LM checkpoint.

    Args:
        config: PosterLLaVA recipe configuration.
        processor: Prompt and JSON layout processor.
        model: Optional upstream causal LM component.
        tokenizer: Optional LLaVA tokenizer component.
        image_processor: Optional CLIP image processor component.

    Examples:
        >>> cfg = PosterLlavaConfig(dataset_name="ad_banner")
        >>> processor = PosterLlavaProcessor.from_config()
        >>> pipe = PosterLlavaPipeline(cfg, processor)
        >>> pipe.config.model_type
        'posterllava'
    """

    config_class: ClassVar[type[PretrainedConfig]] = PosterLlavaConfig
    component_specs: ClassVar[dict[str, PipelineComponentSpec]] = {
        "model": PipelineComponentSpec(
            attribute_name="model",
            loader=_load_model_component,
            config_subfolder_attribute="model_subfolder",
            required=False,
            marker_file="config.json",
        ),
        "tokenizer": PipelineComponentSpec(
            attribute_name="tokenizer",
            loader=_load_tokenizer_component,
            config_subfolder_attribute="tokenizer_subfolder",
            required=False,
            marker_file="tokenizer_config.json",
            save_with_is_main_process=False,
        ),
        "image_processor": PipelineComponentSpec(
            attribute_name="image_processor",
            loader=_load_image_processor_component,
            config_subfolder_attribute="image_processor_subfolder",
            required=False,
            marker_file="preprocessor_config.json",
            save_with_is_main_process=False,
        ),
        "processor": PipelineComponentSpec(
            attribute_name="processor",
            loader=_load_processor_component,
            config_subfolder_attribute="processor_subfolder",
            required=False,
            marker_file="processor_config.json",
            save_with_is_main_process=False,
        ),
    }

    config: PosterLlavaConfig
    processor: PosterLlavaProcessor
    model: PreTrainedModel | None
    tokenizer: PreTrainedTokenizerBase | None
    image_processor: PosterLlavaImageProcessorComponent | None

    def __init__(
        self,
        config: PosterLlavaConfig,
        processor: PosterLlavaProcessor,
        *,
        model: PreTrainedModel | None = None,
        tokenizer: PreTrainedTokenizerBase | None = None,
        image_processor: PosterLlavaImageProcessorComponent | None = None,
    ) -> None:
        """Initialize the PosterLLaVA recipe pipeline."""
        super().__init__(config)
        self.config = config
        self.processor = processor
        self.model = model
        self.tokenizer = tokenizer or processor.tokenizer
        self.image_processor = image_processor or processor.image_processor
        if self.tokenizer is not None:
            self.processor.tokenizer = self.tokenizer
        if self.image_processor is not None:
            self.processor.image_processor = self.image_processor

    @classmethod
    def _from_pretrained_components(
        cls,
        *,
        config: PretrainedConfig,
        components: Mapping[str, PosterLlavaComponent | None],
    ) -> "PosterLlavaPipeline":
        """Build a pipeline from loaded root config and components."""
        cfg = cast(PosterLlavaConfig, config)
        processor = cast(PosterLlavaProcessor | None, components.get("processor"))
        if processor is None:
            processor = PosterLlavaProcessor.from_config(
                dataset_name=cfg.dataset_name,
                id2label=cfg.id2label,
                prompt_template=cfg.prompt_template,
            )
        tokenizer = cast(PreTrainedTokenizerBase | None, components.get("tokenizer"))
        image_processor = components.get("image_processor")
        return cls(
            config=cfg,
            processor=processor,
            model=cast(PreTrainedModel | None, components.get("model")),
            tokenizer=tokenizer,
            image_processor=image_processor,
        )

    def __call__(
        self,
        *,
        images: Image.Image | Sequence[Image.Image] | None = None,
        prompt: str | Sequence[str] | None = None,
        content: Mapping[str, PosterLlavaContentValue]
        | Sequence[Mapping[str, PosterLlavaContentValue]]
        | None = None,
        texts: str | Sequence[str] | Sequence[Sequence[str]] | None = None,
        batch_size: int = 1,
        seed: int | None = None,
        generator: torch.Generator | None = None,
        condition_type: ConditionType | str = ConditionType.content_image,
        labels: Int[torch.Tensor, "batch elements"] | Sequence[str | int] | None = None,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Sequence[Sequence[float]]
        | None = None,
        mask: Bool[torch.Tensor, "batch elements"] | Sequence[bool] | None = None,
        num_elements: int | Sequence[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 | Literal["dataclass", "dict"] = OutputType.dataclass,
        return_intermediates: bool = False,
        max_new_tokens: int | None = None,
        do_sample: bool = True,
        temperature: float | None = None,
        top_p: float | None = 1.0,
        top_k: int | None = None,
        num_beams: int | None = 1,
        conv_mode: ConversationMode | str | None = None,
        domain_name: str = "social media promotion poster with qbposter style",
    ) -> LayoutGenerationOutput | PosterLlavaOutputDict:  # ty: ignore[invalid-method-override]
        """Generate a poster layout from an image-conditioned prompt.

        Args:
            images: Poster/background image or image batch.
            prompt: Optional prompt body override.
            content: Optional payload mapping containing ``image``, ``text``,
                ``texts``, ``num_elements``, or ``json_data``.
            texts: Optional aligned text payload.
            batch_size: Expected batch size when scalar inputs are provided.
            seed: Seed used only when ``generator`` is absent.
            generator: Explicit generator passed to model generation.
            condition_type: Canonical condition type. Only ``content_image`` is
                supported in the first package version.
            labels: Optional initial labels.
            bbox: Optional initial boxes aligned with labels.
            mask: Optional initial valid-element mask.
            num_elements: Requested element count.
            box_format: Public input box format.
            normalized: Whether input boxes are normalized.
            canvas_size: Pixel canvas size for non-normalized input boxes.
            num_inference_steps: Accepted for shared-interface compatibility.
            output_type: Output container mode.
            return_intermediates: Whether raw prompts/text are returned.
            max_new_tokens: Token budget for generation.
            do_sample: Whether to sample from the LLM.
            temperature: Sampling temperature.
            top_p: Nucleus sampling parameter.
            top_k: Top-k sampling parameter.
            num_beams: Beam count.
            conv_mode: Optional LLaVA conversation template override.
            domain_name: Domain phrase inserted into the default prompt.

        Returns:
            Layout output dataclass or dictionary.

        Raises:
            NotImplementedError: If ``condition_type`` is unsupported.
            ValueError: If required image/model/tokenizer components are absent.
        """
        _ = num_inference_steps
        condition = normalize_condition_type(condition_type)
        if condition is not ConditionType.content_image:
            raise NotImplementedError(
                "PosterLLaVA only supports condition_type='content_image'"
            )

        image_list = self._resolve_images(images=images, content=content)
        prompts = self._build_prompts(
            prompt=prompt,
            content=content,
            texts=texts,
            batch_size=batch_size,
            labels=labels,
            bbox=bbox,
            mask=mask,
            num_elements=num_elements,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
            conv_mode=conv_mode,
            domain_name=domain_name,
        )
        if len(image_list) != len(prompts):
            raise ValueError("images and prompts must resolve to the same batch size")

        if self.model is None:
            raise ValueError("model is required for PosterLLaVA generation")

        if self.tokenizer is None:
            raise ValueError("tokenizer is required for PosterLLaVA generation")

        if self.image_processor is None:
            raise ValueError("image_processor is required for PosterLLaVA generation")

        encoded = self.processor(prompts)
        input_ids = cast(Int[torch.Tensor, "batch tokens"], encoded["input_ids"])
        attention_mask = cast(
            Bool[torch.Tensor, "batch tokens"],
            encoded["attention_mask"],
        )
        pixel_values = self._preprocess_images(image_list)
        generation_generator = self.prepare_generator(generator=generator, seed=seed)
        sequences = cast(_CausalLMGenerationModel, self.model).generate(
            input_ids=input_ids,
            images=pixel_values,
            attention_mask=attention_mask,
            max_new_tokens=max_new_tokens or self.config.max_new_tokens,
            do_sample=do_sample,
            temperature=temperature
            if temperature is not None
            else self.config.default_temperature,
            top_p=top_p,
            top_k=top_k,
            num_beams=num_beams,
            generator=generation_generator,
            stopping_criteria=build_stopping_criteria(
                self.tokenizer,
                input_ids=input_ids,
            ),
        )
        generated_text = self.tokenizer.batch_decode(
            sequences[:, input_ids.shape[-1] :],
            skip_special_tokens=True,
        )
        return self.processor.decode_layout(
            generated_text,
            output_type=output_type,
            return_intermediates=return_intermediates,
            sequences=sequences,
            prompts=prompts,
        )

    generate = __call__

    def _resolve_images(
        self,
        *,
        images: Image.Image | Sequence[Image.Image] | None,
        content: Mapping[str, PosterLlavaContentValue]
        | Sequence[Mapping[str, PosterLlavaContentValue]]
        | None,
    ) -> list[Image.Image]:
        if images is None:
            if content is None:
                raise ValueError("images or content['image'] is required")

            content_items = [content] if isinstance(content, Mapping) else list(content)
            raw_images = [item.get("image") for item in content_items]
        else:
            raw_images = [images] if isinstance(images, Image.Image) else list(images)
        if not raw_images or any(
            not isinstance(item, Image.Image) for item in raw_images
        ):
            raise ValueError("PosterLLaVA images must be PIL.Image.Image instances")

        return cast(list[Image.Image], raw_images)

    def _build_prompts(
        self,
        *,
        prompt: str | Sequence[str] | None,
        content: Mapping[str, PosterLlavaContentValue]
        | Sequence[Mapping[str, PosterLlavaContentValue]]
        | None,
        texts: str | Sequence[str] | Sequence[Sequence[str]] | None,
        batch_size: int,
        labels: Int[torch.Tensor, "batch elements"] | Sequence[str | int] | None,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Sequence[Sequence[float]]
        | None,
        mask: Bool[torch.Tensor, "batch elements"] | Sequence[bool] | None,
        num_elements: int | Sequence[int] | Int[torch.Tensor, "batch"] | None,
        box_format: BoxFormat | str,
        normalized: bool,
        canvas_size: tuple[int, int] | None,
        conv_mode: ConversationMode | str | None,
        domain_name: str,
    ) -> list[str]:
        prompt_items = self._broadcast_prompt(prompt, batch_size=batch_size)
        content_items = self._broadcast_content(content, batch_size=len(prompt_items))
        count_items = self._resolve_num_elements(
            num_elements=num_elements,
            content_items=content_items,
            batch_size=len(prompt_items),
        )
        prompts: list[str] = []
        for idx, count in enumerate(count_items):
            content_texts = self._texts_for_index(texts, content_items, idx)
            initial = self.processor.build_initial_json(
                labels=cast(
                    Sequence[str | int] | Int[torch.Tensor, "elements"] | None,
                    self._slice_optional(labels, idx),
                ),
                bbox=cast(
                    Float[torch.Tensor, "elements 4"]
                    | Sequence[Sequence[float]]
                    | None,
                    self._slice_optional(bbox, idx),
                ),
                mask=cast(
                    Bool[torch.Tensor, "elements"] | Sequence[bool] | None,
                    self._slice_optional(mask, idx),
                ),
                box_format=box_format,
                normalized=normalized,
                canvas_size=canvas_size,
            )
            prompts.append(
                self.processor.build_prompt(
                    num_elements=count,
                    canvas_size=canvas_size,
                    elements=initial,
                    domain_name=domain_name,
                    conv_mode=conv_mode or self.config.default_conv_mode,
                    prompt=prompt_items[idx],
                    texts=content_texts,
                )
            )
        return prompts

    def _preprocess_images(
        self,
        images: Sequence[Image.Image],
    ) -> Float[torch.Tensor, "batch channels height width"]:
        image_processor = self.image_processor
        if isinstance(image_processor, PosterLlavaImageProcessor):
            processed = image_processor.preprocess(
                list(images),
                image_aspect_ratio=cast(Literal["pad"], self.config.image_aspect_ratio),
                return_tensors="pt",
            )
        elif hasattr(image_processor, "preprocess"):
            processed = cast(_ImagePreprocessor, image_processor).preprocess(
                list(images),
                return_tensors="pt",
            )
        else:
            raise ValueError("image_processor must provide preprocess")

        return cast(
            Float[torch.Tensor, "batch channels height width"],
            processed["pixel_values"],
        )

    def _broadcast_prompt(
        self,
        prompt: str | Sequence[str] | None,
        *,
        batch_size: int,
    ) -> list[str | None]:
        if prompt is None:
            return [None] * batch_size
        if isinstance(prompt, str):
            return [prompt] * batch_size
        return list(prompt)

    def _broadcast_content(
        self,
        content: Mapping[str, PosterLlavaContentValue]
        | Sequence[Mapping[str, PosterLlavaContentValue]]
        | None,
        *,
        batch_size: int,
    ) -> list[Mapping[str, PosterLlavaContentValue]]:
        if content is None:
            return [{} for _ in range(batch_size)]
        if _is_content_mapping(content):
            return [content for _ in range(batch_size)]
        if _is_content_sequence(content):
            return list(content)
        raise TypeError("content must be a mapping, a sequence of mappings, or None")

    def _resolve_num_elements(
        self,
        *,
        num_elements: int | Sequence[int] | Int[torch.Tensor, "batch"] | None,
        content_items: Sequence[Mapping[str, PosterLlavaContentValue]],
        batch_size: int,
    ) -> list[int]:
        if num_elements is None:
            values = [
                item.get("num_elements") or item.get("elements")
                for item in content_items
            ]
            if any(value is None for value in values):
                raise ValueError("num_elements is required for PosterLLaVA generation")

            return [int(cast(int | str, value)) for value in values]
        if isinstance(num_elements, int):
            return [num_elements] * batch_size
        if isinstance(num_elements, torch.Tensor):
            return [int(value) for value in num_elements.tolist()]
        return [int(value) for value in num_elements]

    def _texts_for_index(
        self,
        texts: str | Sequence[str] | Sequence[Sequence[str]] | None,
        content_items: Sequence[Mapping[str, PosterLlavaContentValue]],
        idx: int,
    ) -> str | Sequence[str] | None:
        content_value = content_items[idx].get("texts", content_items[idx].get("text"))
        if texts is None:
            return cast(str | Sequence[str] | None, content_value)
        if isinstance(texts, str):
            return texts
        item = texts[idx]
        return item

    def _slice_optional(
        self,
        value: Int[torch.Tensor, "batch elements"]
        | Float[torch.Tensor, "batch elements 4"]
        | Bool[torch.Tensor, "batch elements"]
        | Sequence[str | int]
        | Sequence[Sequence[float]]
        | Sequence[bool]
        | None,
        idx: int,
    ) -> (
        Int[torch.Tensor, "elements"]
        | Float[torch.Tensor, "elements 4"]
        | Bool[torch.Tensor, "elements"]
        | Sequence[str | int]
        | Sequence[float]
        | Sequence[Sequence[float]]
        | Sequence[bool]
        | str
        | int
        | bool
        | None
    ):
        if value is None:
            return None
        if isinstance(value, torch.Tensor) and value.ndim >= 2:
            return value[idx]
        if isinstance(value, Sequence) and value and isinstance(value[0], Sequence):
            return value[idx]
        return value

class PosterLlavaProcessor(ProcessorMixin):
    """Build PosterLLaVA prompts and decode generated JSON layouts.

    Args:
        tokenizer: Optional LLaVA tokenizer used to insert the image sentinel.
        image_processor: Optional image processor component.
        dataset_name: Poster/content dataset used for known label metadata.
        canvas_size: Canvas size used when public input boxes are pixel based.
        id2label: Persisted known label map.
        prompt_template: JSON instruction body template.
        default_domain_name: Domain phrase inserted into prompts.

    Examples:
        >>> processor = PosterLlavaProcessor.from_config()
        >>> processor.parse_output("[{'label': 'text', 'box': [0, 0, 1, 1]}]")[0]["label"]
        'text'
    """

    attributes = ["tokenizer", "image_processor"]
    tokenizer_class = "AutoTokenizer"
    image_processor_class = "AutoImageProcessor"

    def __init__(
        self,
        tokenizer: PreTrainedTokenizerBase | None = None,
        image_processor: PosterLlavaImageProcessorComponent | None = None,
        dataset_name: DatasetName | str = DatasetName.ad_banner,
        canvas_size: tuple[int, int] = DEFAULT_CANVAS_SIZE,
        id2label: Mapping[int, str] | Mapping[str, str] | None = None,
        prompt_template: str = DEFAULT_PROMPT_TEMPLATE,
        default_domain_name: str = DEFAULT_DOMAIN_NAME,
    ) -> None:
        """Initialize tokenizer handles and layout metadata."""
        dataset = normalize_dataset_name(dataset_name)
        self.tokenizer = tokenizer
        self.image_processor = image_processor
        self.dataset_name = str(dataset)
        self.canvas_size = canvas_size
        self.id2label = {
            int(key): str(value)
            for key, value in (id2label or id2label_for_dataset(dataset)).items()
        }
        self.prompt_template = prompt_template
        self.default_domain_name = default_domain_name

    @classmethod
    def from_config(
        cls,
        *,
        dataset_name: DatasetName | str = DatasetName.ad_banner,
        canvas_size: tuple[int, int] = DEFAULT_CANVAS_SIZE,
        id2label: Mapping[int, str] | Mapping[str, str] | None = None,
        prompt_template: str = DEFAULT_PROMPT_TEMPLATE,
        default_domain_name: str = DEFAULT_DOMAIN_NAME,
    ) -> PosterLlavaProcessor:
        """Construct a metadata-only processor for tests and local smoke checks.

        Args:
            dataset_name: Poster/content dataset key.
            canvas_size: Canvas size used for pixel input normalization.
            id2label: Optional known label map.
            prompt_template: Prompt body template.
            default_domain_name: Default domain phrase.

        Returns:
            Metadata-only processor.
        """
        return cls(
            tokenizer=None,
            image_processor=None,
            dataset_name=dataset_name,
            canvas_size=canvas_size,
            id2label=id2label,
            prompt_template=prompt_template,
            default_domain_name=default_domain_name,
        )

    @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",
        *,
        subfolder: str | None = None,
        **kwargs: str | int | bool | list[int] | dict[str, str],
    ) -> PosterLlavaProcessor:
        """Load processor metadata from a checkpoint directory.

        Args:
            pretrained_model_name_or_path: Root checkpoint path.
            cache_dir: Accepted for Transformers processor compatibility.
            force_download: Accepted for Transformers processor compatibility.
            local_files_only: Accepted for compatibility with pipeline loaders.
            token: Accepted for Transformers processor compatibility.
            revision: Accepted for Transformers processor compatibility.
            subfolder: Optional processor subfolder.
            kwargs: Metadata overrides.

        Returns:
            Loaded processor.

        Raises:
            FileNotFoundError: If ``processor_config.json`` is absent.
        """
        _ = cache_dir, force_download, local_files_only, token, revision
        root = Path(pretrained_model_name_or_path)
        path = root / subfolder if subfolder is not None else root
        config_path = path / PROCESSOR_CONFIG_NAME
        data = json.loads(config_path.read_text())
        data.update(kwargs)
        canvas_size = cast(list[int], data["canvas_size"])
        if len(canvas_size) != 2:
            raise ValueError("canvas_size must contain width and height")

        return cls.from_config(
            dataset_name=cast(str, data["dataset_name"]),
            canvas_size=(canvas_size[0], canvas_size[1]),
            id2label=cast(dict[str, str], data["id2label"]),
            prompt_template=cast(str, data["prompt_template"]),
            default_domain_name=cast(str, data["default_domain_name"]),
        )

    def save_pretrained(
        self,
        save_directory: str | Path,
        push_to_hub: bool = False,
        **kwargs: str | int | bool | None,
    ) -> None:
        """Save processor metadata and optional component processors.

        Args:
            save_directory: Directory to write.
            push_to_hub: Accepted for Transformers processor compatibility.
            kwargs: Additional save options accepted for compatibility.
        """
        _ = push_to_hub, kwargs
        root = Path(save_directory)
        root.mkdir(parents=True, exist_ok=True)
        data = {
            "processor_class": self.__class__.__name__,
            "dataset_name": self.dataset_name,
            "canvas_size": list(self.canvas_size),
            "id2label": {str(key): value for key, value in self.id2label.items()},
            "prompt_template": self.prompt_template,
            "default_domain_name": self.default_domain_name,
        }
        (root / PROCESSOR_CONFIG_NAME).write_text(
            json.dumps(data, indent=2, sort_keys=True) + "\n"
        )

    def build_initial_json(
        self,
        *,
        labels: Sequence[str | int] | Int[torch.Tensor, "elements"] | None = None,
        bbox: Float[torch.Tensor, "elements 4"]
        | Sequence[Sequence[float]]
        | None = None,
        mask: Bool[torch.Tensor, "elements"] | Sequence[bool] | None = None,
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
    ) -> list[PosterLlavaJsonElement]:
        """Build optional initial layout JSON from public layout inputs.

        Args:
            labels: Known labels as strings or integer ids.
            bbox: Optional boxes aligned with labels.
            mask: Optional valid-element mask.
            box_format: Public box format for ``bbox``.
            normalized: Whether ``bbox`` is already normalized.
            canvas_size: Pixel canvas size when ``normalized=False``.

        Returns:
            Initial JSON elements used in the prompt.

        Raises:
            ValueError: If labels and boxes are inconsistently shaped.
        """
        if labels is None:
            return []
        label_items = (
            [int(item) for item in labels.tolist()]
            if isinstance(labels, torch.Tensor)
            else list(labels)
        )
        if bbox is None:
            return [
                {
                    "label": self.id2label.get(label, str(label))
                    if isinstance(label, int)
                    else str(label),
                    "box": [],
                }
                for label in label_items
            ]
        bbox_t, _, mask_t = prepare_layout_tensors(
            bbox=bbox,
            labels=torch.arange(len(label_items)),
            mask=mask,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size or self.canvas_size,
        )
        if bbox_t.shape[1] != len(label_items):
            raise ValueError("labels and bbox must contain the same element count")

        ltrb = self._xywh_to_ltrb(bbox_t[0])
        elements: list[PosterLlavaJsonElement] = []
        for idx, label in enumerate(label_items):
            if not bool(mask_t[0, idx]):
                continue
            label_text = (
                self.id2label.get(label, str(label))
                if isinstance(label, int)
                else str(label)
            )
            elements.append({"label": label_text, "box": ltrb[idx].tolist()})
        return elements

    def build_prompt(
        self,
        *,
        num_elements: int,
        canvas_size: tuple[int, int] | None = None,
        elements: Sequence[PosterLlavaJsonElement]
        | Sequence[Mapping[str, PosterLlavaJsonValue]] = (),
        domain_name: str | None = None,
        conv_mode: ConversationMode | str = ConversationMode.llava_v0,
        prompt: str | None = None,
        texts: str | Sequence[str] | Sequence[Sequence[str]] | None = None,
    ) -> str:
        """Build the LLaVA conversation prompt.

        Args:
            num_elements: Number of requested layout elements.
            canvas_size: Optional canvas metadata. Stored in prompt text only
                when a custom prompt uses it.
            elements: Optional initial layout JSON.
            domain_name: Domain phrase for the default template.
            conv_mode: LLaVA conversation template.
            prompt: Optional user-supplied prompt body override.
            texts: Optional text payload inserted into the default body.

        Returns:
            Full conversation prompt with the ``<image>`` marker.

        Raises:
            ValueError: If ``num_elements`` is not positive.
        """
        if num_elements <= 0:
            raise ValueError("num_elements must be positive")

        element_list = list(elements)
        initial = ""
        if element_list:
            initial = " Initial layout JSON: " + json.dumps(element_list)
        resolution = list(canvas_size or self.canvas_size)
        text_payload = self._format_texts(texts)
        body = prompt or self.prompt_template.format(
            num_elements=num_elements,
            domain_name=domain_name or self.default_domain_name,
            initial_layout=initial,
            initial_json=json.dumps(element_list),
            canvas_size=resolution,
            resolution=resolution,
            texts=text_payload,
        )
        if text_payload and "{texts}" not in self.prompt_template and prompt is None:
            body = f"{body}\nText payload: {text_payload}"
        return self._wrap_conversation(body, conv_mode=conv_mode)

    def __call__(
        self,
        prompt: str | Sequence[str],
        *,
        return_tensors: Literal["pt"] = "pt",
    ) -> BatchEncoding:
        """Tokenize prompts with LLaVA image-token insertion.

        Args:
            prompt: Prompt string or prompt sequence.
            return_tensors: Only ``"pt"`` is supported.

        Returns:
            Batch encoding with ``input_ids`` and ``prompt_text``.

        Raises:
            ValueError: If tokenizer is absent.
        """
        if self.tokenizer is None:
            raise ValueError("tokenizer is required to encode PosterLLaVA prompts")

        prompts = [prompt] if isinstance(prompt, str) else list(prompt)
        encoded = [
            tokenizer_image_token(item, self.tokenizer, return_tensors=return_tensors)
            for item in prompts
        ]
        input_ids = torch.nn.utils.rnn.pad_sequence(
            encoded,
            batch_first=True,
            padding_value=getattr(self.tokenizer, "pad_token_id", 0) or 0,
        )
        attention_mask = input_ids.ne(getattr(self.tokenizer, "pad_token_id", 0) or 0)
        return BatchEncoding(
            {
                "input_ids": input_ids,
                "attention_mask": attention_mask,
                "prompt_text": prompts,
            }
        )

    def parse_output(self, text: str) -> list[PosterLlavaJsonElement]:
        """Parse the first generated JSON-like array span.

        Args:
            text: Decoded LLaVA generation text.

        Returns:
            Parsed element dictionaries.

        Raises:
            ValueError: If no JSON array span can be parsed.
        """
        span = self._extract_json_array(text)
        raw_items = json.loads(span.replace("'", '"'))
        if not isinstance(raw_items, list):
            raise ValueError("PosterLLaVA output JSON must be a list")

        elements: list[PosterLlavaJsonElement] = []
        for item in raw_items:
            if not isinstance(item, Mapping):
                raise ValueError("PosterLLaVA output elements must be objects")

            label = item.get("label")
            box = item.get("box")
            if not isinstance(label, str):
                raise ValueError("PosterLLaVA output element label must be a string")

            if not isinstance(box, Sequence) or len(box) != 4:
                raise ValueError("PosterLLaVA output element box must have four values")

            elements.append(
                {
                    "label": label,
                    "box": [float(value) for value in box],
                }
            )
        return elements

    def decode_layout(
        self,
        text: str | Sequence[str],
        *,
        output_type: OutputType | Literal["dataclass", "dict"] = OutputType.dataclass,
        return_intermediates: bool = False,
        sequences: Int[torch.Tensor, "batch generated_tokens"] | None = None,
        prompts: Sequence[str] | None = None,
    ) -> LayoutGenerationOutput | PosterLlavaOutputDict:
        """Decode generated text into the shared layout output schema.

        Args:
            text: Generated text or batch of generated texts.
            output_type: Output container mode.
            return_intermediates: Whether to include raw text and parser data.
            sequences: Optional generated token ids.
            prompts: Optional prompt texts.

        Returns:
            Layout output dataclass or dictionary.
        """
        texts = [text] if isinstance(text, str) else list(text)
        parsed_batches = [self.parse_output(item) for item in texts]
        label_names = self._batch_label_names(parsed_batches)
        id2label = dict(enumerate(label_names))
        label2id = {label: idx for idx, label in id2label.items()}
        max_len = max((len(item) for item in parsed_batches), default=0) or 1
        bbox_rows: list[Float[torch.Tensor, "elements 4"]] = []
        label_rows: list[Int[torch.Tensor, "elements"]] = []
        mask_rows: list[Bool[torch.Tensor, "elements"]] = []
        parsed_numeric: list[list[ParsedPosterLlavaElement]] = []
        for parsed in parsed_batches:
            numeric: list[ParsedPosterLlavaElement] = [
                {
                    "label": label2id[item["label"]],
                    "label_text": item["label"],
                    "bbox_ltrb": item["box"],
                }
                for item in parsed
            ]
            parsed_numeric.append(numeric)
            boxes = torch.tensor(
                [item["bbox_ltrb"] for item in numeric], dtype=torch.float32
            )
            labels = torch.tensor([item["label"] for item in numeric], dtype=torch.long)
            mask = torch.ones(len(numeric), dtype=torch.bool)

            if len(numeric) == 0:
                boxes = torch.zeros(max_len, 4, dtype=torch.float32)
                labels = torch.zeros(max_len, dtype=torch.long)
                mask = torch.zeros(max_len, dtype=torch.bool)
            elif len(numeric) < max_len:
                pad = max_len - len(numeric)
                boxes = torch.nn.functional.pad(boxes, (0, 0, 0, pad))
                labels = torch.nn.functional.pad(labels, (0, pad))
                mask = torch.nn.functional.pad(mask, (0, pad))

            bbox_rows.append(boxes)
            label_rows.append(labels)
            mask_rows.append(mask)
        raw_ltrb = torch.stack(bbox_rows)

        bbox = self._ltrb_to_xywh(raw_ltrb).clamp(0.0, 1.0)

        intermediates: PosterLlavaIntermediates | None = None
        if return_intermediates:
            intermediates = {
                "generated_text": texts,
                "parsed_json": parsed_batches,
                "parsed_elements": parsed_numeric,
                "id2label_per_example": [
                    dict(enumerate(self._batch_label_names([batch])))
                    for batch in parsed_batches
                ],
            }
            if prompts is not None:
                intermediates["prompts"] = list(prompts)
        output = LayoutGenerationOutput(
            bbox=bbox.float(),
            labels=torch.stack(label_rows).long(),
            mask=torch.stack(mask_rows).bool(),
            id2label=id2label,
            sequences=sequences,
            intermediates=intermediates,
        )
        mode = normalize_output_type(output_type)
        if mode is OutputType.dict:
            return cast(PosterLlavaOutputDict, dict(output))
        return output

    def _extract_json_array(self, text: str) -> str:
        start = text.find("[")
        if start < 0:
            raise ValueError("PosterLLaVA output does not contain a JSON array")

        depth = 0
        in_string: str | None = None
        escaped = False

        for idx, char in enumerate(text[start:], start=start):
            if in_string is not None:
                if escaped:
                    escaped = False
                elif char == "\\":
                    escaped = True
                elif char == in_string:
                    in_string = None
                continue

            if char in {"'", '"'}:
                in_string = char
            elif char == "[":
                depth += 1

            elif char == "]":
                depth -= 1

                if depth == 0:
                    return text[start : idx + 1]

        raise ValueError("PosterLLaVA output contains an unterminated JSON array")

    def _wrap_conversation(
        self,
        body: str,
        *,
        conv_mode: ConversationMode | str,
    ) -> str:
        mode = normalize_conversation_mode(conv_mode)
        image_body = f"{IMAGE_TOKEN}\n{body}"
        if mode is ConversationMode.llava_v0:
            return (
                "A chat between a curious human and an artificial intelligence "
                "assistant. The assistant gives helpful, detailed, and polite "
                "answers to the human's questions.###Human: "
                f"{image_body}###Assistant:"
            )
        return (
            "A chat between a curious human and an artificial intelligence "
            "assistant. The assistant gives helpful, detailed, and polite "
            "answers to the human's questions. USER: "
            f"{image_body} ASSISTANT:"
        )

    def _format_texts(
        self,
        texts: str | Sequence[str] | Sequence[Sequence[str]] | None,
    ) -> str:
        if texts is None:
            return ""
        if isinstance(texts, str):
            return texts
        values: list[str] = []
        for item in texts:
            if isinstance(item, str):
                values.append(item)
            else:
                values.append(", ".join(str(value) for value in item))
        return "; ".join(values)

    def _batch_label_names(
        self,
        batches: Sequence[Sequence[PosterLlavaJsonElement]],
    ) -> list[str]:
        names: list[str] = []
        seen: set[str] = set()
        for batch in batches:
            for item in batch:
                label = item["label"]
                if label not in seen:
                    seen.add(label)
                    names.append(label)
        return names or ["unknown"]

    def _ltrb_to_xywh(
        self,
        bbox: Float[torch.Tensor, "... 4"],
    ) -> Float[torch.Tensor, "... 4"]:
        left, top, right, bottom = bbox.unbind(dim=-1)
        return torch.stack(
            ((left + right) / 2, (top + bottom) / 2, right - left, bottom - top),
            dim=-1,
        )

    def _xywh_to_ltrb(
        self,
        bbox: Float[torch.Tensor, "... 4"],
    ) -> Float[torch.Tensor, "... 4"]:
        x, y, width, height = bbox.unbind(dim=-1)
        return torch.stack(
            (x - width / 2, y - height / 2, x + width / 2, y + height / 2),
            dim=-1,
        )