commit

Former-commit-id: 74a3d876eb8af5d012d025eb65b2dd9dc3646558

logs/app_debug.log

@@ -1389,3 +1389,78 @@ speed: {'preprocess': 1.9655227661132812, 'inference': 86.20810508728027, 'postp
 2024-03-22 09:47:05,180 [INFO] models - Kandinsky Inpainting Inference
 2024-03-22 09:51:28,523 [INFO] models - Kandinsky Inpainting Inference
 2024-03-22 09:53:18,039 [INFO] models - Kandinsky Inpainting Inference
+2024-03-23 08:26:37,691 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 08:29:23,557 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 08:30:47,545 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 08:35:34,516 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:31:44,362 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:32:16,521 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:32:45,559 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:35:33,253 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:36:06,844 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:43:52,587 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:44:30,749 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:56:41,476 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:57:16,318 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 10:07:32,135 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 11:02:29,843 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 11:05:02,471 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 19:23:38,952 [INFO] clear_memory - Memory Cleared
+2024-03-23 19:42:39,347 [INFO] clear_memory - Memory Cleared
+2024-03-24 12:51:21,021 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:07:46,765 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:14:03,557 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:17:25,009 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:23:11,117 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:28:30,138 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:32:00,626 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:35:18,404 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:38:37,096 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:41:27,631 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:44:00,439 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:47:29,728 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:51:27,224 [INFO] clear_memory - Memory Cleared
+2024-03-24 14:29:41,821 [INFO] clear_memory - Memory Cleared
+2024-03-24 14:32:41,074 [INFO] clear_memory - Memory Cleared
+2024-03-24 14:40:54,554 [INFO] clear_memory - Memory Cleared
+2024-03-24 14:43:46,918 [INFO] clear_memory - Memory Cleared
+2024-03-24 14:46:15,513 [INFO] clear_memory - Memory Cleared
+2024-03-24 14:49:51,695 [INFO] clear_memory - Memory Cleared
+2024-03-24 15:21:03,538 [INFO] clear_memory - Memory Cleared
+2024-03-24 15:24:31,062 [INFO] clear_memory - Memory Cleared
+2024-03-24 18:15:16,198 [INFO] clear_memory - Memory Cleared
+2024-03-24 18:30:54,085 [INFO] clear_memory - Memory Cleared
+2024-03-24 18:31:39,719 [INFO] clear_memory - Memory Cleared
+2024-03-24 18:56:00,754 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:01:30,572 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:05:33,502 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:08:09,191 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:10:09,003 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:10:54,562 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:15:11,151 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:20:54,871 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:23:27,620 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:26:56,507 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:39:00,070 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:44:51,106 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:52:32,346 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:58:14,409 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:05:13,208 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:11:07,106 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:18:40,535 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:24:46,394 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:30:39,800 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:37:12,153 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:42:47,710 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:48:23,383 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:56:31,162 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:03:44,503 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:09:56,651 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:17:23,320 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:23:06,580 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:29:14,870 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:36:09,328 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:40:40,507 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:44:47,907 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:48:43,724 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:52:50,583 [INFO] clear_memory - Memory Cleared

logs/app_info.log

@@ -1389,3 +1389,78 @@ speed: {'preprocess': 1.9655227661132812, 'inference': 86.20810508728027, 'postp
 2024-03-22 09:47:05,180 [INFO] models - Kandinsky Inpainting Inference
 2024-03-22 09:51:28,523 [INFO] models - Kandinsky Inpainting Inference
 2024-03-22 09:53:18,039 [INFO] models - Kandinsky Inpainting Inference
+2024-03-23 08:26:37,691 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 08:29:23,557 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 08:30:47,545 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 08:35:34,516 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:31:44,362 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:32:16,521 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:32:45,559 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:35:33,253 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:36:06,844 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:43:52,587 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:44:30,749 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:56:41,476 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 09:57:16,318 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 10:07:32,135 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 11:02:29,843 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 11:05:02,471 [INFO] mask_generator - Mask generation completed successfully
+2024-03-23 19:23:38,952 [INFO] clear_memory - Memory Cleared
+2024-03-23 19:42:39,347 [INFO] clear_memory - Memory Cleared
+2024-03-24 12:51:21,021 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:07:46,765 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:14:03,557 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:17:25,009 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:23:11,117 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:28:30,138 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:32:00,626 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:35:18,404 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:38:37,096 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:41:27,631 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:44:00,439 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:47:29,728 [INFO] clear_memory - Memory Cleared
+2024-03-24 13:51:27,224 [INFO] clear_memory - Memory Cleared
+2024-03-24 14:29:41,821 [INFO] clear_memory - Memory Cleared
+2024-03-24 14:32:41,074 [INFO] clear_memory - Memory Cleared
+2024-03-24 14:40:54,554 [INFO] clear_memory - Memory Cleared
+2024-03-24 14:43:46,918 [INFO] clear_memory - Memory Cleared
+2024-03-24 14:46:15,513 [INFO] clear_memory - Memory Cleared
+2024-03-24 14:49:51,695 [INFO] clear_memory - Memory Cleared
+2024-03-24 15:21:03,538 [INFO] clear_memory - Memory Cleared
+2024-03-24 15:24:31,062 [INFO] clear_memory - Memory Cleared
+2024-03-24 18:15:16,198 [INFO] clear_memory - Memory Cleared
+2024-03-24 18:30:54,085 [INFO] clear_memory - Memory Cleared
+2024-03-24 18:31:39,719 [INFO] clear_memory - Memory Cleared
+2024-03-24 18:56:00,754 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:01:30,572 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:05:33,502 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:08:09,191 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:10:09,003 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:10:54,562 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:15:11,151 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:20:54,871 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:23:27,620 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:26:56,507 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:39:00,070 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:44:51,106 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:52:32,346 [INFO] clear_memory - Memory Cleared
+2024-03-24 19:58:14,409 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:05:13,208 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:11:07,106 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:18:40,535 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:24:46,394 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:30:39,800 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:37:12,153 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:42:47,710 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:48:23,383 [INFO] clear_memory - Memory Cleared
+2024-03-24 20:56:31,162 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:03:44,503 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:09:56,651 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:17:23,320 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:23:06,580 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:29:14,870 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:36:09,328 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:40:40,507 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:44:47,907 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:48:43,724 [INFO] clear_memory - Memory Cleared
+2024-03-24 21:52:50,583 [INFO] clear_memory - Memory Cleared

scripts/config.py

@@ -8,6 +8,4 @@ mask_dir = '../masks'
 controlnet_adapter_model_name= 'lllyasviel/control_v11p_sd15_inpaint'
 controlnet_base_model_name = "runwayml/stable-diffusion-inpainting"
 kandinsky_model_name = 'kandinsky-community/kandinsky-2-2-decoder-inpaint'
-width = 512
-height = 512
-yolo_model = 'yolov8s-seg.pt'
+

scripts/datautils.py

@@ -1,58 +0,0 @@
-from datasets import load_dataset, Image
-from config import Dataset_Name, DATA_DIR
-from logger import rich_logger as l
-import wandb
-from config import Project_Name, entity
-import pandas as pd
-from tqdm import tqdm
-
-
-
-
-
-
-
-
-class DatasetUtils:
-    """
-    Utility class for working with datasets.
-    """
-    def __init__(self, dataset_name:str,split:str=None):
-        super().__init__()
-        """
-        Initializes a new instance of the DatasetUtils class.
-
-        Args:
-            dataset_name (str): The name of the dataset to use.
-        """
-        self.dataset_name = dataset_name
-        self.split = split
-        self.dataset = load_dataset(self.dataset_name,cache_dir=DATA_DIR,split=split)
-       
-        self.dataset=self.dataset.remove_columns(['id'])
-        l.info(f"Initialized dataset: {self.dataset_name}")
-        l.info(self.dataset.features)
-        
-
-        
-    
-    
-
-    
-
-
-    
-        
-       
-      
-        
-        
-    
-    
-if __name__=="__main__":
-    dataset = DatasetUtils(Dataset_Name,split="train")
-    
-    
-   
-       
-    

scripts/mask_generator.py

@@ -1,35 +1,14 @@
-from PIL import Image
+from PIL import Image, ImageFilter,ImageDraw
 import numpy as np
 from logger import rich_logger as l
 from ultralytics import YOLO
 import cv2
-from config import yolo_model
-from pathlib import Path
 import PIL.ImageOps
 
 
 
 
-def generate_mask(image_path: str) -> np.ndarray:
-    """Method to segment image
-    Args:
-        image_path (str): path to input image
-    Returns:
-        np.ndarray: segmented image mask
-    """
-    model = YOLO(model=yolo_model)  # Initialize YOLO model
-    results = model(image_path)  # Perform object detection
-    for result in results:
-        orig_img = result.orig_img
-        masks = result.masks.xy
-        height, width = result.orig_img.shape[:2]
-        mask_img = np.ones((height, width), dtype=np.uint8) * 255  # Initialize mask with white background
-        
-        for mask in masks:
-            mask = mask.astype(int)
-            cv2.fillPoly(mask_img, [mask], 0)  # Fill mask with detected object areas
-
-    return mask_img
+
 
 def invert_mask(mask_image: Image) -> np.ndarray:
     """Method to invert mask
@@ -41,3 +20,42 @@ def invert_mask(mask_image: Image) -> np.ndarray:
     inverted_mask_image =PIL.ImageOps.invert(mask_image)
     return inverted_mask_image
 
+
+
+
+
+def extend_image(image_path, target_width, target_height, roi_scale=0.5):
+    # Open the original image
+    original_image = Image.open(image_path)
+
+    # Get the dimensions of the original image
+    original_width, original_height = original_image.size
+
+    # Calculate the scale to fit the target resolution while keeping the aspect ratio
+    scale = min(target_width / original_width, target_height / original_height)
+
+    # Calculate the new dimensions of the image
+    new_width = int(original_width * scale * roi_scale)
+    new_height = int(original_height * scale * roi_scale)
+
+    # Resize the original image with keeping the aspect ratio
+    original_image_resized = original_image.resize((new_width, new_height))
+
+    # Create a new image with white background
+    extended_image = Image.new("RGB", (target_width, target_height), "white")
+
+    # Calculate the position to paste the resized image at the center
+    paste_x = (target_width - new_width) // 2
+    paste_y = (target_height - new_height) // 2
+
+    # Paste the resized image onto the new image
+    extended_image.paste(original_image_resized, (paste_x, paste_y))
+
+    return extended_image
+
+
+
+
+
+
+

scripts/models.py

@@ -2,17 +2,15 @@ from logger import rich_logger as l
 from wandb.integration.diffusers import autolog
 from config import Project_Name
 from clear_memory import clear_memory
-from typing import List
 import numpy as np
 import torch
-from PIL import Image,ImageFilter,ImageOps
+from PIL import Image
 from mask_generator import invert_mask
 from diffusers.utils import load_image
-from pipeline import fetch_control_pipeline,fetch_kandinsky_pipeline,fetch_kandinsky_prior_pipeline,fetch_kandinsky_img2img_pipeline
+from pipeline import fetch_kandinsky_pipeline
 from config import controlnet_adapter_model_name,controlnet_base_model_name,kandinsky_model_name
 import cv2 
-
-from transformers import pipeline
+from diffusers import StableDiffusionInpaintPipeline, DPMSolverMultistepScheduler
 
 
 
@@ -23,39 +21,8 @@ autolog(init=dict(project=Project_Name))
 
 
 
-def make_controlnet_condition(image: Image.Image) -> Image.Image:
-    """
-    Applies image processing operations to create a controlnet condition image.
-
-    Args:
-        image (PIL.Image.Image): The input image.
-
-    Returns:
-        PIL.Image.Image: The controlnet condition image.
-    """
-    image = np.array(image)
-    image = cv2.Canny(image, 100, 200)
-    image = image[:, :, None]
-    image = np.concatenate([image, image, image], axis=2)
-    image = Image.fromarray(image)
-    return image
 
 def make_inpaint_condition(init_image, mask_image):
-    """
-    Prepare the initial image for inpainting by applying a mask.
-
-    Args:
-        init_image (PIL.Image.Image): The initial image.
-        mask_image (PIL.Image.Image): The mask image.
-
-    Returns:
-        torch.Tensor: The prepared initial image for inpainting.
-
-    Raises:
-        AssertionError: If the image and mask have different sizes.
-
-    """
-    # Prepare control image
     init_image = np.array(init_image.convert("RGB")).astype(np.float32) / 255.0
     mask_image = np.array(mask_image.convert("L")).astype(np.float32) / 255.0
 
@@ -66,15 +33,6 @@ def make_inpaint_condition(init_image, mask_image):
     return init_image
     
 
-def make_hint(image, depth_estimator):
-    image = depth_estimator(image)["depth"]
-    image = np.array(image)
-    image = image[:, :, None]
-    image = np.concatenate([image, image, image], axis=2)
-    detected_map = torch.from_numpy(image).float() / 255.0
-    hint = detected_map.permute(2, 0, 1)
-    return hint
-
 
     
 
@@ -82,36 +40,7 @@ def make_hint(image, depth_estimator):
 
 
 
-def controlnet_inpainting_inference(prompt,
-                         image,
-                         mask_image,
-                         control_image,
-                         num_inference_steps=200,
-                         guidance_scale=1.2,
-                         strength=5.0,
-                         generator=torch.Generator(device="cpu").manual_seed(1)
-                        ) -> List[Image.Image]:
-    """
-    Perform inpainting inference on an image using the given parameters.
 
-    Args:
-        prompt: The prompt for the inpainting inference.
-        image: The input image to be inpainted.
-        mask_image: The mask image indicating the regions to be inpainted.
-        controlnet_conditioning_image: The conditioning image for the controlnet.
-        num_inference_steps: The number of inference steps to perform (default: 200).
-        guidance_scale: The scale factor for the guidance loss (default: 1.2).
-        strength: The strength of the inpainting (default: 5.0).
-        generator: The random number generator for reproducibility (default: torch.Generator(device="cpu").manual_seed(1)).
-
-    Returns:
-        A list of inpainted images.
-
-    """
-    clear_memory()
-    pipe = fetch_control_pipeline(controlnet_adapter_model_name, controlnet_base_model_name,kandinsky_model_name, control_image)
-    image = pipe(prompt = prompt,num_inference_steps=num_inference_steps, generator=generator, eta=1.0, image=image, mask_image=mask_image,guidance_scale=guidance_scale,strenght=strength, control_image=control_image).images[0]
-    return image
 
 def kandinsky_inpainting_inference(prompt, negative_prompt, image, mask_image):
     """
@@ -126,37 +55,40 @@ def kandinsky_inpainting_inference(prompt, negative_prompt, image, mask_image):
     Returns:
         PIL.Image.Image: The output inpainted image.
     """
+    clear_memory()
     pipe = fetch_kandinsky_pipeline(controlnet_adapter_model_name, controlnet_base_model_name,kandinsky_model_name, image)
-    output_image = pipe(prompt=prompt,negative_prompt=negative_prompt,image=image,mask_image=mask_image,num_inference_steps=200,strength=1.0).images[0]
+    output_image = pipe(prompt=prompt,negative_prompt=negative_prompt,image=image,mask_image=mask_image,num_inference_steps=800,strength=1.0,guidance_scale = 7.8,height = 1472, width = 2560).images[0]
     return output_image
 
     
 
-def kandinsky_controlnet_inpainting_inference(prompt, negative_prompt, image, hint, generator=torch.Generator(device="cuda").manual_seed(43)):
+
+  
+
+def sd2_inpainting_inference(prompt, img, mask, repo_id="stabilityai/stable-diffusion-2-inpainting", revision="fp16"):
     """
-    Perform inpainting inference using the Kandinsky ControlNet model.
+    Generate an image based on a prompt using a pretrained model.
 
     Args:
-        prompt (str): The prompt for the inpainting process.
-        negative_prompt (str): The negative prompt for the inpainting process.
-        image (torch.Tensor): The input image for inpainting.
-        hint (torch.Tensor): The hint for guiding the inpainting process.
-        generator (torch.Generator, optional): The random number generator. Defaults to CUDA generator with seed 43.
+        prompt (str): The prompt for the image generation.
+        img_url (str): The URL of the initial image.
+        mask_url (str): The URL of the mask image.
+        repo_id (str, optional): The ID of the repository of the pretrained model. Defaults to "stabilityai/stable-diffusion-2-inpainting".
+        revision (str, optional): The revision of the pretrained model. Defaults to "fp16".
 
     Returns:
-        torch.Tensor: The inpainted image.
-
+        Image: The generated image.
     """
-    prior_pipe = fetch_kandinsky_prior_pipeline(controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name, image)
-    img_embed = prior_pipe(prompt=prompt, image=image, strength=1.0, generator=generator)
-    negative_embed = prior_pipe(prompt=negative_prompt, image=image, strength=1, generator=generator)
-    controlnet_pipe = fetch_kandinsky_img2img_pipeline(controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name, image)
-    image = controlnet_pipe(image=image, strength=1.0, image_embeds=img_embed.image_embeds, negative_image_embeds=negative_embed.image_embeds, hint=hint, num_inference_steps=200, generator=generator, height=768, width=768).images[0]
+    init_image = load_image(img)
+    mask_image = load_image(mask)
+    pipe = StableDiffusionInpaintPipeline.from_pretrained(
+    repo_id,
+    torch_dtype=torch.float16)
+    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+    pipe = pipe.to("cuda")
+    image = pipe(prompt=prompt, image=init_image, mask_image=mask_image, num_inference_steps=400,guidence_scale=7.5).images[0]
     return image
 
-  
-
-
     
     
     

scripts/pipeline.py

@@ -1,15 +1,11 @@
-from diffusers import ControlNetModel,StableDiffusionControlNetInpaintPipeline,AutoPipelineForInpainting,KandinskyV22ControlnetImg2ImgPipeline,KandinskyV22PriorEmb2EmbPipeline
-from diffusers.utils import load_image
-import torch
-from PIL import Image
-import numpy as np
-import cv2
+from diffusers import ControlNetModel,StableDiffusionControlNetInpaintPipeline,AutoPipelineForInpainting
 import torch
 
 
 
 
 
+
 class PipelineFetcher:
     """
     A class that fetches different pipelines for image processing.
@@ -56,52 +52,28 @@ class PipelineFetcher:
         pipe.to('cuda')
         return pipe
 
-    def KandinskyPriorPipeline(self):
-        """
-        Fetches the Kandinsky prior pipeline.
-
-        Returns:
-            prior_pipeline (KandinskyV22PriorEmb2EmbPipeline): The Kandinsky prior pipeline.
-
-        """
-        prior_pipeline = KandinskyV22PriorEmb2EmbPipeline.from_pretrained(
-            "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16, use_safetensors=False
-        ).to("cuda")
-        return prior_pipeline
 
-    def KandinskyImg2ImgPipeline(self):
-        """
-        Fetches the Kandinsky img2img pipeline.
+def fetch_control_pipeline(controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name, image):
+    """
+    Fetches the control pipeline for image processing.
 
-        Returns:
-            img2img_pipeline (KandinskyV22ControlnetImg2ImgPipeline): The Kandinsky img2img pipeline.
+    Args:
+        controlnet_adapter_model_name (str): The name of the controlnet adapter model.
+        controlnet_base_model_name (str): The name of the controlnet base model.
+        kandinsky_model_name (str): The name of the Kandinsky model.
+        image: The input image for processing.
 
-        """
-        img2img_pipeline = KandinskyV22ControlnetImg2ImgPipeline.from_pretrained(
-            "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.float16, use_safetensors=False
-        ).to("cuda")
-        return img2img_pipeline
-    
-   
-
-def fetch_control_pipeline(controlnet_adapter_model_name,controlnet_base_model_name,kandinsky_model_name,image):
-    pipe_fetcher = PipelineFetcher(controlnet_adapter_model_name,controlnet_base_model_name,kandinsky_model_name,image)
+    Returns:
+        pipe: The control pipeline for image processing.
+    """
+    pipe_fetcher = PipelineFetcher(controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name, image)
     pipe = pipe_fetcher.ControlNetInpaintPipeline()
     return pipe
 
-def fetch_kandinsky_pipeline(controlnet_adapter_model_name,controlnet_base_model_name,kandinsky_model_name,image):
-    pipe_fetcher = PipelineFetcher(controlnet_adapter_model_name,controlnet_base_model_name,kandinsky_model_name,image)
-    pipe = pipe_fetcher.KandinskyPipeline()
-    return pipe
 
-def fetch_kandinsky_prior_pipeline(controlnet_adapter_model_name,controlnet_base_model_name,kandinsky_model_name,image):
-    pipe_fetcher = PipelineFetcher(controlnet_adapter_model_name,controlnet_base_model_name,kandinsky_model_name,image)
-    pipe = pipe_fetcher.KandinskyPriorPipeline()
-    return pipe
-
-def fetch_kandinsky_img2img_pipeline(controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name, image):
+def fetch_kandinsky_pipeline(controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name, image):
     """
-    Fetches the Kandinsky image-to-image pipeline.
+    Fetches the Kandinsky pipeline.
 
     Args:
         controlnet_adapter_model_name (str): The name of the controlnet adapter model.
@@ -110,12 +82,10 @@ def fetch_kandinsky_img2img_pipeline(controlnet_adapter_model_name, controlnet_b
         image: The input image.
 
     Returns:
-        pipe: The Kandinsky image-to-image pipeline.
+        pipe: The Kandinsky pipeline.
     """
     pipe_fetcher = PipelineFetcher(controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name, image)
-    pipe = pipe_fetcher.KandinskyImg2ImgPipeline()
-    return pipe
-def fetch_kandinsky_img2img_pipeline(controlnet_adapter_model_name,controlnet_base_model_name,kandinsky_model_name,image):
-    pipe_fetcher = PipelineFetcher(controlnet_adapter_model_name,controlnet_base_model_name,kandinsky_model_name,image)
-    pipe = pipe_fetcher.KandinskyImg2ImgPipeline()
+    pipe = pipe_fetcher.KandinskyPipeline()
     return pipe
+
+

scripts/run.py

@@ -1,37 +1,49 @@
 import argparse
 import os
-from mask_generator import generate_mask, invert_mask
+from mask_generator import  invert_mask
+from segment_everything import generate_mask_from_bbox
 from models import kandinsky_inpainting_inference, load_image
 from PIL import Image
+from pipeline import fetch_kandinsky_pipeline
+from config import controlnet_adapter_model_name, controlnet_base_model_name, kandinsky_model_name
+from mask_generator import extend_image
+
 
 def main(args):
-    # Generate mask
-    mask = generate_mask(args.image_path)
+    # Create directories if they don't exist
+    os.makedirs(args.output_dir, exist_ok=True)
+    os.makedirs(args.mask_dir, exist_ok=True)
+
+    # Extend image
+    extended_image = extend_image(image_path=args.image_path, target_width=2560, target_height=1440, roi_scale=0.6)
+    extended_image_path = os.path.join(args.output_dir, f'{args.uid}_extended.jpg')
+    extended_image.save(extended_image_path)
+
+    # Generate mask from extended image
+    mask = generate_mask_from_bbox(extended_image_path)
     mask_image = Image.fromarray(mask)
 
     # Save original mask
     original_mask_path = os.path.join(args.mask_dir, f'{args.uid}_original_mask.jpg')
     mask_image.save(original_mask_path)
-    
-    
+
     # Invert mask
     mask_image = load_image(original_mask_path)
     inverted_mask = invert_mask(mask_image)
     inverted_mask_path = os.path.join(args.mask_dir, f'{args.uid}_inverted_mask.jpg')
     inverted_mask.save(inverted_mask_path)
 
-    # Load mask and image
-    invert_mask_image = load_image(inverted_mask_path)
-    image = load_image(args.image_path)
+    # Load mask and extended image
+    mask_image = load_image(inverted_mask_path)
+    extended_image = load_image(extended_image_path)
 
-    # Perform inpainting
-    output_image = kandinsky_inpainting_inference(args.prompt, args.negative_prompt, image, mask_image)
+    # Perform inpainting on extended image
+    output_image = kandinsky_inpainting_inference(args.prompt, args.negative_prompt, extended_image, mask_image)
 
     # Save output image
     output_image_path = os.path.join(args.output_dir, f'{args.uid}_output.jpg')
     output_image.save(output_image_path)
     
-    
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='Perform Kandinsky inpainting on an image.')
     parser.add_argument('--image_path', type=str, required=True, help='Path to the input image.')

scripts/segment_everything.py

@@ -0,0 +1,51 @@
+from transformers import pipeline
+from ultralytics import YOLO
+from transformers import SamModel, SamProcessor
+import torch
+from PIL import Image
+from diffusers.utils import load_image
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+
+
+
+
+def generate_mask_from_bbox(image_path):
+    """
+    Generates a mask from the bounding box of an image using YOLO and SAM-ViT models.
+
+    Args:
+        image_path (str): The path to the input image.
+
+    Returns:
+        numpy.ndarray: The generated mask as a NumPy array.
+    """
+    # Initialize YOLO and Processor
+    yolo = YOLO('yolov8m.pt')
+    processor = SamProcessor.from_pretrained('facebook/sam-vit-large')
+    model = SamModel.from_pretrained("facebook/sam-vit-large").to(device)
+
+    # Generate bounding boxes
+    results = yolo(image_path)
+    bboxes = results[0].boxes.xyxy.tolist()
+    input_boxes = [[[bboxes[0]]]]
+
+    # Process inputs
+    inputs = processor(load_image(image_path), input_boxes=input_boxes, return_tensors="pt").to("cuda")
+
+    # Perform inference
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    # Post-process masks
+    mask = processor.image_processor.post_process_masks(
+        outputs.pred_masks.cpu(),
+        inputs["original_sizes"].cpu(),
+        inputs["reshaped_input_sizes"].cpu()
+    )[0][0][0].numpy()
+    print(mask)
+    return mask
+
+

scripts/yolov8m.pt.REMOVED.git-id

@@ -0,0 +1 @@
+cdb22b922d8e59c557a944f4560a5358770baad5