#!/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()