Add LMSDiscreteSchedulerTest #467

sidthekidder · 2022-09-11T04:25:10Z

PR for #338, still requires some fixes

HuggingFaceDocBuilderDev · 2022-09-11T04:28:42Z

The documentation is not available anymore as the PR was closed or merged.

anton-l

Hi @sidthekidder thanks a lot for diving into these tests! The LMS scheduler differs quite a bit in its API from other schedulers (and needs refactoring, to be honest), so this is quite a challenge!

I've left some suggestions to make the full_loop test work, let me know if something's not clear :)
In general, since this scheduler is unique to Stable Diffusion at the moment, feel free to reference the SD pipeline for how it's used in practice.

anton-l · 2022-09-12T13:26:52Z

tests/test_scheduler.py

+            self.check_over_configs(beta_start=beta_start, beta_end=beta_end)
+
+    def test_schedules(self):
+        for schedule in ["linear"]:


We should check over all shedules that are currently supported / in use

Suggested change

for schedule in ["linear"]:

for schedule in ["linear", "scaled_linear"]:

Updated to add scaled_linear (squaredcos_cap_v2 is not implemented for LMSDiscreteScheduler)

anton-l · 2022-09-12T13:29:34Z

tests/test_scheduler.py

+
+        num_trained_timesteps = len(scheduler)
+
+        model = self.dummy_model()


Suggested change

num_trained_timesteps = len(scheduler)

model = self.dummy_model()

model = self.dummy_model()

num_inference_steps = 10

scheduler.set_timesteps(num_inference_steps)

This scheduler requires a .set_timesteps() call. Reference: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L225

anton-l · 2022-09-12T13:31:30Z

tests/test_scheduler.py

+        num_trained_timesteps = len(scheduler)
+
+        model = self.dummy_model()
+        sample = self.dummy_sample_deter


Suggested change

sample = self.dummy_sample_deter

sample = self.dummy_sample_deter * scheduler.sigmas[0]

The initial sample needs to be multiplied by max sigma to make sure that it's within the supported range. Ref: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L229

Makes sense, updated

anton-l · 2022-09-12T13:33:05Z

tests/test_scheduler.py

+        model = self.dummy_model()
+        sample = self.dummy_sample_deter
+
+        for t in reversed(range(num_trained_timesteps - 1)):


Suggested change

for t in reversed(range(num_trained_timesteps - 1)):

for i, t in enumerate(scheduler.timesteps):

Can't get rid of enumeration yet, need both i and t for the loop in the current implementation 😅

anton-l · 2022-09-12T13:34:33Z

tests/test_scheduler.py

+            # 1. predict noise residual
+            residual = model(sample, t)
+            # print("residual: ")
+            # print(residual)


Suggested change

# 1. predict noise residual

residual = model(sample, t)

# print("residual: ")

# print(residual)

with torch.no_grad():

sigma = scheduler.sigmas[i]

sample = sample / ((sigma**2 + 1) ** 0.5)

model_output = model(sample, t)

Need to rescale the model input here too, to conform to the ODE equation used in K-LMS. Ref: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L246

anton-l · 2022-09-12T13:35:38Z

tests/test_scheduler.py

+            # 2. predict previous mean of sample x_t-1
+            pred_prev_sample = scheduler.step(residual, t, sample).prev_sample
+
+            # if t > 0:
+            #     noise = self.dummy_sample_deter
+            #     variance = scheduler.get_variance(t) ** (0.5) * noise
+            #
+            # sample = pred_prev_sample + variance
+            sample = pred_prev_sample


Suggested change

# 2. predict previous mean of sample x_t-1

pred_prev_sample = scheduler.step(residual, t, sample).prev_sample

# if t > 0:

# noise = self.dummy_sample_deter

# variance = scheduler.get_variance(t) ** (0.5) * noise

#

# sample = pred_prev_sample + variance

sample = pred_prev_sample

output = scheduler.step(model_output, i, sample)

sample = output.prev_sample

The LMS sampler takes i instead of an explicit timestep

anton-l · 2022-09-12T13:36:58Z

tests/test_scheduler.py

+        assert abs(result_sum.item() - 259.0883) < 1e-2
+        assert abs(result_mean.item() - 0.3374) < 1e-3


Suggested change

assert abs(result_sum.item() - 259.0883) < 1e-2

assert abs(result_mean.item() - 0.3374) < 1e-3

assert abs(result_sum.item() - 1006.3885) < 1e-2

assert abs(result_mean.item() - 1.3104) < 1e-3

These are the reference values if all of the above suggestions are applied, make sure they match on your hardware! (if not - we'll need to adjust)

Adjusted values as necessary with the new output

sidthekidder · 2022-09-13T08:49:02Z

Thanks for the detailed comments @anton-l! I didn't realize we could follow the implementation in pipeline_stable_diffusion.py, thanks for linking to the references - helped me understand the flow better.

Let me know if I am missing any test combinations or edge cases.

anton-l

Thank you @sidthekidder! I think all that's left is to implement test_pytorch_equal_numpy(), could you add that? :)
And feel free to modify the scheduler if it fails this test for some reason!

sidthekidder · 2022-09-15T09:31:37Z

Added test_pytorch_equal_numpy similar to the PNDMSchedulerTest.test_pytorch_equal_numpy fn.

patrickvonplaten · 2022-09-16T16:51:20Z

Nice looks good to me :-)

patrickvonplaten · 2022-09-16T16:51:59Z

@anton-l if you could review one more time this would be great :-)

anton-l

I'll just remove the unused past residuals (they're needed only in the PNDM tests) and we're good :)

tests/test_scheduler.py

anton-l · 2022-09-16T16:56:52Z

Thank you @sidthekidder!

* [WEB] CSS changes to the web-ui (huggingface#465) This commit updates UI with styling. Signed-Off-by: Gaurav Shukla <[email protected]> Signed-off-by: Gaurav Shukla <[email protected]> * [WEB] Update the title (huggingface#466) * [WEB] Add support for long prompts (huggingface#467) * [WEB] fix background color Signed-Off-by: Gaurav Shukla * [WEB] Remove long prompts support It removes support to long prompts due to higher lag in loading long prompts. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs> * [WEB] Update nod logo and enable debug feature. Signed-Off-by: Gaurav Shukla <[email protected]> Signed-off-by: Gaurav Shukla <[email protected]> Signed-off-by: Gaurav Shukla Signed-off-by: Gaurav Shukla <gaurav@nod-labs>

anton-l reviewed Sep 12, 2022

View reviewed changes

sidthekidder changed the title ~~[WIP] Add LMSDiscreteSchedulerTest~~ Add LMSDiscreteSchedulerTest Sep 13, 2022

anton-l reviewed Sep 13, 2022

View reviewed changes

sidthekidder added 4 commits September 15, 2022 02:10

[WIP] add LMSDiscreteSchedulerTest

ee15e3d

fixes for comments

c159bca

add torch numpy test

af9701a

rebase

629b934

sidthekidder force-pushed the lms-discrete-testing branch from d340b18 to 629b934 Compare September 15, 2022 09:12

patrickvonplaten approved these changes Sep 16, 2022

View reviewed changes

patrickvonplaten requested a review from anton-l September 16, 2022 16:51

anton-l approved these changes Sep 16, 2022

View reviewed changes

tests/test_scheduler.py Outdated Show resolved Hide resolved

tests/test_scheduler.py Outdated Show resolved Hide resolved

anton-l added 2 commits September 16, 2022 18:56

Update tests/test_scheduler.py

ab880e1

Update tests/test_scheduler.py

b404452

anton-l added 3 commits September 16, 2022 18:59

Merge branch 'main' into lms-discrete-testing

3e74c0b

style

7c31707

return residuals

bd713b7

anton-l merged commit a54cfe6 into huggingface:main Sep 16, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add LMSDiscreteSchedulerTest #467

Add LMSDiscreteSchedulerTest #467

sidthekidder commented Sep 11, 2022 •

edited

Loading

HuggingFaceDocBuilderDev commented Sep 11, 2022 •

edited

Loading

anton-l left a comment

anton-l Sep 12, 2022

sidthekidder Sep 13, 2022

anton-l Sep 12, 2022

sidthekidder Sep 13, 2022

anton-l Sep 12, 2022

sidthekidder Sep 13, 2022

anton-l Sep 12, 2022

anton-l Sep 12, 2022

anton-l Sep 12, 2022

anton-l Sep 12, 2022

sidthekidder Sep 13, 2022

sidthekidder commented Sep 13, 2022 •

edited

Loading

anton-l left a comment •

edited

Loading

sidthekidder commented Sep 15, 2022

patrickvonplaten commented Sep 16, 2022

patrickvonplaten commented Sep 16, 2022

anton-l left a comment •

edited

Loading

anton-l commented Sep 16, 2022

	for schedule in ["linear"]:
	for schedule in ["linear", "scaled_linear"]:


		num_trained_timesteps = len(scheduler)

		model = self.dummy_model()

	sample = self.dummy_sample_deter
	sample = self.dummy_sample_deter * scheduler.sigmas[0]

	for t in reversed(range(num_trained_timesteps - 1)):
	for i, t in enumerate(scheduler.timesteps):

		assert abs(result_sum.item() - 259.0883) < 1e-2
		assert abs(result_mean.item() - 0.3374) < 1e-3

Add LMSDiscreteSchedulerTest #467

Add LMSDiscreteSchedulerTest #467

Conversation

sidthekidder commented Sep 11, 2022 • edited Loading

HuggingFaceDocBuilderDev commented Sep 11, 2022 • edited Loading

anton-l left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sidthekidder commented Sep 13, 2022 • edited Loading

anton-l left a comment • edited Loading

Choose a reason for hiding this comment

sidthekidder commented Sep 15, 2022

patrickvonplaten commented Sep 16, 2022

patrickvonplaten commented Sep 16, 2022

anton-l left a comment • edited Loading

Choose a reason for hiding this comment

anton-l commented Sep 16, 2022

sidthekidder commented Sep 11, 2022 •

edited

Loading

HuggingFaceDocBuilderDev commented Sep 11, 2022 •

edited

Loading

sidthekidder commented Sep 13, 2022 •

edited

Loading

anton-l left a comment •

edited

Loading

anton-l left a comment •

edited

Loading