fix face-hires with control

CHANGELOG.md

@@ -59,6 +59,7 @@
   - variation seed possible NaNs
   - enumerate diffusers model with multiple variants
   - diffusers skip non-models on enum
+  - face-hires compatibility with control modules
   - css fixes
 
 ## Update for 2024-03-19

TODO.md

@@ -31,6 +31,7 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma
 - include reference styles
 - quick apply style
 - lora: sc lora, dora, etc
+- samplers: smea <https://github.com/Koishi-Star/Euler-Smea-Dyn-Sampler>, restart <https://github.com/Newbeeer/diffusion_restart_sampling>
 
 ### Missing
 

modules/processing.py

@@ -332,8 +332,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                 return create_infotext(p, p.prompts, p.seeds, p.subseeds, index=index, all_negative_prompts=p.negative_prompts)
 
             for i, x_sample in enumerate(x_samples_ddim):
-                if hasattr(p, 'recursion'):
-                    continue
                 debug(f'Processing result: index={i+1}/{len(x_samples_ddim)} iteration={n+1}/{p.n_iter}')
                 p.batch_index = i
                 if type(x_sample) == Image.Image:
@@ -347,11 +345,13 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                         images.save_image(Image.fromarray(x_sample), path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=infotext(i), p=p, suffix="-before-face-restore")
                     p.ops.append('face')
                     x_sample = face_restoration.restore_faces(x_sample, p)
-                    image = Image.fromarray(x_sample)
+                    if x_sample is not None:
+                        image = Image.fromarray(x_sample)
                 if p.scripts is not None and isinstance(p.scripts, scripts.ScriptRunner):
                     pp = scripts.PostprocessImageArgs(image)
                     p.scripts.postprocess_image(p, pp)
-                    image = pp.image
+                    if pp.image is not None:
+                        image = pp.image
                 if p.color_corrections is not None and i < len(p.color_corrections):
                     if not p.do_not_save_samples and shared.opts.save_images_before_color_correction:
                         orig = p.color_corrections
@@ -376,7 +376,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
                     image2 = Image.new('RGBa', image.size)
                     mask = images.resize_image(3, p.mask_for_overlay, image.width, image.height).convert('L')
                     image_mask_composite = Image.composite(image1, image2, mask).convert('RGBA')
-                    image_mask_composite.save('/tmp/composite.png')
                     if shared.opts.save_mask:
                         images.save_image(image_mask, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=text, p=p, suffix="-mask")
                     if shared.opts.save_mask_composite:

scripts/face-details.py

@@ -91,6 +91,8 @@ class FaceRestorerYolo(FaceRestoration):
 
     def restore(self, np_image, p: processing.StableDiffusionProcessing = None):
         from modules import devices, processing_class
+        if hasattr(p, 'recursion'):
+            return
         if not hasattr(p, 'facehires'):
             p.facehires = 0
         if np_image is None or p.facehires >= p.batch_size * p.n_iter:
@@ -128,6 +130,10 @@ class FaceRestorerYolo(FaceRestoration):
             'inpaint_full_res_padding': 15,
             'restore_faces': True,
         }
+        if getattr(p, 'is_control', False):
+            from modules.control import run
+            run.restore_pipeline()
+
         p = processing_class.switch_class(p, processing.StableDiffusionProcessingImg2Img, args)
         p.facehires += 1 # set flag to avoid recursion
 
@@ -139,6 +145,7 @@ class FaceRestorerYolo(FaceRestoration):
             p.negative_prompt = orig_p.get('all_negative_prompts', [''])[0]
 
         shared.log.debug(f'Face HiRes: faces={[f.__dict__ for f in faces]} strength={p.denoising_strength} blur={p.mask_blur} padding={p.inpaint_full_res_padding} steps={p.steps}')
+
         for face in faces:
             if face.mask is None:
                 continue