| import sys | |
| import torch | |
| sys.path.append('../') | |
| from Kandinsky.kandinsky3 import get_inpainting_pipeline | |
| from api_utils import ImageAugmentation | |
| from diffusers.utils import load_image | |
| from PIL import Image | |
| device_map = torch.device('cuda:0') | |
| dtype_map = { | |
| 'unet': torch.float16, | |
| 'text_encoder': torch.float16, | |
| 'movq': torch.float32, | |
| } | |
| pipe = get_inpainting_pipeline( | |
| device_map, dtype_map, | |
| ) | |
| augmenter = ImageAugmentation(target_width=2560, target_height=1440) | |
| image = Image.open(image_path='/home/product_diffusion_api/sample_data/example1.jpg') | |
| extended_image = augmenter.extend_image(image) | |
| mask_image = augmenter.generate_mask_from_bbox(extended_image, segmentation_model='facebook/sam-vit-base', detection_model='yolov8s') | |
| mask_image = augmenter.invert_mask(mask_image) | |
| image = pipe( "Product on the Kitchen used for cooking", extended_image, mask_image) | |
| image.save('output.jpg') | |