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
from fastapi import FastAPI
import uvicorn
#urllib3 1.26.4  兼容
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

#初始化model,已经排除了flask部分
main(FakeArgs())


app = FastAPI()
@app.on_event("startup")
async def app_start():
    image_bytes = open('image.jpg', 'rb')
    mask_bytes = open('mask.jpg', '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)
    
    # 从 _io.BytesIO 对象中读取字节数据
    resp_bytes = resp.read()
    # 使用 Image.open() 方法打开图片
    img = Image.open(io.BytesIO(resp_bytes))
    img.show()

if __name__ == '__main__':
    uvicorn.run(app, host='0.0.0.0', port=7860)