Update sd-webui-progressive-growing/scripts/progressive_growing_always.py

sd-webui-progressive-growing/scripts/progressive_growing_always.py

@@ -655,7 +655,17 @@ def _stage_sample_txt2img(p, x, conditioning, unconditional_conditioning,
             _call_script_hooks(p, x)
             _rewrap_scheduler_override_for_stage(p, stage_w, stage_h)
             try:
-                getattr(p.sampler, 'sampler_extra_args', {}).pop('sigmas', None)
+                extra = getattr(p.sampler, 'sampler_extra_args', {})
+                if isinstance(extra, dict):
+                    extra.pop('sigmas', None)
+                    # noise_sampler may have been built from full-size p.x by
+                    # a previous initialize() call (e.g. adept-sampler-v5 stores
+                    # the original and passes it through).  Removing it here
+                    # forces A1111's initialize() to recreate it from the
+                    # stage-sized p.x that _stage_sampler_context set up,
+                    # preventing "size of tensor a (26) must match tensor b (104)"
+                    # in ancestral / TCD steps.
+                    extra.pop('noise_sampler', None)
             except Exception:
                 pass
             _patch_sampler_func_sigmas_to_1d(p)
@@ -676,6 +686,18 @@ def _stage_sample_txt2img(p, x, conditioning, unconditional_conditioning,
         with _stage_sampler_context(p, stage_w, stage_h):
             p.sampler_noise_scheduler_override = None
             _call_script_hooks(p, x)
+            # For standard samplers the sampler object is reused across stages.
+            # A1111's initialize() may skip recreating noise_sampler if it is
+            # already present in sampler_extra_args (built from full-size p.x
+            # on the first call).  Drop it here so initialize() always builds a
+            # fresh one from the stage-sized p.x set by _stage_sampler_context.
+            try:
+                extra = getattr(p.sampler, 'sampler_extra_args', {})
+                if isinstance(extra, dict):
+                    extra.pop('sigmas', None)
+                    extra.pop('noise_sampler', None)
+            except Exception:
+                pass
             return p.sampler.sample(
                 p, x, conditioning, unconditional_conditioning,
                 image_conditioning=image_cond,
@@ -697,7 +719,10 @@ def _stage_sample_img2img(p, samples, noise, conditioning, unconditional_conditi
             _call_script_hooks(p, samples)
             _rewrap_scheduler_override_for_stage(p, stage_w, stage_h)
             try:
-                getattr(p.sampler, 'sampler_extra_args', {}).pop('sigmas', None)
+                extra = getattr(p.sampler, 'sampler_extra_args', {})
+                if isinstance(extra, dict):
+                    extra.pop('sigmas', None)
+                    extra.pop('noise_sampler', None)
             except Exception:
                 pass
             _patch_sampler_func_sigmas_to_1d(p)
@@ -713,6 +738,13 @@ def _stage_sample_img2img(p, samples, noise, conditioning, unconditional_conditi
         with _stage_sampler_context(p, stage_w, stage_h):
             p.sampler_noise_scheduler_override = None
             _call_script_hooks(p, samples)
+            try:
+                extra = getattr(p.sampler, 'sampler_extra_args', {})
+                if isinstance(extra, dict):
+                    extra.pop('sigmas', None)
+                    extra.pop('noise_sampler', None)
+            except Exception:
+                pass
             return p.sampler.sample_img2img(
                 p, samples, noise,
                 conditioning, unconditional_conditioning,