fix: enhance model saving functionality to load multiple JSON configs present in diffusers by davidberenstein1957 · Pull Request #98 · PrunaAI/pruna

davidberenstein1957 · 2025-05-06T11:03:21Z

Updated save_pruna_model_to_hub to load all JSON files in the model directory, ensuring smash_config.json is present.
Modified README template to reflect that model configuration is now stored in *.json files instead of a single config.json file.
Adjusted code examples for clarity and consistency.
Resolved bug where loading saved SmashConfig does not have locally stored smash_config.json in cache dir after loading from hub
Added integration pytest marker

Description

Error reproduction code using another branch

from pruna import PrunaModel

model_id = "PrunaAI/tiny-stable-diffusion-pipe-smashed"
loaded_model = PrunaModel.from_hub(model_id)
loaded_model.save_to_hub(model_id)

Related Issue

Fixes #96

Type of Change

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
This change requires a documentation update

How Has This Been Tested?

Checklist

My code follows the style guidelines of this project
I have performed a self-review of my own code
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

Additional Notes

src/pruna/engine/save.py

src/pruna/engine/handler/handler_utils.py

…nfig files - Updated `save_pruna_model_to_hub` to load all JSON files in the model directory, ensuring `smash_config.json` is present. - Modified README template to reflect that model configuration is now stored in `*.json` files instead of a single `config.json` file. - Adjusted code examples for clarity and consistency.

- Changed the key used to retrieve the smash configuration from the config files to ensure compatibility with the updated file naming convention.

- Modified test cases in `test_load.py` to include new model names. - Updated `test_save.py` to parameterize the `repo_id` for saving models, allowing for testing with multiple model configurations.

- Added functionality to save the old smash configuration, including tokenizer and processor, before saving the model. - Improved handling of model types in the `original_save_fn` to ensure correct loading functions are appended for diffusers and transformers. - Updated tests to parameterize repository IDs for better model saving tests.

- Introduced a new marker "integration" in the pytest configuration to categorize integration tests.

- Updated test cases in `test_load.py` to reflect new repository names for models. - Enhanced `test_save.py` with new tests for saving both LLM and diffusers models to the Hugging Face Hub. - Added debugging breakpoints in `smash_config.py` and `test_load.py` for troubleshooting.

- Removed the debugging breakpoint (pdb.set_trace) from the SmashConfig class to clean up the code after troubleshooting.

… smash_config - Removed the debugging breakpoint from `test_load.py` to clean up the code. - Updated `smash_config` initialization in `test_save.py` to specify the device as "cpu".

- Removed unnecessary model attribute access in `register_inference_handler` to streamline the inference handler registration process.

- Eliminated the debugging breakpoint (pdb.set_trace) from the `register_inference_handler` function to clean up the code and improve readability.

- Removed redundant check for the presence of the smash configuration file in the model directory, simplifying the loading process. - Adjusted the handling of the model attribute in the original save function to improve clarity.

- Enhanced the loading of the smash configuration by directly reading from the file instead of using a temporary dictionary. - Added logic to determine the library name based on the model's module for better clarity in the README. - Updated the README template to include a more informative loading instruction for the model, linking to the appropriate library documentation.

- Revised instructions in the model card template to clarify the use of the {library_name} for loading models, emphasizing that it may not include all optimizations by default. - Added guidance to utilize the pruna library for loading models to ensure all optimizations are applied.

- Modified the `test_pruna_model_from_hub` to include the `force_download=True` parameter when loading models from the hub, ensuring that the latest model versions are retrieved during tests.

gsprochette

Looks good already, but needs an iteration before it's good to go :) I left a couple of suggestions and there is the question about saving in SAVE_BEFORE_SMASH_CACHE_DIR which I didn't get.

tests/engine/test_save.py

tests/engine/test_load.py

src/pruna/engine/save.py

- Removed the saving of the old smash configuration from the `save_pruna_model` function to streamline the saving process. - Updated test cases in `test_load.py` and `test_save.py` to reflect new repository names for models, ensuring consistency in naming conventions.

gsprochette

LGTM! Thanks for adressing all the comments 👍 ✨

begumcig

LGTM!!!

davidberenstein1957 linked an issue May 6, 2025 that may be closed by this pull request

[BUG] PrunaModel.save_to_hub does not work with diffusers #96

Closed

davidberenstein1957 force-pushed the fix/96-bug-prunamodelsave_to_hub-does-not-work-with-diffusers branch from feccf3a to b2a57dc Compare May 6, 2025 11:05

davidberenstein1957 marked this pull request as ready for review May 6, 2025 11:05

davidberenstein1957 changed the title ~~refactor: enhance model saving functionality to load multiple JSON co…~~ fix: enhance model saving functionality to load multiple JSON co… May 6, 2025

davidberenstein1957 requested a review from johannaSommer May 6, 2025 11:58

davidberenstein1957 changed the title ~~fix: enhance model saving functionality to load multiple JSON co…~~ fix: enhance model saving functionality to load multiple JSON configs present in diffusers May 6, 2025

davidberenstein1957 self-assigned this May 6, 2025

davidberenstein1957 added the bug Something isn't working label May 6, 2025

davidberenstein1957 requested review from gsprochette and sharpenb May 7, 2025 03:44

davidberenstein1957 commented May 15, 2025

View reviewed changes

src/pruna/engine/save.py Outdated Show resolved Hide resolved

davidberenstein1957 requested review from begumcig and llcnt and removed request for gsprochette and sharpenb May 15, 2025 08:18

begumcig requested changes May 15, 2025

View reviewed changes

src/pruna/engine/handler/handler_utils.py Outdated Show resolved Hide resolved

davidberenstein1957 added 13 commits May 15, 2025 21:57

fix: add backwards compatibility check for load function in model saving

bd0832a

fix: update key for loading smash_config in model saving

aa37f89

- Changed the key used to retrieve the smash configuration from the config files to ensure compatibility with the updated file naming convention.

test: update model tests for new repository names

b3480c3

- Modified test cases in `test_load.py` to include new model names. - Updated `test_save.py` to parameterize the `repo_id` for saving models, allowing for testing with multiple model configurations.

Update src/pruna/engine/save.py

cda51d7

feat: add integration test marker in pytest configuration

33c81f5

- Introduced a new marker "integration" in the pytest configuration to categorize integration tests.

chore: remove debugging breakpoint from smash_config.py

9663151

- Removed the debugging breakpoint (pdb.set_trace) from the SmashConfig class to clean up the code after troubleshooting.

refactor: update test cases to remove debugging and specify device in…

2d12041

… smash_config - Removed the debugging breakpoint from `test_load.py` to clean up the code. - Updated `smash_config` initialization in `test_save.py` to specify the device as "cpu".

refactor: simplify model handling in inference registration

9af63cd

- Removed unnecessary model attribute access in `register_inference_handler` to streamline the inference handler registration process.

refactor: remove debugging breakpoint from handler_utils.py

f8de80b

- Eliminated the debugging breakpoint (pdb.set_trace) from the `register_inference_handler` function to clean up the code and improve readability.

refactor: streamline smash config loading in model saving

4647bd8

- Removed redundant check for the presence of the smash configuration file in the model directory, simplifying the loading process. - Adjusted the handling of the model attribute in the original save function to improve clarity.

davidberenstein1957 force-pushed the fix/96-bug-prunamodelsave_to_hub-does-not-work-with-diffusers branch from 0c6888d to 4647bd8 Compare May 15, 2025 19:57

davidberenstein1957 requested review from begumcig and gsprochette and removed request for johannaSommer May 15, 2025 19:59

davidberenstein1957 added 3 commits May 16, 2025 12:13

refactor: update model loading in tests to ensure forced download

10f9b6d

- Modified the `test_pruna_model_from_hub` to include the `force_download=True` parameter when loading models from the hub, ensuring that the latest model versions are retrieved during tests.

gsprochette reviewed May 16, 2025

View reviewed changes

tests/engine/test_save.py Outdated Show resolved Hide resolved

tests/engine/test_load.py Outdated Show resolved Hide resolved

src/pruna/engine/save.py Outdated Show resolved Hide resolved

src/pruna/engine/save.py Outdated Show resolved Hide resolved

davidberenstein1957 requested a review from gsprochette May 17, 2025 13:03

gsprochette approved these changes May 20, 2025

View reviewed changes

davidberenstein1957 removed the request for review from llcnt May 22, 2025 08:24

davidberenstein1957 assigned SaboniAmine May 22, 2025

begumcig approved these changes May 22, 2025

View reviewed changes

davidberenstein1957 merged commit d06595f into main May 22, 2025
7 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix: enhance model saving functionality to load multiple JSON configs present in diffusers#98

fix: enhance model saving functionality to load multiple JSON configs present in diffusers#98
davidberenstein1957 merged 17 commits intomainfrom
fix/96-bug-prunamodelsave_to_hub-does-not-work-with-diffusers

davidberenstein1957 commented May 6, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

gsprochette left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gsprochette left a comment

Uh oh!

begumcig left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

davidberenstein1957 commented May 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Type of Change

How Has This Been Tested?

Checklist

Additional Notes

Uh oh!

Uh oh!

Uh oh!

gsprochette left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gsprochette left a comment

Choose a reason for hiding this comment

Uh oh!

begumcig left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

davidberenstein1957 commented May 6, 2025 •

edited

Loading