Fix callback_outputs is None error with callback on step end.

Skquark · Mar 8, 2024 · 943fc6a · 943fc6a
1 parent 28190fd
commit 943fc6a
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 9 deletions.
diff --git a/src/diffusers/pipelines/pia/pipeline_pia.py b/src/diffusers/pipelines/pia/pipeline_pia.py
@@ -1009,7 +1009,7 @@ def __call__(
                         for k in callback_on_step_end_tensor_inputs:
                             callback_kwargs[k] = locals()[k]
                         callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-                        if callback_outputs is None:
+                        if callback_outputs is not None:
                             latents = callback_outputs.pop("latents", latents)
                             prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
                             negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)

diff --git a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py
@@ -435,10 +435,10 @@ def __call__(
                 for k in callback_on_step_end_tensor_inputs:
                     callback_kwargs[k] = locals()[k]
                 callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                latents = callback_outputs.pop("latents", latents)
-                prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-                negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+                if callback_outputs is not None:
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
 
         if output_type not in ["pt", "np", "pil", "latent"]:
             raise ValueError(

diff --git a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py
@@ -593,10 +593,10 @@ def __call__(
                 for k in callback_on_step_end_tensor_inputs:
                     callback_kwargs[k] = locals()[k]
                 callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                latents = callback_outputs.pop("latents", latents)
-                prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-                negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+                if callback_outputs is not None:
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
 
         # Offload all models
         self.maybe_free_model_hooks()