Skip to content

Commit

Permalink
Workaround for upscale with large output tensors.
Browse files Browse the repository at this point in the history
Fixes #10040.
  • Loading branch information
pcuenca committed Dec 1, 2024
1 parent c96bfa5 commit 35085ea
Showing 1 changed file with 6 additions and 0 deletions.
6 changes: 6 additions & 0 deletions src/diffusers/models/upsampling.py
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,12 @@ def forward(self, hidden_states: torch.Tensor, output_size: Optional[int] = None
# if `output_size` is passed we force the interpolation output
# size and do not make use of `scale_factor=2`
if self.interpolate:
# upsample_nearest_nhwc also fails when the number of output elements is large
# https://github.com/pytorch/pytorch/issues/141831
scale_factor = 2 if output_size is None else max(output_size)
if hidden_states.numel() * scale_factor > pow(2, 31):
hidden_states = hidden_states.contiguous()

if output_size is None:
hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
else:
Expand Down

0 comments on commit 35085ea

Please sign in to comment.