33 lines
1.0 KiB
Python
33 lines
1.0 KiB
Python
import argparse
|
|
import os
|
|
import sys
|
|
|
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src')))
|
|
|
|
from mflux.flux import Flux1
|
|
from mflux.config.model_config import ModelConfig
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Save a quantized version of Flux.1 to disk.')
|
|
parser.add_argument('--path', type=str, required=True, help='Local path for loading a model from disk')
|
|
parser.add_argument('--model', "-m", type=str, required=True, choices=["dev", "schnell"], help='The model to use ("schnell" or "dev").')
|
|
parser.add_argument('--quantize', "-q", type=int, choices=[4, 8], default=8, help='Quantize the model (4 or 8, Default is 8)')
|
|
|
|
args = parser.parse_args()
|
|
|
|
print(f"Saving model {args.model} with quantization level {args.quantize}\n")
|
|
|
|
flux = Flux1(
|
|
model_config=ModelConfig.from_alias(args.model),
|
|
quantize=args.quantize,
|
|
)
|
|
|
|
flux.save_model(args.path)
|
|
|
|
print(f"Model saved at {args.path}\n")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|