chore: Update inpainting pipeline configuration and parameters

Former-commit-id: e7d2ea3813700ed8994cf9ada9210b08b1d0d425 [formerly cebfaeaff312b0e8c563bb81b61865d4e8dd186e]
Former-commit-id: ad9d035eaffc7b7de46e960149e59ff2228b2174

configs/inpainting.yaml

@@ -1,13 +1,14 @@
-segmentation_model : 'facebook/sam-vit-huge'
+segmentation_model : 'facebook/sam-vit-base'
 detection_model : 'yolov8l'
 model : 'kandinsky-community/kandinsky-2-2-decoder-inpaint'
-target_width : 1920
-target_height : 1080
-prompt : 'product in the kitchen used in cooking 4k'
-negative_prompt : 'low resolution , bad resolution'
-roi_scale : 0.7
-strength : 0.2
-guidance_scale : 7.0
-num_inference_steps : 800
+target_width : 2560
+target_height : 1472
+prompt : 'Product on the table 4k ultrarealistic'
+negative_prompt : 'low resolution , bad resolution , Deformation , Weird Artifacts, bad quality,blown up image, high brightness , high saturation '
+roi_scale : 0.6
+strength : 0.5
+guidance_scale : 8
+num_inference_steps : 150
 output_path : '../outputs'
+
  

scripts/pipeline.py

@@ -5,7 +5,8 @@ from utils import (accelerator, ImageAugmentation, clear_memory)
 import hydra
 from omegaconf import DictConfig
 from PIL import Image
-
+import lightning.pytorch as pl
+#pl.seed_everything(1234)
 
 
 
@@ -19,13 +20,16 @@ class AutoPaintingPipeline:
         mask_image (Image): The mask image indicating the areas to be inpainted.
     """
     
-    def __init__(self, model_name: str, image: Image, mask_image: Image):
+    def __init__(self, model_name: str, image: Image, mask_image: Image,target_width: int, target_height: int):
         self.model_name = model_name
         self.device = accelerator()
         self.pipeline = AutoPipelineForInpainting.from_pretrained(self.model_name, torch_dtype=torch.float16)
         self.image = load_image(image)
         self.mask_image = load_image(mask_image)
+        self.target_width = target_width
+        self.target_height = target_height
         self.pipeline.to(self.device)
+        
        
         
         
@@ -39,8 +43,7 @@ class AutoPaintingPipeline:
         clear_memory()
         image = load_image(self.image)
         mask_image = load_image(self.mask_image)
-        output = self.pipeline(prompt=prompt,negative_prompt=negative_prompt,image=image,mask_image=mask_image,num_inference_steps=num_inference_steps,strength=strength,guidance_scale=guidance_scale,height = 1080, width = 1920).images[0]
-        
+        output = self.pipeline(prompt=prompt,negative_prompt=negative_prompt,image=image,mask_image=mask_image,num_inference_steps=num_inference_steps,strength=strength,guidance_scale=guidance_scale, height = self.target_height ,width = self.target_width).images[0]
         return output
     
     
@@ -52,22 +55,24 @@ def inference(cfg: DictConfig):
     Args:
         cfg (DictConfig): The configuration file for the inpainting pipeline.
     """
-    augmenter = ImageAugmentation(target_width=cfg.target_width, target_height=cfg.target_height, roi_scale=cfg.roi_scale)
+    augmenter = ImageAugmentation(target_width=cfg.target_width, target_height=cfg.target_height)
     model_name = cfg.model
-    image_path = "../sample_data/example1.jpg"
+    image_path = "../sample_data/example3.jpg"
     image = Image.open(image_path)
     extended_image = augmenter.extend_image(image)
     mask_image = augmenter.generate_mask_from_bbox(extended_image, cfg.segmentation_model, cfg.detection_model)
     mask_image = augmenter.invert_mask(mask_image)
+    #mask_image = cv2.GaussianBlur(np.array(mask_image), (7,7), 0)
     prompt = cfg.prompt
     negative_prompt = cfg.negative_prompt
     num_inference_steps = cfg.num_inference_steps
     strength = cfg.strength
     guidance_scale = cfg.guidance_scale
-    pipeline = AutoPaintingPipeline(model_name=model_name, image=extended_image, mask_image=mask_image)
+    pipeline = AutoPaintingPipeline(model_name=model_name, image = extended_image, mask_image=mask_image, target_height=cfg.target_height, target_width=cfg.target_width)
     output = pipeline.run_inference(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, strength=strength, guidance_scale=guidance_scale)
     output.save(f'{cfg.output_path}/output.jpg')
-    return output
+    mask_image.save(f'{cfg.output_path}/mask.jpg')
+    
     
 if __name__ == "__main__":
     inference()

scripts/utils.py

@@ -111,7 +111,7 @@ class ImageAugmentation:
         return inverted_mask_pil
 
 if __name__ == "__main__":
-    augmenter = ImageAugmentation(target_width=2560, target_height=1440, roi_scale=0.7)
+    augmenter = ImageAugmentation(target_width=1024, target_height=1024, roi_scale=0.5)
     image_path = "../sample_data/example3.jpg"
     image = Image.open(image_path)
     extended_image = augmenter.extend_image(image)