Update app.py

app.py

@@ -71,6 +71,91 @@ def get_image_ext(img_bytes):
 def diffuser_callback(i, t, latents):
     pass
 
+def process(init_image, mask):
+    global model
+    '''
+    input = request.files
+    # RGB
+    origin_image_bytes = input["image"].read()
+    '''
+    
+    # image, alpha_channel = load_img(origin_image_bytes)
+    original_shape = init_image.shape
+    interpolation = cv2.INTER_CUBIC
+    
+    '''
+    form = request.form
+    '''
+    size_limit = 1080 # : Union[int, str] = form.get("sizeLimit", "1080")
+    if size_limit == "Original":
+        size_limit = max(image.shape)
+    else:
+        size_limit = int(size_limit)
+
+    config = Config(
+        '''
+        ldm_steps=form["ldmSteps"],
+        ldm_sampler=form["ldmSampler"],
+        hd_strategy=form["hdStrategy"],
+        zits_wireframe=form["zitsWireframe"],
+        hd_strategy_crop_margin=form["hdStrategyCropMargin"],
+        hd_strategy_crop_trigger_size=form["hdStrategyCropTrigerSize"],
+        hd_strategy_resize_limit=form["hdStrategyResizeLimit"],
+        prompt=form["prompt"],
+        use_croper=form["useCroper"],
+        croper_x=form["croperX"],
+        croper_y=form["croperY"],
+        croper_height=form["croperHeight"],
+        croper_width=form["croperWidth"],
+        sd_mask_blur=form["sdMaskBlur"],
+        sd_strength=form["sdStrength"],
+        sd_steps=form["sdSteps"],
+        sd_guidance_scale=form["sdGuidanceScale"],
+        sd_sampler=form["sdSampler"],
+        sd_seed=form["sdSeed"],
+        cv2_flag=form["cv2Flag"],
+        cv2_radius=form['cv2Radius']
+        '''
+    )
+
+    if config.sd_seed == -1:
+        config.sd_seed = random.randint(1, 999999999)
+
+    logger.info(f"Origin image shape: {original_shape}")
+    image = resize_max_size(image, size_limit=size_limit, interpolation=interpolation)
+    logger.info(f"Resized image shape: {image.shape}")
+
+    mask, _ = load_img(input["mask"].read(), gray=True)
+    mask = resize_max_size(mask, size_limit=size_limit, interpolation=interpolation)
+
+    start = time.time()
+    res_np_img = model(image, mask, config)
+    logger.info(f"process time: {(time.time() - start) * 1000}ms")
+
+    torch.cuda.empty_cache()
+
+    if alpha_channel is not None:
+        if alpha_channel.shape[:2] != res_np_img.shape[:2]:
+            alpha_channel = cv2.resize(
+                alpha_channel, dsize=(res_np_img.shape[1], res_np_img.shape[0])
+            )
+        res_np_img = np.concatenate(
+            (res_np_img, alpha_channel[:, :, np.newaxis]), axis=-1
+        )
+
+    ext = get_image_ext(origin_image_bytes)
+    return ext
+    '''
+    response = make_response(
+        send_file(
+            io.BytesIO(numpy_to_bytes(res_np_img, ext)),
+            mimetype=f"image/{ext}",
+        )
+    )
+    response.headers["X-Seed"] = str(config.sd_seed)
+    return response
+    '''
+    
 model = ModelManager(
         name='lama',
         device=device,
@@ -101,9 +186,11 @@ def read_content(file_path: str) -> str:
     return content
 
 def predict(dict, prompt=""):
-    init_image = dict["image"].convert("RGB").resize((512, 512))
-    mask = dict["mask"].convert("RGB").resize((512, 512))
-    output = pipe(prompt = prompt, image=init_image, mask_image=mask,guidance_scale=7.5)
+    init_image = dict["image"].convert("RGB") #.resize((512, 512))
+    mask = dict["mask"].convert("RGB") #.resize((512, 512))
+    output = process(init_image, mask)
+    # output = pipe(prompt = prompt, image=init_image, mask_image=mask,guidance_scale=7.5)
+    
     return output.images[0], gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)