Skip to content

Commit

Permalink
[fix] FreeInit step index out of bounds (#8969)
Browse files Browse the repository at this point in the history
* fix step index out of bounds

* add test for free_init with different schedulers

* add test to vid2vid and pia
  • Loading branch information
a-r-r-o-w authored Jul 26, 2024
1 parent 1168eaa commit 57a021d
Show file tree
Hide file tree
Showing 4 changed files with 146 additions and 0 deletions.
2 changes: 2 additions & 0 deletions src/diffusers/pipelines/free_init_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,8 @@ def _apply_free_init(
num_inference_steps = max(
1, int(num_inference_steps / self._free_init_num_iters * (free_init_iteration + 1))
)

if num_inference_steps > 0:
self.scheduler.set_timesteps(num_inference_steps, device=device)

return latents, self.scheduler.timesteps
48 changes: 48 additions & 0 deletions tests/pipelines/animatediff/test_animatediff.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
AnimateDiffPipeline,
AutoencoderKL,
DDIMScheduler,
DPMSolverMultistepScheduler,
LCMScheduler,
MotionAdapter,
StableDiffusionPipeline,
UNet2DConditionModel,
Expand Down Expand Up @@ -353,6 +355,52 @@ def test_free_init(self):
"Disabling of FreeInit should lead to results similar to the default pipeline results",
)

def test_free_init_with_schedulers(self):
components = self.get_dummy_components()
pipe: AnimateDiffPipeline = self.pipeline_class(**components)
pipe.set_progress_bar_config(disable=None)
pipe.to(torch_device)

inputs_normal = self.get_dummy_inputs(torch_device)
frames_normal = pipe(**inputs_normal).frames[0]

schedulers_to_test = [
DPMSolverMultistepScheduler.from_config(
components["scheduler"].config,
timestep_spacing="linspace",
beta_schedule="linear",
algorithm_type="dpmsolver++",
steps_offset=1,
clip_sample=False,
),
LCMScheduler.from_config(
components["scheduler"].config,
timestep_spacing="linspace",
beta_schedule="linear",
steps_offset=1,
clip_sample=False,
),
]
components.pop("scheduler")

for scheduler in schedulers_to_test:
components["scheduler"] = scheduler
pipe: AnimateDiffPipeline = self.pipeline_class(**components)
pipe.set_progress_bar_config(disable=None)
pipe.to(torch_device)

pipe.enable_free_init(num_iters=2, use_fast_sampling=False)

inputs = self.get_dummy_inputs(torch_device)
frames_enable_free_init = pipe(**inputs).frames[0]
sum_enabled = np.abs(to_np(frames_normal) - to_np(frames_enable_free_init)).sum()

self.assertGreater(
sum_enabled,
1e1,
"Enabling of FreeInit should lead to results different from the default pipeline results",
)

@unittest.skipIf(
torch_device != "cuda" or not is_xformers_available(),
reason="XFormers attention is only available with CUDA and `xformers` installed",
Expand Down
48 changes: 48 additions & 0 deletions tests/pipelines/animatediff/test_animatediff_video2video.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
AnimateDiffVideoToVideoPipeline,
AutoencoderKL,
DDIMScheduler,
DPMSolverMultistepScheduler,
LCMScheduler,
MotionAdapter,
StableDiffusionPipeline,
UNet2DConditionModel,
Expand Down Expand Up @@ -380,3 +382,49 @@ def test_free_init(self):
1e-4,
"Disabling of FreeInit should lead to results similar to the default pipeline results",
)

def test_free_init_with_schedulers(self):
components = self.get_dummy_components()
pipe: AnimateDiffVideoToVideoPipeline = self.pipeline_class(**components)
pipe.set_progress_bar_config(disable=None)
pipe.to(torch_device)

inputs_normal = self.get_dummy_inputs(torch_device)
frames_normal = pipe(**inputs_normal).frames[0]

schedulers_to_test = [
DPMSolverMultistepScheduler.from_config(
components["scheduler"].config,
timestep_spacing="linspace",
beta_schedule="linear",
algorithm_type="dpmsolver++",
steps_offset=1,
clip_sample=False,
),
LCMScheduler.from_config(
components["scheduler"].config,
timestep_spacing="linspace",
beta_schedule="linear",
steps_offset=1,
clip_sample=False,
),
]
components.pop("scheduler")

for scheduler in schedulers_to_test:
components["scheduler"] = scheduler
pipe: AnimateDiffVideoToVideoPipeline = self.pipeline_class(**components)
pipe.set_progress_bar_config(disable=None)
pipe.to(torch_device)

pipe.enable_free_init(num_iters=2, use_fast_sampling=False)

inputs = self.get_dummy_inputs(torch_device)
frames_enable_free_init = pipe(**inputs).frames[0]
sum_enabled = np.abs(to_np(frames_normal) - to_np(frames_enable_free_init)).sum()

self.assertGreater(
sum_enabled,
1e1,
"Enabling of FreeInit should lead to results different from the default pipeline results",
)
48 changes: 48 additions & 0 deletions tests/pipelines/pia/test_pia.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
from diffusers import (
AutoencoderKL,
DDIMScheduler,
DPMSolverMultistepScheduler,
LCMScheduler,
MotionAdapter,
PIAPipeline,
StableDiffusionPipeline,
Expand Down Expand Up @@ -360,6 +362,52 @@ def test_free_init(self):
"Disabling of FreeInit should lead to results similar to the default pipeline results",
)

def test_free_init_with_schedulers(self):
components = self.get_dummy_components()
pipe: PIAPipeline = self.pipeline_class(**components)
pipe.set_progress_bar_config(disable=None)
pipe.to(torch_device)

inputs_normal = self.get_dummy_inputs(torch_device)
frames_normal = pipe(**inputs_normal).frames[0]

schedulers_to_test = [
DPMSolverMultistepScheduler.from_config(
components["scheduler"].config,
timestep_spacing="linspace",
beta_schedule="linear",
algorithm_type="dpmsolver++",
steps_offset=1,
clip_sample=False,
),
LCMScheduler.from_config(
components["scheduler"].config,
timestep_spacing="linspace",
beta_schedule="linear",
steps_offset=1,
clip_sample=False,
),
]
components.pop("scheduler")

for scheduler in schedulers_to_test:
components["scheduler"] = scheduler
pipe: PIAPipeline = self.pipeline_class(**components)
pipe.set_progress_bar_config(disable=None)
pipe.to(torch_device)

pipe.enable_free_init(num_iters=2, use_fast_sampling=False)

inputs = self.get_dummy_inputs(torch_device)
frames_enable_free_init = pipe(**inputs).frames[0]
sum_enabled = np.abs(to_np(frames_normal) - to_np(frames_enable_free_init)).sum()

self.assertGreater(
sum_enabled,
1e1,
"Enabling of FreeInit should lead to results different from the default pipeline results",
)

@unittest.skipIf(
torch_device != "cuda" or not is_xformers_available(),
reason="XFormers attention is only available with CUDA and `xformers` installed",
Expand Down

0 comments on commit 57a021d

Please sign in to comment.