88 lines
2.9 KiB
Python
88 lines
2.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/serve.py
|
|
|
|
import argparse
|
|
from typing import List, cast
|
|
|
|
from fastvideo.v1.entrypoints.cli import utils
|
|
from fastvideo.v1.entrypoints.cli.cli_types import CLISubcommand
|
|
from fastvideo.v1.fastvideo_args import FastVideoArgs
|
|
from fastvideo.v1.utils import FlexibleArgumentParser
|
|
|
|
|
|
class GenerateSubcommand(CLISubcommand):
|
|
"""The `generate` subcommand for the FastVideo CLI"""
|
|
|
|
def __init__(self) -> None:
|
|
self.name = "generate"
|
|
super().__init__()
|
|
|
|
def cmd(self, args: argparse.Namespace) -> None:
|
|
excluded_args = [
|
|
'subparser', 'config', 'num_gpus', 'master_port',
|
|
'dispatch_function'
|
|
]
|
|
|
|
# Create a filtered dictionary of arguments
|
|
filtered_args = {
|
|
k: v
|
|
for k, v in vars(args).items()
|
|
if k not in excluded_args and v is not None
|
|
}
|
|
|
|
main_args = []
|
|
|
|
for key, value in filtered_args.items():
|
|
# Convert underscores to dashes in argument names
|
|
arg_name = f"--{key.replace('_', '-')}"
|
|
|
|
# Handle boolean flags
|
|
if isinstance(value, bool):
|
|
if value:
|
|
main_args.append(arg_name)
|
|
else:
|
|
main_args.append(arg_name)
|
|
main_args.append(str(value))
|
|
|
|
utils.launch_distributed(args.num_gpus,
|
|
main_args,
|
|
master_port=args.master_port)
|
|
|
|
def validate(self, args: argparse.Namespace) -> None:
|
|
if args.num_gpus is not None and args.num_gpus <= 0:
|
|
raise ValueError("Number of gpus must be positive")
|
|
|
|
if args.master_port is not None and (args.master_port < 1024
|
|
or args.master_port > 65535):
|
|
raise ValueError("Master port must be between 1024 and 65535")
|
|
|
|
def subparser_init(
|
|
self,
|
|
subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
|
|
generate_parser = subparsers.add_parser(
|
|
"generate",
|
|
help="Run inference on a model",
|
|
usage=
|
|
"fastvideo generate --model-path MODEL_PATH_OR_ID --prompt PROMPT [OPTIONS]"
|
|
)
|
|
|
|
generate_parser.add_argument(
|
|
"--config",
|
|
type=str,
|
|
default='',
|
|
required=False,
|
|
help="Read CLI options from a config YAML file.")
|
|
|
|
generate_parser.add_argument("--master-port",
|
|
type=int,
|
|
default=None,
|
|
help="Port for the master process")
|
|
|
|
generate_parser = FastVideoArgs.add_cli_args(generate_parser)
|
|
|
|
return cast(FlexibleArgumentParser, generate_parser)
|
|
|
|
|
|
def cmd_init() -> List[CLISubcommand]:
|
|
return [GenerateSubcommand()]
|