Update app.py

app.py

@@ -318,18 +318,87 @@ def model_process_filepath(input): #image, mask):
     image = Image.open(io.BytesIO(numpy_to_bytes(res_np_img, ext)))
     image.save(f'./result_image.png')
     return  image # image
-    '''
-    ext = get_image_ext(origin_image_bytes)
 
-    response = make_response(
-        send_file(
-            io.BytesIO(numpy_to_bytes(res_np_img, ext)),
-            mimetype=f"image/{ext}",
-        )
+def model_process(image, mask, alpha_channel, ext):
+    global model
+    original_shape = image.shape
+    interpolation = cv2.INTER_CUBIC
+
+    # image_pil = Image.fromarray(image)
+    # mask_pil = Image.fromarray(mask).convert("L")
+    
+    size_limit = "Original"
+    print(f'size_limit_2_ =  {size_limit}')
+    if size_limit == "Original":
+        size_limit = max(image.shape)
+    else:
+        size_limit = int(size_limit)
+    print(f'size_limit_3_ =  {size_limit}')
+
+    config = Config(
+        ldm_steps=25,
+        ldm_sampler='plms',
+        zits_wireframe=True,
+        hd_strategy='Original',
+        hd_strategy_crop_margin=196,
+        hd_strategy_crop_trigger_size=1280,
+        hd_strategy_resize_limit=2048,
+        prompt='',
+        use_croper=False,
+        croper_x=0,
+        croper_y=0,
+        croper_height=512,
+        croper_width=512,
+        sd_mask_blur=5,
+        sd_strength=0.75,
+        sd_steps=50,
+        sd_guidance_scale=7.5,
+        sd_sampler='ddim',
+        sd_seed=42,
+        cv2_flag='INPAINT_NS',
+        cv2_radius=5,
     )
-    response.headers["X-Seed"] = str(config.sd_seed)
-    return response
-    '''
+    
+    print(f'config/alpha_channel/size_limit = {config} / {alpha_channel} / {size_limit}')
+    if config.sd_seed == -1:
+        config.sd_seed = random.randint(1, 999999999)
+
+    logger.info(f"Origin image shape: {original_shape}")
+    print(f"Origin image shape: {original_shape} / {image[250][250]}")
+    image = resize_max_size(image, size_limit=size_limit, interpolation=interpolation)
+    logger.info(f"Resized image shape: {image.shape} / {type(image)}")
+    print(f"Resized image shape: {image.shape} / {image[250][250]}")
+
+    mask = resize_max_size(mask, size_limit=size_limit, interpolation=interpolation)
+    print(f"mask image shape: {mask.shape} / {type(mask)} / {mask[250][250]} / {alpha_channel}")
+
+    if model is None:
+        return None
+        
+    start = time.time()
+    res_np_img = model(image, mask, config)
+    logger.info(f"process time: {(time.time() - start) * 1000}ms, {res_np_img.shape}")
+    print(f"process time_1_: {(time.time() - start) * 1000}ms, {alpha_channel.shape}, {res_np_img.shape} / {res_np_img[250][250]} / {res_np_img.dtype}")
+
+    torch.cuda.empty_cache()
+
+    if alpha_channel is not None:
+        print(f"liuyz_here_10_: {alpha_channel.shape} / {alpha_channel.dtype} / {res_np_img.dtype}") 
+        if alpha_channel.shape[:2] != res_np_img.shape[:2]:
+            print(f"liuyz_here_20_: {alpha_channel.shape} / {res_np_img.shape}")
+            alpha_channel = cv2.resize(
+                alpha_channel, dsize=(res_np_img.shape[1], res_np_img.shape[0])
+            )
+        print(f"liuyz_here_30_: {alpha_channel.shape} / {res_np_img.shape} / {alpha_channel.dtype} / {res_np_img.dtype}")
+        res_np_img = np.concatenate(
+            (res_np_img, alpha_channel[:, :, np.newaxis]), axis=-1
+        )
+        print(f"liuyz_here_40_: {alpha_channel.shape} / {res_np_img.shape} / {alpha_channel.dtype} / {res_np_img.dtype}")
+    
+    print(f"process time_2_: {(time.time() - start) * 1000}ms, {alpha_channel.shape}, {res_np_img.shape} / {res_np_img[250][250]} / {res_np_img.dtype} /{ext}")
+  
+    image = Image.open(io.BytesIO(numpy_to_bytes(res_np_img, ext)))
+    return  image # image
     
 model = ModelManager(
         name='lama',
@@ -377,9 +446,26 @@ def predict(input):
     # mask = dict["mask"] # .convert("RGB") #.resize((512, 512))
     '''   
     if image_type == 'filepath':
-        output = model_process_filepath(input) # dict["image"], dict["mask"])
+        # input: {'image': '/tmp/tmp8mn9xw93.png', 'mask': '/tmp/tmpn5ars4te.png'}
+        origin_image_bytes = read_content(input["image"])
+        print(f'origin_image_bytes = ', type(origin_image_bytes), len(origin_image_bytes))    
+        image, _ = load_img(origin_image_bytes) 
+        mask, _ = load_img(read_content(input["mask"]), gray=True)       
+        alpha_channel = (np.ones((image.shape[0],image.shape[1]))*255).astype(np.uint8)  
+        ext = get_image_ext(origin_image_bytes)
+        
+        output = model_process(image, mask, alpha_channel, ext)
     elif image_type == 'pil':
-        output = model_process_pil(input)
+        # input: {'image': pil, 'mask': pil}
+        image_pil = input['image']
+        mask_pil = input['mask']
+            
+        image = np.array(image_pil)
+        mask = np.array(mask_pil.convert("L"))
+        alpha_channel = (np.ones((image.shape[0],image.shape[1]))*255).astype(np.uint8)  
+        ext = 'png'
+        
+        output = model_process(image, mask, alpha_channel, ext)
     
     # output = mask #output.images[0]
     # output = pipe(prompt = prompt, image=init_image, mask_image=mask,guidance_scale=7.5)