'''
The following code performs an inpainting process using a model obtained from SDXL inpainting model
Inpainting allows you to modify specific areas of an image while keeping the rest of the image unchanged.
This code uses a PNG image downloaded from Bria's RMBG 1.4, which includes an alpha channel representing the areas you wish to modify.

'''
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
from PIL import Image
import requests
from io import BytesIO
import PIL
from PIL import Image
import requests
import torch
from io import BytesIO
from diffusers import StableDiffusionXLInpaintPipeline, DDIMScheduler, UNet2DConditionModel

def download_image(url):
   response = requests.get(url)
   return PIL.Image.open(BytesIO(response.content)).convert("RGB")

image_path = "/home/ubuntu/spring/algo_resarch/Efrat_solutions/inpainting_blog/Efrat.PNG"
init_image = Image.open(image_path).convert("RGBA")

# Resize the image and mask to 1024x1024
init_image = init_image.resize((1024, 1024))

if init_image.mode == "RGBA":
   alpha_channel = init_image.getchannel('A')
   mask_image = alpha_channel.resize((1024, 1024))  # Resize mask as well

   # Invert the mask: what was white becomes black, and vice versa
   inverted_mask = PIL.ImageOps.invert(mask_image)

   # Convert the inverted mask to RGB
   mask_image = Image.merge("RGB", (inverted_mask, inverted_mask, inverted_mask))
   init_image = init_image.convert("RGB")
else:
   raise ValueError("The image does not have an alpha channel.")

# Load the pre-trained inpainting pipeline
pipeline = AutoPipelineForInpainting.from_pretrained(
   "kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch.float16
)
pipeline.enable_model_cpu_offload()

# Remove the following line if xFormers is not installed or you have PyTorch 2.0 or higher installed
pipeline.enable_xformers_memory_efficient_attention()

prompt = "A classic library with tall wooden bookshelves filled with old books, warm lighting, and a cozy, scholarly atmosphere. The setting should evoke a sense of wisdom and tradition."
negative_prompt = "deformed, ugly, disfigured"

image = pipeline(prompt=prompt, negative_prompt=negative_prompt, image=init_image, mask_image=mask_image).images[0]

generator = torch.Generator(device='cuda:0').manual_seed(123456)
image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, generator=generator, guidance_scale=5, strength=1).images[0]
image.save("./library_background_SDXL.png")