chore: Update .gitignore and add new files for inpainting pipeline

Former-commit-id: 98433cc2ee6b23538fc629e82c96c608a7ace21d [formerly e830430e6009854c8f7ef28751729b821abdb4b8]
Former-commit-id: f17fae1701fed57adad6c5861841805ef068dec4

.gitignore

@@ -4,4 +4,5 @@ variables.tf
 .terraform
 config.env
 /scripts/yolov8s*
-/scripts/*jpg
+/scripts/*jpg
+/scripts/outputs

configs/inpainting.yaml

@@ -0,0 +1,13 @@
+segmentation_model : 'facebook/sam-vit-large'
+detection_model : 'yolov8l'
+model : 'kandinsky-community/kandinsky-2-2-decoder-inpaint'
+target_width : 1920
+target_height : 1080
+prompt : 'product is on the kitchen floor , ultrarealistic lighting , commercial, award , winning photography'
+negative_prompt : 'low resolution , bad resolution'
+roi_scale : 0.5
+strength : 0.7
+guidance_scale : 7.5
+num_inference_steps : 600
+output_path : '../outputs'
+ 

scripts/config.py

@@ -7,7 +7,7 @@ PROJECT_NAME = "Product Photography"
 PRODUCTS_10k_DATASET = "VikramSingh178/Products-10k-BLIP-captions"
 CAPTIONING_MODEL_NAME = "Salesforce/blip-image-captioning-base"
 SEGMENTATION_MODEL_NAME = "facebook/sam-vit-large"
-DETECTION_MODEL_NAME = "yolov8s"
+DETECTION_MODEL_NAME = "yolov8l"
 
 
 

scripts/pipeline.py

@@ -0,0 +1,77 @@
+import torch
+from diffusers import AutoPipelineForInpainting
+from diffusers.utils import load_image
+from utils import (accelerator, ImageAugmentation, clear_memory)
+import hydra
+from omegaconf import OmegaConf, DictConfig
+from PIL import Image
+import lightning.pytorch as pl
+pl.seed_everything(42)
+generator = torch.Generator("cuda").manual_seed(92)
+
+class AutoPaintingPipeline:
+    """
+    AutoPaintingPipeline class represents a pipeline for auto painting using an inpainting model from diffusers.
+    
+    Args:
+        model_name (str): The name of the pretrained inpainting model.
+        image (Image): The input image to be processed.
+        mask_image (Image): The mask image indicating the areas to be inpainted.
+    """
+    
+    def __init__(self, model_name: str, image: Image, mask_image: Image):
+        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.pipeline.to(self.device)
+        self.pipeline.unet = torch.compile(self.pipeline.unet, mode="reduce-overhead", fullgraph=True)
+        
+        
+    def run_inference(self, prompt: str, negative_prompt: str, num_inference_steps: int, strength: float, guidance_scale: float):
+        """
+        Runs the inference on the input image using the inpainting pipeline.
+        
+        Returns:
+            Image: The output image after inpainting.
+        """
+        
+        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 = 1472, width = 2560).images[0]
+        clear_memory()
+        return output
+    
+    
+@hydra.main(version_base=None ,config_path="../configs", config_name="inpainting")
+def inference(cfg: DictConfig):
+    """
+    Load the configuration file for the inpainting pipeline.
+    
+    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)
+    model_name = cfg.model
+    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)
+    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)
+    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
+    
+if __name__ == "__main__":
+    inference()
+
+        
+
+  

scripts/utils.py

@@ -5,6 +5,25 @@ import numpy as np
 from PIL import Image, ImageOps
 from config import SEGMENTATION_MODEL_NAME, DETECTION_MODEL_NAME
 from diffusers.utils import load_image
+import gc
+
+
+
+
+
+
+
+def clear_memory():
+    """
+    Clears the memory by collecting garbage and emptying the CUDA cache.
+
+    This function is useful when dealing with memory-intensive operations in Python, especially when using libraries like PyTorch.
+
+   """
+    gc.collect()
+    torch.cuda.empty_cache()
+   
+
 
 
 
@@ -93,7 +112,7 @@ class ImageAugmentation:
 
 if __name__ == "__main__":
     augmenter = ImageAugmentation(target_width=2560, target_height=1440, roi_scale=0.7)
-    image_path = "/home/product_diffusion_api/sample_data/example3.jpg"
+    image_path = "../sample_data/example3.jpg"
     image = Image.open(image_path)
     extended_image = augmenter.extend_image(image)
     mask = augmenter.generate_mask_from_bbox(extended_image, SEGMENTATION_MODEL_NAME, DETECTION_MODEL_NAME)

scripts/yolov8l.pt.REMOVED.git-id

@@ -0,0 +1 @@
+4b11a02a7599e520f9f14a4703f4991237d6ce50