Skip to content

feat: add mlx-vlm support for Gemma 4 models - #1

Open
waylonkenning wants to merge 2 commits into
ammaarreshi:mainfrom
waylonkenning:fix/gemma-4-mlx-vlm
Open

feat: add mlx-vlm support for Gemma 4 models#1
waylonkenning wants to merge 2 commits into
ammaarreshi:mainfrom
waylonkenning:fix/gemma-4-mlx-vlm

Conversation

@waylonkenning

@waylonkenning waylonkenning commented Apr 29, 2026 •
edited
Loading

Copy link
Copy Markdown

Issue: Gemma 4 models crash with ValueError on startup

Description:
When attempting to run Gemma 4 models (e.g., gemma-4-e4b-it-4bit), the MLX server crashes with a ValueError: Received 126 parameters not in model. This is because Gemma 4 is a multi-modal model that includes new architecture components (like k_norm and vision/audio projectors) that the standard mlx-lm library does not yet fully support in its default model loader.

Fix:
This PR updates the MLX runtime management to use mlx-vlm for Gemma 4 models, which provides native support for the full Gemma 4 model family.

Changes:

  • Virtual Environment: Updated installMLX and locateMLX to ensure both mlx-lm>=0.31.0 and mlx-vlm>=0.4.0 are installed and verified.
  • Server Lifecycle: Modified startServer to dynamically switch to mlx_vlm.server when a Gemma 4 model is detected.
  • Chat API: Updated chatStream to support the OpenAI-compatible multi-modal message format (image_url), enabling the app to leverage Gemma 4's vision capabilities.
  • Health Checks: Improved startup log parsing to recognize mlx-vlm's "Started server process" message, preventing the UI from hanging during initialization.

Fixes #2

stamsam
stamsam reviewed Apr 29, 2026
View reviewed changes

@stamsam stamsam left a comment •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for working on this. I’m new to the project, so please take this as a drive-by suggestion rather than anything authoritative. The overall direction makes sense to me: Gemma 4 seems to need a VLM-capable loader, and routing those models through mlx_vlm.server feels like the right shape.

A few things that might make this easier for the maintainer to merge:

  • The README screenshot change looks unrelated to the runtime fix, so it might be cleaner as a separate PR.
  • model.includes("gemma-4") works with the current names, but explicit model/runtime metadata in AVAILABLE_MODELS might age better, something like runtime: "mlx-lm" | "mlx-vlm".
  • Checking for mlx_vlm in every locateMLX() call means setup can fail even for users who only want text/MLX-LM models. Maybe the VLM dependency could be checked/installed only when starting a VLM-backed model, unless the intent is for the managed venv to always include both runtimes.
  • This may overlap with #8, which moves the MLX server off the hardcoded Ollama port. If that lands first, this would probably want to use the selected port rather than MLX_PORT directly.
  • The image payload formatting looks useful, but I’m not sure the current renderer sends images into MLXChatMessage yet. Could be a follow-up unless it is needed for the crash fix.

Again, not trying to block anything here. This PR fixes an important problem; I just wanted to share a couple of review notes in case they help make the change smaller and easier to land.

MindXpansion added a commit to MindXpansion/gemma-chat that referenced this pull request May 26, 2026
@MindXpansion @claude
Patch 20: Gemma Chat's own integrated Neo4j KG (foundation, GOAT/SOTA)
951e48e 
Stands up gemma-chat-memory — a fully isolated, RAG-grade Neo4j 5.x
database living inside the existing kg-arch-enterprise DBMS. Schema
pinned verbatim to docs/research/05-neo4j-voyageai-rag-design.md
(authored 2026-05-17 by Claude-as-neo4j-kg-architect, 1029 lines).

FOUNDATION (this patch, ~2 hrs work):
  • New database gemma-chat-memory created inside kg-arch-enterprise
    (Enterprise multi-DB — same heap/page-cache, fully isolated stores)
  • Dedicated user `gemma-chat` with role `gemma_chat_rw`:
      - GRANT ALL ON DATABASE gemma-chat-memory
      - GRANT ALL ON GRAPH gemma-chat-memory  ← critical, separate from DATABASE
      - DENY ACCESS ON DATABASE neo4j (partnership KG)
      - DENY ACCESS ON DATABASE system
  • Generated password persisted to ~/.gemma-chat.env (mode 600)
  • Full schema applied (idempotent, IF NOT EXISTS throughout):
      - 10 uniqueness constraints (Document.uri, Document.sha256, all
        UUID keys for Chunk/Image/Conversation/Turn/Summary/Observation/
        Pattern + Workspace.id)
      - 5 range indexes on created_at / started_at for time-window
        queries
      - 2 full-text indexes (chunk_text_ft, summary_text_ft) for hybrid
        retrieval (vector + BM25) per research-05 §5
      - 5 vector indexes (chunk/image/summary/observation/pattern
        _embedding) at 1024-dim cosine, native vector-2.0 provider

DEFERS TO LATER PATCHES per research-05 §12 phased rollout:
  • Patch 21 — voyageai-4-large ingestion of the 4 research docs (~$0.04)
  • Patch 22 — aios.recall + reranker + citation packaging
  • Patch 23 — aios.cite / aios.search_kg / aios.index
  • Patch 24 — auto-conversation-ingestion (Turn/Summary autowrites)
  • Patch 25 — backup launchd, ingest queue, UI badges, spend tracker

APP-SIDE WIRING:
  • env-loader.ts: NEO4J_GEMMA_URI/USER/PASSWORD/DATABASE added to
    whitelist, read from ~/.gemma-chat.env at app boot
  • aios-neo4j.ts refactored to support TWO graph targets via
    GraphTarget = 'partnership' | 'gemma':
      - Per-target driver cache, lazy init, failure memo
      - runCypher(target, cypher, params) — accepts both reads and
        writes, surfaces write counters
      - getSchemaSummary(target) — labels/rels/constraints, indicates
        which graph
      - closeNeo4j() drains both drivers cleanly on app quit
  • New tools (symmetric to aios_kg_*):
      - gemma_kg_schema()             — your-own-graph schema
      - gemma_kg_query(cypher, params) — read/write Cypher to gemma-
        chat-memory. Vector search ready once Patch 21 ingests embed
        chunks: `CALL db.index.vector.queryNodes('chunk_embedding', $k,
        $vec) YIELD node, score`
  • System prompt updated: explicit two-graph model — partnership KG
    treads-carefully + own KG write-freely

SETUP SCRIPT (scripts/patch-20-foundation.cjs):
  • Idempotent, re-runnable
  • Reads admin creds from ~/.intelligence_partner/neo4j-creds.env
  • Pre-flight in code AND verified manually before run: enterprise
    edition, single :7687 listener (no Docker/Homebrew duplicate per
    architect anti-pattern ammaarreshi#1), no existing gemma-chat-memory
  • Smoke-tests after schema apply: probe document write+read+cleanup,
    constraint-fires verification (duplicate INSERT correctly rejected)
  • Rotates user password if user already exists

VERIFICATION (executed live, this commit):
  ✓ partnership driver: 69 labels reachable
  ✓ gemma driver: 10 constraints, 5 vector indexes, 2 fulltext, 15 range
  ✓ Isolation: gemma-chat user DENIED on database 'neo4j' (auth layer)

DESIGN DOC: docs/patch-20-gemma-own-kg-design.md distills research-05
into the patch sequence with risk assessment and Bear's three sign-off
decisions (voyage-4-large, demand-driven embedding, Foundation+Ingestion
tonight).

Build: vite clean. Pre-existing index.ts:22 hasMode warning unchanged.

Next: Patch 21 (voyageai-4-large ingestion of the 4 research docs +
minimum-viable recall tool). Will need a VOYAGE_API_KEY in
~/.gemma-chat.env or ~/.zshenv.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
MindXpansion added a commit to MindXpansion/gemma-chat that referenced this pull request May 27, 2026
@MindXpansion @claude
Patch 68: ToM parser — coerce out-of-enum USER_INTENTION + tighten pr…
cc76893 
…ompt

Issue surfaced today: heartbeat-adjacent chat turn produced a clean ToM
output where every field parsed correctly EXCEPT user_intention, which
the model emitted as "clarifying" — not in our closed eight-value
enum. parseToM() silently dropped the field; the required-field guard
then nulled the entire UserMentalModel, so no UMM was written for that
turn, no PSV adaptation drove the response, and the operator only saw
a single `[tom] parse-failed wall=1836ms raw="..."` log line.

Two contributors, both fixed:

1. Prompt-side (proximate). The enum was listed but nothing forbade
   inventing new values. Strengthened the line to
   "pick EXACTLY one: ..." and added a note explicitly calling out the
   three categories the model is most prone to invent
   ("clarifying" -> asking; "thanking" -> celebrating; "agreeing" ->
   other) so the model has somewhere to put them.

2. Parser-side (defense in depth). Even with a tighter prompt, models
   deviate. Instead of silently dropping, coerce unknown verbs to
   "other" and log the original verb so we can see the deviation rate
   in journal/console. This preserves the rest of the analyzer's read
   (emotion, rapport, knowledge_gap, confidence) instead of nulling
   everything because of one bad token.

Net effect: ToM keeps working through model deviation; we get an audit
trail of what the model tried to invent; nothing crashes.

Closes Issue ammaarreshi#1 of today's three-issue triage (parser was the only
fully-diagnosed item — timeout and stutter come next as separate
patches per agreed sequence).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

1 more reviewer
@stamsam stamsam stamsam left review comments
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

ValueError: Received 126 parameters not in model when running Gemma 4

3 participants

@waylonkenning @stamsam @ammaarreshi