# ModelOsmosis

**Universal Model Knowledge Transfer Framework** - absorb, merge, stack, and bridge ML models across frameworks, then package them as trusted runnable artifacts.

ModelOsmosis combines two sides into one product:
- **Model operation engine**: absorb, merge, weight-merge, stack, cross-merge, compare, benchmark
- **Artifact/runtime engine**: build trusted `.osmoimg` containers with compression, trust signing, verification, policy checks, and runtime launch

---

## Status

| | |
|---|---|
| Version | `0.8.4` |
| Maturity | Pre-1.0 (alpha) |
| Python | `>=3.9` (3.9-3.13) |
| License | MIT |

## Install

```bash
pip install model-osmosis
```

Optional extras:

```bash
pip install "model-osmosis[huggingface]"   # transformers, huggingface-hub, accelerate
pip install "model-osmosis[vision]"        # timm, torchvision, pillow
pip install "model-osmosis[audio]"         # librosa, soundfile
pip install "model-osmosis[diffusion]"     # diffusers, accelerate
pip install "model-osmosis[multimodal]"    # multimodal helpers
pip install "model-osmosis[lora]"          # peft
pip install "model-osmosis[onnx]"          # onnx, onnxruntime
pip install "model-osmosis[mlx]"           # mlx, mlx-lm (Apple Silicon)
pip install "model-osmosis[tensorflow]"    # tensorflow
pip install "model-osmosis[jax]"           # jax, jaxlib, flax
pip install "model-osmosis[compression]"   # zstandard
pip install "model-osmosis[trust]"         # trust/signing deps
pip install "model-osmosis[demo]"          # gradio demo app
pip install "model-osmosis[dev]"           # pytest, ruff, mypy
pip install "model-osmosis[all]"           # broad optional stack
```

Core dependencies: `numpy`, `torch`, `scikit-learn`, `safetensors`, `cryptography`.

---

## Run This One Command First

```bash
osmosis demo
```

`demo` downloads `HuggingFaceTB/SmolLM2-135M-Instruct`, builds a `.osmoimg`, runs inspect/verify/smoke checks, and prints PASS/FAIL with artifact path, trust status, codec/profile, RAM estimate, and elapsed time.

First run is usually 10-30 minutes (mostly download time).

No invite/license activation is required.

---

## Unified 4-Command Workflow

```bash
# 1) Make model output (merge/weight-merge/absorb/stack/cross-merge)
osmosis make --mode weight-merge --models pt:./a.pt pt:./b.pt --output-type osmoimg -o merged.osmoimg

# 2) Protect/package pass (optional re-sign/recompress hardening)
osmosis protect merged.osmoimg -o merged.protected.osmoimg --compression-profile lossless-max --trust-sign

# 3) Check trust + self-check (+ optional runtime/release gates)
osmosis check merged.protected.osmoimg --self-check-level smoke

# 4) Launch runtime (run/serve/export)
osmosis launch merged.protected.osmoimg --action run --prompt "Hello"
```

---

## CLI Surface

ModelOsmosis currently exposes:
- **38 top-level commands** (`osmosis --help`)
- **40+ leaf commands** when grouped subcommands are expanded (`trust`, `artifact`)

Core top-level commands:

```text
artifact, benchmark, benchmark-runtime, build, check, compare, convert, demo,
diff, doctor, expand, expand-rollback, expand-status, export, history, images,
info, inspect, launch, make, merge, obfuscate, optimize-reset,
optimize-status, pipeline, protect, pull, push, quantize, register,
release-check, run, serve, transcode, trust, validate, verify, version
```

Grouped subcommands:
- `trust`: `keygen`, `list-keys`, `import-signer`, `revoke-signer`
- `artifact`: `track`, `verify`, `hook install`

---

## Model Operation Engine

### Absorb (knowledge transfer)

`absorb` extracts knowledge from a source and injects it into a target, with optional architecture growth and distillation.

Pipeline: **extract -> grow -> distill -> inject**

```python
import numpy as np
import torch.nn as nn
from model_osmosis import OsmosisPipeline, OsmosisConfig

source = nn.Sequential(nn.Linear(8, 64), nn.ReLU(), nn.Linear(64, 4))
target = nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 4))
sample = np.random.rand(200, 8).astype("float32")

pipeline = OsmosisPipeline(OsmosisConfig(verbose=False))
result = pipeline.absorb(source=source, target=target, sample_data=sample, auto_grow=True)
```

### Growth strategies

| Strategy | Description |
|---|---|
| `WIDER` | Add width to existing layers |
| `DEEPER` | Insert new layers |
| `HYBRID` | Combine wider + deeper |

### Merge strategies (7)

| Strategy | Description |
|---|---|
| `weighted_average` | Weighted tensor blending |
| `slerp` | Spherical interpolation |
| `ties` | Conflict-aware sparse merge |
| `dare` | Drop-and-rescale sparse merge |
| `breadcrumbs` | Sparse variant |
| `consensus` | Agreement-based merge |
| `auto` | Smart recipe selection engine |

### Weight-merge mode (pure merge path)

`pipeline --mode weight-merge` merges only exact-compatible tensors and enforces a coverage floor (`--coverage-floor`, default `0.70`) using minimum per-model merged byte coverage.

### Expand workflow (`osmosis expand`)

`expand` is now a **strict expansion-first** flow with deterministic routing and revisioned outputs.

- Default policy: strict expansion required (`--fallback-mode off`).
- Optional degraded path: `--fallback-mode distill` (explicit only).
- Universal accepted inputs:
  - refs: `hf:`, `pt:`, `onnx:`, `tf:`, `jax:`, `vision:`, `audio:`, `diffusion:`, `multimodal:`, `lora:`, `mlx:`
  - artifacts: `.osmo`, `.osmoimg`
- Route metadata is always reported: `route`, `runtime_source`, `strict_mode`, `fallback_mode`, `degraded`.
- Data planner v2 (default `--data-mode auto`) emits task-aware batches and deterministic planner fields.

```bash
# Expand and save trusted versioned artifact
osmosis expand pt:./model.pt --target-params x1.5 --budget-minutes 5

# Expand directly from an existing artifact
osmosis expand models/base.osmoimg --target-params x1.2

# Force cache-only data planning (no network fetch)
osmosis expand hf:your/model --offline --data-mode auto

# Restrict acceptable dataset licenses
osmosis expand pt:./model.pt --allow-license mit --allow-license apache-2.0

# Allow explicit distill fallback when strict structure expansion is unavailable
osmosis expand onnx:./model.onnx --fallback-mode distill

# Inspect local revision chain / active revision
osmosis expand-status --json

# Roll active revision pointer back
osmosis expand-rollback --to model.r0001
```

### Cross-merge

- `make --mode cross-merge` gives explicit cross-merge flow
- `pipeline --mode merge` can auto-route to cross-merge when direct state-dict merge is not compatible

### MoE stack

`stack` composes experts behind a top-k gate with optional gate training and optimization.

---

## Smart/Adaptive Engines (Local-Only)

ModelOsmosis includes deterministic local optimizer loops (no telemetry upload):

- **SmartMergeEngine** (`smart-merge-mlp-v1`)
  - Chooses merge recipe when strategy is `auto`
  - Controlled by `--strategy-engine {auto,off}`
- **CrossMergeEngine** (`cross-merge-mlp-v1`)
  - Adaptive recipe selection for classic and LLM cross-merge paths
  - Stage control: `--cross-opt-stage {auto,shadow,active}`
- **Compression selector** (`compression-selector-v1`)
  - Chooses tuned zstd level profiles for `lossless-max`
  - Controlled by `--compression-engine {auto,off}`

Optimizer state is stored in:

```text
.model_osmosis/self_optimize/state.json
```

Status/reset:

```bash
osmosis optimize-status --json
osmosis optimize-reset
```

---

## Formats

### `.osmo` (universal interchange format)

Portable model representation for conversion/merge operations.

Logical layout:

```text
model.osmo/
├── manifest.json
├── architecture.json
├── parameters/*.safetensors
├── knowledge/*.npy
├── auxiliary/*
└── fingerprint.json
```

### `.osmoimg` (trusted runnable artifact)

Container format with header + payload:

```text
[8 bytes]  magic: OSMOIMG\0
[4 bytes]  format version (1/2/3)
[4 bytes]  header length
[N bytes]  header JSON (compressed)
[...]      payload bytes
```

Current runtime supports reading v1/v2/v3 (`OSMOIMG_VERSION = 3`).

Compression profiles:

| Profile | Payload codec | Typical container version |
|---|---|---|
| `legacy` | `zip-deflate` | v1 |
| `lossless-max` | `adaptive-zstd-v1` | v2/v3 |

`.osmoimg` includes:
- model parameters + architecture + runtime metadata
- tokenizer bundle with manifest (`kind`, `sha256`, `size_bytes`) and validation
- trust envelope metadata/signature fields

Tokenizer safety:
- non-text tokenizer assets are handled as raw bytes
- text models require real tokenizer assets during packaging
- load/transcode validates tokenizer manifest for corruption

---

## Runtime (Run/Serve/Export)

Runtime profiles:
- `default`
- `cpu-lowram`
- `gpu-fast`

Device policy:
- `auto`, `cpu`, `cuda`

Runtime quantization (runtime-only, no artifact mutation):
- `none`, `int8`

Important behavior:
- Text generation returns **new tokens only by default**
- Legacy behavior is available with `--include-prompt`

Examples:

```bash
osmosis run model.osmoimg --prompt "Write a haiku"
osmosis run model.osmoimg --prompt "Hello" --include-prompt
osmosis serve model.osmoimg --port 8080 --warmup
osmosis export model.osmoimg --format pytorch -o model.pt
```

---

## Trust, Provenance, and Policy

Trust uses offline Ed25519 signatures on `.osmoimg` trust statements.

Verification statuses:
- `trusted`
- `signed_untrusted`
- `unsigned`
- `tampered`

Key commands:

```bash
osmosis trust keygen --name default
osmosis trust list-keys
osmosis trust import-signer --name partner --pubkey partner.pub
osmosis trust revoke-signer --name partner

osmosis verify model.osmoimg --json
osmosis verify model.osmoimg --require-trust --repro-check
```

Trust profiles:
- `standard`
- `production`

`production` enforcement (run/serve/export) requires:
- `--policy-file`
- non-empty `allowed_signer_ids` or `allowed_key_ids`
- baseline policy floor (`min_header_version >= 3`, `required_trust_status = trusted`)

Policy keys supported:
- `min_header_version`
- `required_trust_status`
- `allowed_signer_ids`
- `allowed_key_ids`
- `required_provenance_mode`
- `required_repro_mode`
- `require_transparency`

Signer backends:
- `local-keyring`
- `aws-kms` (currently envelope-simulated backend interface)

---

## Access Model

ModelOsmosis is open by default: no invite key, activation code, or machine lock is required.

---

## Team Artifact Tracking (Git)

Track trusted `.osmoimg` files in a deterministic lock file:

```text
.model_osmosis/artifacts.lock.json
```

Commands:

```bash
osmosis artifact track models/prod.osmoimg
osmosis artifact verify --all
osmosis artifact hook install
```

Recommended team flow:
1. Track artifacts.
2. Commit lock file.
3. Enforce `osmosis artifact verify --all` in hooks/CI.
4. Re-track after intentional artifact updates.

---

## Unified Pipeline (`osmosis pipeline`)

`pipeline` modes (required `--mode`):
- `merge`
- `weight-merge`
- `absorb`
- `stack`

`make` adds one extra explicit mode:
- `cross-merge`

Pipeline output types (11):
- `osmoimg`, `osmo`, `pytorch`, `hf`, `onnx`, `tensorflow`, `keras`, `jax`, `flax`, `safetensors`, `torchscript`

Examples:

```bash
# Merge -> trusted osmoimg
osmosis pipeline --mode merge --output-type osmoimg \
  --models hf:HuggingFaceTB/SmolLM2-135M-Instruct pt:./model_b.pt \
  --compression-profile lossless-max --trust-sign -o merged.osmoimg

# Weight-merge -> pytorch with coverage gate
osmosis pipeline --mode weight-merge --output-type pytorch \
  --models pt:./a.pt pt:./b.pt --merge-strategy auto --coverage-floor 0.70 -o merged.pt

# Absorb -> pytorch
osmosis pipeline --mode absorb --output-type pytorch \
  --source pt:./source.pt --target pt:./target.pt -o absorbed.pt

# Stack -> onnx
osmosis pipeline --mode stack --output-type onnx \
  --models pt:./expert_a.pt pt:./expert_b.pt --onnx-input-shape 1,128 -o stacked.onnx
```

Reliability flags:
- `--dry-run`
- `--report-file`
- `--self-check-level {off,smoke,parity}`
- `--allow-partial-load`
- `--allow-lossy-cast`

---

## Release Gates and CI

`release-check` is the aggregate release gate command.

```bash
osmosis release-check --json
```

Default behavior is CI-aware and strict. Use `--local-only` for offline local-only gate runs.

Required CI workflows for release status:
- `clean-machine-validation`
- `trusted-osmoimg`
- `runtime-benchmark`
- `cross-framework-reliability`

Workflow summary:

| Workflow | Scope |
|---|---|
| `clean-machine-validation.yml` | Fresh install + quick demo on Linux/Windows/macOS |
| `trusted-osmoimg.yml` | Trust/sign/repro tests |
| `runtime-benchmark.yml` | CPU PR gate + GPU nightly/dispatch |
| `cross-framework-reliability.yml` | Ubuntu core parity + Windows/macOS PR smoke + TF/JAX nightly |

---

## Runtime Benchmarking

```bash
osmosis benchmark-runtime model.osmoimg --profile i5_2gb --json
```

Profiles in `benchmarks/runtime_baselines.json`:
- `i5_2gb`
- `ci_standard`
- `ci_gpu_standard`

Reported metrics include:
- `cold_start_seconds`
- `first_token_latency_seconds`
- `warm_first_token_latency_seconds`
- `tokens_per_second`
- `p95_latency_seconds`
- `peak_rss_mb`

---

## Python API Quick Reference

```python
from model_osmosis import OsmosisPipeline, OsmosisConfig

pipeline = OsmosisPipeline(OsmosisConfig(verbose=False, seed=42))

# absorb
result = pipeline.absorb(source, target, sample_data, auto_grow=True)

# merge
merged = pipeline.merge([model_a, model_b], strategy="auto")

# stack
stacked = pipeline.stack([expert_a, expert_b, expert_c], top_k=2, train_gate=True)

# compare / benchmark
metrics = pipeline.compare(model_a, model_b, sample_data)
ranked = pipeline.benchmark([model_a, model_b], eval_data=sample_data)
```

`.osmoimg` API:

```python
from model_osmosis.format.image import OsmoImage

image = OsmoImage.build(
    "HuggingFaceTB/SmolLM2-135M-Instruct",
    name="quick-demo",
    compression_profile="lossless-max",
    trust_sign=True,
    provenance_mode="private",
    repro_mode="strict",
)
image.save("quick-demo.osmoimg")

info = OsmoImage.inspect("quick-demo.osmoimg")
verify = OsmoImage.verify("quick-demo.osmoimg", require_trust=False)
out = image.run("Hello")
```

---

## Framework Adapters (10)

- HuggingFace
- Vision
- Audio
- Diffusion
- Multimodal
- LoRA/PEFT
- ONNX
- MLX
- TensorFlow/Keras
- JAX/Flax

Model reference syntax accepted by pipeline/make loaders:
- `hf:repo/id`
- `pt:path`
- `onnx:path`
- `tf:path`
- `jax:path`
- plain paths (auto-detect)

---

## Testing

The repository currently includes **32 test files** under `tests/`.

Run full tests:

```bash
pytest tests -v
```

Important markers:
- `crossfw_core` (PR-blocking reliability)
- `crossfw_extended` (nightly extended reliability)

---

## Example Scripts

`examples/` currently includes:

1. `basic_neural_transfer.py`
2. `build_tiny_i5_chat_image.py`
3. `cross_merge_code_llms.py`
4. `dynamic_growth_demo.py`
5. `merge_1b_model.py`
6. `merge_hebrew_coder.py`
7. `tree_to_neural.py`

---

## Repository Layout

```text
model_osmosis/
  adapters/        # 10 framework adapters
  absorbers/       # target injection
  core/            # config + exceptions
  extractors/      # source extraction
  format/          # .osmo/.osmoimg, trust, serialization, converters
  growers/         # architecture growth
  knowledge/       # universal knowledge representation
  merging/         # merge engines + strategies
  obfuscation/     # function-preserving obfuscation
  optimizer/       # strategy/compression/cross optimizer state
  services/        # workflow and artifact tracking services
  stacking/        # MoE stack components
  utils/           # determinism, helpers, cards
  pipeline.py      # core orchestrator
  cli.py           # CLI entrypoint
examples/
tests/
.github/workflows/
benchmarks/
scripts/
```

---

## Troubleshooting

| Symptom | Fix |
|---|---|
| `Production trust profile requires a policy file` | Add `--policy-file <policy.json>` with signer/key allowlist |
| ONNX export complains about shape | Add `--onnx-input-shape 1,128` |
| Raw export is not compressed/protected as container | Use `--output-type osmoimg` |
| HF download warnings about symlinks on Windows | Enable Developer Mode (warning affects cache efficiency, not correctness) |
| Trust status is `signed_untrusted` | Import signer public key: `osmosis trust import-signer --name <id> --pubkey <file>` |
| `release-check` cannot query CI | Install/login `gh`, or run `osmosis release-check --local-only --json` |

---

## Known Limitations

- `0.8.x` is pre-1.0, so APIs and defaults may evolve.
- Some adapters require optional dependencies and external model downloads.
- Small demo models may be functionally limited for chat quality.
- TensorFlow/JAX extended reliability runs are nightly/dispatch, not core PR gate.

## License

MIT
