Release NLD-Uni

.gitattributes

@@ -0,0 +1,3 @@
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+*.py[cod]
+validation-output/

LICENSE

@@ -0,0 +1,73 @@
+                                   NVIDIA License
+
+1. Definitions
+
+“Licensor” means any person or entity that distributes its Work.
+
+“Work” means (a) the original work of authorship made available under this license, 
+which may include software, documentation, or other files, and (b) any additions to or 
+derivative works  thereof  that are made available under this license.
+
+The terms “reproduce,” “reproduction,” “derivative works,” and “distribution” have the 
+meaning as provided under U.S. copyright law; provided, however, that for the purposes 
+of this license, derivative works shall not include works that remain separable from, or 
+merely link (or bind by name) to the interfaces of, the Work.
+
+Works are “made available” under this license by including in or with the Work either 
+(a) a copyright notice referencing the applicability of this license to the Work, or 
+(b) a copy of this license.
+
+2. License Grant
+
+2.1 Copyright Grant. Subject to the terms and conditions of this license, each Licensor 
+grants to you a perpetual, worldwide, non-exclusive, royalty-free, copyright license to 
+use, reproduce, prepare derivative works of, publicly display, publicly perform, 
+sublicense and distribute its Work and any resulting derivative works in any form.
+
+3. Limitations
+
+3.1 Redistribution. You may reproduce or distribute the Work only if (a) you do so under 
+this license, (b) you include a complete copy of this license with your distribution, 
+and (c) you retain without modification any copyright, patent, trademark, or attribution 
+notices that are present in the Work.
+
+3.2 Derivative Works. You may specify that additional or different terms apply to the 
+use, reproduction, and distribution of your derivative works of the Work (“Your Terms”) 
+only if (a) Your Terms provide that the use limitation in Section 3.3 applies to your 
+derivative works, and (b) you identify the specific derivative works that are subject 
+to Your Terms. Notwithstanding Your Terms, this license (including the redistribution 
+requirements in Section 3.1) will continue to apply to the Work itself.
+
+3.3 Use Limitation. The Work and any derivative works thereof only may be used or 
+intended for use non-commercially. Notwithstanding the foregoing, NVIDIA Corporation 
+and its affiliates may use the Work and any derivative works commercially. As used 
+herein, “non-commercially” means for research or evaluation purposes only.
+
+3.4 Patent Claims. If you bring or threaten to bring a patent claim against any 
+Licensor (including any claim, cross-claim or counterclaim in a lawsuit) to enforce any 
+patents that you allege are infringed by any Work, then your rights under this license 
+from such Licensor (including the grant in Section 2.1) will terminate immediately.
+
+3.5 Trademarks. This license does not grant any rights to use any Licensor’s or its 
+affiliates’ names, logos, or trademarks, except as necessary to reproduce the notices 
+described in this license.
+
+3.6 Termination. If you violate any term of this license, then your rights under this 
+license (including the grant in Section 2.1) will terminate immediately.
+
+4. Disclaimer of Warranty.
+
+THE WORK IS PROVIDED “AS IS” WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER 
+EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF MERCHANTABILITY, FITNESS FOR 
+A PARTICULAR PURPOSE, TITLE OR NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY 
+ACTIVITIES UNDER THIS LICENSE. 
+
+5. Limitation of Liability.
+
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER 
+IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE SHALL ANY LICENSOR BE LIABLE TO 
+YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL 
+DAMAGES ARISING OUT OF OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK 
+(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION, LOST PROFITS OR 
+DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER DAMAGES OR LOSSES), EVEN IF THE 
+LICENSOR HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.

README.md

@@ -0,0 +1,227 @@
+---
+library_name: transformers
+pipeline_tag: any-to-any
+tags:
+  - custom_code
+  - image-text-to-text
+  - text-to-image
+  - image-to-image
+---
+
+# Nemotron-Labs-Diffusion-Uni-8B
+
+## Description
+
+NLD-Uni is a unified multimodal model for image understanding and image generation. It combines a diffusion language model with specialized visual components and a mixture-of-transformers design. Understanding and generation use distinct transformer parameter paths within a common multimodal framework.
+
+For image understanding, NLD-Uni encodes images into continuous visual features and generates text responses. For image generation, it models discrete visual tokens through iterative diffusion and decodes those tokens into RGB images. The architecture supports autoregressive and blockwise diffusion decoding for text, as well as image-conditioned generation and editing when enabled by the checkpoint and inference implementation.
+
+This card describes the NLD-Uni architecture generally.
+
+### License/Terms of Use
+
+GOVERNING TERMS: Use of this model is governed by the [NVIDIA Open Source License Agreement](https://huggingface.co/nvidia/NL-Diffusion-Image/raw/main/LICENSE).
+
+## Deployment Geography
+
+Global, subject to the governing release terms.
+
+## Use Case
+
+NLD-Uni is intended for research and evaluation of unified multimodal systems, including:
+
+- Visual question answering and image-grounded reasoning.
+- Document, chart, and scene-text understanding.
+- Visual grounding and interface localization using task-specific prompts and output formats.
+- Text-to-image generation.
+- Image-conditioned generation and editing for checkpoints trained for these tasks.
+
+## Release Date
+
+GitLab 10/17/2026 via [NLD-Uni](https://gitlab-master.nvidia.com/shufanl/nld-uni).
+
+## Model Architecture
+
+**Architecture Type:** Transformer
+
+**Network Architecture:** Diffusion language model with understanding and generation transformer paths, a continuous visual encoder for understanding, and a discrete image tokenizer for generation.
+
+**Language-model foundation:** Nemotron diffusion language model.
+
+## Input
+
+**Input Type(s):** Text and images.
+
+**Input Format(s):** Text strings and image objects or tensors accepted by the supplied processor.
+
+**Input Parameters:** One-dimensional text-token sequences and two-dimensional image content, represented internally as visual features or discrete image tokens according to the task.
+
+**Other Properties Related to Input:** Context length, image resolution, visual-token budget, and preprocessing are checkpoint- and task-dependent. The released inference configuration defines the supported limits; a universal 900-token prompt limit is not specified for the architecture.
+
+## Output
+
+**Output Type(s):**
+
+- Text
+- Image
+
+**Output Format(s):**
+
+- **Text:** `.json` (UTF-8 JSON records containing generated text).
+- **Image:** `.png` (RGB images).
+
+These serialized formats are implemented by `validate.py`. The inference API returns text strings or token IDs for understanding and decoded image objects or tensors for generation. Grounding outputs follow the requested coordinate or location-token format.
+
+**Output Parameters:** One-dimensional token sequences or images with height, width, and three color channels. Tensor layout and pixel range are defined by the inference API.
+
+**Other Properties Related to Output:** Text output length and generated image resolution depend on the checkpoint and inference settings. Configurations for 1,024 × 1,024 image generation are implemented; the release should list its validated resolutions.
+
+## Software Integration
+
+NLD-Uni’s native implementation uses PyTorch with CUDA-accelerated execution. The codebase also includes a dedicated SGLang integration for the image-understanding and text-generation path. This does not establish SGLang support for the image-generation path.
+
+**Runtime Engine(s):**
+
+- PyTorch for native multimodal inference.
+- SGLang with the NLD-Uni/LaVida RADIO-H adapter for supported understanding configurations.
+
+**Reference Platform:** Linux with NVIDIA CUDA GPUs.
+
+**Reference Evaluation Hardware:** NVIDIA H100 (Hopper).
+
+## Model Version(s)
+
+- **NLD-Uni:** Unified image-understanding and image-generation architecture based on LaVida-O.
+
+Training stage, checkpoint revision, tokenizer revision, and inference configuration are part of the identity of an evaluated model. Results from earlier LaVida-O or image-only diffusion models should retain their original model names.
+
+# Training and Evaluation Datasets
+
+## Training Dataset
+
+More than 100M image-text pairs.
+
+**Data Collection Method by dataset**
+
+- Hybrid: Automated, Manually-Collected, Synthetic
+
+**Labeling Method by dataset**
+
+- Hybrid: Automated, Manually-Labeled, Synthetic
+
+## Evaluation Datasets
+
+### Documented Functional Validation Dataset
+
+- **Modalities:** Text and RGB images.
+- **Nature of content:** Synthetic everyday-object, animal, and landscape scenes; text-to-image prompts, edit instructions, and image-understanding questions. The generated images are used as inputs to subsequent editing and understanding tasks.
+- **Linguistic characteristics:** English prompts, edit instructions, questions, and recorded answers.
+- **Collection and labeling:** Three scripted scene cases in `validate.py`; model-generated images and answers were manually inspected for prompt alignment and editing fidelity.
+
+Counts and results are recorded in [VALIDATION.md](VALIDATION.md), [validation.json](validation.json), and [examples/results.json](examples/results.json). These are functional smoke tests, not measurements of benchmark accuracy.
+
+### Benchmark Evaluation Properties
+
+The benchmark suite below covers image and text inputs, including mathematical and scientific questions, documents, charts, scene text, interfaces, referring expressions, and image-generation prompts. OCRBenchV2-EN identifies an English subset; the exact language coverage of the other evaluated splits must be documented from the evaluation records.
+
+### Image Understanding
+
+The evaluation suite should identify the following benchmarks and splits explicitly:
+
+| Benchmark | Evaluation target | Primary reporting requirement |
+|---|---|---|
+| MathVista_MINI | Visual mathematical reasoning | Full 1,000-example accuracy; identify answer-extraction judge |
+| MMMU_DEV_VAL | Multidisciplinary visual reasoning | Report development and validation splits separately |
+| AI2D_TEST | Diagram understanding | Accuracy and image/prompt protocol |
+| ChartQA_TEST | Chart question answering | Relaxed accuracy and subset aggregation |
+| DocVQA_VAL | Document question answering | ANLS |
+| TextVQA_VAL | Scene-text question answering | VQA accuracy with specified normalization |
+| RealWorldQA | Real-world visual understanding | Accuracy |
+| InfoVQA_VAL | Infographic question answering | ANLS |
+| ScienceQA_TEST | Science question answering | Identify the image-containing subset when used |
+| ScreenSpot-Pro | Interface grounding | Point-localization accuracy; separate first-pass and crop-refined results |
+| RefCOCO | Referring-expression grounding | IoU ≥ 0.5 accuracy for val, testA, and testB; disclose the native LOC protocol |
+| OCRBenchV2-EN | OCR and document understanding | Official English category aggregation |
+
+### Image Generation
+
+| Benchmark | Evaluation target | Primary metric |
+|---|---|---|
+| GenEval | Object presence, counts, colors, positions, and attribute binding | Overall and category scores |
+| DPG-Bench | Alignment with detailed compositional prompts | Overall and category scores |
+| MJHQ-30K | Generated-image distribution and quality | FID; any additional preference metric must identify its model and version |
+
+Evaluation should identify the exact checkpoint, prompt processing, image preprocessing, decoding parameters, seed policy, sample count, and evaluator version. Report reasoning and non-reasoning results separately. Tokenizer reconstruction evaluation, if performed, should be identified separately from end-to-end image-generation quality.
+
+## Key Considerations
+
+NLD-Uni can produce incorrect text, inaccurate grounding coordinates, and synthetic images that do not fully follow the prompt. Generated content may reflect biases in training data or contain inappropriate material. Image-conditioned outputs may change details that users intended to preserve.
+
+Applications should be evaluated on their intended tasks and data, with appropriate input and output controls. Model outputs should not be treated as verified factual evidence solely because they are fluent or visually plausible.
+
+## Inference
+
+**Native Engine:** PyTorch.
+
+**Understanding Serving Integration:** SGLang with the model-specific adapter, for validated configurations.
+
+**Reference Precision:** BF16 where supported by the selected implementation.
+
+**Reference Hardware:** NVIDIA H100.
+
+Use the model-specific tokenizer, processor, chat template, and task interface. Image understanding uses the continuous visual encoder; image generation uses the discrete image-token pipeline. Runtime optimizations, including CUDA graphs, should be evaluated with the exact checkpoint and decode configuration before publishing accuracy or speed claims.
+
+## Ethical Considerations
+
+NVIDIA believes trustworthy AI is a shared responsibility. Developers should evaluate NLD-Uni for the intended application, including potential errors, bias, privacy concerns, and misuse of generated content.
+
+Users should have appropriate rights and permissions for input images and other supplied content. Generated or edited images may alter identities, text, object proportions, or scene details. Applications should communicate the synthetic nature of generated content where appropriate.
+
+Users are responsible for model inputs and outputs and for implementing appropriate safeguards before deployment.
+
+Report model quality concerns, security vulnerabilities, or NVIDIA AI concerns through [NVIDIA’s reporting channel](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).
+
+## Environment
+
+Install the dependencies in `requirements.txt` with Python 3.13 and CUDA 12.6. Install PyTorch before building FlashAttention:
+
+```bash
+pip install -r requirements.txt
+pip install --no-build-isolation flash-attn==2.8.0.post2
+```
+
+Validated on an NVIDIA H100 80 GB. The eight BF16 weight shards contain all 1,557 active model tensors. `weights_manifest.json` records their hashes.
+
+## Python inference
+
+```python
+import torch
+from transformers import AutoModelForCausalLM
+
+model = AutoModelForCausalLM.from_pretrained(
+    '/path/to/hf_checkpoint',  # or 'organization/model-name'
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+).to('cuda').eval()
+
+image = model.text_to_image('A red ceramic mug on a wooden table.', seed=42)
+image.save('mug.png')
+
+edited = model.image_edit(image, 'Change the mug to blue.', seed=42)
+edited.save('blue_mug.png')
+
+answer = model.generate('What color is the mug?', image=edited)
+print(answer)
+```
+
+Image inputs accept Pillow images or local file paths. Image methods return Pillow images; `generate` returns text by default. Token IDs can be supplied through `input_ids` for text prompts, and `return_token_ids=True` returns the generated token tensor. The `inference.py` command exposes each task.
+
+## Minimal command-line inference
+
+```bash
+python inference.py /path/to/hf_checkpoint --task t2i --prompt 'A red mug.' --output mug.png
+python inference.py /path/to/hf_checkpoint --task edit --image mug.png --prompt 'Make the mug blue.' --output blue.png
+python inference.py /path/to/hf_checkpoint --task understand --image blue.png --prompt 'What color is the mug?'
+```
+
+Replace the model argument with the published Hugging Face model ID for remote loading. The standalone script only uses Transformers and this repository's model code.

attention.py

@@ -0,0 +1,169 @@
+
+from __future__ import annotations
+from typing import Optional, Tuple
+import torch
+import torch.nn.functional as F
+from torch.nn.attention.flex_attention import BlockMask, flex_attention, create_block_mask
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.processing_utils import Unpack
+from transformers.cache_utils import Cache
+from flash_attn import flash_attn_func, flash_attn_varlen_func
+from .core import Ministral3Attention, _MixedTokenRoutingPlan, _create_split_tower_block_mask, _get_llama_4_attn_scale, _routing_mask, apply_rotary_pos_emb, repeat_kv
+try:
+    from flash_attn import flash_attn_func
+except:
+    pass
+try:
+    from flash_attn import flash_attn_varlen_func
+except:
+    flash_attn_varlen_func = None
+
+@torch.compile(fullgraph=True, dynamic=False)
+def fused_flex_attention(q, k, v, block_mask=None):
+    return flex_attention(q, k, v, block_mask=block_mask)
+_flex_attn_debug_printed = False
+
+def _routed_flash_attention(query_states, key_states, value_states, *, query_selection, key_value_selection, causal):
+    if flash_attn_varlen_func is None:
+        raise RuntimeError('partial dual-tower FlashAttention requires flash_attn_varlen_func')
+    query_tokens = query_states.transpose(1, 2)
+    key_tokens = key_states.transpose(1, 2)
+    value_tokens = value_states.transpose(1, 2)
+    query_unpadded = query_tokens[query_selection]
+    key_unpadded = key_tokens[key_value_selection]
+    value_unpadded = value_tokens[key_value_selection]
+    query_lengths = query_selection.sum(dim=1, dtype=torch.int32)
+    key_value_lengths = key_value_selection.sum(dim=1, dtype=torch.int32)
+    zero = torch.zeros(1, dtype=torch.int32, device=query_states.device)
+    cumulative_query_lengths = torch.cat((zero, query_lengths.cumsum(dim=0, dtype=torch.int32)))
+    cumulative_key_value_lengths = torch.cat((zero, key_value_lengths.cumsum(dim=0, dtype=torch.int32)))
+    output_unpadded = flash_attn_varlen_func(query_unpadded, key_unpadded, value_unpadded, cumulative_query_lengths, cumulative_key_value_lengths, query_selection.shape[1], key_value_selection.shape[1], dropout_p=0.0, softmax_scale=None, causal=causal)
+    output = torch.zeros_like(query_tokens)
+    output[query_selection] = output_unpadded
+    return output.transpose(1, 2)
+
+class MinistralFlexAttention(Ministral3Attention):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.max_seq_length = getattr(self.config, 'max_seq_length', 4096)
+        self.block_size_orig = self.config.block_size
+        self.bidirectional_mask = None
+        self.sbd_block_diff_mask = None
+        self.block_size = self.block_size_orig
+        self.mode = 'sbd_block_diff'
+        self._quadratic_block_mask = {}
+        import torch._dynamo.config as dcfg
+        dcfg.cache_size_limit = 512
+
+    def set_attention_mode(self, mode, block_size=None):
+        self.mode = mode
+        self.block_size = block_size
+
+    def forward(self, hidden_states: torch.Tensor, position_embeddings: Tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_values: Optional[Cache]=None, cache_position: Optional[torch.LongTensor]=None, overwrite_block_mask=None, overwrite_attn_impl=None, use_cache: Optional[bool]=False, routing=None, single_stream_image_diffusion: bool=False, **kwargs: Unpack[FlashAttentionKwargs]) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if overwrite_attn_impl == 'base':
+            return super().forward(hidden_states=hidden_states, position_embeddings=position_embeddings, attention_mask=attention_mask, past_key_values=past_key_values, cache_position=cache_position, use_cache=use_cache, routing=routing, **kwargs)
+        (bsz, q_len, _) = hidden_states.size()
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        (query_states, key_states, value_states) = self._project_qkv(hidden_states, routing)
+        query_states = query_states.view(hidden_shape).transpose(1, 2)
+        key_states = key_states.view(hidden_shape).transpose(1, 2)
+        value_states = value_states.view(hidden_shape).transpose(1, 2)
+        (cos, sin) = position_embeddings
+        (query_states, key_states) = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        query_states = query_states * _get_llama_4_attn_scale(cache_position, self.config.rope_parameters.get('llama_4_scaling_beta'), self.config.rope_parameters.get('original_max_position_embeddings')).to(query_states.dtype)
+        current_key_states = key_states
+        current_value_states = value_states
+        cached_routing_metadata = getattr(past_key_values, '_nlduni_dual_tower_routing', None)
+        isolate_generation_flash = overwrite_attn_impl == 'flash_attn' and self.split_tower and (routing is not None or cached_routing_metadata is not None)
+        cached_has_generation = bool(getattr(past_key_values, '_nlduni_dual_tower_has_generation', False))
+        drop_cached_understanding = isolate_generation_flash and routing is True and (not cached_has_generation)
+        if past_key_values is not None:
+            if use_cache:
+                cache_kwargs = {'sin': sin, 'cos': cos, 'cache_position': cache_position}
+                (key_states, value_states) = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+            elif not drop_cached_understanding:
+                (old_k, old_v) = (past_key_values.layers[self.layer_idx].keys, past_key_values.layers[self.layer_idx].values)
+                key_states = torch.cat([old_k, key_states], dim=-2)
+                value_states = torch.cat([old_v, value_states], dim=-2)
+        self_spec_inference_mode = None
+        if overwrite_attn_impl == 'flash_attn':
+            block_mask = None
+        elif overwrite_block_mask is not None:
+            block_mask = overwrite_block_mask
+            if block_mask == 'full':
+                block_mask = None
+        elif self.mode == 'bidirectional':
+            block_mask = None
+            overwrite_attn_impl = 'flash_attn'
+        elif True:
+            raise ValueError('Supply an explicit block mask or use FlashAttention')
+        if overwrite_block_mask is None and overwrite_attn_impl != 'flash_attn' and self.split_tower and (routing is not None or cached_routing_metadata is not None):
+            current_routing = _routing_mask(routing, hidden_states)
+            cached_routing = cached_routing_metadata
+            if cached_routing is None:
+                cached_routing = torch.zeros(bsz, key_states.shape[-2] - q_len, dtype=torch.bool, device=hidden_states.device)
+            block_mask = _create_split_tower_block_mask(block_mask, current_routing=current_routing, cached_routing=cached_routing)
+        if isinstance(block_mask, BlockMask) and overwrite_attn_impl != 'flash_attn':
+            (mask_query_length, mask_key_value_length) = block_mask.seq_lengths
+            query_length = query_states.shape[-2]
+            key_value_length = key_states.shape[-2]
+            if mask_query_length < query_length:
+                raise ValueError(f'FlexAttention block-mask query length is shorter than the projected queries: {mask_query_length} < {query_length}')
+            if mask_key_value_length < key_value_length:
+                raise ValueError(f'FlexAttention block-mask key/value length is shorter than the projected keys: {mask_key_value_length} < {key_value_length}')
+            query_padding = mask_query_length - query_length
+            key_value_padding = mask_key_value_length - key_value_length
+            if query_padding:
+                query_states = F.pad(query_states, (0, 0, 0, query_padding))
+            if key_value_padding:
+                key_states = F.pad(key_states, (0, 0, 0, key_value_padding))
+                value_states = F.pad(value_states, (0, 0, 0, key_value_padding))
+        if overwrite_attn_impl == 'flash_attn':
+            if self.diffusion_lm:
+                causal = False
+            else:
+                causal = True
+            isolate_generation_flash = self.split_tower and (routing is not None or cached_routing_metadata is not None)
+            if isolate_generation_flash and routing is True and (not cached_has_generation):
+                attn_output = flash_attn_func(query_states.transpose(1, 2), current_key_states.transpose(1, 2), current_value_states.transpose(1, 2), dropout_p=0.0, softmax_scale=None, causal=causal).transpose(1, 2)
+            elif isolate_generation_flash:
+                if routing is True:
+                    query_routing = torch.ones(bsz, q_len, dtype=torch.bool, device=hidden_states.device)
+                elif routing is None:
+                    query_routing = torch.zeros(bsz, q_len, dtype=torch.bool, device=hidden_states.device)
+                elif isinstance(routing, _MixedTokenRoutingPlan):
+                    query_routing = routing.mask
+                else:
+                    raise TypeError('split-tower FlashAttention requires canonical routing')
+                cached_routing = cached_routing_metadata
+                cached_length = key_states.shape[-2] - q_len
+                if cached_routing is None:
+                    cached_routing = torch.zeros(bsz, cached_length, dtype=torch.bool, device=hidden_states.device)
+                elif tuple(cached_routing.shape) != (bsz, cached_length):
+                    raise ValueError('cached dual-tower routing does not align with FlashAttention keys')
+                key_value_routing = torch.cat((cached_routing, query_routing), dim=1)
+                if routing is None:
+                    attn_output = _routed_flash_attention(query_states, key_states, value_states, query_selection=~query_routing, key_value_selection=~key_value_routing, causal=causal)
+                elif routing is True:
+                    attn_output = _routed_flash_attention(query_states, key_states, value_states, query_selection=query_routing, key_value_selection=key_value_routing, causal=causal)
+                else:
+                    attn_output = _routed_flash_attention(query_states, key_states, value_states, query_selection=~query_routing, key_value_selection=~key_value_routing, causal=causal)
+                    attn_output = attn_output + _routed_flash_attention(query_states, key_states, value_states, query_selection=query_routing, key_value_selection=key_value_routing, causal=causal)
+            else:
+                key_states = repeat_kv(key_states, self.num_key_value_groups)
+                value_states = repeat_kv(value_states, self.num_key_value_groups)
+                attn_output = flash_attn_func(query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), dropout_p=0.0, softmax_scale=None, causal=causal).transpose(1, 2)
+        else:
+            key_states = repeat_kv(key_states, self.num_key_value_groups)
+            value_states = repeat_kv(value_states, self.num_key_value_groups)
+            global _flex_attn_debug_printed
+            if not _flex_attn_debug_printed:
+                _flex_attn_debug_printed = True
+            attn_output = fused_flex_attention(query_states, key_states, value_states, block_mask=block_mask)
+        if attn_output.shape[-2] != q_len:
+            attn_output = attn_output[..., :q_len, :]
+        attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
+        attn_output = self._project_output(attn_output, routing)
+        return (attn_output, None)

bridge.py

@@ -0,0 +1,82 @@
+
+from __future__ import annotations
+import torch
+from torch import nn
+IMAGE_LOCAL_AR_CHILD_COUNT = 4
+
+class PackedLinearImageCodec(nn.Module):
+
+    def __init__(self, channels: int, *, norm_eps: float=1e-06, device: torch.device | str | None=None, dtype: torch.dtype | None=None) -> None:
+        super().__init__()
+        if isinstance(channels, bool) or not isinstance(channels, int) or channels < 1:
+            raise ValueError('channels must be a positive integer')
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        self.channels = channels
+        self.down_norm = nn.RMSNorm(channels, eps=norm_eps, elementwise_affine=True, **factory_kwargs)
+        self.down_projection = nn.Linear(4 * channels, channels, bias=False, **factory_kwargs)
+        self.up_norm = nn.RMSNorm(channels, eps=norm_eps, elementwise_affine=True, **factory_kwargs)
+        self.up_projection = nn.Linear(channels, 4 * channels, bias=False, **factory_kwargs)
+
+    @staticmethod
+    def _shape(shape: tuple[int, int] | list[int]) -> tuple[int, int]:
+        if not isinstance(shape, (tuple, list)) or len(shape) != 2:
+            raise ValueError('image token shape must contain height and width')
+        (height, width) = shape
+        if any((isinstance(value, bool) or not isinstance(value, int) for value in shape)):
+            raise TypeError('image token height and width must be integers')
+        if height < 1 or width < 1 or height % 2 or width % 2:
+            raise ValueError('packed-linear image grids must have positive even dimensions')
+        return (height, width)
+
+    @staticmethod
+    def _patches(grid: torch.Tensor) -> torch.Tensor:
+        return torch.stack((grid[..., 0::2, 0::2, :], grid[..., 0::2, 1::2, :], grid[..., 1::2, 0::2, :], grid[..., 1::2, 1::2, :]), dim=-2)
+
+    @classmethod
+    def pack_raster_values(cls, child_values: torch.Tensor, shape: tuple[int, int] | list[int]) -> torch.Tensor:
+        (height, width) = cls._shape(shape)
+        if not isinstance(child_values, torch.Tensor):
+            raise TypeError('child_values must be a tensor')
+        if child_values.ndim < 2 or child_values.shape[1] != height * width:
+            raise ValueError(f'child_values must have shape [B, {height * width}, ...]')
+        grid = child_values.reshape(child_values.shape[0], height, width, *child_values.shape[2:])
+        patches = torch.stack((grid[:, 0::2, 0::2, ...], grid[:, 0::2, 1::2, ...], grid[:, 1::2, 0::2, ...], grid[:, 1::2, 1::2, ...]), dim=3)
+        return patches.flatten(1, 2)
+
+    @classmethod
+    def unpack_patch_values(cls, patch_values: torch.Tensor, shape: tuple[int, int] | list[int]) -> torch.Tensor:
+        (height, width) = cls._shape(shape)
+        (parent_height, parent_width) = (height // 2, width // 2)
+        parent_count = parent_height * parent_width
+        if not isinstance(patch_values, torch.Tensor):
+            raise TypeError('patch_values must be a tensor')
+        if patch_values.ndim < 3 or patch_values.shape[1] != parent_count or patch_values.shape[2] != IMAGE_LOCAL_AR_CHILD_COUNT:
+            raise ValueError(f'patch_values must have shape [B, {parent_count}, {IMAGE_LOCAL_AR_CHILD_COUNT}, ...]')
+        patches = patch_values.reshape(patch_values.shape[0], parent_height, parent_width, IMAGE_LOCAL_AR_CHILD_COUNT, *patch_values.shape[3:])
+        grid = patch_values.new_empty(patch_values.shape[0], height, width, *patch_values.shape[3:])
+        grid[:, 0::2, 0::2, ...] = patches[:, :, :, 0, ...]
+        grid[:, 0::2, 1::2, ...] = patches[:, :, :, 1, ...]
+        grid[:, 1::2, 0::2, ...] = patches[:, :, :, 2, ...]
+        grid[:, 1::2, 1::2, ...] = patches[:, :, :, 3, ...]
+        return grid.flatten(1, 2)
+
+    def compress_fixed_grid(self, child_states: torch.Tensor, shape: tuple[int, int] | list[int]) -> torch.Tensor:
+        (height, width) = self._shape(shape)
+        if child_states.ndim != 3 or child_states.shape[1:] != (height * width, self.channels):
+            raise ValueError(f'child_states must have shape [B, {height * width}, {self.channels}]')
+        grid = child_states.reshape(child_states.shape[0], height, width, self.channels)
+        patches = self._patches(grid)
+        normalized = self.down_norm(patches)
+        parents = self.down_projection(normalized.reshape(child_states.shape[0], height // 2, width // 2, 4 * self.channels))
+        return parents.flatten(1, 2)
+
+    def expand_fixed_grid(self, parent_states: torch.Tensor, shape: tuple[int, int] | list[int]) -> torch.Tensor:
+        children = self.expand_fixed_grid_patches(parent_states, shape)
+        return self.unpack_patch_values(children, shape)
+
+    def expand_fixed_grid_patches(self, parent_states: torch.Tensor, shape: tuple[int, int] | list[int]) -> torch.Tensor:
+        (height, width) = self._shape(shape)
+        parent_count = height // 2 * (width // 2)
+        if parent_states.ndim != 3 or parent_states.shape[1:] != (parent_count, self.channels):
+            raise ValueError(f'parent_states must have shape [B, {parent_count}, {self.channels}]')
+        return self.up_projection(self.up_norm(parent_states)).reshape(parent_states.shape[0], parent_count, 4, self.channels)

chat_template.jinja

@@ -0,0 +1,209 @@
+{% macro render_extra_keys(json_dict, handled_keys) %}
+    {%- if json_dict is mapping %}
+        {%- for json_key in json_dict if json_key not in handled_keys %}
+            {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
+                {{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '</' ~ json_key ~ '>' }}
+            {%- else %}
+                {{-'\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '</' ~ json_key ~ '>' }}
+            {%- endif %}
+        {%- endfor %}
+    {%- endif %}
+{% endmacro %}
+{%- set enable_thinking = enable_thinking if enable_thinking is defined else True %}
+{%- set reasoning_budget = reasoning_budget if reasoning_budget is defined else None %}
+{%- set truncate_history_thinking = truncate_history_thinking if truncate_history_thinking is defined else True %}
+
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- set loop_messages = messages %}
+{%- for m in loop_messages %}
+  {%- if m["role"] == "user" %}
+    {%- set ns.last_user_idx = loop.index0 %}
+  {%- endif %}
+{%- endfor %}
+
+{%- if messages[0]["role"] == "system" %}
+    {%- set system_message = messages[0]["content"] %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set system_message = "" %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+{%- if not tools is defined %}
+    {%- set tools = [] %}
+{%- endif %}
+{# Recompute last_user_idx relative to loop_messages after handling system #}
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- for m in loop_messages %}
+  {%- if m["role"] == "user" %}
+    {%- set ns.last_user_idx = loop.index0 %}
+  {%- endif %}
+{%- endfor %}
+{%- if system_message is defined %}
+    {{- "<|im_start|>system\n" + system_message }}
+{%- else %}
+    {%- if tools is iterable and tools | length > 0 %}
+        {{- "<|im_start|>system\n" }}
+    {%- endif %}
+{%- endif %}
+{%- if tools is iterable and tools | length > 0 %}
+    {%- if system_message is defined and system_message | length > 0 %}
+        {{- "\n\n" }}
+    {%- endif %}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n" }}
+    {{- "<tools>" }}
+    {%- for tool in tools %}
+        {%- if tool.function is defined %}
+            {%- set tool = tool.function %}
+        {%- endif %}
+        {{- "\n<function>\n<name>" ~ tool.name ~ "</name>" }}
+        {%- if tool.description is defined %}
+            {{- '\n<description>' ~ (tool.description | trim) ~ '</description>' }}
+        {%- endif %}
+        {{- '\n<parameters>' }}
+        {%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}
+            {%- for param_name, param_fields in tool.parameters.properties|items %}
+                {{- '\n<parameter>' }}
+                {{- '\n<name>' ~ param_name ~ '</name>' }}
+                {%- if param_fields.type is defined %}
+                    {{- '\n<type>' ~ (param_fields.type | string) ~ '</type>' }}
+                {%- endif %}
+                {%- if param_fields.description is defined %}
+                    {{- '\n<description>' ~ (param_fields.description | trim) ~ '</description>' }}
+                {%- endif %}
+                {%- if param_fields.enum is defined %}
+                    {{- '\n<enum>' ~ (param_fields.enum | tojson | safe) ~ '</enum>' }}
+                {%- endif %}
+                {%- set handled_keys = ['name', 'type', 'description', 'enum'] %}
+                {{- render_extra_keys(param_fields, handled_keys) }}
+                {{- '\n</parameter>' }}
+            {%- endfor %}
+        {%- endif %}
+        {% set handled_keys = ['type', 'properties', 'required'] %}
+        {{- render_extra_keys(tool.parameters, handled_keys) }}
+        {%- if tool.parameters is defined and tool.parameters.required is defined %}
+            {{- '\n<required>' ~ (tool.parameters.required | tojson | safe) ~ '</required>' }}
+        {%- endif %}
+        {{- '\n</parameters>' }}
+        {%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}
+        {{- render_extra_keys(tool, handled_keys) }}
+        {{- '\n</function>' }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+{%- endif %}
+
+
+{%- if system_message is defined %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if tools is iterable and tools | length > 0 %}
+        {{- '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+
+{%- for message in loop_messages %}
+    {%- if message.role == "assistant" %}
+        {# Add reasoning content in to content field for unified processing below. #}
+        {%- if message.reasoning_content is defined and message.reasoning_content is string and message.reasoning_content | trim | length > 0 %}
+            {%- set content = "<think>\n" ~ message.reasoning_content ~ "\n</think>\n" ~ (message.content | default('', true)) %}
+        {%- else %}
+            {%- set content = message.content | default('', true) %}
+            {%- if content is string -%}
+                {# Allow downstream logic to to take care of broken thought, only handle coherent reasoning here. #}
+                {%- if '<think>' not in content and '</think>' not in content -%}
+                    {%- set content = "<think></think>" ~ content -%}
+                {%- endif -%}
+            {%- else -%}
+                {%- set content = content -%}
+            {%- endif -%}
+        {%- endif %}
+        {%- if message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %}
+            {# Assistant message has tool calls. #}
+            {{- '<|im_start|>assistant\n' }}
+                {%- set include_content = not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+                {%- if content is string and content | trim | length > 0 %}
+                    {%- if include_content %}
+                        {{- (content | trim) ~ '\n' -}}
+                    {%- else %}
+                        {%- set c = (content | string) %}
+                        {%- if '</think>' in c %}
+                            {# Keep only content after the last closing think. Also generation prompt causes this. #}
+                            {%- set c = c.split('</think>')[-1] %}
+                        {%- elif '<think>' in c %}
+                            {# If <think> was opened but never closed, drop the trailing think segment #}
+                            {%- set c = c.split('<think>')[0] %}
+                        {%- endif %}
+                        {%- set c = "<think></think>" ~ c | trim %}
+                        {%- if c | length > 0 %}
+                            {{- c ~ '\n' -}}
+                        {%- endif %}
+                    {%- endif %}
+                {%- else %}
+                    {{- "<think></think>" -}}
+                {%- endif %}
+                {%- for tool_call in message.tool_calls %}
+                    {%- if tool_call.function is defined %}
+                        {%- set tool_call = tool_call.function %}
+                    {%- endif %}
+                    {{- '<tool_call>\n<function=' ~ tool_call.name ~ '>\n' -}}
+                        {%- if tool_call.arguments is defined %}
+                            {%- for args_name, args_value in tool_call.arguments|items %}
+                                {{- '<parameter=' ~ args_name ~ '>\n' -}}
+                                    {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                                {{- args_value ~ '\n</parameter>\n' -}}
+                            {%- endfor %}
+                        {%- endif %}
+                    {{- '</function>\n</tool_call>\n' -}}
+                {%- endfor %}
+                {{- '<|im_end|>\n' }}
+        {%- else %}
+            {# Assistant message doesn't have tool calls. #}
+            {%- if not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+                {{- '<|im_start|>assistant\n' ~ (content | default('', true) | string | trim) ~ '<|im_end|>\n' }}
+            {%- else %}
+                {%- set c = (content | default('', true) | string) %}
+                {%- if '<think>' in c and '</think>' in c %}
+                    {%- set c = "<think></think>" ~ c.split('</think>')[-1] %}
+                {%- endif %}
+                {%- set c = c | trim %}
+                {%- if c | length > 0 %}
+                    {{- '<|im_start|>assistant\n' ~ c ~ '<|im_end|>\n' }}
+                {%- else %}
+                    {{- '<|im_start|>assistant\n<|im_end|>\n' }}
+                {%- endif %}
+            {%- endif %}
+        {%- endif %}
+    {%- elif message.role == "user" or message.role == "system" %}
+        {{- '<|im_start|>' + message.role + '\n' }}
+        {%- set content = message.content | string %}
+        {%- if message.role == "user" and loop.index0 == ns.last_user_idx and reasoning_budget is not none %}
+            {{- content + '\n\n{thinking token budget: ' + (reasoning_budget | string) + '}' }}
+        {%- else %}
+            {{- content }}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user\n' }}
+        {%- endif %}
+        {{- '<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>\n' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endfor %}
+
+{%- if add_generation_prompt %}
+    {%- if enable_thinking %}
+        {{- '<|im_start|>assistant\n<think>\n' }}
+    {%- else %}
+        {{- '<|im_start|>assistant\n<think></think>' }}
+    {%- endif %}
+{%- endif %}

chat_utils.py

@@ -0,0 +1,342 @@
+
+from __future__ import annotations
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.nn.attention.flex_attention import create_block_mask
+from transformers.utils import ModelOutput
+from dataclasses import dataclass
+from transformers.cache_utils import Cache
+
+@dataclass
+class SimpleOutputWithPast(ModelOutput):
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    causal_logits: torch.FloatTensor | None = None
+    past_key_values: Cache | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+from .utils import maybe_truncate_last_dim, pad_along_last_dim
+from .configuration_nlduni import get_ministral_shift_sbd_logits
+from .sbd_attention import build_image_attention_doc_ids
+
+def _resolve_shift_sbd_logits(model, shift_logits):
+    model_module = model.module if hasattr(model, 'module') else model
+    config = getattr(model_module, 'config', None)
+    configured = get_ministral_shift_sbd_logits(config) if config is not None else None
+    has_alignment_contract = config is not None
+    if shift_logits is None:
+        return bool(configured) if configured is not None else False
+    requested = bool(shift_logits)
+    if has_alignment_contract and requested != configured:
+        raise ValueError('Inference shift_logits must match config.shift_sbd_logits to avoid a one-token alignment error. Update the model config explicitly if this override is intentional.')
+    return requested
+
+def wte(model, x, t2i_inference=False, gen_shape=None, x_gen=None, inputs_embeds_curr=None, new_token_mask=None, x_gen_child_embeddings=None, gen_parent_mask=None, gen_source_null_mask=None, gen_role=None):
+    if t2i_inference:
+        if (x_gen is None) == (x_gen_child_embeddings is None):
+            raise ValueError('provide exactly one of x_gen or x_gen_child_embeddings')
+        if new_token_mask is None:
+            new_token_mask = x >= INT_MAX
+        if gen_parent_mask is not None or gen_source_null_mask is not None or gen_role is not None:
+            raise ValueError('structured generation masks and roles require Infinity BSQ')
+        gen_latents_comp_embeds = model.call_gen_embedding(x_gen, gen_shape, child_embeddings=x_gen_child_embeddings)
+        if inputs_embeds_curr is None:
+            x_txt_only = x.clone()
+            x_txt_only[new_token_mask] = 0
+            inputs_embeds_curr = model.embed_tokens(x_txt_only)
+        inputs_embeds_curr[new_token_mask] = pad_along_last_dim(gen_latents_comp_embeds, inputs_embeds_curr.shape[-1]).view(-1, inputs_embeds_curr.shape[-1])
+    else:
+        inputs_embeds_curr = model.embed_tokens(x)
+        new_token_mask = None
+    return (inputs_embeds_curr, new_token_mask)
+INT_MAX = 1000000
+SBD_FLEX_ATTENTION_PAD_MULTIPLE = 512
+
+def _round_up_to_multiple(length, multiple):
+    if isinstance(multiple, bool) or not isinstance(multiple, int) or multiple < 1:
+        raise ValueError('padding multiple must be a positive integer')
+    return (length + multiple - 1) // multiple * multiple
+
+def single_stream_sbd_sampling_mask(block_size, past_length, b, h, q_idx, kv_idx, query_image_doc_id, key_value_image_doc_id):
+    del block_size, h
+    absolute_q = q_idx + past_length
+    token_causal = absolute_q >= kv_idx
+    image_query = query_image_doc_id[b, q_idx] > 0
+    same_image_doc = image_query & (query_image_doc_id[b, q_idx] == key_value_image_doc_id[b, kv_idx])
+    return token_causal | same_image_doc
+
+def _create_single_stream_sbd_sampling_block_mask(input_embeddings, *, past_key_values, block_size, input_modality_indices, generation_token_mask=None, pad_to_multiple=None):
+    if isinstance(block_size, bool) or not isinstance(block_size, int) or block_size < 1:
+        raise ValueError('SBD diffusion sampling requires a positive integer block_size')
+    query_length = input_embeddings.shape[1]
+    past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+    key_value_length = past_length + query_length
+    padded_query_length = _round_up_to_multiple(query_length, pad_to_multiple) if pad_to_multiple is not None else query_length
+    padded_key_value_length = _round_up_to_multiple(key_value_length, pad_to_multiple) if pad_to_multiple is not None else key_value_length
+    if input_modality_indices is None:
+        raise ValueError('SBD diffusion sampling requires input_modality_indices to route same-block generation attention')
+    if tuple(input_modality_indices.shape) != tuple(input_embeddings.shape[:2]):
+        raise ValueError('SBD diffusion sampling requires input_modality_indices to align with the current input embeddings')
+    query_transformer_routing = input_modality_indices.to(dtype=torch.bool)
+    if generation_token_mask is None:
+        generation_token_mask = query_transformer_routing
+    elif tuple(generation_token_mask.shape) != tuple(input_embeddings.shape[:2]):
+        raise ValueError('SBD diffusion sampling requires generation_token_mask to align with the current input embeddings')
+    query_target_generation_routing = generation_token_mask.to(dtype=torch.bool)
+    cached_generation_routing = torch.zeros(query_transformer_routing.shape[0], past_length, dtype=torch.bool, device=query_transformer_routing.device)
+    padded_generation_routing = torch.zeros(query_transformer_routing.shape[0], padded_key_value_length - key_value_length, dtype=torch.bool, device=query_transformer_routing.device)
+    key_value_target_generation_routing = torch.cat((cached_generation_routing, query_target_generation_routing, padded_generation_routing), dim=1)
+    key_value_valid = (torch.arange(padded_key_value_length, device=query_transformer_routing.device).unsqueeze(0) < key_value_length).expand(query_transformer_routing.shape[0], -1).contiguous()
+    key_value_image_doc_id = build_image_attention_doc_ids(key_value_target_generation_routing, valid_tokens=key_value_valid)
+    query_image_doc_id = torch.cat((key_value_image_doc_id[:, past_length:past_length + query_length], torch.zeros(query_transformer_routing.shape[0], padded_query_length - query_length, dtype=key_value_image_doc_id.dtype, device=query_transformer_routing.device)), dim=1)
+    past_length_tensor = torch.tensor(past_length, dtype=torch.long, device=query_transformer_routing.device)
+    query_valid = (torch.arange(padded_query_length, device=query_transformer_routing.device).unsqueeze(0) < query_length).expand(query_transformer_routing.shape[0], -1).contiguous()
+
+    def mask_mod(b, h, q_idx, kv_idx):
+        logical_attention = single_stream_sbd_sampling_mask(block_size, past_length_tensor, b, h, q_idx, kv_idx, query_image_doc_id, key_value_image_doc_id) & query_valid[b, q_idx] & key_value_valid[b, kv_idx]
+        padded_query_attention = ~query_valid[b, q_idx] & (kv_idx == 0)
+        return logical_attention | padded_query_attention
+    return create_block_mask(mask_mod, B=input_embeddings.shape[0], H=None, Q_LEN=padded_query_length, KV_LEN=padded_key_value_length, device=input_embeddings.device)
+
+def get_logits(model, input_emnbeddings, modality_indices=None, t2i_inference=False, past_key_values=None, gen_shape=None, timesteps=None, input_modality_indices=None, generation_attention_mask=None, sbd_diffusion_sampling=False, full_attention_approx=False, loop_k=None):
+    if t2i_inference:
+        resolved_loop_k = None
+        if input_modality_indices is None:
+            input_modality_indices = modality_indices
+        attention_kwargs = {}
+        if full_attention_approx:
+            attention_kwargs['overwrite_block_mask'] = 'full'
+            attention_kwargs['overwrite_attn_impl'] = 'flash_attn'
+        elif sbd_diffusion_sampling:
+            if generation_attention_mask is None:
+                generation_attention_mask = modality_indices
+            attention_kwargs['overwrite_block_mask'] = _create_single_stream_sbd_sampling_block_mask(input_emnbeddings, past_key_values=past_key_values, block_size=getattr(model.config, 'block_size', None), input_modality_indices=input_modality_indices, generation_token_mask=generation_attention_mask, pad_to_multiple=SBD_FLEX_ATTENTION_PAD_MULTIPLE)
+        else:
+            attention_kwargs['overwrite_attn_impl'] = 'flash_attn'
+        output = model(None, input_embeddings=input_emnbeddings, modality_indices=input_modality_indices, output_hidden_states=True, past_key_values=past_key_values, use_cache=None, single_stream_image_diffusion=True, loop_k=None, **attention_kwargs)
+        hidden_states = getattr(output, 'last_hidden_state', None)
+        if hidden_states is None:
+            hidden_states = output.hidden_states[-1]
+        gen_hidden_states = hidden_states[modality_indices]
+        gen_hidden_states = maybe_truncate_last_dim(gen_hidden_states, model.config.d_model_gen)
+        gen_logits = model.call_gen_predictor(gen_hidden_states, gen_shape, timesteps=timesteps)
+        seq_len_per_img = np.prod(gen_shape)
+        if len(gen_logits.shape) == 2:
+            gen_logits = gen_logits.view(-1, seq_len_per_img, gen_logits.shape[-1])
+        else:
+            gen_logits = gen_logits.view(-1, seq_len_per_img, *gen_logits.shape[-2:])
+        return gen_logits
+    else:
+        raise ValueError('get_logits is the image sampler helper')
+
+def add_gumbel_noise(logits, temperature):
+    if temperature == 0:
+        return logits
+    logits = logits.to(torch.float64)
+    noise = torch.rand_like(logits, dtype=torch.float64)
+    gumbel_noise = (-torch.log(noise)) ** temperature
+    return logits.exp() / gumbel_noise
+
+def get_transfer_index(logits, temperature, remasking, mask_index, x, num_transfer_tokens, threshold=None, neg_entropy=False):
+    logits_with_noise = add_gumbel_noise(logits, temperature=temperature)
+    x0 = torch.argmax(logits_with_noise, dim=-1)
+    if remasking == 'low_confidence':
+        p = F.softmax(logits, dim=-1)
+        x0_p = torch.squeeze(torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)), -1)
+    elif remasking == 'top_p_margin':
+        p = F.softmax(logits, dim=-1)
+        top2 = torch.topk(p, k=2, dim=-1).values
+        margin = top2[..., 0] - top2[..., 1]
+        plus_inf = torch.full_like(margin, 1e309)
+        minus_inf = torch.full_like(margin, -1e309)
+        masked_for_min = torch.where(mask_index, margin, plus_inf)
+        masked_for_max = torch.where(mask_index, margin, minus_inf)
+        row_min = masked_for_min.amin(dim=1, keepdim=True)
+        row_max = masked_for_max.amax(dim=1, keepdim=True)
+        denom = row_max - row_min
+        normalized = torch.zeros_like(margin)
+        nonzero = denom > 0
+        normalized = torch.where(mask_index & nonzero, (margin - row_min) / (denom + 1e-12), normalized)
+        normalized = torch.where(mask_index & ~nonzero, torch.ones_like(normalized), normalized)
+        x0_p = normalized
+    elif remasking == 'random':
+        x0_p = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device)
+    else:
+        raise NotImplementedError(remasking)
+    if neg_entropy:
+        p = F.softmax(logits, dim=-1)
+        epsilon = 1e-10
+        log_probs = torch.log(p + epsilon)
+        confidence_scores = torch.sum(p * log_probs, dim=-1)
+    else:
+        confidence_scores = x0_p
+    x0 = torch.where(mask_index, x0, x)
+    confidence = torch.where(mask_index, confidence_scores, -np.inf)
+    transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device)
+    if threshold is not None:
+        num_transfer_tokens = mask_index.sum(dim=1, keepdim=True)
+    for j in range(confidence.shape[0]):
+        (_, select_index) = torch.topk(confidence[j], k=num_transfer_tokens[j])
+        transfer_index[j, select_index] = True
+        if threshold is not None:
+            for k in range(1, num_transfer_tokens[j]):
+                if confidence[j, select_index[k]] < threshold:
+                    transfer_index[j, select_index[k]] = False
+    return (x0, transfer_index)
+
+def get_num_transfer_tokens(mask_index, steps: int):
+    mask_num = mask_index.sum(dim=1, keepdim=True)
+    base = mask_num // steps
+    remainder = mask_num % steps
+    num_transfer_tokens = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.int64) + base
+    for i in range(mask_num.size(0)):
+        num_transfer_tokens[i, :int(remainder[i])] += 1
+    return num_transfer_tokens
+
+def simple_fwd(model, input_ids=None, inputs_embeds=None, attention_mask=None, position_ids=None, past_key_values=None, **kwargs):
+    enc_out = model.encoder(past_key_values=past_key_values, input_ids=input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, overwrite_attn_impl='flash_attn', **kwargs)
+    logits = model.diffusion_head(enc_out.last_hidden_state)
+    return SimpleOutputWithPast(loss=logits, logits=logits, causal_logits=None, past_key_values=enc_out.past_key_values, hidden_states=None, attentions=None)
+
+@torch.no_grad()
+def generate_with_prefix_cache_block_diff(model, prompt=None, prompt_embeds=None, steps=128, gen_length=128, block_length=128, temperature=0.0, remasking='low_confidence', mask_id=126336, threshold=None, factor=None, shift_logits=None, neg_entropy=False, causal_context=False, eos_token_id=None, max_thinking_tokens=None, end_think_token_id=None):
+    shift_logits = _resolve_shift_sbd_logits(model, shift_logits)
+    dream_style = shift_logits
+    if (prompt is None) == (prompt_embeds is None):
+        raise ValueError('Exactly one of `prompt` or `prompt_embeds` must be provided.')
+    if prompt is not None:
+        prompt_ids = prompt
+        prompt_len = prompt_ids.shape[1]
+        x_accum = prompt_ids.clone()
+        B = prompt_ids.shape[0]
+        token_device = prompt_ids.device
+        token_dtype = prompt_ids.dtype
+    else:
+        prompt_ids = None
+        prompt_len = prompt_embeds.shape[1]
+        B = prompt_embeds.shape[0]
+        token_device = prompt_embeds.device
+        token_dtype = torch.long
+        x_accum = torch.full((B, prompt_len), mask_id, dtype=token_dtype, device=token_device)
+    if gen_length % block_length != 0:
+        gen_length = (gen_length // block_length + 1) * block_length
+    num_blocks = gen_length // block_length
+    assert steps % num_blocks == 0
+    steps_per_block = steps // num_blocks
+    nfe = 0
+    model_module = model.module if hasattr(model, 'module') else model
+    for layer in model_module.encoder.layers:
+        layer.self_attn.mode = 'bidirectional'
+    if causal_context:
+        for layer in model_module.encoder.layers:
+            if hasattr(layer.self_attn, 'diffusion_lm'):
+                layer.self_attn.diffusion_lm = False
+    output = simple_fwd(model, input_ids=prompt_ids, inputs_embeds=prompt_embeds, use_cache=True, use_causal_mask=causal_context)
+    past_key_values = output.past_key_values
+    if causal_context:
+        for layer in model_module.encoder.layers:
+            if hasattr(layer.self_attn, 'diffusion_lm'):
+                layer.self_attn.diffusion_lm = True
+    next_token = None
+    if causal_context:
+        last_logit = output.logits[:, -1, :]
+        if temperature > 0:
+            probs = torch.softmax(last_logit / temperature, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+        else:
+            next_token = torch.argmax(last_logit, dim=-1, keepdim=True)
+    next_logits_context = None
+    if dream_style:
+        next_logits_context = output.logits[:, -1:, :]
+    for num_block in range(num_blocks):
+        mask_block = torch.ones((B, block_length), dtype=token_dtype, device=token_device) * mask_id
+        if causal_context:
+            mask_block[:, 0] = next_token[:, 0]
+        x_accum = torch.cat([x_accum, mask_block], dim=1)
+        current_block_start = prompt_len + num_block * block_length
+        block_slice = slice(current_block_start, current_block_start + block_length)
+        if end_think_token_id is not None and max_thinking_tokens is not None:
+            tokens_before_block = num_block * block_length
+            tokens_after_block = tokens_before_block + block_length
+            if tokens_after_block > max_thinking_tokens:
+                gen_so_far = x_accum[:, prompt_len:current_block_start]
+                has_end_think = (gen_so_far == end_think_token_id).any(dim=1) if gen_so_far.size(1) > 0 else torch.zeros(B, dtype=torch.bool, device=token_device)
+                if not has_end_think.all():
+                    if tokens_before_block < max_thinking_tokens:
+                        offset = max_thinking_tokens - tokens_before_block
+                    else:
+                        offset = 0
+                    inject_pos = current_block_start + offset
+                    for b in range(B):
+                        if not has_end_think[b]:
+                            x_accum[b, inject_pos] = end_think_token_id
+        mask_block_idx0 = x_accum[:, block_slice] == mask_id
+        if dream_style:
+            schedule_mask = mask_block_idx0
+        else:
+            schedule_mask = mask_block_idx0
+        num_transfer_tokens = get_num_transfer_tokens(schedule_mask, steps_per_block)
+        for i in range(steps_per_block):
+            mask_block_idx = x_accum[:, block_slice] == mask_id
+            if mask_block_idx.sum() == 0:
+                break
+            nfe += 1
+            logits_block = simple_fwd(model, x_accum[:, block_slice], past_key_values=past_key_values, use_cache=False).logits
+            if dream_style:
+                if block_length == 1:
+                    logits_use = next_logits_context
+                else:
+                    logits_use = torch.cat([next_logits_context, logits_block[:, :-1, :]], dim=1)
+                mask_use = mask_block_idx
+                x_use = x_accum[:, block_slice]
+                (x0, transfer_idx) = get_transfer_index(logits_use, temperature, remasking, mask_use, x_use, num_transfer_tokens=num_transfer_tokens[:, i], threshold=threshold, neg_entropy=neg_entropy)
+                cur = x_accum[:, block_slice].clone()
+                cur[transfer_idx] = x0[transfer_idx]
+                x_accum[:, block_slice] = cur
+            else:
+                (x0, transfer_idx) = get_transfer_index(logits_block, temperature, remasking, mask_block_idx, x_accum[:, block_slice], num_transfer_tokens=num_transfer_tokens[:, i], threshold=threshold, neg_entropy=neg_entropy)
+                cur = x_accum[:, block_slice].clone()
+                cur[transfer_idx] = x0[transfer_idx]
+                x_accum[:, block_slice] = cur
+            if eos_token_id is not None:
+                block_tokens = x_accum[:, block_slice]
+                eos_mask = block_tokens == eos_token_id
+                any_eos = eos_mask.any(dim=1)
+                if any_eos.any():
+                    after_eos = eos_mask.cumsum(dim=1).bool()
+                    mask_before = (block_tokens == mask_id) & ~after_eos
+                    if (any_eos & ~mask_before.any(dim=1)).any():
+                        break
+        if causal_context:
+            for layer in model_module.encoder.layers:
+                if hasattr(layer.self_attn, 'diffusion_lm'):
+                    layer.self_attn.diffusion_lm = False
+        output = simple_fwd(model, x_accum[:, block_slice], past_key_values=past_key_values, use_cache=True, use_causal_mask=causal_context)
+        past_key_values = output.past_key_values
+        nfe += 1
+        if causal_context:
+            for layer in model_module.encoder.layers:
+                if hasattr(layer.self_attn, 'diffusion_lm'):
+                    layer.self_attn.diffusion_lm = True
+            last_logit = output.logits[:, -1, :]
+            if temperature > 0:
+                probs = torch.softmax(last_logit / temperature, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+            else:
+                next_token = torch.argmax(last_logit, dim=-1, keepdim=True)
+        if dream_style and num_block < num_blocks - 1:
+            next_logits_context = output.logits[:, -1:, :]
+        if eos_token_id is not None:
+            gen_so_far = x_accum[:, prompt_len:]
+            is_eos = gen_so_far == eos_token_id
+            has_eos = is_eos.any(dim=1)
+            if has_eos.all():
+                first_eos_pos = is_eos.to(torch.int64).argmax(dim=1)
+                max_eos = first_eos_pos.max().item()
+                if prompt_ids is None:
+                    return (x_accum[:, prompt_len:prompt_len + max_eos + 1], nfe)
+                return (x_accum[:, :prompt_len + max_eos + 1], nfe)
+    if prompt_ids is None:
+        return (x_accum[:, prompt_len:], nfe)
+    return (x_accum, nfe)

config.json

@@ -0,0 +1,36 @@
+{
+  "hidden_size": 4096,
+  "intermediate_size": 14336,
+  "num_hidden_layers": 34,
+  "num_attention_heads": 32,
+  "num_key_value_heads": 8,
+  "head_dim": 128,
+  "vocab_size": 132101,
+  "max_position_embeddings": 131072,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "beta_fast": 32.0,
+    "beta_slow": 1.0,
+    "factor": 8.0,
+    "llama_4_scaling_beta": 0.1,
+    "mscale": 1.0,
+    "mscale_all_dim": 1.0,
+    "original_max_position_embeddings": 16384,
+    "rope_theta": 1000000.0,
+    "rope_type": "yarn"
+  },
+  "bos_token_id": 1,
+  "eos_token_id": 11,
+  "pad_token_id": 11,
+  "tie_word_embeddings": false,
+  "model_type": "nlduni",
+  "architectures": [
+    "NLDUniForCausalLM"
+  ],
+  "torch_dtype": "bfloat16",
+  "auto_map": {
+    "AutoConfig": "configuration_nlduni.NLDUniConfig",
+    "AutoModel": "modeling_nlduni.NLDUniForCausalLM",
+    "AutoModelForCausalLM": "modeling_nlduni.NLDUniForCausalLM"
+  }
+}

configuration_nlduni.py

@@ -0,0 +1,59 @@
+from transformers import PretrainedConfig
+
+class NLDUniConfig(PretrainedConfig):
+    model_type = 'nlduni'
+    # Architecture invariants; deliberately absent from serialized configuration.
+    d_model_gen = 4096
+    d_model = 4096
+    diffusion_lm = True
+    block_size = 32
+    mm_native_newline_token = '\n'
+    mm_hidden_size = 1280
+    projector_hidden_size = 4096
+    gen_image_fixed_prefix_tokens = 128
+    gen_image_row_markers = True
+    generation_skip_down_sample = False
+    image_aspect_ratio = 'radio_h_native'
+    semantic_tree_token_diffusion = True
+    single_stream_image_diffusion = True
+    shift_sbd_logits = True
+    downsample = True
+    mm_vqvae = 'emu3'
+    image_embedding_refinement_steps = 4
+    lock_source_padding_vq = True
+    generation_embedding_scale = 0.005357294692673304
+    generation_downsample_scale = 0.4498539166224965
+    def __init__(self, hidden_size=4096, intermediate_size=14336, num_hidden_layers=34,
+                 num_attention_heads=32, num_key_value_heads=8, head_dim=128,
+                 vocab_size=132101, max_position_embeddings=131072,
+                 rms_norm_eps=1e-5, rope_parameters=None, **kwargs):
+        super().__init__(**kwargs)
+        self.hidden_size=hidden_size
+        self.intermediate_size=intermediate_size
+        self.num_hidden_layers=num_hidden_layers
+        self.num_attention_heads=num_attention_heads
+        self.num_key_value_heads=num_key_value_heads
+        self.head_dim=head_dim
+        self.vocab_size=vocab_size
+        self.max_position_embeddings=max_position_embeddings
+        self.rms_norm_eps=rms_norm_eps
+        self.hidden_act='silu'
+        self.attention_bias=False
+        self.attention_dropout=0.0
+        self.mlp_bias=False
+        self.sliding_window=None
+        self.initializer_range=0.02
+        self.rope_parameters=rope_parameters or dict(rope_type='yarn',rope_theta=1000000.,factor=8.,original_max_position_embeddings=16384,beta_fast=32.,beta_slow=1.,mscale=1.,mscale_all_dim=1.,llama_4_scaling_beta=.1)
+        self.rope_scaling=self.rope_parameters
+        self.rope_theta=self.rope_parameters['rope_theta']
+
+    def to_dict(self):
+        values = super().to_dict()
+        values.pop('image_token_tree_path', None)
+        values.pop('_name_or_path', None)
+        return values
+
+MinistralDLMConfig = NLDUniConfig
+
+def get_ministral_shift_sbd_logits(config):
+    return True

constants.py

@@ -0,0 +1,17 @@
+
+from __future__ import annotations
+class NemotronMinistralDLMConstant:
+    reserve_id = 18
+    reserve_id_token = '<SPECIAL_18>'
+    reserve_id_enc = 19
+    reserve_id_token_enc = '<SPECIAL_19>'
+    mask_id = 100
+    mask_id_token = '<SPECIAL_100>'
+    fim_id = 20
+    fim_id_token = '<SPECIAL_20>'
+    eos_id = 11
+    eos_id_token = '<|im_end|>'
+    gen_im_start_token = '<SPECIAL_21>'
+    gen_im_end_token = '<SPECIAL_22>'
+    und_im_pad_token = '<SPECIAL_23>'
+    und_im_pad_token_id = 23

conversation.py

@@ -0,0 +1,20 @@
+import copy
+
+class Conversation:
+    roles = ('<|im_start|>user', '<|im_start|>assistant')
+    system = '<|im_start|>system\nYou are a helpful language and vision assistant. You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language.'
+    def __init__(self):
+        self.messages = []
+    def copy(self):
+        return copy.deepcopy(self)
+    def append_message(self, role, message):
+        self.messages.append((role, message))
+    def get_prompt(self):
+        text = self.system + '<|im_end|>\n' if self.system else ''
+        for role, message in self.messages:
+            text += role + '\n'
+            if message:
+                text += message + '<|im_end|>\n'
+        return text
+
+conv_templates = {'ministral': Conversation()}

core.py

@@ -0,0 +1,426 @@
+
+from __future__ import annotations
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Optional, Union
+import torch
+from torch import nn
+from transformers.utils.generic import check_model_inputs
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.integrations import use_kernel_forward_from_hub
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring
+from .configuration_nlduni import MinistralDLMConfig
+
+def rotate_half(x):
+    x1 = x[..., :x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = q * cos + rotate_half(q) * sin
+    k_embed = k * cos + rotate_half(k) * sin
+    return (q_embed, k_embed)
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    (batch, num_key_value_heads, slen, head_dim) = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+def eager_attention_forward(module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: Optional[torch.Tensor], scaling: float, dropout: float=0.0, **kwargs: Unpack[TransformersKwargs]):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, :key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=False)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+    return (attn_output, attn_weights)
+
+def _get_llama_4_attn_scale(positions_ids: torch.Tensor, beta: float, max_position_embeddings: int) -> torch.Tensor:
+    scaling = 1 + beta * torch.log(1 + torch.floor(positions_ids / max_position_embeddings))
+    return scaling.unsqueeze(-1)
+
+@dataclass(frozen=True, slots=True)
+class _MixedTokenRoutingPlan:
+    mask: torch.Tensor
+    understanding_indices: torch.Tensor
+    generation_indices: torch.Tensor
+
+    @classmethod
+    def from_mask(cls, mask: torch.Tensor) -> '_MixedTokenRoutingPlan':
+        flat_mask = mask.reshape(-1)
+        return cls(mask=mask, understanding_indices=(~flat_mask).nonzero(as_tuple=False).flatten(), generation_indices=flat_mask.nonzero(as_tuple=False).flatten())
+TokenRouting = Optional[Union[bool, _MixedTokenRoutingPlan]]
+
+def _resolve_dual_tower_layers(config: MinistralDLMConfig) -> int:
+    num_hidden_layers = config.num_hidden_layers
+    configured = 18
+    if configured is None:
+        return num_hidden_layers
+    if isinstance(configured, bool) or not isinstance(configured, int):
+        raise TypeError('dual_tower_layers must be an integer')
+    if not 0 <= configured <= num_hidden_layers:
+        raise ValueError(f'dual_tower_layers must be between 0 and num_hidden_layers ({num_hidden_layers}), got {configured}')
+    return configured
+
+def _routing_mask(routing: TokenRouting, hidden_states: torch.Tensor) -> torch.Tensor:
+    if routing is None:
+        return torch.zeros(hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device)
+    if routing is True:
+        return torch.ones(hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device)
+    return routing.mask
+
+def _create_split_tower_block_mask(block_mask, *, current_routing: torch.Tensor, cached_routing: torch.Tensor):
+    from torch.nn.attention.flex_attention import BlockMask, create_block_mask
+    if isinstance(block_mask, str):
+        if block_mask != 'full':
+            raise ValueError(f'Unsupported attention block-mask sentinel: {block_mask}')
+        base_mask_mod = None
+        query_length = current_routing.shape[1]
+        key_value_length = cached_routing.shape[1] + query_length
+        block_size = 128
+    elif isinstance(block_mask, BlockMask):
+        base_mask_mod = block_mask.mask_mod
+        (query_length, key_value_length) = block_mask.seq_lengths
+        block_size = block_mask.BLOCK_SIZE
+    else:
+        raise TypeError("partial dual-tower FlexAttention requires an explicit BlockMask or the 'full' sentinel")
+    current_length = current_routing.shape[1]
+    cached_length = cached_routing.shape[1]
+    if query_length < current_length:
+        raise ValueError(f'FlexAttention block-mask query length is shorter than routing: {query_length} < {current_length}')
+    if key_value_length < cached_length + current_length:
+        raise ValueError(f'FlexAttention block-mask key/value length is shorter than cached and current routing: {key_value_length} < {cached_length + current_length}')
+    query_routing = torch.nn.functional.pad(current_routing, (0, query_length - current_length), value=False)
+    key_value_routing = torch.cat((cached_routing, current_routing), dim=1)
+    key_value_routing = torch.nn.functional.pad(key_value_routing, (0, key_value_length - key_value_routing.shape[1]), value=False)
+
+    def split_tower_mask(b, h, q_idx, kv_idx):
+        same_tower = query_routing[b, q_idx] == key_value_routing[b, kv_idx]
+        same_tower = same_tower | (q_idx >= current_length)
+        if base_mask_mod is None:
+            return same_tower
+        return base_mask_mod(b, h, q_idx, kv_idx) & same_tower
+    return create_block_mask(split_tower_mask, B=current_routing.shape[0], H=None, Q_LEN=query_length, KV_LEN=key_value_length, device=current_routing.device, BLOCK_SIZE=block_size)
+
+def _canonicalize_token_routing(routing: Optional[torch.Tensor], hidden_states: torch.Tensor, *, dual_tower: bool) -> TokenRouting:
+    if routing is None:
+        return None
+    if not torch.is_tensor(routing) or routing.ndim != 2:
+        raise TypeError('routing must be a rank-2 tensor with shape [batch, sequence]')
+    if tuple(routing.shape) != tuple(hidden_states.shape[:2]):
+        raise ValueError(f'routing must match the transformer token shape: expected {tuple(hidden_states.shape[:2])}, got {tuple(routing.shape)}')
+    if routing.device != hidden_states.device:
+        raise ValueError('routing and transformer hidden states must share a device')
+    if routing.dtype == torch.bool:
+        normalized = routing
+    elif routing.dtype in {torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64}:
+        if bool(((routing != 0) & (routing != 1)).any()):
+            raise ValueError('routing values must be binary (0 for U, 1 for G)')
+        normalized = routing.to(dtype=torch.bool)
+    else:
+        raise TypeError('routing must have boolean or integer dtype')
+    has_generation_route = bool(normalized.any())
+    if has_generation_route and (not dual_tower):
+        raise ValueError('routing selects G but config.dual_tower is false')
+    if not has_generation_route:
+        return None
+    if bool(normalized.all()):
+        return True
+    return _MixedTokenRoutingPlan.from_mask(normalized)
+
+def _routed_module(hidden_states: torch.Tensor, routing: TokenRouting, understanding_module: Callable[[torch.Tensor], torch.Tensor], generation_module: Optional[Callable[[torch.Tensor], torch.Tensor]]) -> torch.Tensor:
+    if routing is None:
+        return understanding_module(hidden_states)
+    if routing is True:
+        if generation_module is None:
+            raise RuntimeError('G routing requires generation-tower parameters')
+        return generation_module(hidden_states)
+    if generation_module is None:
+        raise RuntimeError('mixed routing requires generation-tower parameters')
+    if not isinstance(routing, _MixedTokenRoutingPlan):
+        raise TypeError('mixed routing must use a canonical dispatch plan')
+    flat_hidden = hidden_states.reshape(-1, hidden_states.shape[-1])
+    understanding_output = understanding_module(flat_hidden.index_select(0, routing.understanding_indices))
+    generation_output = generation_module(flat_hidden.index_select(0, routing.generation_indices))
+    flat_output = understanding_output.new_empty((flat_hidden.shape[0], *understanding_output.shape[1:]))
+    flat_output = flat_output.index_copy(0, routing.understanding_indices, understanding_output)
+    flat_output = flat_output.index_copy(0, routing.generation_indices, generation_output)
+    return flat_output.reshape(*hidden_states.shape[:-1], *flat_output.shape[1:])
+
+class Ministral3Attention(nn.Module):
+
+    def __init__(self, config: MinistralDLMConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.split_tower = layer_idx >= _resolve_dual_tower_layers(config)
+        self.head_dim = getattr(config, 'head_dim', None) or config.hidden_size // config.num_attention_heads
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim ** -0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
+        self.q_proj_gen = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
+        self.k_proj_gen = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj_gen = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj_gen = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
+        self.diffusion_lm = config.diffusion_lm
+
+    def _project_qkv(self, hidden_states: torch.Tensor, routing: TokenRouting) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return (_routed_module(hidden_states, routing, self.q_proj, getattr(self, 'q_proj_gen', None)), _routed_module(hidden_states, routing, self.k_proj, getattr(self, 'k_proj_gen', None)), _routed_module(hidden_states, routing, self.v_proj, getattr(self, 'v_proj_gen', None)))
+
+    def _project_output(self, hidden_states: torch.Tensor, routing: TokenRouting) -> torch.Tensor:
+        return _routed_module(hidden_states, routing, self.o_proj, getattr(self, 'o_proj_gen', None))
+
+    def forward(self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_values: Optional[Cache]=None, cache_position: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=False, routing: TokenRouting=None, **kwargs: Unpack[FlashAttentionKwargs]) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+        (query_states, key_states, value_states) = self._project_qkv(hidden_states, routing)
+        query_states = query_states.view(hidden_shape).transpose(1, 2)
+        key_states = key_states.view(hidden_shape).transpose(1, 2)
+        value_states = value_states.view(hidden_shape).transpose(1, 2)
+        (cos, sin) = position_embeddings
+        (query_states, key_states) = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        query_states = query_states * _get_llama_4_attn_scale(cache_position, self.config.rope_parameters.get('llama_4_scaling_beta'), self.config.rope_parameters.get('original_max_position_embeddings')).to(query_states.dtype)
+        if past_key_values is not None:
+            if use_cache:
+                cache_kwargs = {'sin': sin, 'cos': cos, 'cache_position': cache_position}
+                (key_states, value_states) = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+            else:
+                (old_k, old_v) = (past_key_values.layers[self.layer_idx].keys, past_key_values.layers[self.layer_idx].values)
+                key_states = torch.cat([old_k, key_states], dim=-2)
+                value_states = torch.cat([old_v, value_states], dim=-2)
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != 'eager':
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+        if self.diffusion_lm:
+            (attn_output, attn_weights) = attention_interface(self, query_states, key_states, value_states, None, dropout=0.0, scaling=self.scaling, is_causal=False, **kwargs)
+        else:
+            (attn_output, attn_weights) = attention_interface(self, query_states, key_states, value_states, attention_mask, dropout=0.0, scaling=self.scaling, sliding_window=getattr(self.config, 'sliding_window', None), **kwargs)
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self._project_output(attn_output, routing)
+        return (attn_output, attn_weights)
+
+class Ministral3MLP(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.gate_proj_gen = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj_gen = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj_gen = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def _understanding_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
+
+    def _generation_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.down_proj_gen(self.act_fn(self.gate_proj_gen(hidden_states)) * self.up_proj_gen(hidden_states))
+
+    def forward(self, hidden_states: torch.Tensor, routing: TokenRouting=None) -> torch.Tensor:
+        generation_module = self._generation_forward if hasattr(self, 'gate_proj_gen') else None
+        return _routed_module(hidden_states, routing, self._understanding_forward, generation_module)
+
+@use_kernel_forward_from_hub('RMSNorm')
+class Ministral3RMSNorm(nn.Module):
+
+    def __init__(self, hidden_size, eps=1e-06):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f'{tuple(self.weight.shape)}, eps={self.variance_epsilon}'
+
+class Ministral3DecoderLayer(GradientCheckpointingLayer):
+
+    def __init__(self, config: MinistralDLMConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        from .attention import MinistralFlexAttention
+        attn_class = MinistralFlexAttention
+        self.self_attn = attn_class(config=config, layer_idx=layer_idx)
+        self.mlp = Ministral3MLP(config)
+        self.input_layernorm = Ministral3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Ministral3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.input_layernorm_gen = Ministral3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm_gen = Ministral3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Cache]=None, use_cache: Optional[bool]=False, cache_position: Optional[torch.LongTensor]=None, position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]]=None, routing: TokenRouting=None, **kwargs: Unpack[TransformersKwargs]) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = _routed_module(hidden_states, routing, self.input_layernorm, getattr(self, 'input_layernorm_gen', None))
+        (hidden_states, _) = self.self_attn(hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, routing=routing, **kwargs)
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = _routed_module(hidden_states, routing, self.post_attention_layernorm, getattr(self, 'post_attention_layernorm_gen', None))
+        hidden_states = self.mlp(hidden_states, routing=routing)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+class Ministral3PreTrainedModel(PreTrainedModel):
+    config: MinistralDLMConfig
+    base_model_prefix = 'model'
+    supports_gradient_checkpointing = True
+    _no_split_modules = ['Ministral3DecoderLayer']
+    _skip_keys_device_placement = ['past_key_values']
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {'hidden_states': Ministral3DecoderLayer, 'attentions': Ministral3Attention}
+
+class Ministral3RotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor
+
+    def __init__(self, config: MinistralDLMConfig, device=None):
+        super().__init__()
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+        self.config = config
+        self.rope_type = self.config.rope_parameters['rope_type']
+        rope_init_fn: Callable = self.compute_default_rope_parameters
+        if self.rope_type != 'default':
+            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        (inv_freq, self.attention_scaling) = rope_init_fn(self.config, device)
+        self.register_buffer('inv_freq', inv_freq, persistent=False)
+        self.original_inv_freq = inv_freq
+
+    @staticmethod
+    def compute_default_rope_parameters(config: Optional[MinistralDLMConfig]=None, device: Optional['torch.device']=None, seq_len: Optional[int]=None) -> tuple['torch.Tensor', float]:
+        base = config.rope_parameters['rope_theta']
+        dim = getattr(config, 'head_dim', None) or config.hidden_size // config.num_attention_heads
+        attention_factor = 1.0
+        inv_freq = 1.0 / base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        return (inv_freq, attention_factor)
+
+    @torch.no_grad()
+    @dynamic_rope_update
+    def forward(self, x, position_ids):
+        with torch.autocast(device_type=x.device.type, enabled=False):
+            inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+            position_ids_expanded = position_ids[:, None, :].float()
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+        return (cos.to(dtype=x.dtype), sin.to(dtype=x.dtype))
+
+class Ministral3Model(Ministral3PreTrainedModel):
+
+    def __init__(self, config: MinistralDLMConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.dual_tower_layers = _resolve_dual_tower_layers(config)
+        self._layer_loop_config = None
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList([Ministral3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
+        self.norm = Ministral3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.norm_gen = Ministral3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Ministral3RotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    @check_model_inputs
+    def forward(self, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, position_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Cache]=None, inputs_embeds: Optional[torch.FloatTensor]=None, use_cache: Optional[bool]=None, cache_position: Optional[torch.LongTensor]=None, input_embeddings: Optional[torch.FloatTensor]=None, routing: Optional[torch.Tensor]=None, modality_indices: Optional[torch.Tensor]=None, loop_k: Optional[int]=None, loop_depths: Optional[torch.LongTensor]=None, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutputWithPast:
+        if input_embeddings is not None:
+            inputs_embeds = input_embeddings
+            input_embeddings = None
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError('You must specify exactly one of input_ids or inputs_embeds')
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        single_stream_call = kwargs.get('single_stream_image_diffusion', False)
+        if single_stream_call is None:
+            single_stream_call = False
+        if not isinstance(single_stream_call, bool):
+            raise TypeError('single_stream_image_diffusion must be boolean')
+        use_layer_loop = self._layer_loop_config is not None and single_stream_call
+        if loop_k is not None or loop_depths is not None:
+            raise ValueError('loop_k and loop_depths require an active single-stream looped_layer_execution call')
+        resolved_loop_depths = None
+        loop_visits = 1
+        dual_tower = True
+        partial_dual_tower = dual_tower and self.dual_tower_layers < self.config.num_hidden_layers
+        if dual_tower and modality_indices is not None:
+            modality_routing = _canonicalize_token_routing(modality_indices, inputs_embeds, dual_tower=dual_tower)
+            if routing is None:
+                canonical_routing = modality_routing
+            else:
+                canonical_routing = _canonicalize_token_routing(routing, inputs_embeds, dual_tower=dual_tower)
+                if not torch.equal(routing.to(dtype=torch.bool), modality_indices.to(dtype=torch.bool)):
+                    raise ValueError('routing and deprecated modality_indices alias must select identical U/G tokens when both are provided')
+        else:
+            canonical_routing = _canonicalize_token_routing(routing, inputs_embeds, dual_tower=dual_tower)
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None and (partial_dual_tower or cache_position is None) else 0
+        stored_cached_routing = getattr(past_key_values, '_nlduni_dual_tower_routing', None)
+        track_dual_tower_routing = partial_dual_tower and (canonical_routing is not None or stored_cached_routing is not None)
+        current_routing = _routing_mask(canonical_routing, inputs_embeds) if track_dual_tower_routing else None
+        cached_routing = stored_cached_routing
+        if track_dual_tower_routing and cached_routing is None:
+            cached_routing = torch.zeros(current_routing.shape[0], past_seen_tokens, dtype=torch.bool, device=current_routing.device)
+        elif track_dual_tower_routing and tuple(cached_routing.shape) != (current_routing.shape[0], past_seen_tokens):
+            raise ValueError(f'cached dual-tower routing must match the KV cache shape: expected {(current_routing.shape[0], past_seen_tokens)}, got {tuple(cached_routing.shape)}')
+        elif track_dual_tower_routing and cached_routing.device != current_routing.device:
+            raise ValueError('cached dual-tower routing and inputs must share a device')
+        if track_dual_tower_routing and past_key_values is not None:
+            past_key_values._nlduni_dual_tower_routing = cached_routing
+            past_key_values._nlduni_dual_tower_has_generation = bool(getattr(past_key_values, '_nlduni_dual_tower_has_generation', False))
+        if cache_position is None:
+            cache_position = torch.arange(past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device)
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+        if kwargs.get('use_causal_mask', False):
+            mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask
+            causal_mask = mask_function(config=self.config, input_embeds=inputs_embeds, attention_mask=attention_mask, cache_position=cache_position, past_key_values=past_key_values, position_ids=position_ids)
+        else:
+            causal_mask = None
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+        split_layer_kwargs = None
+        has_split_layers = track_dual_tower_routing
+        if has_split_layers and kwargs.get('overwrite_attn_impl') != 'flash_attn' and (kwargs.get('overwrite_block_mask') is not None):
+            split_layer_kwargs = dict(kwargs)
+            split_layer_kwargs['overwrite_block_mask'] = _create_split_tower_block_mask(kwargs['overwrite_block_mask'], current_routing=current_routing, cached_routing=cached_routing)
+        for (layer_idx, decoder_layer) in enumerate(self.layers[:self.config.num_hidden_layers]):
+            layer_kwargs = split_layer_kwargs if split_layer_kwargs is not None and layer_idx >= self.dual_tower_layers else kwargs
+            hidden_states = decoder_layer(hidden_states, attention_mask=causal_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, routing=canonical_routing, **layer_kwargs)
+        if track_dual_tower_routing and use_cache and (past_key_values is not None):
+            past_key_values._nlduni_dual_tower_routing = torch.cat((cached_routing, current_routing), dim=1)
+            past_key_values._nlduni_dual_tower_has_generation = past_key_values._nlduni_dual_tower_has_generation or canonical_routing is not None
+        hidden_states = _routed_module(hidden_states, canonical_routing, self.norm, getattr(self, 'norm_gen', None))
+        return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=past_key_values if use_cache else None)

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "bos_token_id": 1,
+  "eos_token_id": 11,
+  "pad_token_id": 11,
+  "max_new_tokens": 16000,
+  "do_sample": false
+}

image_utils.py

@@ -0,0 +1,65 @@
+# ADOBE CONFIDENTIAL
+# Copyright 2025 Adobe
+# All Rights Reserved.
+# NOTICE: All information contained herein is, and remains
+# the property of Adobe and its suppliers, if any. The intellectual
+# and technical concepts contained herein are proprietary to Adobe
+# and its suppliers and are protected by all applicable intellectual
+# property laws, including trade secret and copyright laws.
+# Dissemination of this information or reproduction of this material
+# is strictly forbidden unless prior written permission is obtained
+# from Adobe.
+
+from __future__ import annotations
+import torch
+from .token_constants import DEFAULT_IMAGE_PAD_TOKEN, DEFAULT_IMAGE_TOKEN
+import numpy as np
+from PIL import ImageOps
+
+def pad_to_square_and_resize(img, size=1024):
+    padded_img = ImageOps.pad(img, (max(img.size), max(img.size)), color=(0, 0, 0))
+    resized_img = padded_img.resize((size, size))
+    return resized_img
+
+def process_images(images, image_processor, model_cfg):
+    result = []
+    for image in images:
+        pixels = image_processor(images=[image], return_tensors='pt').pixel_values
+        (h, w) = pixels.shape[-2:]
+        result.append((pixels, torch.tensor([[1, h // 16, w // 16]], dtype=torch.long)))
+    return result
+
+def tokenizer_image_token(prompt, tokenizer, image_token_index=-200, return_tensors=None, constant_version=None):
+    prompt = prompt.replace('<image_padding>', '<SPECIAL_23>')
+    ids = torch.tensor(tokenizer(prompt).input_ids, dtype=torch.long)
+    ids[ids == 23] = -500
+    return ids if return_tensors == 'pt' else ids.tolist()
+
+def compute_num_image_tokens(x, image_aspect_ratio, image_processor):
+    grid = x[0][1]
+    return (int(grid[0, 1]) // 2, int(grid[0, 2]) // 2)
+
+def fix_one_prompt(prompt, config, n_image_tokens):
+    image_idx = 0
+    while DEFAULT_IMAGE_TOKEN in prompt:
+        newline_token = config.mm_native_newline_token
+        if newline_token:
+            if len(n_image_tokens[image_idx]) == 2:
+                (hh, ww) = n_image_tokens[image_idx]
+                row = DEFAULT_IMAGE_PAD_TOKEN * ww + newline_token
+                rows = row * hh
+                rows = rows.rstrip(newline_token)
+                prompt = prompt.replace('<image>', rows, 1)
+            elif len(n_image_tokens[image_idx]) == 3:
+                (tt, hh, ww) = (int(value) for value in n_image_tokens[image_idx])
+                row = DEFAULT_IMAGE_PAD_TOKEN * ww + newline_token
+                per_image = row * hh + newline_token + newline_token
+                images = (per_image * tt).rstrip(newline_token + newline_token)
+                prompt = prompt.replace('<image>', images, 1)
+            else:
+                raise NotImplementedError()
+        else:
+            image_tokens = DEFAULT_IMAGE_PAD_TOKEN * np.prod(n_image_tokens[image_idx])
+            prompt = prompt.replace('<image>', image_tokens, 1)
+        image_idx += 1
+    return prompt

inference.py

@@ -0,0 +1,22 @@
+"""Examples: python inference.py /path/to/model --task t2i --prompt 'A red mug'."""
+import argparse
+from pathlib import Path
+import torch
+from transformers import AutoModelForCausalLM
+
+parser=argparse.ArgumentParser()
+parser.add_argument('model')
+parser.add_argument('--task',choices=['t2i','edit','understand'],default='t2i')
+parser.add_argument('--prompt',required=True)
+parser.add_argument('--image')
+parser.add_argument('--output',default='output.png')
+parser.add_argument('--seed',type=int,default=42)
+args=parser.parse_args()
+model=AutoModelForCausalLM.from_pretrained(args.model,trust_remote_code=True,torch_dtype=torch.bfloat16).to('cuda').eval()
+if args.task=='t2i':
+    model.text_to_image(args.prompt,seed=args.seed).save(args.output)
+elif args.task=='edit':
+    if not args.image:parser.error('--image is required for editing')
+    model.image_edit(args.image,args.prompt,seed=args.seed).save(args.output)
+else:
+    print(model.generate(args.prompt,image=args.image))

markers.py

@@ -0,0 +1,182 @@
+
+from __future__ import annotations
+import math
+from numbers import Integral
+import torch
+from torch import nn
+GEN_IMAGE_ROW_MARKERS = 'gen_image_row_markers'
+GEN_IMAGE_FIXED_PREFIX_TOKENS = 'gen_image_fixed_prefix_tokens'
+GENERATION_SKIP_DOWNSAMPLE = 'generation_skip_down_sample'
+ROW_MARKER_OFFSET = 0
+IMAGE_MASK_OFFSET = 1
+IMAGE_END_OFFSET = 2
+FIXED_PREFIX_OFFSET = 3
+RESERVED_GENERATION_ROWS = 256
+MARKERS_PER_RAW_ROW = 2
+UNCOMPRESSED_MARKERS_PER_RAW_ROW = 1
+_SUPPORTED_MARKERS_PER_ROW = (UNCOMPRESSED_MARKERS_PER_RAW_ROW, MARKERS_PER_RAW_ROW)
+
+def _validate_markers_per_row(value) -> int:
+    if isinstance(value, bool) or not isinstance(value, Integral):
+        raise TypeError('markers_per_row must be an integer')
+    value = int(value)
+    if value not in _SUPPORTED_MARKERS_PER_ROW:
+        raise ValueError('markers_per_row must be 1 for an uncompressed image grid or 2 for a grid that will undergo 2x spatial compression')
+    return value
+
+def _validate_fixed_prefix_count(value) -> int:
+    if value is None:
+        return 0
+    if isinstance(value, bool) or not isinstance(value, Integral):
+        raise TypeError(f'config.{GEN_IMAGE_FIXED_PREFIX_TOKENS} must be an integer')
+    value = int(value)
+    max_tokens = RESERVED_GENERATION_ROWS - FIXED_PREFIX_OFFSET
+    if not 0 <= value <= max_tokens:
+        raise ValueError(f'config.{GEN_IMAGE_FIXED_PREFIX_TOKENS} must be between 0 and {max_tokens}')
+    return value
+
+def image_row_markers_enabled(config) -> bool:
+    value = getattr(config, GEN_IMAGE_ROW_MARKERS, False)
+    if value is None:
+        return False
+    if not isinstance(value, bool):
+        raise TypeError(f'config.{GEN_IMAGE_ROW_MARKERS} must be boolean')
+    return value
+
+def resolve_generation_skip_down_sample(config, *, runtime_value=None) -> bool:
+    configured = getattr(config, GENERATION_SKIP_DOWNSAMPLE, None)
+    if configured is not None and (not isinstance(configured, bool)):
+        raise TypeError(f'config.{GENERATION_SKIP_DOWNSAMPLE} must be boolean or null')
+    runtime_configured = runtime_value is not None
+    runtime_enabled = bool(runtime_value)
+    if configured is None:
+        return runtime_enabled
+    if runtime_configured and runtime_enabled != configured:
+        raise ValueError(f'SKIP_DOWN_SAMPLE conflicts with checkpoint-backed config.{GENERATION_SKIP_DOWNSAMPLE}: runtime={runtime_enabled}, checkpoint={configured}')
+    return configured
+
+def image_generation_special_ids(codebook_size: int) -> tuple[int, int, int]:
+    if isinstance(codebook_size, bool) or not isinstance(codebook_size, Integral):
+        raise TypeError('codebook_size must be a positive integer')
+    codebook_size = int(codebook_size)
+    if codebook_size < 1:
+        raise ValueError('codebook_size must be a positive integer')
+    return (codebook_size + ROW_MARKER_OFFSET, codebook_size + IMAGE_MASK_OFFSET, codebook_size + IMAGE_END_OFFSET)
+
+def image_fixed_prefix_token_count(config) -> int:
+    return _validate_fixed_prefix_count(getattr(config, GEN_IMAGE_FIXED_PREFIX_TOKENS, 0))
+
+def image_generation_fixed_prefix_ids(codebook_size: int, count: int) -> tuple[int, ...]:
+    image_generation_special_ids(codebook_size)
+    count = _validate_fixed_prefix_count(count)
+    start = int(codebook_size) + FIXED_PREFIX_OFFSET
+    return tuple(range(start, start + count))
+
+def split_generation_fixed_prefix(generation_placeholder_mask: torch.Tensor, *, codebook_size: int, count: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if not torch.is_tensor(generation_placeholder_mask):
+        raise TypeError('generation_placeholder_mask must be a tensor')
+    if generation_placeholder_mask.ndim != 2:
+        raise ValueError('generation_placeholder_mask must have shape [batch, sequence]')
+    placeholder_mask = generation_placeholder_mask.to(dtype=torch.bool)
+    fixed_ids = image_generation_fixed_prefix_ids(codebook_size, count)
+    if not fixed_ids:
+        return (torch.zeros_like(placeholder_mask), placeholder_mask, torch.empty(0, dtype=torch.long, device=placeholder_mask.device))
+    previous_placeholder = torch.cat((torch.zeros_like(placeholder_mask[:, :1]), placeholder_mask[:, :-1]), dim=1)
+    run_starts = placeholder_mask & ~previous_placeholder
+    positions = torch.arange(placeholder_mask.shape[1], device=placeholder_mask.device, dtype=torch.long).unsqueeze(0).expand_as(placeholder_mask)
+    most_recent_start = torch.where(run_starts, positions, positions.new_full((), -1)).cummax(dim=1).values
+    run_offsets = positions - most_recent_start
+    fixed_mask = placeholder_mask & run_offsets.lt(len(fixed_ids))
+    image_mask = placeholder_mask & ~fixed_mask
+    next_placeholder = torch.cat((placeholder_mask[:, 1:], torch.zeros_like(placeholder_mask[:, :1])), dim=1)
+    run_ends = placeholder_mask & ~next_placeholder
+    run_lengths = run_offsets[run_ends] + 1
+    if bool(run_lengths.le(len(fixed_ids)).any()):
+        shortest = int(run_lengths.min().item())
+        raise ValueError(f'each generated-image placeholder run must contain the fixed prefix plus at least one image slot; shortest run={shortest}, fixed prefix={len(fixed_ids)}')
+    dense_ids = int(codebook_size) + FIXED_PREFIX_OFFSET + run_offsets
+    return (fixed_mask, image_mask, dense_ids[fixed_mask])
+
+def generation_embedding_capacity(generation_model) -> int:
+    embedding = getattr(generation_model, 'gen_embedding', None)
+    if isinstance(embedding, nn.Embedding):
+        return int(embedding.num_embeddings)
+    if isinstance(embedding, nn.Module):
+        for module in embedding.modules():
+            if isinstance(module, nn.Embedding):
+                return int(module.num_embeddings)
+    raise TypeError('generation model does not expose an nn.Embedding lookup')
+
+def image_row_marker_count(generation_model, *, skip_down_sample=False) -> int:
+    compression_available = getattr(generation_model, 'packed_image_codec', None) is not None or getattr(generation_model, 'downsample_gen', None) is not None
+    if compression_available and (not skip_down_sample):
+        return MARKERS_PER_RAW_ROW
+    return UNCOMPRESSED_MARKERS_PER_RAW_ROW
+
+def validate_image_row_marker_support(generation_model, *, skip_down_sample=False):
+    config = generation_model.config
+    if 'emu3' not in 'emu3'.lower():
+        raise NotImplementedError("generation image row markers currently require Emu3's single-codebook tokenizer")
+    image_row_marker_count(generation_model, skip_down_sample=skip_down_sample)
+    codebook_size = int(generation_model.vqvae.config.num_vq_embeddings)
+    (row_id, mask_id, end_id) = image_generation_special_ids(codebook_size)
+    capacity = generation_embedding_capacity(generation_model)
+    if end_id >= capacity:
+        raise ValueError(f'generation embedding has no spare rows for row/end markers: required ID {end_id}, capacity {capacity}')
+    return (row_id, mask_id, end_id)
+
+def validate_image_fixed_prefix_support(generation_model) -> tuple[int, ...]:
+    config = generation_model.config
+    count = image_fixed_prefix_token_count(config)
+    if count == 0:
+        return ()
+    if getattr(config, 'generation_embedding_scale', None) is None:
+        raise ValueError('generation fixed-prefix tokens require factorized generation scaling so uncompressed tokens receive the common input scale')
+    if 'emu3' not in 'emu3'.lower():
+        raise NotImplementedError("generation fixed-prefix tokens currently require Emu3's single-codebook tokenizer")
+    codebook_size = int(generation_model.vqvae.config.num_vq_embeddings)
+    fixed_ids = image_generation_fixed_prefix_ids(codebook_size, count)
+    capacity = generation_embedding_capacity(generation_model)
+    if fixed_ids[-1] >= capacity:
+        raise ValueError(f'generation embedding has no spare rows for the requested fixed prefix: required ID {fixed_ids[-1]}, capacity {capacity}')
+    special_ids = set(image_generation_special_ids(codebook_size))
+    if special_ids.intersection(fixed_ids):
+        raise RuntimeError('generation fixed-prefix IDs overlap image mask/row/end markers')
+    return fixed_ids
+
+def augmented_image_shape(gen_shape, *, markers_per_row: int=MARKERS_PER_RAW_ROW) -> tuple[int, int]:
+    if gen_shape is None or len(gen_shape) != 2:
+        raise ValueError('gen_shape must contain image height and width')
+    markers_per_row = _validate_markers_per_row(markers_per_row)
+    (height, width) = (int(value) for value in gen_shape)
+    if height < 1 or width < 1:
+        raise ValueError('image row markers require a positive image grid')
+    if markers_per_row == MARKERS_PER_RAW_ROW and (height % 2 or width % 2):
+        raise ValueError('two-column image row markers require an even-height, even-width grid for 2x spatial compression')
+    return (height, width + markers_per_row)
+
+def append_image_row_markers(token_ids: torch.Tensor, gen_shape, *, row_marker_id: int, image_end_id: int, markers_per_row: int=MARKERS_PER_RAW_ROW) -> tuple[torch.Tensor, tuple[int, int], torch.Tensor]:
+    if not torch.is_tensor(token_ids) or token_ids.ndim != 2:
+        raise ValueError('token_ids must have shape [batch, raw_image_tokens]')
+    markers_per_row = _validate_markers_per_row(markers_per_row)
+    (height, width) = (int(value) for value in gen_shape)
+    augmented_shape = augmented_image_shape((height, width), markers_per_row=markers_per_row)
+    if token_ids.shape[-1] != height * width:
+        raise ValueError(f'raw image token count does not match gen_shape: {token_ids.shape[-1]} != {height * width}')
+    grid = token_ids.reshape(token_ids.shape[0], height, width)
+    markers = torch.full((token_ids.shape[0], height, markers_per_row), int(row_marker_id), dtype=token_ids.dtype, device=token_ids.device)
+    markers[:, -1, :] = int(image_end_id)
+    augmented = torch.cat((grid, markers), dim=-1).flatten(1)
+    structural_mask = torch.zeros_like(augmented, dtype=torch.bool)
+    structural_mask.reshape(token_ids.shape[0], *augmented_shape)[:, :, -markers_per_row:] = True
+    return (augmented, augmented_shape, structural_mask)
+
+def strip_image_row_markers(token_ids: torch.Tensor, original_gen_shape, *, markers_per_row: int=MARKERS_PER_RAW_ROW) -> torch.Tensor:
+    if not torch.is_tensor(token_ids) or token_ids.ndim != 2:
+        raise ValueError('token_ids must have shape [batch, augmented_tokens]')
+    markers_per_row = _validate_markers_per_row(markers_per_row)
+    (height, width) = (int(value) for value in original_gen_shape)
+    augmented_shape = augmented_image_shape((height, width), markers_per_row=markers_per_row)
+    if token_ids.shape[-1] != math.prod(augmented_shape):
+        raise ValueError(f'augmented image token count does not match marker grid: {token_ids.shape[-1]} != {math.prod(augmented_shape)}')
+    return token_ids.reshape(token_ids.shape[0], *augmented_shape)[:, :, :width].flatten(1)

model-00001-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c4d490e010475dde3e2b57de64261a56d6bf8e9ca7d077754913fd628f7de5f6
+size 4974634832

model-00002-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:933c4b2f44f9cd7be30ac7a40efbb3c58485aa9fbd6396245d8e2337faf5c5b3
+size 4999819672

model-00003-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6893a316ec6655eb3f518be0f3ae3dfc80621922741df7d14f12b3d9fd123bf6
+size 4949454016

model-00004-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:858fab295f100b897634b72824d01e01f368dc1d500de03916854d7b2769986c
+size 4999819744

model-00005-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d31608093622df48c2284fa310170f9a1bfd127d7108bdacbc8106fcf296691
+size 4999819776

model-00006-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4356b837d798b2d3250cd067c15b7f5d56fd19a8b8605bb107fb9d727c1726d0
+size 4999819272

model-00007-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:003f5a6c4c727e5f560e16ad91ee2264ed138c7190108dfc9b68d1b8c8445545
+size 4456023446

model-00008-of-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:299783fae2f540dfd723abbb09f8661b57122210d44b23a4c65f24b254542550
+size 2155913464

modeling_nlduni.py

@@ -0,0 +1,252 @@
+from pathlib import Path
+import math
+import torch
+from torch import nn
+from PIL import Image
+from transformers import AutoTokenizer
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.utils.hub import cached_file
+from .configuration_nlduni import NLDUniConfig
+from .core import Ministral3Model, Ministral3PreTrainedModel
+from .attention import MinistralFlexAttention
+from .bridge import PackedLinearImageCodec
+from .vq import Emu3VQ, Emu3ImageProcessor
+from .vision import RadioVisionTower
+from .projector import RadioHProjector
+from .positions import _build_2d_sincos_position_embedding, GENERATION_2D_SINCOS_SCALE
+from .markers import strip_image_row_markers
+from .chat_utils import generate_with_prefix_cache_block_diff
+from .sampling import text_to_image as sample_image
+from .image_utils import process_images, fix_one_prompt, tokenizer_image_token
+from .conversation import Conversation
+from .tree import BalancedSemanticTree, production_semantic_tree_contract
+
+class NLDUniBackbone(Ministral3Model):
+    def __init__(self, config):
+        super().__init__(config)
+        self.vision_tower=RadioVisionTower()
+        self.mm_projector=RadioHProjector(config)
+        self.image_newline=nn.Parameter(torch.empty(config.hidden_size))
+        self.vqvae=Emu3VQ()
+        self.image_processor_gen=Emu3ImageProcessor()
+        self.gen_embedding=nn.Embedding(131072+256,4096)
+        self.gen_predictor=nn.Linear(4096,131072,bias=False)
+        self.packed_image_codec=PackedLinearImageCodec(4096,norm_eps=1e-6)
+        self.downsample_gen=None
+        self.upsample_gen=None
+        self.downsample_gen_enc=None
+        self.gen_embedding_2=None
+        self.gen_predictor_2=None
+        self.extra_gen_dit=None
+    def get_vision_tower(self):
+        return self.vision_tower
+    def get_gen_tokenizer(self):
+        return self.vqvae
+    def call_gen_token_embedding(self, token_ids):
+        return self.gen_embedding(token_ids)
+    def call_gen_fixed_token_embedding(self, token_ids):
+        return self.gen_embedding(token_ids)*self.config.generation_embedding_scale
+    def call_gen_embedding(self, token_ids=None, gen_shape=None, enc=False, child_embeddings=None,
+                           spatial_compression=True, add_spatial_position=True):
+        values=self.gen_embedding(token_ids) if child_embeddings is None else child_embeddings
+        if spatial_compression:
+            values=self.packed_image_codec.compress_fixed_grid(values,gen_shape)
+        scale=self.config.generation_embedding_scale
+        if spatial_compression:scale*=self.config.generation_downsample_scale
+        values=values*scale
+        if add_spatial_position:
+            h,w=gen_shape
+            if spatial_compression:h,w=h//2,w//2
+            pos=_build_2d_sincos_position_embedding(h,w,values.shape[-1],marker_column=not enc,device=values.device,dtype=values.dtype)
+            values=values+pos.unsqueeze(0)*GENERATION_2D_SINCOS_SCALE
+        return values
+    def call_gen_predictor(self, hidden_states, gen_shape=None, timesteps=None):
+        h,w=gen_shape
+        states=self.packed_image_codec.expand_fixed_grid(hidden_states.reshape(-1,(h//2)*(w//2),4096),gen_shape)
+        return self.gen_predictor(states.flatten(0,1))
+
+class NLDUniForCausalLM(Ministral3PreTrainedModel, GenerationMixin):
+    config_class=NLDUniConfig
+    base_model_prefix=''
+    def __init__(self, config):
+        super().__init__(config)
+        self.encoder=NLDUniBackbone(config)
+        self.diffusion_head=nn.Linear(config.hidden_size,config.vocab_size,bias=False)
+        self.mask_token_id=100
+        self.vocab_size=config.vocab_size
+        self._tokenizer=None
+        self._asset_root=None
+        self.post_init()
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        model=super().from_pretrained(pretrained_model_name_or_path,*args,**kwargs)
+        target=model[0] if isinstance(model,tuple) else model
+        options={k:kwargs[k] for k in ('revision','cache_dir','local_files_only','token','subfolder') if k in kwargs}
+        target._tokenizer=AutoTokenizer.from_pretrained(pretrained_model_name_or_path,**options)
+        tree=cached_file(pretrained_model_name_or_path,'pca_tree.pt',**options)
+        BalancedSemanticTree.from_artifact(production_semantic_tree_contract(tree))
+        target.config.image_token_tree_path=tree
+        target.eval()
+        return model
+    def save_pretrained(self, save_directory, **kwargs):
+        import shutil
+        destination = Path(save_directory)
+        destination.mkdir(parents=True, exist_ok=True)
+        super().save_pretrained(destination, **kwargs)
+        self.tokenizer.save_pretrained(destination)
+        asset = Path(self.config.image_token_tree_path)
+        if asset.resolve() != (destination / 'pca_tree.pt').resolve():
+            shutil.copyfile(asset, destination / 'pca_tree.pt')
+    @property
+    def model(self):return self.encoder
+    def get_model(self):return self.encoder
+    def get_vision_tower(self):return self.encoder.vision_tower
+    def get_input_embeddings(self):return self.encoder.embed_tokens
+    def set_input_embeddings(self, value):self.encoder.embed_tokens=value
+    def get_output_embeddings(self):return self.diffusion_head
+    def set_output_embeddings(self,value):self.diffusion_head=value
+    @property
+    def tokenizer(self):
+        if self._tokenizer is None:raise RuntimeError('Load with from_pretrained before inference')
+        return self._tokenizer
+    def forward(self,input_ids=None,inputs_embeds=None,attention_mask=None,position_ids=None,past_key_values=None,use_cache=False,**kwargs):
+        result=self.encoder(input_ids=input_ids,inputs_embeds=inputs_embeds,attention_mask=attention_mask,position_ids=position_ids,past_key_values=past_key_values,use_cache=use_cache,**kwargs)
+        return CausalLMOutputWithPast(logits=self.diffusion_head(result.last_hidden_state),past_key_values=result.past_key_values)
+    def prepare_inputs_labels_for_multimodal(self, input_ids, position_ids, attention_mask,
+             past_key_values, labels, images, modalities=None, image_sizes=None, return_inputs=False, **kwargs):
+        embeds=self.encoder.embed_tokens(input_ids.clamp_min(0))
+        if images:
+            features=[]
+            for pixels,grid in images:
+                visual=self.encoder.vision_tower(pixels).squeeze(0)
+                features.append(self.encoder.mm_projector(visual,grid.to(visual.device)))
+            visual=torch.cat(features)
+            mask=input_ids.eq(-500)
+            if int(mask.sum())!=len(visual):raise ValueError('Image placeholder count does not match vision features')
+            embeds[mask]=visual.to(embeds)
+        result=(None,position_ids,attention_mask,past_key_values,embeds,labels)
+        return result+(input_ids,) if return_inputs else result
+    def encode_image_gen(self, images, enc=False):
+        ref=next(self.encoder.vqvae.parameters())
+        quant,_,(_,_,ids)=self.encoder.vqvae.encode(images.to(ref))
+        return ids.reshape(len(images),-1),tuple(quant.shape[-2:])
+    def decode_image_gen(self, codes, height, width):
+        shape=(height//16,width//16)
+        if codes.shape[-1] != math.prod(shape):
+            codes=strip_image_row_markers(codes,shape,markers_per_row=2)
+        if (codes>=131072).any():raise ValueError('Unresolved image tokens at decoding')
+        quant=self.encoder.vqvae.quantize.get_codebook_entry(codes)
+        quant=quant.reshape(len(codes),*shape,256).permute(0,3,1,2).contiguous()
+        pixels=self.encoder.vqvae.decode(quant).float().clamp(-1,1)
+        return ((pixels+1)/2*255).permute(0,2,3,1).cpu().numpy().astype('uint8')
+    def _sample(self,prompt,*,image=None,resolution=1024,steps=64,refinement_steps=4,guidance_scale=3.,image_guidance_scale=2.,temperature=.8,seed=None,top_k=0,top_p=1.,shift=5.,show_progress=False):
+        if resolution not in (256,512,1024):raise ValueError('resolution must be 256, 512, or 1024')
+        if steps<1 or not math.isfinite(temperature) or temperature<=0:raise ValueError('steps and temperature must be positive')
+        # A supplied seed is local to this call and does not disturb caller RNG state.
+        devices=[self.device.index if self.device.index is not None else torch.cuda.current_device()] if self.device.type=='cuda' else []
+        with torch.random.fork_rng(devices=devices,enabled=seed is not None):
+            if seed is not None:torch.manual_seed(seed)
+            return sample_image(self,prompt,tokenizer=self.tokenizer,edit_image=image,image_processor=self.encoder.vision_tower.image_processor,
+                image_resolution=resolution,n_tokens=(resolution//16)**2,n_steps=steps,image_embedding_refinement_steps=refinement_steps,
+                guidance_scale=guidance_scale,guidance_scale_image=image_guidance_scale,temperature=temperature,
+                confidence_policy='stratified',shift=shift,schedule='shift',alg_temp=5.,dynamic_temperature=True,
+                min_temperature=.5,schedule_temp='cosine2',top_k=top_k,top_p=top_p,disable_tqdm=not show_progress,
+                micro_cond=f'ORIGINAL WIDTH : {resolution}; ORIGINAL HEIGHT : {resolution}; TOP : 0; LEFT : 0; SCORE : 7.110;',
+                template='Generate an image with the caption:\n <prompt>',edit_mode=1 if image is not None else 0)
+    @torch.inference_mode()
+    def text_to_image(self,prompt,*,resolution=1024,steps=64,refinement_steps=4,guidance_scale=3.,temperature=.8,seed=None,top_k=0,top_p=1.,shift=5.,show_progress=False):
+        """Generate one PIL image using the checkpoint's semantic-tree sampler."""
+        return self._sample(prompt,resolution=resolution,steps=steps,refinement_steps=refinement_steps,guidance_scale=guidance_scale,temperature=temperature,seed=seed,top_k=top_k,top_p=top_p,shift=shift,show_progress=show_progress)
+    @torch.inference_mode()
+    def image_edit(self,image,prompt,*,resolution=1024,steps=64,refinement_steps=4,guidance_scale=1.6,image_guidance_scale=2.,temperature=.8,seed=None,top_k=0,top_p=1.,shift=5.,crop_to_source=True,show_progress=False):
+        """Edit a PIL image or local image path; return a PIL image."""
+        if isinstance(image,(str,Path)):image=Image.open(image)
+        image=image.convert('RGB')
+        result=self._sample(prompt,image=image,resolution=resolution,steps=steps,refinement_steps=refinement_steps,guidance_scale=guidance_scale,image_guidance_scale=image_guidance_scale,temperature=temperature,seed=seed,top_k=top_k,top_p=top_p,shift=shift,show_progress=show_progress)
+        if crop_to_source:
+            w,h=image.size; scale=resolution/max(w,h)
+            cw,ch=round(w*scale),round(h*scale)
+            left,top=(resolution-cw)//2,(resolution-ch)//2
+            result=result.crop((left,top,left+cw,top+ch)).resize((w,h),Image.Resampling.LANCZOS)
+        return result
+    @torch.inference_mode()
+    def generate(self,prompt=None,*,image=None,inputs=None,input_ids=None,max_new_tokens=16000,block_length=32,steps_per_block=32,threshold=.9,temperature=0.,return_token_ids=False,**kwargs):
+        """Visual/text understanding, using greedy block diffusion without thinking."""
+        if kwargs:raise TypeError(f'Unknown inference options: {sorted(kwargs)}')
+        if max_new_tokens<1 or block_length<1 or steps_per_block<1:raise ValueError('Token and step counts must be positive')
+        if not 0<=threshold<=1:raise ValueError('threshold must be between 0 and 1')
+        if isinstance(prompt,torch.Tensor):input_ids,prompt=prompt,None
+        if inputs is not None:input_ids=inputs
+        if input_ids is None:
+            if not isinstance(prompt,str):raise TypeError('prompt must be a string, or provide input_ids')
+            conv=Conversation(); images=[]
+            if image is not None:
+                if isinstance(image,(str,Path)):image=Image.open(image)
+                images=process_images([image.convert('RGB')],self.encoder.vision_tower.image_processor,self.config)
+                prompt='<image>\n'+prompt
+            conv.append_message(conv.roles[0],prompt)
+            conv.append_message(conv.roles[1],None)
+            rendered=conv.get_prompt()
+            if images:
+                grid=images[0][1]
+                rendered=fix_one_prompt(rendered,self.config,[(int(grid[0,1])//2,int(grid[0,2])//2)])
+            input_ids=tokenizer_image_token(rendered,self.tokenizer,return_tensors='pt',constant_version='ministral').unsqueeze(0).to(self.device)
+            embeds=self.prepare_inputs_labels_for_multimodal(input_ids,None,None,None,None,images)[4]
+        else:
+            if image is not None:raise ValueError('Use a string prompt when providing image')
+            input_ids=input_ids.to(self.device)
+            embeds=self.encoder.embed_tokens(input_ids)
+        length=math.ceil(max_new_tokens/block_length)*block_length
+        ids,_=generate_with_prefix_cache_block_diff(self,prompt_embeds=embeds,gen_length=length,steps=steps_per_block*(length//block_length),block_length=block_length,temperature=temperature,mask_id=100,threshold=threshold,shift_logits=True,causal_context=True,eos_token_id=11)
+        ids=ids[:, :max_new_tokens]
+        if return_token_ids:return ids
+        return self.tokenizer.decode(ids[0],skip_special_tokens=True).strip()
+
+from .positions import _generation_grid_shape as _bundle_positions
+
+from .core import rotate_half as _bundle_core
+
+from .refinement import _validate_refinement_steps as _bundle_refinement
+
+from .conversation import Conversation
+from .tree import BalancedSemanticTree, production_semantic_tree_contract as _bundle_conversation
+
+from .constants import NemotronMinistralDLMConstant as _bundle_constants
+
+from .chat_utils import SimpleOutputWithPast as _bundle_chat_utils
+
+from .tree import _plain_positive_int as _bundle_tree
+
+from .schedules import cosine_schedule_2 as _bundle_schedules
+
+from .radio_cls import ClsToken as _bundle_radio_cls
+
+from .padding import resolve_source_padding_vq_lock as _bundle_padding
+
+from .pmj import stratified_random as _bundle_pmj
+
+from .image_utils import pad_to_square_and_resize as _bundle_image_utils
+
+from .sampling import _load_prediction_image as _bundle_sampling
+
+from .markers import _validate_markers_per_row as _bundle_markers
+
+from .radio_patch import ViTPatchGenerator as _bundle_radio_patch
+
+from .vision import RadioBackbone as _bundle_vision
+
+from .sbd_attention import build_image_attention_doc_ids as _bundle_sbd_attention
+
+from .attention import fused_flex_attention as _bundle_attention
+
+from .vq import Codebook as _bundle_vq
+
+from .utils import pad_along_last_dim as _bundle_utils
+
+from .vq_model import nonlinearity as _bundle_vq_model
+
+from .projector import SquaredReLU as _bundle_projector
+
+from .bridge import PackedLinearImageCodec as _bundle_bridge
+from .token_constants import IMAGE_PAD_TOKEN_INDEX as _bundle_token_constants

padding.py

@@ -0,0 +1,90 @@
+
+from __future__ import annotations
+import numpy as np
+import torch
+
+def resolve_source_padding_vq_lock(model_config, override=None, *, has_edit_image):
+    if not isinstance(has_edit_image, bool):
+        raise TypeError('has_edit_image must be boolean')
+    if override is None:
+        resolved = getattr(model_config, 'lock_source_padding_vq', False)
+        if resolved is None:
+            resolved = False
+        if not isinstance(resolved, bool):
+            raise TypeError('config.lock_source_padding_vq must be boolean or null')
+        if resolved and (not has_edit_image):
+            return False
+        return resolved
+    if not isinstance(override, bool):
+        raise TypeError('lock_source_padding_vq override must be boolean or null')
+    if override and (not has_edit_image):
+        raise ValueError('lock_source_padding_vq requires an edit_image')
+    return override
+
+def _resolve_source_native_size(value, *, fallback):
+    resolved = fallback if value is None else value
+    if not isinstance(resolved, (tuple, list)) or len(resolved) != 2 or any((isinstance(item, bool) or not isinstance(item, (int, np.integer)) or int(item) < 1 for item in resolved)):
+        raise ValueError('edit_source_native_size must contain positive integer width and height')
+    return (int(resolved[0]), int(resolved[1]))
+
+def _source_padding_vq_mask(source_native_size, raw_grid_shape, *, device):
+    (width, height) = _resolve_source_native_size(source_native_size, fallback=source_native_size)
+    if not isinstance(raw_grid_shape, (tuple, list)) or len(raw_grid_shape) != 2 or any((isinstance(item, bool) or not isinstance(item, (int, np.integer)) or int(item) < 1 for item in raw_grid_shape)):
+        raise ValueError('raw_grid_shape must contain positive integer height and width')
+    (grid_height, grid_width) = (int(value) for value in raw_grid_shape)
+    if grid_height != grid_width:
+        raise ValueError('source-padding VQ locking requires a square raw VQ grid')
+    side = max(width, height)
+    left = round((side - width) * 0.5)
+    top = round((side - height) * 0.5)
+    right = left + width
+    bottom = top + height
+    columns = torch.arange(grid_width, dtype=torch.long, device=device)
+    rows = torch.arange(grid_height, dtype=torch.long, device=device)
+    horizontally_outside = ((columns + 1) * side <= left * grid_width) | (columns * side >= right * grid_width)
+    vertically_outside = ((rows + 1) * side <= top * grid_height) | (rows * side >= bottom * grid_height)
+    return (vertically_outside[:, None] | horizontally_outside[None, :]).reshape(1, -1)
+
+def _augment_source_padding_vq_mask(state, raw_padding_mask, content_positions):
+    if not torch.is_tensor(state) or state.ndim != 2:
+        raise ValueError('generation state must have shape [batch, tokens]')
+    if not torch.is_tensor(raw_padding_mask) or raw_padding_mask.dtype is not torch.bool or raw_padding_mask.ndim != 2 or (raw_padding_mask.shape[0] != state.shape[0]) or (raw_padding_mask.device != state.device):
+        raise TypeError('raw_padding_mask must be device-local bool with the state batch size')
+    if not torch.is_tensor(content_positions) or content_positions.ndim != 1 or content_positions.dtype is not torch.long or (content_positions.device != state.device) or (content_positions.numel() != raw_padding_mask.shape[1]):
+        raise TypeError('content_positions must be a device-local long vector covering raw cells')
+    if bool((content_positions < 0).any()) or bool((content_positions >= state.shape[1]).any()):
+        raise ValueError('content_positions contains an out-of-range state index')
+    locked_mask = torch.zeros_like(state, dtype=torch.bool)
+    locked_mask[:, content_positions] = raw_padding_mask
+    return locked_mask
+
+def _prefill_source_padding_vq(state, source_codes, raw_padding_mask, content_positions):
+    if not torch.is_tensor(state) or state.ndim != 2:
+        raise ValueError('generation state must have shape [batch, tokens]')
+    if not torch.is_tensor(source_codes) or source_codes.ndim != 2 or source_codes.dtype != state.dtype or (source_codes.device != state.device):
+        raise TypeError('source_codes must match the state integer dtype and device')
+    if source_codes.shape != raw_padding_mask.shape:
+        raise TypeError('raw_padding_mask must be bool and match source_codes')
+    locked_mask = _augment_source_padding_vq_mask(state, raw_padding_mask, content_positions)
+    locked_codes = torch.zeros_like(state)
+    locked_codes[:, content_positions] = source_codes
+    prefilled = state.clone()
+    prefilled[locked_mask] = locked_codes[locked_mask]
+    return (prefilled, locked_mask, locked_codes)
+
+def _restore_source_padding_vq(state, locked_mask, locked_codes):
+    if locked_mask is None:
+        return state
+    if not torch.is_tensor(state) or state.ndim != 2 or (not torch.is_tensor(locked_mask)) or (locked_mask.dtype is not torch.bool) or (locked_mask.shape != state.shape) or (locked_mask.device != state.device) or (not torch.is_tensor(locked_codes)) or (locked_codes.shape != state.shape) or (locked_codes.dtype != state.dtype) or (locked_codes.device != state.device):
+        raise TypeError('locked source-padding tensors must align with sampler state')
+    restored = state.clone()
+    restored[locked_mask] = locked_codes[locked_mask]
+    return restored
+
+def _assert_source_padding_vq_locked(state, locked_mask, locked_codes):
+    if locked_mask is None:
+        return
+    if not torch.is_tensor(locked_mask) or locked_mask.dtype is not torch.bool or locked_mask.shape != state.shape or (not torch.is_tensor(locked_codes)) or (locked_codes.shape != state.shape):
+        raise TypeError('locked source-padding tensors must align with sampler state')
+    if not torch.equal(state[locked_mask], locked_codes[locked_mask]):
+        raise RuntimeError('a locked source-padding VQ code changed during sampling')

pca_tree.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4b7ee94fbb10fd4146a4ad167b595b037e3f591eaa12ecf3f0b411c82384a714
+size 1052237

pmj.py

@@ -0,0 +1,67 @@
+# ADOBE CONFIDENTIAL
+# Copyright 2025 Adobe
+# All Rights Reserved.
+# NOTICE: All information contained herein is, and remains
+# the property of Adobe and its suppliers, if any. The intellectual
+# and technical concepts contained herein are proprietary to Adobe
+# and its suppliers and are protected by all applicable intellectual
+# property laws, including trade secret and copyright laws.
+# Dissemination of this information or reproduction of this material
+# is strictly forbidden unless prior written permission is obtained
+# from Adobe.
+
+from __future__ import annotations
+from typing import List, Tuple, Optional
+import random
+
+def stratified_random(n: int=64, seed: Optional[int]=None, shuffle_blocks: bool=True) -> List[int]:
+    if n <= 0 or n & n - 1 != 0:
+        raise ValueError('n must be a positive power of two (e.g., 64)')
+    rng = random.Random(seed)
+    occupied = [[False] * n for _ in range(n)]
+    seq: List[int] = []
+    blocks: List[Tuple[int, int, int]] = [(0, 0, n)]
+
+    def block_has_sample(x0: int, y0: int, size: int) -> bool:
+        for yy in range(y0, y0 + size):
+            row = occupied[yy]
+            for xx in range(x0, x0 + size):
+                if row[xx]:
+                    return True
+        return False
+
+    def place_random_in_block(x0: int, y0: int, size: int):
+        x = rng.randrange(x0, x0 + size)
+        y = rng.randrange(y0, y0 + size)
+        attempts = 0
+        while occupied[y][x]:
+            x = rng.randrange(x0, x0 + size)
+            y = rng.randrange(y0, y0 + size)
+            attempts += 1
+            if attempts > 10000:
+                raise RuntimeError('Too many attempts to place a sample; logic error?')
+        occupied[y][x] = True
+        seq.append(y * n + x)
+    size = n
+    while size > 1:
+        half = size // 2
+        children: List[Tuple[int, int, int]] = []
+        for (x0, y0, s) in blocks:
+            assert s == size
+            children.extend([(x0, y0, half), (x0 + half, y0, half), (x0, y0 + half, half), (x0 + half, y0 + half, half)])
+        if shuffle_blocks:
+            rng.shuffle(children)
+        for (x0, y0, s) in children:
+            if not block_has_sample(x0, y0, s):
+                place_random_in_block(x0, y0, s)
+        blocks = children
+        size = half
+    remaining: List[int] = []
+    for y in range(n):
+        for x in range(n):
+            if not occupied[y][x]:
+                remaining.append(y * n + x)
+    rng.shuffle(remaining)
+    seq.extend(remaining)
+    assert len(seq) == n * n, (len(seq), n * n)
+    return seq

positions.py

@@ -0,0 +1,74 @@
+# ADOBE CONFIDENTIAL
+# Copyright 2025 Adobe
+# All Rights Reserved.
+# NOTICE: All information contained herein is, and remains
+# the property of Adobe and its suppliers, if any. The intellectual
+# and technical concepts contained herein are proprietary to Adobe
+# and its suppliers and are protected by all applicable intellectual
+# property laws, including trade secret and copyright laws.
+# Dissemination of this information or reproduction of this material
+# is strictly forbidden unless prior written permission is obtained
+# from Adobe.
+
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+
+from __future__ import annotations
+from numbers import Integral
+import torch
+GENERATION_2D_SINCOS_SCALE = 0.0018
+GENERATION_2D_SINCOS_CANONICAL_CONTENT_SIZE = 32
+
+def _generation_grid_shape(gen_shape, *, compressed: bool, encoder_compressed: bool, token_count: int) -> tuple[int, int]:
+    if not isinstance(gen_shape, (tuple, list)) or len(gen_shape) != 2:
+        raise ValueError('gen_enc_add_pos_emb requires gen_shape=(height, width)')
+    if any((isinstance(value, bool) or not isinstance(value, Integral) for value in gen_shape)):
+        raise TypeError('generation image height and width must be integers')
+    (height, width) = (int(value) for value in gen_shape)
+    if height < 1 or width < 1:
+        raise ValueError('generation image height and width must be positive')
+    for (stage_name, stage_enabled) in (('generation', compressed), ('generation-encoder', encoder_compressed)):
+        if not stage_enabled:
+            continue
+        if height % 2 or width % 2:
+            raise ValueError(f'{stage_name} compression requires an even image-token grid')
+        height //= 2
+        width //= 2
+    if height * width != token_count:
+        raise ValueError(f'generation positional grid does not match the transformer image tokens: {height}x{width} != {token_count}')
+    return (height, width)
+
+def _build_2d_sincos_position_embedding(height: int, width: int, embedding_dim: int, *, marker_column: bool=False, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    if embedding_dim % 4:
+        raise ValueError('gen_enc_add_pos_emb requires d_model_gen divisible by 4')
+    if not isinstance(marker_column, bool):
+        raise TypeError('marker_column must be boolean')
+    content_width = width - int(marker_column)
+    if height < 1 or content_width < 1:
+        raise ValueError('generation positional grid must contain positive content height and width')
+    axis_frequency_count = embedding_dim // 4
+    frequencies = torch.arange(axis_frequency_count, device=device, dtype=torch.float32)
+    frequencies = 10000.0 ** (-frequencies / axis_frequency_count)
+    canonical_size = float(GENERATION_2D_SINCOS_CANONICAL_CONTENT_SIZE)
+    column_stride = canonical_size / float(content_width)
+    row_stride = canonical_size / float(height)
+    columns = torch.arange(width, device=device, dtype=torch.float32) * column_stride
+    rows = torch.arange(height, device=device, dtype=torch.float32) * row_stride
+    column_angles = torch.outer(columns, frequencies)
+    row_angles = torch.outer(rows, frequencies)
+    column_embedding = torch.cat((column_angles.sin(), column_angles.cos()), dim=-1)
+    row_embedding = torch.cat((row_angles.sin(), row_angles.cos()), dim=-1)
+    position_embedding = torch.cat((column_embedding.unsqueeze(0).expand(height, -1, -1), row_embedding.unsqueeze(1).expand(-1, width, -1)), dim=-1)
+    return position_embedding.reshape(height * width, embedding_dim).to(dtype=dtype)

projector.py

@@ -0,0 +1,66 @@
+
+from __future__ import annotations
+import torch
+import torch.nn as nn
+
+class SquaredReLU(nn.Module):
+
+    def forward(self, x):
+        return torch.pow(torch.nn.functional.relu(x), 2)
+
+class RMSNorm(nn.Module):
+
+    def __init__(self, hidden_size, eps=1e-05):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        return (self.weight.to(torch.float32) * hidden_states).to(input_dtype)
+
+class RadioHProjector(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.downsample_ratio = 0.5
+        self.ps_version = 'v2'
+        vit_hidden_size = 1280
+        vision_projection_hidden_size = getattr(config, 'projector_hidden_size', config.hidden_size)
+        llm_hidden_size = config.hidden_size
+        in_dim = int(vit_hidden_size * (1 / self.downsample_ratio) ** 2)
+        self.mlp = nn.Sequential(RMSNorm(in_dim, eps=1e-05), nn.Linear(in_dim, vision_projection_hidden_size, bias=False), SquaredReLU(), nn.Linear(vision_projection_hidden_size, llm_hidden_size, bias=False))
+        self.radio_multiply = True
+
+    def pixel_shuffle(self, x, scale_factor=0.5):
+        (n, w, h, c) = x.size()
+        x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))
+        x = x.permute(0, 2, 1, 3).contiguous()
+        x = x.view(n, int(h * scale_factor), int(w * scale_factor), int(c / (scale_factor * scale_factor)))
+        if self.ps_version != 'v1':
+            x = x.permute(0, 2, 1, 3).contiguous()
+        return x
+
+    def forward(self, image_features: torch.Tensor, image_grid_thw: torch.Tensor):
+        device = image_features.device
+        indices = []
+        offset = 0
+        for (_, H, W) in image_grid_thw.tolist():
+            (H, W) = (int(H), int(W))
+            (H2, W2) = (H // 2, W // 2)
+            k = torch.arange(H2 * W2, device=device)
+            if self.ps_version == 'v1':
+                (out_r, out_c) = (k % H2, k // H2)
+            else:
+                (out_r, out_c) = (k // W2, k % W2)
+            tl = offset + 2 * out_r * W + 2 * out_c
+            indices.append(torch.stack([tl, tl + 1, tl + W, tl + W + 1], dim=1).reshape(-1))
+            offset += H * W
+        gather_idx = torch.cat(indices)
+        merged = image_features[gather_idx].view(-1, 4 * image_features.shape[-1])
+        if self.radio_multiply:
+            return 0.01 * self.mlp(merged)
+        return self.mlp(merged)

radio_cls.py

@@ -0,0 +1,44 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+from __future__ import annotations
+from typing import Optional
+import torch
+from torch import nn
+
+class ClsToken(nn.Module):
+
+    def __init__(self, ndim: int, num_tokens: int=1, enabled: bool=True, register_multiple: Optional[int]=None, num_registers: Optional[int]=None):
+        super().__init__()
+        self.ndim = ndim
+        self.enabled = enabled
+        self.num_registers = 0
+        self.num_tokens = num_tokens
+        if enabled:
+            if num_registers:
+                self.num_registers = num_registers
+            elif register_multiple:
+                self.num_registers = register_multiple - num_tokens % register_multiple
+            scale = ndim ** -0.5
+            self.token = nn.Parameter(torch.randn(num_tokens + self.num_registers, ndim) * scale)
+        else:
+            self.token = None
+        self.num_patches = self.num_tokens + self.num_registers
+
+    def disable(self):
+        self.token = None
+        self.enabled = False
+
+    def forward(self, x: torch.Tensor):
+        if self.token is None:
+            return x
+        token = self.token.unsqueeze(0).expand(x.shape[0], -1, -1)
+        x = torch.cat([token, x], dim=1)
+        return x
+
+    def no_weight_decay(self):
+        return ['token']

radio_patch.py

@@ -0,0 +1,167 @@
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from __future__ import annotations
+import math
+from typing import Union, Tuple, Optional
+import torch
+import torch.nn.functional as F
+from torch import nn
+from einops import rearrange
+from .radio_cls import ClsToken
+input_dim_t = Union[int, Tuple[int, int]]
+
+class ViTPatchGenerator(nn.Module):
+
+    def __init__(self, patch_size: int, embed_dim: int, input_dims: input_dim_t, abs_pos: bool=True, normalize_patches: bool=False, cls_token: bool=False, max_input_dims: Optional[input_dim_t]=None, pos_dropout: float=0.0, return_pos_enc: bool=False, num_cls_tokens: int=1, register_multiple: Optional[int]=None, num_registers: Optional[int]=None, patch_bias: bool=False, device=None, dtype=None):
+        super().__init__()
+        if isinstance(input_dims, int):
+            input_dims = (input_dims, input_dims)
+        if max_input_dims is None:
+            max_input_dims = input_dims
+        if isinstance(max_input_dims, int):
+            max_input_dims = (max_input_dims, max_input_dims)
+        max_input_dims = tuple((int(math.ceil(d / patch_size) * patch_size) for d in max_input_dims))
+        self.cpe_mode = max_input_dims != input_dims
+        self.pos_dropout = pos_dropout
+        self.return_pos_enc = return_pos_enc
+        factory = dict(device=device, dtype=dtype)
+        self.patch_size = patch_size
+        self.abs_pos = abs_pos
+        self.embed_dim = embed_dim
+        self.num_rows = max_input_dims[0] // patch_size
+        self.num_cols = max_input_dims[1] // patch_size
+        self.input_dims = tuple((d // patch_size for d in input_dims))
+        self.num_patches = self.num_rows * self.num_cols
+        self.max_input_dims = max_input_dims
+        self.im_to_patches = Im2Patches(patch_size)
+        self.embedder = ViTPatchLinear(patch_size, embed_dim, bias=patch_bias, **factory)
+        if abs_pos:
+            scale = embed_dim ** -0.5
+            self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches, embed_dim, **factory) * scale)
+        self.cls_token = ClsToken(embed_dim, num_tokens=num_cls_tokens, enabled=cls_token, register_multiple=register_multiple, num_registers=num_registers)
+        self.patch_normalizer = nn.LayerNorm(embed_dim) if normalize_patches else nn.Identity()
+        self.num_video_frames = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        patches = self.embed_patches(x)
+        (patches, pos_enc) = self.apply_pos_enc(patches, input_size=x.shape[2:])
+        patches = self.cls_token(patches)
+        patches = self.patch_normalizer(patches)
+        if self.return_pos_enc:
+            return (patches, pos_enc)
+        return patches
+
+    @property
+    def apply_cls_token(self):
+        return self.cls_token.enabled
+
+    @property
+    def num_cls_tokens(self):
+        return self.cls_token.num_tokens
+
+    @property
+    def num_cls_patches(self):
+        return self.cls_token.num_patches
+
+    @property
+    def num_registers(self):
+        return self.cls_token.num_registers
+
+    @property
+    def num_skip(self):
+        return self.num_cls_tokens + self.num_registers
+
+    def no_weight_decay(self):
+        return ['pos_embed']
+
+    def _load_embed(self, src_embed: torch.Tensor, targ_embed: nn.Parameter):
+        if src_embed.shape != targ_embed.shape:
+            src_size = int(math.sqrt(src_embed.shape[1]))
+            assert src_size ** 2 == src_embed.shape[1], 'Unable to interpolate non-square embedding'
+            src_embed = rearrange(src_embed, 'b (h w) c -> b c h w', h=src_size, w=src_size)
+            src_embed = F.interpolate(src_embed, size=(self.num_rows, self.num_cols), mode='bicubic', align_corners=True, antialias=False)
+            src_embed = rearrange(src_embed, 'b c h w -> b (h w) c')
+        targ_embed.data.copy_(src_embed)
+
+    def _load_projection(self, src_proj_weight: torch.Tensor, targ_proj_weight: torch.Tensor):
+        if src_proj_weight.shape != targ_proj_weight.shape:
+            src_patch_size = int(math.sqrt(src_proj_weight.shape[1] // 3))
+            assert src_patch_size ** 2 * 3 == src_proj_weight.shape[1], 'Unable to interpolate non-square patch size'
+            src_proj_weight = rearrange(src_proj_weight, 'b (c h w) -> b c h w', c=3, h=src_patch_size, w=src_patch_size)
+            src_proj_weight = F.interpolate(src_proj_weight, size=(self.patch_size, self.patch_size), mode='bicubic', align_corners=True, antialias=False)
+            src_proj_weight = rearrange(src_proj_weight, 'b c h w -> b (c h w)')
+        targ_proj_weight.data.copy_(src_proj_weight)
+
+    def embed_patches(self, x: torch.Tensor) -> torch.Tensor:
+        patches = self.im_to_patches(x)
+        patches = self.embedder(patches)
+        return patches
+
+    def apply_pos_enc(self, patches: torch.Tensor, patch_idxs: Optional[torch.Tensor]=None, input_size: Optional[Tuple[int, int]]=None) -> torch.Tensor:
+        if not self.abs_pos:
+            return patches
+        pos_enc = self.get_pos_enc(patches.shape[0], patch_idxs, input_size)
+        pos_enc_drop = pos_enc
+        return (patches + pos_enc_drop, pos_enc)
+
+    def get_pos_enc(self, batch_size: int, patch_idxs: Optional[torch.Tensor]=None, input_size: Optional[Tuple[int, int]]=None) -> torch.Tensor:
+        if input_size is None:
+            input_dims = self.input_dims
+        else:
+            input_dims = tuple((d // self.patch_size for d in input_size))
+        pos_embed = self._get_pos_embeddings(batch_size, input_dims)
+        if patch_idxs is None:
+            return pos_embed
+        exp_patch_idxs = patch_idxs.unsqueeze(-1).expand(-1, -1, pos_embed.shape[-1])
+        pos_embed = torch.gather(pos_embed.expand(patch_idxs.shape[0], -1, -1), dim=1, index=exp_patch_idxs)
+        return pos_embed
+
+    def _get_pos_embeddings(self, batch_size: int, input_dims: Tuple[int, int]):
+        if (self.num_rows, self.num_cols) == input_dims:
+            return self.pos_embed
+        pos_embed = self.pos_embed.reshape(1, self.num_rows, self.num_cols, -1).permute(0, 3, 1, 2)
+
+        def window_select(pos_embed):
+            if input_dims[0] < pos_embed.shape[-2]:
+                pos_embed = pos_embed[..., :input_dims[0], :]
+            if input_dims[1] < pos_embed.shape[-1]:
+                pos_embed = pos_embed[..., :, :input_dims[1]]
+            return pos_embed
+        if self.cpe_mode:
+            max_dim = max(input_dims)
+            pos_embed = F.interpolate(pos_embed.float(), size=(max_dim, max_dim), align_corners=False, mode='bilinear').to(pos_embed.dtype)
+            pos_embed = window_select(pos_embed)
+        else:
+            pos_embed = window_select(pos_embed)
+        if pos_embed.shape[-2:] != input_dims:
+            pos_embed = F.interpolate(pos_embed.float(), size=input_dims, align_corners=False, mode='bilinear').to(pos_embed.dtype)
+        pos_embed = pos_embed.flatten(2).permute(0, 2, 1)
+        return pos_embed
+
+class Im2Patches(nn.Module):
+
+    def __init__(self, patch_size: int):
+        super().__init__()
+        self.patch_size = patch_size
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.patch_size == 1:
+            patches = x.flatten(2)
+            patches = patches.permute(0, 2, 1)
+            return patches
+        py = x.shape[-2] // self.patch_size
+        px = x.shape[-1] // self.patch_size
+        patches = rearrange(x, 'b c (py yy) (px xx) -> b (py px) (c yy xx)', py=py, yy=self.patch_size, px=px, xx=self.patch_size)
+        return patches
+
+class ViTPatchLinear(nn.Linear):
+
+    def __init__(self, patch_size: int, embed_dim: int, bias: bool=False, **factory):
+        super().__init__(3 * patch_size ** 2, embed_dim, bias=bias, **factory)
+        self.patch_size = patch_size

refinement.py

@@ -0,0 +1,19 @@
+
+from __future__ import annotations
+from numbers import Integral
+import torch
+
+def _validate_refinement_steps(value: object, *, source: str) -> int:
+    if isinstance(value, bool) or not isinstance(value, Integral):
+        raise TypeError(f'{source} must be a positive integer')
+    value = int(value)
+    if value < 1:
+        raise ValueError(f'{source} must be a positive integer')
+    return value
+
+def refinement_mask_fraction(ages: torch.Tensor, *, steps: int) -> torch.Tensor:
+    steps = _validate_refinement_steps(steps, source='steps')
+    if ages.dtype == torch.bool or ages.is_floating_point():
+        raise TypeError('ages must use an integer dtype')
+    effective_age = ages.clamp(min=0, max=steps).to(dtype=torch.float32)
+    return (1.0 - effective_age / steps).mean(dim=-1)

release_manifest.json

@@ -0,0 +1,54 @@
+{
+  "revision": "generation-only-step-10000",
+  "sha256": {
+    ".gitattributes": "50360a3ded843d95a2667395b3131f6d499b905d941331a4e84e7f0351769023",
+    ".gitignore": "e1900f365321556fa39bfe45527be4e15a1de7fe96b71ca8d294b690a1026a8d",
+    "README.md": "50038fb0e318e290c123382d4f6284ce80699bd60d7fabf0ee6db27419e12b07",
+    "attention.py": "5617bdd53aaa0d6f0e79745c5e811f36335a0c1f458218ab1fc02c4606d79f37",
+    "bridge.py": "0e20f9ad1c81912489dd7c5063507b7fb72ead7d1895467442277e06ee3bc846",
+    "chat_template.jinja": "b80135b85225bef1a5129eb2068d0d5df4975edf4ffb102a50de1fb069c9e78a",
+    "chat_utils.py": "0683087dbadea1ce510d77816eed8ccfba998aa6c64833fd93c4e46c6edbc734",
+    "config.json": "9dd7b2514b477712b9e023dec3f0d3749e8ccd9e27869cb4bf70fe9a8097fff0",
+    "configuration_nlduni.py": "fa6d15657c9fadeb1f2bc142ea33b73c225374d5646565a077b43d6fd3e25f38",
+    "constants.py": "1149095b05a6da911aae371d4899edcaba154352dc30efd2d5b72bb8767f785f",
+    "conversation.py": "5e11082e6f5201978a0eb99523120d1272a6e0b311bfa75b273fa02d927d1002",
+    "core.py": "4dad9ea004daa31236fb1eb628d2ad1b46d310914114bd93c401fd41161d5c66",
+    "generation_config.json": "35e13707e4e588edaa4d7c12dff1d9ccd49efe0e07c99e3392c42b88f9398faa",
+    "image_utils.py": "62755735051898eb909937712e0badcf6614d30405ea51d7d33522f484a319d1",
+    "inference.py": "37e0b19b3d003621a620b226d6a9d39725e9450a7a6431d992ceea49cb2b92b2",
+    "markers.py": "17f69056ff9db548dc223971021a91444ecaef3257df5b5f30fd2e5230bbadc7",
+    "model-00001-of-00008.safetensors": "c4d490e010475dde3e2b57de64261a56d6bf8e9ca7d077754913fd628f7de5f6",
+    "model-00002-of-00008.safetensors": "933c4b2f44f9cd7be30ac7a40efbb3c58485aa9fbd6396245d8e2337faf5c5b3",
+    "model-00003-of-00008.safetensors": "6893a316ec6655eb3f518be0f3ae3dfc80621922741df7d14f12b3d9fd123bf6",
+    "model-00004-of-00008.safetensors": "858fab295f100b897634b72824d01e01f368dc1d500de03916854d7b2769986c",
+    "model-00005-of-00008.safetensors": "0d31608093622df48c2284fa310170f9a1bfd127d7108bdacbc8106fcf296691",
+    "model-00006-of-00008.safetensors": "4356b837d798b2d3250cd067c15b7f5d56fd19a8b8605bb107fb9d727c1726d0",
+    "model-00007-of-00008.safetensors": "003f5a6c4c727e5f560e16ad91ee2264ed138c7190108dfc9b68d1b8c8445545",
+    "model-00008-of-00008.safetensors": "299783fae2f540dfd723abbb09f8661b57122210d44b23a4c65f24b254542550",
+    "model.safetensors.index.json": "f2c7c2353c2a2eb70e7f80889c4fca85caceba217f1ccd9aa6647a7c644f6e30",
+    "modeling_nlduni.py": "d7a04c330d1cd14fc638297a61a674cfdc666d8eb5ee24c41625ff1875040021",
+    "padding.py": "53ba2cc8b37eec837b4fd0a3932a5f2d55051bbe5796c848921ce95171336dce",
+    "pca_tree.pt": "4b7ee94fbb10fd4146a4ad167b595b037e3f591eaa12ecf3f0b411c82384a714",
+    "pmj.py": "33920c0bb48ab99bafb3dde99f0e44eff7e6004596777d2e09ec4d07e039bfa3",
+    "positions.py": "3e0ee23d4bd59059690bc9962fa0948c18f52af13c1295d5322309ba6472f0c4",
+    "projector.py": "b402ae3a9ca1d7867c139e3180f905411d4932901cd7ce19910865db8709691f",
+    "radio_cls.py": "6c04ff0ac364676b7ca2610d8103c53f3b7a26d5d4b6cfd593415556e721c710",
+    "radio_patch.py": "41d63126936a4691d28979f9bb7a7b35c2daa9f10c8e98ef909656930ec861c2",
+    "refinement.py": "b69a5bcfec2ecdbc97d5b5796dbcd6dfa83d29fe8b835017a18b2e96596891ca",
+    "requirements.txt": "a3725eb78eb9ed882a46b2b0370f4b93eafc93e59d6e8db38698963016b9405d",
+    "sampling.py": "6111a83d7d737eab9c702af281f209316e7458a785d0102c4f2bf3fd4e315491",
+    "sbd_attention.py": "828f02ae59d430485507b00c76d0b85158dfa10dfea7888519ec80522798a134",
+    "schedules.py": "243d3fa6924f11397df422c9e50dc37bd4e85d882c518c3a4a75545a3b1a367e",
+    "special_tokens_map.json": "57cc1f9b3ea2a038d1107bdf4f4b3c45b5ed9ce175e7410d70070e07a3f35a76",
+    "token_constants.py": "6515bb3f128730edcd02907cf49334ac9cae2ed373611b0a260b3bdf4ff678f1",
+    "tokenizer.json": "f9ad3482369ba91711ff939f9eeb761d618b3157fcbfc0715f18e34b918eca97",
+    "tokenizer_config.json": "5eb18e9f9e0653added53a9200a8d38fdbc1d0cdb4602f8e354f30391f8cb1f3",
+    "tree.py": "b22d647f4d6def2a8df6f1b336db0b03a9cf0131dccf124f949cfa704e227c81",
+    "tree_metadata.json": "16ccc59386d2d73094ba9b2562feff45b7a171f4f37053f9b5ed756d5c44e3d5",
+    "utils.py": "6a0036a6451918de85fcc9cc3a82f5c195227bda5df58103156dd9941f6da0bc",
+    "vision.py": "375c9cd952135701b71b35581a72c46cc58bb9701392343f9817f9cf1e68bb72",
+    "vq.py": "3029623e15687b8fe95594b25410cfa2c0a82eebaa821b6ab380d1dcc394e537",
+    "vq_model.py": "0f28d77789671b9c05402de7dca7bfccaa57264336ba778a76f3b59966baced7",
+    "weights_manifest.json": "fe3af69d1ff222d3de89efd6305fcb3d396204b3c208491866d4da1137bea6f9"
+  }
+}

requirements.txt

@@ -0,0 +1,13 @@
+# Python 3.13, CUDA 12.6; install torch before building flash-attn.
+torch==2.7.0
+torchvision==0.22.0
+transformers==4.57.6
+huggingface-hub==0.36.2
+tokenizers==0.22.2
+accelerate==1.6.0
+safetensors==0.7.0
+timm==1.0.26
+einops==0.6.1
+numpy==2.4.3
+pillow==12.1.1
+tqdm>=4.66,<5

sampling.py

@@ -0,0 +1,831 @@
+# ADOBE CONFIDENTIAL
+# Copyright 2025 Adobe
+# All Rights Reserved.
+# NOTICE: All information contained herein is, and remains
+# the property of Adobe and its suppliers, if any. The intellectual
+# and technical concepts contained herein are proprietary to Adobe
+# and its suppliers and are protected by all applicable intellectual
+# property laws, including trade secret and copyright laws.
+# Dissemination of this information or reproduction of this material
+# is strictly forbidden unless prior written permission is obtained
+# from Adobe.
+
+
+from __future__ import annotations
+import os
+from .image_utils import compute_num_image_tokens, fix_one_prompt, process_images, tokenizer_image_token
+from .token_constants import IMAGE_TOKEN_INDEX, IMAGE_PAD_TOKEN_INDEX
+from .conversation import conv_templates
+from PIL import Image
+import copy
+import functools
+import math
+import torch
+from .schedules import cosine_schedule_2, logit_normal_schedule
+import time
+from .chat_utils import wte as wte_ministral
+from .chat_utils import get_logits as get_logits_ministral
+from .chat_utils import _resolve_shift_sbd_logits as resolve_ministral_shift_sbd_logits
+from .sbd_attention import generated_image_attention_window
+from tqdm.auto import tqdm
+import numpy as np
+import torch.distributions as dists
+import torch.nn.functional as F
+from .image_utils import pad_to_square_and_resize
+from .utils import pad_along_last_dim
+from .padding import _assert_source_padding_vq_locked, _prefill_source_padding_vq, _restore_source_padding_vq, _resolve_source_native_size, _source_padding_vq_mask, resolve_source_padding_vq_lock
+from .refinement import refinement_mask_fraction
+from .markers import append_image_row_markers, augmented_image_shape, image_fixed_prefix_token_count, image_generation_special_ids, image_row_marker_count, image_row_markers_enabled, resolve_generation_skip_down_sample, split_generation_fixed_prefix, strip_image_row_markers, validate_image_fixed_prefix_support, validate_image_row_marker_support
+from .tree import BalancedSemanticTree, advance_revisable_semantic_tree, build_semantic_tree_embeddings, production_semantic_tree_contract, resolve_semantic_tree_atomic_block_size, resolve_semantic_tree_refinement_gap, resolve_semantic_tree_steps, semantic_tree_activation_counts, semantic_tree_decode_schedule, semantic_tree_noise_for_ages, semantic_tree_refinement_due
+from .pmj import stratified_random
+
+def _load_prediction_image(image):
+    if isinstance(image, (str, os.PathLike)):
+        image = Image.open(image)
+    return image.convert('RGB')
+
+def _validate_text_to_image_batch_size(batch_size):
+    if batch_size != 1:
+        raise ValueError(f'text_to_image currently supports batch_size=1 only; got {batch_size}')
+
+def _prediction_constant_version(model_config):
+    model_type = str(getattr(model_config, 'model_type', ''))
+    return 'ministral'
+    return None
+
+def _move_image_to_device(value, *, device, dtype):
+    if torch.is_tensor(value):
+        return value.to(device=device, dtype=dtype)
+    if isinstance(value, tuple):
+        if not value:
+            return value
+        return (_move_image_to_device(value[0], device=device, dtype=dtype), *value[1:])
+    if isinstance(value, list):
+        return [_move_image_to_device(item, device=device, dtype=dtype) for item in value]
+    raise TypeError(f'Unsupported processed image type: {type(value).__name__}')
+
+def _tokenize_prediction_prompt_with_images(prompt, tokenizer, model, image_processor, image_tensor):
+    image_aspect_ratio = 'radio_h_native'
+    native_image_token_shapes = None
+    if 'native' in image_aspect_ratio:
+        native_images = image_tensor if isinstance(image_tensor, list) else [image_tensor]
+        native_image_token_shapes = [compute_num_image_tokens([item], image_aspect_ratio, image_processor) for item in native_images]
+        prompt = fix_one_prompt(prompt, model.config, native_image_token_shapes)
+    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt', constant_version=_prediction_constant_version(model.config))
+    if native_image_token_shapes is not None:
+        expected_slots = sum((math.prod(shape) for shape in native_image_token_shapes))
+        actual_slots = int(input_ids.eq(IMAGE_PAD_TOKEN_INDEX).sum().item())
+        if actual_slots != expected_slots:
+            raise RuntimeError(f'Native-image T2I prompt/token mismatch: expected {expected_slots} image slots but tokenization produced {actual_slots}.')
+    return (prompt, input_ids)
+
+def _transformer_image_token_count(model, raw_token_count, gen_shape):
+    raw_token_count = int(raw_token_count)
+    raw_grid_tokens = math.prod(gen_shape)
+    if raw_token_count != raw_grid_tokens:
+        raise ValueError(f'Raw generation token count must match gen_shape: {raw_token_count} != {raw_grid_tokens}')
+    generation_model = model.get_model()
+    skip_down_sample = resolve_generation_skip_down_sample(generation_model.config, runtime_value=False)
+    compression_active = (getattr(generation_model, 'packed_image_codec', None) is not None or getattr(generation_model, 'downsample_gen', None) is not None) and (not skip_down_sample)
+    if not compression_active:
+        return raw_token_count
+    (height, width) = (int(value) for value in gen_shape)
+    return height // 2 * (width // 2)
+
+def gumbel_noise(t, generator=None):
+    noise = torch.zeros_like(t).uniform_(0, 1, generator=generator)
+    return -torch.log(-torch.log(noise))
+
+def top_k_top_p_filtering(logits: torch.Tensor, top_k: int=0, top_p: float=1.0, filter_value: float=-1e309, min_tokens_to_keep: int=1) -> torch.Tensor:
+    if top_k > 0:
+        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+    if top_p < 1.0:
+        (sorted_logits, sorted_indices) = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        if min_tokens_to_keep > 1:
+            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+        logits[indices_to_remove] = filter_value
+    return logits
+from .constants import NemotronMinistralDLMConstant
+NC = NemotronMinistralDLMConstant
+
+def _resolve_ministral_sbd_diffusion_sampling(model, requested):
+    model_module = model.module if hasattr(model, 'module') else model
+    config = model_module.config
+    if requested is None:
+        requested = True
+    enabled = bool(requested)
+    return enabled
+
+def _configure_ministral_t2i_attention(model_module, *, sbd_diffusion_sampling):
+    config = model_module.config
+    for layer in model_module.encoder.layers:
+        attention = layer.self_attn
+        if sbd_diffusion_sampling:
+            attention.set_attention_mode('sbd_block_diff', block_size=config.block_size)
+        else:
+            attention.mode = 'bidirectional'
+        if hasattr(attention, 'diffusion_lm'):
+            attention.diffusion_lm = True
+
+def _restore_attention_state_after_prediction(function):
+
+    @functools.wraps(function)
+    def wrapped(model, *args, **kwargs):
+        model_module = model.module if hasattr(model, 'module') else model
+        encoder = getattr(model_module, 'encoder', None)
+        layers = getattr(encoder, 'layers', ())
+        saved_states = []
+        for layer in layers:
+            attention = layer.self_attn
+            saved_states.append((attention, getattr(attention, 'mode', None), getattr(attention, 'block_size', None), getattr(attention, 'diffusion_lm', None)))
+        try:
+            return function(model, *args, **kwargs)
+        finally:
+            for (attention, mode, block_size, diffusion_lm) in saved_states:
+                if hasattr(attention, 'mode'):
+                    attention.mode = mode
+                if hasattr(attention, 'block_size'):
+                    attention.block_size = block_size
+                if hasattr(attention, 'diffusion_lm'):
+                    attention.diffusion_lm = diffusion_lm
+    return wrapped
+
+def _load_semantic_tree_path(model_config, *, cache_owner, device, vocab_size):
+    artifact_path = getattr(model_config, 'image_token_tree_path', None)
+    if not isinstance(artifact_path, str) or not artifact_path.strip():
+        raise ValueError('semantic-tree token sampling requires config.image_token_tree_path')
+    artifact_path = os.path.abspath(os.path.expanduser(artifact_path.strip()))
+    cache_key = (artifact_path, str(torch.device(device)), int(vocab_size))
+    cache_entry = cache_owner.__dict__.get('_semantic_tree_token_path_cache')
+    if isinstance(cache_entry, tuple) and len(cache_entry) == 2 and (cache_entry[0] == cache_key):
+        return cache_entry[1]
+    tree = BalancedSemanticTree.from_artifact(production_semantic_tree_contract(artifact_path), expected_vocab_size=int(vocab_size)).to(device)
+    cache_owner.__dict__['_semantic_tree_token_path_cache'] = (cache_key, tree)
+    return tree
+
+def _select_within_unstarted(scores, unstarted, count, *, multinomial=False):
+    if scores.ndim != 2 or unstarted.shape != scores.shape:
+        raise ValueError('scores and unstarted must have matching [batch, token] shapes')
+    if unstarted.dtype != torch.bool:
+        raise TypeError('unstarted must use torch.bool')
+    if scores.shape[0] != 1:
+        raise ValueError('delayed refinement selection currently requires batch size 1')
+    count = int(count.item()) if torch.is_tensor(count) else int(count)
+    eligible = unstarted[0].nonzero(as_tuple=False).flatten()
+    if count < 0 or count > eligible.numel():
+        raise ValueError('selection count exceeds the unstarted-token set')
+    if count == 0:
+        return eligible[:0]
+    if count == eligible.numel():
+        return eligible
+    eligible_scores = scores[0, eligible]
+    if multinomial:
+        local_indices = torch.multinomial(eligible_scores, num_samples=count)
+    else:
+        local_indices = torch.topk(eligible_scores, k=count).indices
+    return eligible[local_indices]
+
+def _combine_edit_guidance_logits(logits_full, logits_unconditional, logits_vq_dropped, *, guidance_scale, guidance_scale_image, edit_mode, guidance_mode='legacy', logits_image_only=None):
+    if guidance_mode == 'legacy':
+        if edit_mode in [0, 3]:
+            logits_cond = (1 + guidance_scale_image) * logits_full - guidance_scale_image * logits_vq_dropped
+        elif edit_mode in [1, 2]:
+            if guidance_scale_image <= -1:
+                raise ValueError('legacy edit modes 1/2 require guidance_scale_image > -1')
+            logits_cond = (logits_full + guidance_scale_image * logits_vq_dropped) / (1 + guidance_scale_image)
+        else:
+            raise ValueError(f'Unsupported edit_mode: {edit_mode}')
+        return (1 + guidance_scale) * logits_cond - guidance_scale * logits_unconditional
+    if guidance_mode == 'independent':
+        return logits_full + guidance_scale * (logits_full - logits_unconditional) + guidance_scale_image * (logits_full - logits_vq_dropped)
+    if guidance_mode == 'ip2p':
+        if logits_image_only is None:
+            raise ValueError('ip2p guidance requires image-only logits')
+        return logits_unconditional + (1 + guidance_scale_image) * (logits_image_only - logits_unconditional) + (1 + guidance_scale) * (logits_full - logits_image_only)
+    raise ValueError("edit_guidance_mode must be 'legacy', 'independent', or 'ip2p'")
+
+def _resolve_guidance_scale_schedule(guidance_scale_schedule, n_steps):
+    if guidance_scale_schedule is None:
+        return None
+    if not isinstance(guidance_scale_schedule, list):
+        raise TypeError('guidance_scale_schedule must be a list')
+    if isinstance(n_steps, bool) or not isinstance(n_steps, int) or n_steps < 1:
+        raise ValueError('guidance_scale_schedule requires a positive integer n_steps')
+    if len(guidance_scale_schedule) != n_steps:
+        raise ValueError(f'guidance_scale_schedule must contain exactly n_steps values: {len(guidance_scale_schedule)} != {n_steps}')
+    resolved = []
+    for value in guidance_scale_schedule:
+        if isinstance(value, bool) or not isinstance(value, (int, float, np.integer, np.floating)):
+            raise TypeError('guidance_scale_schedule values must be real numbers')
+        value = float(value)
+        if not math.isfinite(value):
+            raise ValueError('guidance_scale_schedule values must be finite')
+        resolved.append(value)
+    return tuple(resolved)
+
+def _resolve_semantic_tree_guidance_scale_range(guidance_scale_range):
+    if guidance_scale_range is None:
+        return None
+    if not isinstance(guidance_scale_range, (list, tuple)):
+        raise TypeError('semantic_tree_guidance_scale_range must be a two-value sequence')
+    if len(guidance_scale_range) != 2:
+        raise ValueError('semantic_tree_guidance_scale_range must contain exactly two values')
+    resolved = []
+    for value in guidance_scale_range:
+        if isinstance(value, bool) or not isinstance(value, (int, float, np.integer, np.floating)):
+            raise TypeError('semantic_tree_guidance_scale_range values must be real numbers')
+        value = float(value)
+        if not math.isfinite(value):
+            raise ValueError('semantic_tree_guidance_scale_range values must be finite')
+        resolved.append(value)
+    return tuple(resolved)
+
+def _resolve_semantic_tree_global_guidance_scale_range(guidance_scale_range):
+    if guidance_scale_range is None:
+        return None
+    if not isinstance(guidance_scale_range, (list, tuple)):
+        raise TypeError('semantic_tree_global_guidance_scale_range must be a two-value sequence')
+    if len(guidance_scale_range) != 2:
+        raise ValueError('semantic_tree_global_guidance_scale_range must contain exactly two values')
+    resolved = []
+    for value in guidance_scale_range:
+        if isinstance(value, bool) or not isinstance(value, (int, float, np.integer, np.floating)):
+            raise TypeError('semantic_tree_global_guidance_scale_range values must be real numbers')
+        value = float(value)
+        if not math.isfinite(value):
+            raise ValueError('semantic_tree_global_guidance_scale_range values must be finite')
+        resolved.append(value)
+    return tuple(resolved)
+
+def _resolve_semantic_tree_cfg_noise_interval(noise_interval):
+    if noise_interval is None:
+        return None
+    if not isinstance(noise_interval, (list, tuple)):
+        raise TypeError('semantic_tree_cfg_noise_interval must be a two-value sequence')
+    if len(noise_interval) != 2:
+        raise ValueError('semantic_tree_cfg_noise_interval must contain exactly two values')
+    resolved = []
+    for value in noise_interval:
+        if isinstance(value, bool) or not isinstance(value, (int, float, np.integer, np.floating)):
+            raise TypeError('semantic_tree_cfg_noise_interval values must be real numbers')
+        value = float(value)
+        if not math.isfinite(value):
+            raise ValueError('semantic_tree_cfg_noise_interval values must be finite')
+        resolved.append(value)
+    (minimum, maximum) = resolved
+    if not 0.0 <= minimum <= maximum <= 1.0:
+        raise ValueError('semantic_tree_cfg_noise_interval must satisfy 0 <= min_q <= max_q <= 1')
+    return (minimum, maximum)
+
+@torch.no_grad()
+@_restore_attention_state_after_prediction
+def text_to_image(model, prompt, confidence_policy='mmada', guidance_scale=5, n_steps=20, tokenizer=None, image_resolution=512, n_tokens=1024, shift=3, alg_temp=1, schedule='shift', min_temperature=0.01, dynamic_temperature=False, micro_cond='ORIGINAL WIDTH : 1024; ORIGINAL HEIGHT : 1024; TOP : 0; LEFT : 0; SCORE : 6.5', temperature=1, schedule_temp='linear', shift_alg=None, top_p=None, top_k=None, unmask_order=None, cfg_interval=[0, 1], edit_image=None, image_processor=None, lock_source_padding_vq=None, edit_source_native_size=None, guidance_scale_image=5, edit_guidance_mode='legacy', edit_mode=0, template='Generate an image with the caption:\n <prompt>', plan=None, conv_template='llada', use_cache=None, cache_prompt=None, causal_context=True, edit_threshold=-1, disable_tqdm=False, ministral_sbd_sampling=None, ministral_full_attention_approx=True, ministral_shift_sbd_logits=None, image_embedding_refinement_steps=None, loop_k=None):
+    _validate_text_to_image_batch_size(1)
+    lock_source_padding_vq = resolve_source_padding_vq_lock(model.get_model().config, lock_source_padding_vq, has_edit_image=edit_image is not None)
+    if edit_source_native_size is not None and (not lock_source_padding_vq):
+        raise ValueError('edit_source_native_size is valid only with lock_source_padding_vq=true')
+    semantic_tree_refinement_gap = resolve_semantic_tree_refinement_gap(0)
+    semantic_tree_atomic_block_size = resolve_semantic_tree_atomic_block_size(1)
+    if shift_alg is None:
+        shift_alg = shift
+    resolved_loop_k = None
+    guidance_scale_schedule = _resolve_guidance_scale_schedule(None, n_steps)
+    semantic_tree_guidance_scale_range = _resolve_semantic_tree_guidance_scale_range(None)
+    semantic_tree_global_guidance_scale_range = _resolve_semantic_tree_global_guidance_scale_range(None)
+    semantic_tree_cfg_noise_interval = _resolve_semantic_tree_cfg_noise_interval(None)
+    guidance_scale_modes = sum((value is not None for value in (None, None, None)))
+    if guidance_scale_modes > 1:
+        raise ValueError('guidance_scale_schedule, semantic_tree_guidance_scale_range, and semantic_tree_global_guidance_scale_range are mutually exclusive')
+    device = model.get_model().device
+    use_image_local_ar_diffusion = False
+    if isinstance(edit_image, (str, os.PathLike)):
+        edit_image = _load_prediction_image(edit_image)
+    source_native_size = None
+    if lock_source_padding_vq:
+        source_native_size = _resolve_source_native_size(edit_source_native_size, fallback=edit_image.size)
+    conv_template = conv_template
+    use_mdm_prime = False
+    use_semantic_tree_tokens = True
+    semantic_tree_noise_cfg_requested = any((value is not None for value in (None, None, None)))
+    if semantic_tree_noise_cfg_requested:
+        if edit_image is not None:
+            raise NotImplementedError('semantic-tree noise-dependent CFG currently supports T2I only')
+    image_embedding_refinement_steps = resolve_semantic_tree_steps(model.get_model().config, image_embedding_refinement_steps)
+    use_refinement_state = True
+    reserve_token = '<|reserved_token_5|>'
+    reserve_token_2 = '<|reserved_token_6|>'
+    reserve_id = 126089
+    reserve_id2 = 126090
+    INT_MAX = 1000000
+    is_emu3 = 'emu3' in 'emu3'.lower()
+    use_image_row_markers = image_row_markers_enabled(model.get_model().config)
+    fixed_prefix_count = image_fixed_prefix_token_count(model.get_model().config)
+    if fixed_prefix_count:
+        validate_image_fixed_prefix_support(model.get_model())
+    row_marker_id = None
+    image_end_id = None
+    markers_per_row = None
+    skip_down_sample = resolve_generation_skip_down_sample(model.get_model().config, runtime_value=False)
+    (row_marker_id, _, image_end_id) = validate_image_row_marker_support(model.get_model(), skip_down_sample=skip_down_sample)
+    markers_per_row = image_row_marker_count(model.get_model(), skip_down_sample=skip_down_sample)
+    image_discrete_path = 'tree_metric'
+    use_tree_metric_path = False
+    use_affine_mixture_path = False
+    use_endpoint_probability_path = False
+    config = model.get_model().config
+    semantic_block_size = getattr(config, 'block_size', None)
+    if isinstance(semantic_block_size, bool) or not isinstance(semantic_block_size, int) or semantic_block_size < 1:
+        raise NotImplementedError('semantic_tree_token_diffusion requires a positive integer block_size')
+    if ministral_sbd_sampling is False:
+        raise NotImplementedError('semantic_tree_token_diffusion cannot disable checkpoint-native Ministral SBD sampling')
+    semantic_edit_controls_requested = edit_threshold > 0
+    if semantic_edit_controls_requested:
+        raise NotImplementedError('semantic_tree_token_diffusion does not support post-commit edit/refinement controls')
+    codebook_size = int(model.get_model().vqvae.config.num_vq_embeddings)
+    (_, img_mask_id, _) = image_generation_special_ids(codebook_size)
+    txt_mask_id = 126336
+    eot_id = 126384
+    model_type = 'llada'
+    model_type = 'ministral'
+    is_ministral = True
+    use_ministral_sbd_sampling = False
+    _wte = wte_ministral
+    img_begin = NC.gen_im_start_token
+    img_end = NC.gen_im_end_token
+    if use_cache is None:
+        use_cache = True
+    if cache_prompt is None:
+        cache_prompt = True
+    conv_template = 'ministral'
+    reserve_id = NC.reserve_id
+    reserve_id2 = NC.reserve_id_enc
+    reserve_token = NC.reserve_id_token
+    reserve_token_2 = NC.reserve_id_token_enc
+    img_begin = NC.gen_im_start_token
+    img_end = NC.gen_im_end_token
+    txt_mask_id = NC.mask_id
+    eot_id = NC.eos_id
+    _wte = wte_ministral
+    model_module = model.module if hasattr(model, 'module') else model
+    use_ministral_sbd_sampling = _resolve_ministral_sbd_diffusion_sampling(model_module, ministral_sbd_sampling)
+    resolve_ministral_shift_sbd_logits(model_module, ministral_shift_sbd_logits)
+    _configure_ministral_t2i_attention(model_module, sbd_diffusion_sampling=use_ministral_sbd_sampling)
+    gen_shape_map = {1024: (64, 64), 512: (32, 32), 256: (16, 16)}
+    content_gen_shape = gen_shape_map[image_resolution]
+    if int(n_tokens) != math.prod(content_gen_shape):
+        raise ValueError(f'Raw generation token count must match the content image grid: {n_tokens} != {math.prod(content_gen_shape)}')
+    gen_shape = augmented_image_shape(content_gen_shape, markers_per_row=markers_per_row)
+    micro_cond = micro_cond
+    if edit_image is not None:
+        micro_cond = ''
+    prompt = f'{prompt} {micro_cond}'
+    question = template.replace('<prompt>', prompt)
+    image_tokens_txt = _transformer_image_token_count(model, math.prod(gen_shape), gen_shape)
+    n_tokens_txt = image_tokens_txt + fixed_prefix_count
+    if plan is None:
+        plan = ''
+    edit_source_vq_codes = None
+    if edit_image is not None:
+        image = edit_image
+        image_tensor = process_images([image], image_processor, model.config)
+        image_tensor = _move_image_to_device(image_tensor, device=model.device, dtype=torch.bfloat16)
+        modalities = ['image']
+        image_sizes = [image.size]
+        conv = copy.deepcopy(conv_templates[conv_template])
+        image = image.convert('RGB')
+        image_1024 = pad_to_square_and_resize(image, image_resolution)
+        (enc_latents, _gen_shape) = model.encode_image_gen(model.model.image_processor_gen.preprocess(image_1024).to(model.device, model.dtype), enc=True)
+        if lock_source_padding_vq:
+            if tuple((int(value) for value in _gen_shape)) != tuple(content_gen_shape):
+                raise RuntimeError(f'edit source VQ grid does not match generation grid: {tuple(_gen_shape)} != {tuple(content_gen_shape)}')
+            if enc_latents.ndim != 2 or tuple(enc_latents.shape) != (1, n_tokens):
+                raise RuntimeError(f'edit source VQ codes must have shape [1, {n_tokens}], got {tuple(enc_latents.shape)}')
+            edit_source_vq_codes = enc_latents.detach().clone()
+        enc_embeddings = model.model.call_gen_embedding(enc_latents, _gen_shape, enc=True)
+        conv.append_message(conv.roles[0], f'<image> {reserve_token_2 * enc_embeddings.shape[1]}\n {prompt} ')
+        conv.append_message(conv.roles[1], f'Sure {plan}{img_begin}{reserve_token * n_tokens_txt}{img_end}')
+        prompt_question = conv.get_prompt()
+        prompt_question.removesuffix('<|start_header_id|>assistant<|end_header_id|>\n\n')
+        (prompt_question, input_ids) = _tokenize_prediction_prompt_with_images(prompt_question, tokenizer, model, image_processor, image_tensor)
+        input_ids = input_ids.unsqueeze(0).to(model.device)
+        attention_mask = None
+        position_ids = None
+        image_sizes = [image.size]
+        modalities = ['image']
+        (inputs, position_ids, attention_mask, _, inputs_embeds, _, raw_input_ids) = model.prepare_inputs_labels_for_multimodal(input_ids, position_ids, attention_mask, None, None, image_tensor, modalities, image_sizes=image_sizes, return_inputs=True)
+        expected_encoder_slots = int(enc_embeddings.shape[0] * enc_embeddings.shape[1])
+        actual_encoder_slots = int(raw_input_ids.eq(reserve_id2).sum().item())
+        if actual_encoder_slots != expected_encoder_slots:
+            raise RuntimeError(f'T2I edit encoder prompt/token mismatch: expected {expected_encoder_slots} slots but tokenization produced {actual_encoder_slots}.')
+        inputs_embeds[raw_input_ids == reserve_id2] = 0
+        _enc = pad_along_last_dim(enc_embeddings, inputs_embeds.shape[-1])
+        inputs_embeds[raw_input_ids == reserve_id2] = _enc.flatten(0, 1)
+        is_eot = torch.where(raw_input_ids == eot_id)[1]
+        assert len(is_eot) == 3
+        prompt_cutoff = is_eot[1]
+        is_prompt = torch.zeros_like(raw_input_ids, dtype=torch.bool)
+        is_prompt[:, :prompt_cutoff + 1] = True
+        is_gen = raw_input_ids == reserve_id
+        is_gen_enc = raw_input_ids == reserve_id2
+    else:
+        conv = copy.deepcopy(conv_templates[conv_template])
+        conv.append_message(conv.roles[0], question)
+        conv.append_message(conv.roles[1], f'Sure {plan}{img_begin}{reserve_token * n_tokens_txt}{img_end}')
+        prompt_question = conv.get_prompt()
+        prompt_question.removesuffix('<|start_header_id|>assistant<|end_header_id|>\n\n')
+        input_ids = tokenizer_image_token(prompt_question, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(device)
+        is_gen = input_ids == reserve_id
+        is_gen_enc = input_ids == reserve_id2
+        is_eot = torch.where(input_ids == eot_id)[1]
+        assert len(is_eot) == 3
+        prompt_cutoff = is_eot[1]
+        is_prompt = torch.zeros_like(input_ids, dtype=torch.bool)
+        is_prompt[:, :prompt_cutoff + 1] = True
+        raw_input_ids = input_ids
+        (inputs_embeds, _) = _wte(model.model, raw_input_ids)
+    generation_placeholder_mask = is_gen
+    generated_prompt_slots = int(generation_placeholder_mask.sum().item())
+    if generated_prompt_slots != n_tokens_txt:
+        raise RuntimeError(f'T2I prompt/token mismatch: expected {n_tokens_txt} generated-image slots but tokenization produced {generated_prompt_slots}.')
+    (fixed_prefix_mask, is_gen, fixed_prefix_ids) = split_generation_fixed_prefix(generation_placeholder_mask, codebook_size=int(model.get_model().vqvae.config.num_vq_embeddings), count=fixed_prefix_count)
+    if int(is_gen.sum().item()) != image_tokens_txt:
+        raise RuntimeError(f'T2I image-token mismatch after removing the fixed prefix: expected {image_tokens_txt}, found {int(is_gen.sum().item())}.')
+    generation_attention_mask = fixed_prefix_mask | is_gen
+    if fixed_prefix_ids.numel():
+        fixed_prefix_embeddings = model.get_model().call_gen_fixed_token_embedding(fixed_prefix_ids)
+        inputs_embeds[fixed_prefix_mask] = pad_along_last_dim(fixed_prefix_embeddings, inputs_embeds.shape[-1]).reshape(-1, inputs_embeds.shape[-1])
+    if plan:
+        is_gen_min = torch.where(raw_input_ids == reserve_id)[1][0]
+        is_prompt[:, :is_gen_min] = True
+    ministral_sbd_attention_window = None
+    use_exact_ministral_flash = bool(ministral_full_attention_approx)
+    use_ministral_prefix_cache = bool(use_cache and cache_prompt)
+    if use_ministral_sbd_sampling:
+        ministral_sbd_attention_window = generated_image_attention_window(generation_attention_mask)
+        use_ministral_prefix_cache = bool(use_cache and cache_prompt and causal_context)
+        if use_exact_ministral_flash and (not use_ministral_prefix_cache):
+            use_exact_ministral_flash = False
+    inputs_embeds_uncond = inputs_embeds.clone()
+    noise_embed = model.model.embed_tokens(torch.tensor([txt_mask_id]).to(model.device))
+    inputs_embeds_uncond[is_prompt] = noise_embed
+    if edit_image is not None:
+        inputs_embeds_image_only = inputs_embeds_uncond.clone()
+        inputs_embeds_image_only[is_gen_enc] = inputs_embeds[is_gen_enc]
+        inputs_embeds_uncond_enc = inputs_embeds.clone()
+        if edit_mode == 0:
+            inputs_embeds_uncond_enc[~is_gen_enc] = noise_embed
+            is_gen_enc_ccc = is_gen_enc
+        elif edit_mode == 1:
+            inputs_embeds_uncond_enc[is_gen_enc] = noise_embed
+            is_gen_enc_ccc = torch.zeros_like(is_gen_enc, dtype=torch.bool)
+        elif edit_mode == 2:
+            inputs_embeds_uncond_enc[is_gen_enc | (raw_input_ids < 0)] = noise_embed
+            is_gen_enc_ccc = torch.zeros_like(is_gen_enc, dtype=torch.bool)
+        elif edit_mode == 3:
+            inputs_embeds_uncond_enc[~is_gen_enc & (raw_input_ids > 0)] = noise_embed
+            is_gen_enc_ccc = is_gen_enc
+        else:
+            raise ValueError('Not Supported edit_mode')
+        inputs_embeds_uncond_enc[fixed_prefix_mask] = inputs_embeds[fixed_prefix_mask]
+    image_gen_latents_offset = torch.zeros(1, n_tokens, dtype=torch.long, device=input_ids.device)
+    is_unitok = False
+    structural_mask = None
+    endpoint_probability_path = None
+    semantic_tree = None
+    semantic_representatives = None
+    image_vocab_size = int(model.get_model().vqvae.config.num_vq_embeddings)
+    semantic_tree = _load_semantic_tree_path(model.get_model().config, cache_owner=model.get_model(), device=input_ids.device, vocab_size=image_vocab_size)
+    if image_embedding_refinement_steps > semantic_tree.depth:
+        raise ValueError(f'semantic-tree delayed commit steps cannot exceed tree depth: {image_embedding_refinement_steps} > {semantic_tree.depth}')
+    semantic_representatives = semantic_tree.sample_representatives(torch.zeros_like(image_gen_latents_offset), levels=0)
+    image_gen_latents_offset[:] = img_mask_id
+    xt = image_gen_latents_offset.clone()
+    (xt, marker_gen_shape, structural_mask) = append_image_row_markers(xt, content_gen_shape, row_marker_id=row_marker_id, image_end_id=image_end_id, markers_per_row=markers_per_row)
+    if marker_gen_shape != gen_shape:
+        raise RuntimeError('inference marker grid does not match the prompt grid')
+    (semantic_representatives, semantic_marker_shape, semantic_structural_mask) = append_image_row_markers(semantic_representatives, content_gen_shape, row_marker_id=row_marker_id, image_end_id=image_end_id, markers_per_row=markers_per_row)
+    if semantic_marker_shape != gen_shape or not torch.equal(semantic_structural_mask, structural_mask):
+        raise RuntimeError('semantic-tree marker layout does not match generation')
+    if structural_mask is None:
+        content_positions = torch.arange(xt.shape[-1], device=xt.device, dtype=torch.long)
+    else:
+        content_positions = (~structural_mask[0]).nonzero(as_tuple=False).flatten()
+    source_padding_raw_mask = None
+    source_padding_lock_mask = None
+    source_padding_locked_codes = None
+    source_padding_locked_embeddings = None
+    if lock_source_padding_vq:
+        if edit_source_vq_codes is None or source_native_size is None:
+            raise RuntimeError('source-padding VQ locking lost its source codes or geometry')
+        source_padding_raw_mask = _source_padding_vq_mask(source_native_size, content_gen_shape, device=xt.device).expand(edit_source_vq_codes.shape[0], -1)
+        (_, source_padding_lock_mask, source_padding_locked_codes) = _prefill_source_padding_vq(xt, edit_source_vq_codes.to(device=xt.device, dtype=xt.dtype), source_padding_raw_mask, content_positions)
+        source_padding_locked_embeddings = model.get_model().call_gen_token_embedding(source_padding_locked_codes)
+    sampling_protected_mask = structural_mask
+    refinement_child_embeddings = None
+    refinement_stage = None
+    semantic_next_refinement_round = None
+    refinement_stage = torch.full_like(xt, -1)
+    semantic_next_refinement_round = torch.zeros_like(refinement_stage)
+    if structural_mask is not None:
+        refinement_stage[structural_mask] = image_embedding_refinement_steps
+    semantic_noise_levels = semantic_tree_noise_for_ages(refinement_stage.clamp_min(0), steps=image_embedding_refinement_steps)
+    refinement_child_embeddings = build_semantic_tree_embeddings(semantic_representatives, semantic_noise_levels, model.get_model().call_gen_token_embedding, img_mask_id)
+    prime_digits = None
+    endpoint_sampling_times = None
+    mask_idx = xt == img_mask_id
+    n_mask = mask_idx.sum()
+    mask_idx_sch = mask_idx
+    semantic_counts = [semantic_tree_activation_counts(int(row_mask.sum().item()), activation_steps=n_steps, schedule=schedule, shift=shift, device=mask_idx_sch.device) for row_mask in mask_idx_sch]
+    if any((row_counts.shape != semantic_counts[0].shape for row_counts in semantic_counts[1:])):
+        raise ValueError('semantic-tree inference requires equal activation-step counts across batch rows')
+    num_transfer_tokens = torch.stack(semantic_counts)
+    activation_steps = num_transfer_tokens.shape[1]
+    decode_num_transfer_tokens = semantic_tree_decode_schedule(num_transfer_tokens, refinement_steps=image_embedding_refinement_steps, refinement_gap=0)
+    sch_temperatures = torch.linspace(0, 1, n_steps, device='cpu').numpy()
+    if schedule_temp == 'linear':
+        sch_temperatures = (1 - sch_temperatures) * (1 - min_temperature) + min_temperature
+    elif schedule_temp == 'cosine2':
+        sch_temperatures = cosine_schedule_2(1 - sch_temperatures) * (1 - min_temperature) + min_temperature
+    elif schedule_temp == 'shift':
+        sch_temperatures = logit_normal_schedule(shift_alg, 1 - sch_temperatures) * (1 - min_temperature) + min_temperature
+    else:
+        raise NotImplementedError(f'Unknown schedule_temp: {schedule_temp}')
+    sch_temperatures_samp = torch.linspace(0, 1, n_steps, device='cpu').numpy()
+    sch_temperatures_samp = (1 - sch_temperatures_samp) * (1 - 1) + 1
+    sch_temperatures = torch.tensor(sch_temperatures, device=device, dtype=torch.float32)
+    sch_temperatures_samp = torch.tensor(sch_temperatures_samp, device=device, dtype=torch.float32)
+    temp_idx = 0
+    cfg_start = int(cfg_interval[0] * n_steps)
+    cfg_end = int(cfg_interval[1] * n_steps)
+    use_optional_guidance_scale = any((value is not None for value in (None, None, None, None)))
+    scheduled_run_cfg_branches = None
+    if isinstance(unmask_order, str):
+        unmask_order = torch.load(unmask_order).tolist()
+    eligible_parent_positions = None
+    if confidence_policy == 'stratified' and unmask_order is None:
+        stratified_tokens = n_tokens
+        _dim = int(np.sqrt(stratified_tokens))
+        if _dim * _dim != stratified_tokens:
+            raise ValueError('stratified local AR sampling requires a square parent grid')
+        seq = stratified_random(n=_dim, seed=42, shuffle_blocks=True)
+        unmask_order = seq
+    stratified_token_count = n_tokens
+    if confidence_policy == 'stratified':
+        assert unmask_order is not None
+    n_c = 0
+    cat_cfg_time = 0
+    post_cfg_time = 0
+    wte_time = 0
+    fwd_time = 0
+    sample_time = 0
+    all_time = 0
+    total_edited = 0
+    edit_transitions = []
+    last_edit_step = torch.full_like(xt, -10000)
+    distribution_trace = {'initial_commits': [], 'edit_proposals': [], 'edits': []}
+    initial_trace_top_ids = None
+    initial_trace_top_logits = None
+    initial_trace_top_probs = None
+    initial_trace_top_sampling_probs = None
+    previous_trace_top_ids = None
+    previous_trace_top_logits = None
+    previous_trace_top_probs = None
+    intermediate_x0s = []
+    for (decode_step_idx, num_transfer) in tqdm(enumerate(decode_num_transfer_tokens[0]), total=decode_num_transfer_tokens.shape[1], disable=disable_tqdm):
+        t0 = time.time()
+        schedule_temp_idx = temp_idx
+        schedule_temp_idx = min(schedule_temp_idx, min(activation_steps, n_steps) - 1)
+        local_temp = sch_temperatures[schedule_temp_idx]
+        local_temp_samp = sch_temperatures_samp[schedule_temp_idx]
+        temp_idx += 1
+        schedule_progress_idx = temp_idx
+        schedule_progress_idx = min(schedule_progress_idx, activation_steps)
+        if schedule_progress_idx / n_steps > 100:
+            confidence_policy = 'mmada'
+        hard_mask_idx = refinement_stage < image_embedding_refinement_steps
+        mask_idx = refinement_stage < 0
+        n_mask = mask_idx.sum()
+        timestep_refinement_stage = refinement_stage[:, content_positions]
+        timesteps = refinement_mask_fraction(timestep_refinement_stage, steps=image_embedding_refinement_steps)
+        semantic_tree_cfg_q = None
+        cfg_step_idx = temp_idx
+        cfg_step_idx = min(cfg_step_idx, activation_steps)
+        guidance_scale_for_step = guidance_scale
+        if scheduled_run_cfg_branches is None:
+            run_cfg_branches = guidance_scale != 0 or guidance_scale_image != 0
+        else:
+            run_cfg_branches = scheduled_run_cfg_branches
+        do_cfg = run_cfg_branches and cfg_start <= cfg_step_idx and (cfg_end >= cfg_step_idx)
+        model_xt = _restore_source_padding_vq(xt, source_padding_lock_mask, source_padding_locked_codes)
+        if run_cfg_branches:
+            if edit_image is not None:
+                if edit_guidance_mode == 'ip2p':
+                    input_embeddings_input = torch.cat([inputs_embeds_uncond, inputs_embeds_image_only, inputs_embeds]).clone()
+                else:
+                    input_embeddings_input = torch.cat([inputs_embeds_uncond, inputs_embeds_uncond_enc, inputs_embeds]).clone()
+                xt_input = torch.cat([model_xt, model_xt, model_xt])
+                new_token_mask = is_gen.repeat(3, 1)
+                fixed_prefix_routing = fixed_prefix_mask.repeat(3, 1)
+                is_gen_enc_null = torch.zeros_like(is_gen_enc, dtype=torch.bool)
+                if edit_guidance_mode == 'ip2p':
+                    is_gen_enc_mask = torch.cat([is_gen_enc_null, is_gen_enc, is_gen_enc])
+                else:
+                    is_gen_enc_mask = torch.cat([is_gen_enc_null, is_gen_enc_ccc, is_gen_enc])
+                timesteps = timesteps.repeat(3)
+            else:
+                input_embeddings_input = torch.cat([inputs_embeds_uncond, inputs_embeds]).clone()
+                xt_input = torch.cat([model_xt, model_xt])
+                new_token_mask = is_gen.repeat(2, 1)
+                fixed_prefix_routing = fixed_prefix_mask.repeat(2, 1)
+                is_gen_enc_mask = is_gen_enc.repeat(2, 1)
+                is_gen_enc_mask[0, :] = False
+                timesteps = timesteps.repeat(2)
+        else:
+            input_embeddings_input = inputs_embeds.clone()
+            new_token_mask = is_gen
+            fixed_prefix_routing = fixed_prefix_mask
+            xt_input = model_xt
+            is_gen_enc_mask = is_gen_enc
+        enc_use_image_branch = True
+        generation_attention_mask_input = fixed_prefix_routing | new_token_mask
+        if enc_use_image_branch:
+            modality_indices = generation_attention_mask_input | is_gen_enc_mask
+        else:
+            modality_indices = generation_attention_mask_input
+        t1 = time.time()
+        x_gen_input = xt_input
+        x_gen_child_embeddings = None
+        if xt_input.shape[0] % refinement_child_embeddings.shape[0] != 0:
+            raise RuntimeError('CFG batch does not align with the refinement state batch')
+        cfg_repeats = xt_input.shape[0] // refinement_child_embeddings.shape[0]
+        step_refinement_child_embeddings = refinement_child_embeddings
+        if source_padding_lock_mask is not None:
+            if source_padding_locked_embeddings.shape != refinement_child_embeddings.shape:
+                raise RuntimeError('source-padding embeddings do not align with the refinement state')
+            step_refinement_child_embeddings = refinement_child_embeddings.clone()
+            step_refinement_child_embeddings[source_padding_lock_mask] = source_padding_locked_embeddings[source_padding_lock_mask]
+        x_gen_input = None
+        x_gen_child_embeddings = step_refinement_child_embeddings.repeat(cfg_repeats, 1, 1)
+        (all_input_embeddings, new_token_mask) = _wte(model.get_model(), None, True, x_gen=x_gen_input, gen_shape=gen_shape, inputs_embeds_curr=input_embeddings_input, new_token_mask=new_token_mask, x_gen_child_embeddings=x_gen_child_embeddings)
+        t2 = time.time()
+        if use_ministral_prefix_cache:
+            if decode_step_idx == 0:
+                cache_cutoff = ministral_sbd_attention_window[0] if ministral_sbd_attention_window is not None else prompt_cutoff
+                if causal_context:
+                    for layer in model_module.encoder.layers:
+                        if hasattr(layer.self_attn, 'diffusion_lm'):
+                            layer.self_attn.diffusion_lm = False
+                if cache_cutoff > 0:
+                    output = model.get_model()(None, input_embeddings=all_input_embeddings[:, :cache_cutoff], modality_indices=modality_indices[:, :cache_cutoff], output_hidden_states=True, past_key_values=None, use_cache=True, overwrite_attn_impl='flash_attn')
+                    past_key_values = output.past_key_values
+                    cache_len = past_key_values.get_seq_length()
+                    if cache_len != cache_cutoff:
+                        raise RuntimeError(f'SBD prefix cache length does not match its generated-image boundary: {cache_len} != {cache_cutoff}')
+                else:
+                    past_key_values = None
+                    cache_len = 0
+                if causal_context:
+                    for layer in model_module.encoder.layers:
+                        if hasattr(layer.self_attn, 'diffusion_lm'):
+                            layer.self_attn.diffusion_lm = True
+            else:
+                pass
+        else:
+            past_key_values = None
+            cache_len = 0
+        sampling_end = ministral_sbd_attention_window[1] if ministral_sbd_attention_window is not None and use_ministral_prefix_cache else all_input_embeddings.shape[1]
+        if not cache_len <= sampling_end <= all_input_embeddings.shape[1]:
+            raise RuntimeError(f'Invalid exact-image SBD sampling window: cache_len={cache_len}, sampling_end={sampling_end}, sequence_length={all_input_embeddings.shape[1]}')
+        logits = get_logits_ministral(model.get_model(), all_input_embeddings[:, cache_len:sampling_end], new_token_mask[:, cache_len:sampling_end], True, gen_shape=gen_shape, input_modality_indices=modality_indices[:, cache_len:sampling_end], timesteps=timesteps, past_key_values=past_key_values, generation_attention_mask=generation_attention_mask_input[:, cache_len:sampling_end], sbd_diffusion_sampling=use_ministral_sbd_sampling, full_attention_approx=use_exact_ministral_flash, loop_k=None)
+        t3 = time.time()
+        if run_cfg_branches:
+            if edit_image is not None:
+                (new_token_mask, _, _) = new_token_mask.chunk(3)
+                (logits_un, logits_middle, logits) = logits.chunk(3)
+                if do_cfg:
+                    logits_is_ninf = logits == -np.inf
+                    logits = _combine_edit_guidance_logits(logits, logits_un, logits_middle if edit_guidance_mode != 'ip2p' else None, guidance_scale=guidance_scale_for_step, guidance_scale_image=guidance_scale_image, edit_mode=edit_mode, guidance_mode=edit_guidance_mode, logits_image_only=logits_middle if edit_guidance_mode == 'ip2p' else None)
+                    logits[logits_is_ninf] = -np.inf
+            else:
+                (new_token_mask, _) = new_token_mask.chunk(2)
+                (logits_un, logits) = logits.chunk(2)
+                if do_cfg:
+                    logits_is_ninf = logits == -np.inf
+                    if use_optional_guidance_scale:
+                        active_guidance_scale = guidance_scale
+                        logits = (1 + active_guidance_scale) * logits - active_guidance_scale * logits_un
+                    else:
+                        logits = (1 + guidance_scale) * logits - guidance_scale * logits_un
+                    logits[logits_is_ninf] = -np.inf
+        t4 = time.time()
+        if top_p is not None or top_k is not None:
+            (_b, _l, _v) = logits.shape
+            logits = logits.view(_b * _l, _v)
+            logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p, filter_value=-np.inf, min_tokens_to_keep=1)
+            logits = logits.view(_b, _l, _v)
+        probs = logits.softmax(dim=-1)
+        trace_top_ids = None
+        trace_top_logits = None
+        trace_top_probs = None
+        _temperature = temperature
+        sampling_distribution = dists.Categorical(logits=logits / _temperature)
+        x0 = sampling_distribution.sample()
+        del sampling_distribution
+        x0_p = torch.gather(probs, -1, x0.long()[..., None]).squeeze(-1)
+        if edit_threshold <= 0:
+            x0 = torch.where(hard_mask_idx, x0, xt)
+        if decode_step_idx >= activation_steps:
+            select_index = torch.empty(0, dtype=torch.long, device=x0.device)
+        elif confidence_policy == 'mask_git':
+            _alg_temp = alg_temp
+            if dynamic_temperature:
+                _alg_temp = _alg_temp * local_temp
+            confidence = x0_p / _alg_temp
+            confidence = torch.where(mask_idx, confidence, -np.inf)
+            confidence = torch.softmax(confidence, dim=-1)
+            select_index = _select_within_unstarted(confidence, mask_idx, num_transfer, multinomial=True)
+        elif confidence_policy == 'mmada':
+            _alg_temp = alg_temp
+            if dynamic_temperature:
+                _alg_temp = _alg_temp * local_temp
+            confidence = torch.log(x0_p.clamp(1e-20)) + _alg_temp * gumbel_noise(x0_p, generator=None)
+            confidence = torch.where(mask_idx, confidence, -np.inf)
+            select_index = _select_within_unstarted(confidence, mask_idx, num_transfer)
+        elif confidence_policy == 'stratified':
+            start = stratified_token_count - int(n_mask.item())
+            select_index = torch.tensor(unmask_order[start:start + int(num_transfer)], device=x0.device, dtype=torch.long)
+            if structural_mask is not None:
+                select_index = content_positions[select_index]
+        else:
+            raise NotImplementedError(f'Unknown confidence_policy: {confidence_policy}')
+        transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device)
+        j = 0
+        transfer_index[j, select_index] = True
+        proposal_child_embeddings = None
+        activation_index = transfer_index & (refinement_stage < 0)
+        refinement_stage[activation_index] = 0
+        semantic_next_refinement_round[activation_index] = decode_step_idx + 0
+        pending_index = (refinement_stage >= 0) & (refinement_stage < image_embedding_refinement_steps)
+        proposal_child_embeddings = model.get_model().call_gen_token_embedding(x0)
+        pending_index &= semantic_tree_refinement_due(refinement_stage, semantic_next_refinement_round, decode_round=decode_step_idx, refinement_steps=image_embedding_refinement_steps)
+        commit_index = torch.zeros_like(pending_index)
+        if pending_index.any():
+            semantic_advance = advance_revisable_semantic_tree(semantic_representatives[pending_index], x0[pending_index], refinement_stage[pending_index], semantic_tree, steps=image_embedding_refinement_steps)
+            semantic_representatives[pending_index] = semantic_advance.representatives
+            refinement_stage[pending_index] = semantic_advance.ages
+            semantic_next_refinement_round[pending_index] = decode_step_idx + 0 + 1
+            commit_index[pending_index] = semantic_advance.committed
+        semantic_noise_levels = semantic_tree_noise_for_ages(refinement_stage.clamp_min(0), steps=image_embedding_refinement_steps)
+        refinement_child_embeddings = build_semantic_tree_embeddings(semantic_representatives, semantic_noise_levels, model.get_model().call_gen_token_embedding, img_mask_id)
+        xt[commit_index] = semantic_representatives[commit_index]
+        transfer_index = commit_index
+        xt_is_mask = xt == img_mask_id
+        edit_progress = (decode_step_idx + 1) / decode_num_transfer_tokens.shape[1]
+        edits_enabled_this_round = edit_threshold > 0 and 0.0 <= edit_progress <= 1.0 and (image_embedding_refinement_steps == 1 or decode_step_idx < activation_steps)
+        if edits_enabled_this_round:
+            editable_position = ~xt_is_mask & ~transfer_index
+            if structural_mask is not None:
+                editable_position &= ~structural_mask
+            editing_confidence = torch.where(editable_position, x0_p, -torch.inf)
+            high_conf_editing = editing_confidence > edit_threshold
+            margin_accepted = torch.ones_like(editable_position)
+            proposal_logit_margin = None
+            token_changed = (x0 != xt) & high_conf_editing & margin_accepted
+            num_edited = token_changed.sum()
+            if num_edited > 0:
+                xt[token_changed] = x0[token_changed]
+                last_edit_step[token_changed] = decode_step_idx
+                refinement_child_embeddings[token_changed] = proposal_child_embeddings[token_changed]
+                refinement_stage[token_changed] = image_embedding_refinement_steps
+                total_edited += num_edited.item()
+        t5 = time.time()
+        cat_cfg_time += t1 - t0
+        post_cfg_time += t2 - t1
+        wte_time += t3 - t2
+        fwd_time += t4 - t3
+        sample_time += t5 - t4
+        all_time += t5 - t0
+        n_c += 1
+    unresolved_refinement = refinement_stage < image_embedding_refinement_steps
+    if unresolved_refinement.any() or (xt == img_mask_id).any():
+        stage_counts = {stage: (refinement_stage == stage).sum().item() for stage in range(-1, image_embedding_refinement_steps + 1)}
+        raise RuntimeError(f'image embedding refinement ended with unresolved tokens: unfinished_stages={unresolved_refinement.sum().item()}, committed_mask_tokens={(xt == img_mask_id).sum().item()}, stage_counts={stage_counts}')
+    xt = _restore_source_padding_vq(xt, source_padding_lock_mask, source_padding_locked_codes)
+    _assert_source_padding_vq_locked(xt, source_padding_lock_mask, source_padding_locked_codes)
+    if structural_mask is not None:
+        expected_row_id = int(row_marker_id)
+        expected_end_id = int(image_end_id)
+        marker_grid = xt.reshape(xt.shape[0], *gen_shape)
+        if not bool((marker_grid[:, :-1, -markers_per_row:] == expected_row_id).all() and (marker_grid[:, -1:, -markers_per_row:] == expected_end_id).all()):
+            raise RuntimeError('structural generation image markers changed during sampling')
+        xt = strip_image_row_markers(xt, content_gen_shape, markers_per_row=markers_per_row)
+    x0_img = xt
+    result = Image.fromarray(model.decode_image_gen(x0_img, image_resolution, image_resolution)[0])
+    return result

sbd_attention.py

@@ -0,0 +1,39 @@
+
+from __future__ import annotations
+import torch
+
+def build_image_attention_doc_ids(image_token_mask: torch.Tensor, *, valid_tokens: torch.Tensor | None=None) -> torch.Tensor:
+    if image_token_mask.ndim != 2:
+        raise ValueError('image_token_mask must be a rank-2 [batch, length] tensor')
+    if valid_tokens is None:
+        valid_tokens = torch.ones_like(image_token_mask, dtype=torch.bool)
+    elif tuple(valid_tokens.shape) != tuple(image_token_mask.shape):
+        raise ValueError('valid_tokens must align with image_token_mask')
+    image_token_mask = image_token_mask.to(dtype=torch.bool)
+    valid_tokens = valid_tokens.to(dtype=torch.bool)
+    image_token_mask = image_token_mask & valid_tokens
+    previous_image = torch.cat((torch.zeros_like(image_token_mask[:, :1]), image_token_mask[:, :-1]), dim=1)
+    document_starts = image_token_mask & ~previous_image
+    return torch.cumsum(document_starts.long(), dim=1) * image_token_mask
+
+def generated_image_attention_window(generated_image_token_mask: torch.Tensor, *, valid_tokens: torch.Tensor | None=None) -> tuple[int, int]:
+    if generated_image_token_mask.ndim != 2:
+        raise ValueError('generated_image_token_mask must be a rank-2 [batch, length] tensor')
+    if valid_tokens is None:
+        valid_tokens = torch.ones_like(generated_image_token_mask, dtype=torch.bool)
+    elif tuple(valid_tokens.shape) != tuple(generated_image_token_mask.shape):
+        raise ValueError('valid_tokens must align with generated_image_token_mask')
+    generated_image_token_mask = generated_image_token_mask.to(dtype=torch.bool) & valid_tokens.to(dtype=torch.bool)
+    windows = []
+    for row in generated_image_token_mask:
+        active = torch.where(row)[0]
+        if active.numel() == 0:
+            raise ValueError('SBD sampling requires at least one image token per row')
+        start = int(active[0].item())
+        end = int(active[-1].item()) + 1
+        if not bool(row[start:end].all()):
+            raise ValueError('Dense SBD FlashAttention requires one contiguous generated-image block')
+        windows.append((start, end))
+    if any((window != windows[0] for window in windows[1:])):
+        raise ValueError('Dense SBD FlashAttention requires the same generated-image block for every CFG batch row')
+    return windows[0]

schedules.py

@@ -0,0 +1,22 @@
+# ADOBE CONFIDENTIAL
+# Copyright 2025 Adobe
+# All Rights Reserved.
+# NOTICE: All information contained herein is, and remains
+# the property of Adobe and its suppliers, if any. The intellectual
+# and technical concepts contained herein are proprietary to Adobe
+# and its suppliers and are protected by all applicable intellectual
+# property laws, including trade secret and copyright laws.
+# Dissemination of this information or reproduction of this material
+# is strictly forbidden unless prior written permission is obtained
+# from Adobe.
+
+from __future__ import annotations
+import numpy as np
+
+def cosine_schedule_2(x):
+    x = 1 - np.clip(x, 0, 1)
+    return np.cos(np.pi * x / 2)
+
+def logit_normal_schedule(shift, sigmas):
+    sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+    return sigmas

special_tokens_map.json

@@ -0,0 +1,1061 @@
+{
+  "additional_special_tokens": [
+    "<LOC_0>",
+    "<LOC_1>",
+    "<LOC_2>",
+    "<LOC_3>",
+    "<LOC_4>",
+    "<LOC_5>",
+    "<LOC_6>",
+    "<LOC_7>",
+    "<LOC_8>",
+    "<LOC_9>",
+    "<LOC_10>",
+    "<LOC_11>",
+    "<LOC_12>",
+    "<LOC_13>",
+    "<LOC_14>",
+    "<LOC_15>",
+    "<LOC_16>",
+    "<LOC_17>",
+    "<LOC_18>",
+    "<LOC_19>",
+    "<LOC_20>",
+    "<LOC_21>",
+    "<LOC_22>",
+    "<LOC_23>",
+    "<LOC_24>",
+    "<LOC_25>",
+    "<LOC_26>",
+    "<LOC_27>",
+    "<LOC_28>",
+    "<LOC_29>",
+    "<LOC_30>",
+    "<LOC_31>",
+    "<LOC_32>",
+    "<LOC_33>",
+    "<LOC_34>",
+    "<LOC_35>",
+    "<LOC_36>",
+    "<LOC_37>",
+    "<LOC_38>",
+    "<LOC_39>",
+    "<LOC_40>",
+    "<LOC_41>",
+    "<LOC_42>",
+    "<LOC_43>",
+    "<LOC_44>",
+    "<LOC_45>",
+    "<LOC_46>",
+    "<LOC_47>",
+    "<LOC_48>",
+    "<LOC_49>",
+    "<LOC_50>",
+    "<LOC_51>",
+    "<LOC_52>",
+    "<LOC_53>",
+    "<LOC_54>",
+    "<LOC_55>",
+    "<LOC_56>",
+    "<LOC_57>",
+    "<LOC_58>",
+    "<LOC_59>",
+    "<LOC_60>",
+    "<LOC_61>",
+    "<LOC_62>",
+    "<LOC_63>",
+    "<LOC_64>",
+    "<LOC_65>",
+    "<LOC_66>",
+    "<LOC_67>",
+    "<LOC_68>",
+    "<LOC_69>",
+    "<LOC_70>",
+    "<LOC_71>",
+    "<LOC_72>",
+    "<LOC_73>",
+    "<LOC_74>",
+    "<LOC_75>",
+    "<LOC_76>",
+    "<LOC_77>",
+    "<LOC_78>",
+    "<LOC_79>",
+    "<LOC_80>",
+    "<LOC_81>",
+    "<LOC_82>",
+    "<LOC_83>",
+    "<LOC_84>",
+    "<LOC_85>",
+    "<LOC_86>",
+    "<LOC_87>",
+    "<LOC_88>",
+    "<LOC_89>",
+    "<LOC_90>",
+    "<LOC_91>",
+    "<LOC_92>",
+    "<LOC_93>",
+    "<LOC_94>",
+    "<LOC_95>",
+    "<LOC_96>",
+    "<LOC_97>",
+    "<LOC_98>",
+    "<LOC_99>",
+    "<LOC_100>",
+    "<LOC_101>",
+    "<LOC_102>",
+    "<LOC_103>",
+    "<LOC_104>",
+    "<LOC_105>",
+    "<LOC_106>",
+    "<LOC_107>",
+    "<LOC_108>",
+    "<LOC_109>",
+    "<LOC_110>",
+    "<LOC_111>",
+    "<LOC_112>",
+    "<LOC_113>",
+    "<LOC_114>",
+    "<LOC_115>",
+    "<LOC_116>",
+    "<LOC_117>",
+    "<LOC_118>",
+    "<LOC_119>",
+    "<LOC_120>",
+    "<LOC_121>",
+    "<LOC_122>",
+    "<LOC_123>",
+    "<LOC_124>",
+    "<LOC_125>",
+    "<LOC_126>",
+    "<LOC_127>",
+    "<LOC_128>",
+    "<LOC_129>",
+    "<LOC_130>",
+    "<LOC_131>",
+    "<LOC_132>",
+    "<LOC_133>",
+    "<LOC_134>",
+    "<LOC_135>",
+    "<LOC_136>",
+    "<LOC_137>",
+    "<LOC_138>",
+    "<LOC_139>",
+    "<LOC_140>",
+    "<LOC_141>",
+    "<LOC_142>",
+    "<LOC_143>",
+    "<LOC_144>",
+    "<LOC_145>",
+    "<LOC_146>",
+    "<LOC_147>",
+    "<LOC_148>",
+    "<LOC_149>",
+    "<LOC_150>",
+    "<LOC_151>",
+    "<LOC_152>",
+    "<LOC_153>",
+    "<LOC_154>",
+    "<LOC_155>",
+    "<LOC_156>",
+    "<LOC_157>",
+    "<LOC_158>",
+    "<LOC_159>",
+    "<LOC_160>",
+    "<LOC_161>",
+    "<LOC_162>",
+    "<LOC_163>",
+    "<LOC_164>",
+    "<LOC_165>",
+    "<LOC_166>",
+    "<LOC_167>",
+    "<LOC_168>",
+    "<LOC_169>",
+    "<LOC_170>",
+    "<LOC_171>",
+    "<LOC_172>",
+    "<LOC_173>",
+    "<LOC_174>",
+    "<LOC_175>",
+    "<LOC_176>",
+    "<LOC_177>",
+    "<LOC_178>",
+    "<LOC_179>",
+    "<LOC_180>",
+    "<LOC_181>",
+    "<LOC_182>",
+    "<LOC_183>",
+    "<LOC_184>",
+    "<LOC_185>",
+    "<LOC_186>",
+    "<LOC_187>",
+    "<LOC_188>",
+    "<LOC_189>",
+    "<LOC_190>",
+    "<LOC_191>",
+    "<LOC_192>",
+    "<LOC_193>",
+    "<LOC_194>",
+    "<LOC_195>",
+    "<LOC_196>",
+    "<LOC_197>",
+    "<LOC_198>",
+    "<LOC_199>",
+    "<LOC_200>",
+    "<LOC_201>",
+    "<LOC_202>",
+    "<LOC_203>",
+    "<LOC_204>",
+    "<LOC_205>",
+    "<LOC_206>",
+    "<LOC_207>",
+    "<LOC_208>",
+    "<LOC_209>",
+    "<LOC_210>",
+    "<LOC_211>",
+    "<LOC_212>",
+    "<LOC_213>",
+    "<LOC_214>",
+    "<LOC_215>",
+    "<LOC_216>",
+    "<LOC_217>",
+    "<LOC_218>",
+    "<LOC_219>",
+    "<LOC_220>",
+    "<LOC_221>",
+    "<LOC_222>",
+    "<LOC_223>",
+    "<LOC_224>",
+    "<LOC_225>",
+    "<LOC_226>",
+    "<LOC_227>",
+    "<LOC_228>",
+    "<LOC_229>",
+    "<LOC_230>",
+    "<LOC_231>",
+    "<LOC_232>",
+    "<LOC_233>",
+    "<LOC_234>",
+    "<LOC_235>",
+    "<LOC_236>",
+    "<LOC_237>",
+    "<LOC_238>",
+    "<LOC_239>",
+    "<LOC_240>",
+    "<LOC_241>",
+    "<LOC_242>",
+    "<LOC_243>",
+    "<LOC_244>",
+    "<LOC_245>",
+    "<LOC_246>",
+    "<LOC_247>",
+    "<LOC_248>",
+    "<LOC_249>",
+    "<LOC_250>",
+    "<LOC_251>",
+    "<LOC_252>",
+    "<LOC_253>",
+    "<LOC_254>",
+    "<LOC_255>",
+    "<LOC_256>",
+    "<LOC_257>",
+    "<LOC_258>",
+    "<LOC_259>",
+    "<LOC_260>",
+    "<LOC_261>",
+    "<LOC_262>",
+    "<LOC_263>",
+    "<LOC_264>",
+    "<LOC_265>",
+    "<LOC_266>",
+    "<LOC_267>",
+    "<LOC_268>",
+    "<LOC_269>",
+    "<LOC_270>",
+    "<LOC_271>",
+    "<LOC_272>",
+    "<LOC_273>",
+    "<LOC_274>",
+    "<LOC_275>",
+    "<LOC_276>",
+    "<LOC_277>",
+    "<LOC_278>",
+    "<LOC_279>",
+    "<LOC_280>",
+    "<LOC_281>",
+    "<LOC_282>",
+    "<LOC_283>",
+    "<LOC_284>",
+    "<LOC_285>",
+    "<LOC_286>",
+    "<LOC_287>",
+    "<LOC_288>",
+    "<LOC_289>",
+    "<LOC_290>",
+    "<LOC_291>",
+    "<LOC_292>",
+    "<LOC_293>",
+    "<LOC_294>",
+    "<LOC_295>",
+    "<LOC_296>",
+    "<LOC_297>",
+    "<LOC_298>",
+    "<LOC_299>",
+    "<LOC_300>",
+    "<LOC_301>",
+    "<LOC_302>",
+    "<LOC_303>",
+    "<LOC_304>",
+    "<LOC_305>",
+    "<LOC_306>",
+    "<LOC_307>",
+    "<LOC_308>",
+    "<LOC_309>",
+    "<LOC_310>",
+    "<LOC_311>",
+    "<LOC_312>",
+    "<LOC_313>",
+    "<LOC_314>",
+    "<LOC_315>",
+    "<LOC_316>",
+    "<LOC_317>",
+    "<LOC_318>",
+    "<LOC_319>",
+    "<LOC_320>",
+    "<LOC_321>",
+    "<LOC_322>",
+    "<LOC_323>",
+    "<LOC_324>",
+    "<LOC_325>",
+    "<LOC_326>",
+    "<LOC_327>",
+    "<LOC_328>",
+    "<LOC_329>",
+    "<LOC_330>",
+    "<LOC_331>",
+    "<LOC_332>",
+    "<LOC_333>",
+    "<LOC_334>",
+    "<LOC_335>",
+    "<LOC_336>",
+    "<LOC_337>",
+    "<LOC_338>",
+    "<LOC_339>",
+    "<LOC_340>",
+    "<LOC_341>",
+    "<LOC_342>",
+    "<LOC_343>",
+    "<LOC_344>",
+    "<LOC_345>",
+    "<LOC_346>",
+    "<LOC_347>",
+    "<LOC_348>",
+    "<LOC_349>",
+    "<LOC_350>",
+    "<LOC_351>",
+    "<LOC_352>",
+    "<LOC_353>",
+    "<LOC_354>",
+    "<LOC_355>",
+    "<LOC_356>",
+    "<LOC_357>",
+    "<LOC_358>",
+    "<LOC_359>",
+    "<LOC_360>",
+    "<LOC_361>",
+    "<LOC_362>",
+    "<LOC_363>",
+    "<LOC_364>",
+    "<LOC_365>",
+    "<LOC_366>",
+    "<LOC_367>",
+    "<LOC_368>",
+    "<LOC_369>",
+    "<LOC_370>",
+    "<LOC_371>",
+    "<LOC_372>",
+    "<LOC_373>",
+    "<LOC_374>",
+    "<LOC_375>",
+    "<LOC_376>",
+    "<LOC_377>",
+    "<LOC_378>",
+    "<LOC_379>",
+    "<LOC_380>",
+    "<LOC_381>",
+    "<LOC_382>",
+    "<LOC_383>",
+    "<LOC_384>",
+    "<LOC_385>",
+    "<LOC_386>",
+    "<LOC_387>",
+    "<LOC_388>",
+    "<LOC_389>",
+    "<LOC_390>",
+    "<LOC_391>",
+    "<LOC_392>",
+    "<LOC_393>",
+    "<LOC_394>",
+    "<LOC_395>",
+    "<LOC_396>",
+    "<LOC_397>",
+    "<LOC_398>",
+    "<LOC_399>",
+    "<LOC_400>",
+    "<LOC_401>",
+    "<LOC_402>",
+    "<LOC_403>",
+    "<LOC_404>",
+    "<LOC_405>",
+    "<LOC_406>",
+    "<LOC_407>",
+    "<LOC_408>",
+    "<LOC_409>",
+    "<LOC_410>",
+    "<LOC_411>",
+    "<LOC_412>",
+    "<LOC_413>",
+    "<LOC_414>",
+    "<LOC_415>",
+    "<LOC_416>",
+    "<LOC_417>",
+    "<LOC_418>",
+    "<LOC_419>",
+    "<LOC_420>",
+    "<LOC_421>",
+    "<LOC_422>",
+    "<LOC_423>",
+    "<LOC_424>",
+    "<LOC_425>",
+    "<LOC_426>",
+    "<LOC_427>",
+    "<LOC_428>",
+    "<LOC_429>",
+    "<LOC_430>",
+    "<LOC_431>",
+    "<LOC_432>",
+    "<LOC_433>",
+    "<LOC_434>",
+    "<LOC_435>",
+    "<LOC_436>",
+    "<LOC_437>",
+    "<LOC_438>",
+    "<LOC_439>",
+    "<LOC_440>",
+    "<LOC_441>",
+    "<LOC_442>",
+    "<LOC_443>",
+    "<LOC_444>",
+    "<LOC_445>",
+    "<LOC_446>",
+    "<LOC_447>",
+    "<LOC_448>",
+    "<LOC_449>",
+    "<LOC_450>",
+    "<LOC_451>",
+    "<LOC_452>",
+    "<LOC_453>",
+    "<LOC_454>",
+    "<LOC_455>",
+    "<LOC_456>",
+    "<LOC_457>",
+    "<LOC_458>",
+    "<LOC_459>",
+    "<LOC_460>",
+    "<LOC_461>",
+    "<LOC_462>",
+    "<LOC_463>",
+    "<LOC_464>",
+    "<LOC_465>",
+    "<LOC_466>",
+    "<LOC_467>",
+    "<LOC_468>",
+    "<LOC_469>",
+    "<LOC_470>",
+    "<LOC_471>",
+    "<LOC_472>",
+    "<LOC_473>",
+    "<LOC_474>",
+    "<LOC_475>",
+    "<LOC_476>",
+    "<LOC_477>",
+    "<LOC_478>",
+    "<LOC_479>",
+    "<LOC_480>",
+    "<LOC_481>",
+    "<LOC_482>",
+    "<LOC_483>",
+    "<LOC_484>",
+    "<LOC_485>",
+    "<LOC_486>",
+    "<LOC_487>",
+    "<LOC_488>",
+    "<LOC_489>",
+    "<LOC_490>",
+    "<LOC_491>",
+    "<LOC_492>",
+    "<LOC_493>",
+    "<LOC_494>",
+    "<LOC_495>",
+    "<LOC_496>",
+    "<LOC_497>",
+    "<LOC_498>",
+    "<LOC_499>",
+    "<LOC_500>",
+    "<LOC_501>",
+    "<LOC_502>",
+    "<LOC_503>",
+    "<LOC_504>",
+    "<LOC_505>",
+    "<LOC_506>",
+    "<LOC_507>",
+    "<LOC_508>",
+    "<LOC_509>",
+    "<LOC_510>",
+    "<LOC_511>",
+    "<LOC_512>",
+    "<LOC_513>",
+    "<LOC_514>",
+    "<LOC_515>",
+    "<LOC_516>",
+    "<LOC_517>",
+    "<LOC_518>",
+    "<LOC_519>",
+    "<LOC_520>",
+    "<LOC_521>",
+    "<LOC_522>",
+    "<LOC_523>",
+    "<LOC_524>",
+    "<LOC_525>",
+    "<LOC_526>",
+    "<LOC_527>",
+    "<LOC_528>",
+    "<LOC_529>",
+    "<LOC_530>",
+    "<LOC_531>",
+    "<LOC_532>",
+    "<LOC_533>",
+    "<LOC_534>",
+    "<LOC_535>",
+    "<LOC_536>",
+    "<LOC_537>",
+    "<LOC_538>",
+    "<LOC_539>",
+    "<LOC_540>",
+    "<LOC_541>",
+    "<LOC_542>",
+    "<LOC_543>",
+    "<LOC_544>",
+    "<LOC_545>",
+    "<LOC_546>",
+    "<LOC_547>",
+    "<LOC_548>",
+    "<LOC_549>",
+    "<LOC_550>",
+    "<LOC_551>",
+    "<LOC_552>",
+    "<LOC_553>",
+    "<LOC_554>",
+    "<LOC_555>",
+    "<LOC_556>",
+    "<LOC_557>",
+    "<LOC_558>",
+    "<LOC_559>",
+    "<LOC_560>",
+    "<LOC_561>",
+    "<LOC_562>",
+    "<LOC_563>",
+    "<LOC_564>",
+    "<LOC_565>",
+    "<LOC_566>",
+    "<LOC_567>",
+    "<LOC_568>",
+    "<LOC_569>",
+    "<LOC_570>",
+    "<LOC_571>",
+    "<LOC_572>",
+    "<LOC_573>",
+    "<LOC_574>",
+    "<LOC_575>",
+    "<LOC_576>",
+    "<LOC_577>",
+    "<LOC_578>",
+    "<LOC_579>",
+    "<LOC_580>",
+    "<LOC_581>",
+    "<LOC_582>",
+    "<LOC_583>",
+    "<LOC_584>",
+    "<LOC_585>",
+    "<LOC_586>",
+    "<LOC_587>",
+    "<LOC_588>",
+    "<LOC_589>",
+    "<LOC_590>",
+    "<LOC_591>",
+    "<LOC_592>",
+    "<LOC_593>",
+    "<LOC_594>",
+    "<LOC_595>",
+    "<LOC_596>",
+    "<LOC_597>",
+    "<LOC_598>",
+    "<LOC_599>",
+    "<LOC_600>",
+    "<LOC_601>",
+    "<LOC_602>",
+    "<LOC_603>",
+    "<LOC_604>",
+    "<LOC_605>",
+    "<LOC_606>",
+    "<LOC_607>",
+    "<LOC_608>",
+    "<LOC_609>",
+    "<LOC_610>",
+    "<LOC_611>",
+    "<LOC_612>",
+    "<LOC_613>",
+    "<LOC_614>",
+    "<LOC_615>",
+    "<LOC_616>",
+    "<LOC_617>",
+    "<LOC_618>",
+    "<LOC_619>",
+    "<LOC_620>",
+    "<LOC_621>",
+    "<LOC_622>",
+    "<LOC_623>",
+    "<LOC_624>",
+    "<LOC_625>",
+    "<LOC_626>",
+    "<LOC_627>",
+    "<LOC_628>",
+    "<LOC_629>",
+    "<LOC_630>",
+    "<LOC_631>",
+    "<LOC_632>",
+    "<LOC_633>",
+    "<LOC_634>",
+    "<LOC_635>",
+    "<LOC_636>",
+    "<LOC_637>",
+    "<LOC_638>",
+    "<LOC_639>",
+    "<LOC_640>",
+    "<LOC_641>",
+    "<LOC_642>",
+    "<LOC_643>",
+    "<LOC_644>",
+    "<LOC_645>",
+    "<LOC_646>",
+    "<LOC_647>",
+    "<LOC_648>",
+    "<LOC_649>",
+    "<LOC_650>",
+    "<LOC_651>",
+    "<LOC_652>",
+    "<LOC_653>",
+    "<LOC_654>",
+    "<LOC_655>",
+    "<LOC_656>",
+    "<LOC_657>",
+    "<LOC_658>",
+    "<LOC_659>",
+    "<LOC_660>",
+    "<LOC_661>",
+    "<LOC_662>",
+    "<LOC_663>",
+    "<LOC_664>",
+    "<LOC_665>",
+    "<LOC_666>",
+    "<LOC_667>",
+    "<LOC_668>",
+    "<LOC_669>",
+    "<LOC_670>",
+    "<LOC_671>",
+    "<LOC_672>",
+    "<LOC_673>",
+    "<LOC_674>",
+    "<LOC_675>",
+    "<LOC_676>",
+    "<LOC_677>",
+    "<LOC_678>",
+    "<LOC_679>",
+    "<LOC_680>",
+    "<LOC_681>",
+    "<LOC_682>",
+    "<LOC_683>",
+    "<LOC_684>",
+    "<LOC_685>",
+    "<LOC_686>",
+    "<LOC_687>",
+    "<LOC_688>",
+    "<LOC_689>",
+    "<LOC_690>",
+    "<LOC_691>",
+    "<LOC_692>",
+    "<LOC_693>",
+    "<LOC_694>",
+    "<LOC_695>",
+    "<LOC_696>",
+    "<LOC_697>",
+    "<LOC_698>",
+    "<LOC_699>",
+    "<LOC_700>",
+    "<LOC_701>",
+    "<LOC_702>",
+    "<LOC_703>",
+    "<LOC_704>",
+    "<LOC_705>",
+    "<LOC_706>",
+    "<LOC_707>",
+    "<LOC_708>",
+    "<LOC_709>",
+    "<LOC_710>",
+    "<LOC_711>",
+    "<LOC_712>",
+    "<LOC_713>",
+    "<LOC_714>",
+    "<LOC_715>",
+    "<LOC_716>",
+    "<LOC_717>",
+    "<LOC_718>",
+    "<LOC_719>",
+    "<LOC_720>",
+    "<LOC_721>",
+    "<LOC_722>",
+    "<LOC_723>",
+    "<LOC_724>",
+    "<LOC_725>",
+    "<LOC_726>",
+    "<LOC_727>",
+    "<LOC_728>",
+    "<LOC_729>",
+    "<LOC_730>",
+    "<LOC_731>",
+    "<LOC_732>",
+    "<LOC_733>",
+    "<LOC_734>",
+    "<LOC_735>",
+    "<LOC_736>",
+    "<LOC_737>",
+    "<LOC_738>",
+    "<LOC_739>",
+    "<LOC_740>",
+    "<LOC_741>",
+    "<LOC_742>",
+    "<LOC_743>",
+    "<LOC_744>",
+    "<LOC_745>",
+    "<LOC_746>",
+    "<LOC_747>",
+    "<LOC_748>",
+    "<LOC_749>",
+    "<LOC_750>",
+    "<LOC_751>",
+    "<LOC_752>",
+    "<LOC_753>",
+    "<LOC_754>",
+    "<LOC_755>",
+    "<LOC_756>",
+    "<LOC_757>",
+    "<LOC_758>",
+    "<LOC_759>",
+    "<LOC_760>",
+    "<LOC_761>",
+    "<LOC_762>",
+    "<LOC_763>",
+    "<LOC_764>",
+    "<LOC_765>",
+    "<LOC_766>",
+    "<LOC_767>",
+    "<LOC_768>",
+    "<LOC_769>",
+    "<LOC_770>",
+    "<LOC_771>",
+    "<LOC_772>",
+    "<LOC_773>",
+    "<LOC_774>",
+    "<LOC_775>",
+    "<LOC_776>",
+    "<LOC_777>",
+    "<LOC_778>",
+    "<LOC_779>",
+    "<LOC_780>",
+    "<LOC_781>",
+    "<LOC_782>",
+    "<LOC_783>",
+    "<LOC_784>",
+    "<LOC_785>",
+    "<LOC_786>",
+    "<LOC_787>",
+    "<LOC_788>",
+    "<LOC_789>",
+    "<LOC_790>",
+    "<LOC_791>",
+    "<LOC_792>",
+    "<LOC_793>",
+    "<LOC_794>",
+    "<LOC_795>",
+    "<LOC_796>",
+    "<LOC_797>",
+    "<LOC_798>",
+    "<LOC_799>",
+    "<LOC_800>",
+    "<LOC_801>",
+    "<LOC_802>",
+    "<LOC_803>",
+    "<LOC_804>",
+    "<LOC_805>",
+    "<LOC_806>",
+    "<LOC_807>",
+    "<LOC_808>",
+    "<LOC_809>",
+    "<LOC_810>",
+    "<LOC_811>",
+    "<LOC_812>",
+    "<LOC_813>",
+    "<LOC_814>",
+    "<LOC_815>",
+    "<LOC_816>",
+    "<LOC_817>",
+    "<LOC_818>",
+    "<LOC_819>",
+    "<LOC_820>",
+    "<LOC_821>",
+    "<LOC_822>",
+    "<LOC_823>",
+    "<LOC_824>",
+    "<LOC_825>",
+    "<LOC_826>",
+    "<LOC_827>",
+    "<LOC_828>",
+    "<LOC_829>",
+    "<LOC_830>",
+    "<LOC_831>",
+    "<LOC_832>",
+    "<LOC_833>",
+    "<LOC_834>",
+    "<LOC_835>",
+    "<LOC_836>",
+    "<LOC_837>",
+    "<LOC_838>",
+    "<LOC_839>",
+    "<LOC_840>",
+    "<LOC_841>",
+    "<LOC_842>",
+    "<LOC_843>",
+    "<LOC_844>",
+    "<LOC_845>",
+    "<LOC_846>",
+    "<LOC_847>",
+    "<LOC_848>",
+    "<LOC_849>",
+    "<LOC_850>",
+    "<LOC_851>",
+    "<LOC_852>",
+    "<LOC_853>",
+    "<LOC_854>",
+    "<LOC_855>",
+    "<LOC_856>",
+    "<LOC_857>",
+    "<LOC_858>",
+    "<LOC_859>",
+    "<LOC_860>",
+    "<LOC_861>",
+    "<LOC_862>",
+    "<LOC_863>",
+    "<LOC_864>",
+    "<LOC_865>",
+    "<LOC_866>",
+    "<LOC_867>",
+    "<LOC_868>",
+    "<LOC_869>",
+    "<LOC_870>",
+    "<LOC_871>",
+    "<LOC_872>",
+    "<LOC_873>",
+    "<LOC_874>",
+    "<LOC_875>",
+    "<LOC_876>",
+    "<LOC_877>",
+    "<LOC_878>",
+    "<LOC_879>",
+    "<LOC_880>",
+    "<LOC_881>",
+    "<LOC_882>",
+    "<LOC_883>",
+    "<LOC_884>",
+    "<LOC_885>",
+    "<LOC_886>",
+    "<LOC_887>",
+    "<LOC_888>",
+    "<LOC_889>",
+    "<LOC_890>",
+    "<LOC_891>",
+    "<LOC_892>",
+    "<LOC_893>",
+    "<LOC_894>",
+    "<LOC_895>",
+    "<LOC_896>",
+    "<LOC_897>",
+    "<LOC_898>",
+    "<LOC_899>",
+    "<LOC_900>",
+    "<LOC_901>",
+    "<LOC_902>",
+    "<LOC_903>",
+    "<LOC_904>",
+    "<LOC_905>",
+    "<LOC_906>",
+    "<LOC_907>",
+    "<LOC_908>",
+    "<LOC_909>",
+    "<LOC_910>",
+    "<LOC_911>",
+    "<LOC_912>",
+    "<LOC_913>",
+    "<LOC_914>",
+    "<LOC_915>",
+    "<LOC_916>",
+    "<LOC_917>",
+    "<LOC_918>",
+    "<LOC_919>",
+    "<LOC_920>",
+    "<LOC_921>",
+    "<LOC_922>",
+    "<LOC_923>",
+    "<LOC_924>",
+    "<LOC_925>",
+    "<LOC_926>",
+    "<LOC_927>",
+    "<LOC_928>",
+    "<LOC_929>",
+    "<LOC_930>",
+    "<LOC_931>",
+    "<LOC_932>",
+    "<LOC_933>",
+    "<LOC_934>",
+    "<LOC_935>",
+    "<LOC_936>",
+    "<LOC_937>",
+    "<LOC_938>",
+    "<LOC_939>",
+    "<LOC_940>",
+    "<LOC_941>",
+    "<LOC_942>",
+    "<LOC_943>",
+    "<LOC_944>",
+    "<LOC_945>",
+    "<LOC_946>",
+    "<LOC_947>",
+    "<LOC_948>",
+    "<LOC_949>",
+    "<LOC_950>",
+    "<LOC_951>",
+    "<LOC_952>",
+    "<LOC_953>",
+    "<LOC_954>",
+    "<LOC_955>",
+    "<LOC_956>",
+    "<LOC_957>",
+    "<LOC_958>",
+    "<LOC_959>",
+    "<LOC_960>",
+    "<LOC_961>",
+    "<LOC_962>",
+    "<LOC_963>",
+    "<LOC_964>",
+    "<LOC_965>",
+    "<LOC_966>",
+    "<LOC_967>",
+    "<LOC_968>",
+    "<LOC_969>",
+    "<LOC_970>",
+    "<LOC_971>",
+    "<LOC_972>",
+    "<LOC_973>",
+    "<LOC_974>",
+    "<LOC_975>",
+    "<LOC_976>",
+    "<LOC_977>",
+    "<LOC_978>",
+    "<LOC_979>",
+    "<LOC_980>",
+    "<LOC_981>",
+    "<LOC_982>",
+    "<LOC_983>",
+    "<LOC_984>",
+    "<LOC_985>",
+    "<LOC_986>",
+    "<LOC_987>",
+    "<LOC_988>",
+    "<LOC_989>",
+    "<LOC_990>",
+    "<LOC_991>",
+    "<LOC_992>",
+    "<LOC_993>",
+    "<LOC_994>",
+    "<LOC_995>",
+    "<LOC_996>",
+    "<LOC_997>",
+    "<LOC_998>",
+    "<LOC_999>",
+    "<LOC_1000>",
+    "<LOC_1001>",
+    "<LOC_1002>",
+    "<LOC_1003>",
+    "<LOC_1004>",
+    "<LOC_1005>",
+    "<LOC_1006>",
+    "<LOC_1007>",
+    "<LOC_1008>",
+    "<LOC_1009>",
+    "<LOC_1010>",
+    "<LOC_1011>",
+    "<LOC_1012>",
+    "<LOC_1013>",
+    "<LOC_1014>",
+    "<LOC_1015>",
+    "<LOC_1016>",
+    "<LOC_1017>",
+    "<LOC_1018>",
+    "<LOC_1019>",
+    "<LOC_1020>",
+    "<LOC_1021>",
+    "<LOC_1022>",
+    "<LOC_1023>",
+    "<LOC_1024>",
+    "<box_p>",
+    "</box_p>",
+    "<LOC_BEGIN>",
+    "<LOC_END>"
+  ],
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

token_constants.py

@@ -0,0 +1,4 @@
+IMAGE_TOKEN_INDEX = -200
+IMAGE_PAD_TOKEN_INDEX = -500
+DEFAULT_IMAGE_TOKEN = '<image>'
+DEFAULT_IMAGE_PAD_TOKEN = '<image_padding>'

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f9ad3482369ba91711ff939f9eeb761d618b3157fcbfc0715f18e34b918eca97
+size 17268792

tree.py

@@ -0,0 +1,414 @@
+
+from __future__ import annotations
+import hashlib
+import math
+import os
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from numbers import Integral, Real
+from pathlib import Path
+import torch
+import torch.nn.functional as F
+PINNED_ARTIFACT_VERSION = 1
+PINNED_ARTIFACT_ALGORITHM = 'recursive_centered_pca_rank_median_v1'
+PINNED_ARTIFACT_METRIC = 'normalized_squared_l2_tree_ultrametric'
+PINNED_PRODUCTION_VOCAB_SIZE = 131072
+PINNED_PRODUCTION_DEPTH = 17
+DEFAULT_SEMANTIC_TREE_STEPS = 4
+DEFAULT_SEMANTIC_TREE_SCHEDULE = 'shift'
+DEFAULT_SEMANTIC_TREE_SHIFT = 5.0
+MAX_SEMANTIC_TREE_REFINEMENT_GAP = 2
+SUPPORTED_SEMANTIC_TREE_INFERENCE_SCHEDULES = frozenset({'linear', 'shift', 'logit_normal', 'cosine', 'cosine2', 'exp'})
+
+def _plain_positive_int(value: object, *, name: str) -> int:
+    if isinstance(value, bool) or not isinstance(value, Integral):
+        raise TypeError(f'{name} must be a positive integer')
+    value = int(value)
+    if value < 1:
+        raise ValueError(f'{name} must be a positive integer')
+    return value
+
+def _plain_nonnegative_int(value: object, *, name: str) -> int:
+    if isinstance(value, bool) or not isinstance(value, Integral):
+        raise TypeError(f'{name} must be a non-negative integer')
+    value = int(value)
+    if value < 0:
+        raise ValueError(f'{name} must be a non-negative integer')
+    return value
+
+def _plain_positive_float(value: object, *, name: str) -> float:
+    if isinstance(value, bool) or not isinstance(value, Real):
+        raise TypeError(f'{name} must be a finite positive number')
+    value = float(value)
+    if not math.isfinite(value) or value <= 0.0:
+        raise ValueError(f'{name} must be a finite positive number')
+    return value
+
+def _is_power_of_two(value: int) -> bool:
+    return value > 0 and value & value - 1 == 0
+
+def resolve_semantic_tree_steps(config, override=None):
+    steps = 4 if override is None else override
+    if isinstance(steps, bool) or steps not in (1, 4, 8, 16):
+        raise ValueError('refinement_steps must be 1, 4, 8, or 16')
+    return int(steps)
+
+@dataclass(frozen=True, slots=True)
+class SemanticTreeArtifactContract:
+    path: str | os.PathLike[str]
+    bytes: int | None = None
+    sha256: str | None = None
+    algorithm: str | None = PINNED_ARTIFACT_ALGORITHM
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.path, (str, os.PathLike)):
+            raise TypeError('semantic-tree artifact path must be path-like')
+        raw_path = os.fspath(self.path)
+        if not isinstance(raw_path, str) or not raw_path.strip():
+            raise ValueError('semantic-tree artifact path must be nonempty')
+        object.__setattr__(self, 'path', Path(raw_path).expanduser())
+        if self.bytes is not None:
+            object.__setattr__(self, 'bytes', _plain_positive_int(self.bytes, name='semantic-tree artifact bytes'))
+        if self.sha256 is not None:
+            if not isinstance(self.sha256, str):
+                raise TypeError('semantic-tree artifact sha256 must be a string')
+            digest = self.sha256.strip().lower()
+            if len(digest) != 64 or any((character not in '0123456789abcdef' for character in digest)):
+                raise ValueError('semantic-tree artifact sha256 must be 64 hexadecimal characters')
+            object.__setattr__(self, 'sha256', digest)
+        if self.algorithm not in (None, PINNED_ARTIFACT_ALGORITHM):
+            raise ValueError(f'Unsupported semantic-tree algorithm: {self.algorithm!r}')
+
+def production_semantic_tree_contract(path: str | os.PathLike[str]) -> SemanticTreeArtifactContract:
+    return SemanticTreeArtifactContract(path=path, sha256='4b7ee94fbb10fd4146a4ad167b595b037e3f591eaa12ecf3f0b411c82384a714', algorithm=PINNED_ARTIFACT_ALGORITHM)
+
+def _sha256_file(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open('rb') as artifact_file:
+        for block in iter(lambda : artifact_file.read(1024 * 1024), b''):
+            digest.update(block)
+    return digest.hexdigest()
+
+class BalancedSemanticTree:
+
+    def __init__(self, leaf_to_token: torch.Tensor) -> None:
+        if not torch.is_tensor(leaf_to_token):
+            raise TypeError('leaf_to_token must be a tensor')
+        if leaf_to_token.ndim != 1:
+            raise ValueError('leaf_to_token must be one-dimensional')
+        if leaf_to_token.dtype == torch.bool or leaf_to_token.is_floating_point() or leaf_to_token.is_complex():
+            raise TypeError('leaf_to_token must contain integer token IDs')
+        vocab_size = int(leaf_to_token.numel())
+        if not _is_power_of_two(vocab_size):
+            raise ValueError('semantic-tree vocabulary size must be a power of two')
+        permutation = leaf_to_token.to(dtype=torch.long).contiguous()
+        expected = torch.arange(vocab_size, device=permutation.device)
+        if not torch.equal(torch.sort(permutation).values, expected):
+            raise ValueError('leaf_to_token must be a permutation of [0, V)')
+        token_to_leaf = torch.empty_like(permutation)
+        token_to_leaf.scatter_(0, permutation, expected)
+        self.vocab_size = vocab_size
+        self.depth = int(math.log2(vocab_size))
+        self.leaf_to_token = permutation
+        self.token_to_leaf = token_to_leaf
+
+    @classmethod
+    def _from_validated(cls, leaf_to_token: torch.Tensor, token_to_leaf: torch.Tensor) -> BalancedSemanticTree:
+        tree = cls.__new__(cls)
+        tree.vocab_size = int(leaf_to_token.numel())
+        tree.depth = int(math.log2(tree.vocab_size))
+        tree.leaf_to_token = leaf_to_token
+        tree.token_to_leaf = token_to_leaf
+        return tree
+
+    @classmethod
+    def from_artifact(cls, contract: SemanticTreeArtifactContract, expected_vocab_size: int=PINNED_PRODUCTION_VOCAB_SIZE) -> BalancedSemanticTree:
+        if not isinstance(contract, SemanticTreeArtifactContract):
+            raise TypeError('contract must be a SemanticTreeArtifactContract')
+        expected_vocab_size = _plain_positive_int(expected_vocab_size, name='expected_vocab_size')
+        if not _is_power_of_two(expected_vocab_size):
+            raise ValueError('expected_vocab_size must be a power of two')
+        path = contract.path
+        assert isinstance(path, Path)
+        if not path.is_file():
+            raise FileNotFoundError(f'semantic-tree artifact is not a file: {path}')
+        actual_bytes = path.stat().st_size
+        if contract.bytes is not None and actual_bytes != contract.bytes:
+            raise ValueError(f'semantic-tree artifact byte count mismatch: expected {contract.bytes}, got {actual_bytes}')
+        if contract.sha256 is not None:
+            actual_sha256 = _sha256_file(path)
+            if actual_sha256 != contract.sha256:
+                raise ValueError(f'semantic-tree artifact sha256 mismatch: expected {contract.sha256}, got {actual_sha256}')
+        artifact = torch.load(path, map_location='cpu', weights_only=True)
+        if not isinstance(artifact, Mapping):
+            raise TypeError('semantic-tree artifact must contain a mapping')
+        algorithm = artifact.get('algorithm')
+        allowed_algorithms = (contract.algorithm,) if contract.algorithm is not None else (PINNED_ARTIFACT_ALGORITHM,)
+        if algorithm not in allowed_algorithms:
+            raise ValueError(f'semantic-tree artifact algorithm mismatch: expected one of {allowed_algorithms!r}, got {algorithm!r}')
+        pinned_values = {'version': PINNED_ARTIFACT_VERSION, 'metric': PINNED_ARTIFACT_METRIC, 'vocab_size': expected_vocab_size, 'depth': int(math.log2(expected_vocab_size)), 'radii_monotone': True}
+        for (field_name, expected_value) in pinned_values.items():
+            if artifact.get(field_name) != expected_value:
+                raise ValueError(f'semantic-tree artifact {field_name} mismatch: expected {expected_value!r}, got {artifact.get(field_name)!r}')
+        leaf_to_token = artifact.get('leaf_to_token')
+        heap_radii = artifact.get('heap_radii')
+        if not torch.is_tensor(leaf_to_token) or leaf_to_token.dtype != torch.int32:
+            raise TypeError('semantic-tree leaf_to_token must be torch.int32')
+        if tuple(leaf_to_token.shape) != (expected_vocab_size,):
+            raise ValueError('semantic-tree leaf_to_token shape disagrees with vocabulary')
+        if not torch.is_tensor(heap_radii) or heap_radii.dtype != torch.float32:
+            raise TypeError('semantic-tree heap_radii must be torch.float32')
+        if tuple(heap_radii.shape) != (expected_vocab_size - 1,):
+            raise ValueError('semantic-tree heap_radii shape disagrees with vocabulary')
+        if not bool(torch.isfinite(heap_radii).all()) or bool((heap_radii < 0).any()):
+            raise ValueError('semantic-tree heap_radii must be finite and nonnegative')
+        return cls(leaf_to_token)
+
+    @property
+    def device(self) -> torch.device:
+        return self.leaf_to_token.device
+
+    def to(self, device: torch.device | str) -> BalancedSemanticTree:
+        resolved = torch.device(device)
+        if resolved == self.device:
+            return self
+        return type(self)._from_validated(self.leaf_to_token.to(resolved), self.token_to_leaf.to(resolved))
+
+    def _check_tokens(self, tokens: torch.Tensor, *, name: str) -> None:
+        if not torch.is_tensor(tokens):
+            raise TypeError(f'{name} must be a tensor')
+        if tokens.dtype == torch.bool or tokens.is_floating_point() or tokens.is_complex():
+            raise TypeError(f'{name} must contain integer code IDs')
+        if tokens.device != self.device:
+            raise ValueError(f'{name} and semantic-tree lookups must share a device')
+        if tokens.numel() and bool(((tokens < 0) | (tokens >= self.vocab_size)).any()):
+            raise ValueError(f'{name} contains code IDs outside [0, V)')
+
+    def _levels(self, levels: int | torch.Tensor, shape: torch.Size, device: torch.device) -> torch.Tensor:
+        if isinstance(levels, bool) or (not isinstance(levels, Integral) and (not torch.is_tensor(levels))):
+            raise TypeError('levels must contain integers')
+        level_tensor = torch.as_tensor(levels, device=device)
+        if level_tensor.dtype == torch.bool or level_tensor.is_floating_point() or level_tensor.is_complex():
+            raise TypeError('levels must contain integers')
+        if level_tensor.ndim < len(shape):
+            level_tensor = level_tensor.reshape(level_tensor.shape + (1,) * (len(shape) - level_tensor.ndim))
+        try:
+            level_tensor = torch.broadcast_to(level_tensor, shape).long()
+        except RuntimeError as error:
+            raise ValueError('levels are not broadcastable to endpoint shape') from error
+        if bool((level_tensor < 0).any()) or bool((level_tensor > self.depth).any()):
+            raise ValueError(f'levels must lie in [0, {self.depth}]')
+        return level_tensor
+
+    def node_ids(self, tokens: torch.Tensor, levels: int | torch.Tensor) -> torch.Tensor:
+        self._check_tokens(tokens, name='tokens')
+        level_tensor = self._levels(levels, tokens.shape, tokens.device)
+        leaves = self.token_to_leaf[tokens.long()]
+        return torch.bitwise_right_shift(leaves, self.depth - level_tensor)
+
+    @torch.no_grad()
+    def sample_representatives(self, endpoints: torch.Tensor, levels: int | torch.Tensor, generator: torch.Generator | None=None) -> torch.Tensor:
+        self._check_tokens(endpoints, name='endpoints')
+        level_tensor = self._levels(levels, endpoints.shape, endpoints.device)
+        leaves = self.token_to_leaf[endpoints.long()]
+        noisy = level_tensor < self.depth
+        if not bool(noisy.any()):
+            return endpoints.clone()
+        noisy_leaves = leaves[noisy]
+        noisy_levels = level_tensor[noisy]
+        shifts = self.depth - noisy_levels
+        subtree_starts = torch.bitwise_left_shift(torch.bitwise_right_shift(noisy_leaves, shifts), shifts)
+        subtree_sizes = torch.bitwise_left_shift(torch.ones_like(shifts), shifts)
+        random_bits = torch.randint(self.vocab_size, (int(noisy.sum().item()),), device=endpoints.device, dtype=torch.long, generator=generator)
+        sampled_leaves = subtree_starts + torch.bitwise_and(random_bits, subtree_sizes - 1)
+        result = endpoints.clone()
+        result[noisy] = self.leaf_to_token[sampled_leaves].to(endpoints.dtype)
+        return result
+
+def _validate_ages(ages: torch.Tensor, *, steps: int) -> tuple[torch.Tensor, int]:
+    steps = _plain_positive_int(steps, name='steps')
+    if not torch.is_tensor(ages):
+        raise TypeError('ages must be a tensor')
+    if ages.dtype == torch.bool or ages.is_floating_point() or ages.is_complex():
+        raise TypeError('ages must contain integers')
+    if bool((ages < 0).any()) or bool((ages > steps).any()):
+        raise ValueError(f'ages must lie in [0, {steps}]')
+    return (ages.long(), steps)
+
+def semantic_tree_levels_for_ages(ages: torch.Tensor, steps: int=DEFAULT_SEMANTIC_TREE_STEPS, depth: int=PINNED_PRODUCTION_DEPTH) -> torch.Tensor:
+    (ages, steps) = _validate_ages(ages, steps=steps)
+    depth = _plain_positive_int(depth, name='depth')
+    if steps > depth:
+        raise ValueError('steps cannot exceed tree depth')
+    return torch.div(ages * depth, steps, rounding_mode='floor')
+
+def semantic_tree_noise_for_ages(ages: torch.Tensor, steps: int=DEFAULT_SEMANTIC_TREE_STEPS) -> torch.Tensor:
+    (ages, steps) = _validate_ages(ages, steps=steps)
+    return 1.0 - ages.to(dtype=torch.float32) / float(steps)
+
+def semantic_tree_activation_counts(content_tokens: int, *, activation_steps: int, schedule: str=DEFAULT_SEMANTIC_TREE_SCHEDULE, shift: float=DEFAULT_SEMANTIC_TREE_SHIFT, device: torch.device | str | None=None) -> torch.Tensor:
+    content_tokens = _plain_positive_int(content_tokens, name='content_tokens')
+    activation_steps = min(_plain_positive_int(activation_steps, name='activation_steps'), content_tokens)
+    if not isinstance(schedule, str):
+        raise TypeError('schedule must be a string')
+    schedule = schedule.strip().lower()
+    if schedule not in SUPPORTED_SEMANTIC_TREE_INFERENCE_SCHEDULES:
+        allowed = ', '.join(sorted(SUPPORTED_SEMANTIC_TREE_INFERENCE_SCHEDULES))
+        raise ValueError(f'unsupported semantic-tree inference schedule {schedule!r}; expected one of: {allowed}')
+    shift = _plain_positive_float(shift, name='shift')
+    sample_times = torch.linspace(0.0, 1.0, activation_steps + 1, dtype=torch.float32)
+    if schedule == 'linear':
+        sigmas = sample_times
+    elif schedule == 'shift':
+        sigmas = shift * sample_times / (1.0 + (shift - 1.0) * sample_times)
+    elif schedule == 'logit_normal':
+        sigmas = 0.5 * (1.0 + torch.erf(torch.logit(sample_times) / math.sqrt(2.0)))
+    elif schedule == 'cosine':
+        sigmas = 1.0 - 0.5 * (1.0 + torch.cos(math.pi * sample_times))
+    elif schedule == 'cosine2':
+        sigmas = torch.sin(0.5 * math.pi * sample_times)
+    else:
+        sigmas = (1.0 - torch.exp(-5.0 * sample_times)) / (1.0 - math.exp(-5.0))
+        sigmas = sigmas.clamp(0.0001, 1.0)
+    boundaries = (sigmas * content_tokens).to(dtype=torch.long)
+    counts = (boundaries[1:] - boundaries[:-1]).clamp_min(1)
+    excess = int(counts.sum().item()) - content_tokens
+    cursor = 0
+    while excess > 0:
+        cursor %= counts.numel()
+        if int(counts[cursor].item()) == 1:
+            cursor += 1
+            continue
+        counts[cursor] -= 1
+        excess -= 1
+        cursor += 1
+    if int(counts.sum().item()) != content_tokens:
+        raise RuntimeError('semantic-tree activation counts do not cover content')
+    return counts.flip(0).to(device=device)
+
+def resolve_semantic_tree_atomic_block_size(value: object=1) -> int:
+    block_size = _plain_positive_int(value, name='semantic_tree_atomic_block_size')
+    if block_size not in {1, 2, 4, 8}:
+        raise ValueError('semantic_tree_atomic_block_size currently supports 1, 2, 4, or 8')
+    return block_size
+
+def resolve_semantic_tree_refinement_gap(value: object=0) -> int:
+    gap = _plain_nonnegative_int(value, name='semantic_tree_refinement_gap')
+    if gap > MAX_SEMANTIC_TREE_REFINEMENT_GAP:
+        raise ValueError('semantic_tree_refinement_gap currently supports only 0, 1, or 2')
+    return gap
+
+def semantic_tree_decode_schedule(activation_counts: torch.Tensor, *, refinement_steps: int, refinement_gap: int=0) -> torch.Tensor:
+    if not torch.is_tensor(activation_counts):
+        raise TypeError('activation_counts must be a tensor')
+    if activation_counts.ndim != 2:
+        raise ValueError('activation_counts must have shape [batch, rounds]')
+    if activation_counts.dtype not in {torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64}:
+        raise TypeError('activation_counts must use an integer dtype')
+    if activation_counts.shape[1] == 0:
+        raise ValueError('activation_counts must contain at least one round')
+    if bool((activation_counts < 0).any()):
+        raise ValueError('activation_counts must be non-negative')
+    refinement_steps = _plain_positive_int(refinement_steps, name='refinement_steps')
+    refinement_gap = resolve_semantic_tree_refinement_gap(refinement_gap)
+    flush_rounds = refinement_steps - 1 + refinement_steps * refinement_gap
+    if flush_rounds == 0:
+        return activation_counts
+    flush = torch.zeros(activation_counts.shape[0], flush_rounds, dtype=activation_counts.dtype, device=activation_counts.device)
+    return torch.cat((activation_counts, flush), dim=1)
+
+def semantic_tree_refinement_due(ages: torch.Tensor, next_refinement_round: torch.Tensor, *, decode_round: int, refinement_steps: int) -> torch.Tensor:
+    if not torch.is_tensor(ages) or not torch.is_tensor(next_refinement_round):
+        raise TypeError('ages and next_refinement_round must be tensors')
+    if ages.shape != next_refinement_round.shape:
+        raise ValueError('ages and next_refinement_round must align')
+    if ages.device != next_refinement_round.device:
+        raise ValueError('ages and next_refinement_round must share a device')
+    for (tensor, name) in ((ages, 'ages'), (next_refinement_round, 'next_refinement_round')):
+        if tensor.dtype == torch.bool or tensor.is_floating_point():
+            raise TypeError(f'{name} must use an integer dtype')
+    decode_round = _plain_nonnegative_int(decode_round, name='decode_round')
+    refinement_steps = _plain_positive_int(refinement_steps, name='refinement_steps')
+    return ages.ge(0) & ages.lt(refinement_steps) & next_refinement_round.le(decode_round)
+EmbeddingSource = torch.Tensor | Callable[[torch.Tensor], torch.Tensor]
+
+def _embed_ids(source: EmbeddingSource, token_ids: torch.Tensor) -> torch.Tensor:
+    if torch.is_tensor(source):
+        if source.ndim != 2 or not source.is_floating_point():
+            raise TypeError('embedding weight must be a floating-point matrix')
+        if source.device != token_ids.device:
+            raise ValueError('embedding weight and token IDs must share a device')
+        if token_ids.numel() and bool(((token_ids < 0) | (token_ids >= source.shape[0])).any()):
+            raise ValueError('token ID lies outside the embedding weight')
+        return F.embedding(token_ids.long(), source)
+    if not callable(source):
+        raise TypeError('embedding source must be a weight tensor or callable')
+    embedded = source(token_ids)
+    if not torch.is_tensor(embedded):
+        raise TypeError('embedding callback must return a tensor')
+    return embedded
+
+def build_semantic_tree_embeddings(representatives: torch.Tensor, noise_levels: float | torch.Tensor, embedding_callback: EmbeddingSource, mask_token_id: int) -> torch.Tensor:
+    if not torch.is_tensor(representatives):
+        raise TypeError('representatives must be a tensor')
+    if representatives.dtype == torch.bool or representatives.is_floating_point() or representatives.is_complex():
+        raise TypeError('representatives must contain integer code IDs')
+    if isinstance(mask_token_id, bool) or not isinstance(mask_token_id, Integral):
+        raise TypeError('mask_token_id must be an integer')
+    mask_token_id = int(mask_token_id)
+    if mask_token_id < 0:
+        raise ValueError('mask_token_id must be nonnegative')
+    representative_embeddings = _embed_ids(embedding_callback, representatives)
+    expected_prefix = tuple(representatives.shape)
+    if representative_embeddings.ndim != representatives.ndim + 1 or tuple(representative_embeddings.shape[:-1]) != expected_prefix or (not representative_embeddings.is_floating_point()):
+        raise ValueError('representative embeddings must have shape representatives.shape + [H]')
+    mask_ids = torch.tensor([mask_token_id], dtype=representatives.dtype, device=representatives.device)
+    mask_embedding = _embed_ids(embedding_callback, mask_ids)
+    if mask_embedding.shape != (1, representative_embeddings.shape[-1]):
+        raise ValueError('MASK embedding must have shape [1, H]')
+    if mask_embedding.device != representative_embeddings.device or mask_embedding.dtype != representative_embeddings.dtype:
+        raise ValueError('representative and MASK embeddings must share dtype/device')
+    noise = torch.as_tensor(noise_levels, device=representatives.device, dtype=torch.float32)
+    if noise.ndim < representatives.ndim:
+        noise = noise.reshape(noise.shape + (1,) * (representatives.ndim - noise.ndim))
+    try:
+        noise = torch.broadcast_to(noise, representatives.shape)
+    except RuntimeError as error:
+        raise ValueError('noise_levels are not broadcastable to representatives') from error
+    if not bool(torch.isfinite(noise).all()) or bool((noise < 0).any()) or bool((noise > 1).any()):
+        raise ValueError('noise_levels must be finite and lie in [0, 1]')
+    noise_in_embedding_dtype = noise.to(dtype=representative_embeddings.dtype)
+    combined = representative_embeddings + noise_in_embedding_dtype.unsqueeze(-1) * mask_embedding
+    return torch.where(noise.eq(0).unsqueeze(-1), representative_embeddings, combined)
+
+@dataclass(frozen=True, slots=True)
+class SemanticTreeAdvance:
+    representatives: torch.Tensor
+    ages: torch.Tensor
+    levels: torch.Tensor
+    noise_levels: torch.Tensor
+    committed: torch.Tensor
+
+@torch.no_grad()
+def advance_revisable_semantic_tree(current_representatives: torch.Tensor, x0_proposals: torch.Tensor, ages: torch.Tensor, tree: BalancedSemanticTree, steps: int=DEFAULT_SEMANTIC_TREE_STEPS, generator: torch.Generator | None=None) -> SemanticTreeAdvance:
+    if not isinstance(tree, BalancedSemanticTree):
+        raise TypeError('tree must be a BalancedSemanticTree')
+    tree._check_tokens(current_representatives, name='current_representatives')
+    tree._check_tokens(x0_proposals, name='x0_proposals')
+    if current_representatives.shape != x0_proposals.shape:
+        raise ValueError('current_representatives and x0_proposals must align')
+    if current_representatives.dtype != x0_proposals.dtype:
+        raise ValueError('current_representatives and x0_proposals must share dtype')
+    if not torch.is_tensor(ages) or ages.shape != current_representatives.shape:
+        raise ValueError('ages must align with representative tokens')
+    (validated_ages, steps) = _validate_ages(ages, steps=steps)
+    if validated_ages.device != current_representatives.device:
+        raise ValueError('ages and representative tokens must share a device')
+    if steps > tree.depth:
+        raise ValueError('steps cannot exceed tree depth')
+    unfinished = validated_ages < steps
+    next_ages = torch.where(unfinished, validated_ages + 1, validated_ages)
+    next_levels = semantic_tree_levels_for_ages(next_ages, steps=steps, depth=tree.depth)
+    proposal_representatives = tree.sample_representatives(x0_proposals, next_levels, generator=generator)
+    next_representatives = torch.where(unfinished, proposal_representatives, current_representatives)
+    return SemanticTreeAdvance(representatives=next_representatives, ages=next_ages.to(dtype=ages.dtype), levels=next_levels, noise_levels=semantic_tree_noise_for_ages(next_ages, steps=steps), committed=next_ages.eq(steps))

tree_metadata.json

@@ -0,0 +1,7 @@
+{
+  "filename": "pca_tree.pt",
+  "algorithm": "recursive_centered_pca_rank_median_v1",
+  "sha256": "4b7ee94fbb10fd4146a4ad167b595b037e3f591eaa12ecf3f0b411c82384a714",
+  "vocab_size": 131072,
+  "depth": 17
+}

utils.py

@@ -0,0 +1,27 @@
+# ADOBE CONFIDENTIAL
+# Copyright 2025 Adobe
+# All Rights Reserved.
+# NOTICE: All information contained herein is, and remains
+# the property of Adobe and its suppliers, if any. The intellectual
+# and technical concepts contained herein are proprietary to Adobe
+# and its suppliers and are protected by all applicable intellectual
+# property laws, including trade secret and copyright laws.
+# Dissemination of this information or reproduction of this material
+# is strictly forbidden unless prior written permission is obtained
+# from Adobe.
+
+from __future__ import annotations
+import torch
+
+def pad_along_last_dim(x, size):
+    pad_size = size - x.shape[-1]
+    if pad_size == 0:
+        return x
+    padding = torch.zeros(*x.shape[:-1], pad_size, dtype=x.dtype, device=x.device)
+    x_padded = torch.cat([x, padding], dim=-1)
+    return x_padded
+
+def maybe_truncate_last_dim(x, size):
+    if size == x.shape[-1]:
+        return x
+    return x[..., :size]

vision.py

@@ -0,0 +1,64 @@
+import math
+from types import SimpleNamespace
+import numpy as np
+import torch
+from torch import nn
+from timm.models.vision_transformer import Block
+from .radio_patch import ViTPatchGenerator
+from .radio_cls import ClsToken as _bundle_radio_cls
+
+class RadioBackbone(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.patch_generator = ViTPatchGenerator(16,1280,224,cls_token=True,max_input_dims=2048,num_cls_tokens=4,register_multiple=10)
+        self.blocks = nn.Sequential(*[Block(1280,16,mlp_ratio=4,qkv_bias=True) for _ in range(32)])
+    def forward(self, pixels):
+        return self.blocks(self.patch_generator(pixels))[:,10:]
+
+class RadioModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = RadioBackbone()
+        self.register_buffer('summary_idxs',torch.tensor([0,1,2],dtype=torch.int64))
+    def forward(self, pixels):
+        return self.model(pixels)
+
+class RadioWrapper(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.radio_model = RadioModel()
+    def forward(self, pixels):
+        return self.radio_model(pixels)
+
+class RadioVisionTower(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.vision_tower = RadioWrapper()
+        self.image_processor = RadioImageProcessor()
+        self.is_loaded = True
+        self.config = SimpleNamespace(hidden_size=1280,patch_size=16)
+    def forward(self, images):
+        reference = next(self.parameters())
+        if isinstance(images, (tuple,list)):
+            return torch.cat([self.vision_tower(x.to(reference)) for x in images], dim=1)
+        return self.vision_tower(images.to(reference))
+
+class RadioImageProcessor:
+    radio_downsample_ratio = 0.5
+    def __call__(self, images, return_tensors='pt'):
+        if len(images)!=1:
+            raise ValueError('Process one image at a time')
+        image = images[0].convert('RGB')
+        ph,pw = round(image.height/16+0.5), round(image.width/16+0.5)
+        factor=min(math.sqrt(12000/(ph*pw)),1.)
+        h,w=math.floor(factor*ph),math.floor(factor*pw)
+        if h*w<1024:
+            up=math.sqrt(1024/(h*w)); h,w=math.ceil(up*h),math.ceil(up*w)
+        if h%2:h=h+1 if (h+1)*w<=12000 else max(2,h-1)
+        if w%2:w=w+1 if h*(w+1)<=12000 else max(2,w-1)
+        pixels=torch.from_numpy(np.asarray(image,dtype=np.uint8).copy()).permute(2,0,1).unsqueeze(0).float()
+        if pixels.shape[-2:]!=(h*16,w*16):
+            pixels=torch.nn.functional.interpolate(pixels,size=(h*16,w*16),mode='bicubic',align_corners=False,antialias=True)
+        mean=torch.tensor([.48145466,.4578275,.40821073]).view(1,3,1,1)
+        std=torch.tensor([.26862954,.26130258,.27577711]).view(1,3,1,1)
+        return SimpleNamespace(pixel_values=(pixels/255.-mean)/std)

vq.py

@@ -0,0 +1,42 @@
+from types import SimpleNamespace
+import numpy as np
+import torch
+from torch import nn
+from .vq_model import Encoder, Decoder
+
+class Codebook(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embedding = nn.Embedding(131072, 256)
+    def get_codebook_entry(self, indices, shape=None):
+        result = self.embedding(indices)
+        if shape is not None:
+            result = result.view(shape).permute(0,3,1,2).contiguous()
+        return result
+    def forward(self, z):
+        # Identical argmax to the IBQ softmax; chunking bounds scratch memory.
+        b,c,h,w = z.shape
+        flat = z.permute(0,2,3,1).reshape(-1,c)
+        ids = torch.cat([(part @ self.embedding.weight.T).argmax(-1) for part in flat.split(256)])
+        quant = self.embedding(ids).reshape(b,h,w,c).permute(0,3,1,2).contiguous()
+        return quant, 0.0, (None,None,ids)
+
+class Emu3VQ(nn.Module):
+    def __init__(self):
+        super().__init__()
+        architecture = dict(double_z=False,z_channels=256,resolution=256,in_channels=3,out_ch=3,ch=256,ch_mult=[1,1,2,2,4],num_res_blocks=4,attn_resolutions=[16],dropout=0.0)
+        self.encoder = Encoder(**architecture)
+        self.decoder = Decoder(**architecture)
+        self.quantize = Codebook()
+        self.quant_conv = nn.Conv2d(256,256,1)
+        self.post_quant_conv = nn.Conv2d(256,256,1)
+        self.config = SimpleNamespace(num_vq_embeddings=131072,vq_embedding_dim=256)
+    def encode(self, images, mini_batch_size=None):
+        return self.quantize(self.quant_conv(self.encoder(images)))
+    def decode(self, quant, mini_batch_size=None):
+        return self.decoder(self.post_quant_conv(quant))
+
+class Emu3ImageProcessor:
+    def preprocess(self, image):
+        pixels = np.asarray(image.convert('RGB'), dtype=np.float32) / 127.5 - 1.0
+        return torch.from_numpy(pixels).permute(2,0,1).unsqueeze(0)

vq_model.py

@@ -0,0 +1,234 @@
+# pytorch_diffusion + derived encoder decoder
+from __future__ import annotations
+import torch
+import torch.nn as nn
+
+def nonlinearity(x):
+    return x * torch.sigmoid(x)
+
+def Normalize(in_channels):
+    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-06, affine=True)
+
+class Upsample(nn.Module):
+
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode='nearest')
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+class Downsample(nn.Module):
+
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x):
+        if self.with_conv:
+            pad = (0, 1, 0, 1)
+            x = torch.nn.functional.pad(x, pad, mode='constant', value=0)
+            x = self.conv(x)
+        else:
+            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+class ResnetBlock(nn.Module):
+
+    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, dropout, temb_channels=512):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if temb_channels > 0:
+            self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = torch.nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+            else:
+                self.nin_shortcut = torch.nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x, temb):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+        if temb is not None:
+            h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+        return x + h
+
+class AttnBlock(nn.Module):
+
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+        (b, c, h, w) = q.shape
+        q = q.reshape(b, c, h * w)
+        q = q.permute(0, 2, 1)
+        k = k.reshape(b, c, h * w)
+        w_ = torch.bmm(q, k)
+        w_ = w_ * int(c) ** -0.5
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+        v = v.reshape(b, c, h * w)
+        w_ = w_.permute(0, 2, 1)
+        h_ = torch.bmm(v, w_)
+        h_ = h_.reshape(b, c, h, w)
+        h_ = self.proj_out(h_)
+        return x + h_
+
+class Encoder(nn.Module):
+
+    def __init__(self, *, ch, out_ch, ch_mult=(1, 2, 4, 8), num_res_blocks, attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, resolution, z_channels, double_z=True, **ignore_kwargs):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.conv_in = torch.nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out, temb_channels=self.temb_ch, dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in, resamp_with_conv)
+                curr_res = curr_res // 2
+            self.down.append(down)
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in, temb_channels=self.temb_ch, dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in, temb_channels=self.temb_ch, dropout=dropout)
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in, 2 * z_channels if double_z else z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        temb = None
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+        h = hs[-1]
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+class Decoder(nn.Module):
+
+    def __init__(self, *, ch, out_ch, ch_mult=(1, 2, 4, 8), num_res_blocks, attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, resolution, z_channels, give_pre_end=False, **ignorekwargs):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.give_pre_end = give_pre_end
+        in_ch_mult = (1,) + tuple(ch_mult)
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+        self.conv_in = torch.nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in, temb_channels=self.temb_ch, dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in, temb_channels=self.temb_ch, dropout=dropout)
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for i_block in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out, temb_channels=self.temb_ch, dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, resamp_with_conv)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z, return_intermediate_feature=False):
+        self.last_z_shape = z.shape
+        temb = None
+        h = self.conv_in(z)
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+        im_feat_list = []
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h, temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+            im_feat_list.append(h)
+        if self.give_pre_end:
+            return h
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        if return_intermediate_feature:
+            return (h, im_feat_list)
+        return h