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

added_tokens.json

@@ -0,0 +1,28 @@
+{
+  "</think>": 151668,
+  "</tool_call>": 151658,
+  "</tool_response>": 151666,
+  "<think>": 151667,
+  "<tool_call>": 151657,
+  "<tool_response>": 151665,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

chat_template.jinja

@@ -0,0 +1,89 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

config.json

@@ -0,0 +1,281 @@
+{
+  "architectures": [
+    "Qwen3ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "dtype": "bfloat16",
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 12288,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 36,
+  "model_type": "qwen3",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 8,
+  "quantization_config": {
+    "config_groups": {
+      "group_0": {
+        "format": "nvfp4-pack-quantized",
+        "input_activations": {
+          "actorder": null,
+          "block_structure": null,
+          "dynamic": "local",
+          "group_size": 16,
+          "num_bits": 4,
+          "observer": "minmax",
+          "observer_kwargs": {},
+          "strategy": "tensor_group",
+          "symmetric": true,
+          "type": "float"
+        },
+        "output_activations": null,
+        "targets": [
+          "Linear"
+        ],
+        "weights": {
+          "actorder": null,
+          "block_structure": null,
+          "dynamic": false,
+          "group_size": 16,
+          "num_bits": 4,
+          "observer": "minmax",
+          "observer_kwargs": {},
+          "strategy": "tensor_group",
+          "symmetric": true,
+          "type": "float"
+        }
+      }
+    },
+    "format": "nvfp4-pack-quantized",
+    "global_compression_ratio": null,
+    "ignore": [
+      "model.layers.0.self_attn.q_proj",
+      "model.layers.0.self_attn.k_proj",
+      "model.layers.0.self_attn.v_proj",
+      "model.layers.0.self_attn.o_proj",
+      "model.layers.0.mlp.gate_proj",
+      "model.layers.0.mlp.up_proj",
+      "model.layers.0.mlp.down_proj",
+      "model.layers.1.self_attn.q_proj",
+      "model.layers.1.self_attn.k_proj",
+      "model.layers.1.self_attn.v_proj",
+      "model.layers.1.self_attn.o_proj",
+      "model.layers.2.self_attn.q_proj",
+      "model.layers.2.self_attn.k_proj",
+      "model.layers.2.self_attn.v_proj",
+      "model.layers.2.self_attn.o_proj",
+      "model.layers.3.self_attn.q_proj",
+      "model.layers.3.self_attn.k_proj",
+      "model.layers.3.self_attn.v_proj",
+      "model.layers.3.self_attn.o_proj",
+      "model.layers.4.self_attn.q_proj",
+      "model.layers.4.self_attn.k_proj",
+      "model.layers.4.self_attn.v_proj",
+      "model.layers.4.self_attn.o_proj",
+      "model.layers.5.self_attn.q_proj",
+      "model.layers.5.self_attn.k_proj",
+      "model.layers.5.self_attn.v_proj",
+      "model.layers.5.self_attn.o_proj",
+      "model.layers.6.self_attn.q_proj",
+      "model.layers.6.self_attn.k_proj",
+      "model.layers.6.self_attn.v_proj",
+      "model.layers.6.self_attn.o_proj",
+      "model.layers.7.self_attn.q_proj",
+      "model.layers.7.self_attn.k_proj",
+      "model.layers.7.self_attn.v_proj",
+      "model.layers.7.self_attn.o_proj",
+      "model.layers.8.self_attn.q_proj",
+      "model.layers.8.self_attn.k_proj",
+      "model.layers.8.self_attn.v_proj",
+      "model.layers.8.self_attn.o_proj",
+      "model.layers.8.mlp.gate_proj",
+      "model.layers.8.mlp.up_proj",
+      "model.layers.8.mlp.down_proj",
+      "model.layers.9.self_attn.q_proj",
+      "model.layers.9.self_attn.k_proj",
+      "model.layers.9.self_attn.v_proj",
+      "model.layers.9.self_attn.o_proj",
+      "model.layers.9.mlp.gate_proj",
+      "model.layers.9.mlp.up_proj",
+      "model.layers.9.mlp.down_proj",
+      "model.layers.10.self_attn.q_proj",
+      "model.layers.10.self_attn.k_proj",
+      "model.layers.10.self_attn.v_proj",
+      "model.layers.10.self_attn.o_proj",
+      "model.layers.11.self_attn.q_proj",
+      "model.layers.11.self_attn.k_proj",
+      "model.layers.11.self_attn.v_proj",
+      "model.layers.11.self_attn.o_proj",
+      "model.layers.12.self_attn.q_proj",
+      "model.layers.12.self_attn.k_proj",
+      "model.layers.12.self_attn.v_proj",
+      "model.layers.12.self_attn.o_proj",
+      "model.layers.13.self_attn.q_proj",
+      "model.layers.13.self_attn.k_proj",
+      "model.layers.13.self_attn.v_proj",
+      "model.layers.13.self_attn.o_proj",
+      "model.layers.14.self_attn.q_proj",
+      "model.layers.14.self_attn.k_proj",
+      "model.layers.14.self_attn.v_proj",
+      "model.layers.14.self_attn.o_proj",
+      "model.layers.15.self_attn.q_proj",
+      "model.layers.15.self_attn.k_proj",
+      "model.layers.15.self_attn.v_proj",
+      "model.layers.15.self_attn.o_proj",
+      "model.layers.16.self_attn.q_proj",
+      "model.layers.16.self_attn.k_proj",
+      "model.layers.16.self_attn.v_proj",
+      "model.layers.16.self_attn.o_proj",
+      "model.layers.17.self_attn.q_proj",
+      "model.layers.17.self_attn.k_proj",
+      "model.layers.17.self_attn.v_proj",
+      "model.layers.17.self_attn.o_proj",
+      "model.layers.17.mlp.gate_proj",
+      "model.layers.17.mlp.up_proj",
+      "model.layers.17.mlp.down_proj",
+      "model.layers.18.self_attn.q_proj",
+      "model.layers.18.self_attn.k_proj",
+      "model.layers.18.self_attn.v_proj",
+      "model.layers.18.self_attn.o_proj",
+      "model.layers.18.mlp.gate_proj",
+      "model.layers.18.mlp.up_proj",
+      "model.layers.18.mlp.down_proj",
+      "model.layers.19.self_attn.q_proj",
+      "model.layers.19.self_attn.k_proj",
+      "model.layers.19.self_attn.v_proj",
+      "model.layers.19.self_attn.o_proj",
+      "model.layers.20.self_attn.q_proj",
+      "model.layers.20.self_attn.k_proj",
+      "model.layers.20.self_attn.v_proj",
+      "model.layers.20.self_attn.o_proj",
+      "model.layers.21.self_attn.q_proj",
+      "model.layers.21.self_attn.k_proj",
+      "model.layers.21.self_attn.v_proj",
+      "model.layers.21.self_attn.o_proj",
+      "model.layers.22.self_attn.q_proj",
+      "model.layers.22.self_attn.k_proj",
+      "model.layers.22.self_attn.v_proj",
+      "model.layers.22.self_attn.o_proj",
+      "model.layers.23.self_attn.q_proj",
+      "model.layers.23.self_attn.k_proj",
+      "model.layers.23.self_attn.v_proj",
+      "model.layers.23.self_attn.o_proj",
+      "model.layers.24.self_attn.q_proj",
+      "model.layers.24.self_attn.k_proj",
+      "model.layers.24.self_attn.v_proj",
+      "model.layers.24.self_attn.o_proj",
+      "model.layers.25.self_attn.q_proj",
+      "model.layers.25.self_attn.k_proj",
+      "model.layers.25.self_attn.v_proj",
+      "model.layers.25.self_attn.o_proj",
+      "model.layers.26.self_attn.q_proj",
+      "model.layers.26.self_attn.k_proj",
+      "model.layers.26.self_attn.v_proj",
+      "model.layers.26.self_attn.o_proj",
+      "model.layers.26.mlp.gate_proj",
+      "model.layers.26.mlp.up_proj",
+      "model.layers.26.mlp.down_proj",
+      "model.layers.27.self_attn.q_proj",
+      "model.layers.27.self_attn.k_proj",
+      "model.layers.27.self_attn.v_proj",
+      "model.layers.27.self_attn.o_proj",
+      "model.layers.27.mlp.gate_proj",
+      "model.layers.27.mlp.up_proj",
+      "model.layers.27.mlp.down_proj",
+      "model.layers.28.self_attn.q_proj",
+      "model.layers.28.self_attn.k_proj",
+      "model.layers.28.self_attn.v_proj",
+      "model.layers.28.self_attn.o_proj",
+      "model.layers.29.self_attn.q_proj",
+      "model.layers.29.self_attn.k_proj",
+      "model.layers.29.self_attn.v_proj",
+      "model.layers.29.self_attn.o_proj",
+      "model.layers.30.self_attn.q_proj",
+      "model.layers.30.self_attn.k_proj",
+      "model.layers.30.self_attn.v_proj",
+      "model.layers.30.self_attn.o_proj",
+      "model.layers.31.self_attn.q_proj",
+      "model.layers.31.self_attn.k_proj",
+      "model.layers.31.self_attn.v_proj",
+      "model.layers.31.self_attn.o_proj",
+      "model.layers.32.self_attn.q_proj",
+      "model.layers.32.self_attn.k_proj",
+      "model.layers.32.self_attn.v_proj",
+      "model.layers.32.self_attn.o_proj",
+      "model.layers.33.self_attn.q_proj",
+      "model.layers.33.self_attn.k_proj",
+      "model.layers.33.self_attn.v_proj",
+      "model.layers.33.self_attn.o_proj",
+      "model.layers.34.self_attn.q_proj",
+      "model.layers.34.self_attn.k_proj",
+      "model.layers.34.self_attn.v_proj",
+      "model.layers.34.self_attn.o_proj",
+      "model.layers.35.self_attn.q_proj",
+      "model.layers.35.self_attn.k_proj",
+      "model.layers.35.self_attn.v_proj",
+      "model.layers.35.self_attn.o_proj",
+      "lm_head",
+      "re:.*self_attn.*",
+      "re:.*layers\\.(0|8|9|17|18|26|27)\\.mlp.*"
+    ],
+    "kv_cache_scheme": null,
+    "quant_method": "compressed-tensors",
+    "quantization_status": "compressed",
+    "sparsity_config": {},
+    "transform_config": {},
+    "version": "0.12.2"
+  },
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.1",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

extras/Flux2KleinBackend.py

@@ -0,0 +1,408 @@
+"""
+Flux2KleinBackend.py
+FLUX.2 Klein 9B KV Nunchaku + NVFP4 Qwen3 Text Encoder Backend.
+
+Optimized for 16GB GPUs (Blackwell sm_120 / sm_110):
+- Distilled 4-step inference with KV-cache for multi-reference editing.
+- Sequential GPU stage execution:
+    Stage 1: Qwen3 NVFP4 Text Encoder on GPU -> offload to CPU.
+    Stage 2: AutoencoderKLFlux2 VAE on GPU (if image conditioning) -> offload to CPU.
+    Stage 3: NunchakuFlux2Transformer2DModel (FP4/INT4 SVDQuant) on GPU -> offload to CPU.
+    Stage 4: AutoencoderKLFlux2 VAE on GPU (decode) -> offload to CPU.
+- All models reside on CPU when idle, leaving the GPU with 0% VRAM allocation
+  for other concurrent workloads.
+"""
+
+from __future__ import annotations
+
+import gc
+import os
+import time
+from typing import Any, List, Optional, Union
+
+import numpy as np
+import PIL.Image
+import torch
+from diffusers import AutoencoderKLFlux2, FlowMatchEulerDiscreteScheduler
+from diffusers.pipelines.flux2.pipeline_flux2_klein_kv import (
+    Flux2KleinKVPipeline,
+    compute_empirical_mu,
+    retrieve_timesteps,
+)
+from diffusers.pipelines.flux2.pipeline_output import Flux2PipelineOutput
+from nunchaku import NunchakuFlux2Transformer2DModel
+from nunchaku.utils import get_precision
+from transformers import AutoTokenizer
+
+from NVFP4TextEncoder import load_nvfp4_text_encoder
+
+
+def _flush():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+class SequentialFlux2KleinPipeline:
+    """Wrapper that runs Flux2KleinKVPipeline stages sequentially on CUDA,
+    offloading each component to CPU immediately after its stage finishes.
+    """
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKLFlux2,
+        text_encoder: torch.nn.Module,
+        tokenizer: AutoTokenizer,
+        transformer: NunchakuFlux2Transformer2DModel,
+        device: str | torch.device = "cuda",
+    ):
+        self.scheduler = scheduler
+        self.vae = vae
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.transformer = transformer
+        self.device = torch.device(device)
+
+        # Standard diffusers pipeline attributes for compatibility
+        self.inner_pipe = Flux2KleinKVPipeline(
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            transformer=transformer,
+            is_distilled=True,
+        )
+        self.image_processor = self.inner_pipe.image_processor
+        self.vae_scale_factor = self.inner_pipe.vae_scale_factor
+
+        # Expose components dictionary for server sleep / memory introspection
+        self.components = {
+            "text_encoder": self.text_encoder,
+            "transformer": self.transformer,
+            "vae": self.vae,
+            "scheduler": self.scheduler,
+            "tokenizer": self.tokenizer,
+        }
+
+        # Keep everything on CPU at initialization
+        self.to("cpu")
+
+    def to(self, device: str | torch.device):
+        """Move components to specified device (typically 'cpu')."""
+        target = str(device)
+        if target == "cpu":
+            for mod in [self.text_encoder, self.transformer, self.vae]:
+                if isinstance(mod, torch.nn.Module):
+                    try:
+                        mod.to("cpu")
+                    except Exception:
+                        pass
+            _flush()
+        return self
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: str | list[str] = "",
+        image: Union[PIL.Image.Image, List[PIL.Image.Image], None] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 4,
+        guidance_scale: float = 1.0,
+        negative_prompt: Optional[str] = None,
+        num_images_per_prompt: int = 1,
+        generator: Optional[torch.Generator] = None,
+        seed: Optional[int] = None,
+        **kwargs,
+    ) -> Flux2PipelineOutput:
+        device = self.device
+
+        if generator is None:
+            generator = torch.Generator(device=device)
+            if seed is not None:
+                generator.manual_seed(seed)
+
+        t_start = time.time()
+
+        # ------------------------------------------------------------------
+        # STAGE 1: Text Encoding (Qwen3 NVFP4 Text Encoder)
+        # ------------------------------------------------------------------
+        prompt_embeds = None
+        text_ids = None
+        try:
+            print("[Flux2Klein] Stage 1: Moving Text Encoder to GPU...")
+            self.text_encoder.to(device)
+            _flush()
+
+            prompt_embeds, text_ids = self.inner_pipe.encode_prompt(
+                prompt=prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=512,
+            )
+            # Store CPU copies until Stage 3 to minimize VRAM
+            prompt_embeds = prompt_embeds.to("cpu")
+            text_ids = text_ids.to("cpu")
+        finally:
+            self.text_encoder.to("cpu")
+            _flush()
+            print(f"[Flux2Klein] Stage 1 complete. Allocated VRAM: {torch.cuda.memory_allocated() / 1e9:.3f} GB")
+
+        # ------------------------------------------------------------------
+        # STAGE 2: VAE Conditioning (if reference images provided)
+        # ------------------------------------------------------------------
+        condition_images = None
+        if image is not None:
+            if not isinstance(image, list):
+                image = [image]
+            condition_images = []
+            for img in image:
+                self.image_processor.check_image_input(img)
+                img_w, img_h = img.size
+                if img_w * img_h > 1024 * 1024:
+                    img = self.image_processor._resize_to_target_area(img, 1024 * 1024)
+                    img_w, img_h = img.size
+                multiple_of = self.vae_scale_factor * 2
+                img_w = (img_w // multiple_of) * multiple_of
+                img_h = (img_h // multiple_of) * multiple_of
+                img_pre = self.image_processor.preprocess(img, height=img_h, width=img_w, resize_mode="crop")
+                condition_images.append(img_pre)
+                height = height or img_h
+                width = width or img_w
+
+        height = height or 1024
+        width = width or 1024
+
+        image_latents = None
+        image_latent_ids = None
+        if condition_images is not None:
+            try:
+                print(f"[Flux2Klein] Stage 2: Moving VAE to GPU to encode {len(condition_images)} reference image(s)...")
+                self.vae.to(device)
+                _flush()
+
+                image_latents, image_latent_ids = self.inner_pipe.prepare_image_latents(
+                    images=condition_images,
+                    batch_size=num_images_per_prompt,
+                    generator=generator,
+                    device=device,
+                    dtype=self.vae.dtype,
+                )
+                image_latents = image_latents.to("cpu")
+                image_latent_ids = image_latent_ids.to("cpu")
+            finally:
+                self.vae.to("cpu")
+                _flush()
+                print(f"[Flux2Klein] Stage 2 complete. Allocated VRAM: {torch.cuda.memory_allocated() / 1e9:.3f} GB")
+
+        # ------------------------------------------------------------------
+        # STAGE 3: Transformer Denoising Loop (Nunchaku FP4/INT4 KV-Cache)
+        # ------------------------------------------------------------------
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents, latent_ids = self.inner_pipe.prepare_latents(
+            batch_size=num_images_per_prompt,
+            num_latents_channels=num_channels_latents,
+            height=height,
+            width=width,
+            dtype=self.transformer.dtype,
+            device=device,
+            generator=generator,
+        )
+
+        sigmas = np.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps)
+        image_seq_len = latents.shape[1]
+        mu = compute_empirical_mu(image_seq_len=image_seq_len, num_steps=num_inference_steps)
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        self.scheduler.set_begin_index(0)
+
+        kv_cache = None
+        try:
+            print("[Flux2Klein] Stage 3: Moving Nunchaku Transformer to GPU...")
+            self.transformer.to(device)
+            _flush()
+
+            prompt_embeds = prompt_embeds.to(device)
+            text_ids = text_ids.to(device)
+            latents = latents.to(device)
+            latent_ids = latent_ids.to(device)
+            if image_latents is not None:
+                image_latents = image_latents.to(device)
+                image_latent_ids = image_latent_ids.to(device)
+
+            for i, t in enumerate(timesteps):
+                t_step = time.time()
+                timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+                if i == 0 and image_latents is not None:
+                    # Step 0: Extract KV cache for reference tokens
+                    latent_model_input = torch.cat([image_latents, latents], dim=1).to(self.transformer.dtype)
+                    latent_image_ids = torch.cat([image_latent_ids, latent_ids], dim=1)
+                    noise_pred, kv_cache = self.transformer(
+                        hidden_states=latent_model_input,
+                        timestep=timestep / 1000,
+                        guidance=None,
+                        encoder_hidden_states=prompt_embeds,
+                        txt_ids=text_ids,
+                        img_ids=latent_image_ids,
+                        joint_attention_kwargs=None,
+                        return_dict=False,
+                        kv_cache_mode="extract",
+                        num_ref_tokens=image_latents.shape[1],
+                    )
+                elif kv_cache is not None:
+                    # Steps 1+: Reuse cached reference KV
+                    noise_pred = self.transformer(
+                        hidden_states=latents.to(self.transformer.dtype),
+                        timestep=timestep / 1000,
+                        guidance=None,
+                        encoder_hidden_states=prompt_embeds,
+                        txt_ids=text_ids,
+                        img_ids=latent_ids,
+                        joint_attention_kwargs=None,
+                        return_dict=False,
+                        kv_cache=kv_cache,
+                        kv_cache_mode="cached",
+                    )[0]
+                else:
+                    # Standard text-to-image forward
+                    noise_pred = self.transformer(
+                        hidden_states=latents.to(self.transformer.dtype),
+                        timestep=timestep / 1000,
+                        guidance=None,
+                        encoder_hidden_states=prompt_embeds,
+                        txt_ids=text_ids,
+                        img_ids=latent_ids,
+                        joint_attention_kwargs=None,
+                        return_dict=False,
+                    )[0]
+
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                print(
+                    f"[Flux2Klein] Step {i+1}/{len(timesteps)} in {time.time() - t_step:.2f}s | "
+                    f"Peak VRAM: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB"
+                )
+
+            if "noise_pred" in locals():
+                del noise_pred
+        finally:
+            if kv_cache is not None:
+                kv_cache.clear()
+            self.transformer.to("cpu")
+            del prompt_embeds, text_ids, image_latents, image_latent_ids
+            _flush()
+            print(f"[Flux2Klein] Stage 3 complete. Allocated VRAM: {torch.cuda.memory_allocated() / 1e9:.3f} GB")
+
+        # ------------------------------------------------------------------
+        # STAGE 4: VAE Decoding
+        # ------------------------------------------------------------------
+        try:
+            print("[Flux2Klein] Stage 4: Moving VAE to GPU to decode latents...")
+            self.vae.to(device)
+            _flush()
+
+            latents = self.inner_pipe._unpack_latents_with_ids(latents, latent_ids)
+            latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
+            latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(
+                latents.device, latents.dtype
+            )
+            latents = latents * latents_bn_std + latents_bn_mean
+            latents = self.inner_pipe._unpatchify_latents(latents)
+            decoded = self.vae.decode(latents, return_dict=False)[0]
+            out_images = self.image_processor.postprocess(decoded, output_type="pil")
+        finally:
+            self.vae.to("cpu")
+            del latents, latent_ids
+            _flush()
+            print(f"[Flux2Klein] Stage 4 complete. Allocated VRAM: {torch.cuda.memory_allocated() / 1e9:.3f} GB")
+
+        print(f"[Flux2Klein] Total generation finished in {time.time() - t_start:.2f}s!")
+        return Flux2PipelineOutput(images=out_images)
+
+
+class Flux2KleinBackend:
+    """Backend factory for FLUX.2 Klein 9B KV Nunchaku."""
+
+    def __init__(
+        self,
+        model_id: str,
+        optimized_model_path: Optional[str] = None,
+        nvfp4_text_encoder_path: Optional[str] = None,
+    ):
+        self.model_id = model_id
+        self.optimized_model_path = optimized_model_path
+        self.nvfp4_text_encoder_path = nvfp4_text_encoder_path
+        self.pipeline = None
+
+    def load(self):
+        print(f"[Flux2KleinBackend] Initializing backend from base model: {self.model_id}...")
+
+        # 1. Resolve optimized transformer path
+        opt_path = self.optimized_model_path
+        if not opt_path:
+            precision = get_precision()
+            opt_path = os.path.join(
+                self.model_id,
+                f"svdq-{precision}_r32-FLUX.2-klein-9b-kv-Nunchaku.safetensors",
+            )
+            if not os.path.exists(opt_path):
+                # Fallback to fp4 if not found directly
+                opt_path = os.path.join(
+                    self.model_id,
+                    "svdq-fp4_r32-FLUX.2-klein-9b-kv-Nunchaku.safetensors",
+                )
+
+        print(f"[Flux2KleinBackend] Loading Nunchaku transformer from {opt_path}...")
+        transformer = NunchakuFlux2Transformer2DModel.from_pretrained(
+            opt_path,
+            torch_dtype=torch.bfloat16,
+        )
+        transformer.to("cpu")
+
+        # 2. Resolve NVFP4 text encoder path
+        enc_path = self.nvfp4_text_encoder_path
+        if not enc_path:
+            # Check if running directly from within the text encoder repository (extras/)
+            candidate_parent = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+            if os.path.exists(os.path.join(candidate_parent, "config.json")):
+                enc_path = candidate_parent
+            elif os.path.exists(os.path.join(self.model_id, "text_encoder")):
+                enc_path = os.path.join(self.model_id, "text_encoder")
+            else:
+                enc_path = "/path/to/FLUX.2-klein-text-encoder-NVFP4"
+
+        print(f"[Flux2KleinBackend] Loading NVFP4 text encoder from {enc_path}...")
+        text_encoder = load_nvfp4_text_encoder(enc_path, device="cpu", dtype=torch.bfloat16)
+        tokenizer = AutoTokenizer.from_pretrained(enc_path)
+
+        # 3. Load VAE and Scheduler
+        print(f"[Flux2KleinBackend] Loading VAE and Scheduler from {self.model_id}...")
+        vae = AutoencoderKLFlux2.from_pretrained(
+            self.model_id,
+            subfolder="vae",
+            torch_dtype=torch.bfloat16,
+        )
+        vae.to("cpu")
+        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            self.model_id,
+            subfolder="scheduler",
+        )
+
+        # 4. Construct Sequential Pipeline
+        self.pipeline = SequentialFlux2KleinPipeline(
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            transformer=transformer,
+            device="cuda",
+        )
+
+        print("[Flux2KleinBackend] Pipeline successfully loaded and ready on CPU!")
+        return self.pipeline, self.pipeline

extras/ImageEditServer.py

@@ -0,0 +1,505 @@
+import argparse
+import base64
+import io
+import time
+import torch
+import uvicorn
+import gc
+import asyncio
+import traceback
+from typing import List, Optional, Union
+from contextlib import asynccontextmanager
+from fastapi import FastAPI, HTTPException, UploadFile, File, Form
+from pydantic import BaseModel
+from PIL import Image, ImageOps
+
+# Argument parsing
+parser = argparse.ArgumentParser(description="Flux Image Edit Server with Nunchaku")
+parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to bind to")
+parser.add_argument("--port", type=int, default=8000, help="Port to bind to")
+parser.add_argument("--model", type=str, default="black-forest-labs/FLUX.1-Kontext-dev", help="Path or Repo ID of the base model")
+parser.add_argument("--optimized-model", type=str, default=None, help="Path to the optimized Nunchaku model safetensors file")
+parser.add_argument("--optimized-edit-model", type=str, default=None, help="Path to the optimized Nunchaku model safetensors file for editing (optional)")
+parser.add_argument("--backend", type=str, default="kontext", choices=["kontext", "flux2", "flux2-klein", "flux2_klein", "qwen", "glm", "zimage"], help="Backend to use: 'kontext', 'flux2', 'flux2-klein', 'qwen', 'glm', or 'zimage'")
+parser.add_argument("--steps", type=int, default=28, help="Default number of inference steps")
+parser.add_argument("--guidance-scale", type=float, default=3.5, help="Default guidance scale")
+parser.add_argument("--qwenimage", action="store_true", help="Use QwenImageBackend (T2I only) instead of full Qwen edit backend")
+parser.add_argument("--uma", action="store_true", help="Enable Unified Memory Architecture mode (load all to GPU, disable offload)")
+parser.add_argument(
+    "--nvfp4-text-encoder",
+    type=str,
+    default=None,
+    help=(
+        "Path to an NVFP4-pack-quantized HuggingFace text encoder "
+        "(compressed-tensors format). Honoured by the zimage and flux2-klein backends; "
+        "swaps in vLLM's W4A4 NVFP4 CUTLASS GEMM for ~4x text-encoder VRAM savings."
+    ),
+)
+args = parser.parse_args()
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup logic
+    load_model()
+    yield
+    # Shutdown logic (if any) could go here
+
+app = FastAPI(lifespan=lifespan)
+
+# Global components
+IMAGE_DIMENSION_ALIGNMENT = 32
+pipeline = None
+edit_pipeline = None
+request_lock = asyncio.Lock()
+is_sleeping_flag = False
+sleep_requested = False
+
+def load_model():
+    global pipeline, edit_pipeline
+    
+    try:
+        if args.backend == "kontext":
+            import KontextBackend
+            print(f"Initializing KontextBackend...")
+            backend = KontextBackend.KontextBackend(args.model, args.optimized_model)
+            pipeline, edit_pipeline = backend.load()
+        elif args.backend == "flux2":
+            import Flux2Backend
+            print(f"Initializing Flux2Backend...")
+            backend = Flux2Backend.Flux2Backend(args.model)
+            pipeline, edit_pipeline = backend.load()
+        elif args.backend == "glm":
+            import GlmBackend
+            print(f"Initializing GlmBackend...")
+            # Use provided model or default to the one in the snippet if args.model is generic
+            # The user might pass the specific GLM model via --model, or we default in GlmBackend.
+            # Let's pass args.model if it's not the default flux one, otherwise let GlmBackend use its default.
+            model_to_use = args.model if args.model != "black-forest-labs/FLUX.1-Kontext-dev" else "Disty0/GLM-Image-SDNQ-4bit-dynamic"
+            backend = GlmBackend.GlmBackend(model_to_use)
+            pipeline, edit_pipeline = backend.load()
+        elif args.backend.startswith("qwen"):
+            if args.qwenimage:
+                import QwenImageBackend
+                print(f"Initializing QwenImageBackend (T2I only)...")
+                backend = QwenImageBackend.QwenImageBackend(args.model, args.optimized_model)
+                pipeline, edit_pipeline = backend.load()
+            else:
+                import QwenBackend
+                print(f"Initializing QwenBackend...")
+                backend = QwenBackend.QwenBackend(args.model, args.optimized_model, optimized_edit_model_path=args.optimized_edit_model, uma=args.uma)
+                pipeline, edit_pipeline = backend.load()
+        elif args.backend in ["flux2-klein", "flux2_klein"]:
+            import Flux2KleinBackend
+            print(f"Initializing Flux2KleinBackend...")
+            backend = Flux2KleinBackend.Flux2KleinBackend(
+                args.model,
+                args.optimized_model,
+                nvfp4_text_encoder_path=args.nvfp4_text_encoder,
+            )
+            pipeline, edit_pipeline = backend.load()
+        elif args.backend == "zimage":
+            import ZImageTurboBackend
+            print(f"Initializing ZImageTurboBackend...")
+            backend = ZImageTurboBackend.ZImageTurboBackend(
+                args.model,
+                args.optimized_model,
+                uma=args.uma,
+                nvfp4_text_encoder_path=args.nvfp4_text_encoder,
+            )
+            pipeline, edit_pipeline = backend.load()
+        else:
+            raise ValueError(f"Unknown backend: {args.backend}")
+            
+    except Exception as e:
+        print(f"Oh no! The model refused to wake up: {e}")
+        raise e
+    
+    # Enable progress bar for diffusers
+    import diffusers.utils.logging
+    diffusers.utils.logging.enable_progress_bar()
+    diffusers.utils.logging.set_verbosity_info()
+    
+    print("Model loaded successfully! Ready for editing quests!")
+
+def flush():
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+class ImageGenerationRequest(BaseModel):
+    prompt: str
+    n: int = 1
+    size: str = "1024x1024"
+    response_format: str = "b64_json"
+    quality: str = "standard"
+    style: str = "vivid"
+    num_inference_steps: Optional[int] = None
+    guidance_scale: Optional[float] = None
+    negative_prompt: Optional[str] = None
+    seed: Optional[int] = None
+
+
+@app.post("/v1/sleep")
+async def sleep_endpoint():
+    global is_sleeping_flag, sleep_requested
+    sleep_requested = True
+    try:
+        async with request_lock:
+            if not is_sleeping_flag and sleep_requested:
+                print("Sleep requested, moving models to CPU...")
+                for p in [pipeline, edit_pipeline]:
+                    if not p: continue
+                    for name, component in p.components.items():
+                        if isinstance(component, torch.nn.Module):
+                            # Special handling for Nunchaku which blocks .to() if offload is True
+                            if hasattr(component, "set_offload") and getattr(component, "offload", False):
+                                component.set_offload(False)
+                                component._nunchaku_was_offloaded = True
+
+                            try:
+                                component.to("cpu")
+                            except Exception as e:
+                                pass
+                flush()
+                is_sleeping_flag = True
+    finally:
+        sleep_requested = False
+    return {"status": "sleep completed", "is_sleeping": is_sleeping_flag}
+
+@app.post("/v1/wake_up")
+async def wake_up_endpoint():
+    global is_sleeping_flag, sleep_requested
+    sleep_requested = False
+    async with request_lock:
+        if is_sleeping_flag:
+            print("Waking up, restoring models to CUDA...")
+            for p in [pipeline, edit_pipeline]:
+                if not p: continue
+                excluded = getattr(p, "_exclude_from_cpu_offload", [])
+                for name, component in p.components.items():
+                    if isinstance(component, torch.nn.Module):
+                        if getattr(component, "_nunchaku_was_offloaded", False):
+                            component.set_offload(True, use_pin_memory=True, num_blocks_on_gpu=8)
+                            for attr in ["img_in", "txt_in", "txt_norm", "time_text_embed", "norm_out", "proj_out"]:
+                                if hasattr(component, attr):
+                                    try:
+                                        getattr(component, attr).to("cuda")
+                                    except Exception:
+                                        pass
+                            component._nunchaku_was_offloaded = False
+                        elif not hasattr(component, "_hf_hook") or name in excluded:
+                            try:
+                                component.to("cuda")
+                            except Exception:
+                                pass
+            is_sleeping_flag = False
+    return {"status": "awoken", "is_sleeping": False}
+
+@app.get("/v1/is_sleeping")
+async def is_sleeping_endpoint():
+    return {"is_sleeping": is_sleeping_flag}
+
+
+@app.get("/v1/memory_stats")
+async def memory_stats_endpoint():
+    """Lightweight introspection endpoint that returns PyTorch's CUDA allocator
+    snapshot. Used to diagnose VRAM/UMA bloat without restarting the server."""
+    stats = {}
+    if torch.cuda.is_available():
+        stats["allocated_gb"] = torch.cuda.memory_allocated() / 1e9
+        stats["reserved_gb"] = torch.cuda.memory_reserved() / 1e9
+        stats["max_allocated_gb"] = torch.cuda.max_memory_allocated() / 1e9
+        stats["max_reserved_gb"] = torch.cuda.max_memory_reserved() / 1e9
+        # Top allocations by size from the allocator snapshot (>=64 MiB)
+        try:
+            snap = torch.cuda.memory_snapshot()
+            blocks = []
+            for seg in snap:
+                for b in seg.get("blocks", []):
+                    if b.get("state") == "active_allocated" and b.get("size", 0) >= 64 * 1024 * 1024:
+                        blocks.append(b["size"])
+            blocks.sort(reverse=True)
+            stats["large_active_blocks_gb"] = [round(s / 1e9, 3) for s in blocks[:20]]
+            stats["large_active_blocks_total_gb"] = round(sum(blocks) / 1e9, 3)
+            stats["large_active_blocks_count"] = len(blocks)
+        except Exception as e:
+            stats["snapshot_error"] = str(e)
+        # Walk Python objects to find big tensors and group them
+        try:
+            import gc as _gc
+            seen = set()
+            big = []
+            for obj in _gc.get_objects():
+                try:
+                    if isinstance(obj, torch.Tensor) and obj.is_cuda:
+                        ptr = obj.data_ptr()
+                        if ptr in seen or ptr == 0:
+                            continue
+                        seen.add(ptr)
+                        sz = obj.element_size() * obj.numel()
+                        if sz >= 16 * 1024 * 1024:
+                            big.append((sz, tuple(obj.shape), str(obj.dtype)))
+                except Exception:
+                    continue
+            big.sort(reverse=True)
+            # Group by (shape, dtype)
+            from collections import Counter
+            grouped = Counter((shape, dtype) for _, shape, dtype in big)
+            stats["big_tensor_groups"] = [
+                {"shape": list(shape), "dtype": dtype, "count": cnt,
+                 "size_gb_each": round(
+                     (1 if shape == () else (lambda l: __import__('functools').reduce(lambda a, b: a*b, l, 1))(shape)) * (
+                         8 if 'int64' in dtype or 'float64' in dtype else
+                         4 if 'int32' in dtype or 'float32' in dtype else
+                         2 if 'bfloat16' in dtype or 'float16' in dtype else 1
+                     ) / 1e9, 4)}
+                for (shape, dtype), cnt in grouped.most_common(30)
+            ]
+            stats["big_tensor_count"] = len(big)
+            stats["big_tensor_total_gb"] = round(sum(s for s, _, _ in big) / 1e9, 3)
+        except Exception as e:
+            stats["walk_error"] = str(e)
+    return stats
+
+@app.post("/v1/images/edits")
+async def edit_image(
+    image: Union[List[UploadFile], UploadFile] = File(...),
+    prompt: str = Form(...),
+    n: int = Form(1),
+    size: str = Form("1024x1024"),
+    response_format: str = Form("b64_json"), # Default to b64_json
+    guidance_scale: Optional[float] = Form(None), 
+    num_inference_steps: Optional[int] = Form(None),
+    negative_prompt: Optional[str] = Form(None),
+    seed: Optional[int] = Form(None)
+):
+    # Use CLI defaults if not provided
+    steps = num_inference_steps if num_inference_steps is not None else args.steps
+    cfg_scale = guidance_scale if guidance_scale is not None else args.guidance_scale
+    neg_prompt = negative_prompt if negative_prompt is not None else "" # Default empty for now, or maybe None?
+    
+    generator = None
+    import random
+    if seed is None:
+        seed = random.randint(0, 2**32 - 1)
+    
+    print(f"Using seed: {seed}")
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+
+    if not edit_pipeline:
+        raise HTTPException(status_code=500, detail="Model not loaded")
+
+    if sleep_requested or is_sleeping_flag:
+        raise HTTPException(status_code=503, detail="Server is sleeping or trying to sleep.")
+
+    async with request_lock:
+        print(f"Received edit request: {prompt}")
+
+        # Processing the input image(s)
+        input_files = image if isinstance(image, list) else [image]
+        init_images = []
+
+        try:
+            for img_file in input_files:
+                await img_file.seek(0)
+                contents = await img_file.read()
+                img = Image.open(io.BytesIO(contents)).convert("RGB")
+                init_images.append(img)
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=f"Invalid image file: {e}")
+        
+        if not init_images:
+             raise HTTPException(status_code=400, detail="No images provided")
+
+        # Parse max target dimensions from requested size
+        try:
+            target_width, target_height = map(int, size.split("x"))
+        except ValueError:
+            target_width, target_height = 1024, 1024
+
+        # Calculate new dimensions preserving aspect ratio based on the first image
+        first_image = init_images[0]
+        orig_width, orig_height = first_image.size
+        scale = min(target_width / orig_width, target_height / orig_height)
+        new_width = int(orig_width * scale)
+        new_height = int(orig_height * scale)
+
+        # Ensure dimensions are aligned to 32 for compatibility (e.g. GLM-Image)
+        width = (new_width // IMAGE_DIMENSION_ALIGNMENT) * IMAGE_DIMENSION_ALIGNMENT
+        height = (new_height // IMAGE_DIMENSION_ALIGNMENT) * IMAGE_DIMENSION_ALIGNMENT
+        
+        # Resize input images to match the calculated target size, padding if necessary
+        resized_images = []
+        for img in init_images:
+            if img.size != (width, height):
+                # Use ImageOps.pad to preserve aspect ratio and center in the target size
+                # This handles cases where subsequent images might have different ARs
+                img = ImageOps.pad(img, (width, height), method=Image.LANCZOS, color=(0, 0, 0))
+            resized_images.append(img)
+            
+        # If single image, pass as item, if multiple, pass as list
+        # GLM pipeline has a bug where it checks len() on the input, so it must be a list
+        if len(resized_images) > 1 or args.backend == "glm":
+            image_input = resized_images
+        else:
+            image_input = resized_images[0]
+
+        response_images = []
+        
+        try:
+            if args.backend.startswith("qwen"):
+                # Qwen specific parameters
+                # guidance_scale maps to true_cfg_scale
+                if args.qwenimage: # QwenImageBackend is T2I only, so it doesn't take an image
+                    generated_images = edit_pipeline(
+                        prompt=prompt,
+                        height=height,
+                        width=width,
+                        num_inference_steps=steps,
+                        true_cfg_scale=cfg_scale,
+                        num_images_per_prompt=n,
+                        generator=generator,
+                    ).images
+                else: # Full Qwen edit backend takes an image (or list of images now)
+                    generated_images = edit_pipeline(
+                        image=image_input,
+                        prompt=prompt,
+                        height=height,
+                        width=width,
+                        negative_prompt=neg_prompt,
+                        num_inference_steps=steps,
+                        true_cfg_scale=cfg_scale,
+                        num_images_per_prompt=n,
+                        generator=generator,
+                    ).images
+            else:
+                # Standard Flux/Kontext or GLM
+                # GLM I2I Fix: Manually move vision encoder to GPU because get_image_features escapes hooks
+                if args.backend == "glm" and hasattr(edit_pipeline, "vision_language_encoder"):
+                     print("Manually moving GLM Vision Encoder to GPU...")
+                     edit_pipeline.vision_language_encoder.to("cuda")
+                
+                try:
+                    generated_images = edit_pipeline(
+                        image=image_input,
+                        prompt=prompt,
+                        height=height,
+                        width=width,
+                        num_inference_steps=steps, 
+                        guidance_scale=cfg_scale,
+                        num_images_per_prompt=n,
+                        generator=generator,
+                    ).images
+                finally:
+                    if args.backend == "glm" and hasattr(edit_pipeline, "vision_language_encoder"):
+                         print("Moving GLM Vision Encoder back to CPU...")
+                         edit_pipeline.vision_language_encoder.to("cpu")
+            
+            for img in generated_images:
+                buffered = io.BytesIO()
+                img.save(buffered, format="PNG")
+                img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+                
+                if response_format == "b64_json":
+                    response_images.append({"b64_json": img_str})
+                else:
+                    # If url is requested we can't really do it without storage, so we fallback or error?
+                    # For now, let's just assume simple b64_json as per request
+                    response_images.append({"b64_json": img_str}) # Fallback
+                
+        except Exception as e:
+            print(f"Error during editing: {e}")
+            print(traceback.format_exc())
+            raise HTTPException(status_code=500, detail=str(e))
+        finally:
+            flush()
+
+        return {
+            "created": int(time.time()),
+            "data": response_images
+        }
+
+
+
+@app.post("/v1/images/generations")
+async def generate_image(request: ImageGenerationRequest):
+    if not pipeline:
+        raise HTTPException(status_code=500, detail="Model not loaded")
+
+    if sleep_requested or is_sleeping_flag:
+        raise HTTPException(status_code=503, detail="Server is sleeping or trying to sleep.")
+
+    async with request_lock:
+        #print(f"Received generation request: {request.prompt}")
+
+        # Parse size
+        try:
+            width, height = map(int, request.size.split("x"))
+        except ValueError:
+            width, height = 1024, 1024
+
+        # Ensure dimensions are aligned to 32
+        width = (width // IMAGE_DIMENSION_ALIGNMENT) * IMAGE_DIMENSION_ALIGNMENT
+        height = (height // IMAGE_DIMENSION_ALIGNMENT) * IMAGE_DIMENSION_ALIGNMENT
+
+        response_images = []
+        
+        try:
+            # Generate images (no image argument for txt2img!)
+            steps = request.num_inference_steps if request.num_inference_steps is not None else args.steps
+            cfg_scale = request.guidance_scale if request.guidance_scale is not None else args.guidance_scale
+            # negative_prompt not in standard request body in original snippet, but we added it to model
+            neg_prompt = request.negative_prompt if request.negative_prompt is not None else ""
+            
+            generator = None
+            import random
+            seed = request.seed
+            if seed is None:
+                seed = random.randint(0, 2**32 - 1)
+            
+            print(f"Using seed: {seed}")
+            generator = torch.Generator(device="cuda").manual_seed(seed)
+
+            if args.backend.startswith("qwen"):
+                 generated_images = pipeline(
+                    prompt=request.prompt,
+                    height=height,
+                    width=width,
+                    num_inference_steps=steps, 
+                    true_cfg_scale=cfg_scale,     
+                    num_images_per_prompt=request.n,
+                    negative_prompt=neg_prompt,
+                    generator=generator,
+                ).images
+            else:
+                generated_images = pipeline(
+                    prompt=request.prompt,
+                    height=height,
+                    width=width,
+                    num_inference_steps=steps, 
+                    guidance_scale=cfg_scale,     
+                    num_images_per_prompt=request.n,
+                    generator=generator,
+                    # Not passing negative_prompt here for generation unless we confirm support in standard Flux pipeline?
+                ).images
+            
+            for img in generated_images:
+                buffered = io.BytesIO()
+                img.save(buffered, format="PNG")
+                img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+                response_images.append({"b64_json": img_str})
+                
+        except Exception as e:
+            print(f"Error during generation: {e}")
+            print(traceback.format_exc())
+            raise HTTPException(status_code=500, detail=str(e))
+        finally:
+            flush()
+
+        return {
+            "created": int(time.time()),
+            "data": response_images
+        }
+
+if __name__ == "__main__":
+    uvicorn.run(app, host=args.host, port=args.port)

extras/NVFP4TextEncoder.py

@@ -0,0 +1,324 @@
+"""
+NVFP4 text encoder loader for diffusers image pipelines.
+
+Loads a compressed-tensors NVFP4-pack-quantized HuggingFace causal LM and wraps
+it so it can be plugged into ``diffusers.ZImagePipeline`` (or any pipeline
+calling ``self.text_encoder(input_ids, attention_mask, output_hidden_states=True)``).
+
+Strategy:
+- Instantiate the HF model on the ``meta`` device (no real allocation).
+- Walk every ``torch.nn.Linear`` and swap it for vLLM's ``ReplicatedLinear`` with
+  ``CompressedTensorsConfig`` derived from the checkpoint's
+  ``quantization_config``. This registers ``weight_packed`` / ``weight_scale`` /
+  ``*_global_scale`` parameters in the exact layout vLLM's
+  ``CompressedTensorsW4A4Fp4`` scheme expects.
+- Materialise remaining (non-Linear) parameters (embeddings, RMSNorm, k/q norms)
+  on the target device & dtype.
+- Stream the safetensors file and dispatch each tensor through the registered
+  vLLM ``weight_loader`` (which handles layout swizzling on
+  ``process_weights_after_loading``).
+- Tie the LM head to the input embedding when ``config.tie_word_embeddings``.
+
+The result is a regular ``nn.Module`` matching the HF model's call signature
+(``forward(input_ids, attention_mask, output_hidden_states)``) -- usable directly
+as ``ZImagePipeline.text_encoder``.
+
+vLLM requires a minimal global context (distributed process group + model
+parallel state + active VllmConfig) even at TP=1 because ``ReplicatedLinear``
+queries the TP world size at construction. We bootstrap that lazily once.
+
+Forced kernel: we set ``VLLM_NVFP4_GEMM_BACKEND=cutlass`` to skip
+flashinfer-cutlass JIT (which needs the ``ninja`` binary on PATH). The vLLM
+CUTLASS kernel is built into the wheel.
+"""
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import Iterator
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+
+# ----------------------------------------------------------------------------
+# One-time vLLM bootstrap (TP=1, no engine, just enough context for ReplicatedLinear)
+# ----------------------------------------------------------------------------
+_VLLM_BOOTSTRAPPED = False
+_VLLM_CONFIG_CTX = None  # holds the entered set_current_vllm_config context manager
+
+
+def _bootstrap_vllm_once() -> None:
+    """Initialise the bits of vLLM that ReplicatedLinear needs at TP=1.
+
+    Idempotent. Uses ``gloo`` so it works without NCCL/CUDA-aware MPI and even
+    when CUDA is busy with the diffusion transformer.
+    """
+    global _VLLM_BOOTSTRAPPED, _VLLM_CONFIG_CTX
+    if _VLLM_BOOTSTRAPPED:
+        return
+
+    # Force CUTLASS to avoid flashinfer-cutlass JIT (requires `ninja` on PATH).
+    os.environ.setdefault("VLLM_NVFP4_GEMM_BACKEND", "cutlass")
+
+    from vllm.config import VllmConfig
+    from vllm.config.vllm import set_current_vllm_config
+    from vllm.distributed import (
+        ensure_model_parallel_initialized,
+        init_distributed_environment,
+    )
+
+    # Pick a free port; world_size=1.
+    import socket
+
+    s = socket.socket()
+    s.bind(("127.0.0.1", 0))
+    port = s.getsockname()[1]
+    s.close()
+
+    if not torch.distributed.is_initialized():
+        init_distributed_environment(
+            world_size=1,
+            rank=0,
+            local_rank=0,
+            distributed_init_method=f"tcp://127.0.0.1:{port}",
+            backend="gloo",
+        )
+
+    # Enter a long-lived VllmConfig context. We never exit it -- the encoder
+    # may construct submodules lazily and ReplicatedLinear calls
+    # get_current_vllm_config() at init.
+    vc = VllmConfig()
+    _VLLM_CONFIG_CTX = set_current_vllm_config(vc)
+    _VLLM_CONFIG_CTX.__enter__()
+
+    ensure_model_parallel_initialized(1, 1)
+    _VLLM_BOOTSTRAPPED = True
+
+
+# ----------------------------------------------------------------------------
+# Module: linear replacement
+# ----------------------------------------------------------------------------
+def _replace_linears_with_replicated(
+    model: nn.Module, quant_config
+) -> None:
+    """Recursively swap every ``nn.Linear`` for vLLM ``ReplicatedLinear``.
+
+    Carries the ``prefix`` so quant_config's ``ignore`` patterns (e.g. ``lm_head``)
+    are correctly applied.
+    """
+    from vllm.model_executor.layers.linear import ReplicatedLinear
+
+    def _walk(parent: nn.Module, prefix: str) -> None:
+        for child_name, child in list(parent.named_children()):
+            qname = f"{prefix}.{child_name}" if prefix else child_name
+            if isinstance(child, nn.Linear):
+                new = ReplicatedLinear(
+                    input_size=child.in_features,
+                    output_size=child.out_features,
+                    bias=child.bias is not None,
+                    quant_config=quant_config,
+                    prefix=qname,
+                    return_bias=False,
+                    params_dtype=torch.bfloat16,
+                )
+                setattr(parent, child_name, new)
+            else:
+                _walk(child, qname)
+
+    _walk(model, prefix="")
+
+
+def _materialize_remaining_meta_params(
+    model: nn.Module, dtype: torch.dtype, device: torch.device
+) -> None:
+    """Replace any ``meta`` parameter with empty real storage.
+
+    Only touches parameters NOT already created on a real device by the
+    ReplicatedLinear swap above (i.e. embeddings, layernorms, biases).
+    """
+    for name, param in list(model.named_parameters(recurse=True)):
+        if param.device.type == "meta":
+            real = nn.Parameter(
+                torch.empty(param.shape, dtype=dtype, device=device),
+                requires_grad=False,
+            )
+            # Replace in the parent module
+            parent = model
+            *path, leaf = name.split(".")
+            for p in path:
+                parent = getattr(parent, p)
+            setattr(parent, leaf, real)
+    # Same for buffers (e.g. rotary inv_freq if registered as buffer on meta)
+    for name, buf in list(model.named_buffers(recurse=True)):
+        if buf.device.type == "meta":
+            real = torch.empty(buf.shape, dtype=buf.dtype, device=device)
+            parent = model
+            *path, leaf = name.split(".")
+            for p in path:
+                parent = getattr(parent, p)
+            parent.register_buffer(leaf, real, persistent=False)
+
+
+# ----------------------------------------------------------------------------
+# Weight loading
+# ----------------------------------------------------------------------------
+def _iter_safetensors(model_dir: str) -> Iterator[tuple[str, torch.Tensor]]:
+    """Yield (name, tensor) pairs from all *.safetensors shards in ``model_dir``."""
+    from safetensors import safe_open
+
+    # Single-file checkpoint or sharded? Prefer ``model.safetensors.index.json``.
+    index_path = os.path.join(model_dir, "model.safetensors.index.json")
+    if os.path.exists(index_path):
+        with open(index_path) as f:
+            index = json.load(f)
+        shards = sorted(set(index["weight_map"].values()))
+    else:
+        # Find all *.safetensors files in dir
+        shards = sorted(
+            fn for fn in os.listdir(model_dir) if fn.endswith(".safetensors")
+        )
+    for shard in shards:
+        path = os.path.join(model_dir, shard)
+        with safe_open(path, framework="pt") as f:
+            for key in f.keys():
+                yield key, f.get_tensor(key)
+
+
+def _load_weights_into_model(model: nn.Module, model_dir: str) -> None:
+    """Stream safetensors into the (already-structured) model.
+
+    Uses each ReplicatedLinear's registered ``weight_loader`` for quantised
+    params (which handles tensor-parallel sharding, even though TP=1 here it
+    keeps casts consistent). Other params (embeddings, layernorms, biases) are
+    copied directly.
+    """
+    # Strip vllm-omni-style "text_encoder." prefix if present; not applicable
+    # here since we load the standalone HF Qwen3 checkpoint where keys start
+    # with "model.layers..." / "model.embed_tokens..." / "lm_head...".
+    name_to_param: dict[str, nn.Parameter] = dict(model.named_parameters(recurse=True))
+    name_to_buffer: dict[str, torch.Tensor] = dict(model.named_buffers(recurse=True))
+
+    missing = set(name_to_param.keys())
+    unexpected = []
+
+    for key, tensor in _iter_safetensors(model_dir):
+        # Skip rotary inv_freq etc that aren't params (rare in modern HF saves)
+        if key in name_to_param:
+            param = name_to_param[key]
+            wl = getattr(param, "weight_loader", None)
+            if wl is not None:
+                wl(param, tensor.to(param.device))
+            else:
+                with torch.no_grad():
+                    param.data.copy_(tensor.to(param.device, dtype=param.dtype))
+            missing.discard(key)
+        elif key in name_to_buffer:
+            with torch.no_grad():
+                name_to_buffer[key].copy_(tensor.to(name_to_buffer[key].device))
+        else:
+            unexpected.append(key)
+
+    # Tied embeddings (lm_head.weight not in checkpoint when tie_word_embeddings=True)
+    cfg = getattr(model, "config", None)
+    if cfg is not None and getattr(cfg, "tie_word_embeddings", False):
+        try:
+            inp_emb = model.get_input_embeddings().weight
+            model.lm_head.weight = inp_emb  # share storage
+            missing.discard("lm_head.weight")
+        except Exception:
+            pass
+
+    if missing:
+        # It's OK if missing entries are *purely* lm_head.weight when tied; we
+        # already handled that above. Anything else is fatal-ish.
+        leftover = sorted(missing)
+        if leftover:
+            print(
+                f"[NVFP4TextEncoder] WARN: {len(leftover)} params missing from checkpoint; "
+                f"first 5: {leftover[:5]}"
+            )
+    if unexpected:
+        print(
+            f"[NVFP4TextEncoder] WARN: {len(unexpected)} keys in checkpoint unused; "
+            f"first 5: {unexpected[:5]}"
+        )
+
+
+def _process_weights_after_loading(model: nn.Module) -> None:
+    """Invoke vLLM's per-layer ``process_weights_after_loading`` for each
+    ReplicatedLinear (renames ``weight_packed`` -> ``weight``, computes ``alpha``,
+    swizzles scales for the CUTLASS kernel, etc.)."""
+    for module in model.modules():
+        qm = getattr(module, "quant_method", None)
+        if qm is not None and hasattr(qm, "process_weights_after_loading"):
+            qm.process_weights_after_loading(module)
+
+
+# ----------------------------------------------------------------------------
+# Public API
+# ----------------------------------------------------------------------------
+def load_nvfp4_text_encoder(
+    model_dir: str,
+    device: str | torch.device = "cuda",
+    dtype: torch.dtype = torch.bfloat16,
+) -> nn.Module:
+    """Load an NVFP4-quantised HuggingFace causal LM as a plug-in text encoder.
+
+    Args:
+        model_dir: path to the checkpoint directory containing ``config.json``
+            and ``model*.safetensors``. The config must carry a
+            ``quantization_config`` block with ``"format": "nvfp4-pack-quantized"``.
+        device: target CUDA device (forwards to ``model.to(device)``-equivalent
+            during materialisation).
+        dtype: activation / non-quantised-param dtype.
+
+    Returns:
+        A ``PreTrainedModel`` whose ``Linear`` layers are NVFP4 inside the vLLM
+        CUTLASS kernel. Activations flow as ``dtype``.
+    """
+    _bootstrap_vllm_once()
+
+    from transformers import AutoConfig, AutoModelForCausalLM
+    from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (
+        CompressedTensorsConfig,
+    )
+    from vllm.model_executor.models.transformers.utils import (
+        init_on_device_without_buffers,
+    )
+
+    hf_config = AutoConfig.from_pretrained(model_dir, local_files_only=True)
+    if not getattr(hf_config, "quantization_config", None):
+        raise ValueError(
+            f"{model_dir}/config.json has no `quantization_config`; "
+            "this loader only handles NVFP4-quantised checkpoints."
+        )
+    quant_config = CompressedTensorsConfig.from_config(hf_config.quantization_config)
+
+    # 1) Build the model skeleton on meta (zero allocation).
+    with init_on_device_without_buffers("meta"):
+        model = AutoModelForCausalLM.from_config(hf_config)
+
+    # 2) Swap Linear -> ReplicatedLinear(quant_config) (creates real CUDA params
+    #    of the quantised shapes).
+    target_device = torch.device(device)
+    _replace_linears_with_replicated(model, quant_config)
+
+    # 3) Materialise any leftover meta parameters (embeddings, RMSNorms, ...)
+    _materialize_remaining_meta_params(model, dtype=dtype, device=target_device)
+
+    # 4) Move newly-created quantised params to target device (ReplicatedLinear
+    #    creates them on the current default device which is usually CPU).
+    model.to(target_device)
+
+    # 5) Load weights via per-param weight_loader.
+    _load_weights_into_model(model, model_dir)
+
+    # 6) Let vLLM swizzle scales / rename weight_packed->weight / compute alpha.
+    _process_weights_after_loading(model)
+
+    # 7) Match HF semantics for downstream pipelines.
+    model.eval()
+    model.config.use_cache = False
+    return model

extras/README.md

@@ -0,0 +1,162 @@
+# FLUX.2 Klein 9B KV Image Generation & Multi-Reference Edit Server
+
+This directory contains the self-contained image generation and editing server integrating **`FLUX.2-klein-9b-kv-Nunchaku`** with this **`FLUX.2-klein-text-encoder-NVFP4`** text encoder.
+
+It enables fast, 4-step distilled text-to-image generation and multi-reference image editing (with KV caching) on single consumer GPUs (e.g., 16 GB RTX 5060 Ti / RTX 4080 / RTX 5090) without running out of memory.
+
+---
+
+## 🧪 Quantization Details & Strategy
+
+The text encoder was quantized using `llmcompressor` / `compressed-tensors` via [`quantize_flux2_klein_text_encoder_nvfp4.py`](file:///auto/home/amano/olegk/Nikola/src/nvfp4/quantize_flux2_klein_text_encoder_nvfp4.py) based on the recipe in [`recipe.yaml`](../recipe.yaml).
+
+### High-Fidelity Representation Preservation Strategy
+
+FLUX.2 Klein conditions its flow transformer on intermediate representations extracted from specific layers of the Qwen3 text encoder (specifically layer indices **9, 18, and 27**). Blanket quantization across all layers degrades prompt adherence and text grounding. To prevent this, the following selective quantization scheme was applied:
+
+1. **Unquantized Self-Attention Projections (`q_proj`, `k_proj`, `v_proj`, `o_proj`)**:
+   - Kept in full **bfloat16** across all 36 layers.
+   - Preserves semantic query-key routing and relational representations without quantization noise.
+
+2. **Unquantized Root & Near-Tap Layers**:
+   - Layer 0 (root) and layers **8, 9, 17, 18, 26, 27** (the critical tap layers and their immediate predecessors) are preserved in full **bfloat16**.
+   - Ensures the hidden state outputs at indices 9, 18, and 27 have near-perfect representation fidelity (cosine similarity > 0.999 vs unquantized reference).
+
+3. **Compressed Intermediate MLP Blocks**:
+   - The MLP blocks (`gate_proj`, `up_proj`, `down_proj`) of intermediate layers (1–7, 10–16, 19–25, 28–35) are compressed to **NVFP4 (W4A4)** using tensor-group scaling (`format: nvfp4-pack-quantized`).
+   - Yields ~4x compression of MLP weights, reducing memory footprint and load times significantly.
+
+4. **Preserved Head & Embeddings**:
+   - `lm_head` and `embed_tokens` are kept unquantized in **bfloat16**.
+
+---
+
+## 📦 Required Packages in Python Environment
+
+The server runs on Python 3.11 or 3.12 with CUDA 12.8 / 13.0 (e.g. Blackwell sm_120 / Thor sm_110):
+
+| Package | Minimum Version | Purpose |
+| :--- | :---: | :--- |
+| `torch` | `2.11.0+` | PyTorch with Blackwell/sm_120 CUDA compute support |
+| `diffusers` | `0.38.0+` | Built with native `Flux2KleinKVPipeline` support |
+| `transformers` | `5.3.0+` | Qwen3 causal language model architecture support |
+| `vllm` | `0.20.0+` | Provides precompiled CUTLASS NVFP4 GEMM kernels |
+| `nunchaku` | `1.2.1+` | SVDQuant FP4/INT4 4-bit transformer execution |
+| `fastapi`, `uvicorn` | Latest | OpenAI-compatible REST server infrastructure |
+| `pillow` | Latest | Image processing and aspect-ratio padding |
+
+> [!IMPORTANT]
+> When installing packages in environments with custom CUDA builds, use `--no-deps` to avoid overwriting PyTorch or flash-attention wheels tuned for your compute capability.
+
+---
+
+## ⚡ Architecture & Sequential GPU Memory Management
+
+On a 16 GB VRAM GPU, loading both the NVFP4 text encoder (~10.1 GB) and the Nunchaku transformer (~5.8 GB) simultaneously would total ~15.9 GB, exceeding available VRAM and causing an out-of-memory crash.
+
+`Flux2KleinBackend.py` solves this using **Sequential GPU Stage Execution**:
+
+```
+[Idle: 0.01 GB VRAM]
+       │
+       ▼
+Stage 1: Load Text Encoder to CUDA ──> Encode Prompt & IDs ──> Offload to CPU [0.01 GB VRAM]
+       │
+       ▼ (if reference images passed)
+Stage 2: Load VAE to CUDA ───────────> Encode Image Latents ──> Offload to CPU [0.01 GB VRAM]
+       │
+       ▼
+Stage 3: Load Transformer to CUDA ───> 4-Step Denoise (KV) ──> Offload to CPU [0.01 GB VRAM]
+       │                                (Step 0: Extract KV)
+       │                                (Steps 1-3: Cached KV)
+       ▼
+Stage 4: Load VAE to CUDA ───────────> Decode Latents to PIL ─> Offload to CPU [0.01 GB VRAM]
+       │
+       ▼
+[Idle: 0.01 GB VRAM] (GPU completely freed for other tasks!)
+```
+
+### Performance & Memory Metrics (RTX 5060 Ti 16 GB)
+
+- **Text-to-Image (1024x1024, 4 steps)**: ~9.5s total (~0.5s/step), **10.43 GB** peak VRAM.
+- **1-Image Edit (1024x1024, 4 steps)**: ~12.5s total, **10.46 GB** peak VRAM.
+- **2-Image Multi-Reference Edit (1024x1024, 4 steps)**: ~15.0s total, **13.18 GB** peak VRAM.
+- **Idle Memory Footprint**: **0.010 GB** (10 MB).
+- **Concurrency**: All requests are serialized with an `asyncio.Lock` queue, preventing concurrent memory spikes.
+
+---
+
+## 🚀 Running the Server
+
+### 1. Launch via Helper Script
+
+```bash
+# Set paths (or edit defaults inside the script)
+export MODEL_PATH="/path/to/FLUX.2-klein-9b-kv-Nunchaku"
+export NVFP4_TEXT_ENCODER_PATH="/path/to/FLUX.2-klein-text-encoder-NVFP4"
+
+./imagegen_flux2_klein.sh 4501
+```
+
+### 2. Launch Manually via Python
+
+```bash
+export VLLM_NVFP4_GEMM_BACKEND=cutlass
+
+python ImageEditServer.py \
+    --port 4501 \
+    --model /path/to/FLUX.2-klein-9b-kv-Nunchaku \
+    --optimized-model /path/to/FLUX.2-klein-9b-kv-Nunchaku/svdq-fp4_r32-FLUX.2-klein-9b-kv-Nunchaku.safetensors \
+    --nvfp4-text-encoder /path/to/FLUX.2-klein-text-encoder-NVFP4 \
+    --backend flux2-klein \
+    --steps 4 \
+    --guidance-scale 1.0
+```
+
+---
+
+## 📡 API Usage Examples
+
+### 1. Text-to-Image Generation (`POST /v1/images/generations`)
+
+```bash
+curl -X POST http://localhost:4501/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "prompt": "A cute cheerful anime girl with a ponytail in a simple village dress smiling, anime style",
+    "size": "1024x1024",
+    "num_inference_steps": 4,
+    "seed": 42
+  }' | jq -r '.data[0].b64_json' | base64 -d > generated.png
+```
+
+### 2. Single-Image Reference Edit (`POST /v1/images/edits`)
+
+```bash
+curl -X POST http://localhost:4501/v1/images/edits \
+  -F "image=@generated.png" \
+  -F "prompt=The anime girl wearing stylish round glasses and a bright red ribbon in her hair, anime style" \
+  -F "num_inference_steps=4" \
+  -F "seed": 42 | jq -r '.data[0].b64_json' | base64 -d > edited_1img.png
+```
+
+### 3. Multi-Reference 2-Image Edit (`POST /v1/images/edits`)
+
+Pass multiple `-F "image=@..."` form fields:
+
+```bash
+curl -X POST http://localhost:4501/v1/images/edits \
+  -F "image=@girl.png" \
+  -F "image=@sofa.png" \
+  -F "prompt=The anime girl from image 1 is sitting comfortably on the sofa from image 2 in the living room, anime style" \
+  -F "num_inference_steps=4" \
+  -F "seed": 42 | jq -r '.data[0].b64_json' | base64 -d > edited_2imgs.png
+```
+
+### 4. Memory Introspection (`GET /v1/memory_stats`)
+
+```bash
+curl http://localhost:4501/v1/memory_stats | jq .
+```
+Expected output when idle: `"allocated_gb": 0.01`, `"reserved_gb": 0.02`.
+

extras/imagegen_flux2_klein.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+# Activate your python virtual environment containing torch, vllm, nunchaku, diffusers
+# source /path/to/venv/bin/activate
+
+# Force vLLM's built-in CUTLASS NVFP4 kernel
+export VLLM_NVFP4_GEMM_BACKEND=cutlass
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+cd "$SCRIPT_DIR"
+
+PORT=${PORT:-${1:-4501}}
+MODEL_PATH=${MODEL_PATH:-"/path/to/FLUX.2-klein-9b-kv-Nunchaku"}
+OPTIMIZED_MODEL_PATH=${OPTIMIZED_MODEL_PATH:-"$MODEL_PATH/svdq-fp4_r32-FLUX.2-klein-9b-kv-Nunchaku.safetensors"}
+NVFP4_TEXT_ENCODER_PATH=${NVFP4_TEXT_ENCODER_PATH:-"/path/to/FLUX.2-klein-text-encoder-NVFP4"}
+
+python ImageEditServer.py \
+    --port "$PORT" \
+    --model "$MODEL_PATH" \
+    --optimized-model "$OPTIMIZED_MODEL_PATH" \
+    --nvfp4-text-encoder "$NVFP4_TEXT_ENCODER_PATH" \
+    --backend flux2-klein \
+    --steps 4 \
+    --guidance-scale 1.0

generation_config.json

@@ -0,0 +1,13 @@
+{
+  "bos_token_id": 151643,
+  "do_sample": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "pad_token_id": 151643,
+  "temperature": 0.6,
+  "top_k": 20,
+  "top_p": 0.95,
+  "transformers_version": "4.57.1"
+}

model-00001-of-00003.safetensors

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

model-00002-of-00003.safetensors

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

model-00003-of-00003.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,668 @@
+{
+  "metadata": {
+    "total_parameters": 6274987182,
+    "total_size": 10086869688
+  },
+  "weight_map": {
+    "lm_head.weight": "model-00003-of-00003.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.down_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.down_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.down_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.down_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.gate_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.gate_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.gate_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.gate_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.up_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.up_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.up_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.up_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.down_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.down_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.down_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.down_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.gate_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.gate_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.gate_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.gate_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.up_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.up_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.up_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.up_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.down_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.down_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.down_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.down_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.gate_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.gate_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.gate_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.gate_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.up_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.up_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.up_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.up_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.mlp.down_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.12.mlp.down_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.12.mlp.down_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.12.mlp.down_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.12.mlp.gate_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.12.mlp.gate_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.12.mlp.gate_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.12.mlp.gate_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.12.mlp.up_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.12.mlp.up_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.12.mlp.up_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.12.mlp.up_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.13.mlp.down_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.13.mlp.down_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.13.mlp.down_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.13.mlp.down_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.13.mlp.gate_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.13.mlp.gate_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.13.mlp.gate_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.13.mlp.gate_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.13.mlp.up_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.13.mlp.up_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.13.mlp.up_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.13.mlp.up_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.13.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.13.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.14.mlp.down_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.14.mlp.down_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.14.mlp.down_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.14.mlp.down_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.14.mlp.gate_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.14.mlp.gate_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.14.mlp.gate_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.14.mlp.gate_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.14.mlp.up_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.14.mlp.up_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.14.mlp.up_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.14.mlp.up_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.14.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.14.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.15.mlp.down_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.15.mlp.down_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.15.mlp.down_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.15.mlp.down_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.15.mlp.gate_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.15.mlp.gate_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.15.mlp.gate_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.15.mlp.gate_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.15.mlp.up_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.15.mlp.up_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.15.mlp.up_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.15.mlp.up_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.15.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.15.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.16.mlp.down_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.16.mlp.down_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.16.mlp.down_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.16.mlp.down_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.16.mlp.gate_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.16.mlp.gate_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.16.mlp.gate_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.16.mlp.gate_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.16.mlp.up_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.16.mlp.up_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.16.mlp.up_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.16.mlp.up_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.16.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.16.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.17.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.down_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.down_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.down_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.down_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.gate_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.gate_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.gate_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.gate_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.up_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.up_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.up_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.up_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.down_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.down_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.down_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.down_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.gate_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.gate_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.gate_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.gate_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.up_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.up_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.up_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.up_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.down_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.down_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.down_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.down_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.gate_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.gate_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.gate_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.gate_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.up_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.up_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.up_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.up_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.down_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.down_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.down_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.down_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.gate_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.gate_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.gate_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.gate_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.up_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.up_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.up_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.up_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.down_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.down_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.down_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.down_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.gate_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.gate_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.gate_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.gate_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.up_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.up_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.up_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.up_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.down_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.down_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.down_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.down_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.gate_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.gate_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.gate_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.gate_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.up_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.up_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.up_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.up_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.down_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.down_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.down_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.down_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.gate_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.gate_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.gate_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.gate_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.up_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.up_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.up_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.up_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.down_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.down_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.down_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.down_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.gate_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.gate_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.gate_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.gate_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.up_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.up_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.up_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.up_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.26.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.28.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.28.mlp.down_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.28.mlp.down_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.28.mlp.down_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.28.mlp.down_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.28.mlp.gate_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.28.mlp.gate_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.28.mlp.gate_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.28.mlp.gate_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.28.mlp.up_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.28.mlp.up_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.28.mlp.up_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.28.mlp.up_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.28.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.28.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.28.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.28.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.28.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.28.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.28.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.29.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.29.mlp.down_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.29.mlp.down_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.29.mlp.down_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.29.mlp.down_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.29.mlp.gate_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.29.mlp.gate_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.29.mlp.gate_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.29.mlp.gate_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.29.mlp.up_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.29.mlp.up_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.29.mlp.up_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.29.mlp.up_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.29.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.29.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.29.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.29.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.29.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.29.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.29.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.down_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.down_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.down_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.down_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.gate_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.gate_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.gate_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.gate_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.up_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.up_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.up_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.up_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.30.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.30.mlp.down_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.30.mlp.down_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.30.mlp.down_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.30.mlp.down_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.30.mlp.gate_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.30.mlp.gate_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.30.mlp.gate_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.30.mlp.gate_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.30.mlp.up_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.30.mlp.up_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.30.mlp.up_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.30.mlp.up_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.30.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.30.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.30.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.30.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.30.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.30.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.30.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.31.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.31.mlp.down_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.31.mlp.down_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.31.mlp.down_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.31.mlp.down_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.31.mlp.gate_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.31.mlp.gate_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.31.mlp.gate_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.31.mlp.gate_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.31.mlp.up_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.31.mlp.up_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.31.mlp.up_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.31.mlp.up_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.31.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.31.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.31.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.31.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.31.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.31.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.31.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.32.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.32.mlp.down_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.32.mlp.down_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.32.mlp.down_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.32.mlp.down_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.32.mlp.gate_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.32.mlp.gate_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.32.mlp.gate_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.32.mlp.gate_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.32.mlp.up_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.32.mlp.up_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.32.mlp.up_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.32.mlp.up_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.32.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.32.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.32.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.32.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.32.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.32.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.32.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.33.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.33.mlp.down_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.33.mlp.down_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.33.mlp.down_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.33.mlp.down_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.33.mlp.gate_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.33.mlp.gate_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.33.mlp.gate_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.33.mlp.gate_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.33.mlp.up_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.33.mlp.up_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.33.mlp.up_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.33.mlp.up_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.33.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.33.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.33.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.33.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.33.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.33.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.33.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.34.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.34.mlp.down_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.34.mlp.down_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.34.mlp.down_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.34.mlp.down_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.34.mlp.gate_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.34.mlp.gate_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.34.mlp.gate_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.34.mlp.gate_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.34.mlp.up_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.34.mlp.up_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.34.mlp.up_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.34.mlp.up_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.34.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.34.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.34.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.34.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.34.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.34.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.34.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.35.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.35.mlp.down_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.35.mlp.down_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.35.mlp.down_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.35.mlp.down_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.35.mlp.gate_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.35.mlp.gate_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.35.mlp.gate_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.35.mlp.gate_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.35.mlp.up_proj.input_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.35.mlp.up_proj.weight_global_scale": "model-00002-of-00003.safetensors",
+    "model.layers.35.mlp.up_proj.weight_packed": "model-00002-of-00003.safetensors",
+    "model.layers.35.mlp.up_proj.weight_scale": "model-00002-of-00003.safetensors",
+    "model.layers.35.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.35.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.35.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.35.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.35.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.35.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.35.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.down_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.down_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.down_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.down_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.gate_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.gate_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.gate_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.gate_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.up_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.up_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.up_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.up_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.down_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.down_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.down_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.down_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.gate_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.gate_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.gate_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.gate_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.up_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.up_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.up_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.up_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.down_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.down_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.down_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.down_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.gate_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.gate_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.gate_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.gate_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.up_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.up_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.up_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.up_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.down_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.down_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.down_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.down_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.gate_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.gate_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.gate_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.gate_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.up_proj.input_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.up_proj.weight_global_scale": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.up_proj.weight_packed": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.up_proj.weight_scale": "model-00001-of-00003.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "model.norm.weight": "model-00002-of-00003.safetensors"
+  }
+}

recipe.yaml

@@ -0,0 +1,6 @@
+default_stage:
+  default_modifiers:
+    QuantizationModifier:
+      targets: [Linear]
+      ignore: [lm_head, 're:.*self_attn.*', 're:.*layers\.(0|8|9|17|18|26|27)\.mlp.*']
+      scheme: NVFP4

special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,239 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}