fix(ddpm): use _execution_device, validate inputs, free hooks (#13649)#13671
fix(ddpm): use _execution_device, validate inputs, free hooks (#13649)#13671Anai-Guo wants to merge 1 commit intohuggingface:mainfrom
Conversation
…gface#13649) Issue 1: replace self.device with self._execution_device so model_cpu_offload's execution device is honored, and call self.maybe_free_model_hooks() before return to satisfy the offload contract. Issue 2: validate that len(generator) == batch_size for list generators, raising ValueError instead of silently mishandling per-sample seeding (matches DDIM/ ConsistencyModel pipelines). Issue 3: validate output_type and add 'pt' tensor output. Previously any value other than 'pil' silently fell through to NumPy. Closes huggingface#13649.
hlky
left a comment
hlky
left a comment
There was a problem hiding this comment.
Thanks @Anai-Guo.
I left a few comments. Some are just confirming that the PR resolves the reported issues, and a couple are broader questions for maintainers about whether these patterns should become review/agent rules:
- validating
output_type - preferring
VaeImageProcessor.postprocess(...)over manual clamp/permute/NumPy/PIL handling
Overall, this looks like it addresses the ddpm findings from #13649.
Also, feel free to tag me on PRs that come from the review issues / meta issue since I can usually confirm whether they address the reported finding. For approval/merge, please also tag an appropriate maintainer since I can review the issue context but cannot approve. In general, I’m happy to help clarify the reports or point out related patterns from #13656 if useful.
| if output_type not in ["pt", "np", "pil"]: | ||
| raise ValueError(f"output_type must be one of ['pt', 'np', 'pil'], got '{output_type}'.") |
There was a problem hiding this comment.
cc @yiyixuxu Should output_type validation be a review rule?
| if isinstance(generator, list) and len(generator) != batch_size: | ||
| raise ValueError( | ||
| f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" | ||
| f" size of {batch_size}. Make sure the batch size matches the length of the generators." | ||
| ) |
| @@ -108,12 +120,12 @@ def __call__( | |||
| else: | |||
| image_shape = (batch_size, self.unet.config.in_channels, *self.unet.config.sample_size) | |||
|
|
|||
| if self.device.type == "mps": | |||
| if device.type == "mps": | |||
| # randn does not work reproducibly on mps | |||
| image = randn_tensor(image_shape, generator=generator, dtype=self.unet.dtype) | |||
| image = image.to(self.device) | |||
| image = image.to(device) | |||
| else: | |||
| image = randn_tensor(image_shape, generator=generator, device=self.device, dtype=self.unet.dtype) | |||
| image = randn_tensor(image_shape, generator=generator, device=device, dtype=self.unet.dtype) | |||
| image = (image / 2 + 0.5).clamp(0, 1) | ||
| image = image.cpu().permute(0, 2, 3, 1).numpy() | ||
| if output_type == "pil": | ||
| image = self.numpy_to_pil(image) | ||
| if output_type != "pt": | ||
| image = image.cpu().permute(0, 2, 3, 1).numpy() | ||
| if output_type == "pil": | ||
| image = self.numpy_to_pil(image) |
There was a problem hiding this comment.
As per additional review in #13663 (comment) this could use VaeImageProcessor, this would fix all output_type's without the if statements. cc @yiyixuxu For awareness, this is another case that supports introducing VaeImageProcessor usage as a review rule.
|
Thanks @hlky for the review! Glad the fix resolves both reported issues. For the line 147 suggestion — happy to refactor to |
2 participants
What does this PR do?
Fixes the three issues called out in the
ddpmmodel/pipeline review (cc @hlky). All three changes live inpipelines/ddpm/pipeline_ddpm.pyand follow the suggested fixes from the issue, with precedents fromDDIMPipelineandConsistencyModelPipeline.Issue 1 —
DDPMPipelinedoes not run latents on the offload execution deviceDDPMPipelinedeclaresmodel_cpu_offload_seq = "unet"but initializes the latent onself.device(which stays CPU underenable_model_cpu_offload()) and never callsself.maybe_free_model_hooks()before returning. Switched toself._execution_device, threaded that into therandn_tensorcall (and thempsbranch), and added themaybe_free_model_hooks()call before the return so the offload contract is honored and the UNet doesn't stay resident on the accelerator.Issue 2 — Generator lists are not validated against
batch_sizeA short generator list either gets silently treated like a single generator or raises a raw
IndexError. Added the sameValueErrorguardDDIMPipelineandConsistencyModelPipelineuse, so users get a clear message whenlen(generator) != batch_size.Issue 3 — Invalid
output_typevalues silently return NumPyAny
output_typeother than"pil"previously fell through to NumPy, including typos and"pt". Added an explicitoutput_type in {"pt", "np", "pil"}check at the top of__call__, made"pt"actually return the tensor (skipping the.cpu().numpy()round-trip), and updated the docstring to mentiontorch.Tensor. Behavior for the documented values ("pil", default;"np") is unchanged.Reproductions
The repros from #13649 are short and self-contained — see the issue body for runnable scripts that demonstrate each failure mode.
Before submitting
Who can review?
@hlky — this addresses all three issues from the
ddpmreview.🤖 Generated with Claude Code