Update README.md

README.md

@@ -31,11 +31,13 @@ from diffusers.utils import load_image
 import torch
 import numpy as np
 
-pipe_prior = KandinskyPriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-1-prior", torch_dtype=torch.float16)
+pipe_prior = KandinskyPriorPipeline.from_pretrained(
+    "kandinsky-community/kandinsky-2-1-prior", torch_dtype=torch.float16
+)
 pipe_prior.to("cuda")
 
 prompt = "a hat"
-image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False)
+prior_output = pipe_prior(prompt)
 
 pipe = KandinskyInpaintPipeline.from_pretrained("kandinsky-community/kandinsky-2-1-inpaint", torch_dtype=torch.float16)
 pipe.to("cuda")
@@ -45,14 +47,14 @@ init_image = load_image(
 )
 
 mask = np.ones((768, 768), dtype=np.float32)
+# Let's mask out an area above the cat's head
 mask[:250, 250:-250] = 0
 
 out = pipe(
     prompt,
     image=init_image,
     mask_image=mask,
-    image_embeds=image_emb,
-    negative_image_embeds=zero_image_emb,
+    **prior_output,
     height=768,
     width=768,
     num_inference_steps=150,