#装lama-cleaner==1.2.4来得到其他环境依赖，但是使用的时候用lama_cleaner这个文件的



import base64
import requests
import os
from typing import List
import io
from PIL import Image
from pydantic import BaseModel
from lama_cleaner.server import process
from lama_cleaner.server import main#先初始化才能用process
#C:\Users\17331\.cache\torch\hub  缓存下载位置,记得删

class FakeArgs(BaseModel):
    host: str = "0.0.0.0"
    port: int = 7860
    model: str = 'lama'
    hf_access_token: str = ""
    sd_enable_xformers: bool = False
    sd_disable_nsfw: bool = False
    sd_cpu_textencoder: bool = True
    sd_controlnet: bool = False
    sd_controlnet_method: str = "control_v11p_sd15_canny"
    sd_local_model_path: str = ""
    sd_run_local: bool = False
    local_files_only: bool = False
    cpu_offload: bool = False
    device: str = "cpu"
    gui: bool = False
    gui_size: List[int] = [1000, 1000]
    input: str = ''
    disable_model_switch: bool = True
    debug: bool = False
    no_half: bool = False
    disable_nsfw: bool = False
    enable_xformers: bool = False
    enable_interactive_seg: bool = True
    interactive_seg_model: str = "vit_b"
    interactive_seg_device: str = "cpu"
    enable_remove_bg: bool = False
    enable_anime_seg: bool = False
    enable_realesrgan: bool = False
    enable_gfpgan: bool = False
    gfpgan_device: str = "cpu"
    enable_restoreformer: bool = False
    enable_gif: bool = False
    quality: int = 95
    model_dir: str = None
    output_dir: str = None


def inpaint(img_path: str, mask_path: str) -> "img content (resp.content)":
    # urllib3 1.26.4  兼容
    image_bytes = open(img_path, 'rb')
    mask_bytes = open(mask_path, 'rb')
    # 将字节数据转换为Base64编码的字符串

    files = {
        "image": image_bytes,
        "mask": mask_bytes
    }
    payload = {
        "ldmSteps": 25,
        "ldmSampler": "plms",
        "zitsWireframe": True,
        "hdStrategy": "Crop",
        "hdStrategyCropMargin": 196,
        "hdStrategyCropTrigerSize": 800,
        "hdStrategyResizeLimit": 2048,
        "prompt": "",
        "negativePrompt": "",
        "croperX": 307,
        "croperY": 544,
        "croperHeight": 512,
        "croperWidth": 512,
        "useCroper": False,
        "sdMaskBlur": 5,
        "sdStrength": 0.75,
        "sdSteps": 50,
        "sdGuidanceScale": 7.5,
        "sdSampler": "uni_pc",
        "sdSeed": -1,
        "sdMatchHistograms": False,
        "sdScale": 1,
        "cv2Radius": 5,
        "cv2Flag": "INPAINT_NS",
        "paintByExampleSteps": 50,
        "paintByExampleGuidanceScale": 7.5,
        "paintByExampleSeed": -1,
        "paintByExampleMaskBlur": 5,
        "paintByExampleMatchHistograms": False,
        "p2pSteps": 50,
        "p2pImageGuidanceScale": 1.5,
        "p2pGuidanceScale": 7.5,
        "controlnet_conditioning_scale": 0.4,
        "controlnet_method": "control_v11p_sd15_canny"
    }  # payload用data

    resp = process(files=files, payload=payload)
    return resp




def save_img(img_content: "要处理的图片数据", new_save_path: "新文件的保存路径(包含后缀)",
             old_img_path: "旧文件路径") -> "void生成新的文件保存 ,传入旧文件路径是为了删除有问题的旧文件":
    print(new_save_path)
    try:
        # 从 _io.BytesIO 对象中读取字节数据
        resp_bytes = img_content.read()
        # 使用 Image.open() 方法打开图片
        img = Image.open(io.BytesIO(resp_bytes))
        # 如果需要指定图像格式，可以在保存时指定
        img.save(new_save_path, format="JPEG")
    except Exception as e:
        # 对于可能异常的图片->比如因为不合规导致resp.content没有正常返回
        print(e, new_save_path, "图片返回有问题,跳过并删除图片.这里的路径是新保存路径")
        os.remove(old_img_path)


# 传入遮罩图片路径和需要去水印的图片路径,将调整大小后的mask保存 ->void
def mask_resize(maskPath:str,removeMarkImagePath:str):
    maskImg = Image.open(maskPath)
    # 定义新的图片大小（宽度，高度）
    new_size = Image.open(removeMarkImagePath).size  # 例如，将图片调整为宽400像素，高300像素
    print("遮罩大小是:", maskImg.size, "不匹配图片大小是:", new_size,"不匹配图片路径是",removeMarkImagePath)
    # 调整图片大小
    resized_img = maskImg.resize(new_size)
    # 保存调整大小后的图片
    resized_img.save(maskPath)

if __name__ == '__main__':
    main(FakeArgs())#初始化model
    # 获取当前目录的子目录的路径
    img_path = 'manga'
    subdir_path = os.path.join(os.getcwd(), img_path)

    # 图片素材获取（包含子目录下所有图片）
    image_files = []
    for root, dirs, files in os.walk(subdir_path):
        for file in files:
            if file.endswith(".jpg") or file.endswith(".png"):
                image_files.append(os.path.relpath(os.path.join(root, file)))

    # 创建处理后的子目录在与image_files同级目录下
    processed_subdir_path = os.path.join(os.path.dirname(subdir_path), f"{img_path}1")
    os.makedirs(processed_subdir_path, exist_ok=True)

    # 对image_files进行某种处理，生成新图片，并保存在处理后的子目录中
    for img_file in image_files:
        # 处理图片的代码（这里仅作示例）
        # 假设处理后的图片为new_img
        img_dir = os.path.dirname(img_file)
        new_img_dir = os.path.join(processed_subdir_path, img_dir)
        os.makedirs(new_img_dir, exist_ok=True)

        new_img_path = os.path.join(new_img_dir, os.path.basename(img_file))

        if not os.path.exists(new_img_path):
            # 如果已经处理过那么跳过
            # 处理图片并保存 ->每次处理请求的时候都要调整mask大小，使其和img大小一致,这样就可以在taskManger的时候减轻判断负担
            mask_resize(maskPath ='mask/0.jpg', removeMarkImagePath=img_file)

            img_inpainted = inpaint(img_path=img_file, mask_path='mask/0.jpg')  # 上传的遮罩保存都是0开始
            save_img(img_content=img_inpainted, new_save_path=new_img_path, old_img_path=img_file)
        else:
            print(f"Skipping {new_img_path} as it already exists.")