1 files changed, 11 insertions, 3 deletions

diff --git a/pipelines/stable_diffusion/vlpn_stable_diffusion.py b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
index c4f7401..242be29 100644
--- a/pipelines/stable_diffusion/vlpn_stable_diffusion.py
+++ b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
@@ -306,13 +306,19 @@ class VlpnStableDiffusion(DiffusionPipeline):
 
         return timesteps
 
-    def prepare_image(self, batch_size, width, height, dtype, device, generator=None):
-        return torch.randn(
+    def prepare_image(self, batch_size, width, height, max_offset, dtype, device, generator=None):
+        offset = (max_offset * (2 * torch.rand(
             (batch_size, 1, 1, 1),
             dtype=dtype,
             device=device,
             generator=generator
-        ).expand(batch_size, 3, width, height)
+        ) - 1)).expand(batch_size, 3, width, height)
+        image = (.1 * torch.normal(
+            mean=offset,
+            std=1,
+            generator=generator
+        )).clamp(-1, 1)
+        return image
 
     def prepare_latents(self, init_image, timestep, batch_size, dtype, device, generator=None):
         init_image = init_image.to(device=device, dtype=dtype)
@@ -376,6 +382,7 @@ class VlpnStableDiffusion(DiffusionPipeline):
         eta: float = 0.0,
         generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
         image: Optional[Union[torch.FloatTensor, PIL.Image.Image]] = None,
+        max_image_offset: float = 1.0,
         output_type: str = "pil",
         return_dict: bool = True,
         callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
@@ -469,6 +476,7 @@ class VlpnStableDiffusion(DiffusionPipeline):
                 batch_size * num_images_per_prompt,
                 width,
                 height,
+                max_image_offset,
                 prompt_embeds.dtype,
                 device,
                 generator