# https://medium.com/@qacheampong/building-and-deploying-a-fastapi-app-with-hugging-face-9210e9b4a713
# https://huggingface.co/spaces/Queensly/FastAPI_in_Docker

from fastapi import FastAPI,Request
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

app = FastAPI()
pipeline = get_pipeline()

#Endpoints 
#Root endpoints
@app.get("/")
def root():
    return {"API": "Sum of 2 Squares"}
    
@app.post("/img2img")
async def predict(url:str,prompt:str):
    MAX_QUEUE_SIZE = 4
    start = time.time()
   
    
    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=256,
        height=256,
        guidance_scale=1,
        num_inference_steps=4,
    )
    output_image = result.images[0]
    end2 = time.time()
    print("测试",output_image)
    print("s生成完成:", end2 - end1)
    
    output_image.save("./imageclm5.png")
    # 将图片对象转换为bytes
    image_bytes = output_image.to_bytes() 
    
    # 对bytes进行base64编码    
    encoded_string = base64.b64encode(image_bytes).decode('utf-8')
    return encoded_string
        
    
@app.post("/predict")
async def predict(request:Request):
  body = await request.body()
  data = json.loads(body)
  prompt = data.get("prompt")
  return f"您好,{prompt}"