Skip to content

Refactor the model chunking script to work with all version of coremltools#354

Merged
TobyRoseman merged 3 commits intoapple:mainfrom
jakesabathia2:eng/PR-henry-refactor-chunk-script
Aug 23, 2024
Merged

Refactor the model chunking script to work with all version of coremltools#354
TobyRoseman merged 3 commits intoapple:mainfrom
jakesabathia2:eng/PR-henry-refactor-chunk-script

Conversation

@jakesabathia2
Copy link
Contributor

@jakesabathia2 jakesabathia2 commented Aug 22, 2024

  • In coremltools 8.0b2, they introduced an util ct.models.utils.bisect_model (Add bisect_model API under ct.models.utils coremltools#2286),
    in which does the same thing as chunk_mlprogram.py, with a more robust implementation.
    For instance, taking the source CoreML model spec version into account.
  • This PR makes the chunk_mlprogram.py works on older coremltools, like 7.2, which has been reported a bug (chunk unet does not work with coremltool>7.1 #333).
  • Also for the coremltools version >=8.0b2, we refactor the script to directly call into bisect_model API.
  • The legancy implementation should be removed after we officially set the dependency on coremltools>=8.0

Testing:
python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --convert-text-encoder --convert-vae-decoder --convert-safety-checker --model-version CompVis/stable-diffusion-v1-4 -o ../../model_zoo --chunk-unet
ran on both coremltools==8.0b2 and coremltools==7.2

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.

3 participants