Commit 24575991 authored by Patrick von Platen's avatar Patrick von Platen
Browse files

make fix copies

parent bdd16116
...@@ -282,13 +282,13 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin): ...@@ -282,13 +282,13 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
https://arxiv.org/abs/2205.11487 https://arxiv.org/abs/2205.11487
""" """
dtype = sample.dtype dtype = sample.dtype
batch_size, channels, *remaining_dims = sample.shape batch_size, channels, height, width = sample.shape
if dtype not in (torch.float32, torch.float64): if dtype not in (torch.float32, torch.float64):
sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half
# Flatten sample for doing quantile calculation along each image # Flatten sample for doing quantile calculation along each image
sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) sample = sample.reshape(batch_size, channels * height * width)
abs_sample = sample.abs() # "a certain percentile absolute pixel value" abs_sample = sample.abs() # "a certain percentile absolute pixel value"
...@@ -300,7 +300,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin): ...@@ -300,7 +300,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0
sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s"
sample = sample.reshape(batch_size, channels, *remaining_dims) sample = sample.reshape(batch_size, channels, height, width)
sample = sample.to(dtype) sample = sample.to(dtype)
return sample return sample
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment