Skip to content

Commit 9602fc5

Browse files
authored
Fix Ideogram 4 Callback Handling and Tests (#14621)
Fix Ideogram 4 callback_kwargs and corresponding callback tests
1 parent 119c339 commit 9602fc5

2 files changed

Lines changed: 3 additions & 24 deletions

File tree

‎src/diffusers/pipelines/ideogram4/pipeline_ideogram4.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -714,7 +714,9 @@ def __call__(
714714
latents = self.scheduler.step(-v, t, latents, return_dict=False)[0]
715715

716716
if callback_on_step_end is not None:
717-
callback_kwargs = {k: locals()[k] for k in callback_on_step_end_tensor_inputs}
717+
callback_kwargs = {}
718+
for k in callback_on_step_end_tensor_inputs:
719+
callback_kwargs[k] = locals()[k]
718720
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
719721
latents = callback_outputs.pop("latents", latents)
720722

‎tests/pipelines/ideogram4/test_pipeline_ideogram4.py‎

Lines changed: 0 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -170,29 +170,6 @@ def test_save_load_float16(self, tmp_path, expected_max_diff=5e-2):
170170
def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=None, atol=1e-4, rtol=1e-4):
171171
pass
172172

173-
# `callback_on_step_end` is unusable on the Python versions this repo supports (`python_requires>=3.10.0`):
174-
# the denoising loop builds `callback_kwargs` with `{k: locals()[k] for k in callback_on_step_end_tensor_inputs}`,
175-
# and before Python 3.12 a comprehension runs in its own frame, so `locals()` never contains `latents` and the
176-
# call raises `KeyError: 'latents'`. PEP 709 inlined comprehensions in 3.12, which is the only reason this passes
177-
# locally on 3.12 while CI (3.10) fails. Every other pipeline in the repo builds the dict with a plain `for` loop,
178-
# which works on all versions.
179-
_CALLBACK_SKIP = (
180-
"`Ideogram4Pipeline` builds `callback_kwargs` in a dict comprehension, so `locals()` cannot see `latents` "
181-
"before Python 3.12 and any `callback_on_step_end` raises `KeyError: 'latents'` on Python 3.10/3.11."
182-
)
183-
184-
@pytest.mark.skip(reason=_CALLBACK_SKIP)
185-
def test_callback_inputs(self):
186-
pass
187-
188-
@pytest.mark.skip(
189-
reason=(
190-
f"{_CALLBACK_SKIP} The body below is the Ideogram4-specific replacement for the shared assertion (which "
191-
"would not apply here either, since the pipeline republishes the step's schedule weight on "
192-
"`_guidance_scale` so a callback's mutation cannot accumulate); drop this marker once the pipeline "
193-
"builds `callback_kwargs` with a plain `for` loop."
194-
)
195-
)
196173
def test_callback_cfg(self):
197174
# Ideogram4 drives guidance from a per-step schedule and republishes the current step's weight on
198175
# `_guidance_scale` before invoking the callback, so a callback's mutation cannot accumulate across steps

0 commit comments

Comments
 (0)