feat: added toml support for sample prompt

library/train_util.py

@@ -3291,8 +3291,18 @@ def sample_images(
     vae.to(device)
 
     # read prompts
-    with open(args.sample_prompts, "rt", encoding="utf-8") as f:
-        prompts = f.readlines()
+
+    # with open(args.sample_prompts, "rt", encoding="utf-8") as f:
+    #     prompts = f.readlines()
+
+    if args.sample_prompts.endswith('.txt'):
+        with open(args.sample_prompts, 'r') as f:
+            lines = f.readlines()
+        prompts = [line.strip() for line in lines if len(line.strip()) > 0 and line[0] != "#"]
+    elif args.sample_prompts.endswith('.toml'):
+        with open(args.sample_prompts, 'r') as f:
+            data = toml.load(f)
+        prompts = [dict(**data['prompt'], **subset) for subset in data['prompt']['subset']]
     
     # schedulerを用意する
     sched_init_args = {}
@@ -3362,9 +3372,19 @@ def sample_images(
             for i, prompt in enumerate(prompts):
                 if not accelerator.is_main_process:
                     continue
-                prompt = prompt.strip()
-                if len(prompt) == 0 or prompt[0] == "#":
-                    continue
+            
+                if isinstance(prompt, dict):
+                    negative_prompt = prompt.get("negative_prompt")
+                    sample_steps = prompt.get("sample_steps", 30)
+                    width = prompt.get("width", 512)
+                    height = prompt.get("height", 512)
+                    scale = prompt.get("scale", 7.5)
+                    seed = prompt.get("seed")
+                    prompt = prompt.get("prompt")
+                else:               
+                    # prompt = prompt.strip()
+                    # if len(prompt) == 0 or prompt[0] == "#":
+                    #     continue
 
                     # subset of gen_img_diffusers
                     prompt_args = prompt.split(" --")