Merge pull request #122 from younger-plinius/restore-multimodal-v2

restore-multimodal CLI command

ci/test-risk-map.json

@@ -2,15 +2,24 @@
   "schema_version": 2,
   "owner": "OBLITERATUS maintainers",
   "source_inventory": {
-    "roots": ["app.py", "obliteratus"],
-    "coverage_roots": ["obliteratus"]
+    "roots": [
+      "app.py",
+      "obliteratus"
+    ],
+    "coverage_roots": [
+      "obliteratus"
+    ]
   },
   "contract_surfaces": [
     {
       "id": "package-entrypoints",
       "owner": "packaging maintainers",
       "description": "Installed package exports, module entry points, console dispatch, and optional application construction",
-      "contract_types": ["package-entrypoint", "public-interface", "operator-ui"],
+      "contract_types": [
+        "package-entrypoint",
+        "public-interface",
+        "operator-ui"
+      ],
       "paths": [
         "app.py",
         "obliteratus/__init__.py",
@@ -33,7 +42,12 @@
       "id": "core-pipeline",
       "owner": "abliteration pipeline maintainers",
       "description": "Model mutation, orchestration, automated search, study execution, sweeps, and tournament lifecycle",
-      "contract_types": ["model-mutation", "model-runtime", "orchestration", "persistence"],
+      "contract_types": [
+        "model-mutation",
+        "model-runtime",
+        "orchestration",
+        "persistence"
+      ],
       "paths": [
         "obliteratus/abliterate.py",
         "obliteratus/persistence_contracts.py",
@@ -44,7 +58,8 @@
         "obliteratus/runner.py",
         "obliteratus/sweep.py",
         "obliteratus/tourney.py",
-        "obliteratus/tourney_contracts.py"
+        "obliteratus/tourney_contracts.py",
+        "obliteratus/restore_multimodal.py"
       ],
       "required_tests": [
         "tests/test_abliterate.py",
@@ -66,7 +81,12 @@
       "id": "configuration-and-selection",
       "owner": "runtime compatibility maintainers",
       "description": "Configuration, defaults, model profiles, architecture detection, device selection, loading, and preset contracts",
-      "contract_types": ["configuration", "architecture-selection", "device-boundary", "model-runtime"],
+      "contract_types": [
+        "configuration",
+        "architecture-selection",
+        "device-boundary",
+        "model-runtime"
+      ],
       "paths": [
         "obliteratus/adaptive_defaults.py",
         "obliteratus/architecture_profiles.py",
@@ -99,7 +119,12 @@
       "id": "public-interface-and-services",
       "owner": "operator interface maintainers",
       "description": "CLI and interactive interfaces, remote execution, service adapters, scheduling, and UI boundaries",
-      "contract_types": ["public-interface", "external-service", "remote-execution", "operator-ui"],
+      "contract_types": [
+        "public-interface",
+        "external-service",
+        "remote-execution",
+        "operator-ui"
+      ],
       "paths": [
         "obliteratus/bestiary_sync.py",
         "obliteratus/cli.py",
@@ -133,7 +158,12 @@
       "id": "research-inputs-and-provenance",
       "owner": "research integrity maintainers",
       "description": "Prompt sources, hard negatives, community records, reproducibility controls, and telemetry schemas",
-      "contract_types": ["research-input", "research-output", "reproducibility", "persistence"],
+      "contract_types": [
+        "research-input",
+        "research-output",
+        "reproducibility",
+        "persistence"
+      ],
       "paths": [
         "obliteratus/community.py",
         "obliteratus/hard_negative.py",
@@ -154,7 +184,12 @@
       "id": "analysis-methods",
       "owner": "mechanistic analysis maintainers",
       "description": "Mechanistic analysis algorithms, numerical invariants, transfer methods, and visualization contracts",
-      "contract_types": ["research-metric", "numerical-invariant", "research-output", "model-runtime"],
+      "contract_types": [
+        "research-metric",
+        "numerical-invariant",
+        "research-output",
+        "model-runtime"
+      ],
       "paths": [
         "obliteratus/analysis/activation_patching.py",
         "obliteratus/analysis/activation_probing.py",
@@ -205,7 +240,12 @@
       "id": "evaluation-and-benchmarks",
       "owner": "evaluation maintainers",
       "description": "Evaluation routing, benchmark execution, plots, statistical metrics, and external adapter results",
-      "contract_types": ["research-metric", "numerical-invariant", "research-output", "external-service"],
+      "contract_types": [
+        "research-metric",
+        "numerical-invariant",
+        "research-output",
+        "external-service"
+      ],
       "paths": [
         "obliteratus/evaluation/advanced_metrics.py",
         "obliteratus/evaluation/baselines.py",
@@ -233,8 +273,14 @@
       "id": "reporting",
       "owner": "research reporting maintainers",
       "description": "Report schemas, deterministic serialization, plotting, and persisted result artifacts",
-      "contract_types": ["research-output", "persistence", "numerical-invariant"],
-      "paths": ["obliteratus/reporting/report.py"],
+      "contract_types": [
+        "research-output",
+        "persistence",
+        "numerical-invariant"
+      ],
+      "paths": [
+        "obliteratus/reporting/report.py"
+      ],
       "required_tests": [
         "tests/test_report.py",
         "tests/test_evaluation_reporting_contracts.py",
@@ -245,7 +291,11 @@
       "id": "ablation-strategies",
       "owner": "ablation strategy maintainers",
       "description": "Strategy registration, enumeration, architecture navigation, mutation, and restoration behavior",
-      "contract_types": ["ablation-strategy", "architecture-selection", "model-mutation"],
+      "contract_types": [
+        "ablation-strategy",
+        "architecture-selection",
+        "model-mutation"
+      ],
       "paths": [
         "obliteratus/strategies/base.py",
         "obliteratus/strategies/embedding_ablation.py",
@@ -268,28 +318,39 @@
       "path": "obliteratus/cli.py",
       "risk_class": "cpu-contract",
       "risk": "public parsing, validation, dispatch, and local/remote option propagation",
-      "required_tests": ["tests/test_cli.py", "tests/test_cli_boundaries.py"],
+      "required_tests": [
+        "tests/test_cli.py",
+        "tests/test_cli_boundaries.py"
+      ],
       "conditional_gates": []
     },
     {
       "path": "obliteratus/config.py",
       "risk_class": "cpu-contract",
       "risk": "configuration defaults, normalization, validation, and serialization",
-      "required_tests": ["tests/test_config.py", "tests/test_config_properties.py"],
+      "required_tests": [
+        "tests/test_config.py",
+        "tests/test_config_properties.py"
+      ],
       "conditional_gates": []
     },
     {
       "path": "obliteratus/architecture_profiles.py",
       "risk_class": "cpu-contract",
       "risk": "architecture detection and projection-path contracts",
-      "required_tests": ["tests/test_architecture_profiles.py", "tests/test_gemma4_support.py"],
+      "required_tests": [
+        "tests/test_architecture_profiles.py",
+        "tests/test_gemma4_support.py"
+      ],
       "conditional_gates": []
     },
     {
       "path": "obliteratus/community.py",
       "risk_class": "cpu-contract",
       "risk": "contribution schema, aggregation, and atomic persistence",
-      "required_tests": ["tests/test_community.py"],
+      "required_tests": [
+        "tests/test_community.py"
+      ],
       "conditional_gates": []
     },
     {
@@ -317,7 +378,10 @@
       "path": "obliteratus/evaluation/metrics.py",
       "risk_class": "cpu-contract",
       "risk": "research metric semantics and numerical invariants",
-      "required_tests": ["tests/test_metrics.py", "tests/test_property_contracts.py"],
+      "required_tests": [
+        "tests/test_metrics.py",
+        "tests/test_property_contracts.py"
+      ],
       "conditional_gates": []
     },
     {
@@ -335,36 +399,57 @@
       "path": "obliteratus/device.py",
       "risk_class": "mixed-runtime",
       "risk": "device and dtype selection across CPU, CUDA, and Apple backends",
-      "required_tests": ["tests/test_device_boundaries.py"],
-      "conditional_gates": ["cuda-runtime", "mps-runtime"]
+      "required_tests": [
+        "tests/test_device_boundaries.py"
+      ],
+      "conditional_gates": [
+        "cuda-runtime",
+        "mps-runtime"
+      ]
     },
     {
       "path": "obliteratus/runtime_contracts.py",
       "risk_class": "cpu-contract",
       "risk": "model request, memory, architecture scale, and dtype capability decisions",
-      "required_tests": ["tests/test_runtime_contracts.py"],
+      "required_tests": [
+        "tests/test_runtime_contracts.py"
+      ],
       "conditional_gates": []
     },
     {
       "path": "obliteratus/models/loader.py",
       "risk_class": "mixed-runtime",
       "risk": "model loading, cache, architecture, quantization, and device-map boundaries",
-      "required_tests": ["tests/test_loader_boundaries.py", "tests/test_offline_integration.py"],
-      "conditional_gates": ["cuda-runtime", "bitsandbytes-runtime"]
+      "required_tests": [
+        "tests/test_loader_boundaries.py",
+        "tests/test_offline_integration.py"
+      ],
+      "conditional_gates": [
+        "cuda-runtime",
+        "bitsandbytes-runtime"
+      ]
     },
     {
       "path": "obliteratus/models/quant_dequant.py",
       "risk_class": "mixed-runtime",
       "risk": "FP8 and NVFP4 format detection, numerical dequantization, Hub revision integrity, and checkpoint materialization",
-      "required_tests": ["tests/test_quant_dequant.py"],
-      "conditional_gates": ["cuda-runtime"]
+      "required_tests": [
+        "tests/test_quant_dequant.py"
+      ],
+      "conditional_gates": [
+        "cuda-runtime"
+      ]
     },
     {
       "path": "obliteratus/mlx_backend.py",
       "risk_class": "mixed-runtime",
       "risk": "MLX model discovery, tensor placement, mutation, and persistence",
-      "required_tests": ["tests/test_mlx_backend_boundaries.py"],
-      "conditional_gates": ["mlx-runtime"]
+      "required_tests": [
+        "tests/test_mlx_backend_boundaries.py"
+      ],
+      "conditional_gates": [
+        "mlx-runtime"
+      ]
     },
     {
       "path": "obliteratus/persistence_contracts.py",
@@ -390,14 +475,21 @@
         "tests/test_projection_math_contracts.py",
         "tests/test_offline_integration.py"
       ],
-      "conditional_gates": ["model-download-runtime"]
+      "conditional_gates": [
+        "model-download-runtime"
+      ]
     },
     {
       "path": "obliteratus/auto_obliterate.py",
       "risk_class": "mixed-runtime",
       "risk": "automated search state, retry, scoring, and checkpoint behavior",
-      "required_tests": ["tests/test_auto_obliterate.py", "tests/test_remaining_cpu_contracts.py"],
-      "conditional_gates": ["model-download-runtime"]
+      "required_tests": [
+        "tests/test_auto_obliterate.py",
+        "tests/test_remaining_cpu_contracts.py"
+      ],
+      "conditional_gates": [
+        "model-download-runtime"
+      ]
     },
     {
       "path": "obliteratus/bayesian_optimizer.py",
@@ -408,7 +500,9 @@
         "tests/test_module_imports.py",
         "tests/conditional/test_model_download_runtime.py"
       ],
-      "conditional_gates": ["model-download-runtime"]
+      "conditional_gates": [
+        "model-download-runtime"
+      ]
     },
     {
       "path": "obliteratus/evaluation/baselines.py",
@@ -419,7 +513,9 @@
         "tests/test_evaluation_reporting_contracts.py",
         "tests/conditional/test_model_download_runtime.py"
       ],
-      "conditional_gates": ["model-download-runtime"]
+      "conditional_gates": [
+        "model-download-runtime"
+      ]
     },
     {
       "path": "obliteratus/evaluation/evaluator.py",
@@ -430,7 +526,9 @@
         "tests/test_evaluation_reporting_contracts.py",
         "tests/test_offline_integration.py"
       ],
-      "conditional_gates": ["model-download-runtime"]
+      "conditional_gates": [
+        "model-download-runtime"
+      ]
     },
     {
       "path": "obliteratus/informed_pipeline.py",
@@ -441,14 +539,21 @@
         "tests/test_informed_pipeline_contracts.py",
         "tests/test_offline_integration.py"
       ],
-      "conditional_gates": ["model-download-runtime"]
+      "conditional_gates": [
+        "model-download-runtime"
+      ]
     },
     {
       "path": "obliteratus/lora_ablation.py",
       "risk_class": "conditional-runtime",
       "risk": "optional adapter construction and validation against live projections",
-      "required_tests": ["tests/test_module_imports.py", "tests/conditional/test_model_download_runtime.py"],
-      "conditional_gates": ["model-download-runtime"]
+      "required_tests": [
+        "tests/test_module_imports.py",
+        "tests/conditional/test_model_download_runtime.py"
+      ],
+      "conditional_gates": [
+        "model-download-runtime"
+      ]
     },
     {
       "path": "obliteratus/sweep.py",
@@ -459,14 +564,21 @@
         "tests/test_sweep_contracts.py",
         "tests/conditional/test_model_download_runtime.py"
       ],
-      "conditional_gates": ["model-download-runtime"]
+      "conditional_gates": [
+        "model-download-runtime"
+      ]
     },
     {
       "path": "obliteratus/evaluation/heretic_eval.py",
       "risk_class": "mixed-runtime",
       "risk": "publication evaluation prompts, classifiers, scoring, and external adapters",
-      "required_tests": ["tests/test_heretic_eval.py", "tests/conditional/test_external_evaluation_runtime.py"],
-      "conditional_gates": ["external-evaluation"]
+      "required_tests": [
+        "tests/test_heretic_eval.py",
+        "tests/conditional/test_external_evaluation_runtime.py"
+      ],
+      "conditional_gates": [
+        "external-evaluation"
+      ]
     },
     {
       "path": "obliteratus/evaluation/lm_eval_integration.py",
@@ -477,56 +589,93 @@
         "tests/test_module_imports.py",
         "tests/conditional/test_external_evaluation_runtime.py"
       ],
-      "conditional_gates": ["external-evaluation"]
+      "conditional_gates": [
+        "external-evaluation"
+      ]
     },
     {
       "path": "obliteratus/tourney.py",
       "risk_class": "mixed-runtime",
       "risk": "multi-model tournament mutation, comparison, and optional publication",
-      "required_tests": ["tests/test_tourney_contracts.py", "tests/test_module_imports.py", "tests/conditional/test_external_evaluation_runtime.py"],
-      "conditional_gates": ["external-evaluation"]
+      "required_tests": [
+        "tests/test_tourney_contracts.py",
+        "tests/test_module_imports.py",
+        "tests/conditional/test_external_evaluation_runtime.py"
+      ],
+      "conditional_gates": [
+        "external-evaluation"
+      ]
     },
     {
       "path": "obliteratus/tourney_contracts.py",
       "risk_class": "cpu-contract",
       "risk": "checkpoint document parsing and version validation",
-      "required_tests": ["tests/test_tourney_contracts.py"],
+      "required_tests": [
+        "tests/test_tourney_contracts.py"
+      ],
       "conditional_gates": []
     },
     {
       "path": "obliteratus/bestiary_sync.py",
       "risk_class": "conditional-runtime",
       "risk": "external catalog synchronization and malformed service responses",
-      "required_tests": ["tests/test_bestiary_contracts.py", "tests/test_module_imports.py", "tests/conditional/test_network_services.py"],
-      "conditional_gates": ["network-services"]
+      "required_tests": [
+        "tests/test_bestiary_contracts.py",
+        "tests/test_module_imports.py",
+        "tests/conditional/test_network_services.py"
+      ],
+      "conditional_gates": [
+        "network-services"
+      ]
     },
     {
       "path": "obliteratus/models_client.py",
       "risk_class": "conditional-runtime",
       "risk": "catalog resolution across operator files and network services",
-      "required_tests": ["tests/test_models_client_contracts.py", "tests/test_module_imports.py", "tests/conditional/test_network_services.py"],
-      "conditional_gates": ["network-services"]
+      "required_tests": [
+        "tests/test_models_client_contracts.py",
+        "tests/test_module_imports.py",
+        "tests/conditional/test_network_services.py"
+      ],
+      "conditional_gates": [
+        "network-services"
+      ]
     },
     {
       "path": "obliteratus/watchtower.py",
       "risk_class": "conditional-runtime",
       "risk": "scheduled scans, queue state, and live service responses",
-      "required_tests": ["tests/test_watchtower_contracts.py", "tests/test_module_imports.py", "tests/conditional/test_network_services.py"],
-      "conditional_gates": ["network-services"]
+      "required_tests": [
+        "tests/test_watchtower_contracts.py",
+        "tests/test_module_imports.py",
+        "tests/conditional/test_network_services.py"
+      ],
+      "conditional_gates": [
+        "network-services"
+      ]
     },
     {
       "path": "obliteratus/service_contracts.py",
       "risk_class": "cpu-contract",
       "risk": "catalog validation, external record normalization, and scheduler lifecycle decisions",
-      "required_tests": ["tests/test_models_client_contracts.py", "tests/test_watchtower_contracts.py"],
+      "required_tests": [
+        "tests/test_models_client_contracts.py",
+        "tests/test_watchtower_contracts.py"
+      ],
       "conditional_gates": []
     },
     {
       "path": "obliteratus/interactive.py",
       "risk_class": "mixed-runtime",
       "risk": "interactive terminal prompts and operator decision flow",
-      "required_tests": ["tests/test_cli.py", "tests/test_interactive_contracts.py", "tests/conditional/test_operator_ui.py"],
-      "conditional_gates": ["operator-ui"]
+      "required_tests": [
+        "tests/test_cli.py",
+        "tests/test_interactive_contracts.py",
+        "tests/conditional/test_operator_ui.py"
+      ],
+      "conditional_gates": [
+        "operator-ui"
+      ]
     },
     {
       "path": "obliteratus/local_ui.py",
@@ -538,28 +687,54 @@
         "tests/test_local_ui_portability.py",
         "tests/conditional/test_operator_ui.py"
       ],
-      "conditional_gates": ["operator-ui"]
+      "conditional_gates": [
+        "operator-ui"
+      ]
     },
     {
       "path": "obliteratus/ui_watchtower.py",
       "risk_class": "mixed-runtime",
       "risk": "service-backed UI tabs and scheduler controls",
-      "required_tests": ["tests/test_ui_watchtower_contracts.py", "tests/test_module_imports.py", "tests/conditional/test_operator_ui.py"],
-      "conditional_gates": ["operator-ui"]
+      "required_tests": [
+        "tests/test_ui_watchtower_contracts.py",
+        "tests/test_module_imports.py",
+        "tests/conditional/test_operator_ui.py"
+      ],
+      "conditional_gates": [
+        "operator-ui"
+      ]
     },
     {
       "path": "obliteratus/remote.py",
       "risk_class": "mixed-runtime",
       "risk": "SSH discovery, quoting, execution, cancellation, and result synchronization",
-      "required_tests": ["tests/test_remote_boundaries.py", "tests/conditional/test_remote_runtime.py"],
-      "conditional_gates": ["remote-execution"]
+      "required_tests": [
+        "tests/test_remote_boundaries.py",
+        "tests/conditional/test_remote_runtime.py"
+      ],
+      "conditional_gates": [
+        "remote-execution"
+      ]
     },
     {
       "path": "obliteratus/remote_contracts.py",
       "risk_class": "cpu-contract",
       "risk": "remote target validation, GPU selection normalization, and shell-safe command construction",
-      "required_tests": ["tests/test_remote_contracts.py", "tests/test_remote_boundaries.py"],
+      "required_tests": [
+        "tests/test_remote_contracts.py",
+        "tests/test_remote_boundaries.py"
+      ],
       "conditional_gates": []
+    },
+    {
+      "path": "obliteratus/restore_multimodal.py",
+      "risk_class": "cpu-contract",
+      "risk": "merge abliterated text weights back into full multimodal model",
+      "required_tests": [
+        "tests/test_restore_multimodal.py"
+      ],
+      "conditional_gates": [],
+      "contract_owner": "OBLITERATUS maintainers"
     }
   ]
 }

obliteratus/cli.py

@@ -413,6 +413,18 @@ def main(argv: list[str] | None = None):
         "--dir", type=str, default="community_results",
         help="Directory containing contribution JSON files",
     )
+
+    # --- restore-multimodal ---
+    restore_parser = subparsers.add_parser(
+        "restore-multimodal",
+        help="Restore vision/MTP/linear-attn components after text-only abliteration",
+    )
+    restore_parser.add_argument("--abliterated", required=True,
+                                help="Text-only abliterated model directory")
+    restore_parser.add_argument("--stock", required=True,
+                                help="Stock (full) model directory or HF repo")
+    restore_parser.add_argument("--output", required=True,
+                                help="Output directory for merged model")
     aggregate_parser.add_argument(
         "--format",
         choices=["summary", "latex"],
@@ -541,6 +553,12 @@ def main(argv: list[str] | None = None):
         _cmd_report(args)
     elif args.command == "aggregate":
         _cmd_aggregate(args)
+    elif args.command == "restore-multimodal":
+        from obliteratus.restore_multimodal import restore_multimodal
+        import logging
+        logging.basicConfig(level=logging.INFO, format="%(message)s")
+        result = restore_multimodal(args.abliterated, args.stock, args.output)
+        print(f"\nDone: {result['replaced']} abliterated + {result['kept']} stock = {result['total']} total")
     elif args.command == "ui":
         _cmd_ui(args)
     elif args.command == "recommend":

obliteratus/restore_multimodal.py

@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+"""Restore multimodal components (vision tower, MTP head, etc.) after text-only abliteration.
+
+OBLITERATUS surgery strips models to text-only weights with simplified tensor
+names (model.layers.* instead of model.language_model.layers.*). This works
+for text inference but breaks GGUF conversion (missing MTP block) and drops
+the vision tower entirely.
+
+This tool merges abliterated text weights back into the full stock model,
+restoring all non-text components from stock while keeping the abliterated
+text weights intact. The merge is lossless — no blending or interpolation.
+
+Usage:
+    obliteratus restore-multimodal \\
+        --abliterated outputs/my-abliterated-model \\
+        --stock Qwen/Qwen3.8-27B \\
+        --output outputs/my-abliterated-model-full
+
+Typical tensor counts for Qwen3.8-27B:
+    Abliterated (text-only):  851 tensors (model.layers.*)
+    Stock (full multimodal): 1199 tensors (model.language_model.*, mtp.*, visual.*)
+    Merged output:           1199 tensors (abliterated text + stock vision/MTP)
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import shutil
+from pathlib import Path
+
+from safetensors.torch import load_file, save_file
+
+logger = logging.getLogger(__name__)
+
+# Common prefix mappings: text-only key -> full multimodal key
+_PREFIX_STRIPS = [
+    ("model.language_model.", "model."),
+]
+
+
+def _normalize_key(key: str) -> str:
+    """Strip multimodal prefixes to match text-only tensor names."""
+    for full_prefix, stripped_prefix in _PREFIX_STRIPS:
+        if key.startswith(full_prefix):
+            return stripped_prefix + key[len(full_prefix):]
+    return key
+
+
+def _build_abliterated_lookup(abl_dir: Path) -> dict:
+    """Build a lookup from tensor key -> (shard_file, original_key)."""
+    index_path = abl_dir / "model.safetensors.index.json"
+    if not index_path.exists():
+        # Single-file model
+        sf_files = list(abl_dir.glob("*.safetensors"))
+        if len(sf_files) == 1:
+            tensors = load_file(str(sf_files[0]))
+            return {k: (sf_files[0].name, k) for k in tensors}
+        raise FileNotFoundError(f"No safetensors index or single file in {abl_dir}")
+
+    with open(index_path) as f:
+        index = json.load(f)
+    return {key: (shard, key) for key, shard in index["weight_map"].items()}
+
+
+def restore_multimodal(
+    abliterated_dir: str | Path,
+    stock_dir: str | Path,
+    output_dir: str | Path,
+    skip_existing: bool = True,
+) -> dict:
+    """Merge abliterated text weights into full stock model.
+
+    Args:
+        abliterated_dir: Path to text-only abliterated model.
+        stock_dir: Path to stock (full multimodal) model.
+        output_dir: Path to write merged output.
+        skip_existing: Skip files that already exist in output_dir.
+
+    Returns:
+        dict with counts: replaced, kept, total.
+    """
+    abl_dir = Path(abliterated_dir)
+    stock_dir = Path(stock_dir)
+    out_dir = Path(output_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # 1. Copy all non-safetensor files from stock
+    skip_names = {".git", "__pycache__", "crc32.txt", "README.md", ".gitattributes"}
+    for f in stock_dir.iterdir():
+        if f.name.endswith(".safetensors") or f.name == "model.safetensors.index.json":
+            continue
+        if f.name in skip_names:
+            continue
+        dst = out_dir / f.name
+        if not dst.exists() or not skip_existing:
+            shutil.copy2(f, dst)
+            logger.info("Copied %s", f.name)
+
+    # Copy abliteration metadata if present
+    for extra in ["abliteration_metadata.json", "hard_negative_residue.json"]:
+        src = abl_dir / extra
+        if src.exists():
+            shutil.copy2(src, out_dir / extra)
+            logger.info("Copied %s", extra)
+
+    # 2. Build abliterated tensor lookup
+    abl_keys = _build_abliterated_lookup(abl_dir)
+    logger.info("Abliterated model: %d tensors", len(abl_keys))
+
+    # 3. Load stock index
+    with open(stock_dir / "model.safetensors.index.json") as f:
+        stock_index = json.load(f)
+
+    stock_shards = sorted(set(stock_index["weight_map"].values()))
+    logger.info("Stock model: %d tensors across %d shards",
+                len(stock_index["weight_map"]), len(stock_shards))
+
+    # 4. Process each stock shard — replace text tensors with abliterated versions
+    abl_shard_cache = {}
+    replaced = 0
+    kept = 0
+    new_weight_map = {}
+
+    for shard_idx, stock_shard in enumerate(stock_shards):
+        logger.info("[%d/%d] %s", shard_idx + 1, len(stock_shards), stock_shard)
+
+        stock_tensors = load_file(str(stock_dir / stock_shard))
+        merged = {}
+
+        for key, tensor in stock_tensors.items():
+            norm_key = _normalize_key(key)
+
+            # Try normalized key first, then direct match
+            match_key = norm_key if norm_key in abl_keys else (key if key in abl_keys else None)
+
+            if match_key is not None:
+                abl_shard, abl_key = abl_keys[match_key]
+                if abl_shard not in abl_shard_cache:
+                    abl_shard_cache[abl_shard] = load_file(str(abl_dir / abl_shard))
+                merged[key] = abl_shard_cache[abl_shard][abl_key]
+                replaced += 1
+            else:
+                merged[key] = tensor
+                kept += 1
+
+            new_weight_map[key] = stock_shard
+
+        save_file(merged, str(out_dir / stock_shard))
+        logger.info("  Saved %s (%d tensors)", stock_shard, len(merged))
+
+        # Free cache periodically
+        if len(abl_shard_cache) > 5:
+            abl_shard_cache.clear()
+
+    # 5. Write merged index
+    new_index = {
+        "metadata": stock_index.get("metadata", {}),
+        "weight_map": new_weight_map,
+    }
+    with open(out_dir / "model.safetensors.index.json", "w") as f:
+        json.dump(new_index, f, indent=2)
+
+    result = {"replaced": replaced, "kept": kept, "total": replaced + kept}
+    logger.info("Replaced: %d tensors (abliterated)", replaced)
+    logger.info("Kept:     %d tensors (from stock)", kept)
+    logger.info("Output:   %s", out_dir)
+    return result
+
+
+def main():
+    import argparse
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
+
+    p = argparse.ArgumentParser(
+        description="Restore vision tower, MTP head, and other multimodal "
+                    "components after text-only abliteration surgery."
+    )
+    p.add_argument("--abliterated", required=True,
+                   help="Path to text-only abliterated model directory")
+    p.add_argument("--stock", required=True,
+                   help="Path to stock (full) model directory or HF repo ID")
+    p.add_argument("--output", required=True,
+                   help="Output directory for merged model")
+    args = p.parse_args()
+
+    result = restore_multimodal(args.abliterated, args.stock, args.output)
+    print(f"\nDone: {result['replaced']} abliterated + {result['kept']} stock = {result['total']} total")
+
+
+if __name__ == "__main__":
+    main()

tests/test_restore_multimodal.py

@@ -0,0 +1,214 @@
+"""Tests for obliteratus.restore_multimodal."""
+
+import json
+from pathlib import Path
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+
+def _make_fake_model(tmpdir: Path, prefix: str, n_layers: int = 2, extra_tensors: dict | None = None):
+    """Create a minimal fake safetensors model directory."""
+    tensors = {}
+    for i in range(n_layers):
+        tensors[f"{prefix}layers.{i}.weight"] = torch.randn(4, 4)
+    if extra_tensors:
+        for k, v in extra_tensors.items():
+            tensors[k] = v
+
+    shard = "model-00001-of-00001.safetensors"
+    save_file(tensors, str(tmpdir / shard))
+
+    index = {
+        "metadata": {"total_size": 0},
+        "weight_map": {k: shard for k in tensors},
+    }
+    (tmpdir / "model.safetensors.index.json").write_text(json.dumps(index))
+    (tmpdir / "config.json").write_text(json.dumps({"model_type": "test"}))
+    return tensors
+
+
+class TestNormalizeKey:
+    def test_strips_language_model_prefix(self):
+        from obliteratus.restore_multimodal import _normalize_key
+
+        assert _normalize_key("model.language_model.layers.0.weight") == "model.layers.0.weight"
+
+    def test_preserves_non_prefixed(self):
+        from obliteratus.restore_multimodal import _normalize_key
+
+        assert _normalize_key("model.layers.0.weight") == "model.layers.0.weight"
+
+    def test_preserves_visual(self):
+        from obliteratus.restore_multimodal import _normalize_key
+
+        assert _normalize_key("visual.conv.weight") == "visual.conv.weight"
+
+
+class TestBuildAbliteratedLookup:
+    def test_reads_index(self, tmp_path):
+        from obliteratus.restore_multimodal import _build_abliterated_lookup
+
+        _make_fake_model(tmp_path, "model.")
+        lookup = _build_abliterated_lookup(tmp_path)
+        assert "model.layers.0.weight" in lookup
+        assert "model.layers.1.weight" in lookup
+
+    def test_single_file_fallback(self, tmp_path):
+        from obliteratus.restore_multimodal import _build_abliterated_lookup
+
+        tensors = {"model.layers.0.weight": torch.randn(4, 4)}
+        save_file(tensors, str(tmp_path / "model.safetensors"))
+        lookup = _build_abliterated_lookup(tmp_path)
+        assert "model.layers.0.weight" in lookup
+
+    def test_missing_raises(self, tmp_path):
+        from obliteratus.restore_multimodal import _build_abliterated_lookup
+
+        with pytest.raises(FileNotFoundError):
+            _build_abliterated_lookup(tmp_path)
+
+
+class TestRestoreMultimodal:
+    def test_merges_text_and_visual(self, tmp_path):
+        from obliteratus.restore_multimodal import restore_multimodal
+
+        abl_dir = tmp_path / "abliterated"
+        abl_dir.mkdir()
+        stock_dir = tmp_path / "stock"
+        stock_dir.mkdir()
+        out_dir = tmp_path / "output"
+
+        # Abliterated: text-only with model.layers.*
+        _make_fake_model(abl_dir, "model.", n_layers=2)
+
+        # Stock: multimodal with model.language_model.layers.* + visual.*
+        _make_fake_model(
+            stock_dir,
+            "model.language_model.",
+            n_layers=2,
+            extra_tensors={"visual.conv.weight": torch.randn(4, 4)},
+        )
+
+        result = restore_multimodal(str(abl_dir), str(stock_dir), str(out_dir))
+
+        assert result["replaced"] == 2  # 2 text layers replaced
+        assert result["kept"] == 1  # visual.conv.weight kept from stock
+        assert result["total"] == 3
+
+        # Verify output has all tensors
+        output_index = json.loads((out_dir / "model.safetensors.index.json").read_text())
+        assert "model.language_model.layers.0.weight" in output_index["weight_map"]
+        assert "visual.conv.weight" in output_index["weight_map"]
+
+    def test_copies_config_files(self, tmp_path):
+        from obliteratus.restore_multimodal import restore_multimodal
+
+        abl_dir = tmp_path / "abliterated"
+        abl_dir.mkdir()
+        stock_dir = tmp_path / "stock"
+        stock_dir.mkdir()
+        out_dir = tmp_path / "output"
+
+        _make_fake_model(abl_dir, "model.", n_layers=1)
+        _make_fake_model(stock_dir, "model.language_model.", n_layers=1)
+
+        # Add extra files including ones that should be skipped
+        (stock_dir / "tokenizer.json").write_text("{}")
+        (stock_dir / "README.md").write_text("skip me")
+        (stock_dir / ".gitattributes").write_text("skip me too")
+        (abl_dir / "abliteration_metadata.json").write_text("{}")
+
+        restore_multimodal(str(abl_dir), str(stock_dir), str(out_dir))
+
+        assert (out_dir / "tokenizer.json").exists()
+        assert (out_dir / "abliteration_metadata.json").exists()
+        assert (out_dir / "config.json").exists()
+        # Skipped files should NOT be copied
+        assert not (out_dir / "README.md").exists()
+        assert not (out_dir / ".gitattributes").exists()
+
+    def test_abliterated_weights_used(self, tmp_path):
+        """Verify the abliterated weights actually replace stock, not just copy stock."""
+        from safetensors.torch import load_file
+
+        from obliteratus.restore_multimodal import restore_multimodal
+
+        abl_dir = tmp_path / "abliterated"
+        abl_dir.mkdir()
+        stock_dir = tmp_path / "stock"
+        stock_dir.mkdir()
+        out_dir = tmp_path / "output"
+
+        # Create distinguishable tensors
+        abl_weight = torch.ones(4, 4) * 42.0
+        stock_weight = torch.zeros(4, 4)
+
+        abl_tensors = {"model.layers.0.weight": abl_weight}
+        save_file(abl_tensors, str(abl_dir / "model-00001-of-00001.safetensors"))
+        (abl_dir / "model.safetensors.index.json").write_text(
+            json.dumps({"metadata": {}, "weight_map": {"model.layers.0.weight": "model-00001-of-00001.safetensors"}})
+        )
+
+        stock_tensors = {"model.language_model.layers.0.weight": stock_weight}
+        save_file(stock_tensors, str(stock_dir / "model-00001-of-00001.safetensors"))
+        (stock_dir / "model.safetensors.index.json").write_text(
+            json.dumps(
+                {
+                    "metadata": {},
+                    "weight_map": {"model.language_model.layers.0.weight": "model-00001-of-00001.safetensors"},
+                }
+            )
+        )
+        (stock_dir / "config.json").write_text("{}")
+
+        restore_multimodal(str(abl_dir), str(stock_dir), str(out_dir))
+
+        merged = load_file(str(out_dir / "model-00001-of-00001.safetensors"))
+        # Should have the abliterated value (42), not stock (0)
+        assert torch.allclose(merged["model.language_model.layers.0.weight"], abl_weight)
+
+
+class TestMain:
+    def test_cli_main(self, tmp_path):
+        from obliteratus.restore_multimodal import main
+
+        abl_dir = tmp_path / "abl"
+        abl_dir.mkdir()
+        stock_dir = tmp_path / "stock"
+        stock_dir.mkdir()
+        out_dir = tmp_path / "out"
+
+        _make_fake_model(abl_dir, "model.", n_layers=1)
+        _make_fake_model(stock_dir, "model.language_model.", n_layers=1)
+
+        import sys
+        old_argv = sys.argv
+        sys.argv = ["prog", "--abliterated", str(abl_dir), "--stock", str(stock_dir), "--output", str(out_dir)]
+        try:
+            main()
+        finally:
+            sys.argv = old_argv
+
+        assert (out_dir / "model.safetensors.index.json").exists()
+
+    def test_cli_dispatch(self, tmp_path):
+        """Test that 'obliteratus restore-multimodal' dispatches correctly."""
+        from obliteratus.cli import main as cli_main
+
+        abl_dir = tmp_path / "abl"
+        abl_dir.mkdir()
+        stock_dir = tmp_path / "stock"
+        stock_dir.mkdir()
+        out_dir = tmp_path / "out"
+
+        _make_fake_model(abl_dir, "model.", n_layers=1)
+        _make_fake_model(stock_dir, "model.language_model.", n_layers=1)
+
+        cli_main(["restore-multimodal",
+                   "--abliterated", str(abl_dir),
+                   "--stock", str(stock_dir),
+                   "--output", str(out_dir)])
+
+        assert (out_dir / "model.safetensors.index.json").exists()