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