Handle missing total_size in safetensors index files (#1956)

49670c86243fbeae396a865ed30be776831ddfd0 · 2026-04-21 16:39:14 +0100 · ciaranbor

## Motivation

Image models fail to load after a mid-download instance deletion and
recreation. The system skips the download and crashes with
`FileNotFoundError: No safetensors files found in .../vae`.

## Changes

- Make `ModelSafetensorsIndexMetadata.total_size` optional (`PositiveInt
| None = None`)
- Add null guard in `fetch_safetensors_size`
- Add regression test

## Why It Works

Exolabs quantized image models have safetensors index files with mflux
metadata (`quantization_level`, `mflux_version`) but no `total_size`.
The required `PositiveInt` field caused Pydantic validation to fail,
which was silently swallowed by `except Exception: continue` in
`_scan_model_directory`. This skipped all weight map checks, making
incomplete models appear complete.

## Test Plan

### Manual Testing

- Hardware: Mac Studio
- Before: `CreateRunner → LoadModel` (crash). After: `CreateRunner →
DownloadModel` (correct).

### Automated Testing

- `test_safetensors_index.py`: 3 cases covering missing, valid, and null
metadata

Files touched

A src/exo/download/tests/test_safetensors_index.py
M src/exo/shared/models/model_cards.py
M src/exo/shared/types/worker/downloads.py

Diff

commit 49670c86243fbeae396a865ed30be776831ddfd0
Author: ciaranbor <81697641+ciaranbor@users.noreply.github.com>
Date:   Tue Apr 21 16:39:14 2026 +0100

    Handle missing total_size in safetensors index files (#1956)
    
    ## Motivation
    
    Image models fail to load after a mid-download instance deletion and
    recreation. The system skips the download and crashes with
    `FileNotFoundError: No safetensors files found in .../vae`.
    
    ## Changes
    
    - Make `ModelSafetensorsIndexMetadata.total_size` optional (`PositiveInt
    | None = None`)
    - Add null guard in `fetch_safetensors_size`
    - Add regression test
    
    ## Why It Works
    
    Exolabs quantized image models have safetensors index files with mflux
    metadata (`quantization_level`, `mflux_version`) but no `total_size`.
    The required `PositiveInt` field caused Pydantic validation to fail,
    which was silently swallowed by `except Exception: continue` in
    `_scan_model_directory`. This skipped all weight map checks, making
    incomplete models appear complete.
    
    ## Test Plan
    
    ### Manual Testing
    
    - Hardware: Mac Studio
    - Before: `CreateRunner → LoadModel` (crash). After: `CreateRunner →
    DownloadModel` (correct).
    
    ### Automated Testing
    
    - `test_safetensors_index.py`: 3 cases covering missing, valid, and null
    metadata
---
 src/exo/download/tests/test_safetensors_index.py | 35 ++++++++++++++++++++++++
 src/exo/shared/models/model_cards.py             |  2 +-
 src/exo/shared/types/worker/downloads.py         |  2 +-
 3 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/src/exo/download/tests/test_safetensors_index.py b/src/exo/download/tests/test_safetensors_index.py
new file mode 100644
index 00000000..6425501c
--- /dev/null
+++ b/src/exo/download/tests/test_safetensors_index.py
@@ -0,0 +1,35 @@
+"""Tests for safetensors index file parsing."""
+
+from exo.shared.types.worker.downloads import ModelSafetensorsIndex
+
+
+def test_safetensors_index_missing_total_size():
+    """Index files from image models (e.g. FLUX/mflux) have metadata without
+    a total_size field — they use quantization_level and mflux_version instead.
+
+    This must parse successfully — a previous bug caused Pydantic validation to
+    fail silently (total_size was a required PositiveInt), which made
+    _scan_model_directory skip weight-map checks and report incomplete models
+    as complete.
+    """
+    raw = '{"metadata": {"quantization_level": "4", "mflux_version": "0.3.0"}, "weight_map": {"layer.safetensors": "layer.safetensors"}}'
+    index = ModelSafetensorsIndex.model_validate_json(raw)
+    assert index.metadata is not None
+    assert index.metadata.total_size is None
+    assert index.weight_map == {"layer.safetensors": "layer.safetensors"}
+
+
+def test_safetensors_index_valid_total_size():
+    """Standard text model index files with a valid total_size should continue
+    to parse correctly."""
+    raw = '{"metadata": {"total_size": 12345}, "weight_map": {"a.safetensors": "a.safetensors"}}'
+    index = ModelSafetensorsIndex.model_validate_json(raw)
+    assert index.metadata is not None
+    assert index.metadata.total_size == 12345
+
+
+def test_safetensors_index_null_metadata():
+    """Index files with null metadata should parse correctly."""
+    raw = '{"metadata": null, "weight_map": {"a.safetensors": "a.safetensors"}}'
+    index = ModelSafetensorsIndex.model_validate_json(raw)
+    assert index.metadata is None
diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py
index a4f8d9ab..7e9f9c30 100644
--- a/src/exo/shared/models/model_cards.py
+++ b/src/exo/shared/models/model_cards.py
@@ -367,7 +367,7 @@ async def fetch_safetensors_size(model_id: ModelId) -> Memory:
         index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
 
     metadata = index_data.metadata
-    if metadata is not None:
+    if metadata is not None and metadata.total_size is not None:
         return Memory.from_bytes(metadata.total_size)
 
     info = model_info(model_id)
diff --git a/src/exo/shared/types/worker/downloads.py b/src/exo/shared/types/worker/downloads.py
index 29540fe1..52036c0b 100644
--- a/src/exo/shared/types/worker/downloads.py
+++ b/src/exo/shared/types/worker/downloads.py
@@ -53,7 +53,7 @@ DownloadProgress = (
 
 
 class ModelSafetensorsIndexMetadata(BaseModel):
-    total_size: PositiveInt
+    total_size: PositiveInt | None = None
 
 
 class ModelSafetensorsIndex(BaseModel):

← fcc3718e Add sampling defaults (#1947) · back to Exo · Ignore HF remote repo changes (temporary fix) (#1958) af673845 →