update i2i pipeline

app.py

@@ -30,7 +30,10 @@ pipe = diffusers.StableDiffusionPipeline.from_pretrained(
 pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(
   model_id,
   torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-  scheduler=scheduler)
+  scheduler=scheduler,
+  vae=vae
+)
+
 
 if torch.cuda.is_available():
   pipe = pipe.to("cuda")