from typing import List, Union

import torch
from diffusers import StableDiffusionInpaintPipeline

from internals.pipelines.commons import AbstractPipeline
from internals.util.commons import disable_safety_checker, download_image


class InPainter(AbstractPipeline):
    def load(self):
        self.pipe = StableDiffusionInpaintPipeline.from_pretrained(
            "jayparmr/icbinp_v8_inpaint_v2",
            torch_dtype=torch.float16,
        ).to("cuda")
        disable_safety_checker(self.pipe)

    @torch.inference_mode()
    def process(
        self,
        image_url: str,
        mask_image_url: str,
        width: int,
        height: int,
        seed: int,
        prompt: Union[str, List[str]],
        negative_prompt: Union[str, List[str]],
    ):
        torch.manual_seed(seed)

        input_img = download_image(image_url).resize((width, height))
        mask_img = download_image(mask_image_url).resize((width, height))

        return self.pipe.__call__(
            prompt=prompt,
            image=input_img,
            mask_image=mask_img,
            height=height,
            width=width,
            negative_prompt=negative_prompt,
        ).images