models.py

import logging
from typing import List

import torch
import numpy as np
from PIL import Image

from helpers import flush, postprocess_image_masking, convolution
from pipelines import  get_inpainting_pipeline

LOGGING = logging.getLogger(__name__)


@torch.inference_mode()
def make_inpainting(positive_prompt: str,
                    image: Image,
                    mask_image: np.ndarray,
                    negative_prompt: str,
                    num_of_images: int,
                    resolution:int ) -> List[Image.Image]:
    print("make_inpainting", positive_prompt, image, mask_image, negative_prompt, num_of_images, resolution)

    """Method to make inpainting
    Args:
        positive_prompt (str): positive prompt string
        image (Image): input image
        mask_image (np.ndarray): mask image
        negative_prompt (str, optional): negative prompt string. Defaults to "".
    Returns:
        List[Image.Image]: list of generated images
    """
    pipe = get_inpainting_pipeline()
    mask_image = Image.fromarray((mask_image * 255).astype(np.uint8))
    mask_image_postproc = convolution(mask_image)

    flush()
    retList=[]
    for x in range(num_of_images):
        
        resp = pipe(image=image,
                        mask_image=mask_image,
                        prompt=positive_prompt,
                        negative_prompt=negative_prompt,
                        num_inference_steps=50,
                        height=resolution,
                        width=resolution,
                        )
        print("RESP !!!!",resp)
        generated_image = resp.images[0]
        generated_image = postprocess_image_masking(generated_image, image, mask_image_postproc)
        retList.append(generated_image)

    return retList