From c779b309a270dba2ac3bebad5a751b3963951ee2 Mon Sep 17 00:00:00 2001 From: Volpeon <git@volpeon.ink> Date: Fri, 6 Jan 2023 09:07:46 +0100 Subject: Add prompt template argument to inference --- infer.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/infer.py b/infer.py index e31cd88..d3d5f1b 100644 --- a/infer.py +++ b/infer.py @@ -48,6 +48,7 @@ default_args = { default_cmds = { "project": "", "scheduler": "dpmsm", + "template": "{}", "prompt": None, "negative_prompt": None, "shuffle": False, @@ -117,6 +118,10 @@ def create_cmd_parser(): type=str, choices=["plms", "ddim", "klms", "dpmsm", "dpmss", "euler_a", "kdpm2", "kdpm2_a"], ) + parser.add_argument( + "--template", + type=str, + ) parser.add_argument( "--prompt", type=str, @@ -243,6 +248,8 @@ def generate(output_dir: Path, pipeline, args): args.batch_size = 1 args.prompt = [keywords_to_prompt(prompt_to_keywords(prompt), shuffle=True) for prompt in args.prompt] + args.prompt = [args.template.format(prompt) for prompt in args.prompt] + now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S") image_dir = [] -- cgit v1.2.3-70-g09d2