mirror of
https://github.com/finegrain-ai/refiners.git
synced 2024-11-14 00:58:13 +00:00
41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
|
import argparse
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from refiners.fluxion.utils import load_tensors, save_to_safetensors
|
||
|
|
from refiners.foundationals.swin.mvanet.converter import convert_weights
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument(
|
||
|
|
"--from",
|
||
|
|
type=str,
|
||
|
|
required=True,
|
||
|
|
dest="source_path",
|
||
|
|
help="A MVANet checkpoint. One can be found at https://github.com/qianyu-dlut/MVANet",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--to",
|
||
|
|
type=str,
|
||
|
|
dest="output_path",
|
||
|
|
default=None,
|
||
|
|
help=(
|
||
|
|
"Path to save the converted model. If not specified, the output path will be the source path with the"
|
||
|
|
" extension changed to .safetensors."
|
||
|
|
),
|
||
|
|
)
|
||
|
|
parser.add_argument("--half", action="store_true", dest="half")
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
src_weights = load_tensors(args.source_path)
|
||
|
|
weights = convert_weights(src_weights)
|
||
|
|
if args.half:
|
||
|
|
weights = {key: value.half() for key, value in weights.items()}
|
||
|
|
if args.output_path is None:
|
||
|
|
args.output_path = f"{Path(args.source_path).stem}.safetensors"
|
||
|
|
save_to_safetensors(path=args.output_path, tensors=weights)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|