From 32d494c28b56ee4a8ead6cae381d2f6ec3e69dac Mon Sep 17 00:00:00 2001
From: hako-mikan <122196982+hako-mikan@users.noreply.github.com>
Date: Fri, 10 Jan 2025 00:20:48 +0900
Subject: [PATCH] "#163"

---
 scripts/lora_block_weight.py | 34 ++++++++++++++++++----------------
 1 file changed, 18 insertions(+), 16 deletions(-)

diff --git a/scripts/lora_block_weight.py b/scripts/lora_block_weight.py
index fe66fc9..2b9f53f 100644
--- a/scripts/lora_block_weight.py
+++ b/scripts/lora_block_weight.py
@@ -1054,27 +1054,29 @@ def draw_origin(grid, text,width,height,width_one):
     return grid_d
 
 def newrun(p, *args):
-        script_index = args[0]
+    txt2img = isinstance(p, modules.processing.StableDiffusionProcessingTxt2Img)
 
-        if args[0] ==0:
-            script = None
-            for obj in scripts.scripts_txt2img.alwayson_scripts:
-                if "lora_block_weight" in obj.filename:
-                    script = obj 
-                    script_args = args[script.args_from:script.args_to]
-        else:
-            script = scripts.scripts_txt2img.selectable_scripts[script_index-1]
-            
-            if script is None:
-                return None
+    script_index = args[0]
 
-            script_args = args[script.args_from:script.args_to]
+    if args[0] ==0:
+        script = None
+        for obj in scripts.scripts_txt2img.alwayson_scripts if txt2img else scripts.scripts_img2img.alwayson_scripts:
+            if "lora_block_weight" in obj.filename:
+                script = obj 
+                script_args = args[script.args_from:script.args_to]
+    else:
+        script = scripts.scripts_txt2img.selectable_scripts[script_index-1] if txt2img else scripts.scripts_img2img.selectable_scripts[script_index-1]
+        
+        if script is None:
+            return None
 
-        processed = script.run(p, *script_args)
+        script_args = args[script.args_from:script.args_to]
 
-        shared.total_tqdm.clear()
+    processed = script.run(p, *script_args)
 
-        return processed
+    shared.total_tqdm.clear()
+
+    return processed
 
 registerd = False