1 files changed, 9 insertions, 5 deletions

diff --git a/pipelines/stable_diffusion/vlpn_stable_diffusion.py b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
index f27be78..f426de1 100644
--- a/pipelines/stable_diffusion/vlpn_stable_diffusion.py
+++ b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
@@ -307,10 +307,14 @@ class VlpnStableDiffusion(DiffusionPipeline):
         return timesteps, num_inference_steps - t_start
 
     def prepare_image(self, batch_size, width, height, dtype, device, generator=None):
-        noise = perlin_noise(
-            batch_size, 1, width, height, res=1, octaves=4, generator=generator, dtype=dtype, device=device
-        ).expand(batch_size, 3, width, height)
-        return (1.4 * noise).clamp(-1, 1)
+        return (1.4 * perlin_noise(
+            (batch_size, 1, width, height),
+            res=1,
+            octaves=4,
+            generator=generator,
+            dtype=dtype,
+            device=device
+        )).clamp(-1, 1).expand(batch_size, 3, width, height)
 
     def prepare_latents_from_image(self, init_image, timestep, batch_size, dtype, device, generator=None):
         init_image = init_image.to(device=device, dtype=dtype)
@@ -390,7 +394,7 @@ class VlpnStableDiffusion(DiffusionPipeline):
         sag_scale: float = 0.75,
         eta: float = 0.0,
         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        image: Optional[Union[torch.FloatTensor, PIL.Image.Image, Literal["noise"]]] = None,
+        image: Optional[Union[torch.FloatTensor, PIL.Image.Image, Literal["noise"]]] = "noise",
         output_type: str = "pil",
         return_dict: bool = True,
         callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,