huggingface / huggingface/diffusers
Type hint for `callback_on_step_end` in pipeline `__call__` is incorrect
- Vorherrschende Sprache
- Python
- Sterne
- 34.5k
- Forks
- 7.3k
- Ø Merge
- 3 T. 3 Std.
- Gemergte PRs (30 T.)
- 91
Beschreibung
### Describe the bug
As far as I can tell, the type hint `callback_on_step_end` for all pipelines (as seen in [this search](https://github.com/search?q=repo%3Ahuggingface%2Fdiffusers+%22callback_on_step_end%3A+Optional%22&type=code)) is incorrect.
Taking the SDXL pipeline as an example, the type hint in `__call__` is:
```python
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None
```
but it is called like this:
```python
if callback_on_step_end is not None:
callback_kwargs = {}
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)
add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds)
negative_pooled_prompt_embeds = callback_outputs.pop(
"negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
)
add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids)
```
That is, the type hint suggests that it's a function with three arguments that returns nothing, but it's actually a function with four arguments that returns a `Dict`. This ends up failing at runtime. Using a four-argument function leads to type-checker errors.
### Reproduction
Using a Python typechecker (Pyright in my case), attempt to use a correctly-defined callback with a SDXL pipeline:
```python
sdxl_pipe = AutoPipelineForText2Image.from_pretrained(
"stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16"
)
sdxl_pipe = typing.cast(StableDiffusionXLPipeline, sdxl_pipe)
def callback_on_step_end(_pipe, step, _timestep, kwargs):
print(step)
return kwargs
output = sdxl_pipe(
prompt="a prompt",
negative_prompt=self.negative_prompt,
num_inference_steps=4,
guidance_scale=0.0,
return_dict=True,
width=512,
height=512,
callback_on_step_end=callback_on_step_end,
)
```
The type-checker will error on `callback_on_step_end`:
```
Argument of type "(_pipe: Unknown, step: Unknown, _timestep: Unknown, kwargs: Unknown) -> Unknown" cannot be assigned to parameter "callback_on_step_end" of type "((int, int, Dict[Unknown, Unknown]) -> None) | None" in function "__call__"
Type "(_pipe: Unknown, step: Unknown, _timestep: Unknown, kwargs: Unknown) -> Unknown" cannot be assigned to type "((int, int, Dict[Unknown, Unknown]) -> None) | None"
Type "(_pipe: Unknown, step: Unknown, _timestep: Unknown, kwargs: Unknown) -> Unknown" cannot be assigned to type "(int, int, Dict[Unknown, Unknown]) -> None"
Function accepts too few positional parameters; expected 4 but received 3
"function" is incompatible with "None"
```
This can be worked around using `# type: ignore`, which is what I'm doing.
### Logs
_No response_
### System Info
N/A
### Who can help?
_No response_
Beitragsleitfaden
Bewertung
Dieses Issue wurde noch nicht bewertet.