update app.py

app.py

@@ -58,7 +58,7 @@ kosmos_enable = False
 if os.environ.get('IS_MY_DEBUG') is not None:
     sam_enable = False
     ram_enable = False
-    inpainting_enable = False
+    # inpainting_enable = False
     kosmos_enable = False
 
 if lama_cleaner_enable:
@@ -79,8 +79,6 @@ from io import BytesIO
 from diffusers import StableDiffusionInpaintPipeline
 from huggingface_hub import hf_hub_download
 
-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
@@ -617,32 +615,61 @@ 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 load_kolors_inpainting(inpaint_prompt, input_image, mask_image):
+    from gradio_client import Client, handle_file
+    import tempfile
+    try:
+        job_image = {}
+        if 'background' in input_image.keys():
+            _, temp_file_path = tempfile.mkstemp(suffix='.png')
+            img = input_image['background'].convert("RGB")
+            img.save(temp_file_path)
+            job_image["background"] = handle_file(temp_file_path)
+        if 'layers' in input_image.keys() and len(input_image['layers']) > 0:
+            _, temp_file_path = tempfile.mkstemp(suffix='.png')
+            img = input_image['layers'][0].convert("RGB")
+            img.save(temp_file_path)
+            job_image["layers"] = [handle_file(temp_file_path)]
+        if 'composite' in input_image.keys():
+            _, temp_file_path = tempfile.mkstemp(suffix='.png')
+            img = input_image['composite'].convert("RGB")
+            img.save(temp_file_path)
+            job_image["composite"] = handle_file(temp_file_path)
+
+        _, temp_file_path = tempfile.mkstemp(suffix='.png')
+        img = mask_image.convert("RGB")
+        img.save(temp_file_path)
+        job_mask_image = handle_file(temp_file_path)
+
+        client = Client("Kwai-Kolors/Kolors-Inpainting")
+        
+        job = client.submit(
+                prompt=inpaint_prompt,
+                image=job_image,
+                mask_image=job_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"
+        )
+        while not job.done():
+            time.sleep(0.1)
+
+        result = job.outputs()[0]
+        im = Image.open(result)
+        if im.mode == "RGBA":
+            im.load()  
+            background = Image.new("RGB", im.size, (255, 255, 255))
+            background.paste(im, mask=im.split()[3]) 
+        return im 
+    except Exception as e:
+        logger.info(f'load_kolors_inpainting_fail_={str(e)}')
+        return None
 
 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):
-
     text_prompt = getTextTrans(text_prompt, source='zh', target='en')
     inpaint_prompt = getTextTrans(inpaint_prompt, source='zh', target='en')
 
@@ -824,7 +851,10 @@ 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)
             
             # 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)
+            image_inpainting = load_kolors_inpainting(inpaint_prompt, input_image, image_mask_for_inpaint)
+            if image_inpainting is None:
+                logger.info(f'load_kolors_inpainting_failed_')
+                return None, None, None, None, None, None, None
         else:
             # remove from mask
             if mask_source_radio == mask_source_segment:
@@ -1034,6 +1064,8 @@ def main_gradio(args):
             DESCRIPTION += f'Kosmos-2 from [Kosmos-2](https://github.com/microsoft/unilm/tree/master/kosmos-2). <br>'
         if ram_enable:
             DESCRIPTION += f'RAM from [RelateAnything](https://github.com/Luodian/RelateAnything). <br>'
+        if inpainting_enable:
+            DESCRIPTION += f'Inpainting from [Kolors-Inpainting](https://huggingface.co/spaces/Kwai-Kolors/Kolors-Inpainting). <br>'
         DESCRIPTION += f'Thanks for their excellent work.'
         DESCRIPTION += f'<p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings. \
                         <a href="https://huggingface.co/spaces/yizhangliu/Grounded-Segment-Anything?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a></p>'