From 71d0e95edb8ea12a5bc791e5c5a998dcec0e3436 Mon Sep 17 00:00:00 2001
From: rangoliu <liuwenran@users.noreply.github.com>
Date: Sat, 9 Dec 2023 10:39:17 +0800
Subject: [PATCH] [Fix] add openxlab link to powerpaint (#2082)

* add openxlab link

* fix path

* update scale

* fix text-guided bug

* use 3.41
---
 projects/powerpaint/README.md            |  2 +-
 projects/powerpaint/gradio_PowerPaint.py | 16 ++++++++++------
 projects/powerpaint/requirements.txt     |  2 +-
 3 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/projects/powerpaint/README.md b/projects/powerpaint/README.md
index f122e47d2..7fd6d2cd0 100644
--- a/projects/powerpaint/README.md
+++ b/projects/powerpaint/README.md
@@ -1,6 +1,6 @@
 # A Task is Worth One Word: Learning with Task Prompts for High-Quality Versatile Image Inpainting
 
-### [Project Page](https://powerpaint.github.io/) | [Paper](https://arxiv.org/abs/2312.03594)
+### [Project Page](https://powerpaint.github.io/) | [Paper](https://arxiv.org/abs/2312.03594) | [OpenXLab](https://openxlab.org.cn/apps/detail/rangoliu/PowerPaint#basic-information)
 
 This README provides a step-by-step guide to download the repository, set up the required virtual environment named "PowerPaint" using conda, and run PowerPaint with or without ControlNet.
 
diff --git a/projects/powerpaint/gradio_PowerPaint.py b/projects/powerpaint/gradio_PowerPaint.py
index 86bf3f7fe..373245687 100644
--- a/projects/powerpaint/gradio_PowerPaint.py
+++ b/projects/powerpaint/gradio_PowerPaint.py
@@ -30,13 +30,13 @@
 add_tokens(
     tokenizer=pipe.tokenizer,
     text_encoder=pipe.text_encoder,
-    placeholder_tokens=['MMcontext', 'MMshape', 'MMobject'],
+    placeholder_tokens=['P_ctxt', 'P_shape', 'P_obj'],
     initialize_tokens=['a', 'a', 'a'],
     num_vectors_per_token=10)
 pipe.unet.load_state_dict(
-    torch.load('./models/diffusion_pytorch_model.bin'), strict=False)
+    torch.load('./models/unet/diffusion_pytorch_model.bin'), strict=False)
 pipe.text_encoder.load_state_dict(
-    torch.load('./models/pytorch_model.bin'), strict=False)
+    torch.load('./models/text_encoder/pytorch_model.bin'), strict=False)
 pipe = pipe.to('cuda')
 
 depth_estimator = DPTForDepthEstimation.from_pretrained(
@@ -258,9 +258,14 @@ def infer(input_image, text_guided_prompt, text_guided_negative_prompt,
     elif task == 'shape-guided':
         prompt = shape_guided_prompt
         negative_prompt = shape_guided_negative_prompt
-    else:
+    elif task == 'object-removal':
         prompt = ''
         negative_prompt = ''
+        scale = 10.0
+    else:
+        task = 'text-guided'
+        prompt = text_guided_prompt
+        negative_prompt = text_guided_negative_prompt
 
     if enable_control:
         return predict_controlnet(input_image, input_control_image,
@@ -373,8 +378,7 @@ def select_tab_shape_guided():
             inpaint_result = gr.Image()
             gr.Markdown('### Mask')
             gallery = gr.Gallery(
-                label='Generated images', show_label=False).style(
-                    grid=[2], height='auto')
+                label='Generated images', show_label=False, columns=2)
 
     run_button.click(
         fn=infer,
diff --git a/projects/powerpaint/requirements.txt b/projects/powerpaint/requirements.txt
index 14755c0ff..8a0cf9b0f 100644
--- a/projects/powerpaint/requirements.txt
+++ b/projects/powerpaint/requirements.txt
@@ -1,6 +1,6 @@
 controlnet-aux==0.0.3
 diffusers==0.23.1
-gradio==3.23.0
+gradio==3.41.0
 mmengine
 opencv-python
 torch