import io
import cv2
import base64
import requests
from PIL import Image

class ControlnetRequestImg2Img:
    def __init__(self, prompt, net_prompt, path, mask):
        self.url = "http://127.0.0.1:57860/sdapi/v1/img2img"
        self.prompt = prompt
        self.neg_prompt = net_prompt
        self.img_path = path
        self.mask = mask
        self.body = None

    def build_body(self):
        img = cv2.imread(self.img_path)
        self.body = {
            "prompt": self.prompt,
            "negative_prompt": self.neg_prompt,
            "sampler_name": "DPM++ 2M Karras",
            "batch_size": 1,
            "steps": 30,
            "width": img.shape[1],
            "height": img.shape[0],
            "cfg_scale": 7,
            "seed": -1,
            "mask_blur": 15,
            "init_images": [
                self.read_image()
            ],
            "inpaint_full_res": True,
            "inpainting_fill": 1,
            "inpainting_mask_invert": 1,
            "mask": self.read_mask(),
            "denoising_strength": 0.4,
            "alwayson_scripts": {
                "controlnet": {
                    "args": [
                        {
                            "enabled": True,
                            "module": "openpose_full",
                            "model": "openpose",
                            "weight": 1.0,
                            "resize_mode": 1,
                            "lowvram": False,
                            "processor_res": 512,
                            "guidance_start": 0.0,
                            "guidance_end": 1.0,
                            "control_mode": 0,
                            "pixel_perfect": True
                        },
                    ]
                },
            }
        }

    def send_request(self):
        response = requests.post(url=self.url, json=self.body)
        return response.json()

    def read_image(self):
        img = cv2.imread(self.img_path)
        retval, bytes = cv2.imencode('.png', img)
        encoded_image = base64.b64encode(bytes).decode('utf-8')
        return encoded_image

    def read_mask(self):
        img = cv2.imread(self.mask)
        retval, bytes = cv2.imencode('.png', img)
        encoded_image = base64.b64encode(bytes).decode('utf-8')
        return encoded_image


if __name__ == '__main__':
    path = '/home/chinatszrn/Downloads/user1_hr.png'
    mask_path = '/home/chinatszrn/Downloads/user1_hr_mask.png'
    prompt = 'a woman with long blonde hair and a blue shirt on a gray background with a gray background and a gray background, lyco art, An Gyeon, realistic face, a character portrait'
    neg_prompt = '(nsfw:1.5), ng_deepnegative_v1_75t, (badhandv4:1.2), (worst quality:2), (low quality:2), (normal quality:2), lowres, bad anatomy, bad hands, ((monochrome)), ((grayscale)) watermark, moles, large breast, big breast, bad_pictures,easynegative'

    control_net = ControlnetRequestImg2Img(prompt, neg_prompt, path, mask_path)
    control_net.build_body()
    output = control_net.send_request()
    result = output['images'][0]
    image = Image.open(io.BytesIO(base64.b64decode(result.split(",", 1)[0])))
    image.save('save2.png')