merge_images_with_sd.py

from PIL import Image
import stable_inference
import torch
from einops import repeat

def image_grid(imgs, rows, cols):
    assert len(imgs) == rows*cols

    w, h = imgs[0].size
    grid = Image.new('RGB', size=(cols*w, rows*h))
    grid_w, grid_h = grid.size
    
    for i, img in enumerate(imgs):
        grid.paste(img, box=(i%cols*w, i//cols*h))
    return grid


engine = stable_inference.StableDiffusionInference(
    checkpoint_loc='../sd-v1-5-inpainting.ckpt',
    config_loc='../v1-inpainting.yaml',
)

frida_pil = Image.open('frida.png')
frida_tensor, (_, _) = stable_inference.util.load_img(img=frida_pil)
frida_tensor = frida_tensor.half().to('cuda')
earring_pil = Image.open('earring.png')
earring_tensor, (_, _) = stable_inference.util.load_img(img=earring_pil)
earring_tensor = earring_tensor.half().to('cuda')

frida = engine.model.get_first_stage_encoding(
    engine.model.encode_first_stage(frida_tensor))
earring = engine.model.get_first_stage_encoding(
    engine.model.encode_first_stage(earring_tensor))

prompt1 = 'painting of a woman surrounded by two black cats, tropical leaves in the background, there is a thorny necklace with a bird. painted in the style of Frida Kahlo'
prompt2 = 'painting young woman stands looking over her shoulder in a dark non-descript room. her head is wrapped in a blue bonnet with a sheet over her shoulder. painted in the style of Johannes Vermeer'

(
    conditioning_frida,
    unconditioning, # Reuse this as it's the same for both
    weighted_subprompts_start,
    _, # Don't need the individual embedding managers
) = engine.compute_conditioning_and_weights(
    prompt1,
    4)

(
    conditioning_earring,
    _,
    weighted_subprompts_end,
    _, # Don't need the individual embedding managers
) = engine.compute_conditioning_and_weights(
    prompt2,
    4)
c = stable_inference.util.slerp(
    0.5,
    conditioning_frida,
    conditioning_earring,
)

weighted_subprompts = stable_inference.util.combine_weighted_subprompts(0.5,
    weighted_subprompts_start,
    weighted_subprompts_end)
print(weighted_subprompts)

prob = 0.5
mask = torch.bernoulli(torch.full(frida.shape, 0.5)).int()
mask = mask.to('cuda')
reverse_mask = torch.ones(frida.shape).int().to('cuda') - mask
reverse_mask = reverse_mask.to('cuda')
blended = frida * mask + earring * reverse_mask

_, extra_data = engine.sample(
    '',
    4,
    'dpmpp_2m',
    12345,
    30,
    conditioning=c,
    init_latent=repeat(blended, '1 ... -> b ...', b=4),
    scale=7.5,
    unconditioning=unconditioning,
    weighted_subprompts=weighted_subprompts,
)
image_grid(extra_data['images'], 2, 2).save('fusion_ha.png')

slerped = stable_inference.util.slerp(0.5, frida, earring)

_, extra_data = engine.sample(
    '',
    4,
    'dpmpp_2m',
    12345,
    30,
    conditioning=c,
    init_latent=repeat(slerped, '1 ... -> b ...', b=4),
    scale=7.5,
    unconditioning=unconditioning,
    weighted_subprompts=weighted_subprompts,
)
image_grid(extra_data['images'], 2, 2).save('fusion_ha_ha.png')