Add prompt template argument to inference

1 files changed, 7 insertions, 0 deletions

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,
@@ -118,6 +119,10 @@ def create_cmd_parser():
         choices=["plms", "ddim", "klms", "dpmsm", "dpmss", "euler_a", "kdpm2", "kdpm2_a"],
     )
     parser.add_argument(
+        "--template",
+        type=str,
+    )
+    parser.add_argument(
         "--prompt",
         type=str,
         nargs="+",
@@ -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 = []