import sys, click import safetensors_worker # This file deals with command line only. If the command line is parsed successfully, # we will call one of the functions in safetensors_worker.py. readonly_input_file=click.argument("input_file", metavar='input_file', type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True)) output_file=click.argument("output_file", metavar='output_file', type=click.Path(file_okay=True, dir_okay=False, writable=True)) force_overwrite_flag=click.option("-f","--force-overwrite",default=False,is_flag=True, show_default=True, help="overwrite existing files") fix_ued_flag=click.option("-pm","--parse-more",default=False,is_flag=True, show_default=True, help="when printing metadata, unescaped doublequotes to make text more readable" ) quiet_flag=click.option("-q","--quiet",default=False,is_flag=True, show_default=True, help="Quiet mode, don't print informational stuff" ) @click.group() @click.version_option(version=7) @quiet_flag @click.pass_context def cli(ctx,quiet:bool): # ensure that ctx.obj exists and is a dict (in case `cli()` is called # by means other than the `if` block below) ctx.ensure_object(dict) ctx.obj['quiet'] = quiet @cli.command(name="header",short_help="print file header") @readonly_input_file @click.pass_context def cmd_header(ctx,input_file:str) -> int: sys.exit( safetensors_worker.PrintHeader(ctx.obj,input_file) ) @cli.command(name="metadata",short_help="print only __metadata__ in file header") @readonly_input_file @fix_ued_flag @click.pass_context def cmd_meta(ctx,input_file:str,parse_more:bool)->int: ctx.obj['parse_more'] = parse_more sys.exit( safetensors_worker.PrintMetadata(ctx.obj,input_file) ) @cli.command(name="listkeys",short_help="print header key names (except __metadata__) as a Python list") @readonly_input_file @click.pass_context def cmd_keyspy(ctx,input_file:str) -> int: sys.exit( safetensors_worker.HeaderKeysToLists(ctx.obj,input_file) ) @cli.command(name="writemd",short_help="read __metadata__ from json and write to safetensors file") @click.argument("in_st_file", metavar='input_st_file', type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True)) @click.argument("in_json_file", metavar='input_json_file', type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True)) @output_file @force_overwrite_flag @click.pass_context def cmd_writemd(ctx,in_st_file:str,in_json_file:str,output_file:str,force_overwrite:bool) -> int: """Read "__metadata__" from json file and write to safetensors header""" ctx.obj['force_overwrite'] = force_overwrite sys.exit( safetensors_worker.WriteMetadataToHeader(ctx.obj,in_st_file,in_json_file,output_file) ) @cli.command(name="extracthdr",short_help="extract file header and save to output file") @readonly_input_file @output_file @force_overwrite_flag @click.pass_context def cmd_extractheader(ctx,input_file:str,output_file:str,force_overwrite:bool) -> int: ctx.obj['force_overwrite'] = force_overwrite sys.exit( safetensors_worker.ExtractHeader(ctx.obj,input_file,output_file) ) @cli.command(name="extractdata",short_help="extract one tensor and save to file") @readonly_input_file @click.argument("key_name", metavar='key_name',type=click.STRING) @output_file @force_overwrite_flag @click.pass_context def cmd_extractheader(ctx,input_file:str,key_name:str,output_file:str,force_overwrite:bool) -> int: ctx.obj['force_overwrite'] = force_overwrite sys.exit( safetensors_worker.ExtractData(ctx.obj,input_file,key_name,output_file) ) @cli.command(name="checklora",short_help="see if input file is a SD 1.x LoRA file") @readonly_input_file @click.pass_context def cmd_checklora(ctx,input_file:str)->int: sys.exit( safetensors_worker.CheckLoRA(ctx.obj,input_file) ) if __name__ == '__main__': sys.stdout.reconfigure(encoding='utf-8') cli(obj={},max_content_width=96)