import os

from Util.CommonUtil import *

# API地址
url = "http://192.168.1.21:7860"


# 图生图函数
def img2img(url, model_id, prompt_id, input_source_img, output_file):
    # JSON文件模板
    templet_file = r'../JSON/' + str(model_id) + '_' + str(prompt_id) + '.json'
    # 模型图片
    input_refer_img = '../Image/' + str(model_id) + '/' + str(prompt_id) + '.jpg'

    if not os.path.exists(input_refer_img):
        input_refer_img = '../Image/' + str(model_id) + '/' + str(prompt_id) + '.png'

    # 读取JSON文件
    with open(templet_file, 'r', encoding='utf-8') as file:
        payload = json.load(file)

    # 用户输入的图片
    source_img = encode_image(input_source_img)
    # 模型固定的参考图片
    refer_img = encode_image(input_refer_img)

    # 在ControlNet的第一个图片中使用用户输入的图片，用此图片的人脸进行换脸操作
    payload["alwayson_scripts"]["ControlNet"]["args"][0]["input_image"] = source_img
    # 图生图，风格参考图是输入
    payload["init_images"] = [refer_img]
    payload["alwayson_scripts"]["ControlNet"]["args"][1]["input_image"] = refer_img
    # 调用SD服务
    response = requests.post(url=f'{url}/sdapi/v1/img2img', json=payload)
    if response.status_code != 200:
        print(response.text)
        exit(0)
    r = response.json()
    # 保存图片
    save_encoded_image(response.json()['images'][0], output_file)


# 用户输入的源图
input_source_img = '../Image/1ae93e2587822c291362abf9504cb49ec.jpeg'
# 输出图片
output_file = "../Out/Result.png"
# 调用图生图
img2img(url, 12, 1, input_source_img, output_file)