Update README.md

README.md

@@ -30,36 +30,41 @@ Unlike the inpaint controlnets used for general scenarios, this model is fine-tu
 ```python
 from diffusers import (
     ControlNetModel,
-    StableDiffusionXLControlNetInpaintPipeline
+    StableDiffusionXLControlNetPipeline,
+    DDPMScheduler
 )
 from diffusers.utils import load_image
 import torch
 from PIL import Image
+import numpy as np
 
 def make_inpaint_condition(init_image, mask_image):
     init_image = np.array(init_image.convert("RGB")).astype(np.float32) / 255.0
     mask_image = np.array(mask_image.convert("L")).astype(np.float32) / 255.0
-
     assert init_image.shape[0:1] == mask_image.shape[0:1], "image and image_mask must have the same image size"
     init_image[mask_image > 0.5] = -1.0  # set as masked pixel
     init_image = np.expand_dims(init_image, 0).transpose(0, 3, 1, 2)
     init_image = torch.from_numpy(init_image)
     return init_image
 
+def add_fg(full_img, fg_img, mask_img):
+    full_img = np.array(full_img).astype(np.float32)
+    fg_img = np.array(fg_img).astype(np.float32)
+    mask_img = np.array(mask_img).astype(np.float32) / 255.
+    full_img = full_img * mask_img + fg_img * (1-mask_img)
+    return Image.fromarray(np.clip(full_img, 0, 255).astype(np.uint8))
 
 controlnet = ControlNetModel.from_pretrained(
-    "alimama-creative/EcomXL_controlnet_inpaint", torch_dtype=torch.float16, use_safetensors=True
+    "alimama-creative/EcomXL_controlnet_inpaint",
+    use_safetensors=True,
 )
 
-pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(
+pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0", 
     controlnet=controlnet, 
-    torch_dtype=torch.float16
 )
-
-pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-# pipe.enable_xformers_memory_efficient_attention()
-pipe.enable_vae_slicing()
+pipe.to("cuda")
+pipe.scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
 
 image = load_image(
     "https://huggingface.co/alimama-creative/EcomXL_controlnet_inpaint/resolve/main/images/inp_0.png"
@@ -69,24 +74,25 @@ mask = load_image(
 )
 mask = Image.fromarray(255 - np.array(mask))
 
-control_image = make_inpaint_condition(img, mask)
+control_image = make_inpaint_condition(image, mask)
 
 prompt="a product on the table"
 
-images = pipe(
+generator = torch.Generator(device="cuda").manual_seed(1234)
+
+res_image = pipe(
     prompt,
-    image=img,
-    mask_image=mask,
-    control_image=control_image,
-    controlnet_conditioning_scale=0.5,
+    image=control_image,
+    num_inference_steps=25,
     guidance_scale=7,
-    strength=0.75,
     width=1024,
     height=1024,
+    controlnet_conditioning_scale=0.5,
+    generator=generator,
 ).images[0]
 
-image.save(f'test_inp.png')
-
+res_image = add_fg(res_image, image, mask)
+res_image.save(f'res.png')
 ```
 The model exhibits good performance when the controlnet weight (controllet_condition_scale) is 0.5.