Merge pull request #8326 from DejitaruJin/xyz-order-fix

Xyz order fix

1 files changed, 77 insertions, 50 deletions

diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py
index 8e822e16..f068598d 100644
--- a/scripts/xyz_grid.py
+++ b/scripts/xyz_grid.py
@@ -218,49 +218,47 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
     ver_texts = [[images.GridAnnotation(y)] for y in y_labels]

     title_texts = [[images.GridAnnotation(z)] for z in z_labels]

 

-    # Temporary list of all the images that are generated to be populated into the grid.

-    # Will be filled with empty images for any individual step that fails to process properly

-    image_cache = [None] * (len(xs) * len(ys) * len(zs))

+    list_size = (len(xs) * len(ys) * len(zs))

 

     processed_result = None

-    cell_mode = "P"

-    cell_size = (1, 1)

 

-    state.job_count = len(xs) * len(ys) * len(zs) * p.n_iter

+    state.job_count = list_size * p.n_iter

 

     def process_cell(x, y, z, ix, iy, iz):

-        nonlocal image_cache, processed_result, cell_mode, cell_size

+        nonlocal processed_result

 

         def index(ix, iy, iz):

             return ix + iy * len(xs) + iz * len(xs) * len(ys)

 

-        state.job = f"{index(ix, iy, iz) + 1} out of {len(xs) * len(ys) * len(zs)}"

+        state.job = f"{index(ix, iy, iz) + 1} out of {list_size}"

 

         processed: Processed = cell(x, y, z)

 

-        try:

-            # this dereference will throw an exception if the image was not processed

-            # (this happens in cases such as if the user stops the process from the UI)

-            processed_image = processed.images[0]

-

-            if processed_result is None:

-                # Use our first valid processed result as a template container to hold our full results

-                processed_result = copy(processed)

-                cell_mode = processed_image.mode

-                cell_size = processed_image.size

-                processed_result.images = [Image.new(cell_mode, cell_size)]

-                processed_result.all_prompts = [processed.prompt]

-                processed_result.all_seeds = [processed.seed]

-                processed_result.infotexts = [processed.infotexts[0]]

-

-            image_cache[index(ix, iy, iz)] = processed_image

-            if include_lone_images:

-                processed_result.images.append(processed_image)

-                processed_result.all_prompts.append(processed.prompt)

-                processed_result.all_seeds.append(processed.seed)

-                processed_result.infotexts.append(processed.infotexts[0])

-        except:

-            image_cache[index(ix, iy, iz)] = Image.new(cell_mode, cell_size)

+        if processed_result is None:

+            # Use our first processed result object as a template container to hold our full results

+            processed_result = copy(processed)

+            processed_result.images = [None] * list_size

+            processed_result.all_prompts = [None] * list_size

+            processed_result.all_seeds = [None] * list_size

+            processed_result.infotexts = [None] * list_size

+            processed_result.index_of_first_image = 1

+

+        idx = index(ix, iy, iz)

+        if processed.images:

+            # Non-empty list indicates some degree of success.

+            processed_result.images[idx] = processed.images[0]

+            processed_result.all_prompts[idx] = processed.prompt

+            processed_result.all_seeds[idx] = processed.seed

+            processed_result.infotexts[idx] = processed.infotexts[0]

+        else:

+            cell_mode = "P"

+            cell_size = (processed_result.width, processed_result.height)

+            if processed_result.images[0] is not None:

+                cell_mode = processed_result.images[0].mode

+                #This corrects size in case of batches:

+                cell_size = processed_result.images[0].size

+            processed_result.images[idx] = Image.new(cell_mode, cell_size)

+

 

     if first_axes_processed == 'x':

         for ix, x in enumerate(xs):

@@ -294,27 +292,37 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
                         process_cell(x, y, z, ix, iy, iz)

 

     if not processed_result:

+        # Should never happen, I've only seen it on one of four open tabs and it needed to refresh.

+        print("Unexpected error: Processing could not begin, you may need to refresh the tab or restart the service.")

+        return Processed(p, [])

+    elif not any(processed_result.images):

         print("Unexpected error: draw_xyz_grid failed to return even a single processed image")

         return Processed(p, [])

 

-    sub_grids = [None] * len(zs)

-    for i in range(len(zs)):

-        start_index = i * len(xs) * len(ys)

+    z_count = len(zs)

+    sub_grids = [None] * z_count

+    for i in range(z_count):

+        start_index = (i * len(xs) * len(ys)) + i

         end_index = start_index + len(xs) * len(ys)

-        grid = images.image_grid(image_cache[start_index:end_index], rows=len(ys))

+        grid = images.image_grid(processed_result.images[start_index:end_index], rows=len(ys))

         if draw_legend:

-            grid = images.draw_grid_annotations(grid, cell_size[0], cell_size[1], hor_texts, ver_texts, margin_size)

-        sub_grids[i] = grid

-        if include_sub_grids and len(zs) > 1:

-            processed_result.images.insert(i+1, grid)

-

-    sub_grid_size = sub_grids[0].size

-    z_grid = images.image_grid(sub_grids, rows=1)

+            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)

+        processed_result.images.insert(i, grid)

+        processed_result.all_prompts.insert(i, processed_result.all_prompts[start_index])

+        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

+    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()]])

-    processed_result.images[0] = z_grid

+    processed_result.images.insert(0, z_grid)

+    #TODO: Deeper aspects of the program rely on grid info being misaligned between metadata arrays, which is not ideal.

+    #processed_result.all_prompts.insert(0, processed_result.all_prompts[0])

+    #processed_result.all_seeds.insert(0, processed_result.all_seeds[0])

+    processed_result.infotexts.insert(0, processed_result.infotexts[0])

 

-    return processed_result, sub_grids

+    return processed_result

 

 

 class SharedSettingsStackHelper(object):

@@ -540,7 +548,7 @@ class Script(scripts.Script):
         # If one of the axes is very slow to change between (like SD model

         # checkpoint), then make sure it is in the outer iteration of the nested

         # `for` loop.

-        first_axes_processed = 'x'

+        first_axes_processed = 'z'

         second_axes_processed = 'y'

         if x_opt.cost > y_opt.cost and x_opt.cost > z_opt.cost:

             first_axes_processed = 'x'

@@ -600,7 +608,7 @@ class Script(scripts.Script):
             return res

 

         with SharedSettingsStackHelper():

-            processed, sub_grids = draw_xyz_grid(

+            processed = draw_xyz_grid(

                 p,

                 xs=xs,

                 ys=ys,

@@ -617,11 +625,30 @@ class Script(scripts.Script):
                 margin_size=margin_size

             )

 

-        if opts.grid_save and len(sub_grids) > 1:

-            for sub_grid in sub_grids:

-                images.save_image(sub_grid, p.outpath_grids, "xyz_grid", info=grid_infotext[0], extension=opts.grid_format, prompt=p.prompt, seed=processed.seed, grid=True, p=p)

+        if not processed.images:

+            # It broke, no further handling needed.

+            return processed

+

+        z_count = len(zs)

+

+        if not include_lone_images:

+            # Don't need sub-images anymore, drop from list:

+            processed.images = processed.images[:z_count+1]

 

         if opts.grid_save:

-            images.save_image(processed.images[0], p.outpath_grids, "xyz_grid", info=grid_infotext[0], extension=opts.grid_format, prompt=p.prompt, seed=processed.seed, grid=True, p=p)

+            # Auto-save main and sub-grids:

+            grid_count = z_count + 1 if z_count > 1 else 1

+            for g in range(grid_count):

+                #TODO: See previous comment about intentional data misalignment.

+                adj_g = g-1 if g > 0 else g

+                images.save_image(processed.images[g], p.outpath_grids, "xyz_grid", info=processed.infotexts[g], extension=opts.grid_format, prompt=processed.all_prompts[adj_g], seed=processed.all_seeds[adj_g], grid=True, p=processed)

+

+        if not include_sub_grids:

+            # Done with sub-grids, drop all related information:

+            for sg in range(z_count):

+                del processed.images[1]

+                del processed.all_prompts[1]

+                del processed.all_seeds[1]

+                del processed.infotexts[1]

 

         return processed