Release Kelix-AEDIT DiT weights and inference demo

.gitattributes

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

README.md

@@ -0,0 +1,117 @@
+---
+language:
+  - en
+tags:
+  - kelix
+  - image-reconstruction
+  - image-tokenizer
+  - diffusion
+  - flow-matching
+  - safetensors
+---
+
+# Kelix-AEDIT
+
+[Paper](https://arxiv.org/abs/2602.09843) | [GitHub](https://github.com/Kuaishou-OneRec/Kelix) | [Kelix-DiT](https://huggingface.co/OpenOneRec/Kelix-DiT)
+
+## Introduction
+
+Kelix-AEDIT reconstructs images from discrete visual tokens using a customized SANA diffusion transformer. This repository releases the DiT weights, an inference script, and its dependencies.
+
+The inference pipeline reuses the image tokenizer from [Kelix-SFT](https://huggingface.co/OpenOneRec/Kelix-SFT) and the [DC-AE VAE from SANA1.5](https://huggingface.co/Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers/tree/main/vae). These components are loaded separately and are not duplicated in this repository.
+
+Unlike Kelix-DiT, which takes Kelix-LLM hidden states as its condition, Kelix-AEDIT operates on quantized image features and supports image reconstruction without the LLM.
+
+![An example image reconstructed with Kelix-AEDIT.](example.png)
+
+*An example reconstruction produced by Kelix-AEDIT.*
+
+## Model Overview
+
+| Item | Value |
+|---|---|
+| Released weights | Customized SANA-DiT |
+| Image tokenizer | Reused from Kelix-SFT |
+| Latent VAE | DC-AE from SANA1.5 |
+| Condition | Quantized image features with mean fusion |
+| Codebooks | 8 × 8,192 entries |
+| Default condition grid | 18 × 18 |
+| Default output resolution | 512 × 512 |
+| Sampling | 20-step Euler |
+| Guidance scale | 4.5 |
+| Flow shift | 3.0 |
+| Inference precision | BF16 |
+
+## Usage
+
+Requires an NVIDIA CUDA GPU and a FlashAttention build compatible with your PyTorch and CUDA environment. The inference script uses the model implementations in the Kelix code repository.
+
+### Installation
+
+```bash
+git clone https://github.com/Kuaishou-OneRec/Kelix.git
+cd Kelix
+pip install huggingface-hub==0.36.0
+hf download OpenOneRec/Kelix-AEDIT --local-dir checkpoints/Kelix-AEDIT
+pip install -r checkpoints/Kelix-AEDIT/requirements.txt
+pip install flash-attn==2.7.4.post1 --no-build-isolation
+```
+
+### Image Reconstruction
+
+Run from the Kelix repository root. The tokenizer and VAE are downloaded automatically on first use. Only the Kelix-SFT shards containing the visual tokenizer are needed; the LLM is not loaded.
+
+```bash
+PYTHONPATH=. python checkpoints/Kelix-AEDIT/infer.py \
+  --ckpt checkpoints/Kelix-AEDIT \
+  --image input.jpg \
+  --output reconstruction.png
+```
+
+To use existing local checkpoints:
+
+```bash
+PYTHONPATH=. python checkpoints/Kelix-AEDIT/infer.py \
+  --ckpt checkpoints/Kelix-AEDIT \
+  --tokenizer /path/to/Kelix-SFT \
+  --vae /path/to/vae \
+  --image input.jpg \
+  --output reconstruction.png
+```
+
+### Decode Image Tokens
+
+```bash
+PYTHONPATH=. python checkpoints/Kelix-AEDIT/infer.py \
+  --ckpt checkpoints/Kelix-AEDIT \
+  --tokens tokens.pt \
+  --grid 18 \
+  --output generated.png
+```
+
+`tokens.pt` must contain an integer tensor saved with `torch.save`:
+
+- Shape: `[324, 8]` for an 18 × 18 grid.
+- Values: local indices in `[0, 8192)` for each codebook.
+- Spatial order: row-major.
+- Dtype: `torch.int32` or `torch.int64`.
+
+Tokens must use the matching image tokenizer and codebooks. Raw autoregressive vocabulary IDs must be converted to local codebook indices before decoding.
+
+This model does not accept text prompts directly. Text-to-image generation requires a compatible autoregressive model to generate image tokens first.
+
+## Citation
+
+```bibtex
+@article{kelix2026,
+  title   = {Kelix Technique Report: Closing the Understanding Gap of Discrete Tokens in Unified Multimodal Models},
+  author  = {Kuaishou Technology},
+  journal = {arXiv preprint arXiv:2602.09843},
+  year    = {2026},
+  url     = {https://arxiv.org/abs/2602.09843}
+}
+```
+
+## License
+
+Please contact the OneRec Team for the license of the Kelix series. The base SANA-DiT, DC-AE, and Keye vision components are subject to their respective original licenses.

config.json

@@ -0,0 +1,42 @@
+{
+  "model_class": "SanaModel",
+  "input_size": 32,
+  "patch_size": 1,
+  "in_channels": 32,
+  "hidden_size": 2240,
+  "depth": 20,
+  "num_heads": 20,
+  "mlp_ratio": 2.5,
+  "caption_channels": 1024,
+  "model_max_length": 2560,
+  "attn_type": "linear",
+  "cross_attn_type": "flash",
+  "linear_head_dim": 32,
+  "qk_norm": true,
+  "cross_norm": true,
+  "use_cross_attn_rope": true,
+  "use_position_scale": false,
+  "cross_attn_x_norm": true,
+  "ffn_type": "glumbconv",
+  "mlp_acts": [
+    "silu",
+    "silu",
+    null
+  ],
+  "pred_sigma": false,
+  "learn_sigma": false,
+  "use_pe": false,
+  "pe_interpolation": 1.0,
+  "y_norm": true,
+  "y_norm_scale_factor": 1.0,
+  "norm_eps": 1e-06,
+  "class_dropout_prob": 0.1,
+  "drop_path": 0.0,
+  "vae_type": "AutoencoderDC",
+  "vae_pretrained": "mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers",
+  "vae_downsample_rate": 32,
+  "text_encoder_name": "google/gemma-2-2b-it",
+  "y_embedding_init_method": "randn",
+  "use_connector": false,
+  "__class__": "SanaConfig"
+}

infer.py

@@ -0,0 +1,199 @@
+"""Single-GPU Kelix-AEDIT inference: an image or local VQ indices -> PNG."""
+import argparse
+import json
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent
+# Support running the demo or a downloaded HF script inside a Kelix checkout.
+for parent in (ROOT, *ROOT.parents):
+    if (parent / 'muse' / '__init__.py').is_file():
+        sys.path.insert(0, str(parent))
+        break
+
+
+def parse_args():
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument('--ckpt', default=str(ROOT) if (ROOT / 'config.json').is_file() else 'OpenOneRec/Kelix-AEDIT', help='Local DiT directory or HF repo ID')
+    p.add_argument('--tokenizer', default='OpenOneRec/Kelix-SFT', help='Kelix-SFT repo/directory or a standalone image tokenizer directory')
+    p.add_argument('--vae', default='Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers/vae', help='Local DC-AE directory or HF repo/subfolder')
+    inputs = p.add_mutually_exclusive_group(required=True)
+    inputs.add_argument('--image', help='Image to encode and reconstruct')
+    inputs.add_argument('--tokens', help='Safely loaded .pt LongTensor [H*W,8], local codebook indices')
+    p.add_argument('--grid', type=int, default=18, help='Condition grid width/height; 18 means 324 tokens')
+    p.add_argument('--output', default='output.png')
+    p.add_argument('--image-size', type=int, default=512)
+    p.add_argument('--steps', type=int, default=20)
+    p.add_argument('--cfg', type=float, default=4.5)
+    p.add_argument('--flow-shift', type=float, default=3.0)
+    p.add_argument('--seed', type=int, default=42)
+    p.add_argument('--device', default='cuda:0')
+    p.add_argument('--dtype', choices=['bfloat16', 'float16', 'float32'], default='bfloat16')
+    return p.parse_args()
+
+
+def positions(h, w, device):
+    import torch
+    ids = torch.arange(h * w, device=device)
+    return {'height': ids // w, 'width': ids % w}
+
+
+def resolve_directory(source, allow_patterns=None):
+    path = Path(source)
+    if path.is_dir():
+        return path
+    from huggingface_hub import snapshot_download
+    parts = str(source).split('/')
+    repo_id = '/'.join(parts[:2])
+    subfolder = '/'.join(parts[2:])
+    patterns = [f'{subfolder}/*'] if subfolder else allow_patterns
+    root = Path(snapshot_download(repo_id, allow_patterns=patterns))
+    return root / subfolder if subfolder else root
+
+
+def load_image_tokenizer(source, device, dtype):
+    """Load the visual tokenizer without instantiating or downloading the LLM."""
+    import torch
+    from accelerate import init_empty_weights
+    from safetensors import safe_open
+    from muse.config.model_config import KeyeTokenizerConfig
+    from muse.models.keye_tokenizer import KeyeImageTokenizer
+
+    folder = resolve_directory(source, ['*.json', '*.txt', '*.py', '*.jinja'])
+    config = json.loads((folder / 'config.json').read_text())
+    if 'tokenizer_config' not in config:
+        tokenizer = KeyeImageTokenizer.from_pretrained(
+            str(folder), strict=True, fusion_type='mean'
+        ).to(device=device, dtype=dtype).eval()
+        return tokenizer, folder
+
+    tokenizer_config = KeyeTokenizerConfig(**config['tokenizer_config'])
+    tokenizer_config.fusion_type = 'mean'
+    with init_empty_weights():
+        tokenizer = KeyeImageTokenizer(tokenizer_config)
+    index = json.loads((folder / 'model.safetensors.index.json').read_text())
+    prefix = 'visual_tokenizer.'
+    selected = {k: v for k, v in index['weight_map'].items() if k.startswith(prefix)}
+    if not selected:
+        raise ValueError('No visual_tokenizer weights found in the checkpoint index')
+    state = {}
+    for shard in sorted(set(selected.values())):
+        path = folder / shard
+        if not path.is_file():
+            if Path(source).is_dir():
+                raise FileNotFoundError(f'Missing tokenizer shard: {shard}')
+            from huggingface_hub import hf_hub_download
+            path = Path(hf_hub_download(source, shard, revision=folder.name))
+        with safe_open(str(path), framework='pt', device='cpu') as weights:
+            for key, filename in selected.items():
+                if filename == shard:
+                    state[key[len(prefix):]] = weights.get_tensor(key)
+    tokenizer.load_state_dict(state, strict=True, assign=True)
+    tokenizer = tokenizer.to(device=device, dtype=dtype).eval()
+    return tokenizer, folder
+
+
+def main(args):
+    import torch
+    import torch.nn.functional as F
+    from PIL import Image
+    from diffusers import AutoencoderDC, FlowMatchEulerDiscreteScheduler
+    from transformers import AutoProcessor
+    from muse.models.base import Model
+    from muse.training.common import set_default_dtype
+
+    if args.grid < 1 or args.steps < 1 or args.image_size < 32 or args.image_size % 32:
+        raise ValueError('grid and steps must be positive; image-size must be a positive multiple of 32')
+    ckpt = resolve_directory(args.ckpt, ['config.json', '*.safetensors', '*.index.json'])
+    dit_dir = ckpt / 'dit' if (ckpt / 'dit').is_dir() else ckpt
+    device, dtype = torch.device(args.device), getattr(torch, args.dtype)
+    if device.type == 'cuda' and not torch.cuda.is_available():
+        raise RuntimeError('CUDA is required for the FlashAttention vision encoder')
+    torch.manual_seed(args.seed)
+    with torch.inference_mode(), set_default_dtype(dtype):
+        print('Loading DiT (strict=True)...', flush=True)
+        model = Model.from_pretrained(str(dit_dir), strict=True).to(device=device, dtype=dtype).eval()
+        print('Loading image tokenizer (strict=True)...', flush=True)
+        tokenizer, tokenizer_dir = load_image_tokenizer(args.tokenizer, device, dtype)
+        vae_dir = resolve_directory(args.vae, ['config.json', '*.safetensors'])
+        vae = AutoencoderDC.from_pretrained(str(vae_dir), torch_dtype=dtype).to(device).eval()
+        length = args.grid ** 2
+        if args.tokens:
+            indices = torch.load(args.tokens, map_location='cpu', weights_only=True)
+            if not isinstance(indices, torch.Tensor) or tuple(indices.shape) != (length, tokenizer.n_q_tokens):
+                raise ValueError(f'Expected integer tensor [{length}, {tokenizer.n_q_tokens}]')
+            if indices.dtype not in (torch.int32, torch.int64):
+                raise ValueError('Token indices must have integer dtype')
+            limit = tokenizer.config.codebook_size // tokenizer.config.split_voc
+            if indices.min() < 0 or indices.max() >= limit:
+                raise ValueError(f'Expected per-codebook indices in [0, {limit}) (not AR vocabulary IDs)')
+            cond, mask, seq_len = tokenizer.tokenize_indices(
+                indices.to(device=device, dtype=torch.long), max_pad_to=length,
+                lengths=[length], return_attention_mask=True,
+            )
+        else:
+            processor = AutoProcessor.from_pretrained(
+                str(tokenizer_dir), trust_remote_code=True, use_fast=False,
+            )
+            # Match the original inference's square resize, then use an exact grid.
+            # A merged Keye token covers 2*14 pixels; 18*28=504 for 324 tokens.
+            image = Image.open(args.image).convert('RGB').resize(
+                (args.image_size, args.image_size), Image.Resampling.LANCZOS,
+            ).resize((args.grid * 28, args.grid * 28), Image.Resampling.BICUBIC)
+            inputs = processor(text=[''], images=[image], return_tensors='pt')
+            actual_grid = inputs['image_grid_thw'][0, 1:] // 2
+            if actual_grid.tolist() != [args.grid, args.grid]:
+                raise ValueError(f'Unexpected processor grid: {actual_grid.tolist()}')
+            encoded = tokenizer(
+                pixel_values=inputs['pixel_values'].to(device=device, dtype=dtype),
+                image_grid_thw=inputs['image_grid_thw'].to(device),
+            )
+            features = encoded.get('token_embeds')
+            if features is None:
+                # Older Kelix releases compute fused features without returning them.
+                features = torch.stack([
+                    projector(z_q) for projector, z_q in
+                    zip(tokenizer.up_projectors, encoded['z_q'])
+                ], dim=1).mean(dim=1)
+            if features.shape[0] != length:
+                raise ValueError(f'Expected {length} image tokens, got {features.shape[0]}')
+            cond, seq_len = features[None], length
+            mask = torch.ones((1, 1, 1, length), device=device)
+        if cond.shape[-1] != model.config.caption_channels:
+            raise ValueError('Tokenizer condition dimension does not match the DiT config')
+        null = model.y_embedder.y_embedding[:seq_len].to(device=device, dtype=dtype)
+        valid = null.shape[0]
+        null = F.pad(null, (0, 0, 0, seq_len - valid))[None]
+        null_mask = torch.zeros_like(mask)
+        null_mask[..., :valid] = 1
+        cond_cfg = torch.cat([null, cond])
+        mask_cfg = torch.cat([null_mask, mask])
+        latent_size = args.image_size // 32
+        generator = torch.Generator(device=device).manual_seed(args.seed)
+        latents = torch.randn((1, model.config.in_channels, latent_size, latent_size),
+                              device=device, dtype=dtype, generator=generator)
+        scheduler = FlowMatchEulerDiscreteScheduler(shift=args.flow_shift)
+        scheduler.set_timesteps(args.steps, device=device)
+        x_pos = positions(latent_size, latent_size, device)
+        y_pos = positions(args.grid, args.grid, device)
+        print(f'Sampling {args.steps} steps...', flush=True)
+        for t in scheduler.timesteps:
+            pred = model.forward_with_dpmsolver(
+                torch.cat([latents, latents]), t.expand(2), cond_cfg, mask=mask_cfg,
+                x_input_pos=x_pos, cond_input_pos=y_pos,
+            )
+            uncond, conditional = pred.chunk(2)
+            pred = uncond + args.cfg * (conditional - uncond)
+            latents = scheduler.step(pred, t, latents, return_dict=False)[0]
+        image = (vae.decode(latents / vae.config.scaling_factor).sample[0] / 2 + 0.5).clamp(0, 1)
+        if not torch.isfinite(image).all():
+            raise RuntimeError('Non-finite pixels encountered')
+        pixels = (image.float().permute(1, 2, 0).cpu().numpy() * 255).round().astype('uint8')
+        output = Path(args.output)
+        output.parent.mkdir(parents=True, exist_ok=True)
+        Image.fromarray(pixels).save(output)
+        print(f'Saved {output.resolve()}')
+
+
+if __name__ == '__main__':
+    main(parse_args())

model-00000-of-00003.safetensors

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

model-00001-of-00003.safetensors

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

model-00002-of-00003.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,485 @@
+{
+  "metadata": {
+    "total_size": 6426759808
+  },
+  "weight_map": {
+    "pos_embed": "model-00000-of-00003.safetensors",
+    "x_embedder.proj.weight": "model-00000-of-00003.safetensors",
+    "x_embedder.proj.bias": "model-00000-of-00003.safetensors",
+    "t_embedder.mlp.0.weight": "model-00000-of-00003.safetensors",
+    "t_embedder.mlp.0.bias": "model-00000-of-00003.safetensors",
+    "t_embedder.mlp.2.weight": "model-00000-of-00003.safetensors",
+    "t_embedder.mlp.2.bias": "model-00000-of-00003.safetensors",
+    "t_block.1.weight": "model-00000-of-00003.safetensors",
+    "t_block.1.bias": "model-00000-of-00003.safetensors",
+    "y_embedder.y_proj.fc1.weight": "model-00000-of-00003.safetensors",
+    "y_embedder.y_proj.fc1.bias": "model-00000-of-00003.safetensors",
+    "y_embedder.y_proj.fc2.weight": "model-00000-of-00003.safetensors",
+    "y_embedder.y_proj.fc2.bias": "model-00000-of-00003.safetensors",
+    "attention_y_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.0.scale_shift_table": "model-00000-of-00003.safetensors",
+    "blocks.0.attn.to_q.weight": "model-00000-of-00003.safetensors",
+    "blocks.0.attn.to_k.weight": "model-00000-of-00003.safetensors",
+    "blocks.0.attn.to_v.weight": "model-00000-of-00003.safetensors",
+    "blocks.0.attn.proj.weight": "model-00000-of-00003.safetensors",
+    "blocks.0.attn.proj.bias": "model-00000-of-00003.safetensors",
+    "blocks.0.attn.q_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.0.attn.k_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.0.cross_attn.q_linear.weight": "model-00000-of-00003.safetensors",
+    "blocks.0.cross_attn.q_linear.bias": "model-00000-of-00003.safetensors",
+    "blocks.0.cross_attn.to_k.weight": "model-00000-of-00003.safetensors",
+    "blocks.0.cross_attn.to_k.bias": "model-00000-of-00003.safetensors",
+    "blocks.0.cross_attn.to_v.weight": "model-00000-of-00003.safetensors",
+    "blocks.0.cross_attn.to_v.bias": "model-00000-of-00003.safetensors",
+    "blocks.0.cross_attn.proj.weight": "model-00000-of-00003.safetensors",
+    "blocks.0.cross_attn.proj.bias": "model-00000-of-00003.safetensors",
+    "blocks.0.cross_attn.q_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.0.cross_attn.k_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.0.mlp.inverted_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.0.mlp.inverted_conv.conv.bias": "model-00000-of-00003.safetensors",
+    "blocks.0.mlp.depth_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.0.mlp.depth_conv.conv.bias": "model-00000-of-00003.safetensors",
+    "blocks.0.mlp.point_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.1.scale_shift_table": "model-00000-of-00003.safetensors",
+    "blocks.1.attn.to_q.weight": "model-00000-of-00003.safetensors",
+    "blocks.1.attn.to_k.weight": "model-00000-of-00003.safetensors",
+    "blocks.1.attn.to_v.weight": "model-00000-of-00003.safetensors",
+    "blocks.1.attn.proj.weight": "model-00000-of-00003.safetensors",
+    "blocks.1.attn.proj.bias": "model-00000-of-00003.safetensors",
+    "blocks.1.attn.q_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.1.attn.k_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.1.cross_attn.q_linear.weight": "model-00000-of-00003.safetensors",
+    "blocks.1.cross_attn.q_linear.bias": "model-00000-of-00003.safetensors",
+    "blocks.1.cross_attn.to_k.weight": "model-00000-of-00003.safetensors",
+    "blocks.1.cross_attn.to_k.bias": "model-00000-of-00003.safetensors",
+    "blocks.1.cross_attn.to_v.weight": "model-00000-of-00003.safetensors",
+    "blocks.1.cross_attn.to_v.bias": "model-00000-of-00003.safetensors",
+    "blocks.1.cross_attn.proj.weight": "model-00000-of-00003.safetensors",
+    "blocks.1.cross_attn.proj.bias": "model-00000-of-00003.safetensors",
+    "blocks.1.cross_attn.q_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.1.cross_attn.k_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.1.mlp.inverted_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.1.mlp.inverted_conv.conv.bias": "model-00000-of-00003.safetensors",
+    "blocks.1.mlp.depth_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.1.mlp.depth_conv.conv.bias": "model-00000-of-00003.safetensors",
+    "blocks.1.mlp.point_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.2.scale_shift_table": "model-00000-of-00003.safetensors",
+    "blocks.2.attn.to_q.weight": "model-00000-of-00003.safetensors",
+    "blocks.2.attn.to_k.weight": "model-00000-of-00003.safetensors",
+    "blocks.2.attn.to_v.weight": "model-00000-of-00003.safetensors",
+    "blocks.2.attn.proj.weight": "model-00000-of-00003.safetensors",
+    "blocks.2.attn.proj.bias": "model-00000-of-00003.safetensors",
+    "blocks.2.attn.q_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.2.attn.k_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.2.cross_attn.q_linear.weight": "model-00000-of-00003.safetensors",
+    "blocks.2.cross_attn.q_linear.bias": "model-00000-of-00003.safetensors",
+    "blocks.2.cross_attn.to_k.weight": "model-00000-of-00003.safetensors",
+    "blocks.2.cross_attn.to_k.bias": "model-00000-of-00003.safetensors",
+    "blocks.2.cross_attn.to_v.weight": "model-00000-of-00003.safetensors",
+    "blocks.2.cross_attn.to_v.bias": "model-00000-of-00003.safetensors",
+    "blocks.2.cross_attn.proj.weight": "model-00000-of-00003.safetensors",
+    "blocks.2.cross_attn.proj.bias": "model-00000-of-00003.safetensors",
+    "blocks.2.cross_attn.q_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.2.cross_attn.k_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.2.mlp.inverted_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.2.mlp.inverted_conv.conv.bias": "model-00000-of-00003.safetensors",
+    "blocks.2.mlp.depth_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.2.mlp.depth_conv.conv.bias": "model-00000-of-00003.safetensors",
+    "blocks.2.mlp.point_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.3.scale_shift_table": "model-00000-of-00003.safetensors",
+    "blocks.3.attn.to_q.weight": "model-00000-of-00003.safetensors",
+    "blocks.3.attn.to_k.weight": "model-00000-of-00003.safetensors",
+    "blocks.3.attn.to_v.weight": "model-00000-of-00003.safetensors",
+    "blocks.3.attn.proj.weight": "model-00000-of-00003.safetensors",
+    "blocks.3.attn.proj.bias": "model-00000-of-00003.safetensors",
+    "blocks.3.attn.q_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.3.attn.k_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.3.cross_attn.q_linear.weight": "model-00000-of-00003.safetensors",
+    "blocks.3.cross_attn.q_linear.bias": "model-00000-of-00003.safetensors",
+    "blocks.3.cross_attn.to_k.weight": "model-00000-of-00003.safetensors",
+    "blocks.3.cross_attn.to_k.bias": "model-00000-of-00003.safetensors",
+    "blocks.3.cross_attn.to_v.weight": "model-00000-of-00003.safetensors",
+    "blocks.3.cross_attn.to_v.bias": "model-00000-of-00003.safetensors",
+    "blocks.3.cross_attn.proj.weight": "model-00000-of-00003.safetensors",
+    "blocks.3.cross_attn.proj.bias": "model-00000-of-00003.safetensors",
+    "blocks.3.cross_attn.q_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.3.cross_attn.k_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.3.mlp.inverted_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.3.mlp.inverted_conv.conv.bias": "model-00000-of-00003.safetensors",
+    "blocks.3.mlp.depth_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.3.mlp.depth_conv.conv.bias": "model-00000-of-00003.safetensors",
+    "blocks.3.mlp.point_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.4.scale_shift_table": "model-00000-of-00003.safetensors",
+    "blocks.4.attn.to_q.weight": "model-00000-of-00003.safetensors",
+    "blocks.4.attn.to_k.weight": "model-00000-of-00003.safetensors",
+    "blocks.4.attn.to_v.weight": "model-00000-of-00003.safetensors",
+    "blocks.4.attn.proj.weight": "model-00000-of-00003.safetensors",
+    "blocks.4.attn.proj.bias": "model-00000-of-00003.safetensors",
+    "blocks.4.attn.q_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.4.attn.k_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.4.cross_attn.q_linear.weight": "model-00000-of-00003.safetensors",
+    "blocks.4.cross_attn.q_linear.bias": "model-00000-of-00003.safetensors",
+    "blocks.4.cross_attn.to_k.weight": "model-00000-of-00003.safetensors",
+    "blocks.4.cross_attn.to_k.bias": "model-00000-of-00003.safetensors",
+    "blocks.4.cross_attn.to_v.weight": "model-00000-of-00003.safetensors",
+    "blocks.4.cross_attn.to_v.bias": "model-00000-of-00003.safetensors",
+    "blocks.4.cross_attn.proj.weight": "model-00000-of-00003.safetensors",
+    "blocks.4.cross_attn.proj.bias": "model-00000-of-00003.safetensors",
+    "blocks.4.cross_attn.q_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.4.cross_attn.k_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.4.mlp.inverted_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.4.mlp.inverted_conv.conv.bias": "model-00000-of-00003.safetensors",
+    "blocks.4.mlp.depth_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.4.mlp.depth_conv.conv.bias": "model-00000-of-00003.safetensors",
+    "blocks.4.mlp.point_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.5.scale_shift_table": "model-00000-of-00003.safetensors",
+    "blocks.5.attn.to_q.weight": "model-00000-of-00003.safetensors",
+    "blocks.5.attn.to_k.weight": "model-00000-of-00003.safetensors",
+    "blocks.5.attn.to_v.weight": "model-00000-of-00003.safetensors",
+    "blocks.5.attn.proj.weight": "model-00000-of-00003.safetensors",
+    "blocks.5.attn.proj.bias": "model-00000-of-00003.safetensors",
+    "blocks.5.attn.q_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.5.attn.k_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.5.cross_attn.q_linear.weight": "model-00000-of-00003.safetensors",
+    "blocks.5.cross_attn.q_linear.bias": "model-00000-of-00003.safetensors",
+    "blocks.5.cross_attn.to_k.weight": "model-00000-of-00003.safetensors",
+    "blocks.5.cross_attn.to_k.bias": "model-00000-of-00003.safetensors",
+    "blocks.5.cross_attn.to_v.weight": "model-00000-of-00003.safetensors",
+    "blocks.5.cross_attn.to_v.bias": "model-00000-of-00003.safetensors",
+    "blocks.5.cross_attn.proj.weight": "model-00000-of-00003.safetensors",
+    "blocks.5.cross_attn.proj.bias": "model-00000-of-00003.safetensors",
+    "blocks.5.cross_attn.q_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.5.cross_attn.k_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.5.mlp.inverted_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.5.mlp.inverted_conv.conv.bias": "model-00000-of-00003.safetensors",
+    "blocks.5.mlp.depth_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.5.mlp.depth_conv.conv.bias": "model-00000-of-00003.safetensors",
+    "blocks.5.mlp.point_conv.conv.weight": "model-00000-of-00003.safetensors",
+    "blocks.6.scale_shift_table": "model-00000-of-00003.safetensors",
+    "blocks.6.attn.to_q.weight": "model-00000-of-00003.safetensors",
+    "blocks.6.attn.to_k.weight": "model-00000-of-00003.safetensors",
+    "blocks.6.attn.to_v.weight": "model-00000-of-00003.safetensors",
+    "blocks.6.attn.proj.weight": "model-00000-of-00003.safetensors",
+    "blocks.6.attn.proj.bias": "model-00000-of-00003.safetensors",
+    "blocks.6.attn.q_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.6.attn.k_norm.weight": "model-00000-of-00003.safetensors",
+    "blocks.6.cross_attn.q_linear.weight": "model-00000-of-00003.safetensors",
+    "blocks.6.cross_attn.q_linear.bias": "model-00001-of-00003.safetensors",
+    "blocks.6.cross_attn.to_k.weight": "model-00001-of-00003.safetensors",
+    "blocks.6.cross_attn.to_k.bias": "model-00001-of-00003.safetensors",
+    "blocks.6.cross_attn.to_v.weight": "model-00001-of-00003.safetensors",
+    "blocks.6.cross_attn.to_v.bias": "model-00001-of-00003.safetensors",
+    "blocks.6.cross_attn.proj.weight": "model-00001-of-00003.safetensors",
+    "blocks.6.cross_attn.proj.bias": "model-00001-of-00003.safetensors",
+    "blocks.6.cross_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.6.cross_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.6.mlp.inverted_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.6.mlp.inverted_conv.conv.bias": "model-00001-of-00003.safetensors",
+    "blocks.6.mlp.depth_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.6.mlp.depth_conv.conv.bias": "model-00001-of-00003.safetensors",
+    "blocks.6.mlp.point_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.7.scale_shift_table": "model-00001-of-00003.safetensors",
+    "blocks.7.attn.to_q.weight": "model-00001-of-00003.safetensors",
+    "blocks.7.attn.to_k.weight": "model-00001-of-00003.safetensors",
+    "blocks.7.attn.to_v.weight": "model-00001-of-00003.safetensors",
+    "blocks.7.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "blocks.7.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "blocks.7.attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.7.attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.7.cross_attn.q_linear.weight": "model-00001-of-00003.safetensors",
+    "blocks.7.cross_attn.q_linear.bias": "model-00001-of-00003.safetensors",
+    "blocks.7.cross_attn.to_k.weight": "model-00001-of-00003.safetensors",
+    "blocks.7.cross_attn.to_k.bias": "model-00001-of-00003.safetensors",
+    "blocks.7.cross_attn.to_v.weight": "model-00001-of-00003.safetensors",
+    "blocks.7.cross_attn.to_v.bias": "model-00001-of-00003.safetensors",
+    "blocks.7.cross_attn.proj.weight": "model-00001-of-00003.safetensors",
+    "blocks.7.cross_attn.proj.bias": "model-00001-of-00003.safetensors",
+    "blocks.7.cross_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.7.cross_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.7.mlp.inverted_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.7.mlp.inverted_conv.conv.bias": "model-00001-of-00003.safetensors",
+    "blocks.7.mlp.depth_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.7.mlp.depth_conv.conv.bias": "model-00001-of-00003.safetensors",
+    "blocks.7.mlp.point_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.8.scale_shift_table": "model-00001-of-00003.safetensors",
+    "blocks.8.attn.to_q.weight": "model-00001-of-00003.safetensors",
+    "blocks.8.attn.to_k.weight": "model-00001-of-00003.safetensors",
+    "blocks.8.attn.to_v.weight": "model-00001-of-00003.safetensors",
+    "blocks.8.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "blocks.8.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "blocks.8.attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.8.attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.8.cross_attn.q_linear.weight": "model-00001-of-00003.safetensors",
+    "blocks.8.cross_attn.q_linear.bias": "model-00001-of-00003.safetensors",
+    "blocks.8.cross_attn.to_k.weight": "model-00001-of-00003.safetensors",
+    "blocks.8.cross_attn.to_k.bias": "model-00001-of-00003.safetensors",
+    "blocks.8.cross_attn.to_v.weight": "model-00001-of-00003.safetensors",
+    "blocks.8.cross_attn.to_v.bias": "model-00001-of-00003.safetensors",
+    "blocks.8.cross_attn.proj.weight": "model-00001-of-00003.safetensors",
+    "blocks.8.cross_attn.proj.bias": "model-00001-of-00003.safetensors",
+    "blocks.8.cross_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.8.cross_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.8.mlp.inverted_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.8.mlp.inverted_conv.conv.bias": "model-00001-of-00003.safetensors",
+    "blocks.8.mlp.depth_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.8.mlp.depth_conv.conv.bias": "model-00001-of-00003.safetensors",
+    "blocks.8.mlp.point_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.9.scale_shift_table": "model-00001-of-00003.safetensors",
+    "blocks.9.attn.to_q.weight": "model-00001-of-00003.safetensors",
+    "blocks.9.attn.to_k.weight": "model-00001-of-00003.safetensors",
+    "blocks.9.attn.to_v.weight": "model-00001-of-00003.safetensors",
+    "blocks.9.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "blocks.9.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "blocks.9.attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.9.attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.9.cross_attn.q_linear.weight": "model-00001-of-00003.safetensors",
+    "blocks.9.cross_attn.q_linear.bias": "model-00001-of-00003.safetensors",
+    "blocks.9.cross_attn.to_k.weight": "model-00001-of-00003.safetensors",
+    "blocks.9.cross_attn.to_k.bias": "model-00001-of-00003.safetensors",
+    "blocks.9.cross_attn.to_v.weight": "model-00001-of-00003.safetensors",
+    "blocks.9.cross_attn.to_v.bias": "model-00001-of-00003.safetensors",
+    "blocks.9.cross_attn.proj.weight": "model-00001-of-00003.safetensors",
+    "blocks.9.cross_attn.proj.bias": "model-00001-of-00003.safetensors",
+    "blocks.9.cross_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.9.cross_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.9.mlp.inverted_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.9.mlp.inverted_conv.conv.bias": "model-00001-of-00003.safetensors",
+    "blocks.9.mlp.depth_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.9.mlp.depth_conv.conv.bias": "model-00001-of-00003.safetensors",
+    "blocks.9.mlp.point_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.10.scale_shift_table": "model-00001-of-00003.safetensors",
+    "blocks.10.attn.to_q.weight": "model-00001-of-00003.safetensors",
+    "blocks.10.attn.to_k.weight": "model-00001-of-00003.safetensors",
+    "blocks.10.attn.to_v.weight": "model-00001-of-00003.safetensors",
+    "blocks.10.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "blocks.10.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "blocks.10.attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.10.attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.10.cross_attn.q_linear.weight": "model-00001-of-00003.safetensors",
+    "blocks.10.cross_attn.q_linear.bias": "model-00001-of-00003.safetensors",
+    "blocks.10.cross_attn.to_k.weight": "model-00001-of-00003.safetensors",
+    "blocks.10.cross_attn.to_k.bias": "model-00001-of-00003.safetensors",
+    "blocks.10.cross_attn.to_v.weight": "model-00001-of-00003.safetensors",
+    "blocks.10.cross_attn.to_v.bias": "model-00001-of-00003.safetensors",
+    "blocks.10.cross_attn.proj.weight": "model-00001-of-00003.safetensors",
+    "blocks.10.cross_attn.proj.bias": "model-00001-of-00003.safetensors",
+    "blocks.10.cross_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.10.cross_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.10.mlp.inverted_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.10.mlp.inverted_conv.conv.bias": "model-00001-of-00003.safetensors",
+    "blocks.10.mlp.depth_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.10.mlp.depth_conv.conv.bias": "model-00001-of-00003.safetensors",
+    "blocks.10.mlp.point_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.11.scale_shift_table": "model-00001-of-00003.safetensors",
+    "blocks.11.attn.to_q.weight": "model-00001-of-00003.safetensors",
+    "blocks.11.attn.to_k.weight": "model-00001-of-00003.safetensors",
+    "blocks.11.attn.to_v.weight": "model-00001-of-00003.safetensors",
+    "blocks.11.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "blocks.11.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "blocks.11.attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.11.attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.11.cross_attn.q_linear.weight": "model-00001-of-00003.safetensors",
+    "blocks.11.cross_attn.q_linear.bias": "model-00001-of-00003.safetensors",
+    "blocks.11.cross_attn.to_k.weight": "model-00001-of-00003.safetensors",
+    "blocks.11.cross_attn.to_k.bias": "model-00001-of-00003.safetensors",
+    "blocks.11.cross_attn.to_v.weight": "model-00001-of-00003.safetensors",
+    "blocks.11.cross_attn.to_v.bias": "model-00001-of-00003.safetensors",
+    "blocks.11.cross_attn.proj.weight": "model-00001-of-00003.safetensors",
+    "blocks.11.cross_attn.proj.bias": "model-00001-of-00003.safetensors",
+    "blocks.11.cross_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.11.cross_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.11.mlp.inverted_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.11.mlp.inverted_conv.conv.bias": "model-00001-of-00003.safetensors",
+    "blocks.11.mlp.depth_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.11.mlp.depth_conv.conv.bias": "model-00001-of-00003.safetensors",
+    "blocks.11.mlp.point_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.12.scale_shift_table": "model-00001-of-00003.safetensors",
+    "blocks.12.attn.to_q.weight": "model-00001-of-00003.safetensors",
+    "blocks.12.attn.to_k.weight": "model-00001-of-00003.safetensors",
+    "blocks.12.attn.to_v.weight": "model-00001-of-00003.safetensors",
+    "blocks.12.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "blocks.12.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "blocks.12.attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.12.attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.12.cross_attn.q_linear.weight": "model-00001-of-00003.safetensors",
+    "blocks.12.cross_attn.q_linear.bias": "model-00001-of-00003.safetensors",
+    "blocks.12.cross_attn.to_k.weight": "model-00001-of-00003.safetensors",
+    "blocks.12.cross_attn.to_k.bias": "model-00001-of-00003.safetensors",
+    "blocks.12.cross_attn.to_v.weight": "model-00001-of-00003.safetensors",
+    "blocks.12.cross_attn.to_v.bias": "model-00001-of-00003.safetensors",
+    "blocks.12.cross_attn.proj.weight": "model-00001-of-00003.safetensors",
+    "blocks.12.cross_attn.proj.bias": "model-00001-of-00003.safetensors",
+    "blocks.12.cross_attn.q_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.12.cross_attn.k_norm.weight": "model-00001-of-00003.safetensors",
+    "blocks.12.mlp.inverted_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.12.mlp.inverted_conv.conv.bias": "model-00001-of-00003.safetensors",
+    "blocks.12.mlp.depth_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.12.mlp.depth_conv.conv.bias": "model-00001-of-00003.safetensors",
+    "blocks.12.mlp.point_conv.conv.weight": "model-00001-of-00003.safetensors",
+    "blocks.13.scale_shift_table": "model-00001-of-00003.safetensors",
+    "blocks.13.attn.to_q.weight": "model-00001-of-00003.safetensors",
+    "blocks.13.attn.to_k.weight": "model-00001-of-00003.safetensors",
+    "blocks.13.attn.to_v.weight": "model-00001-of-00003.safetensors",
+    "blocks.13.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "blocks.13.attn.proj.bias": "model-00002-of-00003.safetensors",
+    "blocks.13.attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.13.attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.13.cross_attn.q_linear.weight": "model-00002-of-00003.safetensors",
+    "blocks.13.cross_attn.q_linear.bias": "model-00002-of-00003.safetensors",
+    "blocks.13.cross_attn.to_k.weight": "model-00002-of-00003.safetensors",
+    "blocks.13.cross_attn.to_k.bias": "model-00002-of-00003.safetensors",
+    "blocks.13.cross_attn.to_v.weight": "model-00002-of-00003.safetensors",
+    "blocks.13.cross_attn.to_v.bias": "model-00002-of-00003.safetensors",
+    "blocks.13.cross_attn.proj.weight": "model-00002-of-00003.safetensors",
+    "blocks.13.cross_attn.proj.bias": "model-00002-of-00003.safetensors",
+    "blocks.13.cross_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.13.cross_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.13.mlp.inverted_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.13.mlp.inverted_conv.conv.bias": "model-00002-of-00003.safetensors",
+    "blocks.13.mlp.depth_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.13.mlp.depth_conv.conv.bias": "model-00002-of-00003.safetensors",
+    "blocks.13.mlp.point_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.14.scale_shift_table": "model-00002-of-00003.safetensors",
+    "blocks.14.attn.to_q.weight": "model-00002-of-00003.safetensors",
+    "blocks.14.attn.to_k.weight": "model-00002-of-00003.safetensors",
+    "blocks.14.attn.to_v.weight": "model-00002-of-00003.safetensors",
+    "blocks.14.attn.proj.weight": "model-00002-of-00003.safetensors",
+    "blocks.14.attn.proj.bias": "model-00002-of-00003.safetensors",
+    "blocks.14.attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.14.attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.14.cross_attn.q_linear.weight": "model-00002-of-00003.safetensors",
+    "blocks.14.cross_attn.q_linear.bias": "model-00002-of-00003.safetensors",
+    "blocks.14.cross_attn.to_k.weight": "model-00002-of-00003.safetensors",
+    "blocks.14.cross_attn.to_k.bias": "model-00002-of-00003.safetensors",
+    "blocks.14.cross_attn.to_v.weight": "model-00002-of-00003.safetensors",
+    "blocks.14.cross_attn.to_v.bias": "model-00002-of-00003.safetensors",
+    "blocks.14.cross_attn.proj.weight": "model-00002-of-00003.safetensors",
+    "blocks.14.cross_attn.proj.bias": "model-00002-of-00003.safetensors",
+    "blocks.14.cross_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.14.cross_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.14.mlp.inverted_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.14.mlp.inverted_conv.conv.bias": "model-00002-of-00003.safetensors",
+    "blocks.14.mlp.depth_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.14.mlp.depth_conv.conv.bias": "model-00002-of-00003.safetensors",
+    "blocks.14.mlp.point_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.15.scale_shift_table": "model-00002-of-00003.safetensors",
+    "blocks.15.attn.to_q.weight": "model-00002-of-00003.safetensors",
+    "blocks.15.attn.to_k.weight": "model-00002-of-00003.safetensors",
+    "blocks.15.attn.to_v.weight": "model-00002-of-00003.safetensors",
+    "blocks.15.attn.proj.weight": "model-00002-of-00003.safetensors",
+    "blocks.15.attn.proj.bias": "model-00002-of-00003.safetensors",
+    "blocks.15.attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.15.attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.15.cross_attn.q_linear.weight": "model-00002-of-00003.safetensors",
+    "blocks.15.cross_attn.q_linear.bias": "model-00002-of-00003.safetensors",
+    "blocks.15.cross_attn.to_k.weight": "model-00002-of-00003.safetensors",
+    "blocks.15.cross_attn.to_k.bias": "model-00002-of-00003.safetensors",
+    "blocks.15.cross_attn.to_v.weight": "model-00002-of-00003.safetensors",
+    "blocks.15.cross_attn.to_v.bias": "model-00002-of-00003.safetensors",
+    "blocks.15.cross_attn.proj.weight": "model-00002-of-00003.safetensors",
+    "blocks.15.cross_attn.proj.bias": "model-00002-of-00003.safetensors",
+    "blocks.15.cross_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.15.cross_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.15.mlp.inverted_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.15.mlp.inverted_conv.conv.bias": "model-00002-of-00003.safetensors",
+    "blocks.15.mlp.depth_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.15.mlp.depth_conv.conv.bias": "model-00002-of-00003.safetensors",
+    "blocks.15.mlp.point_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.16.scale_shift_table": "model-00002-of-00003.safetensors",
+    "blocks.16.attn.to_q.weight": "model-00002-of-00003.safetensors",
+    "blocks.16.attn.to_k.weight": "model-00002-of-00003.safetensors",
+    "blocks.16.attn.to_v.weight": "model-00002-of-00003.safetensors",
+    "blocks.16.attn.proj.weight": "model-00002-of-00003.safetensors",
+    "blocks.16.attn.proj.bias": "model-00002-of-00003.safetensors",
+    "blocks.16.attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.16.attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.16.cross_attn.q_linear.weight": "model-00002-of-00003.safetensors",
+    "blocks.16.cross_attn.q_linear.bias": "model-00002-of-00003.safetensors",
+    "blocks.16.cross_attn.to_k.weight": "model-00002-of-00003.safetensors",
+    "blocks.16.cross_attn.to_k.bias": "model-00002-of-00003.safetensors",
+    "blocks.16.cross_attn.to_v.weight": "model-00002-of-00003.safetensors",
+    "blocks.16.cross_attn.to_v.bias": "model-00002-of-00003.safetensors",
+    "blocks.16.cross_attn.proj.weight": "model-00002-of-00003.safetensors",
+    "blocks.16.cross_attn.proj.bias": "model-00002-of-00003.safetensors",
+    "blocks.16.cross_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.16.cross_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.16.mlp.inverted_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.16.mlp.inverted_conv.conv.bias": "model-00002-of-00003.safetensors",
+    "blocks.16.mlp.depth_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.16.mlp.depth_conv.conv.bias": "model-00002-of-00003.safetensors",
+    "blocks.16.mlp.point_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.17.scale_shift_table": "model-00002-of-00003.safetensors",
+    "blocks.17.attn.to_q.weight": "model-00002-of-00003.safetensors",
+    "blocks.17.attn.to_k.weight": "model-00002-of-00003.safetensors",
+    "blocks.17.attn.to_v.weight": "model-00002-of-00003.safetensors",
+    "blocks.17.attn.proj.weight": "model-00002-of-00003.safetensors",
+    "blocks.17.attn.proj.bias": "model-00002-of-00003.safetensors",
+    "blocks.17.attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.17.attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.17.cross_attn.q_linear.weight": "model-00002-of-00003.safetensors",
+    "blocks.17.cross_attn.q_linear.bias": "model-00002-of-00003.safetensors",
+    "blocks.17.cross_attn.to_k.weight": "model-00002-of-00003.safetensors",
+    "blocks.17.cross_attn.to_k.bias": "model-00002-of-00003.safetensors",
+    "blocks.17.cross_attn.to_v.weight": "model-00002-of-00003.safetensors",
+    "blocks.17.cross_attn.to_v.bias": "model-00002-of-00003.safetensors",
+    "blocks.17.cross_attn.proj.weight": "model-00002-of-00003.safetensors",
+    "blocks.17.cross_attn.proj.bias": "model-00002-of-00003.safetensors",
+    "blocks.17.cross_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.17.cross_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.17.mlp.inverted_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.17.mlp.inverted_conv.conv.bias": "model-00002-of-00003.safetensors",
+    "blocks.17.mlp.depth_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.17.mlp.depth_conv.conv.bias": "model-00002-of-00003.safetensors",
+    "blocks.17.mlp.point_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.18.scale_shift_table": "model-00002-of-00003.safetensors",
+    "blocks.18.attn.to_q.weight": "model-00002-of-00003.safetensors",
+    "blocks.18.attn.to_k.weight": "model-00002-of-00003.safetensors",
+    "blocks.18.attn.to_v.weight": "model-00002-of-00003.safetensors",
+    "blocks.18.attn.proj.weight": "model-00002-of-00003.safetensors",
+    "blocks.18.attn.proj.bias": "model-00002-of-00003.safetensors",
+    "blocks.18.attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.18.attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.18.cross_attn.q_linear.weight": "model-00002-of-00003.safetensors",
+    "blocks.18.cross_attn.q_linear.bias": "model-00002-of-00003.safetensors",
+    "blocks.18.cross_attn.to_k.weight": "model-00002-of-00003.safetensors",
+    "blocks.18.cross_attn.to_k.bias": "model-00002-of-00003.safetensors",
+    "blocks.18.cross_attn.to_v.weight": "model-00002-of-00003.safetensors",
+    "blocks.18.cross_attn.to_v.bias": "model-00002-of-00003.safetensors",
+    "blocks.18.cross_attn.proj.weight": "model-00002-of-00003.safetensors",
+    "blocks.18.cross_attn.proj.bias": "model-00002-of-00003.safetensors",
+    "blocks.18.cross_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.18.cross_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.18.mlp.inverted_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.18.mlp.inverted_conv.conv.bias": "model-00002-of-00003.safetensors",
+    "blocks.18.mlp.depth_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.18.mlp.depth_conv.conv.bias": "model-00002-of-00003.safetensors",
+    "blocks.18.mlp.point_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.19.scale_shift_table": "model-00002-of-00003.safetensors",
+    "blocks.19.attn.to_q.weight": "model-00002-of-00003.safetensors",
+    "blocks.19.attn.to_k.weight": "model-00002-of-00003.safetensors",
+    "blocks.19.attn.to_v.weight": "model-00002-of-00003.safetensors",
+    "blocks.19.attn.proj.weight": "model-00002-of-00003.safetensors",
+    "blocks.19.attn.proj.bias": "model-00002-of-00003.safetensors",
+    "blocks.19.attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.19.attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.19.cross_attn.q_linear.weight": "model-00002-of-00003.safetensors",
+    "blocks.19.cross_attn.q_linear.bias": "model-00002-of-00003.safetensors",
+    "blocks.19.cross_attn.to_k.weight": "model-00002-of-00003.safetensors",
+    "blocks.19.cross_attn.to_k.bias": "model-00002-of-00003.safetensors",
+    "blocks.19.cross_attn.to_v.weight": "model-00002-of-00003.safetensors",
+    "blocks.19.cross_attn.to_v.bias": "model-00002-of-00003.safetensors",
+    "blocks.19.cross_attn.proj.weight": "model-00002-of-00003.safetensors",
+    "blocks.19.cross_attn.proj.bias": "model-00002-of-00003.safetensors",
+    "blocks.19.cross_attn.q_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.19.cross_attn.k_norm.weight": "model-00002-of-00003.safetensors",
+    "blocks.19.mlp.inverted_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.19.mlp.inverted_conv.conv.bias": "model-00002-of-00003.safetensors",
+    "blocks.19.mlp.depth_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "blocks.19.mlp.depth_conv.conv.bias": "model-00002-of-00003.safetensors",
+    "blocks.19.mlp.point_conv.conv.weight": "model-00002-of-00003.safetensors",
+    "final_layer.scale_shift_table": "model-00002-of-00003.safetensors",
+    "final_layer.linear.weight": "model-00002-of-00003.safetensors",
+    "final_layer.linear.bias": "model-00002-of-00003.safetensors",
+    "y_embedder.y_embedding": "model-00002-of-00003.safetensors"
+  }
+}

requirements.txt

@@ -0,0 +1,14 @@
+# Tested with Python 3.12, PyTorch 2.6.0 and CUDA 12.9.
+torch==2.6.0
+torchvision==0.21.0
+diffusers==0.36.0
+transformers==4.51.1
+safetensors==0.5.3
+einops==0.8.1
+accelerate>=1.9
+pydantic==2.11.7
+numpy==2.2.6
+Pillow==11.3.0
+huggingface-hub==0.36.0
+PyYAML>=6.0
+# Install separately AFTER torch: pip install flash-attn==2.7.4.post1 --no-build-isolation