Skip to content

fix(diarizer): honor configuration.computeUnits in OfflineDiarizerModels.load#743

Open
chadneal wants to merge 1 commit into
FluidInference:mainfrom
apresai:honor-offline-computeunits
Open

fix(diarizer): honor configuration.computeUnits in OfflineDiarizerModels.load#743
chadneal wants to merge 1 commit into
FluidInference:mainfrom
apresai:honor-offline-computeunits

Conversation

@chadneal

Copy link
Copy Markdown

Fixes #742.

OfflineDiarizerModels.load(from:configuration:progressHandler:) accepts a configuration: MLModelConfiguration? but ignored it, always loading the segmentation / embedding / PLDA-rho models with .all. The streaming DiarizerModels loader already honors a caller-supplied configuration (testModelLoadingCustomConfig); this brings the offline loader in line.

Change

- let inferenceComputeUnits: MLComputeUnits = .all
+ let inferenceComputeUnits: MLComputeUnits = configuration?.computeUnits ?? .all
  • Default unchanged (.all) — no behavior change for existing callers.
  • Callers that need CPU/GPU-only inference (e.g. to avoid cross-device ANE numeric variance in the WeSpeaker embeddings) can now pass MLModelConfiguration(computeUnits:).
  • The fbank front-end intentionally stays on .cpuOnly.

Test

Adds testOfflineModelLoadingCustomConfig mirroring the existing streaming testModelLoadingCustomConfig, asserting the segmentation / embedding / PLDA-rho models load with the requested compute units.

swift build --build-tests passes locally.

…els.load

OfflineDiarizerModels.load(from:configuration:progressHandler:) accepted a
`configuration: MLModelConfiguration?` but never read it, hardcoding the
inference compute units to `.all`. The streaming `DiarizerModels` loader
already honors a caller-supplied configuration (covered by
`testModelLoadingCustomConfig`); this brings the offline loader in line.

The default stays `.all`, so existing callers are unaffected. Callers that
need CPU/GPU-only inference (e.g. to avoid cross-device ANE numeric variance
in the WeSpeaker embeddings) can now pass `MLModelConfiguration(computeUnits:)`.
The fbank front-end intentionally stays on `.cpuOnly`.

Adds `testOfflineModelLoadingCustomConfig` mirroring the existing streaming test.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

OfflineDiarizerModels.load(configuration:) ignores configuration.computeUnits (hardcodes .all)

1 participant