Skip to content

Commit

Permalink
Fix depth pipeline for latest Diffusers
Browse files Browse the repository at this point in the history
  • Loading branch information
carson-katri committed Jun 2, 2024
1 parent db13378 commit 1aab036
Showing 1 changed file with 4 additions and 5 deletions.
9 changes: 4 additions & 5 deletions generator_process/actions/depth_to_image.py
Original file line number Diff line number Diff line change
Expand Up @@ -194,7 +194,7 @@ def __call__(
width = width or self.unet.config.sample_size * self.vae_scale_factor

# 1. Check inputs
self.check_inputs(prompt, height, width, strength, callback_steps)
self.check_inputs(prompt=prompt, image=image, mask_image=depth_image, height=height, width=width, strength=strength, callback_steps=callback_steps, output_type=output_type)

# 2. Define call parameters
batch_size = 1 if isinstance(prompt, str) else len(prompt)
Expand Down Expand Up @@ -366,11 +366,10 @@ def __call__(

# Inference
with torch.inference_mode() if device not in ('mps', "dml") else nullcontext():
def callback(pipe, step, timestep, callback_kwargs):
def callback(step, _, latents):
if future.check_cancelled():
raise InterruptedError()
future.add_response(step_latents(pipe, step_preview_mode, callback_kwargs["latents"], generator, step, steps))
return callback_kwargs
future.add_response(step_latents(pipe, step_preview_mode, latents, generator, step, steps))
try:
result = pipe(
prompt=prompt,
Expand All @@ -383,7 +382,7 @@ def callback(pipe, step, timestep, callback_kwargs):
num_inference_steps=steps,
guidance_scale=cfg_scale,
generator=generator,
callback_on_step_end=callback,
callback=callback,
callback_steps=1,
output_type="np"
)
Expand Down

0 comments on commit 1aab036

Please sign in to comment.