Update app.py

app.py

@@ -6,6 +6,20 @@ import os
 from PIL import Image
 hf_token = os.environ.get("HF_TOKEN")
 from diffusers import StableDiffusionXLInpaintPipeline, DDIMScheduler, UNet2DConditionModel
+from diffusers import (
+    AutoencoderKL,
+    LCMScheduler,
+)
+from pipeline_controlnet_sd_xl import StableDiffusionXLControlNetPipeline
+from controlnet import ControlNetModel, ControlNetConditioningEmbedding
+import torch
+import numpy as np
+from PIL import Image
+import requests
+import PIL
+from io import BytesIO
+from torchvision import transforms
+
 
 ratios_map =  {
     0.5:{"width":704,"height":1408},
@@ -28,6 +42,30 @@ ratios_map =  {
 }
 ratios = np.array(list(ratios_map.keys()))
 
+image_transforms = transforms.Compose(
+    [
+        transforms.ToTensor(),
+    ]
+)
+
+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"
+
+
+def get_masked_image(image, image_mask, width, height):
+    image_mask = image_mask # inpaint area is white
+    image_mask = image_mask.resize((width, height)) # object to remove is white (1)
+    image_mask_pil = image_mask
+    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
+
+    
 def get_size(init_image):
     w,h=init_image.size
     curr_ratio = w/h
@@ -40,26 +78,33 @@ def get_size(init_image):
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-unet = UNet2DConditionModel.from_pretrained(
-    "briaai/BRIA-2.2-Inpainting",
-    subfolder="unet",
-    torch_dtype=torch.float16,
-)
 
-scheduler = DDIMScheduler.from_pretrained("briaai/BRIA-2.3", subfolder="scheduler",clip_sample=False)
+# Load, init model
+controlnet = ControlNetModel().from_config('briaai/DEV-ControlNetInpaintingFast', torch_dtype=torch.float16)          
+controlnet.controlnet_cond_embedding = ControlNetConditioningEmbedding(
+    conditioning_embedding_channels=320,
+    conditioning_channels = 5
+)   
 
-pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
-    "briaai/BRIA-2.3",
-    unet=unet,
-    scheduler=scheduler,
-    torch_dtype=torch.float16,
-    force_zeros_for_empty_prompt=False
-)
+controlnet = ControlNetModel().from_pretrained("briaai/DEV-ControlNetInpaintingFast", torch_dtype=torch.float16)
+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()
 
-pipe = pipe.to(device)
-pipe.force_zeros_for_empty_prompt = False
+generator = torch.Generator(device='cuda:0').manual_seed(123456)
 
-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"
+vae = pipe.vae
+
+
+# pipe.force_zeros_for_empty_prompt = False
+
+# 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"
 
 
 def read_content(file_path: str) -> str:
@@ -70,26 +115,74 @@ def read_content(file_path: str) -> str:
 
     return content
 
-def predict(dict, prompt="", negative_prompt="", guidance_scale=5, steps=30, strength=1.0):
+def predict(dict, prompt="", negative_prompt = default_negative_prompt, guidance_scale=1.2, steps=12, strength=1.0):
     if negative_prompt == "":
         negative_prompt = None
 
     
     init_image = dict["image"].convert("RGB")#.resize((1024, 1024))
-    mask = dict["mask"].convert("RGB")#.resize((1024, 1024))
+    mask = dict["mask"].convert("L")#.resize((1024, 1024))
 
-    w,h = get_size(init_image)
+    width, height = get_size(init_image)
 
-    init_image = init_image.resize((w, h))
-    mask = mask.resize((w, h))
+    init_image = init_image.resize((width, height))
+    mask = mask.resize((width, height))
     
     # Resize to nearest ratio ?
     
-    mask = np.array(mask)
-    mask[mask>0]=255
-    mask = Image.fromarray(mask)
+    # mask = np.array(mask)
+    # mask[mask>0]=255
+    # mask = Image.fromarray(mask)
+    
+
+    masked_image, image_mask, masked_image_to_present = get_masked_image(init_image, mask, 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")
+    
+    control_latents = vae.encode(  
+            masked_image_tensor[:, :3, :, :].to(vae.dtype)
+        ).latent_dist.sample()
     
-    output = pipe(prompt = prompt,width=w,height=h, negative_prompt=negative_prompt, image=init_image, mask_image=mask, guidance_scale=guidance_scale, num_inference_steps=int(steps), strength=strength)
+    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)       
+    
+    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)
+
+
+    output = pipe(prompt = prompt,
+                  width=width,
+                  height=height,
+                  negative_prompt=negative_prompt,
+                  image = masked_image, # control image V
+                  init_image = init_image,
+                  mask_image=mask_tensor,
+                  guidance_scale=guidance_scale,
+                  num_inference_steps=int(steps),
+                  strength=strength,
+                  generator=generator,
+                  controlnet_conditioning_sale=1.0, )
+
+    # 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 = init_image,     
+    #         mask_image = mask_tensor,
+    #         guidance_scale = 1.2,
+    #         generator=generator).images[0]
+
     
     return output.images[0] #, gr.update(visible=True)