update prompt parser and image size

cli/nvidia-smi.py

@@ -0,0 +1,31 @@
+import os
+import json
+import shutil
+import subprocess
+import xmltodict
+from rich import print # pylint: disable=redefined-builtin
+from util import log, Map
+
+def get_nvidia_smi(output='dict'):
+    smi = shutil.which('nvidia-smi')
+    if smi is None:
+        log.error("nvidia-smi not found")
+        return None
+    result = subprocess.run(f'"{smi}" -q -x', shell=True, check=False, env=os.environ, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    xml = result.stdout.decode(encoding="utf8", errors="ignore")
+    d = xmltodict.parse(xml)
+    if 'nvidia_smi_log' in d:
+        d = d['nvidia_smi_log']
+    if 'gpu' in d and 'supported_clocks' in d['gpu']:
+        del d['gpu']['supported_clocks']
+    if output == 'dict':
+        return d
+    elif output == 'class' or output == 'map':
+        d = Map(d)
+        return d
+    elif output == 'json':
+        return json.dumps(d, indent=4)
+
+if __name__ == "__main__":
+    res = get_nvidia_smi(output='dict')
+    print(type(res), res)

installer.py

@@ -452,6 +452,7 @@ def install_requirements():
 # set environment variables controling the behavior of various libraries
 def set_environment():
     log.info('Setting environment tuning')
+    os.environ.setdefault('USE_TORCH', '1')
     os.environ.setdefault('TF_CPP_MIN_LOG_LEVEL', '2')
     os.environ.setdefault('ACCELERATE', 'True')
     os.environ.setdefault('FORCE_CUDA', '1')

modules/images.py

@@ -19,6 +19,17 @@ from modules import sd_samplers, shared, script_callbacks, errors, paths
 LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
 
 
+def check_grid_size(imgs):
+    mp = 0
+    for img in imgs:
+        mp += img.width * img.height
+    mp = round(mp / 1000000)
+    ok = mp <= shared.opts.img_max_size_mp
+    if not ok:
+        shared.log.warning(f'Maximum image size exceded: size={mp} maximum={shared.opts.img_max_size_mp} MPixels')
+    return ok
+
+
 def image_grid(imgs, batch_size=1, rows=None):
     if rows is None:
         if shared.opts.n_rows > 0:
@@ -419,10 +430,6 @@ def atomically_save_image():
     Image.MAX_IMAGE_PIXELS = None # disable check in Pillow and rely on check below to allow large custom image sizes
     while True:
         image, filename, extension, params, exifinfo_data, txt_fullfn = save_queue.get()
-        mp = round(image.width * image.height / 1000000)
-        if mp > shared.opts.img_max_size_mp:
-            shared.log.warning(f'Maximum image size exceded: size={image.size} maximum={shared.opts.img_max_size_mp} MPixels')
-            return
         fn = filename + extension
         image_format = Image.registered_extensions()[extension]
         shared.log.debug(f'Saving image: {image_format} {fn} {image.size}')
@@ -433,9 +440,6 @@ def atomically_save_image():
                 pnginfo_data.add_text(k, str(v))
             image.save(fn, format=image_format, quality=shared.opts.jpeg_quality, pnginfo=pnginfo_data)
         elif image_format == 'JPEG':
-            if image.height > 65500 or image.width > 65500:
-                shared.log.warning(f'Maximum image size exceded: size={image.size} maximum=65550 pixels')
-                return
             if image.mode == 'RGBA':
                 shared.log.warning('Saving RGBA image as JPEG: Alpha channel will be lost')
                 image = image.convert("RGB")
@@ -513,6 +517,8 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='jpg', i
     if image is None:
         shared.log.warning('Image is none')
         return None, None
+    if not check_grid_size([image]):
+        return None, None
     if path is None: # set default path to avoid errors when functions are triggered manually or via api and param is not set
         path = shared.opts.outdir_save
     if save_to_dirs is None:

modules/processing.py

@@ -687,15 +687,16 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
         index_of_first_image = 0
         unwanted_grid_because_of_img_count = len(output_images) < 2 and opts.grid_only_if_multiple
         if (opts.return_grid or opts.grid_save) and not p.do_not_save_grid and not unwanted_grid_because_of_img_count:
-            grid = images.image_grid(output_images, p.batch_size)
-            if opts.return_grid:
-                text = infotext()
-                infotexts.insert(0, text)
-                grid.info["parameters"] = text
-                output_images.insert(0, grid)
-                index_of_first_image = 1
-            if opts.grid_save:
-                images.save_image(grid, p.outpath_grids, "grid", p.all_seeds[0], p.all_prompts[0], opts.grid_format, info=infotext(), short_filename=not opts.grid_extended_filename, p=p, grid=True)
+            if images.check_grid_size(output_images):
+                grid = images.image_grid(output_images, p.batch_size)
+                if opts.return_grid:
+                    text = infotext()
+                    infotexts.insert(0, text)
+                    grid.info["parameters"] = text
+                    output_images.insert(0, grid)
+                    index_of_first_image = 1
+                if opts.grid_save:
+                    images.save_image(grid, p.outpath_grids, "grid", p.all_seeds[0], p.all_prompts[0], opts.grid_format, info=infotext(), short_filename=not opts.grid_extended_filename, p=p, grid=True)
 
     if not p.disable_extra_networks and extra_network_data:
         extra_networks.deactivate(p, extra_network_data)

modules/prompt_parser.py

@@ -1,4 +1,10 @@
 # pylint: disable=anomalous-backslash-in-string
+
+import os
+import sys
+from rich import print
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+
 import re
 from collections import namedtuple
 from typing import List
@@ -16,7 +22,7 @@ from modules.shared import log, opts
 # [100, 'fantasy landscape with a lake and a christmas tree in background masterful']
 
 round_bracket_multiplier = 1.1
-square_bracket_multiplier = 0.9
+square_bracket_multiplier = 1.0 / 1.1
 re_AND = re.compile(r"\bAND\b")
 re_weight = re.compile(r"^(.*?)(?:\s*:\s*([-+]?(?:\d+\.?|\d*\.\d+)))?\s*$")
 ScheduledPromptConditioning = namedtuple("ScheduledPromptConditioning", ["end_at_step", "cond"])
@@ -306,7 +312,7 @@ def parse_prompt_attention(text):
         re_attention = re_attention_v1
         whitespace = ''
     else:
-        re_attention = re_attention_v2
+        re_attention = re_attention_v1
         text = text.replace('\\n', ' ')
         whitespace = ' '
 
@@ -326,9 +332,6 @@ def parse_prompt_attention(text):
             square_brackets.append(len(res))
         elif weight is not None and len(round_brackets) > 0:
             multiply_range(round_brackets.pop(), float(weight))
-        elif weight is not None and len(square_brackets) > 0:
-            if opts.prompt_attention == 'Full parser':
-                multiply_range(square_brackets.pop(), float(weight))
         elif text == ')' and len(round_brackets) > 0:
             multiply_range(round_brackets.pop(), round_bracket_multiplier)
         elif text == ']' and len(square_brackets) > 0:
@@ -362,11 +365,21 @@ def parse_prompt_attention(text):
     return res
 
 if __name__ == "__main__":
-    # import os
-    # import sys
-    # sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
-    # input_text = "(upzero) (upone:1.1), ((uptwo:1.2)), [downzero], [downone:0.9], [[downtwo:0.8]], this is a test"
-    input_text = 'a (white (lion:1.4)), cat [mouse] [tiger:0.8], ##, (high) in a jungle'
-    output_list = parse_prompt_attention(input_text)
-    print('INPUT', input_text)
-    print('OUTPUT', output_list)
+    input_text = '[black] [[grey]] (white) ((gray)) ((orange:1.1) yellow) ((purple) and [dark] red:1.1) [mouse:0.2] [(cat:1.1):0.5]'
+    print(f'Prompt: {input_text}')
+    schedules = get_learned_conditioning_prompt_schedules([input_text], 100)[0]
+    print('Schedules', schedules)
+    for schedule in schedules:
+        print('Schedule', schedule[0])
+        opts.data['prompt_attention'] = 'Fixed attention'
+        output_list = parse_prompt_attention(schedule[1])
+        print('  Fixed:', output_list)
+        opts.data['prompt_attention'] = 'Compel parser'
+        output_list = parse_prompt_attention(schedule[1])
+        print('  Compel:', output_list)
+        opts.data['prompt_attention'] = 'A1111 parser'
+        output_list = parse_prompt_attention(schedule[1])
+        print('  A1111:', output_list)
+        opts.data['prompt_attention'] = 'Full parser'
+        output_list = parse_prompt_attention(schedule[1])
+        print('  Full :', output_list)

scripts/prompt_matrix.py

@@ -27,10 +27,12 @@ def draw_xy_grid(xs, ys, x_label, y_label, cell):
 
             res.append(processed.images[0])
 
-    grid = images.image_grid(res, rows=len(ys))
-    grid = images.draw_grid_annotations(grid, res[0].width, res[0].height, hor_texts, ver_texts)
-
-    first_processed.images = [grid]
+    if images.check_grid_size(res):
+        grid = images.image_grid(res, rows=len(ys))
+        grid = images.draw_grid_annotations(grid, res[0].width, res[0].height, hor_texts, ver_texts)
+        first_processed.images = [grid]
+    else:
+        first_processed.images = res
 
     return first_processed
 
@@ -94,13 +96,13 @@ class Script(scripts.Script):
         p.prompt_for_display = positive_prompt
         processed = process_images(p)
 
-        grid = images.image_grid(processed.images, p.batch_size, rows=1 << ((len(prompt_matrix_parts) - 1) // 2))
-        grid = images.draw_prompt_matrix(grid, processed.images[0].width, processed.images[0].height, prompt_matrix_parts, margin_size)
-        processed.images.insert(0, grid)
-        processed.index_of_first_image = 1
-        processed.infotexts.insert(0, processed.infotexts[0])
-
-        if opts.grid_save:
-            images.save_image(processed.images[0], p.outpath_grids, "prompt_matrix", extension=opts.grid_format, prompt=original_prompt, seed=processed.seed, grid=True, p=p)
+        if images.check_grid_size(processed.images):
+            grid = images.image_grid(processed.images, p.batch_size, rows=1 << ((len(prompt_matrix_parts) - 1) // 2))
+            grid = images.draw_prompt_matrix(grid, processed.images[0].width, processed.images[0].height, prompt_matrix_parts, margin_size)
+            processed.images.insert(0, grid)
+            processed.index_of_first_image = 1
+            processed.infotexts.insert(0, processed.infotexts[0])
+            if opts.grid_save:
+                images.save_image(processed.images[0], p.outpath_grids, "prompt_matrix", extension=opts.grid_format, prompt=original_prompt, seed=processed.seed, grid=True, p=p)
 
         return processed

scripts/xyz_grid.py

@@ -322,7 +322,7 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
     for i in range(z_count):
         start_index = (i * len(xs) * len(ys)) + i
         end_index = start_index + len(xs) * len(ys)
-        if not no_grid:
+        if not no_grid and images.check_grid_size(processed_result.images[start_index:end_index]):
             grid = images.image_grid(processed_result.images[start_index:end_index], rows=len(ys))
             if draw_legend:
                 grid = images.draw_grid_annotations(grid, processed_result.images[start_index].size[0], processed_result.images[start_index].size[1], hor_texts, ver_texts, margin_size)
@@ -331,7 +331,7 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
         processed_result.all_seeds.insert(i, processed_result.all_seeds[start_index])
         processed_result.infotexts.insert(i, processed_result.infotexts[start_index])
     sub_grid_size = processed_result.images[0].size
-    if not no_grid:
+    if not no_grid and images.check_grid_size(processed_result.images[:z_count]):
         z_grid = images.image_grid(processed_result.images[:z_count], rows=1)
         if draw_legend:
             z_grid = images.draw_grid_annotations(z_grid, sub_grid_size[0], sub_grid_size[1], title_texts, [[images.GridAnnotation()]])