pipeline restore args after batch and restore pipeline after base

Signed-off-by: Vladimir Mandic <[email protected]>

CHANGELOG.md

@@ -1,6 +1,6 @@
 # Change Log for SD.Next
 
-## Update for 2025-01-08
+## Update for 2025-01-09
 
 - [Allegro Video](https://huggingface.co/rhymes-ai/Allegro)  
   - optimizations: full offload and quantization support  
@@ -53,6 +53,8 @@
   - lora diffusers method apply only once  
   - lora diffusers method set prompt tags and metadata  
   - flux support on-the-fly quantization for bnb of unet only  
+  - control restore pipeline before running hires  
+  - restore args after batch run  
 
 ## Update for 2024-12-31
 

modules/control/run.py

@@ -188,7 +188,7 @@ def check_enabled(p, unit_type, units, active_model, active_strength, active_sta
             selected_models = None
         elif len(active_model) == 1:
             selected_models = active_model[0].model if active_model[0].model is not None else None
-            p.is_tile = p.is_tile or 'tile' in active_model[0].model_id.lower()
+            p.is_tile = p.is_tile or 'tile' in (active_model[0].model_id or '').lower()
             has_models = selected_models is not None
             control_conditioning = active_strength[0] if len(active_strength) > 0 else 1 # strength or list[strength]
             control_guidance_start = active_start[0] if len(active_start) > 0 else 0
@@ -687,6 +687,7 @@ def control_run(state: str = '',
                     if pipe is not None: # run new pipeline
                         if not hasattr(pipe, 'restore_pipeline') and video is None:
                             pipe.restore_pipeline = restore_pipeline
+                            shared.sd_model.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)}')

modules/processing_args.py

@@ -1,6 +1,7 @@
 import typing
 import os
 import re
+import copy
 import math
 import time
 import inspect
@@ -122,7 +123,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
 
     if debug_enabled:
         debug_log(f'Diffusers pipeline possible: {possible}')
-    prompts, negative_prompts, prompts_2, negative_prompts_2 = fix_prompts(prompts, negative_prompts, prompts_2, negative_prompts_2)
+    prompts, negative_prompts, prompts_2, negative_prompts_2 = fix_prompts(p, prompts, negative_prompts, prompts_2, negative_prompts_2)
     steps = kwargs.get("num_inference_steps", None) or len(getattr(p, 'timesteps', ['1']))
     clip_skip = kwargs.pop("clip_skip", 1)
 
@@ -278,7 +279,10 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
         args['callback'] = diffusers_callback_legacy
 
     if 'image' in kwargs:
-        p.init_images = kwargs['image'] if isinstance(kwargs['image'], list) else [kwargs['image']]
+        if isinstance(kwargs['image'], list) and isinstance(kwargs['image'][0], Image.Image):
+            p.init_images = kwargs['image']
+        if isinstance(kwargs['image'], Image.Image):
+            p.init_images = [kwargs['image']]
 
     # handle remaining args
     for arg in kwargs:
@@ -360,4 +364,6 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
         shared.log.debug(f'Profile: pipeline args: {t1-t0:.2f}')
     if debug_enabled:
         debug_log(f'Diffusers pipeline args: {args}')
-    return args
+
+    _args = copy.deepcopy(args) # pipeline may modify underlying args
+    return _args

modules/processing_diffusers.py

@@ -189,7 +189,7 @@ def process_hires(p: processing.StableDiffusionProcessing, output):
 
         # hires
         if p.hr_force and strength == 0:
-            shared.log.warning('HiRes skip: denoising=0')
+            shared.log.warning('Hires skip: denoising=0')
             p.hr_force = False
         if p.hr_force:
             shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, sd_models.DiffusersTaskType.IMAGE_2_IMAGE)

modules/processing_helpers.py

@@ -428,11 +428,16 @@ def resize_hires(p, latents): # input=latents output=pil if not latent_upscaler
     return resized_images
 
 
-def fix_prompts(prompts, negative_prompts, prompts_2, negative_prompts_2):
+def fix_prompts(p, prompts, negative_prompts, prompts_2, negative_prompts_2):
     if type(prompts) is str:
         prompts = [prompts]
     if type(negative_prompts) is str:
         negative_prompts = [negative_prompts]
+    if hasattr(p, '[init_images]') and p.init_images is not None and len(p.init_images) > 1:
+        while len(prompts) < len(p.init_images):
+            prompts.append(prompts[-1])
+        while len(negative_prompts) < len(p.init_images):
+            negative_prompts.append(negative_prompts[-1])
     while len(negative_prompts) < len(prompts):
         negative_prompts.append(negative_prompts[-1])
     while len(prompts) < len(negative_prompts):

modules/sd_models.py

@@ -1243,6 +1243,7 @@ def set_diffuser_pipe(pipe, new_pipe_type):
     image_encoder = getattr(pipe, "image_encoder", None)
     feature_extractor = getattr(pipe, "feature_extractor", None)
     mask_processor = getattr(pipe, "mask_processor", None)
+    restore_pipeline = getattr(pipe, "restore_pipeline", None)
 
     if new_pipe is None:
         if hasattr(pipe, 'config'): # real pipeline which can be auto-switched
@@ -1292,6 +1293,8 @@ def set_diffuser_pipe(pipe, new_pipe_type):
         new_pipe.feature_extractor = feature_extractor
     if mask_processor is not None:
         new_pipe.mask_processor = mask_processor
+    if restore_pipeline is not None:
+        new_pipe.restore_pipeline = restore_pipeline
     if new_pipe.__class__.__name__ in ['FluxPipeline', 'StableDiffusion3Pipeline']:
         new_pipe.register_modules(image_encoder = image_encoder)
         new_pipe.register_modules(feature_extractor = feature_extractor)