Upload 50 files

.gitattributes

@@ -1,35 +1,49 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/architecture.webp filter=lfs diff=lfs merge=lfs -text
+assets/decoupled-dmd.webp filter=lfs diff=lfs merge=lfs -text
+assets/DMDR.webp filter=lfs diff=lfs merge=lfs -text
+assets/leaderboard.png filter=lfs diff=lfs merge=lfs -text
+assets/reasoning.png filter=lfs diff=lfs merge=lfs -text
+assets/showcase_editing.png filter=lfs diff=lfs merge=lfs -text
+assets/showcase_realistic.png filter=lfs diff=lfs merge=lfs -text
+assets/showcase_rendering.png filter=lfs diff=lfs merge=lfs -text
+assets/showcase.jpg filter=lfs diff=lfs merge=lfs -text
+assets/Z-Image-Gallery.pdf filter=lfs diff=lfs merge=lfs -text
+example/depth_cat.png filter=lfs diff=lfs merge=lfs -text
+example/depth_man.png filter=lfs diff=lfs merge=lfs -text
+tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+z_image_turbo_control_unified_q4_k_m.gguf filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,14 @@
+__pycache__/
+*.py[cod]
+mod_tests/
+/.vs
+.vscode/
+.idea/
+venv/
+.venv/
+*.log
+.DS_Store
+.gradio
+download.py
+bk
+outputs/

README.md

@@ -1,3 +1,81 @@
 ---
 license: apache-2.0
+tags:
+- text-to-image
+- image-to-image
+- controlnet
+- diffusers
+- gguf
+- z-image-turbo
+pipeline_tag: text-to-image
 ---
+
+# Z-Image Turbo Control Unified
+
+This repository hosts the **Z-Image Turbo Control Unified** model. This is a specialized architecture that unifies the powerful **Z-Image Turbo** base transformer with **ControlNet** capabilities into a single, cohesive architecture.
+
+Unlike traditional pipelines where ControlNet is an external add-on, this model integrates control layers directly into the transformer structure. This enables **Unified GGUF Quantization**, allowing the entire merged architecture (Base + Control) to be quantized (e.g., Q4_K_M) and run on consumer hardware with limited VRAM.
+
+## 📥 Installation
+
+To set up the environment, simply install the dependencies using the provided requirements file:
+
+```bash
+pip install -r requirements.txt
+```
+
+*Note: This repository contains a `diffusers_local` folder with custom pipelines required to run this specific architecture.*
+
+## 🚀 Usage
+
+We provide two ready-to-use scripts for inference, depending on your hardware capabilities and requirements.
+
+### Option 1: Low VRAM (GGUF) - Recommended
+**Script:** `infer_gguf.py`
+
+Use this version if you have limited VRAM (e.g., 6GB - 8GB) or want to save memory. It loads the model from the quantized **GGUF** file (`z_image_turbo_control_unified_q4_k_m.gguf`).
+
+To run:
+```bash
+python infer_gguf.py
+```
+
+**Key Features of this mode:**
+*   Loads the unified transformer from a single 4-bit quantized file.
+*   Uses `GGUFQuantizationConfig` for efficient computation.
+*   Enables aggressive group offloading to fit large models in consumer GPUs.
+
+### Option 2: High Precision (Diffusers/BF16)
+**Script:** `infer_pretrained.py`
+
+Use this version if you have ample VRAM (e.g., 24GB+) and want to run the model in standard **BFloat16** precision without quantization.
+
+To run:
+```bash
+python infer_pretrained.py
+```
+
+**Key Features of this mode:**
+*   Loads the model using the standard `from_pretrained` directory structure.
+*   maintains full floating-point precision.
+
+---
+
+## 🛠️ Model Configuration
+
+The inference scripts are pre-configured with parameters optimized for the **Turbo** nature of this model:
+
+*   **Inference Steps:** 9 steps (Fast generation).
+*   **Guidance Scale:** 0.0 (Turbo models do not use CFG).
+*   **Conditioning Scale:** 0.7 (Recommended strength for ControlNet).
+*   **Shift:** 3.0 (Scheduler shift parameter).
+
+## 📂 Repository Structure
+
+*   `z_image_turbo_control_unified_q4_k_m.gguf`: The unified, quantized model weights.
+*   `infer_gguf.py`: Script for running GGUF inference.
+*   `infer_pretrained.py`: Script for running standard Diffusers inference.
+*   `diffusers_local/`: Custom pipeline code (`ZImageControlUnifiedPipeline`) and transformer logic.
+*   `requirements.txt`: Python dependencies.
+
+

assets/Z-Image-Gallery.pdf

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

diffusers_local/__init__.py

@@ -0,0 +1 @@
+from .z_image_control_transformer_2d import Transformer2DModelOutput, ZImageControlTransformer2DModel

diffusers_local/patch.py

@@ -0,0 +1,246 @@
+import importlib
+import os
+import diffusers.pipelines.pipeline_loading_utils as pipe_loading_utils
+import diffusers.loaders.single_file_model as single_file_model
+from diffusers.utils import (
+    _maybe_remap_transformers_class,
+    get_class_from_dynamic_module,
+)
+from diffusers.loaders.single_file_utils import (
+    convert_animatediff_checkpoint_to_diffusers,
+    convert_auraflow_transformer_checkpoint_to_diffusers,
+    convert_autoencoder_dc_checkpoint_to_diffusers,
+    convert_chroma_transformer_checkpoint_to_diffusers,
+    convert_controlnet_checkpoint,
+    convert_cosmos_transformer_checkpoint_to_diffusers,
+    convert_flux2_transformer_checkpoint_to_diffusers,
+    convert_flux_transformer_checkpoint_to_diffusers,
+    convert_hidream_transformer_to_diffusers,
+    convert_hunyuan_video_transformer_to_diffusers,
+    convert_ldm_unet_checkpoint,
+    convert_ldm_vae_checkpoint,
+    convert_ltx_transformer_checkpoint_to_diffusers,
+    convert_ltx_vae_checkpoint_to_diffusers,
+    convert_lumina2_to_diffusers,
+    convert_mochi_transformer_checkpoint_to_diffusers,
+    convert_sana_transformer_to_diffusers,
+    convert_sd3_transformer_checkpoint_to_diffusers,
+    convert_stable_cascade_unet_single_file_to_diffusers,
+    convert_wan_transformer_to_diffusers,
+    convert_wan_vae_to_diffusers,
+    convert_z_image_transformer_checkpoint_to_diffusers,
+    create_controlnet_diffusers_config_from_ldm,
+    create_unet_diffusers_config_from_ldm,
+    create_vae_diffusers_config_from_ldm,
+)
+import torch
+def convert_z_image_control_transformer_checkpoint_to_diffusers(checkpoint, **kwargs):
+    Z_IMAGE_KEYS_RENAME_DICT = {
+        "final_layer.": "all_final_layer.2-1.",
+        "x_embedder.": "all_x_embedder.2-1.",
+        ".attention.out.bias": ".attention.to_out.0.bias",
+        ".attention.k_norm.weight": ".attention.norm_k.weight",
+        ".attention.q_norm.weight": ".attention.norm_q.weight",
+        ".attention.out.weight": ".attention.to_out.0.weight",
+        "control_x_embedder.": "control_all_x_embedder.2-1.",
+    }
+
+    def convert_z_image_fused_attention(key: str, state_dict: dict[str, object]) -> None:
+        if ".attention.qkv.weight" not in key:
+            return
+
+        fused_qkv_weight = state_dict.pop(key)
+        to_q_weight, to_k_weight, to_v_weight = torch.chunk(fused_qkv_weight, 3, dim=0)
+        new_q_name = key.replace(".attention.qkv.weight", ".attention.to_q.weight")
+        new_k_name = key.replace(".attention.qkv.weight", ".attention.to_k.weight")
+        new_v_name = key.replace(".attention.qkv.weight", ".attention.to_v.weight")
+
+        state_dict[new_q_name] = to_q_weight
+        state_dict[new_k_name] = to_k_weight
+        state_dict[new_v_name] = to_v_weight
+        return
+
+    TRANSFORMER_SPECIAL_KEYS_REMAP = {
+        ".attention.qkv.weight": convert_z_image_fused_attention,
+    }
+
+    def update_state_dict(state_dict: dict[str, object], old_key: str, new_key: str) -> None:
+        state_dict[new_key] = state_dict.pop(old_key)
+
+    converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())}
+
+    # Handle single file --> diffusers key remapping via the remap dict
+    for key in list(converted_state_dict.keys()):
+        new_key = key[:]
+        for replace_key, rename_key in Z_IMAGE_KEYS_RENAME_DICT.items():
+            new_key = new_key.replace(replace_key, rename_key)
+
+        update_state_dict(converted_state_dict, key, new_key)
+
+    # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in
+    # special_keys_remap
+    for key in list(converted_state_dict.keys()):
+        for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items():
+            if special_key not in key:
+                continue
+            handler_fn_inplace(key, converted_state_dict)
+
+    return converted_state_dict
+
+SINGLE_FILE_LOADABLE_CLASSES = {
+    "StableCascadeUNet": {
+        "checkpoint_mapping_fn": convert_stable_cascade_unet_single_file_to_diffusers,
+    },
+    "UNet2DConditionModel": {
+        "checkpoint_mapping_fn": convert_ldm_unet_checkpoint,
+        "config_mapping_fn": create_unet_diffusers_config_from_ldm,
+        "default_subfolder": "unet",
+        "legacy_kwargs": {
+            "num_in_channels": "in_channels",  # Legacy kwargs supported by `from_single_file` mapped to new args
+        },
+    },
+    "AutoencoderKL": {
+        "checkpoint_mapping_fn": convert_ldm_vae_checkpoint,
+        "config_mapping_fn": create_vae_diffusers_config_from_ldm,
+        "default_subfolder": "vae",
+    },
+    "ControlNetModel": {
+        "checkpoint_mapping_fn": convert_controlnet_checkpoint,
+        "config_mapping_fn": create_controlnet_diffusers_config_from_ldm,
+    },
+    "SD3Transformer2DModel": {
+        "checkpoint_mapping_fn": convert_sd3_transformer_checkpoint_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "MotionAdapter": {
+        "checkpoint_mapping_fn": convert_animatediff_checkpoint_to_diffusers,
+    },
+    "SparseControlNetModel": {
+        "checkpoint_mapping_fn": convert_animatediff_checkpoint_to_diffusers,
+    },
+    "FluxTransformer2DModel": {
+        "checkpoint_mapping_fn": convert_flux_transformer_checkpoint_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "ChromaTransformer2DModel": {
+        "checkpoint_mapping_fn": convert_chroma_transformer_checkpoint_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "LTXVideoTransformer3DModel": {
+        "checkpoint_mapping_fn": convert_ltx_transformer_checkpoint_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "AutoencoderKLLTXVideo": {
+        "checkpoint_mapping_fn": convert_ltx_vae_checkpoint_to_diffusers,
+        "default_subfolder": "vae",
+    },
+    "AutoencoderDC": {"checkpoint_mapping_fn": convert_autoencoder_dc_checkpoint_to_diffusers},
+    "MochiTransformer3DModel": {
+        "checkpoint_mapping_fn": convert_mochi_transformer_checkpoint_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "HunyuanVideoTransformer3DModel": {
+        "checkpoint_mapping_fn": convert_hunyuan_video_transformer_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "AuraFlowTransformer2DModel": {
+        "checkpoint_mapping_fn": convert_auraflow_transformer_checkpoint_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "Lumina2Transformer2DModel": {
+        "checkpoint_mapping_fn": convert_lumina2_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "SanaTransformer2DModel": {
+        "checkpoint_mapping_fn": convert_sana_transformer_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "WanTransformer3DModel": {
+        "checkpoint_mapping_fn": convert_wan_transformer_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "WanVACETransformer3DModel": {
+        "checkpoint_mapping_fn": convert_wan_transformer_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "AutoencoderKLWan": {
+        "checkpoint_mapping_fn": convert_wan_vae_to_diffusers,
+        "default_subfolder": "vae",
+    },
+    "HiDreamImageTransformer2DModel": {
+        "checkpoint_mapping_fn": convert_hidream_transformer_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "CosmosTransformer3DModel": {
+        "checkpoint_mapping_fn": convert_cosmos_transformer_checkpoint_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "QwenImageTransformer2DModel": {
+        "checkpoint_mapping_fn": lambda x: x,
+        "default_subfolder": "transformer",
+    },
+    "Flux2Transformer2DModel": {
+        "checkpoint_mapping_fn": convert_flux2_transformer_checkpoint_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "ZImageTransformer2DModel": {
+        "checkpoint_mapping_fn": convert_z_image_transformer_checkpoint_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "ZImageControlTransformer2DModel": {
+        "checkpoint_mapping_fn": convert_z_image_control_transformer_checkpoint_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+}
+
+def get_class_obj_and_candidates(
+    library_name, class_name, importable_classes, pipelines, is_pipeline_module, component_name=None, cache_dir=None
+):
+    """Simple helper method to retrieve class object of module as well as potential parent class objects"""
+    component_folder = os.path.join(cache_dir, component_name) if component_name and cache_dir else None
+
+    if is_pipeline_module:
+        pipeline_module = getattr(pipelines, library_name)
+
+        class_obj = getattr(pipeline_module, class_name)
+        class_candidates = dict.fromkeys(importable_classes.keys(), class_obj)
+    elif component_folder and os.path.isfile(os.path.join(component_folder, library_name + ".py")):
+        # load custom component
+        class_obj = get_class_from_dynamic_module(
+            component_folder, module_file=library_name + ".py", class_name=class_name
+        )
+        class_candidates = dict.fromkeys(importable_classes.keys(), class_obj)
+    else:
+        # else we just import it from the library.
+        library = importlib.import_module(library_name)
+
+        # Handle deprecated Transformers classes
+        if library_name == "transformers":
+            class_name = _maybe_remap_transformers_class(class_name) or class_name
+        
+        try:
+            class_obj = getattr(library, class_name)
+        except:
+            module = importlib.import_module("diffusers_local")
+            class_obj = getattr(module, class_name)
+        class_candidates = {c: getattr(library, c, None) for c in importable_classes.keys()}
+
+    return class_obj, class_candidates
+
+def _get_single_file_loadable_mapping_class(cls):
+    diffusers_module = importlib.import_module("diffusers")
+    class_name_str = cls.__name__
+    for loadable_class_str in SINGLE_FILE_LOADABLE_CLASSES:
+        try:
+            loadable_class = getattr(diffusers_module, loadable_class_str)
+        except:
+            module = importlib.import_module("diffusers_local")
+            loadable_class = getattr(module, loadable_class_str)
+        if issubclass(cls, loadable_class):
+            return loadable_class_str
+
+    return class_name_str
+
+pipe_loading_utils.get_class_obj_and_candidates = get_class_obj_and_candidates
+single_file_model.SINGLE_FILE_LOADABLE_CLASSES = SINGLE_FILE_LOADABLE_CLASSES
+single_file_model._get_single_file_loadable_mapping_class = _get_single_file_loadable_mapping_class

diffusers_local/pipeline_z_image_control_unified.py

@@ -0,0 +1,288 @@
+# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import List, Optional, Union
+import torch
+from PIL import Image
+
+from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, DiffusionPipeline
+from diffusers.loaders import FromSingleFileMixin, ZImageLoraLoaderMixin
+from diffusers.image_processor import VaeImageProcessor 
+from diffusers.utils import logging
+from diffusers.pipelines.z_image.pipeline_z_image import calculate_shift
+from diffusers.pipelines.z_image.pipeline_output import ZImagePipelineOutput
+from diffusers_local.z_image_control_transformer_2d import ZImageControlTransformer2DModel
+from transformers import AutoTokenizer, PreTrainedModel
+from diffusers.utils.torch_utils import randn_tensor
+
+logger = logging.get_logger(__name__)
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class ZImageControlUnifiedPipeline(DiffusionPipeline, ZImageLoraLoaderMixin, FromSingleFileMixin):
+    _model_cpu_offload_seq = "text_encoder->transformer->vae"
+    _optional_components = []
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+    
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: PreTrainedModel,
+        tokenizer: AutoTokenizer,
+        transformer: ZImageControlTransformer2DModel,
+    ):      
+        self.register_modules(
+            vae=vae, text_encoder=text_encoder, tokenizer=tokenizer,
+            transformer=transformer, scheduler=scheduler
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+
+    def _encode_prompt(self, prompt: str, device: torch.device, max_sequence_length: int) -> torch.Tensor:
+        messages = [{"role": "user", "content": prompt}]
+        if hasattr(self.tokenizer, "apply_chat_template"):
+            prompt_formatted = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=True)
+        else:
+            prompt_formatted = prompt
+            
+        text_inputs = self.tokenizer(prompt_formatted, padding="max_length", max_length=max_sequence_length, truncation=True, return_tensors="pt").to(device)
+        prompt_masks = text_inputs.attention_mask.bool()
+        with torch.no_grad():
+            prompt_embeds = self.text_encoder(input_ids=text_inputs.input_ids, attention_mask=prompt_masks, output_hidden_states=True).hidden_states[-2]
+        return prompt_embeds[0][prompt_masks[0]]
+
+    def prepare_latents(self, batch_size, num_channels, height, width, dtype, device, generator, latents=None):
+        shape = (batch_size, num_channels, height // self.vae_scale_factor, width // self.vae_scale_factor)
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+        return latents * self.scheduler.init_noise_sigma if hasattr(self.scheduler, "init_noise_sigma") else latents
+    
+    def prepare_control_image(self, image, width, height, batch_size, num_images_per_prompt, device, dtype):
+        image = self.image_processor.preprocess(image, height=height, width=width).to(device=device, dtype=dtype)
+                
+        image_batch_size = image.shape[0]
+        if image_batch_size == 1:
+            repeat_by = batch_size
+        else:
+            repeat_by = num_images_per_prompt
+        image = image.repeat_interleave(repeat_by, dim=0)
+        return image
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        image: Union[torch.Tensor, Image.Image],
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 0.0,
+        controlnet_conditioning_scale: float = 1.0,
+        num_images_per_prompt: int = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        **kwargs,
+    ):
+        device = self._execution_device
+        height = height or image.height
+        width = width or image.width
+        
+        # 1. Prompt adjustment and batch size
+        if isinstance(prompt, str): prompt = [prompt]
+        if isinstance(negative_prompt, str): negative_prompt = [negative_prompt]
+        
+        batch_size = len(prompt) * num_images_per_prompt
+        do_cfg = guidance_scale > 0.0
+
+        # 2. Encode text
+        # Repeat embeddings if num_images_per_prompt > 1
+        prompt_embeds_list = []
+        for p in prompt:
+            embed = self._encode_prompt(p, device, 512)
+            for _ in range(num_images_per_prompt):
+                prompt_embeds_list.append(embed)
+        
+        if do_cfg:
+            if negative_prompt is None: negative_prompt = [""] * len(prompt)
+            neg_embeds_list = []
+            for np in negative_prompt:
+                embed = self._encode_prompt(np, device, 512)
+                for _ in range(num_images_per_prompt):
+                    neg_embeds_list.append(embed)
+            
+            prompt_input = neg_embeds_list + prompt_embeds_list
+        else:
+            prompt_input = prompt_embeds_list
+
+       
+        # 3. Control image preparation
+        control_tensor = self.prepare_control_image(
+            image, width, height, batch_size, num_images_per_prompt, device, self.vae.dtype
+        )
+        
+        if len(control_tensor.shape) == 3:
+            control_tensor = control_tensor.unsqueeze(0)
+            
+        with torch.no_grad():
+            # Encode to latents
+            control_latents = self.vae.encode(control_tensor).latent_dist.mode()
+            control_latents = control_latents * self.vae.config.scaling_factor
+
+        # Channel fix: 4 channels -> 16 channels
+        if control_latents.shape[1] == 4 and self.transformer.in_channels == 16:
+            control_latents = control_latents.repeat(1, 4, 1, 1)  # [B, 16, H, W]
+        
+        control_latents = control_latents.to(dtype=self.transformer.dtype)
+        
+        # Fix dimension: frame dimension [B, 16, 1, H, W]
+        control_latents = control_latents.unsqueeze(2)                
+        control_context = list(control_latents.unbind(0))
+        
+        # Expansion for CFG
+        if do_cfg:
+            control_context_input = control_context * 2
+        else:
+            control_context_input = control_context
+       
+        # 4. Initial latents
+        latents = self.prepare_latents(
+            batch_size, self.transformer.in_channels, height, width, 
+            prompt_embeds_list[0].dtype, device, generator
+        )
+        latents = latents.to(self.transformer.dtype)
+
+        # 5. Denoising loop
+        image_seq_len = (height // (self.vae_scale_factor)) * (width // (self.vae_scale_factor))
+        mu = calculate_shift(image_seq_len)
+        self.scheduler.set_timesteps(num_inference_steps, device=device, mu=mu)
+        
+        for t in self.progress_bar(self.scheduler.timesteps):
+            t_input = t.expand(len(prompt_input))
+            timestep_norm = (1000.0 - t_input) / 1000.0
+            
+            latents_input = torch.cat([latents] * 2) if do_cfg else latents
+            
+            # List of [16, 1, H, W]
+            latent_list = list(latents_input.unsqueeze(2).unbind(dim=0))
+            
+            model_out_list = self.transformer(
+                x=latent_list,
+                t=timestep_norm,
+                cap_feats=prompt_input,
+                control_context=control_context_input,
+                conditioning_scale=controlnet_conditioning_scale,
+            )[0]
+            
+            model_out = torch.stack(model_out_list, dim=0).squeeze(2)
+
+            if do_cfg:
+                neg_out, pos_out = model_out.chunk(2)
+                noise_pred = neg_out + guidance_scale * (pos_out - neg_out)
+            else:
+                noise_pred = model_out
+            
+            noise_pred = -noise_pred
+            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+            
+        # 6. Decode
+        if not output_type == "latent":
+            # Pass 16 channels to VAE
+            latents_for_vae = latents.to(self.vae.dtype)
+            latents_for_vae = (latents_for_vae / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+            
+            image = self.vae.decode(latents_for_vae, return_dict=False)[0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+        else:
+            image = latents
+        
+        self.maybe_free_model_hooks()
+        return ZImagePipelineOutput(images=image)

diffusers_local/z_image_control_transformer_2d.py

@@ -0,0 +1,815 @@
+# Copyright 2025 Alibaba Z-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils.rnn import pad_sequence
+
+from diffusers.configuration_utils import ConfigMixin,register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.models.attention_processor import Attention
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import RMSNorm
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+
+
+ADALN_EMBED_DIM = 256
+SEQ_MULTI_OF = 32
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, out_size, mid_size=None, frequency_embedding_size=256):
+        super().__init__()
+        if mid_size is None:
+            mid_size = out_size
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, mid_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(mid_size, out_size, bias=True),
+        )
+
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        with torch.amp.autocast("cuda", enabled=False):
+            half = dim // 2
+            freqs = torch.exp(
+                -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half
+            )
+            args = t[:, None].float() * freqs[None]
+            embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+            if dim % 2:
+                embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+            return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        weight_dtype = self.mlp[0].weight.dtype
+        compute_dtype = getattr(self.mlp[0], "compute_dtype", None)
+        if weight_dtype.is_floating_point:
+            t_freq = t_freq.to(weight_dtype)
+        elif compute_dtype is not None:
+            t_freq = t_freq.to(compute_dtype)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class ZSingleStreamAttnProcessor:
+    """
+    Processor for Z-Image single stream attention that adapts the existing Attention class to match the behavior of the
+    original Z-ImageAttention module.
+    """
+
+    _attention_backend = None
+    _parallel_config = None
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "ZSingleStreamAttnProcessor requires PyTorch 2.0. To use it, please upgrade PyTorch to version 2.0 or higher."
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        freqs_cis: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        query = query.unflatten(-1, (attn.heads, -1))
+        key = key.unflatten(-1, (attn.heads, -1))
+        value = value.unflatten(-1, (attn.heads, -1))
+
+        # Apply Norms
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # Apply RoPE
+        def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+            with torch.amp.autocast("cuda", enabled=False):
+                x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2))
+                freqs_cis = freqs_cis.unsqueeze(2)
+                x_out = torch.view_as_real(x * freqs_cis).flatten(3)
+                return x_out.type_as(x_in)  # todo
+
+        if freqs_cis is not None:
+            query = apply_rotary_emb(query, freqs_cis)
+            key = apply_rotary_emb(key, freqs_cis)
+
+        # Cast to correct dtype
+        dtype = query.dtype
+        query, key = query.to(dtype), key.to(dtype)
+
+        # From [batch, seq_len] to [batch, 1, 1, seq_len] -> broadcast to [batch, heads, seq_len, seq_len]
+        if attention_mask is not None and attention_mask.ndim == 2:
+            attention_mask = attention_mask[:, None, None, :]
+
+        # Compute joint attention
+        hidden_states = dispatch_attention_fn(
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            dropout_p=0.0,
+            is_causal=False,
+            backend=self._attention_backend,
+            parallel_config=self._parallel_config,
+        )
+
+        # Reshape back
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states.to(dtype)
+
+        output = attn.to_out[0](hidden_states)
+        if len(attn.to_out) > 1:  # dropout
+            output = attn.to_out[1](output)
+
+        return output
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int):
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def _forward_silu_gating(self, x1, x3):
+        return F.silu(x1) * x3
+
+    def forward(self, x):
+        return self.w2(self._forward_silu_gating(self.w1(x), self.w3(x)))
+
+def zero_module(module):
+    for p in module.parameters():
+        nn.init.zeros_(p)
+    return module
+
+@maybe_allow_in_graph
+class ZImageTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        layer_id: int,
+        dim: int,
+        n_heads: int,
+        n_kv_heads: int,
+        norm_eps: float,
+        qk_norm: bool,
+        modulation=True,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.head_dim = dim // n_heads
+
+        # Refactored to use diffusers Attention with custom processor
+        # Original Z-Image params: dim, n_heads, n_kv_heads, qk_norm
+        self.attention = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            dim_head=dim // n_heads,
+            heads=n_heads,
+            qk_norm="rms_norm" if qk_norm else None,
+            eps=1e-5,
+            bias=False,
+            out_bias=False,
+            processor=ZSingleStreamAttnProcessor(),
+        )
+
+        self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8))
+        self.layer_id = layer_id
+
+        self.attention_norm1 = RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm1 = RMSNorm(dim, eps=norm_eps)
+
+        self.attention_norm2 = RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm2 = RMSNorm(dim, eps=norm_eps)
+
+        self.modulation = modulation
+        if modulation:
+            self.adaLN_modulation = nn.Sequential(nn.Linear(min(dim, ADALN_EMBED_DIM), 4 * dim, bias=True))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        adaln_input: Optional[torch.Tensor] = None,
+    ):
+        if self.modulation:
+            assert adaln_input is not None
+            scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input).unsqueeze(1).chunk(4, dim=2)
+            gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
+            scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
+
+            # Attention block
+            attn_out = self.attention(
+                self.attention_norm1(x) * scale_msa, attention_mask=attn_mask, freqs_cis=freqs_cis
+            )
+            x = x + gate_msa * self.attention_norm2(attn_out)
+
+            # FFN block
+            x = x + gate_mlp * self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp))
+        else:
+            # Attention block
+            attn_out = self.attention(self.attention_norm1(x), attention_mask=attn_mask, freqs_cis=freqs_cis)
+            x = x + self.attention_norm2(attn_out)
+
+            # FFN block
+            x = x + self.ffn_norm2(self.feed_forward(self.ffn_norm1(x)))
+
+        return x
+    
+@maybe_allow_in_graph
+class ZImageControlTransformerBlock(ZImageTransformerBlock):
+    def __init__(
+        self,
+        layer_id: int,
+        dim: int,
+        n_heads: int,
+        n_kv_heads: int,
+        norm_eps: float,
+        qk_norm: bool,
+        modulation=True,
+        block_id=0,
+    ):
+        super().__init__(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm, modulation)
+        self.block_id = block_id
+        if block_id == 0:
+            self.before_proj = zero_module(nn.Linear(self.dim, self.dim))
+        self.after_proj = zero_module(nn.Linear(self.dim, self.dim))
+
+    def forward(
+        self,
+        c: torch.Tensor,
+        x: torch.Tensor,
+        attn_mask: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        adaln_input: Optional[torch.Tensor] = None,
+    ):
+        if self.block_id == 0:
+            c = self.before_proj(c) + x
+            all_c = []
+        else:
+            all_c = list(torch.unbind(c))
+            c = all_c.pop(-1)
+
+        c = super().forward(c, attn_mask, freqs_cis, adaln_input)
+        c_skip = self.after_proj(c)
+        all_c += [c_skip, c]
+        c = torch.stack(all_c)
+        return c
+    
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(min(hidden_size, ADALN_EMBED_DIM), hidden_size, bias=True),
+        )
+
+    def forward(self, x, c):
+        scale = 1.0 + self.adaLN_modulation(c)
+        x = self.norm_final(x) * scale.unsqueeze(1)
+        x = self.linear(x)
+        return x
+
+
+class RopeEmbedder:
+    def __init__(
+        self,
+        theta: float = 256.0,
+        axes_dims: List[int] = (16, 56, 56),
+        axes_lens: List[int] = (64, 128, 128),
+    ):
+        self.theta = theta
+        self.axes_dims = axes_dims
+        self.axes_lens = axes_lens
+        assert len(axes_dims) == len(axes_lens), "axes_dims and axes_lens must have the same length"
+        self.freqs_cis = None
+
+    @staticmethod
+    def precompute_freqs_cis(dim: List[int], end: List[int], theta: float = 256.0):
+        with torch.device("cpu"):
+            freqs_cis = []
+            for i, (d, e) in enumerate(zip(dim, end)):
+                freqs = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d))
+                timestep = torch.arange(e, device=freqs.device, dtype=torch.float64)
+                freqs = torch.outer(timestep, freqs).float()
+                freqs_cis_i = torch.polar(torch.ones_like(freqs), freqs).to(torch.complex64)  # complex64
+                freqs_cis.append(freqs_cis_i)
+
+            return freqs_cis
+
+    def __call__(self, ids: torch.Tensor):
+        assert ids.ndim == 2
+        assert ids.shape[-1] == len(self.axes_dims)
+        device = ids.device
+
+        if self.freqs_cis is None:
+            self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, self.axes_lens, theta=self.theta)
+            self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
+        else:
+            # Ensure freqs_cis are on the same device as ids
+            if self.freqs_cis[0].device != device:
+                self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
+
+        result = []
+        for i in range(len(self.axes_dims)):
+            index = ids[:, i]
+            result.append(self.freqs_cis[i][index])
+        return torch.cat(result, dim=-1)
+
+
+class ZImageControlTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["ZImageTransformerBlock", "ZImageControlTransformerBlock"]
+    _repeated_blocks = ["ZImageTransformerBlock", "ZImageControlTransformerBlock"]
+    _skip_layerwise_casting_patterns = ["t_embedder", "cap_embedder"]  # precision sensitive layers
+
+    @register_to_config
+    def __init__(
+        self,
+        all_patch_size=(2,),
+        all_f_patch_size=(1,),
+        in_channels=16,
+        dim=3840,
+        n_layers=30,
+        n_refiner_layers=2,
+        n_heads=30,
+        n_kv_heads=30,
+        norm_eps=1e-5,
+        qk_norm=True,
+        cap_feat_dim=2560,
+        rope_theta=256.0,
+        t_scale=1000.0,
+        axes_dims=[32, 48, 48],
+        axes_lens=[1024, 512, 512],
+        control_layers_places: List[int] = [0, 5, 10, 15, 20, 25],
+        control_in_dim=16,
+    ) -> None:
+        super().__init__()
+       
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.all_patch_size = all_patch_size
+        self.all_f_patch_size = all_f_patch_size
+        self.dim = dim
+        self.n_heads = n_heads
+
+        self.rope_theta = rope_theta
+        self.t_scale = t_scale
+        self.gradient_checkpointing = False
+
+        assert len(all_patch_size) == len(all_f_patch_size)
+
+        all_x_embedder = {}
+        all_final_layer = {}
+        for patch_idx, (patch_size, f_patch_size) in enumerate(zip(all_patch_size, all_f_patch_size)):
+            x_embedder = nn.Linear(f_patch_size * patch_size * patch_size * in_channels, dim, bias=True)
+            all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder
+
+            final_layer = FinalLayer(dim, patch_size * patch_size * f_patch_size * self.out_channels)
+            all_final_layer[f"{patch_size}-{f_patch_size}"] = final_layer
+
+        self.all_x_embedder = nn.ModuleDict(all_x_embedder)
+        self.all_final_layer = nn.ModuleDict(all_final_layer)
+        self.noise_refiner = nn.ModuleList(
+            [
+                ZImageTransformerBlock(
+                    1000 + layer_id,
+                    dim,
+                    n_heads,
+                    n_kv_heads,
+                    norm_eps,
+                    qk_norm,
+                    modulation=True,
+                )
+                for layer_id in range(n_refiner_layers)
+            ]
+        )
+        self.context_refiner = nn.ModuleList(
+            [
+                ZImageTransformerBlock(
+                    layer_id,
+                    dim,
+                    n_heads,
+                    n_kv_heads,
+                    norm_eps,
+                    qk_norm,
+                    modulation=False,
+                )
+                for layer_id in range(n_refiner_layers)
+            ]
+        )
+        self.t_embedder = TimestepEmbedder(min(dim, ADALN_EMBED_DIM), mid_size=1024)
+        self.cap_embedder = nn.Sequential(RMSNorm(cap_feat_dim, eps=norm_eps), nn.Linear(cap_feat_dim, dim, bias=True))
+
+        self.x_pad_token = nn.Parameter(torch.empty((1, dim)))
+        self.cap_pad_token = nn.Parameter(torch.empty((1, dim)))
+
+        self.layers = nn.ModuleList(
+            [
+                ZImageTransformerBlock(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm)
+                for layer_id in range(n_layers)
+            ]
+        )
+        head_dim = dim // n_heads
+        assert head_dim == sum(axes_dims)
+        self.axes_dims = axes_dims
+        self.axes_lens = axes_lens
+
+        self.rope_embedder = RopeEmbedder(theta=rope_theta, axes_dims=axes_dims, axes_lens=axes_lens)
+
+        self.control_layers_places = [i for i in range(0, self.n_layers, 2)] if control_layers_places is None else control_layers_places
+        self.control_in_dim = self.dim if control_in_dim is None else control_in_dim
+
+        assert 0 in self.control_layers_places
+
+        # control blocks
+        self.control_layers = nn.ModuleList(
+            [
+                ZImageControlTransformerBlock(i, dim, n_heads, n_kv_heads, norm_eps, qk_norm, block_id=i)
+                for i in self.control_layers_places
+            ]
+        )
+
+        # control patch embeddings
+        all_x_embedder = {}
+        for patch_idx, (patch_size, f_patch_size) in enumerate(zip(all_patch_size, all_f_patch_size)):
+            x_embedder = nn.Linear(f_patch_size * patch_size * patch_size * self.control_in_dim, dim, bias=True)
+            all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder
+
+        self.control_all_x_embedder = nn.ModuleDict(all_x_embedder)
+        self.control_noise_refiner = nn.ModuleList(
+            [
+                ZImageTransformerBlock(
+                    1000 + layer_id,
+                    dim,
+                    n_heads,
+                    n_kv_heads,
+                    norm_eps,
+                    qk_norm,
+                    modulation=True,
+                )
+                for layer_id in range(n_refiner_layers)
+            ]
+        )
+
+    def unpatchify(self, x: List[torch.Tensor], size: List[Tuple], patch_size, f_patch_size) -> List[torch.Tensor]:
+        pH = pW = patch_size
+        pF = f_patch_size
+        bsz = len(x)
+        assert len(size) == bsz
+        for i in range(bsz):
+            F, H, W = size[i]
+            ori_len = (F // pF) * (H // pH) * (W // pW)
+            # "f h w pf ph pw c -> c (f pf) (h ph) (w pw)"
+            x[i] = (
+                x[i][:ori_len]
+                .view(F // pF, H // pH, W // pW, pF, pH, pW, self.out_channels)
+                .permute(6, 0, 3, 1, 4, 2, 5)
+                .reshape(self.out_channels, F, H, W)
+            )
+        return x
+
+    @staticmethod
+    def create_coordinate_grid(size, start=None, device=None):
+        if start is None:
+            start = (0 for _ in size)
+
+        axes = [torch.arange(x0, x0 + span, dtype=torch.int32, device=device) for x0, span in zip(start, size)]
+        grids = torch.meshgrid(axes, indexing="ij")
+        return torch.stack(grids, dim=-1)
+
+    def patchify_and_embed(
+        self,
+        all_image: List[torch.Tensor],
+        all_cap_feats: List[torch.Tensor],
+        patch_size: int,
+        f_patch_size: int,
+    ):
+        pH = pW = patch_size
+        pF = f_patch_size
+        device = all_image[0].device
+
+        all_image_out = []
+        all_image_size = []
+        all_image_pos_ids = []
+        all_image_pad_mask = []
+        all_cap_pos_ids = []
+        all_cap_pad_mask = []
+        all_cap_feats_out = []
+
+        for i, (image, cap_feat) in enumerate(zip(all_image, all_cap_feats)):
+            ### Process Caption
+            cap_ori_len = len(cap_feat)
+            cap_padding_len = (-cap_ori_len) % SEQ_MULTI_OF
+            # padded position ids
+            cap_padded_pos_ids = self.create_coordinate_grid(
+                size=(cap_ori_len + cap_padding_len, 1, 1),
+                start=(1, 0, 0),
+                device=device,
+            ).flatten(0, 2)
+            all_cap_pos_ids.append(cap_padded_pos_ids)
+            # pad mask
+            cap_pad_mask = torch.cat(
+                [
+                    torch.zeros((cap_ori_len,), dtype=torch.bool, device=device),
+                    torch.ones((cap_padding_len,), dtype=torch.bool, device=device),
+                ],
+                dim=0,
+            )
+            all_cap_pad_mask.append(
+                cap_pad_mask if cap_padding_len > 0 else torch.zeros((cap_ori_len,), dtype=torch.bool, device=device)
+            )
+
+            # padded feature
+            cap_padded_feat = torch.cat([cap_feat, cap_feat[-1:].repeat(cap_padding_len, 1)], dim=0)
+            all_cap_feats_out.append(cap_padded_feat)
+
+            ### Process Image
+            C, F, H, W = image.size()
+            all_image_size.append((F, H, W))
+            F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW
+
+            image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW)
+            # "c f pf h ph w pw -> (f h w) (pf ph pw c)"
+            image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape(F_tokens * H_tokens * W_tokens, pF * pH * pW * C)
+
+            image_ori_len = len(image)
+            image_padding_len = (-image_ori_len) % SEQ_MULTI_OF
+
+            image_ori_pos_ids = self.create_coordinate_grid(
+                size=(F_tokens, H_tokens, W_tokens),
+                start=(cap_ori_len + cap_padding_len + 1, 0, 0),
+                device=device,
+            ).flatten(0, 2)
+            image_padded_pos_ids = torch.cat(
+                [
+                    image_ori_pos_ids,
+                    self.create_coordinate_grid(size=(1, 1, 1), start=(0, 0, 0), device=device)
+                    .flatten(0, 2)
+                    .repeat(image_padding_len, 1),
+                ],
+                dim=0,
+            )
+            all_image_pos_ids.append(image_padded_pos_ids if image_padding_len > 0 else image_ori_pos_ids)
+            # pad mask
+            image_pad_mask = torch.cat(
+                [
+                    torch.zeros((image_ori_len,), dtype=torch.bool, device=device),
+                    torch.ones((image_padding_len,), dtype=torch.bool, device=device),
+                ],
+                dim=0,
+            )
+            all_image_pad_mask.append(
+                image_pad_mask
+                if image_padding_len > 0
+                else torch.zeros((image_ori_len,), dtype=torch.bool, device=device)
+            )
+            # padded feature
+            image_padded_feat = torch.cat(
+                [image, image[-1:].repeat(image_padding_len, 1)],
+                dim=0,
+            )
+            all_image_out.append(image_padded_feat if image_padding_len > 0 else image)
+
+        return (
+            all_image_out,
+            all_cap_feats_out,
+            all_image_size,
+            all_image_pos_ids,
+            all_cap_pos_ids,
+            all_image_pad_mask,
+            all_cap_pad_mask,
+        )
+
+    def patchify(
+        self,
+        all_image: List[torch.Tensor],
+        patch_size: int,
+        f_patch_size: int,
+    ):
+        pH = pW = patch_size
+        pF = f_patch_size
+        all_image_out = []
+
+        for i, image in enumerate(all_image):
+            ### Process Image
+            C, F, H, W = image.size()
+            F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW
+
+            image = image.view(C, F_tokens, pF, H_tokens, pH, W_tokens, pW)
+            # "c f pf h ph w pw -> (f h w) (pf ph pw c)"
+            image = image.permute(1, 3, 5, 2, 4, 6, 0).reshape(F_tokens * H_tokens * W_tokens, pF * pH * pW * C)
+
+            image_ori_len = len(image)
+            image_padding_len = (-image_ori_len) % SEQ_MULTI_OF
+
+            # padded feature
+            image_padded_feat = torch.cat([image, image[-1:].repeat(image_padding_len, 1)], dim=0)
+            all_image_out.append(image_padded_feat)
+
+        return all_image_out
+
+    def forward(
+        self,
+        x: List[torch.Tensor],
+        t,
+        cap_feats: List[torch.Tensor],
+        patch_size=2,
+        f_patch_size=1,
+        control_context: Optional[List[torch.Tensor]] = None,
+        conditioning_scale: float = 1.0,
+        return_dict: bool = True,
+    ):
+        assert patch_size in self.all_patch_size
+        assert f_patch_size in self.all_f_patch_size
+
+        bsz = len(x)
+        device = x[0].device
+        t = t * self.t_scale
+        t = self.t_embedder(t)
+
+        (
+            x,
+            cap_feats,
+            x_size,
+            x_pos_ids,
+            cap_pos_ids,
+            x_inner_pad_mask,
+            cap_inner_pad_mask,
+        ) = self.patchify_and_embed(x, cap_feats, patch_size, f_patch_size)
+
+        # x embed & refine
+        x_item_seqlens = [len(_) for _ in x]
+        assert all(_ % SEQ_MULTI_OF == 0 for _ in x_item_seqlens)
+        x_max_item_seqlen = max(x_item_seqlens)
+
+        x = torch.cat(x, dim=0)
+        x = self.all_x_embedder[f"{patch_size}-{f_patch_size}"](x)
+
+        # Match t_embedder output dtype to x for layerwise casting compatibility
+        adaln_input = t.type_as(x)
+        x[torch.cat(x_inner_pad_mask)] = self.x_pad_token.to(x.dtype)
+        x = list(x.split(x_item_seqlens, dim=0))
+        x_freqs_cis = list(self.rope_embedder(torch.cat(x_pos_ids, dim=0)).split([len(_) for _ in x_pos_ids], dim=0))
+
+        x = pad_sequence(x, batch_first=True, padding_value=0.0)
+        x_freqs_cis = pad_sequence(x_freqs_cis, batch_first=True, padding_value=0.0)
+        # Clarify the length matches to satisfy Dynamo due to "Symbolic Shape Inference" to avoid compilation errors
+        x_freqs_cis = x_freqs_cis[:, : x.shape[1]]
+
+        x_attn_mask = torch.zeros((bsz, x_max_item_seqlen), dtype=torch.bool, device=device)
+        for i, seq_len in enumerate(x_item_seqlens):
+            x_attn_mask[i, :seq_len] = 1
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for layer in self.noise_refiner:
+                x = self._gradient_checkpointing_func(layer, x, x_attn_mask, x_freqs_cis, adaln_input)
+        else:
+            for layer in self.noise_refiner:
+                x = layer(x, x_attn_mask, x_freqs_cis, adaln_input)
+
+        # cap embed & refine
+        cap_item_seqlens = [len(_) for _ in cap_feats]
+        cap_max_item_seqlen = max(cap_item_seqlens)
+
+        cap_feats = torch.cat(cap_feats, dim=0)
+        cap_feats = self.cap_embedder(cap_feats)        
+        cap_feats[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token.to(dtype=cap_feats.dtype)
+        cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0))
+        cap_freqs_cis = list(
+            self.rope_embedder(torch.cat(cap_pos_ids, dim=0)).split([len(_) for _ in cap_pos_ids], dim=0)
+        )
+
+        cap_feats = pad_sequence(cap_feats, batch_first=True, padding_value=0.0)
+        cap_freqs_cis = pad_sequence(cap_freqs_cis, batch_first=True, padding_value=0.0)
+        # Clarify the length matches to satisfy Dynamo due to "Symbolic Shape Inference" to avoid compilation errors
+        cap_freqs_cis = cap_freqs_cis[:, : cap_feats.shape[1]]
+
+        cap_attn_mask = torch.zeros((bsz, cap_max_item_seqlen), dtype=torch.bool, device=device)
+        for i, seq_len in enumerate(cap_item_seqlens):
+            cap_attn_mask[i, :seq_len] = 1
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for layer in self.context_refiner:
+                cap_feats = self._gradient_checkpointing_func(layer, cap_feats, cap_attn_mask, cap_freqs_cis)
+        else:
+            for layer in self.context_refiner:
+                cap_feats = layer(cap_feats, cap_attn_mask, cap_freqs_cis)
+
+        # unified
+        unified = []
+        unified_freqs_cis = []
+        for i in range(bsz):
+            x_len = x_item_seqlens[i]
+            cap_len = cap_item_seqlens[i]
+            unified.append(torch.cat([x[i][:x_len], cap_feats[i][:cap_len]]))
+            unified_freqs_cis.append(torch.cat([x_freqs_cis[i][:x_len], cap_freqs_cis[i][:cap_len]]))
+        unified_item_seqlens = [a + b for a, b in zip(cap_item_seqlens, x_item_seqlens)]
+        assert unified_item_seqlens == [len(_) for _ in unified]
+        unified_max_item_seqlen = max(unified_item_seqlens)
+
+        unified = pad_sequence(unified, batch_first=True, padding_value=0.0)
+        unified_freqs_cis = pad_sequence(unified_freqs_cis, batch_first=True, padding_value=0.0)
+        unified_attn_mask = torch.zeros((bsz, unified_max_item_seqlen), dtype=torch.bool, device=device)
+        for i, seq_len in enumerate(unified_item_seqlens):
+            unified_attn_mask[i, :seq_len] = 1
+
+        ## ControlNet start
+
+        controlnet_block_samples = None
+        if control_context is not None:
+            control_context = self.patchify(control_context, patch_size, f_patch_size)
+            control_context = torch.cat(control_context, dim=0)
+            control_context = self.control_all_x_embedder[f"{patch_size}-{f_patch_size}"](control_context)
+
+            control_context[torch.cat(x_inner_pad_mask)] = self.x_pad_token
+            control_context = list(control_context.split(x_item_seqlens, dim=0))
+
+            control_context = pad_sequence(control_context, batch_first=True, padding_value=0.0)
+
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                for layer in self.control_noise_refiner:
+                    control_context = self._gradient_checkpointing_func(
+                        layer, control_context, x_attn_mask, x_freqs_cis, adaln_input
+                    )
+            else:
+                for layer in self.control_noise_refiner:
+                    control_context = layer(control_context, x_attn_mask, x_freqs_cis, adaln_input)
+
+            # unified
+            control_context_unified = []
+            for i in range(bsz):
+                x_len = x_item_seqlens[i]
+                cap_len = cap_item_seqlens[i]
+                control_context_unified.append(torch.cat([control_context[i][:x_len], cap_feats[i][:cap_len]]))
+            control_context_unified = pad_sequence(control_context_unified, batch_first=True, padding_value=0.0)
+
+            for layer in self.control_layers:
+                if torch.is_grad_enabled() and self.gradient_checkpointing:
+                    control_context_unified = self._gradient_checkpointing_func(
+                        layer, control_context_unified, unified, unified_attn_mask, unified_freqs_cis, adaln_input
+                    )
+                else:
+                    control_context_unified = layer(
+                        control_context_unified, unified, unified_attn_mask, unified_freqs_cis, adaln_input
+                    )
+
+            hints = torch.unbind(control_context_unified)[:-1]
+            controlnet_block_samples = {
+                layer_idx: hints[idx] * conditioning_scale for idx, layer_idx in enumerate(self.control_layers_places)
+            }
+
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for layer_idx, layer in enumerate(self.layers):
+                unified = self._gradient_checkpointing_func(
+                    layer, unified, unified_attn_mask, unified_freqs_cis, adaln_input
+                )
+                if controlnet_block_samples is not None:
+                    if layer_idx in controlnet_block_samples:
+                        unified = unified + controlnet_block_samples[layer_idx]
+        else:
+            for layer_idx, layer in enumerate(self.layers):
+                unified = layer(unified, unified_attn_mask, unified_freqs_cis, adaln_input)
+                if controlnet_block_samples is not None:
+                    if layer_idx in controlnet_block_samples:
+                        unified = unified + controlnet_block_samples[layer_idx]
+
+        unified = self.all_final_layer[f"{patch_size}-{f_patch_size}"](unified, adaln_input)
+        unified = list(unified.unbind(dim=0))
+        x = self.unpatchify(unified, x_size, patch_size, f_patch_size)
+
+        if not return_dict:
+            return (x,)
+
+        return Transformer2DModelOutput(sample=x)

infer_gguf.py

@@ -0,0 +1,103 @@
+import torch
+import time
+import os
+from PIL import Image
+
+# 1. Import all necessary components
+from diffusers_local.pipeline_z_image_control_unified import ZImageControlUnifiedPipeline
+from diffusers_local.z_image_control_transformer_2d import ZImageControlTransformer2DModel
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers import (
+    AutoencoderKL,
+    FlowMatchEulerDiscreteScheduler,
+    GGUFQuantizationConfig
+)
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+import diffusers_local.patch #apply patches required
+
+def main():    
+    
+     # 1. Set params ---
+    BASE_MODEL_ID = "."    
+    GGUF_FILENAME = "z_image_turbo_control_unified_q4_k_m.gguf"           
+    prompt = "a man"
+    negative_prompt = None # Set to None, as it won't be used with guidance_scale=0
+    
+    target_height, target_width = 1024, 1024
+    num_inference_steps = 9
+    guidance_scale = 0.0 #for Turbo
+    controlnet_conditioning_scale = 0.7 
+    seed = 42
+    shift = 3.0
+    generator = torch.Generator("cuda").manual_seed(seed)
+    
+    print("Loading Model Components...")
+    vae = AutoencoderKL.from_pretrained(BASE_MODEL_ID, subfolder="vae", torch_dtype=torch.bfloat16)
+    text_encoder = AutoModelForCausalLM.from_pretrained(BASE_MODEL_ID, subfolder="text_encoder", torch_dtype=torch.bfloat16)
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, subfolder="tokenizer")
+    scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=shift)
+
+    print(f"Loading Main Transformer from GGUF...")
+    transformer = ZImageControlTransformer2DModel.from_single_file(
+       os.path.join(BASE_MODEL_ID, GGUF_FILENAME),             
+       quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
+       low_cpu_mem_usage=True,
+       torch_dtype=torch.bfloat16      
+    )
+    
+    print("Loading Pipeline...")      
+    pipe = ZImageControlUnifiedPipeline(
+        vae=vae, text_encoder=text_encoder, tokenizer=tokenizer,
+        scheduler=scheduler, transformer=transformer
+    )    
+   
+    pipe.enable_group_offload(onload_device="cuda", 
+                              offload_device="cpu",
+                              offload_type="leaf_level",
+                              low_cpu_mem_usage=True,
+                              use_stream=True
+                              )
+    
+    
+    print("\nRunning Inference...")
+    input_image = Image.open("example/depth_cat.png").convert("RGB")
+        
+    prompt = "a cat"
+    negative_prompt = None # Set to None, as it won't be used with guidance_scale=0
+    
+    target_height, target_width = 1024, 1024
+    num_inference_steps = 9
+    guidance_scale = 0.0 # Correct setting for Turbo models
+    controlnet_conditioning_scale = 0.7 
+    seed = 42
+    generator = torch.Generator("cuda").manual_seed(seed)
+    
+    input_image_resized = input_image.resize((target_width, target_height), Image.Resampling.LANCZOS)
+    
+    start_inference_time = time.time()
+
+    generated_image = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        image=input_image_resized,
+        height=target_height,
+        width=target_width,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        controlnet_conditioning_scale=controlnet_conditioning_scale,
+        generator=generator,
+    ).images[0]
+    
+    end_inference_time = time.time()
+    print(f"\nGeneration finished in {end_inference_time - start_inference_time:.2f} seconds.")
+
+    # Save Output
+    if not os.path.exists("outputs"): os.makedirs("outputs")
+    output_filename = "outputs/z_image_controlnet_result_gguf.png"
+    generated_image.save(output_filename)
+    print(f"Image successfully saved as '{output_filename}'")
+    generated_image.show()
+
+if __name__ == "__main__":
+    main()

infer_pretrained.py

@@ -0,0 +1,74 @@
+import torch
+import time
+import os
+from PIL import Image
+
+# 1. Import all necessary components
+from diffusers_local.pipeline_z_image_control_unified import ZImageControlUnifiedPipeline
+from diffusers import FlowMatchEulerDiscreteScheduler
+
+import diffusers_local.patch #apply patches required
+
+def main():
+    
+    # 1. Set params ---
+    BASE_MODEL_ID = "."    
+    prompt = "a man"
+    negative_prompt = None # Set to None, as it won't be used with guidance_scale=0
+    
+    target_height, target_width = 1024, 1024
+    num_inference_steps = 9
+    guidance_scale = 0.0 #for Turbo
+    controlnet_conditioning_scale = 0.7 
+    seed = 42
+    shift = 3.0
+    generator = torch.Generator("cuda").manual_seed(seed)
+    
+    print("Loading Pipeline...")        
+    scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=shift)
+
+    pipe = ZImageControlUnifiedPipeline.from_pretrained(
+        BASE_MODEL_ID
+    )
+    pipe.scheduler = scheduler
+    
+    #Apply optimization (Optional)
+    pipe.enable_group_offload(
+        onload_device="cuda", 
+        offload_device="cpu",
+        offload_type="leaf_level",
+        low_cpu_mem_usage=True,
+        use_stream=True
+    )
+        
+    print("\nRunning Inference...")
+    
+    input_image = Image.open("example/depth_man.png").convert("RGB")
+    input_image_resized = input_image.resize((target_width, target_height), Image.Resampling.LANCZOS)    
+    start_inference_time = time.time()
+    
+    with torch.inference_mode():
+        generated_image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=input_image_resized,
+            height=target_height,
+            width=target_width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            generator=generator,
+        ).images[0]
+    
+    end_inference_time = time.time()
+    print(f"\nGeneration finished in {end_inference_time - start_inference_time:.2f} seconds.")
+
+    # Save Output
+    if not os.path.exists("outputs"): os.makedirs("outputs")
+    output_filename = "outputs/z_image_controlnet_result.png"
+    generated_image.save(output_filename)
+    print(f"Image successfully saved as '{output_filename}'")
+    generated_image.show()
+
+if __name__ == "__main__":
+    main()

model_index.json

@@ -0,0 +1,24 @@
+{
+    "_class_name": "ZImagePipeline",
+    "_diffusers_version": "0.36.0.dev0",
+    "scheduler": [
+        "diffusers",
+        "FlowMatchEulerDiscreteScheduler"
+    ],
+    "text_encoder": [
+        "transformers",
+        "Qwen3Model"
+    ],
+    "tokenizer": [
+        "transformers",
+        "Qwen2Tokenizer"
+    ],
+    "transformer": [
+        "diffusers",
+        "ZImageControlTransformer2DModel"
+    ],
+    "vae": [
+        "diffusers",
+        "AutoencoderKL"
+    ]
+}

requirements.txt

@@ -0,0 +1,18 @@
+--extra-index-url https://download.pytorch.org/whl/cu126
+
+diffusers
+torch==2.8.0+cu126
+torchvision==0.23.0+cu126
+torchaudio==2.8.0+cu126
+transformers==4.56.0
+bitsandbytes==0.48.1
+xformers==0.0.32.post2
+hf_xet
+gguf
+accelerate
+protobuf
+sacremoses
+sentencepiece
+scipy
+triton-windows<3.5; sys_platform == 'win32'
+triton==3.4.0; sys_platform != 'win32'

scheduler/scheduler_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.36.0.dev0",
+  "num_train_timesteps": 1000,
+  "use_dynamic_shifting": false,
+  "shift": 3.0
+}

text_encoder/config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "Qwen3ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2560,
+  "initializer_range": 0.02,
+  "intermediate_size": 9728,
+  "max_position_embeddings": 40960,
+  "max_window_layers": 36,
+  "model_type": "qwen3",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

text_encoder/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.51.0"
+}

text_encoder/model-00001-of-00003.safetensors

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

text_encoder/model-00002-of-00003.safetensors

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

text_encoder/model-00003-of-00003.safetensors

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

text_encoder/model.safetensors.index.json

@@ -0,0 +1,405 @@
+{
+  "metadata": {
+    "total_size": 8044936192
+  },
+  "weight_map": {
+    "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.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "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.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "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.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "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.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "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.weight": "model-00001-of-00003.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "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.weight": "model-00001-of-00003.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "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-00002-of-00003.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00002-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-00002-of-00003.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00002-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-00002-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-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00002-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.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "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.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "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.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "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.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "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.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "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.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "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.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "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.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "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.weight": "model-00002-of-00003.safetensors",
+    "model.layers.28.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.28.mlp.up_proj.weight": "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.weight": "model-00002-of-00003.safetensors",
+    "model.layers.29.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.29.mlp.up_proj.weight": "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.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "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.weight": "model-00002-of-00003.safetensors",
+    "model.layers.30.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.30.mlp.up_proj.weight": "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.weight": "model-00002-of-00003.safetensors",
+    "model.layers.31.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.31.mlp.up_proj.weight": "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.weight": "model-00002-of-00003.safetensors",
+    "model.layers.32.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.32.mlp.up_proj.weight": "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.weight": "model-00002-of-00003.safetensors",
+    "model.layers.33.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.33.mlp.up_proj.weight": "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.weight": "model-00002-of-00003.safetensors",
+    "model.layers.34.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.34.mlp.up_proj.weight": "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-00003-of-00003.safetensors",
+    "model.layers.35.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.35.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "model.layers.35.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "model.layers.35.post_attention_layernorm.weight": "model-00003-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.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "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.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "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.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "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.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "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-00003-of-00003.safetensors"
+  }
+}

tokenizer/tokenizer.json

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

tokenizer/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,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# 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>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\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\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- 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>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if message.content is string %}\n        {%- set content = message.content %}\n    {%- else %}\n        {%- set content = '' %}\n    {%- endif %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is string %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in content %}\n                {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n                {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if loop.index0 > ns.last_query_index %}\n            {%- if loop.last or (not loop.last and reasoning_content) %}\n                {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n            {%- else %}\n                {{- '<|im_start|>' + message.role + '\\n' + content }}\n            {%- endif %}\n        {%- else %}\n            {{- '<|im_start|>' + message.role + '\\n' + content }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

transformer/config.json

@@ -0,0 +1,31 @@
+{
+  "_class_name": "ZImageControlTransformer2DModel",
+  "_diffusers_version": "0.36.0.dev0",
+  "all_f_patch_size": [
+    1
+  ],
+  "all_patch_size": [
+    2
+  ],
+  "axes_dims": [
+    32,
+    48,
+    48
+  ],
+  "axes_lens": [
+    1536,
+    512,
+    512
+  ],
+  "cap_feat_dim": 2560,
+  "dim": 3840,
+  "in_channels": 16,
+  "n_heads": 30,
+  "n_kv_heads": 30,
+  "n_layers": 30,
+  "n_refiner_layers": 2,
+  "norm_eps": 1e-05,
+  "qk_norm": true,
+  "rope_theta": 256.0,
+  "t_scale": 1000.0
+}

transformer/diffusion_pytorch_model-00001-of-00004.safetensors

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

transformer/diffusion_pytorch_model-00002-of-00004.safetensors

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

transformer/diffusion_pytorch_model-00003-of-00004.safetensors

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

transformer/diffusion_pytorch_model-00004-of-00004.safetensors

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

transformer/diffusion_pytorch_model.safetensors.index.json

@@ -0,0 +1,664 @@
+{
+  "metadata": {
+    "total_size": 27721207352
+  },
+  "weight_map": {
+    "all_final_layer.2-1.adaLN_modulation.1.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "all_final_layer.2-1.adaLN_modulation.1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "all_final_layer.2-1.linear.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "all_final_layer.2-1.linear.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "all_x_embedder.2-1.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "all_x_embedder.2-1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "cap_embedder.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "cap_embedder.1.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "cap_embedder.1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "cap_pad_token": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.0.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.0.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.0.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.0.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.0.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.0.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.0.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.0.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.0.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.0.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.0.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.0.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.0.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.1.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.1.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.1.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.1.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.1.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.1.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.1.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.1.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.1.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.1.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.1.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.1.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "context_refiner.1.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.0.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.0.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.0.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.0.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.0.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.0.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.0.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.0.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.0.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.0.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.0.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.0.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.0.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.0.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.0.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.1.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.1.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.1.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.1.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.1.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.1.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.1.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.1.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.1.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.1.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.1.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.1.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.1.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.1.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.1.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.10.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.10.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.10.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.10.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.10.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.10.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.10.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.10.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.10.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.10.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.10.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.10.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.10.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.10.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.10.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.11.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.11.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.11.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.11.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.11.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.11.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.11.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.11.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.11.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.11.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.11.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.11.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.11.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.11.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.11.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.12.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.12.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.12.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.12.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.12.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.12.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.12.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.12.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.12.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.12.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.12.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.12.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.12.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.12.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.12.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.13.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.13.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.13.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.13.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.13.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.13.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.13.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.13.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.13.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.13.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.13.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.13.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.13.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.13.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.13.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.14.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.14.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.14.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.14.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.14.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.14.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.14.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.14.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.14.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.14.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.14.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.14.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.14.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.14.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.14.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.15.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.15.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.15.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.15.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.15.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.15.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.15.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.15.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.15.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.15.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.15.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.15.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.15.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.15.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.15.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.16.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.16.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.16.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.16.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.16.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.16.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.16.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.16.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.16.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.16.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.16.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.16.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.16.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.16.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.16.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.17.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.17.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.17.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.17.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.17.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.17.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.17.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.17.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.17.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.17.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.17.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.17.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.17.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.17.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.17.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.18.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.18.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.18.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.18.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.18.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.18.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.18.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.18.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.18.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.18.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.18.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.18.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.18.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.18.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.18.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.19.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.19.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.19.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.19.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.19.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.19.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.19.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.19.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.19.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.19.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.19.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.19.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.19.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.19.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.19.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.2.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.2.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.2.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.2.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.2.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.2.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.2.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.2.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.2.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.2.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.2.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.2.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.2.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.2.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.2.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.20.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.20.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.20.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.20.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.20.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.20.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.20.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.20.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.20.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.20.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.20.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.20.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.20.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.20.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.20.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.21.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.21.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.21.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.21.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.21.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.21.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.21.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.21.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.21.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.21.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.21.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.21.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.21.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.21.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.21.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.22.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.22.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.22.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.22.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.22.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.22.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.22.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.22.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.22.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.22.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.22.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.22.feed_forward.w2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.22.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.22.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.22.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.23.adaLN_modulation.0.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.23.adaLN_modulation.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.23.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.23.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.23.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.23.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.23.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.23.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.23.attention_norm1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.23.attention_norm2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.23.feed_forward.w1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.23.feed_forward.w2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.23.feed_forward.w3.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.23.ffn_norm1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.23.ffn_norm2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.24.adaLN_modulation.0.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.24.adaLN_modulation.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.24.attention.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.24.attention.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.24.attention.to_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.24.attention.to_out.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.24.attention.to_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.24.attention.to_v.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.24.attention_norm1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.24.attention_norm2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.24.feed_forward.w1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.24.feed_forward.w2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.24.feed_forward.w3.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.24.ffn_norm1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.24.ffn_norm2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.25.adaLN_modulation.0.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.25.adaLN_modulation.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.25.attention.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.25.attention.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.25.attention.to_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.25.attention.to_out.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.25.attention.to_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.25.attention.to_v.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.25.attention_norm1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.25.attention_norm2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.25.feed_forward.w1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.25.feed_forward.w2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.25.feed_forward.w3.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.25.ffn_norm1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.25.ffn_norm2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.26.adaLN_modulation.0.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.26.adaLN_modulation.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.26.attention.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.26.attention.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.26.attention.to_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.26.attention.to_out.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.26.attention.to_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.26.attention.to_v.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.26.attention_norm1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.26.attention_norm2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.26.feed_forward.w1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.26.feed_forward.w2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.26.feed_forward.w3.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.26.ffn_norm1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.26.ffn_norm2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.27.adaLN_modulation.0.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.27.adaLN_modulation.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.27.attention.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.27.attention.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.27.attention.to_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.27.attention.to_out.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.27.attention.to_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.27.attention.to_v.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.27.attention_norm1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.27.attention_norm2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.27.feed_forward.w1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.27.feed_forward.w2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.27.feed_forward.w3.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.27.ffn_norm1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.27.ffn_norm2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.28.adaLN_modulation.0.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.28.adaLN_modulation.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.28.attention.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.28.attention.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.28.attention.to_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.28.attention.to_out.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.28.attention.to_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.28.attention.to_v.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.28.attention_norm1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.28.attention_norm2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.28.feed_forward.w1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.28.feed_forward.w2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.28.feed_forward.w3.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.28.ffn_norm1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.28.ffn_norm2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.29.adaLN_modulation.0.bias": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.29.adaLN_modulation.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.29.attention.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.29.attention.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.29.attention.to_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.29.attention.to_out.0.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.29.attention.to_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.29.attention.to_v.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.29.attention_norm1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.29.attention_norm2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.29.feed_forward.w1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.29.feed_forward.w2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.29.feed_forward.w3.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.29.ffn_norm1.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.29.ffn_norm2.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "layers.3.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.3.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.3.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.3.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.3.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.3.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.3.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.3.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.3.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.3.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.3.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.3.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.3.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.3.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.3.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.4.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.4.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.4.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.4.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.4.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.4.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.4.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.4.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.4.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.4.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.4.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.4.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.4.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.4.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.4.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.5.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.5.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.5.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.5.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.5.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.5.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.5.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.5.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.5.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.5.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.5.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.5.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.5.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.5.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.5.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.6.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.6.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.6.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.6.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.6.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.6.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.6.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.6.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.6.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.6.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.6.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.6.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.6.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.6.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.6.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.7.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.7.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.7.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.7.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.7.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.7.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.7.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.7.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.7.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.7.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.7.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.7.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.7.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.7.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.7.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.8.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.8.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.8.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.8.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.8.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.8.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.8.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.8.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.8.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.8.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.8.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.8.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.8.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.8.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.8.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.9.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.9.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.9.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.9.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.9.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.9.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.9.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.9.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.9.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.9.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.9.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.9.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "layers.9.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.9.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "layers.9.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "noise_refiner.0.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.0.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.0.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.0.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.0.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.0.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.0.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.0.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.0.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.0.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.0.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.0.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.0.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.0.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.0.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.1.adaLN_modulation.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.1.adaLN_modulation.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.1.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.1.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.1.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.1.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.1.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.1.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.1.attention_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.1.attention_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.1.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.1.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.1.feed_forward.w3.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.1.ffn_norm1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "noise_refiner.1.ffn_norm2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "t_embedder.mlp.0.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "t_embedder.mlp.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "t_embedder.mlp.2.bias": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "t_embedder.mlp.2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "x_pad_token": "diffusion_pytorch_model-00001-of-00004.safetensors",
+	"control_all_x_embedder.2-1.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_all_x_embedder.2-1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.after_proj.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.after_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.before_proj.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.before_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.0.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.after_proj.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.after_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.1.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.after_proj.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.after_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.2.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.after_proj.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.after_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.3.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.after_proj.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.after_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.4.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.after_proj.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.after_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_layers.5.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.0.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.0.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.0.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.0.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.0.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.0.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.0.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.0.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.0.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.0.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.0.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.0.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.0.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.0.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.0.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.1.adaLN_modulation.0.bias": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.1.adaLN_modulation.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.1.attention.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.1.attention.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.1.attention.to_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.1.attention.to_out.0.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.1.attention.to_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.1.attention.to_v.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.1.attention_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.1.attention_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.1.feed_forward.w1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.1.feed_forward.w2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.1.feed_forward.w3.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.1.ffn_norm1.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "control_noise_refiner.1.ffn_norm2.weight": "diffusion_pytorch_model-00004-of-00004.safetensors"
+  }
+}

vae/config.json

@@ -0,0 +1,38 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.36.0.dev0",
+  "_name_or_path": "flux-dev",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 16,
+  "latents_mean": null,
+  "latents_std": null,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 1024,
+  "scaling_factor": 0.3611,
+  "shift_factor": 0.1159,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": false,
+  "use_quant_conv": false
+}

vae/diffusion_pytorch_model.safetensors

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

z_image_turbo_control_unified_q4_k_m.gguf

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