Flux2Pipeline loader ignores compiled assets for the text encoder, decoder and optional encoder, and hangs on AOT bundles

[james-333i]

Model name

Flux.2 Klein

Command run

Flux2Pipeline(from:) resolves the transformer through ModelBundle.resolveAssetURL, so a compiled .aimodelc is found when the bundle was produced with coreai-build compile. The text encoder and VAE decoder URLs are built with a raw appendingPathComponent(path) from the metadata instead, so on an iOS bundle the loader asks AIModel(contentsOf:) for TextEncoder.aimodel, which does not exist next to TextEncoder.aimodelc. That call never returns and never throws: the app sits at 0% CPU with no log output, which took a while to diagnose on device.

The optional VAE encoder has the same raw path and is not existence-checked, so a text-only export (no img2img components) still gets a non-nil encoder pointing at a missing file, and the first img2img request fails at runtime with a "Missing hash file" error instead of the pipeline reporting that img2img is unavailable.

Repro: export FLUX.2 Klein for iOS, compile the components with coreai-build compile --platform iOS, remove the .aimodel IR directories as a shipping app would, and call Flux2Pipeline(from:) on device.  A ready-made bundle is at https://huggingface.co/333i/flux2-klein-coreai-ios (the ios/ folder, compiled for h18p with --preferred-compute gpu, IR directories omitted). Pointing Flux2Pipeline(from:) at that folder on an iPhone 17 Pro reproduces the hang on unpatched main and generates normally with the branch above.

Fix, two commits, on my fork since PR creation is limited to collaborators: https://github.com/apple/coreai-models/compare/main...james-333i:coreai-models:fix-flux-loader-compiled-assets

The text encoder and decoder go through resolveAssetURL, and the encoder is resolved and existence-checked, becoming nil when absent so supportsImageToImage reports the truth. Verified on an iPhone 17 Pro with both a text-only and a full bundle. Happy to open it as a PR if that becomes possible.

macOS / iOS target

iOS 27.0 Beta 8

Xcode version

Xcode 27 Beta 6

Python / uv version

uv 0.12.9, python 3.11.2

Full error output

Model not found

Anything else?

No response

[stikves]

Thanks for the bug report.

The resolution helper for stale aimodel to aimodelc replacement is a convenience method to help when the metadata was not properly updated after compilation.

But, yes, we can extend this to VAE, too.

And, yes, we can have earlier reporting of missing components.

[james-333i]

@stikves Thanks, that makes sense. I will update the bundle metadata to name the compiled assets directly so the helper is only a fallback, and the VAE extension covers the remaining case for bundles that ship without the optional encoder. Appreciated.