update app.py

app.py

@@ -45,7 +45,7 @@ plt = matplotlib.pyplot
 
 groundingdino_enable = True
 sam_enable = True
-inpainting_enable = False #True
+inpainting_enable = True
 ram_enable = False
 
 lama_cleaner_enable = True
@@ -79,11 +79,13 @@ from io import BytesIO
 from diffusers import StableDiffusionInpaintPipeline
 from huggingface_hub import hf_hub_download
 
-from huggingface_hub import snapshot_download
-from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256_inpainting import StableDiffusionXLInpaintPipeline
-from kolors.models.modeling_chatglm import ChatGLMModel
-from kolors.models.tokenization_chatglm import ChatGLMTokenizer
-from diffusers import AutoencoderKL, EulerDiscreteScheduler, UNet2DConditionModel
+from gradio_client import Client, handle_file
+
+# from huggingface_hub import snapshot_download
+# from kolors.pipelines.pipeline_stable_diffusion_xl_chatglm_256_inpainting import StableDiffusionXLInpaintPipeline
+# from kolors.models.modeling_chatglm import ChatGLMModel
+# from kolors.models.tokenization_chatglm import ChatGLMTokenizer
+# from diffusers import AutoencoderKL, EulerDiscreteScheduler, UNet2DConditionModel
 
 from util_computer import computer_info
 
@@ -329,6 +331,7 @@ def load_sd_model(device):
     global sd_model
     logger.info(f"initialize stable-diffusion-inpainting...")
     sd_model = None
+    '''
     if os.environ.get('IS_MY_DEBUG') is None:
         # sd_model = StableDiffusionInpaintPipeline.from_pretrained(
         #         "runwayml/stable-diffusion-inpainting", 
@@ -355,6 +358,7 @@ def load_sd_model(device):
             
         sd_model.to(device)
         sd_model.enable_attention_slicing()
+    '''
 
 def load_lama_cleaner_model(device):
     # initialize lama_cleaner
@@ -613,6 +617,29 @@ def get_time_cost(run_task_time, time_cost_str):
     run_task_time = now_time
     return run_task_time, time_cost_str
 
+def load_kolors_inpainting(inpaint_prompt, image, mask_image):
+    # sd_model(prompt=inpaint_prompt, image=image_source_for_inpaint, mask_image=image_mask_for_inpaint).images[0]
+
+    client = Client("Kwai-Kolors/Kolors-Inpainting")
+    result = client.predict(
+            prompt=inpaint_prompt,
+            image=image,
+            mask_image = mask_image,
+            negative_prompt="broken fingers, deformed fingers, deformed hands, stumps, blurriness, low quality",
+            seed=0,
+            randomize_seed=True,
+            guidance_scale=6,
+            num_inference_steps=25,
+            api_name="/infer"
+    )
+    logger.info(f'load_kolors_inpainting_result={result}')
+    im = Image.open(result)
+    if im.mode == "RGBA":
+        im.load()  # required for png.split()
+        background = Image.new("RGB", im.size, (255, 255, 255))
+        background.paste(im, mask=im.split()[3]) 
+    return result
+
 def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_threshold, text_threshold, 
             iou_threshold, inpaint_mode, mask_source_radio, remove_mode, remove_mask_extend, num_relation, kosmos_input, cleaner_size_limit=1080):
 
@@ -624,6 +651,7 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
         
     # logger.info(f"input_image==={input_image}")
+    ori_input_image = input_image
     if 'background' in input_image.keys():
         input_image['image'] = input_image['background'].convert("RGB")
         if len(input_image['layers']) > 0:
@@ -794,7 +822,9 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
                 image_mask_for_inpaint = Image.fromarray(255*img_arr.astype('uint8'))
                 output_images.append(image_mask_for_inpaint.convert("RGB"))
                 run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
-            image_inpainting = sd_model(prompt=inpaint_prompt, image=image_source_for_inpaint, mask_image=image_mask_for_inpaint).images[0]
+            
+            # image_inpainting = sd_model(prompt=inpaint_prompt, image=image_source_for_inpaint, mask_image=image_mask_for_inpaint).images[0]
+            image_inpainting = load_kolors_inpainting(ori_input_image, image_source_for_inpaint, image_mask_for_inpaint).images[0])
         else:
             # remove from mask
             if mask_source_radio == mask_source_segment: