Update app.py

app.py

@@ -17,7 +17,7 @@ def resize(value,img):
 
 def infer(source_img, prompt, negative_prompt, guide, steps, seed):
     generator = torch.Generator(device).manual_seed(seed) 
-    imageio.imwrite("data.png", source_img['mask'])
+    imageio.imwrite("data.png", source_img)
     src = resize(768, 'data.png')
     image = pipe(prompt, negative_prompt=negative_prompt, image=src, strength=1, guidance_scale=guide, num_inference_steps=steps).images[0]
     return image