entanglement.py

import os
import torch
from PIL import Image
from kandinsky2 import get_kandinsky2
import numpy as np
from diffusers import KandinskyV22InpaintPipeline
from diffusers.models import UNet2DConditionModel

def make_masks(image, target_img_size):

    img_size = image.size
    mask_size = (min(target_img_size), min(target_img_size))
    image = np.asarray(image)

    delta_width = target_img_size[0] - img_size[0]
    delta_height = target_img_size[1] - img_size[1]
    pad_width = delta_width // 2
    pad_height = delta_height // 2

    mask_1 = Image.new("RGB", mask_size, color=(255,255,255))
    init_img_1 = Image.new("RGB", mask_size, color=(255,255,255))
    for i in range(pad_width, mask_size[0]):
        for j in range(pad_height, mask_size[1]):
            mask_1.putpixel((i,j), (0, 0, 0))
            init_img_1.putpixel((i,j), tuple(image[j-pad_height, i-pad_width]))

    mask_2 = Image.new("RGB", mask_size, color=(255,255,255))
    init_img_2 = Image.new("RGB", mask_size, color=(255,255,255))
    for i in range(0, mask_size[0]- pad_width):
        for j in range(0, mask_size[1] - pad_height):
            mask_2.putpixel((i,j), (0, 0, 0))
            init_img_2.putpixel((i,j), tuple(image[j+target_img_size[1]-mask_size[1]-pad_height-1, i+target_img_size[0]-mask_size[0]-pad_width-1]))

    return [init_img_1, init_img_2], [mask_1, mask_2]

def get_img_embedding(model, cond, num_images_per_prompt=1):
    cond = (
            model.prior.image_processor(cond, return_tensors="pt")
            .pixel_values[0]
            .unsqueeze(0)
            .to(dtype=model.prior.image_encoder.dtype, device=model.device)
            )

    return model.prior.image_encoder(cond)["image_embeds"].repeat(num_images_per_prompt, 1).unsqueeze(0)

def outpaint_img(img, target_img_size, reconstructed_vector, negative_emb, decoder, generator=None):
    # load base and mask image
    img.thumbnail(target_img_size)
    init_images, mask_images = make_masks(img, target_img_size)

    # Generate image 1
    gen_1 = decoder(image_embeds=reconstructed_vector,
                    negative_image_embeds=negative_emb,
                    num_inference_steps=50, 
                    height=1024,
                    width=1024, 
                    guidance_scale=4,
                    image=init_images[0], 
                    mask_image=mask_images[0],
                    generator=generator).images[0].resize(mask_images[0].size)

    # Generate image 2
    gen_2 = decoder(image_embeds=reconstructed_vector, 
                    negative_image_embeds=negative_emb,
                    num_inference_steps=50, 
                    height=1024,
                    width=1024, 
                    guidance_scale=4,
                    image=init_images[1], 
                    mask_image=mask_images[1],
                    generator=generator).images[0].resize(mask_images[1].size)

    # stack images
    gen_images = [gen_1, gen_2]

    # outpainted image
    combined_image = Image.new("RGB", target_img_size)

    # Calculate the position to paste the first image
    x_offset = 0
    y_offset = 0

    # Paste the first image onto the combined image
    combined_image.paste(gen_images[0], (x_offset, y_offset))

    # Calculate the position to paste the second image
    x_offset = combined_image.width - gen_images[1].width
    y_offset = combined_image.height - gen_images[1].height

    # Paste the second image onto the combined image
    combined_image.paste(gen_images[1], (x_offset, y_offset))

    return combined_image

if __name__ == '__main__':
    # make models
    unet = UNet2DConditionModel.from_pretrained('kandinsky-community/kandinsky-2-2-decoder-inpaint', subfolder='unet').to(torch.float16).to('cuda')
    decoder = KandinskyV22InpaintPipeline.from_pretrained('kandinsky-community/kandinsky-2-2-decoder-inpaint', unet=unet, torch_dtype=torch.float16).to('cuda')    
    model = get_kandinsky2('cuda', task_type='text2img', model_version='2.2', use_flash_attention=False)
    generator = None

    # number of embeddings to get to make the PCA space 
    n_images = 1000

    emb_list = []
    emb_path = 'entanglement/emb_background.pt'

    if not os.path.exists(emb_path):
        for i in range(n_images):
            image_emb = model.prior(prompt='background of a painting', 
                                    num_inference_steps=25, 
                                    num_images_per_prompt=1,
                                    guidance_scale=4, 
                                    negative_prompt='',
                                    generator=generator).image_embeds
            
            emb_list.append(image_emb)

        embeddings = torch.cat(emb_list, dim=0).type(torch.float32)
        torch.save(embeddings, emb_path)

    else:
        embeddings = torch.load(emb_path)

    # Normalize embeddings
    emb_mean = embeddings.mean()
    emb_std = embeddings.std()
    embeddings_normalized = (embeddings - emb_mean) / emb_std

    # Perform Singular Value Decomposition
    U, S, V = torch.svd(embeddings_normalized)
    k = 100
    principal_components = V[:, :k]
    print("\nSelected Principal Components: ", principal_components.size())


    ### Get the image embedding
    image_path = 'dog.jpg'
    img = Image.open(image_path).convert('RGB')
    out_path = 'entanglement/og.png'
    img.save(out_path)

    # embedd image
    img_emb = get_img_embedding(model, img).type(torch.float32)[0]
    # normalize embedding
    norm_img_emb = (img_emb - emb_mean) / emb_std

    projected_data = torch.matmul(norm_img_emb, principal_components)
    reconstructed_vector = torch.matmul(projected_data, principal_components.t())

    # Denormalize
    reconstructed_vector = reconstructed_vector * emb_std + emb_mean

    print("\nReconstructed Vector:")
    print(reconstructed_vector.size())

    n_prompt = 'foreground'
    negative_emb = model.prior(prompt=n_prompt, num_inference_steps=25, num_images_per_prompt=1, guidance_scale=4, generator=generator)

    if n_prompt == '':
        negative_emb = negative_emb.negative_image_embeds

    else:
        negative_emb = negative_emb.image_embeds

    # Generate background
    reconstructed_image = model.decoder(
        image_embeds=reconstructed_vector,
        negative_image_embeds=negative_emb,
        num_inference_steps=50, 
        height=1024,
        width=1024, 
        guidance_scale=4,
        generator=generator).images[0]

    out_path = 'entanglement/reconstructed_background.png'
    reconstructed_image.save(out_path)

    # Generate foreground
    diff = img_emb - reconstructed_vector
    print(img_emb.size(), reconstructed_vector.size(), diff.size())
    reconstructed_image = model.decoder(
        image_embeds= diff,
        negative_image_embeds=negative_emb,
        num_inference_steps=50, 
        height=1024,
        width=1024, 
        guidance_scale=4,
        generator=generator).images[0]

    out_path = 'entanglement/reconstructed_foreground.png'
    reconstructed_image.save(out_path)

    # Outpaint
    target_img_size = (1280, 720) 
    combined_image = outpaint_img(img, target_img_size, reconstructed_vector, negative_emb, decoder)

    out_path = 'entanglement/outpainted.png'
    combined_image.save(out_path)