fix restore pipeline

javascript/sdnext.css

@@ -28,7 +28,7 @@ input[type='color'] { width: 64px; height: 32px; }
 .gradio-button.secondary-down { background: var(--button-secondary-background-fill); color: var(--button-secondary-text-color); }
 .gradio-button.secondary-down, .gradio-button.secondary-down:hover { box-shadow: 1px 1px 1px rgba(0,0,0,0.25) inset, 0px 0px 3px rgba(0,0,0,0.15) inset; }
 .gradio-button.secondary-down:hover { background: var(--button-secondary-background-fill-hover); color: var(--button-secondary-text-color-hover); }
-.gradio-button.tool { max-width: min-content; min-width: min-content !important; align-self: end; font-size: 20px !important; color: var(--body-text-color) !important; align-self: center; }
+.gradio-button.tool { max-width: min-content; min-width: min-content !important; align-self: end; font-size: 20px !important; color: var(--body-text-color) !important; align-self: end; margin-bottom: 8px; }
 .gradio-checkbox { margin: 0.75em 1.5em 0 0; align-self: center; }
 .gradio-column { min-width: min(160px, 100%) !important; }
 .gradio-container { max-width: unset !important; padding: var(--block-label-padding) !important; }

modules/control/run.py

@@ -19,15 +19,19 @@ from modules.processing_class import StableDiffusionProcessingControl
 debug = shared.log.trace if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
 debug('Trace: CONTROL')
 pipe = None
+instance = None
 original_pipeline = None
 
 
 def restore_pipeline():
-    global pipe # pylint: disable=global-statement
-    pipe = None
+    global pipe, instance # pylint: disable=global-statement
+    if instance is not None and hasattr(instance, 'restore'):
+        instance.restore()
     if original_pipeline is not None:
         shared.sd_model = original_pipeline
-        debug(f'Control restored pipeline: class={shared.sd_model.__class__.__name__}')
+        shared.log.debug(f'Control restored pipeline: class={shared.sd_model.__class__.__name__}')
+    pipe = None
+    instance = None
     devices.torch_gc()
 
 
@@ -43,7 +47,7 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
                 video_skip_frames, video_type, video_duration, video_loop, video_pad, video_interpolate,
                 *input_script_args # pylint: disable=unused-argument
         ):
-    global pipe, original_pipeline # pylint: disable=global-statement
+    global instance, pipe, original_pipeline # pylint: disable=global-statement
     debug(f'Control: type={unit_type} input={inputs} init={inits} type={input_type}')
     if inputs is None or (type(inputs) is list and len(inputs) == 0):
         inputs = [None]
@@ -244,6 +248,7 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
     index = 0
     frames = 0
 
+    # set pipeline
     original_pipeline = shared.sd_model
     shared.sd_model = pipe
     sd_models.move_model(shared.sd_model, shared.device)
@@ -473,6 +478,7 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
                     # pipeline
                     output = None
                     if pipe is not None: # run new pipeline
+                        pipe.restore_pipeline = restore_pipeline
                         debug(f'Control exec pipeline: task={sd_models.get_diffusers_task(pipe)} class={pipe.__class__}')
                         debug(f'Control exec pipeline: p={vars(p)}')
                         debug(f'Control exec pipeline: args={p.task_args} image={p.task_args.get("image", None)} control={p.task_args.get("control_image", None)} mask={p.task_args.get("mask_image", None) or p.image_mask} ref={p.task_args.get("ref_image", None)}')
@@ -526,10 +532,6 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
         shared.log.error(f'Control pipeline failed: type={unit_type} units={len(active_model)} error={e}')
         errors.display(e, 'Control')
 
-    shared.sd_model = original_pipeline
-    pipe = None
-    devices.torch_gc()
-
     if len(output_images) == 0:
         output_images = None
         image_txt = 'images=None'
@@ -543,8 +545,6 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
         image_txt = f'| Frames {len(output_images)} | Size {output_images[0].width}x{output_images[0].height}'
 
     image_txt += f' | {util.dict2str(p.extra_generation_params)}'
-    if hasattr(instance, 'restore'):
-        instance.restore()
     restore_pipeline()
     debug(f'Control ready: {image_txt}')
     if is_generator:

modules/processing.py

@@ -316,6 +316,10 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             def infotext(index): # pylint: disable=function-redefined # noqa: F811
                 return create_infotext(p, p.prompts, p.seeds, p.subseeds, index=index, all_negative_prompts=p.negative_prompts)
 
+            if hasattr(shared.sd_model, 'restore_pipeline') and shared.sd_model.restore_pipeline is not None:
+                print('HERE')
+                shared.sd_model.restore_pipeline()
+
             for i, x_sample in enumerate(x_samples_ddim):
                 p.batch_index = i
                 if type(x_sample) == Image.Image:

modules/ui_control_helpers.py

@@ -126,9 +126,7 @@ def select_input(input_mode, input_image, init_image, init_type, input_resize, i
     if isinstance(selected_input, Image.Image): # image via upload -> image
         if input_mode == 'Outpaint':
             masking.opts.invert = True
-            selected_input, outpaint_mask = masking.outpaint(input_image=selected_input)
-        if outpaint_mask is not None:
-            input_mask = outpaint_mask
+            selected_input, input_mask = masking.outpaint(input_image=selected_input)
         input_source = [selected_input]
         input_type = 'PIL.Image'
         status = f'Control input | Image | Size {selected_input.width}x{selected_input.height} | Mode {selected_input.mode}'