from fastapi import FastAPI,Body

import uvicorn
import json

from PIL import Image
import time
from constants import DESCRIPTION, LOGO
from model import get_pipeline
from utils import replace_background
from diffusers.utils import load_image
import base64
import io
from datetime import datetime

app = FastAPI(name="mutilParam")
pipeline = get_pipeline()

#Endpoints 
#Root endpoints
@app.get("/")
def root():
    return {"API": "Sum of 2 Squares"}
    
@app.post("/img2img")
async def predict(prompt=Body(...),imgbase64data=Body(...)):
    MAX_QUEUE_SIZE = 4
    start = time.time()
    pipeline = get_pipeline()
    
    url = "https://img2.baidu.com/it/u=1845675188,2679793929&fm=253&fmt=auto&app=138&f=JPEG?w=667&h=500"
    prompt = "a nice Comfortable and clean. According to Baidu Education Information, the adjectives for a room include: comfortable, clean, beautiful, spacious, warm, quiet, luxurious, pleasant, exquisite, and warm ,colorful, light room width sofa,8k"
    
    init_image = load_image(url).convert("RGB")
    # image1 = replace_background(init_image.resize((256, 256)))
    w, h = init_image.size
    newW = 512
    newH = int(h * newW / w)
    
    img = init_image.resize((newW, newH))  
    end1 = time.time()
    print("加载管道:", end1 - start)
    result = pipeline(
        prompt=prompt,
        image=img,
        strength=0.6,
        seed=10,
        width=512,
        height=512,
        guidance_scale=1,
        num_inference_steps=4,
    )
    output_image = result.images[0]
    end2 = time.time()
    print("测试",output_image)
    print("s生成完成:", end2 - end1)
    end2 = time.time()
    print("测试",output_image)
    print("s生成完成:", end2 - end1)    
    # 将图片对象转换为bytes
    end3 = time.time()
    output_image_base64 = base64.b64encode(output_image.tobytes()).decode()
    print("完成的图片:", output_image_base64) 
    print("图像转换时间:", end3 - end2)
    return output_image_base64
        
    
@app.post("/predict")
async def predict(prompt=Body(...)):
  return f"您好,{prompt}"