Upload MarkupDMForCausalLM

config.json

@@ -1,4 +1,5 @@
 {
+  "_name_or_path": "/data/sc1_7b/",
   "architectures": [
     "MarkupDMForCausalLM"
   ],
@@ -109,7 +110,7 @@
     "validate_runner_input": true,
     "vocab_size": 49156
   },
-  "torch_dtype": "bfloat16",
+  "torch_dtype": "float32",
   "transformers_version": "4.47.1",
   "vision_model": {
     "_attn_implementation_autoset": true,

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.47.1"
+}

loss_utils.py

@@ -0,0 +1,70 @@
+import torch
+import torch.nn.functional as F
+
+
+def fixed_cross_entropy(
+    source,
+    target,
+    num_items_in_batch: int | None = None,
+    ignore_index: int = -100,
+    weight=None,
+    **kwargs,
+):
+    reduction = "sum" if num_items_in_batch is not None else "mean"
+    loss = F.cross_entropy(
+        source,
+        target,
+        ignore_index=ignore_index,
+        reduction=reduction,
+        weight=weight,
+    )
+    if reduction == "sum":
+        loss = loss / num_items_in_batch
+    return loss
+
+
+def WeightedCausalLMLoss(
+    logits,
+    labels,
+    image_vocab_size: int,
+    image_loss_weight: float = 1.0,
+    image_token_ratio: float = 2.4,
+    num_items_in_batch: int | None = None,
+    ignore_index: int = -100,
+    **kwargs,
+):
+    # Upcast to float if we need to compute the loss to avoid potential precision issues
+    logits = logits.float()
+    labels = labels.to(logits.device)
+    # Shift so that tokens < n predict n
+    labels = F.pad(labels, (0, 1), value=ignore_index)
+    shift_labels = labels[..., 1:].contiguous()
+
+    # Compute loss weight
+    if image_loss_weight != 1.0:
+        weight = torch.ones(logits.size(-1), device=logits.device)
+        weight[-image_vocab_size:] = image_loss_weight
+    else:
+        weight = None
+
+    # Flatten the tokens
+    logits = logits.view(-1, logits.size(-1))
+    shift_labels = shift_labels.view(-1)
+    # Enable model parallelism
+    shift_labels = shift_labels.to(logits.device)
+    loss = fixed_cross_entropy(
+        logits,
+        shift_labels,
+        num_items_in_batch,
+        ignore_index,
+        weight=weight,
+        **kwargs,
+    )
+
+    # Scale the loss
+    if image_loss_weight != 1.0:
+        denom = 1.0 + (image_token_ratio * image_loss_weight)
+        scale = (1.0 + image_token_ratio) / denom
+        loss = scale * loss
+
+    return loss

model-00001-of-00007.safetensors

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

model-00002-of-00007.safetensors

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

model-00003-of-00007.safetensors

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

model-00004-of-00007.safetensors

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

model-00005-of-00007.safetensors

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

model-00006-of-00007.safetensors

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

model-00007-of-00007.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,862 @@
+{
+  "metadata": {
+    "total_size": 29856215152
+  },
+  "weight_map": {
+    "proj_vt.bias": "model-00006-of-00007.safetensors",
+    "proj_vt.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.0.attn.c_attn.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.0.attn.c_attn.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.0.attn.c_proj.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.0.attn.c_proj.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.0.ln_1.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.0.ln_1.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.0.ln_2.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.0.ln_2.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.0.mlp.c_fc.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.0.mlp.c_fc.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.0.mlp.c_proj.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.0.mlp.c_proj.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.1.attn.c_attn.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.1.attn.c_attn.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.1.attn.c_proj.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.1.attn.c_proj.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.1.ln_1.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.1.ln_1.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.1.ln_2.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.1.ln_2.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.1.mlp.c_fc.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.1.mlp.c_fc.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.1.mlp.c_proj.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.1.mlp.c_proj.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.10.attn.c_attn.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.10.attn.c_attn.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.10.attn.c_proj.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.10.attn.c_proj.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.10.ln_1.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.10.ln_1.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.10.ln_2.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.10.ln_2.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.10.mlp.c_fc.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.10.mlp.c_fc.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.10.mlp.c_proj.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.10.mlp.c_proj.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.11.attn.c_attn.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.11.attn.c_attn.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.11.attn.c_proj.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.11.attn.c_proj.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.11.ln_1.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.11.ln_1.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.11.ln_2.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.11.ln_2.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.11.mlp.c_fc.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.11.mlp.c_fc.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.11.mlp.c_proj.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.11.mlp.c_proj.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.12.attn.c_attn.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.12.attn.c_attn.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.12.attn.c_proj.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.12.attn.c_proj.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.12.ln_1.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.12.ln_1.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.12.ln_2.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.12.ln_2.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.12.mlp.c_fc.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.12.mlp.c_fc.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.12.mlp.c_proj.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.12.mlp.c_proj.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.13.attn.c_attn.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.13.attn.c_attn.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.13.attn.c_proj.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.13.attn.c_proj.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.13.ln_1.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.13.ln_1.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.13.ln_2.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.13.ln_2.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.13.mlp.c_fc.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.13.mlp.c_fc.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.13.mlp.c_proj.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.13.mlp.c_proj.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.14.attn.c_attn.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.14.attn.c_attn.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.14.attn.c_proj.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.14.attn.c_proj.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.14.ln_1.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.14.ln_1.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.14.ln_2.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.14.ln_2.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.14.mlp.c_fc.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.14.mlp.c_fc.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.14.mlp.c_proj.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.14.mlp.c_proj.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.15.attn.c_attn.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.15.attn.c_attn.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.15.attn.c_proj.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.15.attn.c_proj.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.15.ln_1.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.15.ln_1.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.15.ln_2.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.15.ln_2.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.15.mlp.c_fc.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.15.mlp.c_fc.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.15.mlp.c_proj.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.15.mlp.c_proj.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.16.attn.c_attn.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.16.attn.c_attn.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.16.attn.c_proj.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.16.attn.c_proj.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.16.ln_1.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.16.ln_1.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.16.ln_2.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.16.ln_2.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.16.mlp.c_fc.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.16.mlp.c_fc.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.16.mlp.c_proj.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.16.mlp.c_proj.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.17.attn.c_attn.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.17.attn.c_attn.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.17.attn.c_proj.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.17.attn.c_proj.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.17.ln_1.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.17.ln_1.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.17.ln_2.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.17.ln_2.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.17.mlp.c_fc.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.17.mlp.c_fc.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.17.mlp.c_proj.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.17.mlp.c_proj.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.18.attn.c_attn.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.18.attn.c_attn.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.18.attn.c_proj.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.18.attn.c_proj.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.18.ln_1.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.18.ln_1.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.18.ln_2.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.18.ln_2.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.18.mlp.c_fc.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.18.mlp.c_fc.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.18.mlp.c_proj.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.18.mlp.c_proj.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.19.attn.c_attn.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.19.attn.c_attn.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.19.attn.c_proj.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.19.attn.c_proj.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.19.ln_1.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.19.ln_1.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.19.ln_2.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.19.ln_2.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.19.mlp.c_fc.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.19.mlp.c_fc.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.19.mlp.c_proj.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.19.mlp.c_proj.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.2.attn.c_attn.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.2.attn.c_attn.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.2.attn.c_proj.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.2.attn.c_proj.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.2.ln_1.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.2.ln_1.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.2.ln_2.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.2.ln_2.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.2.mlp.c_fc.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.2.mlp.c_fc.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.2.mlp.c_proj.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.2.mlp.c_proj.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.20.attn.c_attn.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.20.attn.c_attn.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.20.attn.c_proj.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.20.attn.c_proj.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.20.ln_1.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.20.ln_1.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.20.ln_2.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.20.ln_2.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.20.mlp.c_fc.bias": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.20.mlp.c_fc.weight": "model-00003-of-00007.safetensors",
+    "text_model.transformer.h.20.mlp.c_proj.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.20.mlp.c_proj.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.21.attn.c_attn.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.21.attn.c_attn.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.21.attn.c_proj.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.21.attn.c_proj.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.21.ln_1.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.21.ln_1.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.21.ln_2.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.21.ln_2.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.21.mlp.c_fc.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.21.mlp.c_fc.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.21.mlp.c_proj.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.21.mlp.c_proj.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.22.attn.c_attn.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.22.attn.c_attn.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.22.attn.c_proj.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.22.attn.c_proj.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.22.ln_1.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.22.ln_1.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.22.ln_2.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.22.ln_2.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.22.mlp.c_fc.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.22.mlp.c_fc.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.22.mlp.c_proj.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.22.mlp.c_proj.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.23.attn.c_attn.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.23.attn.c_attn.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.23.attn.c_proj.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.23.attn.c_proj.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.23.ln_1.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.23.ln_1.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.23.ln_2.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.23.ln_2.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.23.mlp.c_fc.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.23.mlp.c_fc.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.23.mlp.c_proj.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.23.mlp.c_proj.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.24.attn.c_attn.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.24.attn.c_attn.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.24.attn.c_proj.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.24.attn.c_proj.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.24.ln_1.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.24.ln_1.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.24.ln_2.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.24.ln_2.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.24.mlp.c_fc.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.24.mlp.c_fc.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.24.mlp.c_proj.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.24.mlp.c_proj.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.25.attn.c_attn.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.25.attn.c_attn.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.25.attn.c_proj.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.25.attn.c_proj.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.25.ln_1.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.25.ln_1.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.25.ln_2.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.25.ln_2.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.25.mlp.c_fc.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.25.mlp.c_fc.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.25.mlp.c_proj.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.25.mlp.c_proj.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.26.attn.c_attn.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.26.attn.c_attn.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.26.attn.c_proj.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.26.attn.c_proj.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.26.ln_1.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.26.ln_1.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.26.ln_2.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.26.ln_2.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.26.mlp.c_fc.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.26.mlp.c_fc.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.26.mlp.c_proj.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.26.mlp.c_proj.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.27.attn.c_attn.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.27.attn.c_attn.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.27.attn.c_proj.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.27.attn.c_proj.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.27.ln_1.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.27.ln_1.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.27.ln_2.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.27.ln_2.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.27.mlp.c_fc.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.27.mlp.c_fc.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.27.mlp.c_proj.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.27.mlp.c_proj.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.28.attn.c_attn.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.28.attn.c_attn.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.28.attn.c_proj.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.28.attn.c_proj.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.28.ln_1.bias": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.28.ln_1.weight": "model-00004-of-00007.safetensors",
+    "text_model.transformer.h.28.ln_2.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.28.ln_2.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.28.mlp.c_fc.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.28.mlp.c_fc.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.28.mlp.c_proj.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.28.mlp.c_proj.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.29.attn.c_attn.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.29.attn.c_attn.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.29.attn.c_proj.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.29.attn.c_proj.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.29.ln_1.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.29.ln_1.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.29.ln_2.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.29.ln_2.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.29.mlp.c_fc.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.29.mlp.c_fc.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.29.mlp.c_proj.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.29.mlp.c_proj.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.3.attn.c_attn.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.3.attn.c_attn.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.3.attn.c_proj.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.3.attn.c_proj.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.3.ln_1.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.3.ln_1.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.3.ln_2.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.3.ln_2.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.3.mlp.c_fc.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.3.mlp.c_fc.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.3.mlp.c_proj.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.3.mlp.c_proj.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.30.attn.c_attn.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.30.attn.c_attn.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.30.attn.c_proj.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.30.attn.c_proj.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.30.ln_1.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.30.ln_1.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.30.ln_2.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.30.ln_2.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.30.mlp.c_fc.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.30.mlp.c_fc.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.30.mlp.c_proj.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.30.mlp.c_proj.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.31.attn.c_attn.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.31.attn.c_attn.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.31.attn.c_proj.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.31.attn.c_proj.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.31.ln_1.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.31.ln_1.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.31.ln_2.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.31.ln_2.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.31.mlp.c_fc.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.31.mlp.c_fc.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.31.mlp.c_proj.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.31.mlp.c_proj.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.32.attn.c_attn.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.32.attn.c_attn.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.32.attn.c_proj.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.32.attn.c_proj.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.32.ln_1.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.32.ln_1.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.32.ln_2.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.32.ln_2.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.32.mlp.c_fc.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.32.mlp.c_fc.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.32.mlp.c_proj.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.32.mlp.c_proj.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.33.attn.c_attn.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.33.attn.c_attn.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.33.attn.c_proj.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.33.attn.c_proj.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.33.ln_1.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.33.ln_1.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.33.ln_2.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.33.ln_2.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.33.mlp.c_fc.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.33.mlp.c_fc.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.33.mlp.c_proj.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.33.mlp.c_proj.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.34.attn.c_attn.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.34.attn.c_attn.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.34.attn.c_proj.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.34.attn.c_proj.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.34.ln_1.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.34.ln_1.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.34.ln_2.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.34.ln_2.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.34.mlp.c_fc.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.34.mlp.c_fc.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.34.mlp.c_proj.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.34.mlp.c_proj.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.35.attn.c_attn.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.35.attn.c_attn.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.35.attn.c_proj.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.35.attn.c_proj.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.35.ln_1.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.35.ln_1.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.35.ln_2.bias": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.35.ln_2.weight": "model-00005-of-00007.safetensors",
+    "text_model.transformer.h.35.mlp.c_fc.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.35.mlp.c_fc.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.35.mlp.c_proj.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.35.mlp.c_proj.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.36.attn.c_attn.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.36.attn.c_attn.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.36.attn.c_proj.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.36.attn.c_proj.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.36.ln_1.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.36.ln_1.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.36.ln_2.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.36.ln_2.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.36.mlp.c_fc.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.36.mlp.c_fc.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.36.mlp.c_proj.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.36.mlp.c_proj.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.37.attn.c_attn.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.37.attn.c_attn.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.37.attn.c_proj.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.37.attn.c_proj.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.37.ln_1.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.37.ln_1.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.37.ln_2.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.37.ln_2.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.37.mlp.c_fc.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.37.mlp.c_fc.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.37.mlp.c_proj.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.37.mlp.c_proj.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.38.attn.c_attn.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.38.attn.c_attn.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.38.attn.c_proj.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.38.attn.c_proj.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.38.ln_1.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.38.ln_1.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.38.ln_2.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.38.ln_2.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.38.mlp.c_fc.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.38.mlp.c_fc.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.38.mlp.c_proj.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.38.mlp.c_proj.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.39.attn.c_attn.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.39.attn.c_attn.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.39.attn.c_proj.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.39.attn.c_proj.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.39.ln_1.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.39.ln_1.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.39.ln_2.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.39.ln_2.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.39.mlp.c_fc.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.39.mlp.c_fc.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.39.mlp.c_proj.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.39.mlp.c_proj.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.4.attn.c_attn.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.4.attn.c_attn.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.4.attn.c_proj.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.4.attn.c_proj.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.4.ln_1.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.4.ln_1.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.4.ln_2.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.4.ln_2.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.4.mlp.c_fc.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.4.mlp.c_fc.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.4.mlp.c_proj.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.4.mlp.c_proj.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.40.attn.c_attn.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.40.attn.c_attn.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.40.attn.c_proj.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.40.attn.c_proj.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.40.ln_1.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.40.ln_1.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.40.ln_2.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.40.ln_2.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.40.mlp.c_fc.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.40.mlp.c_fc.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.40.mlp.c_proj.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.40.mlp.c_proj.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.41.attn.c_attn.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.41.attn.c_attn.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.41.attn.c_proj.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.41.attn.c_proj.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.41.ln_1.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.41.ln_1.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.41.ln_2.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.41.ln_2.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.41.mlp.c_fc.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.41.mlp.c_fc.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.41.mlp.c_proj.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.41.mlp.c_proj.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.h.5.attn.c_attn.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.5.attn.c_attn.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.5.attn.c_proj.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.5.attn.c_proj.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.5.ln_1.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.5.ln_1.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.5.ln_2.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.5.ln_2.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.5.mlp.c_fc.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.5.mlp.c_fc.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.5.mlp.c_proj.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.5.mlp.c_proj.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.6.attn.c_attn.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.6.attn.c_attn.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.6.attn.c_proj.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.6.attn.c_proj.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.6.ln_1.bias": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.6.ln_1.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.h.6.ln_2.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.6.ln_2.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.6.mlp.c_fc.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.6.mlp.c_fc.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.6.mlp.c_proj.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.6.mlp.c_proj.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.7.attn.c_attn.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.7.attn.c_attn.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.7.attn.c_proj.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.7.attn.c_proj.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.7.ln_1.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.7.ln_1.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.7.ln_2.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.7.ln_2.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.7.mlp.c_fc.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.7.mlp.c_fc.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.7.mlp.c_proj.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.7.mlp.c_proj.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.8.attn.c_attn.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.8.attn.c_attn.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.8.attn.c_proj.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.8.attn.c_proj.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.8.ln_1.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.8.ln_1.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.8.ln_2.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.8.ln_2.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.8.mlp.c_fc.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.8.mlp.c_fc.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.8.mlp.c_proj.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.8.mlp.c_proj.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.9.attn.c_attn.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.9.attn.c_attn.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.9.attn.c_proj.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.9.attn.c_proj.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.9.ln_1.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.9.ln_1.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.9.ln_2.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.9.ln_2.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.9.mlp.c_fc.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.9.mlp.c_fc.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.9.mlp.c_proj.bias": "model-00002-of-00007.safetensors",
+    "text_model.transformer.h.9.mlp.c_proj.weight": "model-00002-of-00007.safetensors",
+    "text_model.transformer.ln_f.bias": "model-00006-of-00007.safetensors",
+    "text_model.transformer.ln_f.weight": "model-00006-of-00007.safetensors",
+    "text_model.transformer.wpe.weight": "model-00001-of-00007.safetensors",
+    "text_model.transformer.wte.weight": "model-00001-of-00007.safetensors",
+    "vis_head.bias": "model-00007-of-00007.safetensors",
+    "vis_head.weight": "model-00007-of-00007.safetensors",
+    "vision_model.model.decoder.conv_in.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.conv_in.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.conv_out.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.conv_out.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.attn_1.k.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.attn_1.k.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.attn_1.norm.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.attn_1.norm.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.attn_1.proj_out.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.attn_1.proj_out.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.attn_1.q.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.attn_1.q.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.attn_1.v.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.attn_1.v.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_1.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_1.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_1.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_1.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_1.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_1.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_1.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_1.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_2.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_2.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_2.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_2.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_2.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_2.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_2.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.mid.block_2.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.norm_out.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.norm_out.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.0.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.0.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.0.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.0.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.0.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.0.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.0.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.0.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.1.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.1.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.1.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.1.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.1.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.1.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.1.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.1.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.2.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.2.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.2.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.2.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.2.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.2.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.2.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.0.block.2.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.0.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.0.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.0.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.0.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.0.nin_shortcut.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.0.nin_shortcut.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.0.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.0.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.0.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.0.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.1.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.1.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.1.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.1.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.1.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.1.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.1.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.1.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.2.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.2.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.2.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.2.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.2.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.2.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.2.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.block.2.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.upsample.conv.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.1.upsample.conv.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.0.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.0.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.0.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.0.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.0.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.0.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.0.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.0.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.1.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.1.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.1.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.1.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.1.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.1.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.1.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.1.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.2.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.2.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.2.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.2.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.2.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.2.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.2.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.block.2.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.upsample.conv.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.2.upsample.conv.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.0.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.0.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.0.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.0.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.0.nin_shortcut.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.0.nin_shortcut.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.0.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.0.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.0.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.0.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.1.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.1.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.1.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.1.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.1.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.1.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.1.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.1.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.2.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.2.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.2.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.2.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.2.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.2.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.2.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.block.2.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.upsample.conv.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.3.upsample.conv.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.0.k.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.0.k.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.0.norm.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.0.norm.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.0.proj_out.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.0.proj_out.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.0.q.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.0.q.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.0.v.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.0.v.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.1.k.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.1.k.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.1.norm.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.1.norm.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.1.proj_out.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.1.proj_out.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.1.q.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.1.q.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.1.v.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.1.v.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.2.k.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.2.k.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.2.norm.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.2.norm.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.2.proj_out.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.2.proj_out.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.2.q.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.2.q.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.2.v.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.attn.2.v.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.0.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.0.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.0.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.0.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.0.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.0.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.0.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.0.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.1.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.1.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.1.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.1.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.1.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.1.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.1.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.1.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.2.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.2.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.2.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.2.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.2.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.2.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.2.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.block.2.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.upsample.conv.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.decoder.up.4.upsample.conv.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.conv_in.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.conv_in.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.conv_out.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.conv_out.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.0.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.0.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.0.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.0.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.0.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.0.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.0.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.0.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.1.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.1.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.1.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.1.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.1.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.1.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.1.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.block.1.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.downsample.conv.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.0.downsample.conv.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.0.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.0.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.0.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.0.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.0.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.0.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.0.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.0.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.1.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.1.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.1.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.1.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.1.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.1.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.1.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.block.1.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.downsample.conv.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.1.downsample.conv.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.0.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.0.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.0.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.0.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.0.nin_shortcut.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.0.nin_shortcut.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.0.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.0.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.0.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.0.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.1.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.1.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.1.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.1.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.1.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.1.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.1.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.block.1.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.downsample.conv.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.2.downsample.conv.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.0.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.0.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.0.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.0.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.0.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.0.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.0.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.0.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.1.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.1.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.1.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.1.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.1.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.1.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.1.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.block.1.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.downsample.conv.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.3.downsample.conv.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.0.k.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.0.k.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.0.norm.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.0.norm.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.0.proj_out.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.0.proj_out.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.0.q.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.0.q.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.0.v.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.0.v.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.1.k.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.1.k.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.1.norm.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.1.norm.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.1.proj_out.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.1.proj_out.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.1.q.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.1.q.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.1.v.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.attn.1.v.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.0.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.0.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.0.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.0.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.0.nin_shortcut.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.0.nin_shortcut.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.0.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.0.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.0.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.0.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.1.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.1.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.1.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.1.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.1.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.1.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.1.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.down.4.block.1.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.attn_1.k.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.attn_1.k.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.attn_1.norm.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.attn_1.norm.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.attn_1.proj_out.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.attn_1.proj_out.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.attn_1.q.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.attn_1.q.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.attn_1.v.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.attn_1.v.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_1.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_1.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_1.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_1.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_1.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_1.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_1.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_1.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_2.conv1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_2.conv1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_2.conv2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_2.conv2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_2.norm1.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_2.norm1.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_2.norm2.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.mid.block_2.norm2.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.norm_out.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.encoder.norm_out.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.post_quant_conv.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.post_quant_conv.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.quant_conv.bias": "model-00006-of-00007.safetensors",
+    "vision_model.model.quant_conv.weight": "model-00006-of-00007.safetensors",
+    "vision_model.model.quantize.embedding.weight": "model-00006-of-00007.safetensors"
+  }
+}

modeling_markupdm.py

@@ -0,0 +1,291 @@
+"""PyTorch MarkupDM model."""
+
+import contextlib
+import math
+import os
+from typing import Any
+
+import rff.layers
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import (
+    AutoModel,
+    AutoModelForCausalLM,
+    GenerationMixin,
+    PreTrainedModel,
+)
+from transformers.loss.loss_utils import LOSS_MAPPING
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.utils import logging
+
+from .configuration_markupdm import MarkupDMConfig
+from .loss_utils import WeightedCausalLMLoss
+
+logger = logging.get_logger(__name__)
+
+LOSS_MAPPING["WeightedCausalLMLoss"] = WeightedCausalLMLoss
+
+
+class MarkupDMForCausalLM(PreTrainedModel, GenerationMixin):  # type: ignore
+    config: MarkupDMConfig
+    config_class = MarkupDMConfig
+
+    supports_gradient_checkpointing = True
+    _supports_flash_attn_2 = True
+
+    def __init__(
+        self,
+        config: MarkupDMConfig,
+        text_model: PreTrainedModel,
+        vision_model: PreTrainedModel,
+    ) -> None:
+        if not isinstance(config, self.config_class):
+            raise ValueError(f"Config: {config} has to be of type {self.config_class}")
+
+        # Initialize with config
+        logger.info(f"MarkupDM config: {config}")
+        super().__init__(config)
+
+        self.text_model = text_model.train()
+        self.vision_model = vision_model.eval().requires_grad_(False)
+
+        if self.text_model.config.to_dict() != self.config.text_model.to_dict():
+            logger.warning(
+                f"Config of the text model: {self.text_model.__class__} is"
+                f"overwritten by shared text config: {self.config.text_model}"
+            )
+        if self.vision_model.config.to_dict() != self.config.vision_model.to_dict():
+            logger.warning(
+                f"Config of the vision model: {self.vision_model.__class__} is"
+                f"overwritten by shared vision config: {self.config.vision_model}"
+            )
+
+        # Make sure that the individual model's config refers to the shared config
+        # so that the updates to the config will be synced
+        self.text_model.config = self.config.text_model
+        self.vision_model.config = self.config.vision_model
+
+        # Resize embedding layer
+        base_size = self.text_model.config.vocab_size
+        if base_size < self.config.vocab_size:
+            self.text_model.resize_token_embeddings(self.config.vocab_size)
+            new_size = self.text_model.get_input_embeddings().num_embeddings
+            logger.info(f"Resize embedding layer from {base_size} to {new_size} tokens")
+
+        d_text = self.text_model.config.hidden_size
+        assert self.vision_model.config.model_type == "vqmodel"
+        d_vision = self.vision_model.model.embed_dim
+        image_pos_size = self.config.image_pos_size
+        sigma = self.config.image_pos_sigma
+        m = math.ceil(image_pos_size / 2)  # (sin, cos)
+        self.image_vocab_size = self.vision_model.model.n_embed
+
+        # Define additional layers
+        self.proj_vpos = rff.layers.PositionalEncoding(sigma, m)
+        self.proj_vt = nn.Linear(d_vision + image_pos_size, d_text)
+        self.vis_head = nn.Linear(d_text, self.image_vocab_size)
+
+        # Compute num_image_tokens
+        scale_factor = 2 ** (vision_model.model.encoder.num_resolutions - 1)
+        latent_size = self.config.image_size // scale_factor
+        self.num_image_tokens = latent_size**2
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+        # Freeze text embeddings if needed
+        if config.freeze_text_embeddings:
+            self.text_model.get_input_embeddings().requires_grad_(False)
+
+    def tie_weights(self) -> None:
+        self.text_model.tie_weights()
+
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any) -> "MarkupDMForCausalLM":
+        assert "config" in kwargs, "Config must be provided"
+        config = kwargs["config"]
+        torch_dtype = kwargs.get("torch_dtype", None)
+
+        # Initialize text model
+        text_model = AutoModelForCausalLM.from_config(
+            config.text_model,
+            torch_dtype=torch_dtype,
+            attn_implementation=config._attn_implementation,
+        )
+
+        # Initialize vision model
+        with contextlib.redirect_stdout(open(os.devnull, "w")):
+            vision_model = AutoModel.from_config(
+                config.vision_model,
+                trust_remote_code=True,
+                torch_dtype=torch_dtype,
+            )
+
+        return super().from_pretrained(  # type: ignore
+            *args,
+            **kwargs,
+            text_model=text_model,
+            vision_model=vision_model,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        inputs_embeds: torch.Tensor | None = None,
+        image_mask: torch.Tensor | None = None,
+        image_pos_ids: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        past_key_values: tuple[tuple[torch.Tensor]] | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        cache_position: torch.Tensor | None = None,
+        num_items_in_batch: int | None = None,
+        **kwargs: Any,
+    ) -> CausalLMOutputWithPast:
+        for key in kwargs.keys():
+            if kwargs[key] is not None:
+                raise ValueError(f"Unknown argument: {key}={kwargs[key]}")
+
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        if image_mask is None:
+            image_mask = input_ids >= self.config.vocab_size
+
+        # Embed inputs
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(
+                input_ids,
+                image_mask=image_mask,
+                image_pos_ids=image_pos_ids,
+            )
+
+        # Core forward pass
+        fwd_kwargs = {
+            "inputs_embeds": inputs_embeds,
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
+            "past_key_values": past_key_values,
+            "use_cache": use_cache,
+            "output_hidden_states": True,
+            "output_attentions": output_attentions,
+        }
+        if self.config.text_model.model_type == "starcoder2":
+            fwd_kwargs["cache_position"] = cache_position
+        outputs = self.text_model(**fwd_kwargs)
+
+        # text_logits: (B, L, V)
+        text_logits = outputs.logits[:, :, : self.config.vocab_size]
+
+        # vision_logits: (B, L, C)
+        last_hidden_states = outputs.hidden_states[-1]
+        vision_logits = self.vis_head(last_hidden_states)
+
+        if labels is not None:
+            # Mask logits with shifted image mask
+            shift_mask = F.pad(image_mask[:, 1:], (0, 1), value=False)
+            text_logits[shift_mask] = -float("inf")
+            vision_logits[~shift_mask] = -float("inf")
+
+        # Concatenate text and vision logits
+        logits = torch.cat([text_logits, vision_logits], dim=-1)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                image_vocab_size=self.image_vocab_size,
+                image_loss_weight=self.config.image_loss_weight,
+                num_items_in_batch=num_items_in_batch,
+                **kwargs,
+            )
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
+            attentions=outputs.attentions,
+        )
+
+    def embed_tokens(
+        self,
+        input_ids: torch.Tensor,
+        image_mask: torch.Tensor | None = None,
+        image_pos_ids: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if image_mask is None:
+            return self.text_embed(input_ids)  # type: ignore
+
+        # Prepare placeholders
+        size = input_ids.size() + (self.text_model.config.hidden_size,)
+        inputs_embeds = torch.zeros(size, device=self.device, dtype=self.dtype)
+
+        # Embed text ids
+        text_embeds = self.text_embed(input_ids[~image_mask])
+        inputs_embeds[~image_mask] = text_embeds
+
+        # Embed image ids
+        image_embeds = self.vis_embed(input_ids[image_mask] - self.config.vocab_size)
+
+        # Concatenate positional embeddings
+        assert image_pos_ids is not None
+        image_pos = image_pos_ids / self.num_image_tokens
+        image_pos = self.proj_vpos(image_pos.unsqueeze(-1)).to(image_embeds)
+        image_pos = image_pos[image_mask][:, : self.config.image_pos_size]
+        image_embeds = torch.cat([image_embeds, image_pos], dim=-1)  # type: ignore
+
+        # Project image features and update inputs_embeds
+        image_embeds = self.proj_vt(image_embeds)
+        inputs_embeds[image_mask] = image_embeds
+
+        return inputs_embeds
+
+    def text_embed(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.text_model.get_input_embeddings()(input_ids)  # type: ignore
+
+    def vis_embed(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.vision_model.model.quantize.embedding(input_ids)  # type: ignore
+
+    def prepare_inputs_for_generation(
+        self, input_ids: torch.Tensor, **model_kwargs: Any
+    ) -> dict:
+        # Prepare inputs with the default function
+        default_prepare_inputs = self.text_model.prepare_inputs_for_generation
+        inputs = default_prepare_inputs(input_ids, **model_kwargs)
+
+        # Compute image_pos_ids
+        base_ids = torch.arange(self.num_image_tokens, device=self.device)
+        image_pos_ids = torch.zeros_like(input_ids)
+        image_mask_all = input_ids >= self.config.vocab_size
+        for i_batch, image_mask in enumerate(image_mask_all):
+            N = sum(image_mask)
+            pos_ids = base_ids.repeat(N // self.num_image_tokens + 1)
+            image_pos_ids[i_batch, image_mask] = pos_ids[:N]
+        length = inputs["input_ids"].size(1)
+        inputs["image_pos_ids"] = image_pos_ids[:, -length:]
+
+        inputs["image_mask"] = inputs["input_ids"] >= self.config.vocab_size
+
+        return inputs  # type: ignore