-
Notifications
You must be signed in to change notification settings - Fork 4.8k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
v-prediction training support #1455
Changes from all commits
9aac7d7
65afd4e
a1996d5
24f6bc9
f2d6c71
576e0c1
415a96a
cd20d7d
6841a18
431d1e5
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -355,5 +355,25 @@ def add_noise( | |
noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise | ||
return noisy_samples | ||
|
||
def get_velocity( | ||
self, sample: torch.FloatTensor, noise: torch.FloatTensor, timesteps: torch.IntTensor | ||
) -> torch.FloatTensor: | ||
# Make sure alphas_cumprod and timestep have same device and dtype as sample | ||
self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device, dtype=sample.dtype) | ||
timesteps = timesteps.to(sample.device) | ||
|
||
sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5 | ||
sqrt_alpha_prod = sqrt_alpha_prod.flatten() | ||
while len(sqrt_alpha_prod.shape) < len(sample.shape): | ||
sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) | ||
|
||
sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5 | ||
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() | ||
while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape): | ||
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) | ||
|
||
velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample | ||
return velocity | ||
Comment on lines
+358
to
+376
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This looks very similar to There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yes, could refactor it in a follow-up PR , wdyt @patrickvonplaten There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Since we use both There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Agree ! I think it's clear to add it directly to There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. scratch that it doesn't work as expected There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. In favour of keeping
|
||
|
||
def __len__(self): | ||
return self.config.num_train_timesteps |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -345,5 +345,25 @@ def add_noise( | |
noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise | ||
return noisy_samples | ||
|
||
def get_velocity( | ||
self, sample: torch.FloatTensor, noise: torch.FloatTensor, timesteps: torch.IntTensor | ||
) -> torch.FloatTensor: | ||
# Make sure alphas_cumprod and timestep have same device and dtype as sample | ||
self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device, dtype=sample.dtype) | ||
timesteps = timesteps.to(sample.device) | ||
|
||
sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5 | ||
sqrt_alpha_prod = sqrt_alpha_prod.flatten() | ||
while len(sqrt_alpha_prod.shape) < len(sample.shape): | ||
sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) | ||
|
||
sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5 | ||
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() | ||
while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape): | ||
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) | ||
|
||
velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample | ||
return velocity | ||
|
||
Comment on lines
+348
to
+367
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Same comment as in DDIM. |
||
def __len__(self): | ||
return self.config.num_train_timesteps |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
👍