Check if callback_outputs is not None
Browse files- pipeline.py +3 -2
pipeline.py
CHANGED
@@ -954,8 +954,9 @@ class FluxDifferentialImg2ImgPipeline(DiffusionPipeline, FluxLoraLoaderMixin):
|
|
954 |
callback_kwargs[k] = locals()[k]
|
955 |
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
|
956 |
|
957 |
-
|
958 |
-
|
|
|
959 |
|
960 |
# call the callback, if provided
|
961 |
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
|
|
|
954 |
callback_kwargs[k] = locals()[k]
|
955 |
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
|
956 |
|
957 |
+
if callback_outputs is not None:
|
958 |
+
latents = callback_outputs.pop("latents", latents)
|
959 |
+
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
|
960 |
|
961 |
# call the callback, if provided
|
962 |
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
|