import os
import requests
import base64

base_url = 'http://192.168.2.157:7860'

def txt2img(prompt, neg_prompt, id=None):
    preset_prompt = 'realistic, masterpiece, best quality, highres, 4k, extremely detailed, '
    preset_neg_prompt = '(worst quality:2), (low quality:2), (normal quality:2), lowres, normal quality, ((monochrome)),((grayscale)), skin spots, acnes, skin blemishes, age spot, (ugly:1.331), (duplicate:1.331), (morbid:1.21), (mutilated:1.21), (tranny:1.331), mutated hands, (poorly drawn hands:1.5), blurry, (bad anatomy:1.21), (bad proportions:1.331), extra limbs, (disfigured:1.331), (missing arms:1.331), (extra legs:1.331), (fused fingers:1.5), (too many fingers:1.5), (unclear eyes:1.331), lowers, bad hands, missing fingers, extra digit,bad hands, missing fingers, (((extra arms and legs))),'

    # 定义请求URL
    url = "http://192.168.2.157:7860/sdapi/v1/txt2img"

    # 定义请求体
    data = {
        "prompt": f"{preset_prompt} ${prompt}",
        "negative_prompt": f"{preset_neg_prompt}, ${neg_prompt}",
        "sampler_index": "DPM++ 2M",
        "save_images": "true",
        "width": "512", 
        "height": "720",
        "batch_size": 4,
        # "upscaler": "latent",  # Specify the upscaler type
        # "upscale_factor": 2  # Set the upscaling factor to 2
        # "hr_scale": 2,
        # "hr_upscaler": "latent"
    }

    # 发送POST请求
    response = requests.post(url, json=data)

    # 检查请求是否成功
    if response.status_code == 200:
        print("Request was successful.")
        print("Response:")
        print(response.json())
            # 提取base64编码的图像
        response_data = response.json()
        if 'images' in response_data:
            for i, image_base64 in enumerate(response_data['images']):
                # 解码base64字符串
                image_data = base64.b64decode(image_base64)
                
                # 将图像保存到本地文件
                if id == None:
                    id = 0
                image_filename = f"{id}_output_image_{i}.png"
                dir_path = os.path.join('.', 'regen_flow')
                # if dir_path does not exist, then create the dir
                if not os.path.exists(dir_path):
                    os.mkdir(dir_path)

                with open(os.path.join(dir_path, image_filename), 'wb') as image_file:
                    image_file.write(image_data)
                
                print(f"Image {i} saved as {image_filename}")
        else:
            print("No images found in the response.")
    else:
        print(f"Request failed with status code {response.status_code}")
        print("Response:")
        print(response.text)


# txt2img("an asian girl, kneel", "")