Update README.md

README.md

@@ -74,18 +74,17 @@ By submitting the form above, you agree to BRIA’s [Privacy policy](https://bri
 ```python
 from diffusers import (
     AutoencoderKL,
-    StableDiffusionXLControlNetInpaintPipeline,
     LCMScheduler,
 )
 from pipeline_controlnet_sd_xl import StableDiffusionXLControlNetPipeline
 from controlnet import ControlNetModel, ControlNetConditioningEmbedding
-import os
-from torchvision import transforms
 import torch
-from tqdm import tqdm
 import numpy as np
-import pandas as pd
 from PIL import Image
+import requests
+import PIL
+from io import BytesIO
+from torchvision import transforms
 
 
 def download_image(url):
@@ -93,12 +92,10 @@ def download_image(url):
     return PIL.Image.open(BytesIO(response.content)).convert("RGB")
 
 
-def get_masked_image(path_to_images_dir, image_name, image, image_mask, width, height):
+def get_masked_image(image, image_mask, width, height):
     image_mask = image_mask # inpaint area is white
-    image_mask = add_margins_to_ratio(image_mask, 1.5)
     image_mask = image_mask.resize((width, height)) # object to remove is white (1)
     image_mask_pil = image_mask
-    orig_image = np.array(image.convert("RGB")).astype(np.uint8)
     image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
     image_mask = np.array(image_mask_pil.convert("L")).astype(np.float32) / 255.0
     assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
@@ -109,6 +106,14 @@ def get_masked_image(path_to_images_dir, image_name, image, image_mask, width, h
     masked_image_to_present = Image.fromarray((masked_image_to_present * 255.0).astype(np.uint8))
     return image, image_mask_pil, masked_image_to_present
 
+image_transforms = transforms.Compose(
+    [
+        transforms.ToTensor(),
+        # transforms.Normalize([0.5], [0.5]),
+    ]
+)
+
+default_negative_prompt = "Logo,Watermark,Text,Ugly,Morbid,Extra fingers,Poorly drawn hands,Mutation,Blurry,Extra limbs,Gross proportions,Missing arms,Mutated hands,Long neck,Duplicate,Mutilated,Mutilated hands,Poorly drawn face,Deformed,Bad anatomy,Cloned face,Malformed limbs,Missing legs,Too many fingers"
 
 img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
 mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
@@ -116,6 +121,11 @@ mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data
 init_image = download_image(img_url).resize((1024, 1024))
 mask_image = download_image(mask_url).resize((1024, 1024))
 
+mask_image = mask_image.convert("L")
+
+width, height = init_image.size
+
+
 # Load, init model
 controlnet = ControlNetModel().from_config('briaai/DEV-ControlNetInpaintingFast', torch_dtype=torch.float16)          
 controlnet.controlnet_cond_embedding = ControlNetConditioningEmbedding(
@@ -123,25 +133,22 @@ controlnet.controlnet_cond_embedding = ControlNetConditioningEmbedding(
     conditioning_channels = 5
 )   
     
-controlnet.load_state_dict(torch.load(local_ckpt_dir + local_ckpt_dir_suffix))
+# controlnet.load_state_dict(torch.load('briaai/DEV-ControlNetInpaintingFast'))
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 pipe = StableDiffusionXLControlNetPipeline.from_pretrained("briaai/BRIA-2.3", controlnet=controlnet.to(dtype=torch.float16), torch_dtype=torch.float16, vae=vae) #force_zeros_for_empty_prompt=False, # vae=vae)
 
 pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-        pipe.load_lora_weights("briaai/BRIA-2.3-FAST-LORA")
-        pipe.fuse_lora()
+pipe.load_lora_weights("briaai/BRIA-2.3-FAST-LORA")
+pipe.fuse_lora()
 
 pipe = pipe.to('cuda:0')
 pipe.enable_xformers_memory_efficient_attention()
 
 generator = torch.Generator(device='cuda:0').manual_seed(123456)
 
-
-
-
 vae = pipe.vae
 
-masked_image, image_mask, masked_image_to_present = get_masked_image(path_to_images_dir, image_name, image_mask, img, width, height)
+masked_image, image_mask, masked_image_to_present = get_masked_image(init_image, mask_image, width, height)
 masked_image_tensor = image_transforms(masked_image)
 masked_image_tensor = (masked_image_tensor - 0.5) / 0.5
 
@@ -158,32 +165,26 @@ mask_tensor = torch.tensor(image_mask, dtype=torch.float32)[None, ...]
 # binarize the mask
 mask_tensor = torch.where(mask_tensor > 128.0, 255.0, 0)       
 
-if normalize_mask_to_0_1:
-    mask_tensor = mask_tensor / 255.0
+mask_tensor = mask_tensor / 255.0
 
 mask_tensor = mask_tensor.to(device="cuda")
 mask_resized = torch.nn.functional.interpolate(mask_tensor[None, ...], size=(control_latents.shape[2], control_latents.shape[3]), mode='nearest')
 # mask_resized = mask_resized.to(torch.float16)
 masked_image = torch.cat([control_latents, mask_resized], dim=1)
 
-gen_img = pipe(negative_prompt=default_negative_prompt, prompt=caption, 
+prompt = "A park bench"
+
+gen_img = pipe(negative_prompt=default_negative_prompt, prompt=prompt, 
             controlnet_conditioning_sale=1.0, 
             num_inference_steps=12, 
             height=height, width=width, 
             image = masked_image, # control image
-            init_image = img,     
+            init_image = init_image,     
             mask_image = mask_tensor,
             guidance_scale = 1.2,
             generator=generator).images[0]
 
 
-
-
-
-
-prompt = "A park bench"
-generator = torch.Generator(device='cuda:0').manual_seed(123456)
-image = pipe(prompt=prompt, image=init_image, mask_image=mask_image,generator=generator,guidance_scale=5,strength=1).images[0]
-image.save("./a_park_bench.png")
+gen_img.save("./a_park_bench.png")
 ```