Create README.md

README.md

@@ -0,0 +1,118 @@
+### How To Use
+
+```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
+
+
+def download_image(url):
+    response = requests.get(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):
+    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"
+    masked_image_to_present = image.copy()
+    masked_image_to_present[image_mask > 0.5] = (0.5,0.5,0.5)  # set as masked pixel
+    image[image_mask > 0.5] = 0.5  # set as masked pixel - s.t. will be grey 
+    image = Image.fromarray((image * 255.0).astype(np.uint8))
+    masked_image_to_present = Image.fromarray((masked_image_to_present * 255.0).astype(np.uint8))
+    return image, image_mask_pil, masked_image_to_present
+
+
+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"
+
+init_image = download_image(img_url).resize((1024, 1024))
+mask_image = download_image(mask_url).resize((1024, 1024))
+
+# Load, init model
+controlnet = ControlNetModel().from_config('/home/ubuntu/spring/Infra/project_x/bria2_controlnet_inpainting/config_controlnet_vae.json', torch_dtype=torch.float16)          
+controlnet.controlnet_cond_embedding = ControlNetConditioningEmbedding(
+    conditioning_embedding_channels=320,
+    conditioning_channels = 5
+)   
+    
+controlnet.load_state_dict(torch.load(local_ckpt_dir + local_ckpt_dir_suffix))
+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 = 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_tensor = image_transforms(masked_image)
+masked_image_tensor = (masked_image_tensor - 0.5) / 0.5
+
+masked_image_tensor = masked_image_tensor.unsqueeze(0).to(device="cuda")
+# masked_image_tensor = masked_image_tensor.permute((0,3,1,2))
+control_latents = vae.encode(  
+        masked_image_tensor[:, :3, :, :].to(vae.dtype)
+    ).latent_dist.sample()   
+control_latents = control_latents * vae.config.scaling_factor 
+
+
+image_mask = np.array(image_mask)[:,:]
+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.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, 
+            controlnet_conditioning_sale=1.0, 
+            num_inference_steps=12, 
+            height=height, width=width, 
+            image = masked_image, # control image
+            init_image = img,     
+            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")
+```
+