paste inpainted image back to original image to make sure size is unchanged

tasks/ref_in.py

@@ -37,7 +37,7 @@ def crop_image(input_image):
 def referring_inpainting(model, image, texts, inpainting_text, *args, **kwargs):
     model.model.metadata = metadata
     texts = [[texts if texts.strip().endswith('.') else (texts.strip() + '.')]]
-    image_ori = crop_image(transform(image))
+    image_ori = transform(image)
 
     with torch.no_grad():
         width = image_ori.size[0]
@@ -58,20 +58,20 @@ def referring_inpainting(model, image, texts, inpainting_text, *args, **kwargs):
     
     if inpainting_text not in ['no', '']:
         # if we want to do inpainting
-        image_ori = image_ori.convert('RGB')
+        image_crop = crop_image(image_ori.convert('RGB'))
         struct2 = ndimage.generate_binary_structure(2, 2)
         mask_dilated = ndimage.binary_dilation(grd_mask[0], structure=struct2, iterations=3).astype(grd_mask[0].dtype)
-        mask = Image.fromarray(mask_dilated * 255).convert('RGB')
-        # image_ori = pad_image(image_ori)
-        # mask = pad_image(Image.fromarray(grd_mask[0] * 255).convert('RGB'))
+        mask = crop_image(Image.fromarray(mask_dilated * 255).convert('RGB'))
         image_and_mask = {
-            "image": image_ori,
+            "image": image_crop,
             "mask": mask,
         }
-        width = image_ori.size[0]; height = image_ori.size[1]
-        images_inpainting = pipe(prompt = inpainting_text.strip(), image=image_and_mask['image'], mask_image=image_and_mask['mask'], height=height, width=width).images
+        width = image_crop.size[0]; height = image_crop.size[1]
+        images_inpainting = pipe(prompt = inpainting_text.strip(), image=image_and_mask['image'], mask_image=image_and_mask['mask'], height=height, width=width).images[0]
+        # put images_inpainting back to original image
+        image_ori.paste(images_inpainting)        
         torch.cuda.empty_cache()
-        return Image.fromarray(res) ,'' , images_inpainting[0]
+        return Image.fromarray(res) ,'' , image_ori
     else:
         torch.cuda.empty_cache()
         return image_ori, 'text', Image.fromarray(res)