sampler flow shift options and fix img2img

Signed-off-by: Vladimir Mandic <mandic00@live.com>

CHANGELOG.md

@@ -8,10 +8,12 @@
 - Add legacy option to use old LoRA loader in *settings -> networks*  
 - Add sigma calculation to VAE preview, thanks @Disty0  
 - Fix live preview image sizes in modern and standard UI  
+- Fix image width/height calculation when doing img2img  
 - HunyuanVideo optimizations: full offload, quantization and tiling support  
 - LTXVideo optimizations: full offload, quantization and tiling support  
 - Do not show disabled networks  
 - CSS optimizations when log view is disabled  
+- Sampler options: add flow shift and separate dynamic thresholding from dynamic shifting  
 
 ## Update for 2024-12-24
 

modules/images_grid.py

@@ -42,10 +42,11 @@ def image_grid(imgs, batch_size=1, rows=None):
     imgs = [i for i in imgs if i is not None] if imgs is not None else []
     if len(imgs) == 0:
         return None
-    w, h = max(i.width for i in imgs), max(i.height for i in imgs)
+    w, h = max(i.width for i in imgs if i is not None), max(i.height for i in imgs if i is not None)
     grid = Image.new('RGB', size=(params.cols * w, params.rows * h), color=shared.opts.grid_background)
     for i, img in enumerate(params.imgs):
-        grid.paste(img, box=(i % params.cols * w, i // params.cols * h))
+        if img is not None:
+            grid.paste(img, box=(i % params.cols * w, i // params.cols * h))
     return grid
 
 

modules/processing_args.py

@@ -22,7 +22,7 @@ def task_specific_kwargs(p, model):
     if len(getattr(p, 'init_images', [])) > 0:
         if isinstance(p.init_images[0], str):
             p.init_images = [helpers.decode_base64_to_image(i, quiet=True) for i in p.init_images]
-        p.init_images = [i.convert('RGB') if i.mode != 'RGB' else i for i in p.init_images]
+        p.init_images = [i.convert('RGB') if i.mode != 'RGB' else i for i in p.init_images if i is not None]
     if (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE or len(getattr(p, 'init_images', [])) == 0) and not is_img2img_model:
         p.ops.append('txt2img')
         if hasattr(p, 'width') and hasattr(p, 'height'):
@@ -262,6 +262,9 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2
     elif 'callback' in possible:
         args['callback'] = diffusers_callback_legacy
 
+    if 'image' in kwargs and len(getattr(p, 'init_images', [])) == 0:
+        p.init_images = kwargs['image'] if isinstance(kwargs['image'], list) else [kwargs['image']]
+
     # handle remaining args
     for arg in kwargs:
         if arg in possible: # add kwargs

modules/processing_diffusers.py

@@ -201,7 +201,8 @@ def process_hires(p: processing.StableDiffusionProcessing, output):
             if p.is_control and hasattr(p, 'task_args') and p.task_args.get('image', None) is not None:
                 if hasattr(shared.sd_model, "vae") and output.images is not None and len(output.images) > 0:
                     output.images = processing_vae.vae_decode(latents=output.images, model=shared.sd_model, full_quality=p.full_quality, output_type='pil', width=p.hr_upscale_to_x, height=p.hr_upscale_to_y) # controlnet cannnot deal with latent input
-                    p.task_args['image'] = output.images # replace so hires uses new output
+                    p.init_images = output.images # replace so hires uses new output
+                    # p.task_args['image'] = output.images # replace so hires uses new output
             update_sampler(p, shared.sd_model, second_pass=True)
             orig_denoise = p.denoising_strength
             p.denoising_strength = strength
@@ -290,8 +291,9 @@ def process_refine(p: processing.StableDiffusionProcessing, output):
                 image = processing_vae.vae_decode(latents=image, model=shared.sd_model, full_quality=p.full_quality, output_type='pil', width=p.width, height=p.height)
                 p.extra_generation_params['Noise level'] = noise_level
                 output_type = 'np'
-            if hasattr(p, 'task_args') and p.task_args.get('image', None) is not None and output is not None: # replace input with output so it can be used by hires/refine
-                p.task_args['image'] = image
+            if p.task_args.get('image', None) is not None and output is not None: # replace input with output so it can be used by hires/refine
+                # p.task_args['image'] = image
+                p.init_images = [image]
             shared.log.info(f'Refiner: class={shared.sd_refiner.__class__.__name__}')
             update_sampler(p, shared.sd_refiner, second_pass=True)
             refiner_args = set_pipeline_args(

modules/sd_samplers_diffusers.py

@@ -235,16 +235,9 @@ class DiffusionSampler:
         if 'beta_end' in self.config and shared.opts.schedulers_beta_end > 0:
             self.config['beta_end'] = shared.opts.schedulers_beta_end
         if 'shift' in self.config:
-            if shared.opts.schedulers_shift == 0:
-                if 'StableDiffusion3' in model.__class__.__name__:
-                    self.config['shift'] = 3
-                if 'Flux' in model.__class__.__name__:
-                    self.config['shift'] = 1
-            else:
-                self.config['shift'] = shared.opts.schedulers_shift
+            self.config['shift'] = shared.opts.schedulers_shift if shared.opts.schedulers_shift > 0 else 3
         if 'use_dynamic_shifting' in self.config:
-            if 'Flux' in model.__class__.__name__:
-                self.config['use_dynamic_shifting'] = shared.opts.schedulers_dynamic_shift
+            self.config['use_dynamic_shifting'] = True if shared.opts.schedulers_shift <= 0 else shared.opts.schedulers_dynamic_shift
         if 'use_beta_sigmas' in self.config and 'sigma_schedule' in self.config:
             self.config['use_beta_sigmas'] = 'StableDiffusion3' in model.__class__.__name__
         if 'rescale_betas_zero_snr' in self.config:

modules/shared.py

@@ -822,8 +822,8 @@ options_templates.update(options_section(('sampler-params', "Sampler Settings"),
     'schedulers_beta_start': OptionInfo(0, "Beta start", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.00001, "visible": native}),
     'schedulers_beta_end': OptionInfo(0, "Beta end", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.00001, "visible": native}),
     'schedulers_timesteps_range': OptionInfo(1000, "Timesteps range", gr.Slider, {"minimum": 250, "maximum": 4000, "step": 1, "visible": native}),
-    'schedulers_shift': OptionInfo(0, "Sampler shift", gr.Slider, {"minimum": 0.1, "maximum": 10, "step": 0.1, "visible": native}),
-    'schedulers_dynamic_shift': OptionInfo(True, "Sampler dynamic shift", gr.Checkbox, {"visible": native}),
+    'schedulers_shift': OptionInfo(3, "Sampler shift", gr.Slider, {"minimum": 0.1, "maximum": 10, "step": 0.1, "visible": False}),
+    'schedulers_dynamic_shift': OptionInfo(True, "Sampler dynamic shift", gr.Checkbox, {"visible": False}),
 
     # managed from ui.py for backend original k-diffusion
     "always_batch_cond_uncond": OptionInfo(False, "Disable conditional batching", gr.Checkbox, {"visible": not native}),

modules/ui_sections.py

@@ -217,7 +217,8 @@ def create_sampler_options(tabname):
         shared.opts.save(shared.config_filename, silent=True)
 
     def set_sampler_options(sampler_options):
-        shared.opts.data['schedulers_use_thresholding'] = 'dynamic' in sampler_options
+        shared.opts.data['schedulers_dynamic_shift'] = 'dynamic' in sampler_options
+        shared.opts.data['schedulers_use_thresholding'] = 'thresholding' in sampler_options
         shared.opts.data['schedulers_use_loworder'] = 'low order' in sampler_options
         shared.opts.data['schedulers_rescale_betas'] = 'rescale' in sampler_options
         shared.log.debug(f'Sampler set options: {sampler_options}')
@@ -253,6 +254,11 @@ def create_sampler_options(tabname):
         shared.opts.schedulers_beta_schedule = sampler_beta
         shared.opts.save(shared.config_filename, silent=True)
 
+    def set_sampler_shift(sampler_shift):
+        shared.log.debug(f'Sampler set options: shift={sampler_shift}')
+        shared.opts.schedulers_shift = sampler_shift
+        shared.opts.save(shared.config_filename, silent=True)
+
     # 'linear', 'scaled_linear', 'squaredcos_cap_v2'
     def set_sampler_preset(preset):
         if preset == 'AYS SD15':
@@ -286,10 +292,13 @@ def create_sampler_options(tabname):
             sampler_timesteps = gr.Textbox(label='Timesteps override', elem_id=f"{tabname}_sampler_timesteps", value=shared.opts.schedulers_timesteps)
         with gr.Row(elem_classes=['flex-break']):
             sampler_order = gr.Slider(minimum=0, maximum=5, step=1, label="Sampler order", value=shared.opts.schedulers_solver_order, elem_id=f"{tabname}_sampler_order")
-            options = ['low order', 'dynamic', 'rescale']
+            sampler_shift = gr.Slider(minimum=0, maximum=10, step=0.1, label="Flow shift", value=shared.opts.schedulers_shift, elem_id=f"{tabname}_sampler_shift")
+        with gr.Row(elem_classes=['flex-break']):
+            options = ['low order', 'thresholding', 'dynamic', 'rescale']
             values = []
             values += ['low order'] if shared.opts.data.get('schedulers_use_loworder', True) else []
-            values += ['dynamic'] if shared.opts.data.get('schedulers_use_thresholding', False) else []
+            values += ['thresholding'] if shared.opts.data.get('schedulers_use_thresholding', False) else []
+            values += ['dynamic'] if shared.opts.data.get('schedulers_dynamic_shift', False) else []
             values += ['rescale'] if shared.opts.data.get('schedulers_rescale_betas', False) else []
             sampler_options = gr.CheckboxGroup(label='Options', elem_id=f"{tabname}_sampler_options", choices=options, value=values, type='value')
 
@@ -300,6 +309,7 @@ def create_sampler_options(tabname):
         sampler_beta.change(fn=set_sampler_beta, inputs=[sampler_beta], outputs=[])
         sampler_prediction.change(fn=set_sampler_prediction, inputs=[sampler_prediction], outputs=[])
         sampler_order.change(fn=set_sampler_order, inputs=[sampler_order], outputs=[])
+        sampler_shift.change(fn=set_sampler_shift, inputs=[sampler_shift], outputs=[])
         sampler_options.change(fn=set_sampler_options, inputs=[sampler_options], outputs=[])
 
 

scripts/xyz_grid_draw.py

@@ -109,7 +109,7 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend
         idx0 = (i * len(xs) * len(ys)) + i # starting index of images in subgrid
         idx1 = (len(xs) * len(ys)) + idx0  # ending index of images in subgrid
         to_process = processed_result.images[idx0:idx1]
-        w, h = max(i.width for i in to_process), max(i.height for i in to_process if i is not None)
+        w, h = max(i.width for i in to_process if i is not None), max(i.height for i in to_process if i is not None)
         if w is None or h is None or w == 0 or h == 0:
             shared.log.error("XYZ grid: failed get valid image")
             continue