Update main.py

main.py

@@ -12,7 +12,7 @@ from model import get_pipeline
 from utils import replace_background
 from diffusers.utils import load_image
 import base64
-
+import io
 app = FastAPI()
 pipeline = get_pipeline()
 
@@ -23,21 +23,18 @@ def root():
     return {"API": "Sum of 2 Squares"}
     
 @app.post("/img2img")
-async def predict(url:str,prompt:str):
+async def predict(prompt=Body(...),imgbase64data=Body(...)):
     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
+    print("参数",imgbase64data,prompt)
+    image_data = base64.b64decode(imgbase64data)
+    image1 = Image.open(io.BytesIO(image_data))
+    w, h = image1.size
     newW = 512
     newH = int(h * newW / w)
-    img = init_image.resize((newW, newH))  
+    img = image1.resize((newW, newH))  
     end1 = time.time()
+    print("图像:", img.size)
     print("加载管道:", end1 - start)
     result = pipeline(
         prompt=prompt,
@@ -52,21 +49,13 @@ async def predict(url:str,prompt:str):
     output_image = result.images[0]
     end2 = time.time()
     print("测试",output_image)
-    print("s生成完成:", end2 - end1)
-    
-    output_image.save("./imageclm5.png")
+    print("s生成完成:", end2 - end1)    
     # 将图片对象转换为bytes
-    image_bytes = output_image.to_bytes() 
-    
-    # 对bytes进行base64编码    
-    encoded_string = base64.b64encode(image_bytes).decode('utf-8')
-    return encoded_string
+    output_image_base64 = base64.b64encode(output_image.tobytes()).decode()
+    return output_image_base64
         
     
 @app.post("/predict")
-async def predict(request:Request):
-  body = await request.body()
-  data = json.loads(body)
-  prompt = data.get("prompt")
+async def predict(prompt=Body(...)):
   return f"您好,{prompt}"