Updated To actually work

prompts_from_file_kareeem.py

@@ -6,6 +6,8 @@ import modules.scripts as scripts
 from modules import sd_samplers, errors, sd_models
 from modules.processing import Processed, process_images
 from modules.shared import state
+import modules.shared as shared
+
 
 def process_string_tag(tag):
     return tag
@@ -51,6 +53,21 @@ prompt_tags = {
 }
 
 
+def sort_checkpoint(checkpoints):
+    current = shared.opts.sd_model_checkpoint
+    # print("Debug", current)
+    # print("Pre Sorted",checkpoints)
+    if (checkpoints[0] == current):
+        # print("Already Sorted")
+        return checkpoints
+    if (not (current in checkpoints)):
+        return checkpoints
+    checkpoints.remove(current)
+    checkpoints.insert(0, current)
+    # print("Sorted",checkpoints)
+    return checkpoints
+
+
 def cmdargs(line, sampler):
     args = shlex.split(line)
     pos = 0
@@ -60,7 +77,8 @@ def cmdargs(line, sampler):
         arg = args[pos]
 
         assert arg.startswith("--"), f'must start with "--": {arg}'
-        assert pos+1 < len(args), f'missing argument for command line option {arg}'
+        assert pos + \
+            1 < len(args), f'missing argument for command line option {arg}'
 
         tag = arg[2:]
 
@@ -75,7 +93,6 @@ def cmdargs(line, sampler):
             res[tag] = prompt
             continue
 
-
         func = prompt_tags.get(tag, None)
         assert func, f'unknown commandline option: {arg}'
 
@@ -86,52 +103,71 @@ def cmdargs(line, sampler):
 
         res[tag] = func(val)
 
-
         pos += 2
-    res["sd_model"] = sampler
+    if (sampler is not None):
+        res["sd_model"] = sampler
     return res
 
-def cmdargs(line):
-    return cmdargs(line, None)
 
 def load_prompt_file(file):
     if file is None:
         return None, gr.update(), gr.update(lines=7)
     else:
-        lines = [x.strip() for x in file.decode('utf8', errors='ignore').split("\n")]
+        lines = [x.strip() for x in file.decode(
+            'utf8', errors='ignore').split("\n")]
         return None, "\n".join(lines), gr.update(lines=7)
 
 
+def apply_checkpoint(p, x):
+    if x == shared.opts.sd_model_checkpoint:
+        return
+    info = sd_models.get_closet_checkpoint_match(x)
+    if info is None:
+        shared.log.warning(
+            f"XYZ grid: apply checkpoint unknown checkpoint: {x}")
+    else:
+        sd_models.reload_model_weights(shared.sd_model, info)
+        p.override_settings['sd_model_checkpoint'] = info.name
+
+
 class Script(scripts.Script):
     def title(self):
         return "Prompts from File Kareem Edition"
 
     def ui(self, is_img2img):
-        checkbox_iterate = gr.Checkbox(label="Iterate seed every line", value=False, elem_id=self.elem_id("checkbox_iterate"))
+        checkbox_iterate = gr.Checkbox(
-        checkbox_iterate_batch = gr.Checkbox(label="Use same random seed for all lines", value=False, elem_id=self.elem_id("checkbox_iterate_batch"))
+            label="Iterate seed every line", value=False, elem_id=self.elem_id("checkbox_iterate"))
-        checkpoint = gr.Dropdown(sorted(sd_models.checkpoints_list),label="Checkpoints", multiselect=True, value=None, lem_id=self.elem_id("checkpoint"))
+        checkbox_iterate_batch = gr.Checkbox(
-        prompt_txt = gr.Textbox(label="List of prompt inputs", lines=1, elem_id=self.elem_id("prompt_txt"))
+            label="Use same random seed for all lines", value=False, elem_id=self.elem_id("checkbox_iterate_batch"))
-        file = gr.File(label="Upload prompt inputs", type='binary', elem_id=self.elem_id("file"))
+        checkpoints = gr.Dropdown(sorted(sd_models.checkpoints_list), label="Checkpoints",
+                                  multiselect=True, value=None, lem_id=self.elem_id("checkpoints"))
+        prompt_txt = gr.Textbox(
+            label="List of prompt inputs", lines=1, elem_id=self.elem_id("prompt_txt"))
+        file = gr.File(label="Upload prompt inputs",
+                       type='binary', elem_id=self.elem_id("file"))
 
-        file.change(fn=load_prompt_file, inputs=[file], outputs=[file, prompt_txt, prompt_txt], show_progress=False)
+        file.change(fn=load_prompt_file, inputs=[file], outputs=[
+                    file, prompt_txt, prompt_txt], show_progress=False)
 
         # We start at one line. When the text changes, we jump to seven lines, or two lines if no \n.
         # We don't shrink back to 1, because that causes the control to ignore [enter], and it may
         # be unclear to the user that shift-enter is needed.
-        prompt_txt.change(lambda tb: gr.update(lines=7) if ("\n" in tb) else gr.update(lines=2), inputs=[prompt_txt], outputs=[prompt_txt], show_progress=False)
+        prompt_txt.change(lambda tb: gr.update(lines=7) if ("\n" in tb) else gr.update(
-        return [checkbox_iterate, checkbox_iterate_batch, prompt_txt, checkpoint]
+            lines=2), inputs=[prompt_txt], outputs=[prompt_txt], show_progress=False)
+        return [checkbox_iterate, checkbox_iterate_batch, prompt_txt, checkpoints]
 
-    def run(self, p, checkbox_iterate, checkbox_iterate_batch, prompt_txt: str, checkpoint):
+    def run(self, p, checkbox_iterate, checkbox_iterate_batch, prompt_txt: str, checkpoints):
+        current = shared.opts.sd_model_checkpoint
         lines = [x.strip() for x in prompt_txt.splitlines()]
         lines = [x for x in lines if len(x) > 0]
-
+        checkpoints = sort_checkpoint(checkpoints)
         job_count = 0
         jobs = []
-        if (checkpoint is None):
+        if (checkpoints is None):
             for line in lines:
                 if "--" in line:
                     try:
-                        args = cmdargs(line)
+                        args = cmdargs(line, None)
                     except Exception as e:
                         errors.display(e, f'parsing prompts: {line}')
                         args = {"prompt": line}
@@ -142,19 +178,19 @@ class Script(scripts.Script):
 
                 jobs.append(args)
         else:
-            for sampler in checkpoint:
+            for checkpoint in checkpoints:
                 for line in lines:
                     if "--" in line:
                         try:
-                            args = cmdargs(line, sampler)
+                            args = cmdargs(line, checkpoint)
                         except Exception as e:
                             errors.display(e, f'parsing prompts: {line}')
                             args = {"prompt": line}
                     else:
-                        args = {"prompt": line}
+                        args = {"prompt": line,
+                                "sd_model_checkpoint": checkpoint}
 
                     job_count += args.get("n_iter", p.n_iter)
-
                     jobs.append(args)
 
         print(f"Will process {len(lines)} lines in {job_count} jobs.")
@@ -172,7 +208,8 @@ class Script(scripts.Script):
             copy_p = copy.copy(p)
             for k, v in args.items():
                 setattr(copy_p, k, v)
-
+            if (checkpoints is not None):
+                apply_checkpoint(copy_p, copy_p.sd_model_checkpoint)
             proc = process_images(copy_p)
             images += proc.images
 
@@ -180,5 +217,5 @@ class Script(scripts.Script):
                 p.seed = p.seed + (p.batch_size * p.n_iter)
             all_prompts += proc.all_prompts
             infotexts += proc.infotexts
-
+        apply_checkpoint(copy_p, current)
         return Processed(p, images, p.seed, "", all_prompts=all_prompts, infotexts=infotexts)