Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

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,300 @@
+---
+license: other
+license_name: nvidia-license
+license_link: LICENSE
+---
+
+# Model Overview
+
+## Description
+
+NL-Diffusion-Image introduces a new paradigm for high-resolution text-to-image generation via LLM based on masked discrete diffusion over tokenized image patches.
+Each image is encoded into a sequence of discrete tokens (using a 128K codebook/vocabulary), and generation proceeds through iterative parallel unmasking - similar to Diffusion LLMs.
+We finetune from [Nemotron-Labs-Diffusion](https://huggingface.co/nvidia/Nemotron-Labs-Diffusion-8B) and introduce 2 key components:
+
+* A token-editing mechanism that allows the model to revise already-unmasked tokens during inference.
+* Grouped Cross-Entropy (GCE) objective to handle large-vocabulary training efficiently.
+
+This model is ready for research or non-commercial evaluation.
+
+### 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
+
+## Use Case
+
+This model is intended for text-to-image generation.
+
+## Release Date
+
+Hugging Face: 07/01/2026 via [HuggingFace](https://huggingface.co/nvidia/NL-Diffusion-Image).
+
+## References
+
+* NL-Diffusion-Image Paper: Shufan Li et al., "Nemotron-Labs-Diffusion-Image: Advancing Masked Discrete Diffusion for High-Resolution Image Synthesis,".
+* Nemotron-Labs-Diffusion Paper: Yonggan Fu et al., "Nemotron-Labs-Diffusion: A Tri-Mode Language Model Unifying Autoregressive, Diffusion, and Self-Speculation Decoding".
+* Emu3.5 Paper: Emu3.5 team, "Emu3.5: Native Multimodal Models are World Learners".
+
+## Model Architecture
+
+**Architecture Type:** Neural Network <br>
+**Network Architecture:** Masked Diffusion Transformer, IBQ tokenizer for visual encoding/decoding <br>
+**Number of model parameters:** ~8B parameters <br>
+
+We encode 16x16 image patches using a pretrained discrete tokenizer from Emu3.5, with a codebook size of 128k token IDs.
+We expand the Nemotron-Labs-Diffusion vocabulary with a corresponding number of randomly-initialized embeddings, and fine-tune the model on a dataset of image/caption pairs.
+
+## Input
+
+**Input Type(s):** Text <br>
+**Input Format(s):** Characters <br>
+**Other Properties Related to Input:** Maximum prompt length is 900 tokens. <br>
+
+## Output
+
+**Output Type(s):** Images <br>
+**Output Format:** Tensor (3xHxW) <br>
+**Other Properties Related to Output:** The output represents an RGB image. <br>
+
+## Software Integration
+
+**Runtime Engine(s):**
+* PyTorch <br>
+
+
+**Supported Hardware Microarchitecture Compatibility:** <br>
+* NVIDIA Ampere <br>
+* NVIDIA Blackwell <br>
+* NVIDIA Jetson <br>
+* NVIDIA Hopper <br>
+* NVIDIA Lovelace <br>
+* NVIDIA Pascal <br>
+* NVIDIA Turing <br>
+* NVIDIA Volta <br>
+
+**[Preferred/Supported] Operating System(s):** <br>
+* Linux
+* Linux 4 Tegra
+* QNX
+* Windows
+
+The integration of foundation and fine-tuned models into AI systems requires additional testing using use-case-specific data to ensure safe and effective deployment.
+Following the V-model methodology, iterative testing and validation at both unit and system levels are essential to mitigate risks, meet technical and functional requirements, and ensure compliance with safety and ethical standards before deployment.
+
+This AI model can be embedded as an Application Programming Interface (API) call into the software environment described above.
+
+## Model Version(s)
+
+* NL-Diffusion-Image (8B).
+
+**Links:**
+
+* https://huggingface.co/nvidia/NL-Diffusion-Image
+
+# Training and Evaluation Datasets
+
+## Training Dataset
+
+**LAION-115M-Clean Recaptioned**
+
+**Data Modality:**
+* Multimodal (image, caption)
+
+**Image Training Data Size:**
+* 115M samples
+
+**Data Collection Method by dataset:**
+* Web scraping
+
+**Labeling Method by dataset:**
+* Subset of 8M images recaptioned using Qwen3-VL
+
+**MidJourney v6 520k Recaptioned**
+
+**Data Modality:**
+* Multimodal (image, caption)
+
+**Image Training Data Size:**
+* 520k samples
+
+**Data Collection Method by dataset:**
+* Automated
+
+**Labeling Method by dataset:**
+* Images recaptioned using Qwen3-VL
+
+**COYO700M Recaptioned**
+
+**Data Modality:**
+* Multimodal (image, caption)
+
+**Image Training Data Size:**
+* 700M samples
+
+**Data Collection Method by dataset:**
+* Automated
+
+**Labeling Method by dataset:**
+* Subset of 24M images recaptioned using Qwen3-VL
+
+**BLIP3o-60k Recaptioned**
+
+**Data Modality:**
+* Multimodal (image, caption)
+
+**Image Training Data Size:**
+* 520k samples
+
+**Data Collection Method by dataset:**
+* Automated
+
+**Labeling Method by dataset:**
+* Images recaptioned using Qwen3-VL
+
+## Evaluation Datasets
+
+**ImageNet**
+
+**Link:**
+* [ImageNet](https://www.image-net.org/)
+
+**Data Collection:**
+* Automated
+
+**Labeling Method:**
+* Manually-Collected
+
+**Training Images:**
+* 1,281,167
+
+**Validation Images:**
+* 50,000
+
+**GenEval**
+
+**Link:**
+* [GenEval](https://github.com/djghosh13/geneval)
+
+**Data Collection:**
+* Manually-Collected
+
+**Labeling Method:**
+* Manually-Collected
+
+**Captions/annotations:**
+* 553 samples
+
+**DPGBench**
+
+**Link:**
+* [DPGBench](https://github.com/TencentQQGYLab/ELLA)
+
+**Data Collection:**
+* Manually-Collected
+
+**Labeling Method:**
+* Manually-Collected
+
+**Captions/annotations:**
+* 1065 samples
+
+**MJHQ-30K**
+
+**Link:**
+* [MJHQ-30K](https://huggingface.co/datasets/playgroundai/MJHQ-30K/blob/main/README.md)
+
+**Data Collection:**
+* Manually-Collected
+
+**Labeling Method:**
+* Automated
+
+**CaptionsImages:**
+* 30k samples
+
+## GenEval Benchmark
+
+
+| Model      | Params | Single Object | Two Objects | Counting | Colors | Position | Color Attri. | Overall |
+|------------|---------|-------------|-----------|-------------|--------|----------|--------------|---------|
+| Qwen-Image-2507 | 20B | 0.99 | 0.92 | 0.89 | 0.88 | 0.76 | 0.77 | 0.87 |
+| Nemotron-Labs-Diffusion-Image | 8B | 0.98 | 0.93 | 0.83 | 0.94 | 0.88 | 0.82 | 0.90 |
+
+## Text-to-Image Generation Performance on DPG Benchmark and MJHQ-30k Dataset
+
+| Model      | Params | Codebook | DPG | MJHQ FID | MJHQ  HPSv3 |
+|------------|--------|-----------|----|----------|-------------|
+| MMaDa      | 8B     | 8,192 | 53.4 | 32.85 | 5.43 |
+| LaViDa-O   | 10B    | 8,192 | 81.8 | 6.68 | 8.81 |
+| Nemotron-Labs-Diffusion-Image | 8B | 131,072 | 85.2 | 6.46 | 9.61|
+| Nemotron-Labs-Diffusion-Image*</sup>  | 8B | 131,072 | 86.9 | 12.23 | 10.76 |
+
+<sup>*</sup> Finetuned on 6M synthetic data for better image quality
+
+
+## Inference
+
+**Acceleration Engine:** TensorRT, TensorRT-LLM <br>
+**Engine:** PyTorch <br>
+**Test Hardware:** NVIDIA Hopper (H100) <br>
+
+## Ethical Considerations
+
+NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications. Developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse.
+
+Please make sure you have proper rights and permissions for all input image and video content; if image or video includes people, personal health information, or intellectual property, the image or video generated will not blur or maintain proportions of image subjects included.
+
+For more detailed information on ethical considerations for this model, please see the Model Card++ Explainability, Bias, Safety & Security, and Privacy Subcards below.
+
+Please report model quality, risk, security vulnerabilities or NVIDIA AI Concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).
+
+### Bias
+
+Field                                                                                               |  Response
+:---------------------------------------------------------------------------------------------------|:---------------
+Participation considerations from adversely impacted groups [protected classes](https://www.senate.ca.gov/content/protected-classes) in model design and testing:  |  None
+Measures taken to mitigate against unwanted bias:                                                   |  None
+Bias Metric (If Measured): | None
+
+### Explainability
+
+Field                                                                                                  |  Response
+:------------------------------------------------------------------------------------------------------|:---------------------------------------------------------------------------------
+Intended Task/Domain:                                                                                  |  Text-to-image generation
+Model Type:                                                                                            |  Masked Diffusion Model
+Intended Users:                                                                                        |  Research
+Output:                                                                                                |  Images
+Describe how the model works:                                                                          |  The model takes a caption as input and generated an image.
+Name the adversely impacted groups this has been tested to deliver comparable outcomes regardless of:  |  Not Applicable
+Technical Limitations:                                                                                 |  The model generates images in a single resolution of 1024x1024 pixels.
+Verified to have met prescribed NVIDIA quality standards:  |  Yes
+Performance Metrics:                                                                                   |  GenEVal, DPG, MJHQ.
+Potential Known Risks:                                                                                 |  This model may not perform well on visual domains that are not represented in the training data. The generated images might fail to disambiguate differences in prompts that appear evident to humans. Domain-specific evaluation is required for the target application.
+Licensing:                                                                                             |  [NVIDIA Open Source License](https://huggingface.co/nvidia/NL-Diffusion-Image/raw/main/LICENSE)
+
+
+### Privacy
+
+Field                                                                                                                              |  Response
+:----------------------------------------------------------------------------------------------------------------------------------|:-----------------------------------------------
+Generatable or reverse engineerable personal data?                                                                               |  No
+Personal data used to create this model?                                                                                       |  No
+How often is dataset reviewed?                                                                                                     |  Before Every Release
+Is there provenance for all datasets used in training?                                                                                |  Yes
+Does data labeling (annotation, metadata) comply with privacy laws?                                                                |  Yes
+Is data compliant with data subject requests for data correction or removal, if such a request was made?                           |  Yes
+Was data from user interactions with the AI model (e.g. user input and prompts) used to train the model?                           | No
+Applicable Privacy Policy                                                                          | https://www.nvidia.com/en-us/about-nvidia/privacy-policy/
+
+### Safety
+
+Field                                               |  Response
+:---------------------------------------------------|:----------------------------------
+Model Application Field(s):                               |  Generation of images
+Describe the life critical impact (if present).   |  Not Applicable
+Use Case Restrictions:                              |  Research/evaluation only, non-commercial applications.
+Model and dataset restrictions:            |  The Principle of least privilege (PoLP) is applied limiting access for dataset generation and model development. Restrictions enforce dataset access during training, and dataset license constraints adhered to.

chat_template.jinja

@@ -0,0 +1,204 @@
+{% 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 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 %}
+        {{- content }}
+        {{- '<|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,448 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from transformers.utils import ModelOutput
+from dataclasses import dataclass
+from transformers.cache_utils import Cache, DynamicCache
+@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 .nemotron_diffusion_image_utils import maybe_truncate_last_dim, pad_along_last_dim
+
+
+def wte(model,x,t2i_inference=False,gen_shape=None,x_gen=None,inputs_embeds_curr=None,new_token_mask=None):
+
+    if t2i_inference:
+        assert x_gen is not None
+        if new_token_mask is None:
+            new_token_mask = x >= INT_MAX
+        # if x_gen is  None:
+        #     x_gen = x[new_token_mask] - OFFSET
+        # else:
+        #     x_gen = x_gen - OFFSET
+
+        gen_latents_comp_embeds = model.call_gen_embedding(x_gen,gen_shape)
+        if inputs_embeds_curr is None:
+            x_txt_only = x.clone()
+            
+            # replace consequtent [1] * 4096 to [1] * 1024
+
+            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 = 1_000_000
+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):
+    if t2i_inference:
+        if input_modality_indices is None:
+            input_modality_indices =modality_indices
+        output = model(None,input_embeddings=input_emnbeddings,modality_indices=input_modality_indices,output_hidden_states=True,past_key_values=past_key_values,
+                        is_training=False,
+                        overwrite_attn_impl='flash_attn'
+        )
+        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) # * 8 D
+        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:])
+            # N L 8 D
+        return gen_logits
+        
+        
+        final_logits = torch.zeros(*gen_logits.shape[:-1],OFFSET+gen_logits.shape[-1],dtype=output.logits.dtype,device=output.logits.device)
+        final_logits[:] = float('-inf')
+        final_logits[...,OFFSET:] = gen_logits
+        # breakpoint()
+        # inal_logits = torch.zeros(*hidden_states.shape[:-1],OFFSET+gen_logits.shape[-1],dtype=output.logits.dtype,device=output.logits.device)
+        
+        # final_logits = final_logits + float('-inf')
+        # final_logits[...,:output.logits.shape[-1]] = output.logits
+        # final_logits[modality_indices] = float('-inf')
+        # local = final_logits[modality_indices]
+
+        # local[...,OFFSET:] = gen_logits
+        # final_logits[modality_indices] = local
+
+        logits = final_logits
+        return logits
+    else:
+        modality_indices = torch.zeros(input_emnbeddings.shape[:-1],device=input_emnbeddings.device,dtype=torch.bool)
+        logits = model(None,input_embeddings=input_emnbeddings,modality_indices=modality_indices,past_key_values=past_key_values).logits
+    return logits
+
+def add_gumbel_noise(logits, temperature):
+    '''
+    The Gumbel max is a method for sampling categorical distributions.
+    According to arXiv:2409.02908, for MDM, low-precision Gumbel Max improves perplexity score but reduces generation quality.
+    Thus, we use float64.
+    '''
+    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.to(torch.float64), dim=-1)
+        p = F.softmax(logits, dim=-1)
+        x0_p = torch.squeeze(
+            torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)), -1) # b, l
+    elif remasking == 'top_p_margin':
+        # Compute probabilities
+        p = F.softmax(logits, dim=-1)                       # (B, L, V)
+        # Top-2 per position
+        top2 = torch.topk(p, k=2, dim=-1).values            # (B, L, 2)
+        margin = top2[..., 0] - top2[..., 1]                # (B, L)
+
+        # Normalize margin to [0,1] over MASKED positions per row
+        plus_inf  = torch.full_like(margin, float('inf'))
+        minus_inf = torch.full_like(margin, float('-inf'))
+        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)  # (B, 1)
+        row_max = masked_for_max.amax(dim=1, keepdim=True)  # (B, 1)
+        denom = (row_max - row_min)
+
+        # If denom==0 (all equal), set normalized=1 on masked; 0 elsewhere by default
+        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  # ∈ [0,1] on masked positions
+    elif remasking == 'random':
+        x0_p = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device)
+    else:
+        raise NotImplementedError(remasking)
+    
+    # Calculate negative entropy if requested
+    if neg_entropy:
+        # p = F.softmax(logits.to(torch.float64), dim=-1)
+        p = F.softmax(logits, dim=-1)
+        epsilon = 1e-10
+        log_probs = torch.log(p + epsilon)
+        confidence_scores = torch.sum(p * log_probs, dim=-1)  # negative entropy per position
+    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)
+    # print(f'confidence: {confidence}')
+    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,
+        is_training=False,
+        overwrite_attn_impl='flash_attn',
+        # overwrite_attn_impl='flash_attn',
+        # overwrite_block_mask='full',
+        **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.,
+    remasking='low_confidence',
+    mask_id=126336,
+    threshold=None,
+    factor=None,
+    shift_logits=False,
+    neg_entropy=False,
+    causal_context=False,
+    eos_token_id=None,
+    max_thinking_tokens=None,
+    end_think_token_id=None,
+):
+    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
+        # Keep prefix slots so block slicing by prompt_len stays identical.
+        x_accum = torch.full((B, prompt_len), mask_id, dtype=token_dtype, device=token_device)
+
+    assert gen_length % block_length == 0
+    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
+
+    # Compute KV cache for the prompt initially
+    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
+
+    # Causal prefill: next token from last position (same as linear_spec_generate).
+    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)
+
+    # For dream_style: store the "next token logit" of the context
+    next_logits_context = None
+    if dream_style:
+        next_logits_context = output.logits[:, -1:, :]  # (B, 1, V)
+
+    for num_block in range(num_blocks):
+        # Create a new block with mask tokens; under causal context, seed position 0
+        # with the next-token prediction from the previous causal forward (prefill or
+        # post-block encode), matching linear_spec_generate.
+        mask_block = torch.ones(
+            (B, block_length),
+            dtype=token_dtype,
+            device=token_device,
+        ) * mask_id
+        if causal_context:
+            mask_block[:, 0] = next_token[:, 0]
+
+        # Append the block of masks
+        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)
+
+        # ---- thinking budget enforcement ----
+        # If we've generated >= max_thinking_tokens without a </think>, inject one.
+        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
+
+        # Build the initial mask for this block
+        mask_block_idx0 = (x_accum[:, block_slice] == mask_id)  # (B, Lb)
+
+        # Precompute the transfer schedule for this block
+        if dream_style:
+            # masked positions only (position 0 may be causal-seeded, not mask_id)
+            schedule_mask = mask_block_idx0
+        else:
+            schedule_mask = mask_block_idx0
+
+        num_transfer_tokens = get_num_transfer_tokens(schedule_mask, steps_per_block)  # (B, steps)
+
+        # Denoise the current block
+        for i in range(steps_per_block):
+            mask_block_idx = (x_accum[:, block_slice] == mask_id)  # (B, Lb)
+            if mask_block_idx.sum() == 0:
+                break
+
+            nfe += 1
+
+            # Forward only the current noisy block using cached context
+            logits_block = simple_fwd(model,
+                x_accum[:, block_slice],
+                past_key_values=past_key_values,
+                use_cache=False
+            ).logits
+
+            if dream_style:
+                # Align logits so that each masked position has a predictor:
+                # prepend context-next logit, then use logits_block[:-1]
+                if block_length == 1:
+                    logits_use = next_logits_context              # (B, 1, V)
+                else:
+                    logits_use = torch.cat(
+                        [next_logits_context, logits_block[:, :-1, :]],
+                        dim=1
+                    )  # (B, Lb, V)
+
+                mask_use = mask_block_idx                        # (B, Lb)
+                x_use   = x_accum[:, block_slice]                # (B, Lb)
+
+                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:
+                # non-AR (same-position) case
+                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]              # (B, Lb)
+                eos_mask = (block_tokens == eos_token_id)           # (B, Lb)
+                any_eos = eos_mask.any(dim=1)                       # (B,)
+                if any_eos.any():
+                    after_eos = eos_mask.cumsum(dim=1).bool()       # (B, Lb)
+                    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
+
+        # after block is fully denoised, update KV cache
+        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
+            # Next block's first position = greedy/sampled next token from this causal encode
+            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:
+            # refresh context-next logit for the next block
+            next_logits_context = output.logits[:, -1:, :]  # (B, 1, V)
+
+        if eos_token_id is not None:
+            gen_so_far = x_accum[:, prompt_len:]                        # (B, gen_len_so_far)
+            is_eos = (gen_so_far == eos_token_id)                       # (B, gen_len_so_far)
+            has_eos = is_eos.any(dim=1)                                 # (B,)
+            if has_eos.all():
+                first_eos_pos = is_eos.to(torch.int64).argmax(dim=1)    # (B,)
+                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,131 @@
+{
+  "ada_dlm_loss_ratio": null,
+  "ada_perm_ratio_global": null,
+  "ada_perm_ratio_per_block": null,
+  "adaptive_mask_rate": false,
+  "add_faster_video": false,
+  "add_time_instruction": false,
+  "ar_loss_weight": 1.0,
+  "architectures": [
+    "NemotronLabsDiffusionImageForMaskedDiffusion"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attn_implementation": null,
+  "auto_map": {
+    "AutoConfig": "configuration_nemotron_labs_diffusion_image.NemotronLabsDiffusionImageConfig",
+    "AutoModel": "modeling_nemotron_labs_diffusion_image.NemotronLabsDiffusionImageForMaskedDiffusion",
+    "AutoModelForCausalLM": "modeling_nemotron_labs_diffusion_image.NemotronLabsDiffusionImageForMaskedDiffusion"
+  },
+  "block_size": 32,
+  "bos_token_id": 1,
+  "d_model": 4096,
+  "d_model_gen": 4096,
+  "diff_loss_weight": 1,
+  "dlm_arch": "encoder",
+  "dlm_loss_weight": null,
+  "dlm_paradigm": "bidirectional",
+  "dlm_type": "llada",
+  "downsample": true,
+  "dp_varying_mask_ratio": false,
+  "dtype": "bfloat16",
+  "dual_tower": true,
+  "dual_tower_layers": 16,
+  "enable_self_spec": false,
+  "enforce_mask": false,
+  "eos_token_id": 11,
+  "faster_token_stride": 10,
+  "flip_ratio": 0.2,
+  "force_sample": false,
+  "gen_edit_loss_weight": 0.2,
+  "global_loss_avg": false,
+  "group_ce_weight": "{16384:0.2}",
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "image_aspect_ratio": "t2i_only",
+  "image_crop_resolution": null,
+  "image_grid_pinpoints": [
+    [384, 768],
+    [768, 384],
+    [768, 768],
+    [1152, 384],
+    [384, 1152]
+  ],
+  "image_split_resolution": null,
+  "include_bias": false,
+  "initializer_range": 0.02,
+  "intermediate_size": 14336,
+  "mask_token_id": 100,
+  "max_position_embeddings": 262144,
+  "mlp_bias": false,
+  "mlp_hidden_size_gen": 14336,
+  "mm_hidden_size": 3584,
+  "mm_newline_position": "grid",
+  "mm_patch_merge_type": "spatial_unpad",
+  "mm_pooler_ratio": 2,
+  "mm_projector_lr": null,
+  "mm_projector_type": "mlp2x_gelu",
+  "mm_resampler_type": "none",
+  "mm_spatial_pool_mode": "conv",
+  "mm_spatial_pool_out_channels": 3584,
+  "mm_spatial_pool_stride": 1,
+  "mm_tunable_parts": "mm_language_model",
+  "mm_use_im_patch_token": false,
+  "mm_use_im_start_end": false,
+  "mm_vision_select_feature": "patch",
+  "mm_vision_select_layer": -2,
+  "mm_vision_tower": null,
+  "mm_vision_tower_lr": null,
+  "mm_vqvae": "emu3_vqvae",
+  "model_type": "nemotron_labs_diffusion_image",
+  "multi_sampling": null,
+  "num_ar_layers": 0,
+  "num_attention_heads": 32,
+  "num_diffusion_layers": 0,
+  "num_hidden_layers": 34,
+  "num_key_value_heads": 8,
+  "num_skip_loss_tokens": 0,
+  "pad_token_id": 11,
+  "pos_skipping_range": 4096,
+  "prefix_ratio": 0.8,
+  "prompt_drop_rate": 0.1,
+  "random_length_prob": 0,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "beta_fast": 32.0,
+    "beta_slow": 1.0,
+    "factor": 16.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",
+    "type": "yarn"
+  },
+  "rope_scaling": {
+    "beta_fast": 32.0,
+    "beta_slow": 1.0,
+    "factor": 16.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",
+    "type": "yarn"
+  },
+  "rope_theta": 1000000.0,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "tok_mask_half_life_ratio": null,
+  "tokenizer_model_max_length": 4096,
+  "tokenizer_padding_side": "right",
+  "transformers_version": "4.57.6",
+  "use_cache": false,
+  "use_mm_proj": false,
+  "use_pos_skipping": false,
+  "vision_tower_pretrained": null,
+  "vocab_size": 132101
+}

configuration_ministral_dlm.py

@@ -0,0 +1,264 @@
+# coding=utf-8
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Ministral DLM model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+try:
+    from transformers.modeling_rope_utils import rope_config_validation
+except ImportError:
+    rope_config_validation = None
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class MinistralDLMConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Ministral3Model`] for diffusion language models.
+    It is used to instantiate a Ministral model according to the specified arguments, defining the model architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 131072):
+            Vocabulary size of the Ministral model.
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 14336):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 34):
+            Number of hidden layers in the Transformer decoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer.
+        num_key_value_heads (`int`, *optional*, defaults to 8):
+            Number of key_value heads for Grouped Query Attention.
+        head_dim (`int`, *optional*, defaults to 128):
+            The attention head dimension.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function.
+        max_position_embeddings (`int`, *optional*, defaults to 262144):
+            The maximum sequence length.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-05):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied.
+        rope_theta (`float`, *optional*, defaults to 1000000.0):
+            The base period of the RoPE embeddings.
+        rope_parameters (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings.
+            Default uses YaRN scaling with factor=16, original_max_position_embeddings=16384.
+        attention_bias (`bool`, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        mlp_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use a bias in up_proj, down_proj and gate_proj layers.
+        sliding_window (`int`, *optional*, defaults to None):
+            Sliding window attention size.
+        mask_token_id (`int`, *optional*, defaults to -1):
+            Token ID for masking in diffusion.
+        dlm_type (`str`, *optional*, defaults to 'llada'):
+            Type of diffusion language model ('llada', 'dream').
+        random_length_prob (`float`, *optional*):
+            Probability of using random lengths during training.
+        num_ar_layers (`int`, *optional*, defaults to 0):
+            Number of autoregressive layers.
+        num_diffusion_layers (`int`, *optional*, defaults to 0):
+            Number of diffusion layers.
+        diff_loss_weight (`float`, *optional*, defaults to 1):
+            Weight for diffusion loss.
+        enforce_mask (`bool`, *optional*, defaults to False):
+            Whether to enforce masking.
+        prefix_ratio (`float`, *optional*, defaults to 0.8):
+            Ratio for prefix in prefix_bidirectional mode.
+        dlm_paradigm (`str`, *optional*, defaults to 'bidirectional'):
+            Paradigm for diffusion ('bidirectional', 'autoregressive', 'prefix_bidirectional', 'efficient_block_diff', 'block_diff', 'sbd_block_diff').
+        dlm_arch (`str`, *optional*, defaults to 'encoder'):
+            Architecture type ('encoder', 'encoder_decoder').
+        block_size (`int`, *optional*, defaults to 32):
+            Block size for block diffusion paradigms.
+        tok_mask_half_life_ratio (`float`, *optional*):
+            Half-life ratio for token masking.
+        adaptive_mask_rate (`bool`, *optional*, defaults to False):
+            Whether to use adaptive mask rate.
+        multi_sampling (`int`, *optional*):
+            Number of samples for multi-sampling.
+        num_skip_loss_tokens (`int`, *optional*, defaults to 0):
+            Number of tokens to skip in loss calculation.
+        dlm_loss_weight (`float`, *optional*):
+            Weight for diffusion LM loss.
+        ar_loss_weight (`float`, *optional*, defaults to 1.0):
+            Weight for autoregressive loss in sbd_block_diff paradigm. Use 10000 to only use AR loss.
+        global_loss_avg (`bool`, *optional*, defaults to False):
+            Whether to use global loss average.
+        dp_varying_mask_ratio (`bool`, *optional*, defaults to False):
+            Whether to use varying mask ratio for each DP rank during sampling.
+        ada_perm_ratio_per_block (`float`, *optional*):
+            Adaptive permutation ratio for each block.
+        ada_perm_ratio_global (`float`, *optional*):
+            Adaptive permutation ratio for global.
+        enable_self_spec (`bool`, *optional*, defaults to `False`):
+            Force MinistralFlexAttention for all paradigms (including bidirectional/autoregressive).
+            Required for self speculative generation; leave False for standard eval to use faster SDPA kernels.
+    """
+
+    model_type = "ministral_dlm"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    # Default tensor parallel plan for base model `Ministral`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size=131072,
+        hidden_size=4096,
+        intermediate_size=14336,
+        num_hidden_layers=34,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        head_dim=128,
+        hidden_act="silu",
+        max_position_embeddings=262144,
+        initializer_range=0.02,
+        rms_norm_eps=1e-05,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        rope_theta=1000000.0,
+        rope_parameters=None,
+        rope_scaling=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        mlp_bias=False,
+        sliding_window=None,
+        attn_implementation="sdpa",
+        mask_token_id=None,
+        dlm_type='llada',
+        random_length_prob=None,
+        num_ar_layers=0,
+        num_diffusion_layers=0,
+        diff_loss_weight=1,
+        enforce_mask=False,
+        prefix_ratio=0.8,
+        dlm_paradigm='bidirectional',
+        dlm_arch='encoder',
+        block_size=32,
+        tok_mask_half_life_ratio=None,
+        adaptive_mask_rate=False,
+        multi_sampling=None,
+        num_skip_loss_tokens=0,
+        dlm_loss_weight=None,
+        ar_loss_weight=1.0,
+        global_loss_avg=False,
+        dp_varying_mask_ratio=False,
+        ada_perm_ratio_per_block=None,
+        ada_perm_ratio_global=None,
+        ada_dlm_loss_ratio=None,
+        enable_self_spec=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        if rope_parameters is None and rope_scaling is not None:
+            rope_parameters = dict(rope_scaling)
+        # llama_4_scaling_beta is used directly by the attention layer; do not strip it.
+        self.rope_parameters = rope_parameters
+        self.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.mlp_bias = mlp_bias
+        self.sliding_window = sliding_window
+        
+        self.attn_implementation = attn_implementation
+        
+        self.mask_token_id = mask_token_id
+        self.dlm_type = dlm_type
+        self.random_length_prob = random_length_prob
+        self.num_ar_layers = num_ar_layers
+        self.num_diffusion_layers = num_diffusion_layers
+        self.diff_loss_weight = diff_loss_weight
+        self.enforce_mask = enforce_mask
+        self.prefix_ratio = prefix_ratio
+        self.dlm_paradigm = dlm_paradigm
+        self.dlm_arch = dlm_arch
+        self.block_size = block_size
+        self.tok_mask_half_life_ratio = tok_mask_half_life_ratio
+        self.adaptive_mask_rate = adaptive_mask_rate
+        self.multi_sampling = multi_sampling
+        self.num_skip_loss_tokens = num_skip_loss_tokens
+        self.dlm_loss_weight = dlm_loss_weight
+        self.ar_loss_weight = ar_loss_weight
+        self.global_loss_avg = global_loss_avg
+        self.dp_varying_mask_ratio = dp_varying_mask_ratio
+        self.ada_perm_ratio_per_block = ada_perm_ratio_per_block
+        self.ada_perm_ratio_global = ada_perm_ratio_global
+        self.ada_dlm_loss_ratio = ada_dlm_loss_ratio
+        self.enable_self_spec = enable_self_spec
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+        # Transformers>=4.57 expects standardized/validated rope_parameters.
+        if hasattr(self, "standardize_rope_params"):
+            self.standardize_rope_params()
+        if hasattr(self, "validate_rope"):
+            self.validate_rope()
+        elif rope_config_validation is not None:
+            rope_config_validation(self)
+
+
+__all__ = ["MinistralDLMConfig"]
+

configuration_nemotron_labs_diffusion_image.py

@@ -0,0 +1,33 @@
+from .configuration_ministral_dlm import MinistralDLMConfig
+
+
+class NemotronLabsDiffusionImageConfig(MinistralDLMConfig):
+    model_type = "nemotron_labs_diffusion_image"
+
+    def __init__(
+        self,
+        max_position_embeddings=262144,
+        rope_parameters=None,
+        rope_scaling=None,
+        **kwargs,
+    ):
+        # Newer transformers standardize RoPE params during PretrainedConfig init.
+        # Make these attributes available early and normalize legacy rope_scaling.
+        if rope_parameters is None and rope_scaling is not None:
+            rope_parameters = dict(rope_scaling)
+
+        self.max_position_embeddings = max_position_embeddings
+        if rope_parameters is not None:
+            self.rope_parameters = rope_parameters
+        if rope_scaling is not None:
+            self.rope_scaling = rope_scaling
+
+        super().__init__(
+            max_position_embeddings=max_position_embeddings,
+            rope_parameters=rope_parameters,
+            rope_scaling=rope_scaling,
+            **kwargs,
+        )
+
+        if getattr(self, "rope_parameters", None) is None and getattr(self, "rope_scaling", None) is not None:
+            self.rope_parameters = dict(self.rope_scaling)

demo_inference_release.py

@@ -0,0 +1,249 @@
+import argparse
+import os
+from pathlib import Path
+
+import torch
+from transformers import AutoModelForCausalLM, PreTrainedTokenizerFast
+
+# Keep behavior aligned with the existing demo entrypoints.
+os.environ.setdefault("DEBUG_FIX_PADDING", "1")
+os.environ.setdefault("NOT_ALWASY_DO_2DPOOL", "1")
+
+DEFAULT_CKPT = "/lustre/fsw/portfolios/nvr/users/shufanl/projects/ministral-output/nemotron-labs-diffuion-image-8B-release"
+DEFAULT_PROMPT = (
+    "A full-body shot of hyper-realistic female cyborg, human facial skin seamlessly integrated with a glossy white mechanical head shell. "
+    "Features a realistic human ear, blue eyes. bright, outdoor, background with blue sky, illuminated by striking bright white studio lighting, "
+    "casting soft shadows. Cyberpunk aesthetic, high-tech minimalism, shot on 85mm lens, photorealistic, Unreal Engine 5 render, intricately detailed, "
+    "8k resolution, high dynamic range, chest with whit armor plate, cute, beautiful, sexy, glossy surface, reflective, Artstation, pixiv, no hair, "
+    "3D render, stylized eyesz"
+)
+DEFAULT_OUTPUT = "/lustre/fsw/portfolios/nvr/users/shufanl/code/LaVida-O/outputs/demo_inference_release.webp"
+
+SCHEDULE_CHOICES = ["shift"]
+CONFIDENCE_POLICY_CHOICES = ["mask_git", "mmada", "stratified"]
+SCHEDULE_TEMP_CHOICES = ["linear", "cosine2", "shift", "exp"]
+RESOLUTION_CHOICES = [256, 512, 1024]
+
+# Match the default generation setup in gradio_t2i_demo.py.
+DEFAULT_GENERATION_CONFIG = {
+    "guidance_scale": 5.0,
+    "n_steps": 64,
+    "shift": 5,
+    "schedule": "shift",
+    "alg_temp": 1.0,
+    "dynamic_temperature": False,
+    "min_temperature": 0.01,
+    "schedule_temp": "linear",
+    "temperature": 0.86,
+    "confidence_policy": "mmada",
+    "micro_cond": "ORIGINAL WIDTH : 1024; ORIGINAL HEIGHT : 1024; TOP : 0; LEFT : 0; SCORE : 6.520; HPS: 3.220",
+    "template": "Generate an image with the caption:\n <prompt>",
+    "edit_threshold": 0.6,
+}
+
+
+def load_release_model_and_tokenizer(model_path: str, device: str):
+    tokenizer = PreTrainedTokenizerFast.from_pretrained(model_path)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=False,
+    )
+    model.to(device)
+    model.eval()
+    model.requires_grad_(False)
+
+    return tokenizer, model
+
+
+def n_tokens_from_resolution(image_resolution: int) -> int:
+    return (image_resolution // 16) * (image_resolution // 16)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Single-image LaVida-O text-to-image inference using the release package defaults.",
+        formatter_class=argparse.RawTextHelpFormatter,
+    )
+    parser.add_argument(
+        "--pretrained",
+        type=str,
+        default=DEFAULT_CKPT,
+        help="Path to the model directory.\nChoices: any local HF-style checkpoint or release directory.",
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default=DEFAULT_PROMPT,
+        help="Prompt text for generation.\nChoices: any text string.",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default=DEFAULT_OUTPUT,
+        help="Output image path.\nChoices: any writable file path; extension should match a Pillow-supported format such as .webp, .png, or .jpg.",
+    )
+    parser.add_argument(
+        "--image-resolution",
+        type=int,
+        choices=RESOLUTION_CHOICES,
+        default=1024,
+        help="Output resolution in pixels.\nChoices: 256, 512, 1024.",
+    )
+    parser.add_argument(
+        "--guidance-scale",
+        type=float,
+        default=DEFAULT_GENERATION_CONFIG["guidance_scale"],
+        help="Classifier-free guidance strength.\nChoices: any positive float; default is the inline release demo setting.",
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=DEFAULT_GENERATION_CONFIG["temperature"],
+        help="Sampling temperature for token draws.\nChoices: any positive float; lower is more conservative.",
+    )
+    parser.add_argument(
+        "--n-steps",
+        type=int,
+        default=DEFAULT_GENERATION_CONFIG["n_steps"],
+        help="Number of denoising steps.\nChoices: any positive integer; default is the inline release demo setting.",
+    )
+    parser.add_argument(
+        "--schedule",
+        type=str,
+        choices=SCHEDULE_CHOICES,
+        default=DEFAULT_GENERATION_CONFIG["schedule"],
+        help="Token transfer schedule.\nChoices: shift.",
+    )
+    parser.add_argument(
+        "--shift",
+        type=int,
+        default=DEFAULT_GENERATION_CONFIG["shift"],
+        help="Shift parameter used by the shift schedule.\nChoices: any non-negative integer; default is the inline release demo setting.",
+    )
+    parser.add_argument(
+        "--confidence-policy",
+        type=str,
+        choices=CONFIDENCE_POLICY_CHOICES,
+        default=DEFAULT_GENERATION_CONFIG["confidence_policy"],
+        help="Policy for selecting which masked tokens to reveal next.\nChoices: mask_git, mmada, stratified.",
+    )
+    parser.add_argument(
+        "--schedule-temp",
+        type=str,
+        choices=SCHEDULE_TEMP_CHOICES,
+        default=DEFAULT_GENERATION_CONFIG["schedule_temp"],
+        help="Temperature schedule shape across denoising steps.\nChoices: linear, cosine2, shift, exp.",
+    )
+    parser.add_argument(
+        "--alg-temp",
+        type=float,
+        default=DEFAULT_GENERATION_CONFIG["alg_temp"],
+        help="Confidence-ranking temperature used by the reveal policy.\nChoices: any positive float; default is the inline release demo setting.",
+    )
+    parser.add_argument(
+        "--dynamic-temperature",
+        action="store_true",
+        default=DEFAULT_GENERATION_CONFIG["dynamic_temperature"],
+        help="Enable dynamic temperature scaling over time.\nChoices: enabled with --dynamic-temperature, disabled with --no-dynamic-temperature.",
+    )
+    parser.add_argument(
+        "--no-dynamic-temperature",
+        dest="dynamic_temperature",
+        action="store_false",
+        help="Disable dynamic temperature scaling.",
+    )
+    parser.add_argument(
+        "--edit-threshold",
+        type=float,
+        default=DEFAULT_GENERATION_CONFIG["edit_threshold"],
+        help="Post-sampling token editing threshold.\nChoices: any float in practice; use -1 to disable edit-based replacement as in the Gradio demo semantics.",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed.\nChoices: any integer; use -1 to sample a fresh random seed.",
+    )
+    parser.add_argument(
+        "--micro-cond",
+        type=str,
+        default=DEFAULT_GENERATION_CONFIG["micro_cond"],
+        help="Micro-conditioning string injected into the prompt template.\nChoices: any string matching the model's expected metadata style.",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda",
+        help="Torch device for model execution.\nChoices: typically cuda or cpu; cuda is expected for practical inference speed.",
+    )
+    parser.add_argument(
+        "--use-cache",
+        action="store_true",
+        help="Enable KV-cache prefill path during denoising.\nChoices: set flag to enable, omit to disable.",
+    )
+    parser.add_argument(
+        "--is-legacy",
+        action="store_true",
+        help="Use legacy generation behavior expected by older checkpoints.\nChoices: set flag to enable, omit to disable.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    gen_cfg = dict(DEFAULT_GENERATION_CONFIG)
+    gen_cfg.update(
+        micro_cond=args.micro_cond,
+        guidance_scale=args.guidance_scale,
+        temperature=args.temperature,
+        edit_threshold=args.edit_threshold,
+        n_steps=args.n_steps,
+        schedule=args.schedule,
+        shift=args.shift,
+        confidence_policy=args.confidence_policy,
+        schedule_temp=args.schedule_temp,
+        alg_temp=args.alg_temp,
+        dynamic_temperature=args.dynamic_temperature,
+        block_policy=2,
+    )
+
+    if args.seed < 0:
+        args.seed = int(torch.seed() % (2**31 - 1))
+    torch.manual_seed(args.seed)
+
+    tokenizer, model = load_release_model_and_tokenizer(args.pretrained, args.device)
+    model.config.dlm_paradigm = "bidirectional"
+
+    with torch.no_grad():
+        with torch.inference_mode():
+            image = model.text_to_image(
+                args.prompt,
+                tokenizer=tokenizer,
+                **gen_cfg,
+                image_resolution=args.image_resolution,
+                n_tokens=n_tokens_from_resolution(args.image_resolution),
+                is_legacy=args.is_legacy,
+                use_cache=args.use_cache,
+                disable_tqdm=False,
+                return_intermediate_steps=False,
+            )
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    image.save(output_path)
+
+    print(f"Saved image to {output_path}")
+    print(f"Seed: {args.seed}")
+    print(f"Resolution: {args.image_resolution}")
+    print(f"Checkpoint: {args.pretrained}")
+
+
+if __name__ == "__main__":
+    main()

emu3_vqvae/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

emu3_vqvae/README.md

@@ -0,0 +1,266 @@
+---
+license: apache-2.0
+---
+<div align='center'>
+<h1>Emu3.5: Native Multimodal Models are World Learners</h1>
+
+Emu3.5 Team, BAAI
+
+[Project Page](https://emu.world/pages/web/landingPage) | [🤗HF Models](https://huggingface.co/collections/BAAI/emu35) | [Paper](https://arxiv.org/pdf/2510.26583) | [App](https://emu.world/pages/web/home?route=index)
+</div>
+
+
+> 🔔 **Latest**: Emu3.5 Web & Mobile Apps and vLLM offline inference are live — see [🔥 News](#news) for details.
+
+
+<div align='center'>
+<img src="https://github.com/baaivision/Emu3.5/blob/main/assets/arch.png?raw=True" class="interpolation-image" alt="arch." height="100%" width="100%" />
+</div>
+
+
+<div align='center'>
+<img src="https://github.com/baaivision/Emu3.5/blob/main/assets/co.png?raw=True" class="interpolation-image" alt="arch." height="90%" width="90%" />
+</div>
+
+
+|  🔹 | **Core Concept**                         | **Description**                                                                                                                            |
+| :-: | :--------------------------------------- | :----------------------------------------------------------------------------------------------------------------------------------------- |
+|  🧠 | **Unified World Modeling**               | Predicts the **next state jointly across vision and language**, enabling coherent **world modeling** and **generation**.              |
+|  🧩 | **End-to-End Pretraining**               | Trained with a **unified next-token prediction** objective over **interleaved vision–language sequences**.                                 |
+|  📚 | **Over 10T+ Multimodal Tokens**               | Pre-trained on **over 10 trillion interleaved tokens** from **video frames** and **transcripts**, capturing **spatiotemporal structure**.       |
+|  🔄 | **Native Multimodal I/O**                | Processes and generates **interleaved visual–text sequences** without **modality adapters** or **task-specific heads**.                    |
+|  🎯 | **RL Post-Training**                     | Large-scale **reinforcement learning** enhances **reasoning**, **compositionality**, and **generation quality**.                           |
+|  ⚡  | **Discrete Diffusion Adaptation (DiDA)** | Converts **sequential decoding → bidirectional parallel prediction**, achieving **≈20× faster inference without performance loss**.      |
+| 🖼️ | **Versatile Generation**                 | Excels in **long-horizon vision–language generation**, **any-to-image (X2I)** synthesis, and **text-rich image creation**.                 |
+|  🌐 | **Generalizable World Modeling**         | Enables **spatiotemporally consistent world exploration**, and **open-world embodied manipulation** across diverse scenarios.          |
+|  🏆 | **Performance Benchmark**                | Matches **Gemini 2.5 Flash Image (Nano Banana)** on **image generation/editing**, and **outperforms** on **interleaved generation tasks**. |
+
+
+<a id="news"></a>
+
+## 🔥 News
+
+- **2025-11-28 · 🌐 Emu3.5 Web & Mobile Apps Live** — Official product experience is **now available** on the web at [zh.emu.world](https://zh.emu.world) (Mainland China) and [emu.world](https://emu.world) (global) 🎉 The new homepage highlights featured cases and a “Get Started” entry, while the workspace and mobile apps bring together creation, inspiration feed, history, profile, and language switch across web, Android APK, and H5. *([See more details](#official-web--mobile-apps) below.)*
+- **2025-11-19 · 🚀 vLLM Offline Inference Released** — Meet `inference_vllm.py` with a new cond/uncond batch scheduler, delivering **4–5× faster end-to-end generation** on vLLM 0.11.0 across Emu3.5 tasks. Jump to [#Run Inference with vLLM](#run-inference-with-vllm) for setup guidance and see PR [#47](https://github.com/baaivision/Emu3.5/pull/47) for full details.
+- **2025-11-17 · 🎛️ Gradio Demo (Transformers Backend)** — Introduced `gradio_demo_image.py` and `gradio_demo_interleave.py` presets for the standard Transformers runtime, providing turnkey T2I/X2I and interleaved generation experiences with streaming output. Try the commands in [#Gradio Demo](#3-gradio-demo) to launch both UIs locally.
+
+## Table of Contents
+
+1. [Model & Weights](#1-model--weights)
+2. [Quick Start](#2-quick-start)
+3. [Gradio Demo](#3-gradio-demo)
+4. [Schedule](#4-schedule)
+5. [Citation](#5-citation)
+
+## 1. Model & Weights
+
+| Model name               | HF Weight |
+| ------------------------ | --------- |
+| Emu3.5               | [🤗 HF link](https://huggingface.co/BAAI/Emu3.5/tree/main) |
+| Emu3.5-Image                | [🤗 HF link](https://huggingface.co/BAAI/Emu3.5-Image/tree/main) |
+| Emu3.5-VisionTokenizer     | [🤗 HF link](https://huggingface.co/BAAI/Emu3.5-VisionTokenizer/tree/main) |
+
+
+*Note:*  
+- **Emu3.5** supports general-purpose multimodal predictions, including interleaved image-text generation and single-image generation (T2I/X2I) tasks.
+- **Emu3.5-Image** is a model focused on T2I/X2I tasks for best performance on these scenarios.
+- Both models are pure next-token predictors without DiDA acceleration (each image may take several minutes to generate).  
+- ⚡ **Stay tuned for DiDA-accelerated weights.**
+
+> 💡 **Usage tip:**  
+> For **interleaved image-text generation**, use **Emu3.5**.  
+> For **single-image generation** (T2I and X2I), use **Emu3.5-Image** for the best quality.
+
+
+
+## 2. Quick Start
+
+### Environment Setup
+
+```bash
+# Requires Python 3.12 or higher.
+git clone https://github.com/baaivision/Emu3.5
+cd Emu3.5
+pip install -r requirements/transformers.txt
+pip install flash_attn==2.8.3 --no-build-isolation
+```
+### Configuration
+
+Edit `configs/config.py` to set:
+
+- Paths: `model_path`, `vq_path`
+- Task template: `task_type in {t2i, x2i, howto, story, explore, vla}`
+- Input image: `use_image` (True to provide reference images, controls <|IMAGE|> token); set `reference_image` in each prompt to specify the image path. For x2i task, we recommand using `reference_image` as a list containing single/multiple image paths to be compatible with multi-image input.
+- Sampling: `sampling_params` (classifier_free_guidance, temperature, top_k/top_p, etc.)
+- Aspect Ratio (for t2i task): `aspect_ratio` ("4:3", "21:9", "1:1", "auto" etc..)
+
+### Run Inference
+
+```bash
+python inference.py --cfg configs/config.py
+```
+
+
+#### Example Configurations by Task
+Below are example commands for different tasks.
+Make sure to set CUDA_VISIBLE_DEVICES according to your available GPUs.
+
+
+```bash
+# 🖼️ Text-to-Image (T2I) task
+CUDA_VISIBLE_DEVICES=0 python inference.py --cfg configs/example_config_t2i.py
+
+# 🔄 Any-to-Image (X2I) task
+CUDA_VISIBLE_DEVICES=0,1 python inference.py --cfg configs/example_config_x2i.py
+
+# 🎯 Visual Guidance task
+CUDA_VISIBLE_DEVICES=0,1 python inference.py --cfg configs/example_config_visual_guidance.py
+
+# 📖 Visual Narrative task
+CUDA_VISIBLE_DEVICES=0,1 python inference.py --cfg configs/example_config_visual_narrative.py
+
+
+# After running inference, the model will generate results in protobuf format (.pb files) for each input prompt.
+```
+
+
+Protobuf outputs are written to `outputs/<exp_name>/proto/`. For better throughput, we recommend ≥2 GPUs.
+
+
+### Run Inference with vLLM
+
+#### vLLM Enviroment Setup
+
+1. [Optional Recommendation] Use a new virtual environment for vLLM backend.
+```bash
+conda create -n Emu3p5 python=3.12
+```
+
+2. Install vLLM and apply the patch files.
+```bash
+# Requires Python 3.12 or higher.
+# Recommended: CUDA 12.8.
+pip install -r requirements/vllm.txt
+pip install flash_attn==2.8.3 --no-build-isolation
+
+cd Emu3.5
+python src/patch/apply.py
+```
+
+#### Example Configurations by Task
+
+```bash
+# 🖼️ Text-to-Image (T2I) task
+CUDA_VISIBLE_DEVICES=0,1 python inference_vllm.py --cfg configs/example_config_t2i.py
+
+# 🔄 Any-to-Image (X2I) task
+CUDA_VISIBLE_DEVICES=0,1 python inference_vllm.py --cfg configs/example_config_x2i.py
+
+# 🎯 Visual Guidance task
+CUDA_VISIBLE_DEVICES=0,1 python inference_vllm.py --cfg configs/example_config_visual_guidance.py
+
+# 📖 Visual Narrative task
+CUDA_VISIBLE_DEVICES=0,1 python inference_vllm.py --cfg configs/example_config_visual_narrative.py
+```
+
+
+### Visualize Protobuf Outputs
+
+To visualize generated protobuf files (--video: Generate video visualizations for interleaved output):
+
+```bash
+python src/utils/vis_proto.py --input <input_proto_path> [--output <output_dir>] [--video]
+```
+
+- `--input`: supports a single `.pb` file or a directory; directories are scanned recursively.
+- `--output`: optional; defaults to `<input_dir>/results/<file_stem>` for files, or `<parent_dir_of_input>/results` for directories.
+
+Expected output directory layout (example):
+
+```text
+results/<pb_name>/
+├── 000_question.txt
+├── 000_global_cot.txt
+├── 001_text.txt
+├── 001_00_image.png
+├── 001_00_image_cot.txt
+├── 002_text.txt
+├── 002_00_image.png
+├── ...
+└── video.mp4              # only when --video is enabled
+```
+
+Each `*_text.txt` stores decoded segments, `*_image.png` stores generated frames, and matching `*_image_cot.txt` keeps image-level chain-of-thought notes when available.
+
+## 3. Gradio Demo
+
+We provide two Gradio Demos for different application scenarios:
+
+ Emu3.5-Image Demo —— Interactive interface optimized for Text-to-Image (T2I) and Any-to-Image (X2I) tasks:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1 python gradio_demo_image.py --host 0.0.0.0 --port 7860
+```
+
+Emu3.5-Interleave Demo —— Launch Emu3.5 Interleave Tasks (Visual Guidance and Visual Narrate) Gradio Demo
+```bash
+CUDA_VISIBLE_DEVICES=0,1 python gradio_demo_interleave.py --host 0.0.0.0 --port 7860
+```
+
+### Features
+
+- Image Generation: Support Text-to-Image Generation and Multimodal Image Generation
+- Interleaved Generation: Support long-sequence creation with alternating image and text generation
+- Multiple Aspect Ratios for T2I: 9 preset aspect ratios (4:3, 16:9, 1:1, etc.) plus auto mode
+- Chain-of-Thought Display: Automatically parse and format model's internal thinking process
+- Real-time Streaming: Stream text and image generation with live updates
+
+### Official Web & Mobile Apps
+
+- **Web**: Production-ready Emu3.5 experience is available at [zh.emu.world](https://zh.emu.world) (Mainland China) and [emu.world](https://emu.world) (global), featuring a curated homepage, “Create” workspace, inspiration feed, history, personal profile, and language switching.
+- **Mobile (Android APK & H5)**: Mobile clients provide the same core flows — prompt-based creation, “inspiration” gallery, personal center, and feedback & privacy entrypoints — with automatic UI language selection based on system settings.
+- **Docs**: For product usage details, see the **Emu3.5 AI 使用指南 (Chinese)** and **Emu3.5 AI User Guide (English)**:  
+  - CN: [Emu3.5 AI 使用指南](https://jwolpxeehx.feishu.cn/wiki/BKuKwkzZOi4pdRkVV13csI0FnIg?from=from_copylink)  
+  - EN: [Emu3.5 AI User Guide](https://jwolpxeehx.feishu.cn/wiki/Gcxtw9XHhisUu8kBEaac6s6xnhc?from=from_copylink)
+
+#### Mobile App Download (QR Codes)
+
+<div align='center'>
+  <table>
+    <tr>
+      <td align="center">
+        <img src="https://github.com/baaivision/Emu3.5/blob/main/assets/qr_zh.png?raw=True" alt="Emu3.5 Mobile App (Mainland China)" width="220" />
+        <br />
+        <sub><b>Emu3.5 Mobile · Mainland China</b></sub>
+      </td>
+      <td align="center">
+        <img src="https://github.com/baaivision/Emu3.5/blob/main/assets/qr.png?raw=True" alt="Emu3.5 Mobile App (Global)" width="220" />
+        <br />
+        <sub><b>Emu3.5 Mobile · Global</b></sub>
+      </td>
+    </tr>
+  </table>
+</div>
+
+
+
+## 4. Schedule
+
+- [x] Inference Code (NTP Version)
+- [ ] Advanced Image Decoder
+- [ ] Discrete Diffusion Adaptation (DiDA) Inference & Weights
+
+
+## 5. Citation
+
+```bibtex
+@misc{cui2025emu35nativemultimodalmodels,
+      title={Emu3.5: Native Multimodal Models are World Learners}, 
+      author={Yufeng Cui and Honghao Chen and Haoge Deng and Xu Huang and Xinghang Li and Jirong Liu and Yang Liu and Zhuoyan Luo and Jinsheng Wang and Wenxuan Wang and Yueze Wang and Chengyuan Wang and Fan Zhang and Yingli Zhao and Ting Pan and Xianduo Li and Zecheng Hao and Wenxuan Ma and Zhuo Chen and Yulong Ao and Tiejun Huang and Zhongyuan Wang and Xinlong Wang},
+      year={2025},
+      eprint={2510.26583},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2510.26583}, 
+}
+```

emu3_vqvae/config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "Emu3p5VisionVQModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_emu3p5visionvq.Emu3p5VisionVQConfig",
+    "AutoModel": "modeling_emu3p5visionvq.Emu3p5VisionVQModel"
+  },
+  "attn_resolutions": [
+    16
+  ],
+  "ch": 256,
+  "ch_mult": [
+    1,
+    1,
+    2,
+    2,
+    4
+  ],
+  "codebook_size": 131072,
+  "double_z": false,
+  "dropout": 0.0,
+  "embed_dim": 256,
+  "in_channels": 3,
+  "model_type": "Emu3p5VisionVQ",
+  "num_res_blocks": 4,
+  "out_ch": 3,
+  "resolution": 256,
+  "torch_dtype": "float32",
+  "transformers_version": "4.51.0",
+  "z_channels": 256
+}

emu3_vqvae/config.yaml

@@ -0,0 +1,14 @@
+ddconfig:
+  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
+
+n_embed: 131072
+embed_dim: 256

emu3_vqvae/configuration_emu3p5visionvq.py

@@ -0,0 +1,101 @@
+# coding=utf-8
+# Copyright 2025 The Emu team, BAAI and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+""" Emu3p5VisionVQ model configuration """
+
+from typing import List
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class Emu3p5VisionVQConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Emu3p5VisionVQ`]. It is used to instantiate an video movq
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+    defaults will yield a configuration to the VQ model presented in Emu3p5 paper.
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+    Args:
+        codebook_size (`int`, *optional*, defaults to 32768):
+            Codebook size of the VQ model.
+        embed_dim (`int`, *optional*, defaults to 4):
+            Dimension of the quantized vector in codebook.
+        z_channels (`int`, *optional*, defaults to 4):
+            Dimension of the output channel of encoder and the input channel of decoder
+        double_z (`bool`, *optional*, defaults to False):
+            Whether double the output dim of the encoder.
+        in_channels (`int`, *optional*, defaults to 3):
+            Input channel of encoder.
+        out_channels (`int`, *optional*, defaults to 3):
+            Output channel of decoder.
+        temporal_downsample_factor (`int`, *optional*, defaults to 4):
+            Temporal downsample factor.
+        ch (`int`, *optional*, defaults to 256):
+            Basic channel number of the intermediate blocks.
+        ch_mult (`List[int]`, *optional*, defaults to `[1, 2, 2, 4]`):
+            Channel scaling factor of the intermediate blocks.
+        num_res_blocks (`int`, *optional*, defaults to 2):
+            Residual block number in each stage.
+        attn_resolutions (`List[int]`, *optional*, defaults to 3):
+            Stage indices to apply attention.
+        dropout (`float`, *optional*, defaults to 0.0):
+            Dropout probability.
+    ```python
+    >>> from configuration_emu3p5visionvq import Emu3VisionVQConfig
+    >>> from modeling_emu3p5visionvq import Emu3VisionVQ
+    >>> # Initializing a video VQ model of Emu3 configuration
+    >>> configuration = Emu3VisionVQConfig()
+    >>> # Initializing a model from the Emu3 VQ model style configuration
+    >>> model = Emu3VisionVQModel(configuration)
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "Emu3p5VisionVQ"
+
+    def __init__(
+        self,
+        double_z: bool = False,
+        z_channels: int = 256,
+        resolution: int = 256,
+        in_channels: int = 3,
+        out_ch: int = 3,
+        ch: int = 256,
+        ch_mult: List[int] = [1, 1, 2, 2, 4],
+        num_res_blocks: int = 4,
+        attn_resolutions: List[int] = [16],
+        dropout: float = 0.0,
+        codebook_size: int = 131072,
+        embed_dim: int = 256,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.double_z = double_z
+        self.z_channels = z_channels
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.out_ch = out_ch
+        self.ch = ch
+        self.ch_mult = ch_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_resolutions = attn_resolutions
+        self.dropout = dropout
+
+        self.codebook_size = codebook_size
+        self.embed_dim = embed_dim

emu3_vqvae/kmeans_16384_centroids.pt

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

emu3_vqvae/kmeans_4096_centroids.pt

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

emu3_vqvae/kmeans_8192_centroids.pt

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

emu3_vqvae/model.ckpt

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

emu3_vqvae/model.safetensors

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

emu3_vqvae/modeling_emu3p5visionvq.py

@@ -0,0 +1,497 @@
+# coding=utf-8
+# Copyright 2025 The Emu team, BAAI and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+""" Emu3p5VisionVQ model """
+
+
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+from torch import nn, einsum
+from torch.nn import functional as F
+from transformers.modeling_utils import PreTrainedModel
+
+from .configuration_emu3p5visionvq import Emu3p5VisionVQConfig
+
+
+def nonlinearity(x):
+    # swish
+    return x * torch.sigmoid(x)
+
+
+def Emu3p5VisionVQNormalize(in_channels):
+    return nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class Emu3p5VisionVQUpsample(nn.Module):
+
+    def __init__(self, in_channels):
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_channels,
+            in_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
+        x = self.conv(x)
+        return x
+
+
+class Emu3p5VisionVQDownsample(nn.Module):
+
+    def __init__(self, in_channels):
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_channels,
+            in_channels,
+            kernel_size=3,
+            stride=2,
+            padding=0,
+        )
+
+    def forward(self, x):
+        pad = (0, 1, 0, 1)
+        x = F.pad(x, pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+
+
+class Emu3p5VisionVQResnetBlock(nn.Module):
+
+    def __init__(
+        self,
+        *,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        conv_shortcut: bool = False,
+        dropout: float = 0.0
+    ):
+        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 = Emu3p5VisionVQNormalize(in_channels)
+        self.conv1 = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+        self.norm2 = Emu3p5VisionVQNormalize(out_channels)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = 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 = nn.Conv2d(
+                    in_channels,
+                    out_channels,
+                    kernel_size=3,
+                    stride=1,
+                    padding=1,
+                )
+            else:
+                self.nin_shortcut = nn.Conv2d(
+                    in_channels,
+                    out_channels,
+                    kernel_size=1,
+                    stride=1,
+                    padding=0,
+                )
+
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        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 Emu3p5VisionVQAttnBlock(nn.Module):
+
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Emu3p5VisionVQNormalize(in_channels)
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = 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_)
+
+        # compute attention
+        b,c,h,w = q.shape
+        q = q.reshape(b, c, h * w)
+        q = q.permute(0, 2, 1)   # b,hw,c
+        k = k.reshape(b, c, h * w) # b,c,hw
+        w_ = torch.bmm(q, k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = F.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b, c, h * w)
+        w_ = w_.permute(0, 2, 1)   # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v, w_)     # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+        h_ = h_.reshape(b, c, h, w)
+
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+
+class Emu3p5VisionVQEncoder(nn.Module):
+
+    def __init__(self, config: Emu3p5VisionVQConfig):
+        super().__init__()
+        self.ch = config.ch
+        self.num_resolutions = len(config.ch_mult)
+        self.num_res_blocks = config.num_res_blocks
+        self.in_channels = config.in_channels
+        self.resolution = config.resolution
+
+        # downsampling
+        self.conv_in = nn.Conv2d(
+            self.in_channels,
+            self.ch,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+        curr_res = self.resolution
+
+        in_ch_mult = (1, ) + tuple(config.ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = config.ch * in_ch_mult[i_level]
+            block_out = config.ch * config.ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(
+                    Emu3p5VisionVQResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        dropout=config.dropout,
+                    ),
+                )
+                block_in = block_out
+                if curr_res in config.attn_resolutions:
+                    attn.append(Emu3p5VisionVQAttnBlock(block_in))
+
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Emu3p5VisionVQDownsample(block_in)
+                curr_res = curr_res // 2
+
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = Emu3p5VisionVQResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            dropout=config.dropout,
+        )
+        self.mid.attn_1 = Emu3p5VisionVQAttnBlock(block_in)
+        self.mid.block_2 = Emu3p5VisionVQResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            dropout=config.dropout,
+        )
+
+        # end
+        self.norm_out = Emu3p5VisionVQNormalize(block_in)
+        self.conv_out = nn.Conv2d(
+            block_in,
+            2 * config.z_channels if config.double_z else config.z_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+
+    def forward(self, x):
+        # downsampling
+        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])
+                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]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Emu3p5VisionVQDecoder(nn.Module):
+
+    def __init__(self, config: Emu3p5VisionVQConfig):
+        super().__init__()
+        self.ch = config.ch
+        self.num_resolutions = len(config.ch_mult)
+        self.num_res_blocks = config.num_res_blocks
+
+        self.resolution = config.resolution
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        in_ch_mult = (1, ) + tuple(config.ch_mult)
+        block_in = config.ch * config.ch_mult[self.num_resolutions-1]
+
+        curr_res = config.resolution // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, config.z_channels, curr_res, curr_res)
+
+        # z to block_in
+        self.conv_in = nn.Conv2d(
+            config.z_channels,
+            block_in,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = Emu3p5VisionVQResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            dropout=config.dropout,
+        )
+        self.mid.attn_1 = Emu3p5VisionVQAttnBlock(block_in)
+        self.mid.block_2 = Emu3p5VisionVQResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            dropout=config.dropout,
+        )
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = config.ch * config.ch_mult[i_level]
+            for i_block in range(self.num_res_blocks + 1):
+                block.append(
+                    Emu3p5VisionVQResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        dropout=config.dropout,
+                    ),
+                )
+                block_in = block_out
+                if curr_res in config.attn_resolutions:
+                    attn.append(Emu3p5VisionVQAttnBlock(block_in))
+
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Emu3p5VisionVQUpsample(block_in)
+                curr_res = curr_res * 2
+            self.up.insert(0, up) # prepend to get consistent order
+
+        # end
+        self.norm_out = Emu3p5VisionVQNormalize(block_in)
+        self.conv_out = nn.Conv2d(
+            block_in,
+            config.out_ch,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+    def forward(self, z):
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        # upsampling
+        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)
+                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)
+
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+
+        return h
+
+
+class Emu3p5VisionVQVectorQuantizer(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+
+        self.n_e = config.codebook_size
+        self.e_dim = config.embed_dim
+
+        self.embedding = nn.Embedding(self.n_e, self.e_dim)
+
+    def forward(self, z):
+        # z: [b, d, h, w]
+        embedding = self.embedding.weight  # [n, d]
+
+        # cal similarity
+        logits = torch.einsum("b d h w, n d -> b n h w", z, embedding)
+
+        # get max indices
+        ind = logits.argmax(dim=1)  # [b, h, w]
+
+        # lookup embedding
+        z_q = embedding[ind]  # [b, h, w, d]
+        z_q = z_q.permute(0, 3, 1, 2).contiguous()  # -> [b, d, h, w]
+
+        return z_q, ind.flatten()
+
+    def get_codebook_entry(self, indices, shape=None):
+        # get quantized latent vectors
+        z_q = self.embedding(indices)
+
+        # shape should in B H W
+        if shape is not None:
+            if len(shape) == 3:
+                shape = shape + (self.e_dim, )
+
+            z_q = z_q.view(shape)
+
+            # reshape back to match original input shape
+            # b h w c -> b c h w
+            z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+        return z_q
+
+
+class Emu3p5VisionVQPretrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = Emu3p5VisionVQConfig
+    base_model_prefix = "emu3p5visionvq"
+    main_input_name = "pixel_values"
+    _no_split_modules = ["Emu3p5VisionVQResnetBlock", "Emu3p5VisionVQAttnBlock"]
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Conv2d):
+            nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
+        # copied from the `reset_parameters` method of `class Linear(Module)` in `torch`.
+        elif isinstance(module, nn.Linear):
+            nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5))
+            if module.bias is not None:
+                fan_in, _ = nn.init._calculate_fan_in_and_fan_out(module.weight)
+                bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+                nn.init.uniform_(module.bias, -bound, bound)
+        elif isinstance(module, (nn.BatchNorm2d, nn.GroupNorm)):
+            nn.init.constant_(module.weight, 1)
+            nn.init.constant_(module.bias, 0)
+
+
+class Emu3p5VisionVQModel(Emu3p5VisionVQPretrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+
+        self.encoder = Emu3p5VisionVQEncoder(config)
+        self.decoder = Emu3p5VisionVQDecoder(config)
+        self.quantize = Emu3p5VisionVQVectorQuantizer(config)
+
+        self.quant_conv = nn.Conv2d(config.z_channels, config.embed_dim, 1)
+        self.post_quant_conv = nn.Conv2d(config.embed_dim, config.z_channels, 1)
+
+        self.post_init()
+
+    def encode(self, x: torch.Tensor):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        quant_embed, token_ids = self.quantize(h)
+        return quant_embed, None, (None, None, token_ids)
+
+    def decode(self, x: torch.Tensor):
+        quant = self.post_quant_conv(x)
+        dec = self.decoder(quant)
+        return dec
+
+    def decode_code(self, code_b, shape=None):
+        # shape specifying (batch, height, width, channel)
+        quant_b = self.quantize.get_codebook_entry(code_b, shape=shape)
+        dec = self.decode(quant_b)
+        return dec
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": [
+    11
+  ],
+  "pad_token_id": 11,
+  "transformers_version": "4.57.6",
+  "use_cache": false
+}

model-00001-of-00006.safetensors

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

model-00002-of-00006.safetensors

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

model-00003-of-00006.safetensors

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

model-00004-of-00006.safetensors

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

model-00005-of-00006.safetensors

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

model-00006-of-00006.safetensors

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

modeling_ministral.py

@@ -0,0 +1,658 @@
+from collections.abc import Callable
+from typing import Optional, Union
+
+import torch
+from torch import nn
+from torch.utils.checkpoint import checkpoint as torch_checkpoint
+
+from transformers.utils.generic import check_model_inputs
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+# from transformers.integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
+from transformers.integrations import use_kernel_forward_from_hub
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask, ALL_MASK_ATTENTION_FUNCTIONS
+try:
+    from transformers.masking_utils import sdpa_mask_older_torch
+except ImportError:
+    sdpa_mask_older_torch = None
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_layers import (
+    GenericForQuestionAnswering,
+    GenericForSequenceClassification,
+    GenericForTokenClassification,
+    GradientCheckpointingLayer,
+)
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+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, can_return_tuple
+# from transformers.utils.generic import maybe_autocast
+from .configuration_ministral_dlm import MinistralDLMConfig
+from functools import partial
+#ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa'] = sdpa_mask_older_torch
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+# @use_kernel_func_from_hub("rotary_pos_emb")
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    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:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    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=module.training)
+    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, max_position_embeddings) -> torch.Tensor:
+    if not beta or not max_position_embeddings:
+        return torch.ones(1, device=positions_ids.device, dtype=torch.float32)
+    scaling = 1 + beta * torch.log(1 + torch.floor(positions_ids / max_position_embeddings))
+    return scaling.unsqueeze(-1)
+
+
+# @use_kernelized_func(apply_rotary_pos_emb)
+class Ministral3Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: MinistralDLMConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        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.diffusion_lm = getattr(config, 'diffusion_lm', False)
+
+    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,
+        **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 = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_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:
+                # sin and cos are specific to RoPE models; cache_position needed for the static 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:  ## if use_cache == False, do not update cache
+                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 if not self.training else self.attention_dropout,
+                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 if not self.training else self.attention_dropout,
+                scaling=self.scaling,
+                sliding_window=getattr(self.config, "sliding_window", None),  # main diff with Llama
+                **kwargs,
+            )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        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.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Ministral3RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        Ministral3RMSNorm is equivalent to T5LayerNorm
+        """
+        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}"
+
+# hack
+try:
+    import deepspeed
+except ImportError:
+    deepspeed = None
+
+
+class DeepSpeedGradientCheckpointingLayer(GradientCheckpointingLayer):
+    """Base class for layers with gradient checkpointing.
+
+    This class enables gradient checkpointing functionality for a layer. By default, gradient checkpointing is disabled
+    (`gradient_checkpointing = False`). When `model.set_gradient_checkpointing()` is called, gradient checkpointing is
+    enabled by setting `gradient_checkpointing = True` and assigning a checkpointing function to `_gradient_checkpointing_func`.
+
+    Important:
+
+        When using gradient checkpointing with `use_reentrant=True`, inputs that require gradients (e.g. hidden states)
+        must be passed as positional arguments (`*args`) rather than keyword arguments to properly propagate gradients.
+
+        Example:
+
+            ```python
+            >>> # Correct - hidden_states passed as positional arg
+            >>> out = self.layer(hidden_states, attention_mask=attention_mask)
+
+            >>> # Incorrect - hidden_states passed as keyword arg
+            >>> out = self.layer(hidden_states=hidden_states, attention_mask=attention_mask)
+            ```
+    """
+
+    gradient_checkpointing = False
+
+    def __call__(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,
+        **kwargs: Unpack[TransformersKwargs]
+        ):
+        if self.gradient_checkpointing and self.training:
+            do_warn = False
+            layer_name = self.__class__.__name__
+            message = f"Caching is incompatible with gradient checkpointing in {layer_name}. Setting"
+
+            if "use_cache" in kwargs and kwargs["use_cache"]:
+                kwargs["use_cache"] = False
+                message += " `use_cache=False`,"
+                do_warn = True
+
+            if "past_key_value" in kwargs and kwargs["past_key_value"] is not None:
+                kwargs["past_key_value"] = None
+                message += " `past_key_value=None`,"
+                do_warn = True
+
+            if "past_key_values" in kwargs and kwargs["past_key_values"] is not None:
+                kwargs["past_key_values"] = None
+                message += " `past_key_values=None`,"
+                do_warn = True
+
+            if "layer_past" in kwargs and kwargs["layer_past"] is not None:
+                kwargs["layer_past"] = None
+                message += " `layer_past=None`,"
+                do_warn = True
+
+            # warn if anything was changed
+            if do_warn:
+                message = message.rstrip(",") + "."
+                print(message)
+            # breakpoint()
+            assert not any([isinstance(x,torch.Tensor) for x in kwargs.values()])
+            checkpoint_fn = deepspeed.checkpointing.checkpoint if deepspeed is not None else torch_checkpoint
+            return checkpoint_fn(
+                partial(super().__call__, **kwargs),
+                hidden_states,
+                attention_mask,
+                position_ids,
+                past_key_values,
+                use_cache,
+                cache_position,
+                position_embeddings,
+            )
+        return super().__call__(
+            hidden_states,attention_mask,position_ids,past_key_values,use_cache,cache_position,
+            position_embeddings, **kwargs
+        )
+
+
+class Ministral3DecoderLayer(DeepSpeedGradientCheckpointingLayer):
+    def __init__(self, config: MinistralDLMConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+    
+        if hasattr(config, 'attn_class'):
+            attn_class = config.attn_class
+        else:
+            attn_class = Ministral3Attention
+
+        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)
+
+    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,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        # Self Attention
+        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,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+@auto_docstring
+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  # fix linting for `register_buffer`
+
+    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]:
+        """
+        Computes the inverse frequencies according to the original RoPE implementation
+        Args:
+            config ([`~transformers.PreTrainedConfig`]):
+                The model configuration.
+            device (`torch.device`):
+                The device to use for initialization of the inverse frequencies.
+            seq_len (`int`, *optional*):
+                The current sequence length. Unused for this type of RoPE.
+        Returns:
+            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+        """
+        base = config.rope_parameters["rope_theta"]
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+        attention_factor = 1.0  # Unused in this type of RoPE
+
+        # Compute the inverse frequencies
+        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  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        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()
+
+        # device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        # with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
+
+        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)
+
+
+@auto_docstring
+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.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.rotary_emb = Ministral3RotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @check_model_inputs
+    @auto_docstring
+    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,  # for compatibility with LLaVA-1.5 input
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        """
+        Args:
+            input_embeddings (`torch.FloatTensor`, *optional*):
+                Alias for `inputs_embeds` kept for backward compatibility with older LLaVA-style callers.
+                If provided, it overrides `inputs_embeds`.
+        """
+        if input_embeddings is not None:
+            inputs_embeds = input_embeddings
+            input_embeddings = None
+        # breakpoint()
+        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)
+
+        if use_cache and past_key_values is None:
+            # past_key_values = DynamicCache(config=self.config)
+            past_key_values = DynamicCache()
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            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)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            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,
+                **kwargs,
+            )
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+
+@auto_docstring
+class Ministral3ForCausalLM(Ministral3PreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = Ministral3Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    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,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        r"""
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, Ministral3ForCausalLM
+
+        >>> model = Ministral3ForCausalLM.from_pretrained("meta-ministral3/Ministral3-2-7b-hf")
+        >>> tokenizer = AutoTokenizer.from_pretrained("meta-ministral3/Ministral3-2-7b-hf")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class Ministral3ForTokenClassification(GenericForTokenClassification, Ministral3PreTrainedModel):
+    pass
+
+
+class Ministral3ForSequenceClassification(GenericForSequenceClassification, Ministral3PreTrainedModel):
+    pass
+
+
+class Ministral3ForQuestionAnswering(GenericForQuestionAnswering, Ministral3PreTrainedModel):
+    pass
+
+
+__all__ = [
+    "Ministral3ForCausalLM",
+    "Ministral3ForQuestionAnswering",
+    "Ministral3Model",
+    "Ministral3PreTrainedModel",
+    "Ministral3ForSequenceClassification",
+    "Ministral3ForTokenClassification",
+]

modeling_ministral_dlm.py

@@ -0,0 +1,1495 @@
+import copy
+from dataclasses import dataclass
+from typing import Callable, Optional, Tuple, Union
+import random
+import os
+import sys
+import json
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutput
+from transformers.utils import ModelOutput
+
+from torch.nn.attention.flex_attention import BlockMask, flex_attention, create_block_mask, or_masks
+
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+
+from transformers.processing_utils import Unpack
+
+from transformers.cache_utils import Cache, DynamicCache
+
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from transformers.generation import GenerationMixin
+
+import math
+
+from .chat_utils import generate_with_prefix_cache_block_diff
+from .modeling_ministral import Ministral3Model, Ministral3PreTrainedModel, Ministral3Attention, apply_rotary_pos_emb, repeat_kv, _get_llama_4_attn_scale
+from .configuration_ministral_dlm import MinistralDLMConfig
+
+try:
+    from flash_attn import flash_attn_func
+except:
+    print("flash attention not found, please install flash attention for better performance.")
+__all__ = ["MinistralDiffEncoderModel", "MinistralFlexAttention"]
+
+@dataclass
+class MinistralDiffOutputWithPast(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
+
+
+# @torch.compile(dynamic=True, mode="reduce-overhead")
+# @torch.compile(mode="default")
+# @torch.compile(fullgraph=True, mode="reduce-overhead", dynamic=False)
+@torch.compile(fullgraph=True, mode="max-autotune-no-cudagraphs", dynamic=False)
+def fused_flex_attention(q, k, v, block_mask=None):
+    return flex_attention(q, k, v, block_mask=block_mask)
+
+
+def _crop_dynamic_cache(past_key_values: DynamicCache, max_length: int):
+    """Crop a DynamicCache to max_length, compatible with both old and new transformers."""
+    if hasattr(past_key_values, 'crop'):
+        past_key_values.crop(max_length)
+    else:
+        for layer_idx in range(len(past_key_values)):
+            past_key_values.key_cache[layer_idx] = past_key_values.key_cache[layer_idx][:, :, :max_length]
+            past_key_values.value_cache[layer_idx] = past_key_values.value_cache[layer_idx][:, :, :max_length]
+        past_key_values._seen_tokens = max_length
+
+
+def _extract_draft_kv_cache(past_key_values: DynamicCache, clean_len: int, block_length: int):
+    """After quadratic decoding, extract only draft tokens (first of each block) from cache."""
+    for layer_idx in range(len(past_key_values)):
+        if hasattr(past_key_values, 'layers'):
+            layer_cache = past_key_values.layers[layer_idx]
+            k, v = layer_cache.keys, layer_cache.values
+        else:
+            k = past_key_values.key_cache[layer_idx]
+            v = past_key_values.value_cache[layer_idx]
+
+        clean_k, draft_k = k[:, :, :clean_len], k[:, :, clean_len::block_length + 1]
+        clean_v, draft_v = v[:, :, :clean_len], v[:, :, clean_len::block_length + 1]
+        new_k = torch.cat([clean_k, draft_k], dim=2)
+        new_v = torch.cat([clean_v, draft_v], dim=2)
+
+        if hasattr(past_key_values, 'layers'):
+            layer_cache.keys = new_k
+            layer_cache.values = new_v
+        else:
+            past_key_values.key_cache[layer_idx] = new_k
+            past_key_values.value_cache[layer_idx] = new_v
+
+    past_key_values._seen_tokens = clean_len + block_length
+
+# with reference to https://github.com/pytorch-labs/attention-gym/blob/main/examples/flex_attn.ipynb
+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
+        if self.config.dlm_paradigm == 'bidirectional':
+            self.bidirectional_mask = self.compute_block_mask(mode='bidirectional')
+        elif self.config.dlm_paradigm == 'autoregressive':
+            self.autoregressive_mask = self.compute_block_mask(mode='autoregressive')
+        elif self.config.dlm_paradigm == 'block_diff':
+            self.block_diff_mask = None
+        elif self.config.dlm_paradigm == 'sbd_block_diff':
+            self.sbd_block_diff_mask = None
+        else:
+            raise ValueError(f"Unknown attention mode: {self.config.dlm_paradigm}")
+
+        self.block_size = self.block_size_orig
+        self.mode = self.config.dlm_paradigm
+        self._quadratic_block_mask = {}
+
+        import torch._dynamo.config as dcfg
+        dcfg.cache_size_limit = 512
+
+
+    def _get_sbd_inference_quadratic_decoding_block_mask(self, block_length: int):
+        if block_length not in self._quadratic_block_mask:
+            draft_len = block_length * (block_length + 1)
+
+            def quadratic(b, h, q_idx, kv_idx):
+                first_clean = torch.logical_and(
+                    kv_idx % (block_length + 1) == 0,
+                    kv_idx < draft_len,
+                )
+                first_clean = torch.logical_and(first_clean, q_idx >= kv_idx)
+                block_q = q_idx // (block_length + 1)
+                block_kv = kv_idx // (block_length + 1)
+                same_block = torch.logical_and(block_q == block_kv, q_idx < draft_len)
+                same_block_except_first = torch.logical_and(
+                    same_block,
+                    q_idx % (block_length + 1) != 0,
+                )
+                draft_part = torch.logical_or(first_clean, same_block_except_first)
+                clean_part = kv_idx >= draft_len
+                return torch.logical_or(draft_part, clean_part)
+
+            block_mask = create_block_mask(
+                quadratic,
+                B=None,
+                H=None,
+                Q_LEN=draft_len,
+                KV_LEN=draft_len + self.config.max_position_embeddings,
+                device="cuda",
+            )
+
+            self._quadratic_block_mask[block_length] = block_mask
+
+        return self._quadratic_block_mask[block_length]
+
+
+    def set_attention_mode(self, mode, block_size=None):
+        self.mode = mode
+        self.block_size = block_size
+
+    def compute_block_mask(self, mode, q_len=None, block_size=None):
+
+        def bidirectional_mask(b, h, q, kv): 
+            return (q >= kv) | (q < kv)
+        
+        def autoregressive_mask(b, h, q, kv):
+            return (q >= kv)
+
+        def block_diff_mask(block_size, b, h, q_idx, kv_idx, n):
+            x0_flag_q = (q_idx >= n)
+            x0_flag_kv = (kv_idx >= n)
+
+            # Compute block indices
+            block_q = torch.where(x0_flag_q == 1,
+                                    (q_idx - n) // block_size,
+                                    q_idx // block_size)
+            block_kv = torch.where(x0_flag_kv == 1,
+                                    (kv_idx - n) // block_size,
+                                    kv_idx // block_size)
+
+            # **1. Block Diagonal Mask (M_BD) **
+            block_diagonal = (block_q == block_kv) & (x0_flag_q == x0_flag_kv)
+
+            # **2. Offset Block-Causal Mask (M_OBC) **
+            offset_block_causal = (
+                (block_q > block_kv)
+                & (x0_flag_kv == 1)
+                & (x0_flag_q == 0)
+            )
+
+            # **3. Block-Causal Mask (M_BC) **
+            block_causal = (block_q >= block_kv) & (x0_flag_kv == 1) & (x0_flag_q == 1)
+
+            # **4. Combine Masks **
+            return block_diagonal | offset_block_causal | block_causal
+        
+
+        def sbd_block_diff_mask(block_size, b, h, q_idx, kv_idx, n):
+            x0_flag_q = (q_idx >= n)
+            x0_flag_kv = (kv_idx >= n)
+
+            # Compute block indices
+            block_q = torch.where(x0_flag_q == 1,
+                                    (q_idx - n) // block_size,
+                                    q_idx // block_size)
+            block_kv = torch.where(x0_flag_kv == 1,
+                                    (kv_idx - n) // block_size,
+                                    kv_idx // block_size)
+
+            # **1. Block Diagonal Mask (M_BD) **
+            block_diagonal = (block_q == block_kv) & (x0_flag_kv == 0) & (x0_flag_q == 0)
+
+            # **2. Offset Block-Causal Mask (M_OBC) **
+            offset_block_causal = (
+                (block_q > block_kv)
+                & (x0_flag_kv == 1)
+                & (x0_flag_q == 0)
+            )
+
+            # **3. Fully Causal Mask (M_BC) **
+            fully_causal = (q_idx >= kv_idx) & (x0_flag_kv == 1) & (x0_flag_q == 1)
+
+            # **4. Combine Masks **
+            return block_diagonal | offset_block_causal | fully_causal
+        
+        def modality_indices_based_mask(block_size, b, h, q_idx, kv_idx, image_doc_id):
+            return  (image_doc_id[b, q_idx] > 0) & (image_doc_id[b, q_idx] == image_doc_id[b, kv_idx])
+
+        if mode == 'bidirectional':
+            attn_mask = bidirectional_mask
+        elif mode == 'autoregressive':
+            attn_mask = autoregressive_mask
+        elif mode == 'block_diff':
+            assert block_size is not None
+            attn_mask = lambda b, h, q, kv: block_diff_mask(block_size, b, h, q, kv, self.max_seq_length)
+        elif mode == 'sbd_block_diff':
+            assert block_size is not None
+            attn_mask = lambda b, h, q, kv: sbd_block_diff_mask(block_size, b, h, q, kv, self.max_seq_length)
+        else:
+            raise ValueError(f"Unknown attention mode: {mode}")
+
+        if q_len is not None:
+            Q_LEN = q_len
+        else:
+            if mode in ['block_diff', 'sbd_block_diff']:
+                Q_LEN = self.max_seq_length * 2
+            else:
+                Q_LEN = self.max_seq_length
+
+        block_mask = create_block_mask(
+            attn_mask, B=None, H=None, Q_LEN=Q_LEN, KV_LEN=Q_LEN
+        )
+
+        return block_mask
+
+
+    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,
+        is_training: bool = True,
+        overwrite_block_mask = None,
+        overwrite_attn_impl = None,
+        use_cache: Optional[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,
+                is_training=is_training,
+                use_cache=use_cache,
+                **kwargs,
+            )
+        bsz, q_len, _ = hidden_states.size()
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+
+        if self.mode in ['block_diff', 'sbd_block_diff'] and is_training:
+            # Split query and key states in half along sequence length dimension
+            q1, q2 = query_states.chunk(2, dim=2)
+            k1, k2 = key_states.chunk(2, dim=2)
+            
+            # Apply RoPE independently to each half
+            q1, k1 = apply_rotary_pos_emb(q1, k1, cos, sin)
+            q2, k2 = apply_rotary_pos_emb(q2, k2, cos, sin)
+            
+            # Recombine the halves
+            query_states = torch.cat([q1, q2], dim=2)
+            key_states = torch.cat([k1, k2], dim=2)
+        else:
+            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:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            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:  ## if use_cache == False, do not update cache
+                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 = getattr(self.config, "self_spec_inference_mode", None)
+        if self_spec_inference_mode is not None:
+            if self_spec_inference_mode == "quadratic":
+                block_length = getattr(self.config, "block_length", None) or getattr(self.config, "block_size", None)
+                if block_length is None:
+                    raise ValueError("SBD quadratic decoding requires block_length in config.")
+                if past_key_values is not None:
+                    seq_len = key_states.shape[2]
+                    draft_len = block_length * (block_length + 1)
+
+                    clean_keys = key_states[:, :, :-draft_len]
+                    draft_keys = key_states[:, :, -draft_len:]
+                    clean_values = value_states[:, :, :-draft_len]
+                    draft_values = value_states[:, :, -draft_len:]
+                    key_states = torch.cat([draft_keys, clean_keys], dim=2)
+                    value_states = torch.cat([draft_values, clean_values], dim=2)
+
+                    block_mask: BlockMask = self._get_sbd_inference_quadratic_decoding_block_mask(
+                        block_length=block_length
+                    )
+                    block_mask.seq_lengths = (draft_len, seq_len)
+                else:
+                    seq_len = query_states.shape[2]
+                    draft_len = block_length * (block_length + 1)
+                    clean_len = seq_len - draft_len
+
+                    def _causal_mask(b, h, q_idx, kv_idx):
+                        return torch.logical_and(q_idx >= kv_idx, q_idx < clean_len)
+
+                    def _draft2clean_mask(b, h, q_idx, kv_idx):
+                        full_clean = torch.logical_and(q_idx >= clean_len, kv_idx <= clean_len)
+                        first_clean = torch.logical_and(
+                            q_idx >= clean_len, (kv_idx - clean_len) % (block_length + 1) == 0
+                        )
+                        first_clean = torch.logical_and(first_clean, q_idx >= kv_idx)
+                        return torch.logical_or(full_clean, first_clean)
+
+                    def _draft_mask(b, h, q_idx, kv_idx):
+                        block_q = (q_idx - clean_len) // (block_length + 1)
+                        block_kv = (kv_idx - clean_len) // (block_length + 1)
+                        quadrant = torch.logical_and(q_idx >= clean_len, kv_idx >= clean_len)
+                        same_block = torch.logical_and(block_q == block_kv, quadrant)
+                        same_block_except_first = torch.logical_and(
+                            same_block,
+                            (q_idx - clean_len) % (block_length + 1) != 0,
+                        )
+                        return torch.logical_and(block_q == block_kv, same_block_except_first)
+
+                    mask = or_masks(_causal_mask, _draft2clean_mask)
+                    mask = or_masks(mask, _draft_mask)
+
+                    block_mask = create_block_mask(
+                        mask, B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len,
+                    )
+
+                key_states = repeat_kv(key_states, self.num_key_value_groups)
+                value_states = repeat_kv(value_states, self.num_key_value_groups)
+                attn_output = flex_attention(query_states, key_states, value_states, block_mask=block_mask)
+                attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
+                attn_output = self.o_proj(attn_output)
+                return attn_output, None
+
+            elif self_spec_inference_mode == "default":
+                block_length = getattr(self.config, "block_length", None) or getattr(self.config, "block_size", None)
+                if block_length is None:
+                    raise ValueError("SBD default decoding requires block_length in config.")
+                seq_len = query_states.shape[2]
+                prefix_len = seq_len - block_length
+
+                def _clean_q_mask(b, h, q_idx, kv_idx):
+                    return torch.logical_and(q_idx >= kv_idx, q_idx < prefix_len)
+
+                def _noisy_q_mask(b, h, q_idx, kv_idx):
+                    return q_idx >= prefix_len
+
+                block_mask = create_block_mask(
+                    or_masks(_clean_q_mask, _noisy_q_mask),
+                    B=None,
+                    H=None,
+                    Q_LEN=seq_len,
+                    KV_LEN=seq_len,
+                )
+
+                key_states = repeat_kv(key_states, self.num_key_value_groups)
+                value_states = repeat_kv(value_states, self.num_key_value_groups)
+                attn_output = flex_attention(query_states, key_states, value_states, block_mask=block_mask)
+                attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
+                attn_output = self.o_proj(attn_output)
+                return attn_output, None
+        
+        else:
+            key_states = repeat_kv(key_states, self.num_key_value_groups)
+            value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+            if overwrite_block_mask is not None:
+                block_mask = overwrite_block_mask
+                if block_mask == 'full':
+                    block_mask = None
+            else:
+                if self.mode == 'bidirectional':
+                    block_mask = None
+                    overwrite_attn_impl = 'flash_attn'
+                    # if self.bidirectional_mask is None or q_len != self.bidirectional_mask.shape[-2]:
+                    #     block_mask = self.compute_block_mask(mode='bidirectional', q_len=q_len)
+                    # else:
+                    #     block_mask = self.bidirectional_mask
+
+                elif self.mode == 'autoregressive':
+                    if self.autoregressive_mask is None or q_len != self.autoregressive_mask.shape[-2]:
+                        block_mask = self.compute_block_mask(mode='autoregressive', q_len=q_len)
+                    else:
+                        block_mask = self.autoregressive_mask
+
+                elif self.mode == 'block_diff':
+                    if self.block_diff_mask is None or self.block_size != self.block_size_orig or q_len != self.block_diff_mask.shape[-2]:
+                        block_mask = self.compute_block_mask(mode='block_diff', block_size=self.block_size, q_len=q_len)
+                    else:
+                        block_mask = self.block_diff_mask
+                elif self.mode == 'sbd_block_diff':
+                    if self.sbd_block_diff_mask is None or self.block_size != self.block_size_orig or q_len != self.sbd_block_diff_mask.shape[-2]:
+                        block_mask = self.compute_block_mask(mode='sbd_block_diff', block_size=self.block_size, q_len=q_len)
+                    else:
+                        block_mask = self.sbd_block_diff_mask
+                else:
+                    raise ValueError(f"Unknown attention mode: {self.mode}")
+            if overwrite_attn_impl == 'flash_attn':
+                
+    
+                # FlashAttention expects (batch, seqlen, nheads, headdim)
+                # Ensure your tensors are in this layout or permute them here
+                #print(query_states.shape,key_states.shape,value_states.shape)
+                if self.diffusion_lm:
+                    causal = False
+                else:
+                    causal = True
+                attn_output = flash_attn_func(
+                    query_states.transpose(1,2), 
+                    key_states.transpose(1,2), 
+                    value_states.transpose(1,2), 
+                    dropout_p=0.0,      # Set your dropout probability
+                    softmax_scale=None, # Defaults to 1/sqrt(head_dim)
+                    causal=causal         # Set to True if using a causal block_mask logic
+                ).transpose(1,2)
+                
+            else:
+                attn_output = fused_flex_attention(query_states, key_states, value_states, block_mask=block_mask)
+            attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
+
+            attn_output = self.o_proj(attn_output)
+
+            return attn_output, None
+
+
+def gumbel_topk(log_w: torch.Tensor, k: int) -> torch.Tensor:
+    """Return a Bool mask of length len(log_w) with exactly k True."""
+    g = -torch.log(-torch.log(torch.rand_like(log_w) + 1e-9) + 1e-9)
+    topk = torch.topk(log_w + g, k).indices
+    mask = torch.zeros_like(log_w, dtype=torch.bool)
+    mask[topk] = True
+    return mask
+
+
+class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
+    """
+    A single model with:
+      - a bidirectional encoder + diffusion‐LM head over A
+      - a causal decoder + LM head over B, conditioned on F_A
+    """
+
+    # Shared/tied tensors that can appear dynamically based on config.
+    # Registering these patterns lets save_pretrained() deduplicate safely.
+    # _dynamic_tied_weights_keys = [
+    #     r"encoder\.embed_tokens\.weight",
+    #     r"diffusion_head\.weight",
+    #     r"encoder\.vision_tower(?:\.vision_tower)?\.visual_bridge_model\.quantizer\.quantize\.codebooks\.\d+\.(?:embed|embed_ema|cluster_size_ema)",
+    # ]
+
+    def __init__(self, config: MinistralDLMConfig):
+        super().__init__(config)
+
+        self.mask_token_id = config.mask_token_id
+
+        diffusion_config = copy.deepcopy(config)
+        diffusion_config.diffusion_lm = True
+
+        use_flex = getattr(config, 'enable_self_spec', False)
+
+        if config.dlm_paradigm in ['block_diff', 'sbd_block_diff']:
+            diffusion_config.attn_class = MinistralFlexAttention
+        elif config.dlm_paradigm in ['bidirectional', 'autoregressive']:
+            diffusion_config.attn_class = MinistralFlexAttention if use_flex else Ministral3Attention
+            if config.dlm_paradigm == 'autoregressive':
+                diffusion_config.diffusion_lm = False
+        else:
+            raise ValueError(f"Unsupported DLM paradigm: {config.dlm_paradigm}")
+        
+        self.encoder = Ministral3Model(diffusion_config)
+        self.diffusion_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.vocab_size = config.vocab_size
+
+        self.current_iter_ratio = None
+
+        self.post_init()
+
+
+    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, new_embeddings):
+        self.diffusion_head = new_embeddings
+
+
+    def forward_process(self, input_ids, eps=1e-3, block_size=None, loss_mask=None):
+        b, l = input_ids.shape
+        device = input_ids.device
+
+        if self.config.dp_varying_mask_ratio:
+            # Enable different random seeds for each DP rank during sampling
+            import torch.distributed as dist
+            dp_rank = 0
+            if dist.is_initialized():
+                try:
+                    dp_rank = dist.get_rank()
+                except Exception:
+                    dp_rank = 0
+            # Use a local generator to avoid affecting global RNG state
+            generator = torch.Generator(device=device)
+            generator.manual_seed(torch.seed() + dp_rank)
+        else:
+            generator = None
+            
+        if self.config.adaptive_mask_rate:
+            assert block_size is not None
+
+            # --- simple linear window mapping ---
+            bs_min = getattr(self.config, "t_bs_min", 16)
+            bs_max = getattr(self.config, "t_bs_max", 128)
+            w = getattr(self.config, "t_window_width", 0.6)  # fixed width
+
+            # fraction in [0,1] (unclamped first)
+            frac = (float(block_size) - float(bs_min)) / max(1.0, float(bs_max - bs_min))
+            # upper bound decreases linearly from 1.0 -> 0.5
+            u_max = 1.0 - w * frac
+            # clamp to [0.6, 1.0] to handle bs outside [bs_min, bs_max]
+            u_max = max(0.6, min(1.0, u_max))
+            u_min = u_max - w  # ensures width = w
+
+            # sample t ~ Uniform(u_min, u_max)
+            t = u_min + (u_max - u_min) * torch.rand(b, device=device, generator=generator)
+        else:
+            t = torch.rand(b, device=device, generator=generator)
+        
+        p_mask = (1 - eps) * t + eps  # shape: (b,)
+        p_mask = p_mask[:, None].expand(-1, l)  # shape: (b, l)
+
+        masked_indices = torch.rand((b, l), device=device) < p_mask
+
+        if loss_mask is not None:
+            masked_indices[loss_mask == 0] = 0
+
+        noisy_batch = torch.where(masked_indices, self.mask_token_id, input_ids)        
+
+        return noisy_batch, masked_indices, p_mask
+
+
+    def forward_process_exp(
+        self,
+        input_ids: torch.Tensor,
+        eps: float = 1e-3,
+        block_size: int | None = None,
+        half_life_ratio: float = 0.25, # λ = ln 2 / (half_life_ratio·L)
+        loss_mask: Optional[torch.Tensor] = None,
+    ):
+        """
+        Two-stage corruption with optional per-block sampling.
+        • Stage 1:  m ~ U(eps, 1)   →   k = round(m · len)  (exact budget).
+        • Stage 2:  sample exactly k positions with weights
+                    w_i(m) = exp[ λ · (1−m) · i ]   (late-heavy when m→0,
+                                                     uniform when m→1).
+          If `block_size` is given, the procedure is run *independently*
+          inside each contiguous block of that length (last block may be shorter).
+          When block_size is provided, m is sampled per-block and p_mask is per-block.
+        Args
+        ----
+        input_ids : (B, L)  LongTensor
+        eps       : minimum corruption ratio
+        block_size: if not None, operate block-wise with per-block m sampling
+        half_life_ratio : controls steepness when m→0
+        """
+        B, L = input_ids.shape
+        device = input_ids.device
+        dtype  = torch.float32
+
+        masked_indices = torch.zeros((B, L), dtype=torch.bool, device=device)
+        p_mask = torch.zeros((B, L), dtype=dtype, device=device)
+
+        # ---------- Stage 1 & 2: whole-sentence or block-wise -------------------
+        for b in range(B):
+            if block_size is None:
+                # ---------- Per-batch sampling (original behavior) ----------
+                m = eps + (1.0 - eps) * torch.rand(1, device=device).item()   # scalar
+                k_tot = int(round(m * L))
+                k_tot = max(1, min(k_tot, L))  # clamp to [1, L]
+                
+                # Fill p_mask for this batch
+                p_mask[b, :] = m
+                
+                slope = 1.0 - m          # ∈ [0,1]; 0 ⇒ uniform, 1 ⇒ late-heavy
+                
+                # ------- single pool over the whole sentence -------------
+                lam_base = math.log(2.0) / (half_life_ratio * L) # base decay rate (λ when slope=1)
+
+                pos   = torch.arange(L, device=device, dtype=dtype)
+                log_w = (lam_base * slope * pos).clone()
+
+                masked_indices[b] = gumbel_topk(log_w, k_tot)
+
+            else:
+                # ---------- Per-block sampling ----------
+                num_blocks = math.ceil(L / block_size)
+                lam_base = math.log(2.0) / (half_life_ratio * block_size) # base decay rate (λ when slope=1)
+
+                for blk in range(num_blocks):
+                    start = blk * block_size
+                    end   = min((blk + 1) * block_size, L)
+                    blk_len = end - start
+
+                    # Sample m per block
+                    m_blk = eps + (1.0 - eps) * torch.rand(1, device=device).item()
+                    
+                    # Fill p_mask for this block
+                    p_mask[b, start:end] = m_blk
+                    
+                    # per-block budget
+                    k_blk = int(round(m_blk * blk_len))
+                    k_blk = max(0, min(k_blk, blk_len))
+                    if k_blk == 0:
+                        continue
+
+                    slope = 1.0 - m_blk          # ∈ [0,1]; 0 ⇒ uniform, 1 ⇒ late-heavy
+
+                    pos   = torch.arange(blk_len, device=device, dtype=dtype)
+                    log_w = lam_base * slope * pos
+
+                    blk_mask = gumbel_topk(log_w, k_blk)
+                    masked_indices[b, start:end] = blk_mask
+
+        if loss_mask is not None:
+            masked_indices[loss_mask == 0] = 0
+
+        noisy_batch = torch.where(masked_indices, self.mask_token_id, input_ids)
+        return noisy_batch, masked_indices, p_mask
+    
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor]   = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        labels: Optional[torch.LongTensor]       = None,
+        split_len: Optional[int]                 = None,
+        past_key_values: Optional[Cache]         = None,
+        block_size: Optional[int]                = None,
+        block_diff_ppl: bool                     = False,
+        eps: float                               = 1e-3,
+        is_teacher: bool                        = False,
+        masked_indices: Optional[torch.Tensor]   = None,
+        p_mask: Optional[torch.Tensor]           = None,
+        teacher_logits: Optional[torch.Tensor]   = None,
+        masked_indices_teacher: Optional[torch.Tensor] = None,
+        loss_mask: Optional[torch.Tensor] = None,
+        ce_loss_weight: float = 1.0,
+        output_last_hidden_states_only: bool = False,
+        skip_loss: bool = False,
+        inputs_embeds: torch.Tensor = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+
+        if input_ids is None:
+            if inputs_embeds is None:
+                raise ValueError("Either `input_ids` or `inputs_embeds` must be provided.")
+            batch_size, seq_len = inputs_embeds.shape[:2]
+            if labels is not None:
+                raise ValueError("`labels` training path requires `input_ids`.")
+        else:
+            batch_size, seq_len = input_ids.shape
+
+
+        if self.config.dlm_paradigm == 'bidirectional' or self.config.dlm_paradigm == 'autoregressive':
+            if labels is not None and torch.rand(1) < self.config.random_length_prob:
+                raise NotImplementedError("Random length training not yet implemented for bidirectional/autoregressive paradigms.")
+                random_length = torch.randint(2, input_ids.shape[1] + 1, (1,))
+                input_ids = input_ids[:, :random_length]
+                labels = labels[:, :random_length]
+                
+                if attention_mask is not None:
+                    attention_mask = attention_mask[:, :random_length]
+                if position_ids is not None:
+                    position_ids = position_ids[:, :random_length]
+                if loss_mask is not None:
+                    loss_mask = loss_mask[:, :random_length]
+
+        elif self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']:
+            if labels is not None and block_size is None:
+                if torch.rand(1) < self.config.random_length_prob:
+                    block_size = torch.randint(1, 8, (1,)).item() * 4  ## [4, 32] divisible by 4
+                else:
+                    block_size = self.config.block_size
+
+        else:
+            raise ValueError(f"Unknown dLM paradigm: {self.config.dlm_paradigm}")
+
+        if labels is not None and self.config.dlm_paradigm != 'autoregressive':
+            if masked_indices is not None:
+                # assert p_mask is not None
+
+                if loss_mask is not None:
+                    masked_indices[loss_mask == 0] = 0
+
+                noisy_inputs = torch.where(masked_indices, self.mask_token_id, input_ids)
+
+            else:
+                if self.config.tok_mask_half_life_ratio is not None:
+                    noisy_inputs, masked_indices, p_mask = self.forward_process_exp(input_ids, eps=eps, block_size=block_size, half_life_ratio=self.config.tok_mask_half_life_ratio, loss_mask=loss_mask)
+                else:
+                    noisy_inputs, masked_indices, p_mask = self.forward_process(input_ids, eps=eps, block_size=block_size, loss_mask=loss_mask)
+
+        else:
+            noisy_inputs = input_ids
+            masked_indices = None
+            p_mask = None
+
+        if self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']:
+            for layer in self.encoder.layers:
+                if hasattr(layer.self_attn, 'set_attention_mode'):
+                    layer.self_attn.set_attention_mode(self.config.dlm_paradigm, block_size=block_size)
+
+        input_ids_len = noisy_inputs.shape[1] if noisy_inputs is not None else seq_len
+        if labels is not None and self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']:
+            if position_ids is None:
+                position_ids = torch.arange(input_ids_len, device=noisy_inputs.device).unsqueeze(0)
+            noisy_inputs = torch.cat([noisy_inputs, input_ids], dim=1)
+
+        if block_diff_ppl:
+            if position_ids is None:
+                position_ids = torch.arange(input_ids_len // 2, device=noisy_inputs.device).unsqueeze(0)
+
+        enc_out  = self.encoder(
+            past_key_values=past_key_values,
+            input_ids=noisy_inputs,
+            inputs_embeds=inputs_embeds if noisy_inputs is None else None,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            is_training=(labels is not None) or (block_diff_ppl),
+            **kwargs,
+        )
+
+        if output_last_hidden_states_only:
+            return BaseModelOutput(last_hidden_state=enc_out.last_hidden_state)
+
+        logits = self.diffusion_head(enc_out.last_hidden_state)  # (batch, len_B, vocab)
+        causal_logits = None
+
+        if labels is not None and self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']:
+            if self.config.dlm_paradigm == 'sbd_block_diff':
+                causal_logits = logits[:, input_ids_len:]
+            else:
+                causal_logits = None
+
+            logits = logits[:, :input_ids_len]
+
+        loss = None
+        if labels is not None and not skip_loss:
+            if self.config.dlm_paradigm == 'autoregressive':
+                shift_logits = logits[..., :-1, :].contiguous()
+                shift_labels = labels[..., 1:].contiguous()
+                
+                if loss_mask is None:
+                    loss_fct = CrossEntropyLoss()
+                    shift_logits = shift_logits.view(-1, shift_logits.size(-1))
+                    shift_labels = shift_labels.view(-1)
+                    loss = loss_fct(shift_logits, shift_labels)
+
+                else:
+                    loss_mask = loss_mask[..., 1:].contiguous()
+
+                    loss_fct = CrossEntropyLoss(reduction='none')
+                    shift_logits = shift_logits.view(-1, shift_logits.size(-1))
+                    shift_labels = shift_labels.view(-1)
+                    shift_labels = shift_labels.to(shift_logits.device)
+                    
+                    token_losses = loss_fct(shift_logits, shift_labels)
+                                    
+                    flat_loss_mask = loss_mask.reshape(-1)
+                    loss = token_losses[flat_loss_mask == 1].sum() / flat_loss_mask.sum()
+
+            else:
+                # Handle DREAM vs LLADA style losses
+                if hasattr(self.config, 'dlm_type') and self.config.dlm_type == 'dream':
+                    logits = logits[..., :-1, :].contiguous()
+                    labels = labels[..., 1:].contiguous()
+                    masked_indices = masked_indices[:, 1:]
+                    p_mask = p_mask[:, 1:]
+
+                if self.config.ada_perm_ratio_per_block is not None:
+                    # Only compute loss for the top ada_perm_ratio_per_block tokens by confidence within each block
+                    block_size = self.config.block_size
+                    batch_size, seq_len = masked_indices.shape
+                    num_blocks = seq_len // block_size
+                    
+                    # Get the max logit (confidence) for each position
+                    confidence = logits.max(dim=-1).values.detach()  # (batch_size, seq_len)
+                    
+                    # Create a mask for tokens to include in loss
+                    selected_mask = torch.zeros_like(masked_indices, dtype=torch.bool)
+                    
+                    for blk in range(num_blocks):
+                        start = blk * block_size
+                        end = min((blk + 1) * block_size, seq_len)
+                        
+                        # Get masked indices within this block
+                        block_masked = masked_indices[:, start:end]  # (batch_size, block_len)
+                        block_confidence = confidence[:, start:end]  # (batch_size, block_len)
+                        
+                        for b in range(batch_size):
+                            # Get positions that are masked in this block for this batch
+                            masked_positions = torch.where(block_masked[b])[0]
+                            num_masked = len(masked_positions)
+                            
+                            if num_masked > 0:
+                                # Number of tokens to keep (top by confidence)
+                                k = min(max(1, int(block_size * self.config.ada_perm_ratio_per_block)), num_masked)
+                                
+                                # Get confidence values for masked positions
+                                masked_confidence = block_confidence[b, masked_positions]
+                                
+                                # Get indices of top-k confident tokens
+                                _, topk_indices = torch.topk(masked_confidence, k)
+                                selected_positions = masked_positions[topk_indices]
+                                
+                                # Mark these positions in the selected mask
+                                selected_mask[b, start + selected_positions] = True
+                    
+                    # Calculate loss only for selected positions
+                    token_loss = torch.nn.functional.cross_entropy(
+                        logits[selected_mask],
+                        labels[selected_mask],
+                        reduction='none'
+                    ) / p_mask[selected_mask]
+
+                    num_mask_tokens = selected_mask.sum()
+
+                else:
+                    # Calculate token-wise cross entropy loss for masked positions in B
+                    token_loss = torch.nn.functional.cross_entropy(
+                        logits[masked_indices], 
+                        labels[masked_indices], 
+                        reduction='none'
+                    ) / p_mask[masked_indices]
+
+                    num_mask_tokens = masked_indices.sum()
+
+                if self.config.global_loss_avg:
+                    loss = token_loss.sum()
+                else:
+                    loss = token_loss.sum() / num_mask_tokens
+                
+                if self.config.ada_dlm_loss_ratio is not None:
+                    assert self.current_iter_ratio is not None
+                    assert self.config.dlm_loss_weight is not None
+
+                    dlm_loss_weight = min(self.config.dlm_loss_weight, self.current_iter_ratio / self.config.ada_dlm_loss_ratio * self.config.dlm_loss_weight)
+                    loss = dlm_loss_weight * loss
+
+                elif self.config.dlm_loss_weight is not None:
+                    loss = self.config.dlm_loss_weight * loss
+
+                if self.config.dlm_paradigm == 'sbd_block_diff':
+                    causal_logits = causal_logits[..., :-1, :].contiguous()
+                    causal_logits = causal_logits.view(-1, causal_logits.size(-1))
+
+                    if hasattr(self.config, 'dlm_type') and self.config.dlm_type == 'dream':
+                        causal_labels = labels.view(-1)
+                    else: 
+                        causal_labels = labels[..., 1:].contiguous().view(-1)
+                    
+                    if self.config.global_loss_avg:
+                        loss_fct = CrossEntropyLoss(reduction='sum')
+                        ar_loss = loss_fct(causal_logits, causal_labels)
+
+                        self.loss_diffusion = loss.detach().item() / num_mask_tokens
+                        self.loss_ar = ar_loss.detach().item() / seq_len
+
+                        loss = loss + self.config.ar_loss_weight * ar_loss
+                    else:
+                        loss_fct = CrossEntropyLoss()
+                        ar_loss = loss_fct(causal_logits, causal_labels)
+
+                        self.loss_diffusion = loss.detach().item()
+                        self.loss_ar = ar_loss.detach().item()
+
+                        loss = loss + self.config.ar_loss_weight * ar_loss
+                
+                if self.config.global_loss_avg:
+                    if self.config.dlm_paradigm == 'sbd_block_diff':
+                        loss = (loss, num_mask_tokens + int(self.config.ar_loss_weight * seq_len))
+                    else: 
+                        loss = (loss, num_mask_tokens)
+
+        return MinistralDiffOutputWithPast(
+            loss=loss if not is_teacher else logits,
+            logits=logits,
+            causal_logits=causal_logits,
+            past_key_values=enc_out.past_key_values,
+            hidden_states=None,
+            attentions=None,
+        )
+
+
+    def generate_diffusion(self, prompt_ids, max_new_tokens=512, steps=512, block_length=32, shift_logits=False, threshold=0.9, causal_context=True, temperature=0, eos_token_id=None, max_thinking_tokens=None, end_think_token_id=None, step_ratio=None,prompt_embeds=None,**kwargs):
+        if prompt_embeds is None and prompt_ids is not None and torch.is_floating_point(prompt_ids):
+            prompt_embeds = prompt_ids
+            prompt_ids = None
+
+        if (prompt_ids is None) == (prompt_embeds is None):
+            raise ValueError("Exactly one of `prompt_ids` or `prompt_embeds` must be provided.")
+
+        if eos_token_id is None:
+            eos_token_id = getattr(self.config, 'eos_token_id', None)
+        if step_ratio is not None:
+            steps_per_block = int(block_length * step_ratio)
+            num_blocks = max_new_tokens // block_length
+            steps = steps_per_block * num_blocks
+        out_ids, nfe = generate_with_prefix_cache_block_diff(
+                        model=self,
+                        prompt=prompt_ids,
+                        prompt_embeds=prompt_embeds,
+                        gen_length=max_new_tokens,
+                        steps=steps,
+                        block_length=block_length,
+                        remasking="low_confidence",
+                        temperature=temperature,
+                        mask_id=self.mask_token_id,
+                        threshold=threshold,
+                        shift_logits=shift_logits,
+                        neg_entropy=False,
+                        causal_context=causal_context,
+                        eos_token_id=eos_token_id,
+                        max_thinking_tokens=max_thinking_tokens,
+                        end_think_token_id=end_think_token_id,
+                    )
+
+        return out_ids, nfe
+
+
+    @torch.no_grad()
+    def sbd_inference_diffusion_quadratic(
+        self,
+        clean_input_ids: Optional[torch.Tensor],
+        draft_input_ids: torch.Tensor,
+        block_length: int,
+        draft_only: bool = False,
+        past_key_values: Optional[Cache] = None,
+        use_cache: bool = False,
+    ):
+        enc_config = self.encoder.config
+        enc_config.use_sbd_objective = True
+        enc_config.block_length = block_length
+
+        if draft_only:
+            assert clean_input_ids is not None
+
+            if use_cache and past_key_values is None:
+                past_key_values = DynamicCache()
+
+            enc_config.self_spec_inference_mode = "default"
+            input_ids = torch.cat([clean_input_ids, draft_input_ids], dim=-1)
+            outputs = self.encoder(
+                input_ids=input_ids,
+                position_ids=None,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                is_training=False,
+            )
+
+            hidden_states = outputs.last_hidden_state
+            logits = self.diffusion_head(hidden_states)
+
+            past_key_values = getattr(outputs, "past_key_values", None)
+            if use_cache and past_key_values is not None:
+                _crop_dynamic_cache(past_key_values, clean_input_ids.shape[1])
+
+            return logits, past_key_values
+        else:
+            enc_config.self_spec_inference_mode = "quadratic"
+
+            draft_len = block_length * (block_length + 1)
+            draft_input_ids = torch.cat(
+                [
+                    draft_input_ids.view(-1, block_length, 1),
+                    torch.full(
+                        (draft_input_ids.shape[0], block_length, block_length),
+                        fill_value=self.config.mask_token_id,
+                        device=draft_input_ids.device,
+                    ),
+                ],
+                dim=-1,
+            ).view(-1, draft_len)
+
+            if use_cache:
+                assert past_key_values is not None, (
+                    "Past key values should be provided when using cache, e.g. run draft_only=True first."
+                )
+                assert clean_input_ids is None, (
+                    "Clean input ids should already be in cache, thus none should be provided."
+                )
+                clean_len = past_key_values.get_seq_length()
+                input_ids = draft_input_ids
+            else:
+                clean_len = clean_input_ids.shape[1]
+                input_ids = torch.cat([clean_input_ids, draft_input_ids], dim=-1)
+
+            per_block_position_ids = torch.arange(
+                clean_len, clean_len + block_length + 1, device=draft_input_ids.device
+            )[None,].repeat(block_length, 1)
+            per_block_position_ids += torch.arange(block_length, device=draft_input_ids.device).view(-1, 1)
+
+            if use_cache:
+                position_ids = per_block_position_ids.view(-1)[None,]
+            else:
+                clean_position_ids = torch.arange(clean_len, device=draft_input_ids.device)
+                position_ids = torch.cat([clean_position_ids, per_block_position_ids.view(-1)], dim=-1)[None,]
+
+            outputs = self.encoder(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                is_training=False,
+            )
+
+            hidden_states = outputs.last_hidden_state
+            logits = self.diffusion_head(hidden_states)
+            past_key_values = getattr(outputs, "past_key_values", None)
+
+            if use_cache and past_key_values is not None:
+                _extract_draft_kv_cache(past_key_values, clean_len, block_length)
+
+            return logits, past_key_values
+
+
+    @torch.no_grad()
+    def ar_generate(
+        self,
+        prompt_ids: torch.Tensor,
+        max_new_tokens: int = 128,
+        temperature: float = 0.0,
+        eos_token_id: Optional[int] = None,
+        max_thinking_tokens: Optional[int] = None,
+        end_think_token_id: Optional[int] = None,
+    ) -> tuple:
+        """Autoregressive generation calling the encoder directly (injected by build_hf_tidar_repo).
+
+        Bypasses MinistralDiffEncoderModel.forward() to avoid diffusion-specific
+        code paths. Calls self.encoder (Ministral3Model) with explicit cache_position,
+        position_ids, and use_cache so the KV cache and causal masking behave
+        identically to MistralForCausalLM / vLLM.
+
+        Returns:
+            (output_ids, nfe) where output_ids includes the prompt.
+        """
+        for layer in self.encoder.layers:
+            if hasattr(layer.self_attn, 'diffusion_lm'):
+                layer.self_attn.diffusion_lm = False
+
+        if eos_token_id is None:
+            eos_token_id = getattr(self.config, 'eos_token_id', None)
+
+        device = prompt_ids.device
+        batch_size, prompt_len = prompt_ids.shape
+
+        past_key_values = DynamicCache()
+        cache_position = torch.arange(prompt_len, device=device)
+        position_ids = cache_position.unsqueeze(0).expand(batch_size, -1)
+
+        enc_out = self.encoder(
+            input_ids=prompt_ids,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=True,
+            cache_position=cache_position,
+        )
+        past_key_values = enc_out.past_key_values
+        next_logit = self.diffusion_head(enc_out.last_hidden_state[:, -1:, :]).squeeze(1)
+
+        generated_tokens = []
+        nfe = 0
+
+        for step in range(max_new_tokens):
+            nfe += 1
+
+            if temperature > 0:
+                probs = torch.softmax(next_logit / temperature, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+            else:
+                next_token = torch.argmax(next_logit, dim=-1, keepdim=True)
+
+            # ---- thinking budget enforcement ----
+            if end_think_token_id is not None and max_thinking_tokens is not None:
+                if step >= max_thinking_tokens:
+                    if generated_tokens:
+                        gen_tensor = torch.cat(generated_tokens, dim=1)
+                        has_end_think = (gen_tensor == end_think_token_id).any(dim=1)
+                    else:
+                        has_end_think = torch.zeros(batch_size, dtype=torch.bool, device=device)
+                    for b in range(batch_size):
+                        if not has_end_think[b]:
+                            next_token[b] = end_think_token_id
+
+            generated_tokens.append(next_token)
+
+            if eos_token_id is not None and (next_token == eos_token_id).all():
+                break
+
+            if step < max_new_tokens - 1:
+                cur_pos = prompt_len + step
+                step_cache_pos = torch.tensor([cur_pos], device=device)
+                step_pos_ids = step_cache_pos.unsqueeze(0).expand(batch_size, -1)
+
+                enc_out = self.encoder(
+                    input_ids=next_token,
+                    position_ids=step_pos_ids,
+                    past_key_values=past_key_values,
+                    use_cache=True,
+                    cache_position=step_cache_pos,
+                )
+                past_key_values = enc_out.past_key_values
+                next_logit = self.diffusion_head(enc_out.last_hidden_state[:, -1:, :]).squeeze(1)
+
+        all_generated = torch.cat(generated_tokens, dim=1)
+        output_ids = torch.cat([prompt_ids, all_generated], dim=1)
+        return output_ids, nfe
+
+
+    @torch.no_grad()
+    def self_spec_generate(
+        self,
+        prompt_ids: torch.Tensor,
+        max_new_tokens: int = 128,
+        steps: int = 128,
+        block_length: int = 16,
+        ar_mix_weight: Optional[float] = None,
+        temperature: float = 0.0,
+        mask_token_id: Optional[int] = None,
+        eos_token_id: Optional[int] = None,
+        max_thinking_tokens: Optional[int] = None,
+        end_think_token_id: Optional[int] = None,
+    ):
+        self.config.use_sbd_objective = True
+        self.config.dlm_paradigm = "sbd"
+
+        if prompt_ids.shape[0] != 1:
+            raise ValueError("Self speculation quadratic decoding currently requires batch_size == 1")
+
+        token_mask_id = mask_token_id if mask_token_id is not None else self.config.mask_token_id
+        if eos_token_id is None:
+            eos_token_id = getattr(self.config, "eos_token_id", None)
+
+        x = torch.full(
+            (1, prompt_ids.shape[1] + max_new_tokens + block_length * 2),
+            token_mask_id,
+            dtype=torch.long,
+            device=prompt_ids.device,
+        )
+        x[:, : prompt_ids.shape[1]] = prompt_ids.clone()
+
+        if max_new_tokens % block_length != 0:
+            raise ValueError("max_new_tokens must be divisible by block_length")
+        num_blocks = max_new_tokens // block_length
+        if steps % num_blocks != 0:
+            raise ValueError("steps must be divisible by (max_new_tokens // block_length)")
+
+        prompt_len = prompt_ids.shape[1]
+        nfe = 0
+        nfe += 1
+        logits, past_key_values = self.sbd_inference_diffusion_quadratic(
+            clean_input_ids=x[:, :prompt_len],
+            draft_input_ids=x[:, prompt_len : prompt_len + block_length],
+            block_length=block_length,
+            draft_only=True,
+            use_cache=True,
+        )
+
+        logits_proposal = logits[:, prompt_len - 1 : prompt_len + block_length]
+        logits_proposal[:, 1] = logits_proposal[:, 0]
+        logits_proposal = logits_proposal[:, 1:]
+        x0_proposal = torch.argmax(logits_proposal, dim=-1)
+        x[:, prompt_len : prompt_len + block_length] = x0_proposal
+
+        total_accept_token = 0
+        while True:
+            nfe += 1
+            block_start = prompt_len + total_accept_token
+            block_end = block_start + block_length
+            draft_input_ids = x[:, block_start:block_end]
+
+            logits, past_key_values = self.sbd_inference_diffusion_quadratic(
+                clean_input_ids=None,
+                draft_input_ids=draft_input_ids,
+                block_length=block_length,
+                draft_only=False,
+                past_key_values=past_key_values,
+                use_cache=True,
+            )
+
+            useful_token_logits = logits.view(1, block_length, block_length + 1, -1)
+            if ar_mix_weight is None:
+                useful_token_logits[:, :, 1] = useful_token_logits[:, :, 0]
+            else:
+                if not (0.0 <= ar_mix_weight <= 1.0):
+                    raise ValueError("ar_mix_weight must be between 0 and 1")
+                mix_logits = useful_token_logits[:, :, 0] * ar_mix_weight + useful_token_logits[:, :, 1] * (1 - ar_mix_weight)
+                useful_token_logits[:, :, 0] = mix_logits
+                useful_token_logits[:, :, 1] = mix_logits
+
+            if temperature > 0:
+                useful_token_logits = useful_token_logits / temperature
+
+            useful_token_pred = torch.argmax(useful_token_logits, dim=-1)
+            new_draft_input_ids = useful_token_pred[:, 0, 1:]
+            accept_cnt = 1
+
+            while accept_cnt < block_length:
+                if useful_token_pred[:, accept_cnt - 1, 0].item() != draft_input_ids[:, accept_cnt].item():
+                    break
+                new_draft_input_ids = useful_token_pred[:, accept_cnt, 1:]
+                accept_cnt += 1
+
+            x[:, block_start : block_start + accept_cnt] = draft_input_ids[:, :accept_cnt]
+
+            # EoS early stopping: all accepted tokens are finalized left-to-right,
+            # so if any is EoS we can truncate and return immediately.
+            if eos_token_id is not None:
+                accepted = x[0, block_start : block_start + accept_cnt]
+                eos_positions = (accepted == eos_token_id).nonzero(as_tuple=True)[0]
+                if len(eos_positions) > 0:
+                    first_eos_rel = eos_positions[0].item()
+                    total_accept_token += first_eos_rel + 1
+                    output_end = prompt_len + total_accept_token
+                    return x[:, :output_end], nfe
+
+            x[:, block_start + accept_cnt : block_start + accept_cnt + block_length] = new_draft_input_ids
+            past_key_values.crop(block_start + accept_cnt)
+
+            # ---- thinking budget enforcement ----
+            # Insert end_think as the first token of the next draft block,
+            # shifting all subsequent tokens right by 1 (discarding the last).
+            # The first draft token is always accepted unconditionally, so
+            # end_think is guaranteed to be finalized in the next iteration
+            # without needing to re-encode or touch the KV cache.
+            if end_think_token_id is not None and max_thinking_tokens is not None:
+                tokens_so_far = total_accept_token + accept_cnt
+                if tokens_so_far > max_thinking_tokens:
+                    gen_so_far = x[0, prompt_len : prompt_len + tokens_so_far]
+                    has_end_think = (gen_so_far == end_think_token_id).any()
+                    if not has_end_think:
+                        insert_pos = block_start + accept_cnt
+                        x[0, insert_pos + 1:] = x[0, insert_pos:-1].clone()
+                        x[0, insert_pos] = end_think_token_id
+
+            total_accept_token += accept_cnt
+
+            if total_accept_token >= max_new_tokens:
+                break
+
+        return x[:, : -(block_length * 2)], nfe
+
+
+    @torch.no_grad()
+    def linear_spec_generate(
+        self,
+        prompt_ids: torch.Tensor,
+        max_new_tokens: int = 128,
+        block_length: int = 32,
+        temperature: float = 0.0,
+        mask_token_id: Optional[int] = None,
+        eos_token_id: Optional[int] = None,
+        max_thinking_tokens: Optional[int] = None,
+        end_think_token_id: Optional[int] = None,
+        threshold: float = 0.0,
+    ):
+        """Linear speculative decoding: diffusion draft + AR verification.
+
+        Each step:
+          1. Draft: forward [last_accepted, mask, ...] with bidirectional attention
+             (diffusion_lm=True, use_cache=False).  Shift AR logits to get
+             per-position predictions; apply confidence filtering.
+          2. Verify: forward the drafted block with causal attention
+             (diffusion_lm=False, use_cache=True, use_causal_mask=True).
+             Accept consecutive AR-matching tokens plus one bonus token.
+
+        Args:
+            prompt_ids: Input token IDs of shape (1, prompt_len).
+            max_new_tokens: Maximum number of tokens to generate.
+            block_length: Number of tokens per draft/verify block.
+            temperature: Sampling temperature (0 = greedy).
+            mask_token_id: Override for config.mask_token_id.
+            eos_token_id: Override for config.eos_token_id.
+            max_thinking_tokens: Budget for thinking tokens before forcing end_think.
+            end_think_token_id: Token ID inserted when thinking budget is exceeded.
+            threshold: Confidence threshold for accepting draft predictions.
+
+        Returns:
+            (output_ids, nfe): output_ids includes the prompt; nfe is the number
+            of forward evaluations (matching self_spec_generate interface).
+        """
+        if prompt_ids.shape[0] != 1:
+            raise ValueError("Linear speculative decoding requires batch_size == 1")
+
+        token_mask_id = mask_token_id if mask_token_id is not None else self.config.mask_token_id
+        if eos_token_id is None:
+            eos_token_id = getattr(self.config, "eos_token_id", None)
+
+        device = prompt_ids.device
+        prompt_len = prompt_ids.shape[1]
+        dream_style = getattr(self.config, 'dlm_type', 'llada') == 'dream'
+
+        def _set_diffusion_lm(val: bool):
+            for layer in self.encoder.layers:
+                if hasattr(layer.self_attn, 'diffusion_lm'):
+                    layer.self_attn.diffusion_lm = val
+
+        # ===== Prefill (causal) =====
+        _set_diffusion_lm(False)
+
+        enc_out = self.encoder(
+            input_ids=prompt_ids,
+            past_key_values=DynamicCache(),
+            use_cache=True,
+            use_causal_mask=True,
+        )
+        past_key_values = enc_out.past_key_values
+        last_logit = self.diffusion_head(enc_out.last_hidden_state[:, -1:, :]).squeeze(1)
+        nfe = 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 eos_token_id is not None and next_token.item() == eos_token_id:
+            output_ids = torch.cat([prompt_ids, next_token], dim=1)
+            return output_ids, nfe
+
+        generated = [next_token]
+        total_gen = 1
+
+        # ===== Main loop =====
+        while total_gen < max_new_tokens:
+            cache_len = past_key_values.get_seq_length()
+
+            block = torch.full(
+                (1, block_length), token_mask_id, dtype=torch.long, device=device
+            )
+            block[0, 0] = next_token.item()
+
+            # -------- Draft (bidirectional, don't update cache) --------
+            _set_diffusion_lm(True)
+            enc_out = self.encoder(
+                input_ids=block,
+                past_key_values=past_key_values,
+                use_cache=False,
+            )
+            nfe += 1
+
+            draft_logits = self.diffusion_head(enc_out.last_hidden_state)
+            if dream_style:
+                # DREAM: logit[i] predicts position i+1 → shift to self-prediction
+                draft_logits = torch.cat(
+                    [draft_logits[:, :1, :], draft_logits[:, :-1, :]], dim=1
+                )
+            # LLaDA: logit[i] already predicts position i → no shift needed
+
+            if temperature > 0:
+                draft_probs = torch.softmax(draft_logits / temperature, dim=-1)
+                draft_tokens = torch.multinomial(
+                    draft_probs.view(-1, draft_probs.shape[-1]), num_samples=1
+                ).view(1, block_length)
+            else:
+                draft_tokens = draft_logits.argmax(dim=-1)
+                draft_probs = torch.softmax(draft_logits, dim=-1)
+
+            draft_conf = torch.gather(
+                draft_probs, -1, draft_tokens.unsqueeze(-1)
+            ).squeeze(-1)
+
+            is_mask = block == token_mask_id
+            draft_conf = torch.where(is_mask, draft_conf, -torch.inf)
+            unmask = draft_conf > threshold
+
+            if unmask.sum() > 0:
+                block[unmask] = draft_tokens[unmask]
+            else:
+                raise AssertionError(
+                    "No mask token above threshold for prediction"
+                )
+
+            # -------- Verify (causal, update cache) --------
+            _set_diffusion_lm(False)
+            enc_out = self.encoder(
+                input_ids=block,
+                past_key_values=past_key_values,
+                use_cache=True,
+                use_causal_mask=True,
+            )
+            past_key_values = enc_out.past_key_values
+            nfe += 1
+
+            verify_logits = self.diffusion_head(enc_out.last_hidden_state)
+            if temperature > 0:
+                verify_probs = torch.softmax(verify_logits / temperature, dim=-1)
+                ar_tokens = torch.multinomial(
+                    verify_probs.view(-1, verify_probs.shape[-1]), num_samples=1
+                ).view(1, block_length)
+            else:
+                ar_tokens = verify_logits.argmax(dim=-1)
+
+            accepted = 0
+            for i in range(block_length - 1):
+                if ar_tokens[0, i].item() == block[0, i + 1].item():
+                    accepted += 1
+                else:
+                    break
+            accepted += 1  # bonus token from AR verification
+
+            accepted_toks = ar_tokens[:, :accepted]
+            generated.append(accepted_toks)
+            total_gen += accepted
+
+            _crop_dynamic_cache(past_key_values, cache_len + accepted)
+
+            next_token = ar_tokens[:, accepted - 1 : accepted]
+
+            # -------- EOS check --------
+            if eos_token_id is not None:
+                eos_pos = (accepted_toks[0] == eos_token_id).nonzero(as_tuple=True)[0]
+                if len(eos_pos) > 0:
+                    first_eos = eos_pos[0].item()
+                    generated[-1] = accepted_toks[:, : first_eos + 1]
+                    total_gen = total_gen - accepted + first_eos + 1
+                    break
+
+            # -------- Thinking budget enforcement --------
+            if end_think_token_id is not None and max_thinking_tokens is not None:
+                if total_gen > max_thinking_tokens:
+                    all_gen = torch.cat(generated, dim=1)
+                    if not (all_gen == end_think_token_id).any():
+                        next_token = torch.tensor(
+                            [[end_think_token_id]], device=device
+                        )
+
+            if total_gen >= max_new_tokens:
+                break
+
+        all_generated = torch.cat(generated, dim=1)
+        output_ids = torch.cat([prompt_ids, all_generated], dim=1)
+        
+        return output_ids, nfe

modeling_nemotron_labs_diffusion_image.py

@@ -0,0 +1,840 @@
+import copy
+import math
+import random
+import time
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.distributions as dists
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.resnet import Downsample2D, Upsample2D
+from einops import rearrange
+from PIL import Image
+from tqdm.auto import tqdm
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
+from transformers.generation.utils import GenerateOutput
+
+from .configuration_nemotron_labs_diffusion_image import NemotronLabsDiffusionImageConfig
+from .modeling_ministral import Ministral3Model
+from .modeling_ministral_dlm import MinistralDiffEncoderModel
+# The imports below are not used directly but MUST stay here so that HF's
+# dynamic-module cache scanner (regex: r"from\.X import") copies every
+# transitive dependency into the hash directory.
+from .chat_utils import generate_with_prefix_cache_block_diff as _gcbd  # noqa: F401
+from .nemotron_diffusion_image_utils import maybe_truncate_last_dim as _mtld  # noqa: F401
+from .configuration_ministral_dlm import MinistralDLMConfig as _MinistralDLMConfig  # noqa: F401
+
+
+def _resolve_local_path(path_value: str) -> Path:
+    base_dir = Path(__file__).resolve().parent
+    candidate = Path(path_value)
+    if candidate.is_absolute():
+        return candidate
+    return (base_dir / candidate).resolve()
+
+
+def _load_vqvae_from_local(vqvae_path: Path):
+    """Load Emu3p5VisionVQModel directly from local files.
+
+    Bypasses AutoModel.from_pretrained because newer huggingface_hub versions
+    validate the path argument as a HF repo ID, rejecting absolute local paths.
+    """
+    import importlib.util
+    import json
+    import sys
+    import types
+
+    from safetensors.torch import load_file
+
+    pkg = f"_emu3_vqvae_{vqvae_path.name}"
+
+    # Create a package namespace so relative imports inside the vqvae files work
+    pkg_mod = types.ModuleType(pkg)
+    pkg_mod.__path__ = [str(vqvae_path)]
+    pkg_mod.__package__ = pkg
+    sys.modules[pkg] = pkg_mod
+
+    def _load_mod(mod_name, filename):
+        spec = importlib.util.spec_from_file_location(
+            f"{pkg}.{mod_name}",
+            vqvae_path / filename,
+            submodule_search_locations=[str(vqvae_path)],
+        )
+        mod = importlib.util.module_from_spec(spec)
+        mod.__package__ = pkg
+        sys.modules[f"{pkg}.{mod_name}"] = mod
+        spec.loader.exec_module(mod)
+        return mod
+
+    cfg_mod = _load_mod("configuration_emu3p5visionvq", "configuration_emu3p5visionvq.py")
+    mdl_mod = _load_mod("modeling_emu3p5visionvq", "modeling_emu3p5visionvq.py")
+
+    with open(vqvae_path / "config.json") as f:
+        cfg_data = json.load(f)
+
+    # PretrainedConfig accepts and stores arbitrary kwargs, so pass everything
+    vqvae_config = cfg_mod.Emu3p5VisionVQConfig(**cfg_data)
+    model = mdl_mod.Emu3p5VisionVQModel(vqvae_config)
+
+    sf_path = vqvae_path / "model.safetensors"
+    state_dict = load_file(str(sf_path))
+    model.load_state_dict(state_dict)
+
+    return model
+
+
+def _preprocess_emu3_image(image):
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    image = np.asarray(image, dtype=np.float32)
+    image = image / 127.5 - 1.0
+    return torch.from_numpy(image).permute(2, 0, 1).float()
+
+
+class Emu3ImageProcessor:
+    def preprocess(self, image):
+        return _preprocess_emu3_image(image).unsqueeze(0)
+
+
+# ---------------------------------------------------------------------------
+# T2I helpers (inlined — no llava imports required)
+# ---------------------------------------------------------------------------
+
+class _NC:
+    """Token constants for the Ministral diffusion model."""
+    reserve_id           = 18
+    reserve_id_token     = '<SPECIAL_18>'
+    reserve_id_enc       = 19
+    reserve_id_token_enc = '<SPECIAL_19>'
+    mask_id              = 100
+    eos_id               = 11
+    gen_im_start_token   = '<SPECIAL_21>'
+    gen_im_end_token     = '<SPECIAL_22>'
+
+
+def _pad_along_last_dim(tensor: torch.Tensor, size: int) -> torch.Tensor:
+    pad_size = size - tensor.shape[-1]
+    if pad_size <= 0:
+        return tensor
+    padding = torch.zeros(*tensor.shape[:-1], pad_size,
+                          dtype=tensor.dtype, device=tensor.device)
+    return torch.cat([tensor, padding], dim=-1)
+
+
+def _maybe_truncate_last_dim(tensor: torch.Tensor, size: int) -> torch.Tensor:
+    if size >= tensor.shape[-1]:
+        return tensor
+    return tensor[..., :size]
+
+
+_INT_MAX = 1_000_000
+
+
+def _t2i_wte(model, x, gen_shape=None, x_gen=None,
+             inputs_embeds_curr=None, new_token_mask=None):
+    """Embed text tokens and splice in gen-token embeddings."""
+    assert x_gen is not None
+    if new_token_mask is None or not torch.any(new_token_mask):
+        if inputs_embeds_curr is None:
+            return model.embed_tokens(x), new_token_mask
+        return inputs_embeds_curr, new_token_mask
+    gen_latents_comp_embeds = model.call_gen_embedding(x_gen, gen_shape)
+    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])
+    )
+    return inputs_embeds_curr, new_token_mask
+
+
+def _t2i_get_logits(model, input_embeddings, modality_indices,
+                    past_key_values=None, gen_shape=None, timesteps=None,
+                    input_modality_indices=None):
+    """Forward pass returning generation logits only."""
+    if input_modality_indices is None:
+        input_modality_indices = modality_indices
+    output = model(
+        None,
+        input_embeddings=input_embeddings,
+        modality_indices=input_modality_indices,
+        past_key_values=past_key_values,
+        is_training=False,
+        overwrite_attn_impl='flash_attn',
+    )
+    hidden_states = output.last_hidden_state
+    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 = int(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
+
+
+def _cosine_schedule_2(x):
+    x = 1.0 - np.clip(x, 0.0, 1.0)
+    return np.cos(np.pi * x / 2.0)
+
+
+def _exp_schedule(x):
+    z = (1.0 - np.exp(-5.0 * x)) / (1.0 - np.exp(-5.0))
+    return np.clip(z, 0.0001, 1.0)
+
+
+def _logit_normal_schedule(shift, sigmas):
+    return shift * sigmas / (1.0 + (shift - 1.0) * sigmas)
+
+
+def _get_num_transfer_tokens(mask_index: torch.Tensor, steps: int,
+                              schedule: str = 'shift',
+                              shift: int = 3) -> torch.Tensor:
+    mask_num = mask_index.sum(dim=1, keepdim=True)
+    steps = int(min(steps, mask_num[0]))
+    t = torch.linspace(0, 1, steps + 1)
+    sigmas = _logit_normal_schedule(shift, t)
+    sigmas = sigmas.to(mask_num.device)
+    num_transfer_tokens = torch.zeros(mask_num.size(0), steps,
+                                      device=mask_index.device, dtype=torch.int64)
+    for i in range(mask_num.size(0)):
+        sigmas_sample = (sigmas * mask_num[i]).to(torch.int64)
+        sigmas_sample = sigmas_sample[1:] - sigmas_sample[:-1]
+        sigmas_sample = torch.clamp(sigmas_sample, 1, None)
+        delta = sigmas_sample.sum() - mask_num[i]
+        assert delta >= 0
+        j = 0
+        while delta > 0:
+            j = j % len(sigmas_sample)
+            if sigmas_sample[j] == 1:
+                j += 1
+                continue
+            delta -= 1
+            sigmas_sample[j] -= 1
+            j += 1
+        assert sigmas_sample.sum() == mask_num[i]
+        num_transfer_tokens[i] = sigmas_sample
+    return num_transfer_tokens.flip(-1)
+
+
+class _MinistralConv:
+    """Minimal CHATML conversation template for the Ministral model."""
+    _SYSTEM = (
+        "<|im_start|>system\n"
+        "You 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."
+    )
+    _SEP = "<|im_end|>"
+    _ROLES = ("<|im_start|>user", "<|im_start|>assistant")
+
+    def __init__(self):
+        self.messages: List[Tuple[str, Optional[str]]] = []
+
+    def append_message(self, role: str, message: Optional[str]) -> None:
+        self.messages.append((role, message))
+
+    def get_prompt(self) -> str:
+        ret = self._SYSTEM + self._SEP + "\n"
+        for role, message in self.messages:
+            if message is not None:
+                ret += role + "\n" + message + self._SEP + "\n"
+            else:
+                ret += role + "\n"
+        return ret
+
+    @property
+    def roles(self):
+        return self._ROLES
+
+
+_IMAGE_TOKEN_INDEX = -200
+
+
+def _tokenizer_image_token(prompt: str, tokenizer,
+                            return_tensors: str = "pt") -> torch.Tensor:
+    """Tokenise a prompt that may contain <image> placeholder tokens."""
+    prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split("<image>")]
+
+    def _insert_sep(X, sep):
+        return [e for pair in zip(X, [sep] * len(X)) for e in pair][:-1]
+
+    input_ids: List[int] = []
+    offset = 0
+    if (prompt_chunks and prompt_chunks[0]
+            and prompt_chunks[0][0] == tokenizer.bos_token_id):
+        offset = 1
+        input_ids.append(prompt_chunks[0][0])
+    for x in _insert_sep(prompt_chunks, [_IMAGE_TOKEN_INDEX] * (offset + 1)):
+        input_ids.extend(x[offset:])
+    ids = torch.tensor(input_ids, dtype=torch.long)
+    if return_tensors == "pt":
+        return ids
+    return ids.tolist()
+
+
+def _stratified_random(n: int = 64, seed: Optional[int] = None,
+                        shuffle_blocks: bool = True) -> List[int]:
+    """Progressive Multi-Jittered ordering over an n×n integer grid."""
+    if n <= 0 or (n & (n - 1)) != 0:
+        raise ValueError("n must be a positive power of two")
+    rng = random.Random(seed)
+    occupied = [[False] * n for _ in range(n)]
+    seq: List[int] = []
+    blocks: List[Tuple[int, int, int]] = [(0, 0, n)]
+
+    def _has(x0, y0, size):
+        for yy in range(y0, y0 + size):
+            for xx in range(x0, x0 + size):
+                if occupied[yy][xx]:
+                    return True
+        return False
+
+    def _place(x0, y0, size):
+        x, y, attempts = rng.randrange(x0, x0 + size), rng.randrange(y0, y0 + size), 0
+        while occupied[y][x]:
+            x, y = rng.randrange(x0, x0 + size), rng.randrange(y0, y0 + size)
+            attempts += 1
+            if attempts > 10000:
+                raise RuntimeError("placement failed")
+        occupied[y][x] = True
+        seq.append(y * n + x)
+
+    size = n
+    while size > 1:
+        half = size // 2
+        children = [(x0 + dx, y0 + dy, half)
+                    for (x0, y0, _) in blocks
+                    for dx, dy in [(0, 0), (half, 0), (0, half), (half, half)]]
+        if shuffle_blocks:
+            rng.shuffle(children)
+        for (x0, y0, s) in children:
+            if not _has(x0, y0, s):
+                _place(x0, y0, s)
+        blocks = children
+        size = half
+
+    remaining = [y * n + x for y in range(n) for x in range(n) if not occupied[y][x]]
+    rng.shuffle(remaining)
+    seq.extend(remaining)
+    return seq
+
+
+def _gumbel_noise(t: torch.Tensor) -> torch.Tensor:
+    noise = torch.zeros_like(t).uniform_(0, 1)
+    return -torch.log(-torch.log(noise))
+
+
+class SimpleUVitBlock(nn.Module):
+    def __init__(self, channels, downsample: bool, upsample: bool):
+        super().__init__()
+        self.downsample = None
+        self.upsample = None
+        if downsample:
+            self.downsample = Downsample2D(
+                channels,
+                use_conv=True,
+                padding=0,
+                name="Conv2d_0",
+                kernel_size=2,
+                norm_type="rms_norm",
+                eps=1e-6,
+                elementwise_affine=True,
+                bias=False,
+                out_channels=channels,
+            )
+        if upsample:
+            self.upsample = Upsample2D(
+                channels,
+                use_conv_transpose=True,
+                kernel_size=2,
+                padding=0,
+                name="conv",
+                norm_type="rms_norm",
+                eps=1e-6,
+                elementwise_affine=True,
+                bias=False,
+                interpolate=False,
+                out_channels=channels,
+            )
+
+    def forward(self, hidden_states, size):
+        hidden_states = rearrange(hidden_states, "b (h w) d -> b d h w", h=size[0], w=size[1])
+        if self.downsample is not None:
+            hidden_states = self.downsample(hidden_states)
+        if self.upsample is not None:
+            hidden_states = self.upsample(hidden_states)
+        return rearrange(hidden_states, "b d h w -> b (h w) d")
+
+
+class NemotronLabsDiffusionImageModel(Ministral3Model):
+    config_class = NemotronLabsDiffusionImageConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.build_vqvae(config)
+        self.build_gen_embedding(config)
+        self.image_newline = nn.Parameter(torch.empty(config.hidden_size))
+
+    def build_vqvae(self, config):
+        mm_vqvae = getattr(config, "mm_vqvae", "emu3_vqvae")
+        # Prefer model_dir/_name_or_path so this works both from the release dir
+        # and when loaded via trust_remote_code (where __file__ is the HF cache).
+        model_dir = Path(getattr(config, "_name_or_path", ""))
+        if model_dir.is_dir():
+            vqvae_path = (model_dir / mm_vqvae).resolve()
+        else:
+            vqvae_path = _resolve_local_path(mm_vqvae)
+        self.vqvae = _load_vqvae_from_local(vqvae_path)
+        self.vqvae.eval()
+        self.vqvae.requires_grad_(False)
+        self.image_processor_gen = Emu3ImageProcessor()
+
+    def build_gen_embedding(self, config):
+        self.downsample_gen = SimpleUVitBlock(config.d_model_gen, downsample=True, upsample=False) if config.downsample else None
+        self.upsample_gen = SimpleUVitBlock(config.d_model_gen, downsample=False, upsample=True) if config.downsample else None
+        self.gen_embedding = nn.Embedding(self.vqvae.config.codebook_size + 256, config.d_model_gen)
+        self.gen_predictor = nn.Linear(config.d_model_gen, self.vqvae.config.codebook_size, bias=config.include_bias)
+        self.gen_embedding_2 = None
+        self.gen_predictor_2 = None
+
+    def call_gen_embedding(self, token_ids, gen_shape=None, enc=False):
+        del enc
+        hidden_states = self.gen_embedding(token_ids)
+        if self.downsample_gen is not None:
+            hidden_states = self.downsample_gen(hidden_states, gen_shape)
+        return hidden_states
+
+    def call_gen_predictor(self, gen_hidden_states, gen_shape=None, timesteps=None, labels=None):
+        del timesteps, labels
+        if self.upsample_gen is not None:
+            seq_len_per_image = (gen_shape[0] // 2) * (gen_shape[1] // 2)
+            gen_hidden_states = self.upsample_gen(
+                gen_hidden_states.view(-1, seq_len_per_image, gen_hidden_states.shape[-1]),
+                (gen_shape[0] // 2, gen_shape[1] // 2),
+            )
+            gen_hidden_states = gen_hidden_states.flatten(0, 1)
+        return self.gen_predictor(gen_hidden_states)
+
+    def encode_image_gen(self, images, enc=False):
+        batch_size = images.shape[0]
+        # Emu3p5VisionVQModel.encode does not accept mini_batch_size;
+        # implement manual chunking for large images.
+        if images.shape[2] > 256 and batch_size > 2:
+            mini_bs = 2
+            qs, idxs = [], []
+            for i in range(0, batch_size, mini_bs):
+                q, _, (_, _, idx) = self.vqvae.encode(images[i:i + mini_bs])
+                qs.append(q)
+                idxs.append(idx)
+            quantized = torch.cat(qs, dim=0)
+            indices = torch.cat(idxs, dim=0)
+        else:
+            quantized, _, (_, _, indices) = self.vqvae.encode(images)
+        latent_height, latent_width = quantized.shape[-2], quantized.shape[-1]
+        return indices.reshape(batch_size, -1), (latent_height, latent_width)
+
+    @torch.no_grad()
+    def decode_image_gen(self, images_to_decode, height, width):
+        vae_scale_factor = 16
+        indices = self.vqvae.quantize.get_codebook_entry(images_to_decode)
+        indices = rearrange(
+            indices,
+            "b (h w) d -> b d h w",
+            h=height // vae_scale_factor,
+            w=width // vae_scale_factor,
+        )
+        # Emu3p5VisionVQModel.decode does not accept mini_batch_size;
+        # implement manual chunking for large images.
+        if height > 256 and len(indices) > 2:
+            mini_bs = 2
+            chunks = [self.vqvae.decode(indices[i:i + mini_bs])
+                      for i in range(0, len(indices), mini_bs)]
+            images = torch.cat(chunks, dim=0).float()
+        else:
+            images = self.vqvae.decode(indices).float()
+        images = images.clamp(-1, 1)
+        images = (images + 1) / 2
+        images = (images * 255).permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
+        return images
+
+
+class NemotronLabsDiffusionImageForMaskedDiffusion(MinistralDiffEncoderModel):
+    config_class = NemotronLabsDiffusionImageConfig
+    supports_gradient_checkpointing = True
+    base_model_prefix = ""
+
+    def __init__(self, config: NemotronLabsDiffusionImageConfig, **kwargs):
+        del kwargs
+        config.d_model = config.hidden_size
+        config.include_bias = config.mlp_bias
+        if not hasattr(config, "d_model_gen") or config.d_model_gen < 0:
+            config.d_model_gen = config.d_model
+        if not hasattr(config, "mlp_hidden_size_gen") or config.mlp_hidden_size_gen < 0:
+            config.mlp_hidden_size_gen = config.intermediate_size
+        if not hasattr(config, "downsample"):
+            config.downsample = False
+        super().__init__(config)
+        self.encoder = NemotronLabsDiffusionImageModel(self.config)
+        self.post_init()
+
+    @property
+    def model(self):
+        return self.encoder
+
+    def get_model(self):
+        return self.encoder
+
+    @torch.no_grad()
+    def generate(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        images: Optional[torch.Tensor] = None,
+        image_sizes: Optional[torch.Tensor] = None,
+        modalities: Optional[List[str]] = None,
+        return_nfe: bool = False,
+        **kwargs,
+    ) -> Union[GenerateOutput, torch.LongTensor]:
+        del image_sizes, modalities
+        if images is not None:
+            raise NotImplementedError("This public release only supports text-to-image generation without multimodal image inputs.")
+        if "inputs_embeds" in kwargs:
+            raise NotImplementedError("inputs_embeds is not supported")
+        if self.config.dlm_paradigm == "bidirectional":
+            kwargs.setdefault("causal_context", False)
+        inputs_embeds = self.get_model().embed_tokens(inputs)
+        output, nfe = MinistralDiffEncoderModel.generate_diffusion(
+            self,
+            prompt_ids=None,
+            prompt_embeds=inputs_embeds,
+            **kwargs,
+        )
+        if return_nfe:
+            return output, nfe
+        return output
+
+    def encode_image_gen(self, images, enc=False):
+        return self.encoder.encode_image_gen(images, enc=enc)
+
+    def decode_image_gen(self, images_to_decode, height, width):
+        return self.encoder.decode_image_gen(images_to_decode, height, width)
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
+        return super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+    @torch.no_grad()
+    def text_to_image(
+        self,
+        prompt: str,
+        tokenizer,
+        sample_policy: str = 'multinomial',
+        confidence_policy: str = 'mmada',
+        guidance_scale: float = 5.0,
+        n_steps: int = 20,
+        batch_size: int = 1,
+        image_resolution: int = 512,
+        n_tokens: int = 1024,
+        shift: int = 3,
+        alg_temp: float = 1.0,
+        min_temperature: float = 0.01,
+        dynamic_temperature: bool = False,
+        micro_cond: str = 'ORIGINAL WIDTH : 1024; ORIGINAL HEIGHT : 1024; TOP : 0; LEFT : 0; SCORE : 6.5',
+        temperature: float = 1.0,
+        schedule_temp: str = 'linear',
+        shift_alg=None,
+        top_p=None,
+        top_k=None,
+        unmask_order=None,
+        cfg_interval=(0, 1),
+        order_cutoff: float = 100,
+        template: str = 'Generate an image with the caption:\n <prompt>',
+        use_cache=None,
+        cache_prompt=None,
+        causal_context: bool = True,
+        is_legacy: bool = False,
+        edit_threshold: float = -1,
+        disable_tqdm: bool = False,
+        return_intermediate_steps: bool = False,
+        **kwargs,
+    ):
+        """Generate an image from a text prompt using masked diffusion."""
+        if shift_alg is None:
+            shift_alg = shift
+
+        NC = _NC
+        device = self.get_model().device
+
+        reserve_token  = NC.reserve_id_token
+        reserve_id     = NC.reserve_id
+        img_mask_id    = 131073   # Emu3 VQ mask token
+        txt_mask_id    = NC.mask_id
+        eot_id         = NC.eos_id
+        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
+        if self.config.dlm_paradigm == 'bidirectional':
+            causal_context = False
+            cache_prompt = False
+            use_cache = False
+
+        if is_legacy:
+            img_begin = img_end = ''
+
+        model_module = self.module if hasattr(self, "module") else self
+        for layer in model_module.encoder.layers:
+            layer.self_attn.mode = 'bidirectional'
+        for layer in model_module.encoder.layers:
+            if hasattr(layer.self_attn, 'diffusion_lm'):
+                layer.self_attn.diffusion_lm = True
+
+        gen_shape_map = {1024: (64, 64), 512: (32, 32), 256: (16, 16)}
+        gen_shape = gen_shape_map[image_resolution]
+        n_tokens_txt = 1024 if image_resolution == 1024 else n_tokens
+
+        prompt_full = f"{prompt} {micro_cond}"
+        question = template.replace('<prompt>', prompt_full)
+
+        conv = _MinistralConv()
+        conv.append_message(conv.roles[0], question)
+        conv.append_message(conv.roles[1],
+                            f"Sure {img_begin}{reserve_token * n_tokens_txt}{img_end}")
+        prompt_question = conv.get_prompt()
+        print(prompt_question.replace(reserve_token, '*'))
+
+        input_ids = _tokenizer_image_token(
+            prompt_question, tokenizer, return_tensors="pt"
+        ).unsqueeze(0).to(device)
+
+        is_gen = input_ids == reserve_id
+        is_gen_enc = input_ids == NC.reserve_id_enc
+        is_eot = torch.where(input_ids == eot_id)[1]
+        assert len(is_eot) == 3, f"Expected 3 EOT tokens, got {len(is_eot)}"
+        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
+
+        # Standard text embedding (no gen tokens yet)
+        inputs_embeds = self.get_model().embed_tokens(raw_input_ids)
+
+        inputs_embeds_uncond = inputs_embeds.clone()
+        noise_embed = self.get_model().embed_tokens(
+            torch.tensor([txt_mask_id], device=device)
+        )
+        inputs_embeds_uncond[is_prompt] = noise_embed
+
+        xt = torch.full((batch_size, n_tokens), img_mask_id,
+                        dtype=torch.long, device=device)
+
+        mask_idx = xt == img_mask_id
+        num_transfer_tokens = _get_num_transfer_tokens(
+            mask_idx, n_steps, schedule='shift', shift=shift
+        )
+        print(num_transfer_tokens)
+
+        sch_t = np.linspace(0, 1, n_steps)
+        if schedule_temp == 'linear':
+            sch_temperatures = (1.0 - sch_t) * (1.0 - min_temperature) + min_temperature
+        elif schedule_temp == 'cosine2':
+            sch_temperatures = _cosine_schedule_2(1.0 - sch_t) * (1.0 - min_temperature) + min_temperature
+        elif schedule_temp == 'shift':
+            sch_temperatures = _logit_normal_schedule(shift_alg, 1.0 - sch_t) * (1.0 - min_temperature) + min_temperature
+        elif schedule_temp == 'exp':
+            sch_temperatures = _exp_schedule(1.0 - sch_t) * (1.0 - min_temperature) + min_temperature
+        else:
+            raise NotImplementedError(f"Unknown schedule_temp: {schedule_temp}")
+        sch_temperatures = torch.tensor(sch_temperatures, device=device, dtype=torch.float32)
+
+        cfg_start = int(cfg_interval[0] * n_steps)
+        cfg_end   = int(cfg_interval[1] * n_steps)
+
+        if confidence_policy == 'stratified' and unmask_order is None:
+            _dim = int(math.sqrt(n_tokens))
+            unmask_order = _stratified_random(n=_dim, seed=42, shuffle_blocks=True)
+
+        total_edited = 0
+        intermediate_x0s = []
+        temp_idx = 0
+        past_key_values = None
+        cache_len = 0
+
+        for decode_step_idx, num_transfer in tqdm(
+            enumerate(num_transfer_tokens[0]),
+            total=num_transfer_tokens.shape[1],
+            disable=disable_tqdm,
+        ):
+            local_temp = sch_temperatures[temp_idx]
+            temp_idx += 1
+            if temp_idx / n_steps > order_cutoff:
+                confidence_policy = 'mmada'
+
+            mask_idx = xt == img_mask_id
+            n_mask = mask_idx.sum()
+            timesteps = (n_mask / mask_idx.numel()).view(1)
+
+            do_cfg = guidance_scale > 0 and cfg_start <= temp_idx <= cfg_end
+            if do_cfg:
+                input_embeddings_input = torch.cat([inputs_embeds_uncond, inputs_embeds]).clone()
+                xt_input = torch.cat([xt, xt])
+                new_token_mask = is_gen.repeat(2, 1)
+                is_gen_enc_mask = is_gen_enc.repeat(2, 1)
+                is_gen_enc_mask[0, :] = False
+                timesteps_in = timesteps.repeat(2)
+            else:
+                input_embeddings_input = inputs_embeds.clone()
+                new_token_mask = is_gen
+                xt_input = xt
+                is_gen_enc_mask = is_gen_enc
+                timesteps_in = timesteps
+
+            all_input_embeddings, new_token_mask = _t2i_wte(
+                self.get_model(), None, gen_shape=gen_shape,
+                x_gen=xt_input,
+                inputs_embeds_curr=input_embeddings_input,
+                new_token_mask=new_token_mask,
+            )
+
+            if use_cache and cache_prompt:
+                if decode_step_idx == 0:
+                    if causal_context:
+                        for layer in model_module.encoder.layers:
+                            if hasattr(layer.self_attn, 'diffusion_lm'):
+                                layer.self_attn.diffusion_lm = False
+                    output = self.get_model()(
+                        None,
+                        input_embeddings=all_input_embeddings[:, :prompt_cutoff],
+                        modality_indices=new_token_mask[:, :prompt_cutoff],
+                        output_hidden_states=True,
+                        past_key_values=None,
+                        is_training=False,
+                        use_cache=True,
+                        overwrite_attn_impl='flash_attn',
+                    )
+                    past_key_values = output.past_key_values
+                    cache_len = past_key_values.get_seq_length()
+                    if causal_context:
+                        for layer in model_module.encoder.layers:
+                            if hasattr(layer.self_attn, 'diffusion_lm'):
+                                layer.self_attn.diffusion_lm = True
+            else:
+                past_key_values = None
+                cache_len = 0
+
+            logits = _t2i_get_logits(
+                self.get_model(),
+                all_input_embeddings[:, cache_len:],
+                new_token_mask[:, cache_len:],
+                past_key_values=past_key_values,
+                gen_shape=gen_shape,
+                input_modality_indices=new_token_mask[:, cache_len:],
+                timesteps=timesteps_in,
+            )
+
+            if do_cfg:
+                new_token_mask, _ = new_token_mask.chunk(2)
+                logits_un, logits = logits.chunk(2)
+                logits_is_ninf = logits == -np.inf
+                logits = (1.0 + guidance_scale) * logits - guidance_scale * logits_un
+                logits[logits_is_ninf] = -np.inf
+
+            if top_p is not None or top_k is not None:
+                _b, _l, _v = logits.shape
+                logits_flat = logits.view(_b * _l, _v)
+                if top_k and top_k > 0:
+                    topk = min(top_k, logits_flat.size(-1))
+                    idx_rm = logits_flat < torch.topk(logits_flat, topk)[0][..., -1, None]
+                    logits_flat[idx_rm] = -np.inf
+                if top_p and top_p < 1.0:
+                    sl, si = torch.sort(logits_flat, descending=True)
+                    cp = torch.cumsum(F.softmax(sl, dim=-1), dim=-1)
+                    si_rm = cp > top_p
+                    si_rm[..., 1:] = si_rm[..., :-1].clone()
+                    si_rm[..., 0] = 0
+                    logits_flat[si_rm.scatter(1, si, si_rm)] = -np.inf
+                logits = logits_flat.view(_b, _l, _v)
+
+            probs = logits.softmax(dim=-1)
+            if sample_policy == 'multinomial':
+                x0 = dists.Categorical(logits=logits / temperature).sample()
+                x0_p = torch.gather(probs, -1, x0.long()[..., None]).squeeze(-1)
+            elif sample_policy == 'argmax':
+                x0 = logits.argmax(-1)
+                x0_p = torch.gather(probs, -1, x0.long()[..., None]).squeeze(-1)
+            else:
+                raise NotImplementedError(f"Unknown sample_policy: {sample_policy}")
+
+            if edit_threshold <= 0:
+                x0 = torch.where(mask_idx, x0, xt)
+
+            if confidence_policy == 'mask_git':
+                _alg_t = alg_temp * local_temp if dynamic_temperature else alg_temp
+                confidence = torch.where(mask_idx, x0_p / _alg_t, torch.tensor(-np.inf, device=device))
+                confidence = torch.softmax(confidence, dim=-1)
+                select_index = torch.multinomial(confidence, num_samples=num_transfer)
+            elif confidence_policy == 'mmada':
+                _alg_t = alg_temp * local_temp if dynamic_temperature else alg_temp
+                confidence = torch.log(x0_p.clamp(1e-20)) + _alg_t * _gumbel_noise(x0_p)
+                confidence = torch.where(mask_idx, confidence, torch.tensor(-np.inf, device=device))
+                _, select_index = torch.topk(confidence[0], k=num_transfer)
+            elif confidence_policy == 'stratified':
+                assert unmask_order is not None
+                start = n_tokens - n_mask
+                select_index = torch.tensor(
+                    unmask_order[start: start + num_transfer],
+                    device=x0.device, dtype=torch.long,
+                )
+            else:
+                raise NotImplementedError(f"Unknown confidence_policy: {confidence_policy}")
+
+            transfer_index = torch.zeros_like(x0, dtype=torch.bool)
+            transfer_index[0, select_index] = True
+            xt[transfer_index] = x0[transfer_index]
+
+            xt_is_mask = xt == img_mask_id
+            if edit_threshold > 0:
+                editable = (~xt_is_mask) & (~transfer_index)
+                hi_conf = torch.where(editable, x0_p, torch.tensor(-torch.inf, device=device)) > edit_threshold
+                changed = (x0 != xt) & hi_conf
+                if changed.sum() > 0:
+                    xt[changed] = x0[changed]
+                    total_edited += changed.sum().item()
+
+            if return_intermediate_steps:
+                x0_inter = xt.clone()
+                x0_inter[xt_is_mask] = x0[xt_is_mask]
+                intermediate_x0s.append(x0_inter.cpu())
+
+        xt = x0.clone()
+        xt[xt == img_mask_id] = x0[xt == img_mask_id]
+        x0_img = xt
+        print(f"Total edited tokens: {total_edited}")
+
+        if return_intermediate_steps:
+            images_npy = self.decode_image_gen(
+                torch.cat(intermediate_x0s).to(x0_img.device),
+                image_resolution, image_resolution,
+            )
+            return [Image.fromarray(x) for x in images_npy]
+        return Image.fromarray(
+            self.decode_image_gen(x0_img, image_resolution, image_resolution)[0]
+        )
+
+
+AutoConfig.register("nemotron_labs_diffusion_image", NemotronLabsDiffusionImageConfig)
+AutoModel.register(NemotronLabsDiffusionImageConfig, NemotronLabsDiffusionImageForMaskedDiffusion)
+AutoModelForCausalLM.register(NemotronLabsDiffusionImageConfig, NemotronLabsDiffusionImageForMaskedDiffusion)

nemotron_diffusion_image_utils.py

@@ -0,0 +1,16 @@
+import torch
+
+
+def pad_along_last_dim(tensor, size):
+    pad_size = size - tensor.shape[-1]
+    if pad_size <= 0:
+        return tensor
+    padding = torch.zeros(*tensor.shape[:-1], pad_size, dtype=tensor.dtype, device=tensor.device)
+    return torch.cat([tensor, padding], dim=-1)
+
+
+def maybe_truncate_last_dim(tensor, size):
+    if size >= tensor.shape[-1]:
+        return tensor
+    return tensor[..., :size]
+    return tensor[..., :size]

tokenizer.json

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