Add export of Flux-dev transformer and uploading to Azure #717
Conversation
There was a slight difference in the Flux schnell and dev variants. Namely, dev has a guidance layer and schenll does not. Fixed some tensor argument element value types as they were always passed as f32 while some of the should use the models dtype. Refactored a bit the Flux transformer export boilerplate. Added a script that uploads models to Azure. Right now it uploads the Flux transformer models. This can become a part of the CI jobs at some point.
sharktank/requirements.txt
Outdated
| azure-identity>=1.19 | ||
| azure-storage-blob>=12.24 |
There was a problem hiding this comment.
Let's put these new requirements in a dev requirements file or an "extras" so they don't spill into the packages we publish unconditionally.
There was a problem hiding this comment.
The sharktank package currently has two requirements files, one for the standard install and one for the sharktank[testing] extra:
shark-ai/sharktank/pyproject.toml
Lines 33 to 37 in 5e3f5e6
Maybe testing is a good fit for this? We can also add a new category as needed
There was a problem hiding this comment.
This is not required by any tests right now. I added a dev package component. And put them in another file requirements file.
| @@ -0,0 +1,51 @@ | |||
| from ..utils.azure import upload_all_models | |||
There was a problem hiding this comment.
copyright headers at the tops of files please
| from ..utils.azure import upload_all_models | |
| # Copyright 2024 Advanced Micro Devices, Inc. | |
| # | |
| # Licensed under the Apache License v2.0 with LLVM Exceptions. | |
| # See https://llvm.org/LICENSE.txt for license information. | |
| # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | |
| from ..utils.azure import upload_all_models |
sharktank/requirements.txt
Outdated
| @@ -12,3 +12,6 @@ datasets | |||
| # Serving deps. | |||
| fastapi>=0.112.2 | |||
| uvicorn>=0.30.6 | |||
|
|
|||
| azure-identity>=1.19 | |||
There was a problem hiding this comment.
I don't think we should add azure as a requirement here. It is very nice for us developers, but our users will likely not have a use for it
There was a problem hiding this comment.
Just saw @ScottTodd's comment on this above. I agree that maybe we should add another dev or extras requirements file
There was a problem hiding this comment.
I answered on the other thread.
Add missed license header.
There was a slight difference in the Flux schnell and dev variants. Namely, dev has a guidance layer and schenll does not.
Fixed some tensor argument element value types as they were always passed as f32 while some of them should use the model's dtype.
Refactored a bit the Flux transformer export boilerplate.
Added a script that uploads models to Azure. Right now it uploads the Flux transformer models.
This can become a part of the CI jobs at some point.