sag prototype and allow xyz continue on error

CHANGELOG.md

@@ -25,15 +25,18 @@ And it also includes fixes for all reported issues so far
   - fix correct image mode  
   - fix batch/folder/video modes  
 - **Improvements**  
-  - allow deployment without git clone  
+  - **offline deployment**: allow deployment without git clone  
     for example, you can now deploy a zip of the sdnext folder  
-  - hypertile: enable vae tiling  
-  - hypertile: add autodetect optimial value  
-    set tile size to 0 to use autodetected value  
-  - cli: sdapi.py allow manual api invoke  
+  - **xyz grid**: continue on error  
+    now you can use xyz grid with different params and test which ones work and which dont  
+  - **hypertile**  
+    - enable vae tiling  
+    - add autodetect optimial value  
+      set tile size to 0 to use autodetected value  
+  - **cli**: sdapi.py allow manual api invoke  
     example: `python cli/sdapi.py /sdapi/v1/sd-models`  
-  - memory: add ram usage monitoring in addition to gpu memory usage monitoring  
-  - vae: enable taesd batch decode  
+  - **memory**: add ram usage monitoring in addition to gpu memory usage monitoring  
+  - **vae**: enable taesd batch decode  
     enable/disable with settings -> diffusers > vae slicing  
   - updated core requirements  
   - major internal ui module refactoring  

modules/control/run.py

@@ -64,7 +64,7 @@ def restore_pipeline():
 def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_generator: bool, input_type: int,
                 prompt, negative, styles, steps, sampler_index,
                 seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w,
-                cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, full_quality, restore_faces, tiling,
+                cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, full_quality, restore_faces, tiling,
                 hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry,
                 resize_mode, resize_name, width, height, scale_by, selected_scale_tab, resize_time,
                 denoising_strength, batch_count, batch_size,
@@ -103,6 +103,7 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
         clip_skip = clip_skip,
         image_cfg_scale = image_cfg_scale,
         diffusers_guidance_rescale = diffusers_guidance_rescale,
+        sag_scale = sag_scale,
         full_quality = full_quality,
         restore_faces = restore_faces,
         tiling = tiling,

modules/img2img.py

@@ -121,6 +121,7 @@ def img2img(id_task: str, mode: int,
             n_iter, batch_size,
             cfg_scale, image_cfg_scale,
             diffusers_guidance_rescale,
+            sag_scale,
             refiner_start,
             clip_skip,
             denoising_strength,
@@ -229,6 +230,7 @@ def img2img(id_task: str, mode: int,
         denoising_strength=denoising_strength,
         image_cfg_scale=image_cfg_scale,
         diffusers_guidance_rescale=diffusers_guidance_rescale,
+        sag_scale=sag_scale,
         refiner_start=refiner_start,
         inpaint_full_res=inpaint_full_res != 0,
         inpaint_full_res_padding=inpaint_full_res_padding,
@@ -258,5 +260,5 @@ def img2img(id_task: str, mode: int,
         if processed is None:
             processed = processing.process_images(p)
     p.close()
-    generation_info_js = processed.js()
+    generation_info_js = processed.js() if processed is not None else ''
     return processed.images, generation_info_js, processed.info, plaintext_to_html(processed.comments)

modules/processing.py

@@ -137,7 +137,7 @@ class StableDiffusionProcessing:
     """
     The first set of paramaters: sd_models -> do_not_reload_embeddings represent the minimum required to create a StableDiffusionProcessing
     """
-    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt: str = "", styles: List[str] = None, seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1, seed_enable_extras: bool = True, sampler_name: str = None, hr_sampler_name: str = None, batch_size: int = 1, n_iter: int = 1, steps: int = 50, cfg_scale: float = 7.0, image_cfg_scale: float = None, clip_skip: int = 1, width: int = 512, height: int = 512, full_quality: bool = True, restore_faces: bool = False, tiling: bool = False, do_not_save_samples: bool = False, do_not_save_grid: bool = False, extra_generation_params: Dict[Any, Any] = None, overlay_images: Any = None, negative_prompt: str = None, eta: float = None, do_not_reload_embeddings: bool = False, denoising_strength: float = 0, diffusers_guidance_rescale: float = 0.7, resize_mode: int = 0, resize_name: str = 'None', scale_by: float = 0, selected_scale_tab: int = 0, hdr_clamp: bool = False, hdr_boundary: float = 4.0, hdr_threshold: float = 3.5, hdr_center: bool = False, hdr_channel_shift: float = 0.8, hdr_full_shift: float = 0.8, hdr_maximize: bool = False, hdr_max_center: float = 0.6, hdr_max_boundry: float = 1.0, override_settings: Dict[str, Any] = None, override_settings_restore_afterwards: bool = True, sampler_index: int = None, script_args: list = None): # pylint: disable=unused-argument
+    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt: str = "", styles: List[str] = None, seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1, seed_enable_extras: bool = True, sampler_name: str = None, hr_sampler_name: str = None, batch_size: int = 1, n_iter: int = 1, steps: int = 50, cfg_scale: float = 7.0, image_cfg_scale: float = None, clip_skip: int = 1, width: int = 512, height: int = 512, full_quality: bool = True, restore_faces: bool = False, tiling: bool = False, do_not_save_samples: bool = False, do_not_save_grid: bool = False, extra_generation_params: Dict[Any, Any] = None, overlay_images: Any = None, negative_prompt: str = None, eta: float = None, do_not_reload_embeddings: bool = False, denoising_strength: float = 0, diffusers_guidance_rescale: float = 0.7, sag_scale: float = 0.0, resize_mode: int = 0, resize_name: str = 'None', scale_by: float = 0, selected_scale_tab: int = 0, hdr_clamp: bool = False, hdr_boundary: float = 4.0, hdr_threshold: float = 3.5, hdr_center: bool = False, hdr_channel_shift: float = 0.8, hdr_full_shift: float = 0.8, hdr_maximize: bool = False, hdr_max_center: float = 0.6, hdr_max_boundry: float = 1.0, override_settings: Dict[str, Any] = None, override_settings_restore_afterwards: bool = True, sampler_index: int = None, script_args: list = None): # pylint: disable=unused-argument
         self.outpath_samples: str = outpath_samples
         self.outpath_grids: str = outpath_grids
         self.prompt: str = prompt
@@ -159,6 +159,7 @@ class StableDiffusionProcessing:
         self.scale_by: float = scale_by
         self.image_cfg_scale = image_cfg_scale
         self.diffusers_guidance_rescale = diffusers_guidance_rescale
+        self.sag_scale = sag_scale
         if devices.backend == "ipex" and width == 1024 and height == 1024 and os.environ.get('DISABLE_IPEX_1024_WA', None) is None:
             width = 1080
             height = 1080

modules/processing_diffusers.py

@@ -5,6 +5,7 @@ import inspect
 import typing
 import torch
 import torchvision.transforms.functional as TF
+import diffusers
 import modules.devices as devices
 import modules.shared as shared
 import modules.sd_samplers as sd_samplers
@@ -368,6 +369,13 @@ def process_diffusers(p: StableDiffusionProcessing):
             # TODO extra_generation_params add sampler options
             # p.extra_generation_params['Sampler options'] = ''
 
+    def update_pipeline(sd_model, p: StableDiffusionProcessing):
+        if p.sag_scale > 0 and is_txt2img():
+            sd_model = sd_models.switch_diffuser_pipe(sd_model, diffusers.StableDiffusionSAGPipeline)
+            p.extra_generation_params["SAG scale"] = p.sag_scale
+            p.task_args['sag_scale'] = p.sag_scale
+        return sd_model
+
     if len(getattr(p, 'init_images', [])) > 0:
         while len(p.init_images) < len(p.prompts):
             p.init_images.append(p.init_images[-1])
@@ -431,6 +439,7 @@ def process_diffusers(p: StableDiffusionProcessing):
         debug_steps(f'Steps: type=refiner input={p.refiner_steps} output={steps} start={p.refiner_start} denoise={p.denoising_strength}')
         return max(2, int(steps))
 
+    shared.sd_model = update_pipeline(shared.sd_model, p)
     base_args = set_pipeline_args(
         model=shared.sd_model,
         prompts=p.prompts,

modules/sd_models.py

@@ -15,6 +15,7 @@ from enum import Enum
 from rich import progress # pylint: disable=redefined-builtin
 import torch
 import safetensors.torch
+import diffusers
 from omegaconf import OmegaConf
 import tomesd
 from transformers import logging as transformers_logging
@@ -25,11 +26,6 @@ from modules.timer import Timer
 from modules.memstats import memory_stats
 from modules.paths import models_path, script_path
 
-try:
-    import diffusers
-except Exception as ex:
-    shared.log.error(f'Failed to import diffusers: {ex}')
-
 
 transformers_logging.set_verbosity_error()
 model_dir = "Stable-diffusion"
@@ -692,7 +688,6 @@ def copy_diffuser_options(new_pipe, orig_pipe):
     new_pipe.is_sd1 = getattr(orig_pipe, 'is_sd1', True)
 
 
-
 def set_diffuser_options(sd_model, vae = None, op: str = 'model'):
     if sd_model is None:
         shared.log.warning(f'{op} is not loaded')
@@ -1032,6 +1027,47 @@ def get_diffusers_task(pipe: diffusers.DiffusionPipeline) -> DiffusersTaskType:
         return DiffusersTaskType.TEXT_2_IMAGE
 
 
+def switch_diffuser_pipe(pipeline, cls):
+    try:
+        new_pipe = None
+        if isinstance(pipeline, cls):
+            return pipeline
+        elif isinstance(pipeline, diffusers.StableDiffusionXLPipeline):
+            new_pipe = cls(
+                vae=pipeline.vae,
+                text_encoder=pipeline.text_encoder,
+                text_encoder_2=pipeline.text_encoder_2,
+                tokenizer=pipeline.tokenizer,
+                tokenizer_2=pipeline.tokenizer_2,
+                unet=pipeline.unet,
+                scheduler=pipeline.scheduler,
+                feature_extractor=getattr(pipeline, 'feature_extractor', None),
+            ).to(pipeline.device)
+        elif isinstance(pipeline, diffusers.StableDiffusionPipeline):
+            new_pipe = cls(
+                vae=pipeline.vae,
+                text_encoder=pipeline.text_encoder,
+                tokenizer=pipeline.tokenizer,
+                unet=pipeline.unet,
+                scheduler=pipeline.scheduler,
+                feature_extractor=getattr(pipeline, 'feature_extractor', None),
+                requires_safety_checker=False,
+                safety_checker=None,
+            ).to(pipeline.device)
+        else:
+            shared.log.error(f'Pipeline switch error: {pipeline.__class__.__name__} unrecognized')
+            return pipeline
+        if new_pipe is not None:
+            copy_diffuser_options(new_pipe, pipeline)
+            shared.log.debug(f'Pipeline switch: from={pipeline.__class__.__name__} to={new_pipe.__class__.__name__}')
+            return new_pipe
+        else:
+            shared.log.error(f'Pipeline switch error: from={pipeline.__class__.__name__} to={cls.__name__} empty pipeline')
+    except Exception as e:
+        shared.log.error(f'Pipeline switch error: from={pipeline.__class__.__name__} to={cls.__name__} {e}')
+    return pipeline
+
+
 def set_diffuser_pipe(pipe, new_pipe_type):
     sd_checkpoint_info = getattr(pipe, "sd_checkpoint_info", None)
     sd_model_checkpoint = getattr(pipe, "sd_model_checkpoint", None)

modules/txt2img.py

@@ -14,7 +14,7 @@ def txt2img(id_task,
             steps, sampler_index, hr_sampler_index,
             full_quality, restore_faces, tiling,
             n_iter, batch_size,
-            cfg_scale, image_cfg_scale, diffusers_guidance_rescale,
+            cfg_scale, image_cfg_scale, diffusers_guidance_rescale, sag_scale,
             clip_skip,
             seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w,
             height, width,
@@ -58,6 +58,7 @@ def txt2img(id_task,
         cfg_scale=cfg_scale,
         image_cfg_scale=image_cfg_scale,
         diffusers_guidance_rescale=diffusers_guidance_rescale,
+        sag_scale=sag_scale,
         clip_skip=clip_skip,
         width=width,
         height=height,
@@ -89,5 +90,5 @@ def txt2img(id_task,
     p.close()
     if processed is None:
         return [], '', '', 'Error: processing failed'
-    generation_info_js = processed.js()
+    generation_info_js = processed.js() if processed is not None else ''
     return processed.images, generation_info_js, processed.info, plaintext_to_html(processed.comments)

modules/ui_control.py

@@ -337,7 +337,7 @@ def create_ui(_blocks: gr.Blocks=None):
 
                 batch_count, batch_size = ui_sections.create_batch_inputs('control')
                 seed, _reuse_seed, subseed, _reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w = ui_sections.create_seed_inputs('control', reuse_visible=False)
-                cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_advanced_inputs('control')
+                cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_advanced_inputs('control')
 
                 with gr.Accordion(open=False, label="Video", elem_id="control_video", elem_classes=["small-accordion"]):
                     with gr.Row():
@@ -707,7 +707,7 @@ def create_ui(_blocks: gr.Blocks=None):
                     prompt, negative, styles,
                     steps, sampler_index,
                     seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w,
-                    cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry,
+                    cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry,
                     resize_mode, resize_name, width, height, scale_by, selected_scale_tab, resize_time,
                     denoising_strength, batch_count, batch_size,
                     video_skip_frames, video_type, video_duration, video_loop, video_pad, video_interpolate,

modules/ui_img2img.py

@@ -142,7 +142,7 @@ def create_ui():
                             denoising_strength = gr.Slider(minimum=0.0, maximum=0.99, step=0.01, label='Denoising strength', value=0.50, elem_id="img2img_denoising_strength")
                             refiner_start = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Denoise start', value=0.0, elem_id="img2img_refiner_start")
 
-                    cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_advanced_inputs('img2img')
+                    cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_advanced_inputs('img2img')
 
                     with FormGroup(elem_id="inpaint_controls", visible=False) as inpaint_controls:
                         with FormRow():
@@ -195,7 +195,7 @@ def create_ui():
                 full_quality, restore_faces, tiling,
                 batch_count, batch_size,
                 cfg_scale, image_cfg_scale,
-                diffusers_guidance_rescale,
+                diffusers_guidance_rescale, sag_scale,
                 refiner_start,
                 clip_skip,
                 denoising_strength,

modules/ui_sections.py

@@ -105,10 +105,13 @@ def create_advanced_inputs(tab):
         with gr.Group():
             with FormRow():
                 cfg_scale = gr.Slider(minimum=0.0, maximum=30.0, step=0.1, label='CFG scale', value=6.0, elem_id=f"{tab}_cfg_scale")
-                clip_skip = gr.Slider(label='CLIP skip', value=1, minimum=1, maximum=14, step=1, elem_id=f"{tab}_clip_skip", interactive=True)
-            with FormRow():
                 image_cfg_scale = gr.Slider(minimum=0.0, maximum=30.0, step=0.1, label='Secondary CFG scale', value=6.0, elem_id=f"{tab}_image_cfg_scale")
+            with FormRow():
                 diffusers_guidance_rescale = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Guidance rescale', value=0.7, elem_id=f"{tab}_image_cfg_rescale", visible=shared.backend == shared.Backend.DIFFUSERS)
+                # diffusers_sag_scale = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Self-attention guidance', value=0.0, elem_id=f"{tab}_image_sag_scale", visible=shared.backend == shared.Backend.DIFFUSERS) # TODO enable SAG once fixed in diffusers
+                diffusers_sag_scale = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Self-attention guidance', value=0.0, elem_id=f"{tab}_image_sag_scale", visible=False)
+            with FormRow():
+                clip_skip = gr.Slider(label='CLIP skip', value=1, minimum=1, maximum=14, step=1, elem_id=f"{tab}_clip_skip", interactive=True)
         with gr.Group():
             with FormRow():
                 full_quality = gr.Checkbox(label='Full quality', value=True, elem_id=f"{tab}_full_quality")
@@ -127,7 +130,7 @@ def create_advanced_inputs(tab):
                 hdr_maximize = gr.Checkbox(label='HDR maximize', value=False, elem_id=f"{tab}_hdr_maximize")
                 hdr_max_center = gr.Slider(minimum=0.0, maximum=2.0, step=0.1, value=0.6,  label='Center', elem_id=f"{tab}_hdr_max_center")
                 hdr_max_boundry = gr.Slider(minimum=0.5, maximum=2.0, step=0.1, value=1.0,  label='Range', elem_id=f"{tab}_hdr_max_boundry")
-    return cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry
+    return cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, diffusers_sag_scale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry
 
 
 def create_sampler_and_steps_selection(choices, tabname):

modules/ui_txt2img.py

@@ -44,7 +44,7 @@ def create_ui():
                     steps, sampler_index = ui_sections.create_sampler_inputs('txt2img')
                     batch_count, batch_size = ui_sections.create_batch_inputs('txt2img')
                     seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w = ui_sections.create_seed_inputs('txt2img')
-                    cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_advanced_inputs('txt2img')
+                    cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_advanced_inputs('txt2img')
                     enable_hr, hr_sampler_index, denoising_strength, hr_final_resolution, hr_upscaler, hr_force, hr_second_pass_steps, hr_scale, hr_resize_x, hr_resize_y, refiner_steps, refiner_start, refiner_prompt, refiner_negative = ui_sections.create_hires_inputs('txt2img')
                     override_settings = ui_common.create_override_inputs('txt2img')
 
@@ -71,7 +71,7 @@ def create_ui():
                 steps, sampler_index, hr_sampler_index,
                 full_quality, restore_faces, tiling,
                 batch_count, batch_size,
-                cfg_scale, image_cfg_scale, diffusers_guidance_rescale,
+                cfg_scale, image_cfg_scale, diffusers_guidance_rescale, sag_scale,
                 clip_skip,
                 seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w,
                 height, width,

scripts/xyz_grid.py

@@ -10,10 +10,7 @@ from io import StringIO
 from PIL import Image
 import numpy as np
 import gradio as gr
-import modules.scripts as scripts
-import modules.shared as shared
-from modules import images, sd_samplers, processing, sd_models, sd_vae
-from modules.processing import process_images, Processed, StableDiffusionProcessingTxt2Img
+from modules import shared, errors, scripts, images, sd_samplers, processing, sd_models, sd_vae
 from modules.ui_components import ToolButton
 import modules.ui_symbols as symbols
 
@@ -133,11 +130,11 @@ def apply_vae(p, x, xs):
     sd_vae.reload_vae_weights(shared.sd_model, vae_file=find_vae(x))
 
 
-def apply_styles(p: StableDiffusionProcessingTxt2Img, x: str, _):
+def apply_styles(p: processing.StableDiffusionProcessingTxt2Img, x: str, _):
     p.styles.extend(x.split(','))
 
 
-def apply_upscaler(p: StableDiffusionProcessingTxt2Img, opt, x):
+def apply_upscaler(p: processing.StableDiffusionProcessingTxt2Img, opt, x):
     p.enable_hr = True
     p.hr_force = True
     p.denoising_strength = 0.0
@@ -289,7 +286,7 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
             return ix + iy * len(xs) + iz * len(xs) * len(ys)
 
         shared.state.job = 'grid'
-        processed: Processed = cell(x, y, z, ix, iy, iz)
+        processed: processing.Processed = cell(x, y, z, ix, iy, iz)
         if processed_result is None:
             processed_result = copy(processed)
             processed_result.images = [None] * list_size
@@ -298,7 +295,7 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
             processed_result.infotexts = [None] * list_size
             processed_result.index_of_first_image = 1
         idx = index(ix, iy, iz)
-        if processed.images:
+        if processed is not None and processed.images:
             processed_result.images[idx] = processed.images[0]
             processed_result.all_prompts[idx] = processed.prompt
             processed_result.all_seeds[idx] = processed.seed
@@ -344,10 +341,10 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
 
     if not processed_result:
         shared.log.error("XYZ grid: Failed to initialize processing")
-        return Processed(p, [])
+        return processing.Processed(p, [])
     elif not any(processed_result.images):
         shared.log.error("XYZ grid: Failed to return processed image")
-        return Processed(p, [])
+        return processing.Processed(p, [])
 
     z_count = len(zs)
     for i in range(z_count):
@@ -617,7 +614,7 @@ class Script(scripts.Script):
             total_steps = sum(zs) * len(xs) * len(ys)
         else:
             total_steps = p.steps * len(xs) * len(ys) * len(zs)
-        if isinstance(p, StableDiffusionProcessingTxt2Img) and p.enable_hr:
+        if isinstance(p, processing.StableDiffusionProcessingTxt2Img) and p.enable_hr:
             if x_opt.label == "Hires steps":
                 total_steps += sum(xs) * len(ys) * len(zs)
             elif y_opt.label == "Hires steps":
@@ -660,14 +657,19 @@ class Script(scripts.Script):
 
         def cell(x, y, z, ix, iy, iz):
             if shared.state.interrupted:
-                return Processed(p, [], p.seed, "")
+                return processing.Processed(p, [], p.seed, "")
             pc = copy(p)
             pc.override_settings_restore_afterwards = False
             pc.styles = pc.styles[:]
             x_opt.apply(pc, x, xs)
             y_opt.apply(pc, y, ys)
             z_opt.apply(pc, z, zs)
-            res = process_images(pc)
+            try:
+                res = processing.process_images(pc)
+            except Exception as e:
+                shared.log.error(f"XYZ grid: Failed to process image: {e}")
+                errors.display(e, 'XYZ grid')
+                res = None
             subgrid_index = 1 + iz # Sets subgrid infotexts
             if grid_infotext[subgrid_index] is None and ix == 0 and iy == 0:
                 pc.extra_generation_params = copy(pc.extra_generation_params)