-
Notifications
You must be signed in to change notification settings - Fork 270
Add popular model IO configs #2304
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
lintrunner found more than 20 potential problems in the proposed changes. Check the Files changed tab for more details.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull request overview
This PR adds a comprehensive internal implementation of HuggingFace model IO configurations to replace the dependency on the optimum library. The changes include adding support for popular model types and diffusers pipeline components (SD, SDXL, SD3, Flux, Sana), implementing ONNX export configurations, and updating the DiffusersModelVariant enum from "sd15" to "sd".
Changes:
- Adds a new
olive/common/hf/io_config/module with base classes, model configs, input generators, normalized configs, and task definitions for ONNX export - Updates
DiffusersModelVariantenum: renames SD15 to SD, adds SD3 and SANA variants, and moves toolive/constants.py - Adds
DiffusersComponentenum for pipeline component names - Refactors diffusers model detection to use config files instead of model name patterns
- Updates ONNX conversion to use new io_config instead of optimum's
main_export - Updates all tests and CLI to use the new variant names
Reviewed changes
Copilot reviewed 27 out of 29 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| olive/constants.py | Adds DiffusersModelVariant and DiffusersComponent enums |
| olive/common/hf/io_config/ | New module with 1000+ lines implementing ONNX export configs for 80+ model types |
| olive/model/handler/diffusers.py | Updates variant detection logic and adds exportable component methods |
| olive/passes/onnx/conversion.py | Refactors diffusers export to use new io_config instead of optimum |
| olive/common/hf/model_io.py | Simplifies to use new io_config module |
| test/ files | Updates all tests to use new variant names and adds comprehensive test coverage |
| olive/cli/ | Updates CLI to support new variants |
| if model_type not in _ONNX_CONFIG_REGISTRY: | ||
| return False | ||
| return task in _ONNX_CONFIG_REGISTRY[model_type] |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Rewrite
return model_type in _ONNX_CONFIG_REGISTRY and task in _ONNX_CONFIG_REGISTRY[model_type]
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done!
olive/model/handler/diffusers.py
Outdated
| if variant == DiffusersModelVariant.SD: | ||
| return [ | ||
| DiffusersComponent.TEXT_ENCODER, | ||
| DiffusersComponent.UNET, | ||
| DiffusersComponent.VAE_ENCODER, | ||
| DiffusersComponent.VAE_DECODER, | ||
| ] | ||
| elif variant == DiffusersModelVariant.SDXL: | ||
| return [ | ||
| DiffusersComponent.TEXT_ENCODER, | ||
| DiffusersComponent.TEXT_ENCODER_2, | ||
| DiffusersComponent.UNET, | ||
| DiffusersComponent.VAE_ENCODER, | ||
| DiffusersComponent.VAE_DECODER, | ||
| ] | ||
| elif variant == DiffusersModelVariant.SD3: | ||
| return [ | ||
| DiffusersComponent.TEXT_ENCODER, | ||
| DiffusersComponent.TEXT_ENCODER_2, | ||
| DiffusersComponent.TEXT_ENCODER_3, | ||
| DiffusersComponent.TRANSFORMER, | ||
| DiffusersComponent.VAE_ENCODER, | ||
| DiffusersComponent.VAE_DECODER, | ||
| ] | ||
| elif variant == DiffusersModelVariant.FLUX: | ||
| return [ | ||
| DiffusersComponent.TEXT_ENCODER, | ||
| DiffusersComponent.TEXT_ENCODER_2, | ||
| DiffusersComponent.TRANSFORMER, | ||
| DiffusersComponent.VAE_ENCODER, | ||
| DiffusersComponent.VAE_DECODER, | ||
| ] | ||
| elif variant == DiffusersModelVariant.SANA: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Worth converting this if/else to a map, may be even a constant map?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Using map now.
Describe your changes
This change is verified by running capture-onnx-graph CLI for 70+ huggingface models.
Checklist before requesting a review
lintrunner -a(Optional) Issue link