handler.py

from typing import Dict, List, Any
from diffusers import DPMSolverMultistepScheduler, StableDiffusionXLPipeline
from PIL import Image
from io import BytesIO
import base64
import torch

def merge_images(original, new_image, offset, direction):
    if direction in ["left", "right"]:
        merged_image = np.zeros((original.shape[0], original.shape[1] + offset, 3), dtype=np.uint8)
    elif direction in ["top", "bottom"]:
        merged_image = np.zeros((original.shape[0] + offset, original.shape[1], 3), dtype=np.uint8)

    if direction == "left":
        merged_image[:, offset:] = original
        merged_image[:, : new_image.shape[1]] = new_image
    elif direction == "right":
        merged_image[:, : original.shape[1]] = original
        merged_image[:, original.shape[1] + offset - new_image.shape[1] : original.shape[1] + offset] = new_image
    elif direction == "top":
        merged_image[offset:, :] = original
        merged_image[: new_image.shape[0], :] = new_image
    elif direction == "bottom":
        merged_image[: original.shape[0], :] = original
        merged_image[original.shape[0] + offset - new_image.shape[0] : original.shape[0] + offset, :] = new_image

    return merged_image


def slice_image(image):
    height, width, _ = image.shape
    slice_size = min(width // 2, height // 3)

    slices = []

    for h in range(3):
        for w in range(2):
            left = w * slice_size
            upper = h * slice_size
            right = left + slice_size
            lower = upper + slice_size

            if w == 1 and right > width:
                left -= right - width
                right = width
            if h == 2 and lower > height:
                upper -= lower - height
                lower = height

            slice = image[upper:lower, left:right]
            slices.append(slice)

    return slices


def process_image(
    image,
    fill_color=(0, 0, 0),
    mask_offset=50,
    blur_radius=500,
    expand_pixels=256,
    direction="left",
    inpaint_mask_color=50,
    max_size=1024,
):
    height, width = image.shape[:2]

    new_height = height + (expand_pixels if direction in ["top", "bottom"] else 0)
    new_width = width + (expand_pixels if direction in ["left", "right"] else 0)

    if new_height > max_size:
        # If so, crop the image from the opposite side
        if direction == "top":
            image = image[:max_size, :]
        elif direction == "bottom":
            image = image[new_height - max_size :, :]
        new_height = max_size

    if new_width > max_size:
        # If so, crop the image from the opposite side
        if direction == "left":
            image = image[:, :max_size]
        elif direction == "right":
            image = image[:, new_width - max_size :]
        new_width = max_size

    height, width = image.shape[:2]

    new_image = np.full((new_height, new_width, 3), fill_color, dtype=np.uint8)
    mask = np.full_like(new_image, 255, dtype=np.uint8)
    inpaint_mask = np.full_like(new_image, 0, dtype=np.uint8)

    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
    inpaint_mask = cv2.cvtColor(inpaint_mask, cv2.COLOR_BGR2GRAY)

    if direction == "left":
        new_image[:, expand_pixels:] = image[:, : max_size - expand_pixels]
        mask[:, : expand_pixels + mask_offset] = inpaint_mask_color
        inpaint_mask[:, :expand_pixels] = 255
    elif direction == "right":
        new_image[:, :width] = image
        mask[:, width - mask_offset :] = inpaint_mask_color
        inpaint_mask[:, width:] = 255
    elif direction == "top":
        new_image[expand_pixels:, :] = image[: max_size - expand_pixels, :]
        mask[: expand_pixels + mask_offset, :] = inpaint_mask_color
        inpaint_mask[:expand_pixels, :] = 255
    elif direction == "bottom":
        new_image[:height, :] = image
        mask[height - mask_offset :, :] = inpaint_mask_color
        inpaint_mask[height:, :] = 255

    # mask blur
    if blur_radius % 2 == 0:
        blur_radius += 1
    mask = cv2.GaussianBlur(mask, (blur_radius, blur_radius), 0)

    # telea inpaint
    _, mask_np = cv2.threshold(inpaint_mask, 128, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    inpaint = cv2.inpaint(new_image, mask_np, 3, cv2.INPAINT_TELEA)

    # convert image to tensor
    inpaint = cv2.cvtColor(inpaint, cv2.COLOR_BGR2RGB)
    inpaint = torch.from_numpy(inpaint).permute(2, 0, 1).float()
    inpaint = inpaint / 127.5 - 1
    inpaint = inpaint.unsqueeze(0).to("cuda")

    # convert mask to tensor
    mask = torch.from_numpy(mask)
    mask = mask.unsqueeze(0).float() / 255.0
    mask = mask.to("cuda")

    return inpaint, mask


def image_resize(image, new_size=1024):
    height, width = image.shape[:2]

    aspect_ratio = width / height
    new_width = new_size
    new_height = new_size

    if aspect_ratio != 1:
        if width > height:
            new_height = int(new_size / aspect_ratio)
        else:
            new_width = int(new_size * aspect_ratio)

    image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_LANCZOS4)

    return image


class EndpointHandler():
    def __init__(self, path=""):
        self.pipeline = StableDiffusionXLPipeline.from_pretrained(
            "SG161222/RealVisXL_V4.0",
            torch_dtype=torch.float16,
            variant="fp16",
            custom_pipeline="pipeline_stable_diffusion_xl_differential_img2img",
        ).to("cuda")
        self.pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config, use_karras_sigmas=True)

        self.pipeline.load_ip_adapter(
            "h94/IP-Adapter",
            subfolder="sdxl_models",
            weight_name=[
                "ip-adapter-plus_sdxl_vit-h.safetensors",
            ],
            image_encoder_folder="models/image_encoder",
        )
        self.pipeline.set_ip_adapter_scale(0.1)
    
    def generate_image(prompt, negative_prompt, image, mask, ip_adapter_image, seed: int = None):
        if seed is None:
            seed = random.randint(0, 2**32 - 1)
    
        generator = torch.Generator(device="cpu").manual_seed(seed)
    
        image = self.pipeline(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=1024,
            height=1024,
            guidance_scale=4.0,
            num_inference_steps=25,
            original_image=image,
            image=image,
            strength=1.0,
            map=mask,
            generator=generator,
            ip_adapter_image=[ip_adapter_image],
            output_type="np",
        ).images[0]
    
        image = (image * 255).astype(np.uint8)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
        return image

    def __call__(self, data: Dict[str, Any]):

        prompt = ""
        negative_prompt = ""
        # direction = "right"  # left, right, top, bottom
        inpaint_mask_color = 50  # lighter use more of the Telea inpainting
        # expand_pixels = 256  # I recommend to don't go more than half of the picture so it has context
        # times_to_expand = 4
        
        inputs = data.pop("inputs", data)

        # decode base64 image to PIL
        original = Image.open(BytesIO(base64.b64decode(inputs['image'])))
        mask = Image.open(BytesIO(base64.b64decode(inputs['mask'])))
        original = numpy.array(original)
        
        image = image_resize(original)
        expand_pixels_to_square = 1024 - image.shape[1]  # image.shape[1] for horizontal, image.shape[0] for vertical
        image, mask = process_image(
            image, expand_pixels=expand_pixels_to_square, direction=direction, inpaint_mask_color=inpaint_mask_color
        )
        
        ip_adapter_image = []
        for index, part in enumerate(slice_image(original)):
            ip_adapter_image.append(part)
        
        generated = generate_image(prompt, negative_prompt, image, mask, ip_adapter_image)
        final_image = generated
        
        for i in range(times_to_expand):
            image, mask = process_image(
                final_image, direction=direction, expand_pixels=expand_pixels, inpaint_mask_color=inpaint_mask_color
            )
        
            ip_adapter_image = []
            for index, part in enumerate(slice_image(generated)):
                ip_adapter_image.append(part)
        
            generated = generate_image(prompt, negative_prompt, image, mask, ip_adapter_image)
            final_image = merge_images(final_image, generated, 256, direction)