update

README.md

@@ -0,0 +1,44 @@
+Model from: https://huggingface.co/JunhaoZhuang/PowerPaint_v2
+
+Tokens (P_ctxt, P_shape, P_obj) added by PowerPaint has been integrated into the text_encoder and tokenizer.
+
+Unlike PowerPaint_v1, PowerPaint_v2 uses a method similar to BrushNet, so it can be applied to any sd1.5 type basic model.
+
+Clone demo code and models:
+
+```bash
+git lfs install
+git clone https://huggingface.co/Sanster/PowerPaint_v2
+```
+
+Run `main.py`:
+
+```bash
+python3 main.py runwayml/stable-diffusion-v1-5
+```
+
+The demo code will generate following results:
+
+| Original Image                                                                                                                                 | Mask                                                                                                                                            |
+| ---------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------- |
+| ![original_image](https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png) | ![mask_image](https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png) |
+
+**Object Removal Task**
+
+![object-removal](./object-removal_result.png)
+
+**Shape Guided Task**
+
+![shape-guided](./shape-guided_result.png)
+
+**Context aware Task**
+
+![context-aware](./context-aware_result.png)
+
+**Inpaint Task**
+
+![inpaint](./inpaint_result.png)
+
+**Outpaint Task**
+
+![outpaint](./image-outpainting_result.png)

main.py

@@ -1,11 +1,11 @@
+import sys
 import cv2
 import numpy as np
 import torch
-from PIL import Image, ImageFilter, ImageOps
+from PIL import Image, ImageOps
 from transformers import CLIPTextModel, CLIPTokenizer
 from diffusers.utils import load_image
 from diffusers import DPMSolverMultistepScheduler
-from safetensors.torch import load_model
 
 from powerpaint_v2.BrushNet_CA import BrushNetModel
 from powerpaint_v2.pipeline_PowerPaint_Brushnet_CA import (
@@ -90,24 +90,17 @@ def predict(
         height=W,
     ).images[0]
     return result
-    m_img = (
-        input_image["mask"].convert("RGB").filter(ImageFilter.GaussianBlur(radius=3))
-    )
-    m_img = np.asarray(m_img) / 255.0
-    img_np = np.asarray(input_image["image"].convert("RGB")) / 255.0
-    ours_np = np.asarray(result) / 255.0
-    ours_np = ours_np * m_img + (1 - m_img) * img_np
-    result_paste = Image.fromarray(np.uint8(ours_np * 255))
-    return result_paste
 
 
+# base_model_name = "runwayml/stable-diffusion-v1-5"
+base_model_name = sys.argv[1]
 text_encoder_brushnet = CLIPTextModel.from_pretrained(
     "text_encoder_brushnet",
     variant="fp16",
     torch_dtype=torch.float16,
 )
 unet = UNet2DConditionModel.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
+    base_model_name,
     subfolder="unet",
     variant="fp16",
     torch_dtype=torch.float16,
@@ -118,7 +111,7 @@ brushnet = BrushNetModel.from_pretrained(
     torch_dtype=torch.float16,
 )
 pipe = StableDiffusionPowerPaintBrushNetPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
+    base_model_name,
     torch_dtype=torch.float16,
     safety_checker=None,
     unet=unet,
@@ -170,7 +163,7 @@ tasks = [
     {
         "task": "image-outpainting",
         "guidance_scale": 7.5,
-        "prompt": "empty scene",
+        "prompt": "",
         "negative_prompt": negative_prompt,
     },
 ]