Update README.md

README.md

@@ -85,6 +85,21 @@ import requests
 import PIL
 from io import BytesIO
 from torchvision import transforms
+import pandas as pd 
+import os 
+
+
+def resize_image_to_retain_ratio(image):
+    pixel_number = 1024*1024
+    granularity_val = 8
+    ratio = image.size[0] / image.size[1]
+    width = int((pixel_number * ratio) ** 0.5)
+    width = width - (width % granularity_val)
+    height = int(pixel_number / width)
+    height = height - (height % granularity_val)
+
+    image = image.resize((width, height))
+    return image
 
 
 def download_image(url):
@@ -106,6 +121,7 @@ def get_masked_image(image, image_mask, width, height):
     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(),
@@ -120,7 +136,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")
+
+init_image = resize_image_to_retain_ratio(init_image)
+width, height = init_image.size
+
+mask_image = mask_image.convert("L").resize(init_image.size)
 
 width, height = init_image.size
 
@@ -132,18 +152,21 @@ pipe = StableDiffusionXLControlNetPipeline.from_pretrained("briaai/BRIA-2.3", co
 pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
 pipe.load_lora_weights("briaai/BRIA-2.3-FAST-LORA")
 pipe.fuse_lora()
+pipe = pipe.to(device="cuda")
 
-pipe = pipe.to('cuda:0')
-pipe.enable_xformers_memory_efficient_attention()
+# pipe.enable_xformers_memory_efficient_attention()
 
-generator = torch.Generator(device='cuda:0').manual_seed(123456)
+generator = torch.Generator(device="cuda").manual_seed(123456)
 
 vae = pipe.vae
 
+
 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
 
+
 masked_image_tensor = masked_image_tensor.unsqueeze(0).to(device="cuda")
 control_latents = vae.encode(  
         masked_image_tensor[:, :3, :, :].to(vae.dtype)
@@ -160,13 +183,13 @@ 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)
 
-prompt = "A park bench"
+prompt = ""
 
 gen_img = pipe(negative_prompt=default_negative_prompt, prompt=prompt, 
-            controlnet_conditioning_sale=1.0, 
+            controlnet_conditioning_scale=1.0, 
             num_inference_steps=12, 
             height=height, width=width, 
             image = masked_image, # control image
@@ -175,6 +198,5 @@ gen_img = pipe(negative_prompt=default_negative_prompt, prompt=prompt,
             guidance_scale = 1.2,
             generator=generator).images[0]
 
-
 ```