Update main.py

main.py

@@ -27,35 +27,42 @@ def root():
 async def predict(prompt=Body(...),imgbase64data=Body(...)):
     MAX_QUEUE_SIZE = 4
     start = time.time()
-    print("参数",imgbase64data,prompt)
-    image_data = base64.b64decode(imgbase64data)
-    image1 = Image.open(io.BytesIO(image_data))
-    w, h = image1.size
-    newW = 256
+    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 = image1.resize((newW, newH))  
+    
+    img = init_image.resize((newW, newH))  
     end1 = time.time()
-    now = datetime.now()
-    print(now)
-    print("图像:", img.size)
     print("加载管道:", end1 - start)
     result = pipeline(
         prompt=prompt,
-        image=image1,
+        image=img,
         strength=0.6,
         seed=10,
-        width=256,
-        height=256,
+        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