Inference

LMSDiscreteScheduler.py

@@ -0,0 +1,97 @@
+import warnings
+from typing import Tuple, Union
+
+import torch
+from diffusers.schedulers.scheduling_lms_discrete import \
+    LMSDiscreteScheduler as _LMSDiscreteScheduler
+from diffusers.schedulers.scheduling_lms_discrete import \
+    LMSDiscreteSchedulerOutput
+
+
+class LMSDiscreteScheduler(_LMSDiscreteScheduler):
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        step_index: int,
+        sample: torch.FloatTensor,
+        order: int = 4,
+        return_dict: bool = True,
+    ) -> Union[LMSDiscreteSchedulerOutput, Tuple]:
+        """
+        Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Args:
+            model_output (`torch.FloatTensor`): direct output from learned diffusion model.
+            timestep (`float`): current timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                current instance of sample being created by diffusion process.
+            order: coefficient for multi-step inference.
+            return_dict (`bool`): option for returning tuple rather than LMSDiscreteSchedulerOutput class
+
+        Returns:
+            [`~schedulers.scheduling_utils.LMSDiscreteSchedulerOutput`] or `tuple`:
+            [`~schedulers.scheduling_utils.LMSDiscreteSchedulerOutput`] if `return_dict` is True, otherwise a `tuple`.
+            When returning a tuple, the first element is the sample tensor.
+
+        """
+        if not self.is_scale_input_called:
+            warnings.warn(
+                "The `scale_model_input` function should be called before `step` to ensure correct denoising. "
+                "See `StableDiffusionPipeline` for a usage example."
+            )
+
+        sigma = self.sigmas[step_index]
+
+        # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
+        if self.config.prediction_type == "epsilon":
+            pred_original_sample = sample - sigma * model_output
+        elif self.config.prediction_type == "v_prediction":
+            # * c_out + input * c_skip
+            pred_original_sample = model_output * \
+                (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1))
+        else:
+            raise ValueError(
+                f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`"
+            )
+
+        # 2. Convert to an ODE derivative
+        derivative = (sample - pred_original_sample) / sigma
+        self.derivatives.append(derivative)
+        if len(self.derivatives) > order:
+            self.derivatives.pop(0)
+
+        # 3. Compute linear multistep coefficients
+        order = min(step_index + 1, order)
+        lms_coeffs = [self.get_lms_coefficient(
+            order, step_index, curr_order) for curr_order in range(order)]
+
+        # 4. Compute previous sample based on the derivatives path
+        prev_sample = sample + sum(
+            coeff * derivative for coeff, derivative in zip(lms_coeffs, reversed(self.derivatives))
+        )
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return LMSDiscreteSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)
+
+    def scale_model_input(
+        self, 
+        sample: torch.FloatTensor,
+        iteration: int
+    ) -> torch.FloatTensor:
+        """
+        Scales the denoising model input by `(sigma**2 + 1) ** 0.5` to match the K-LMS algorithm.
+
+        Args:
+            sample (`torch.FloatTensor`): input sample
+            timestep (`float` or `torch.FloatTensor`): the current timestep in the diffusion chain
+
+        Returns:
+            `torch.FloatTensor`: scaled input sample
+        """
+        sample = sample / ((self.sigmas[iteration]**2 + 1) ** 0.5)
+        self.is_scale_input_called = True
+        return sample

StableDiffuser.py

@@ -0,0 +1,276 @@
+import argparse
+
+import torch
+from baukit import TraceDict
+from diffusers import AutoencoderKL, UNet2DConditionModel
+from PIL import Image
+from tqdm.auto import tqdm
+from transformers import CLIPTextModel, CLIPTokenizer
+
+import util
+from LMSDiscreteScheduler import LMSDiscreteScheduler
+
+
+def default_parser():
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('prompts', type=str, nargs='+')
+    parser.add_argument('outpath', type=str)
+
+    parser.add_argument('--images', type=str, nargs='+', default=None)
+    parser.add_argument('--nsteps', type=int, default=1000)
+    parser.add_argument('--nimgs', type=int, default=1)
+    parser.add_argument('--start_itr', type=int, default=0)
+    parser.add_argument('--return_steps', action='store_true', default=False)
+    parser.add_argument('--pred_x0', action='store_true', default=False)
+    parser.add_argument('--device', type=str, default='cuda:0')
+    parser.add_argument('--seed', type=int, default=42)
+
+    return parser
+
+
+class StableDiffuser(torch.nn.Module):
+
+    def __init__(self,
+                seed=None
+        ):
+
+        super().__init__()
+
+        self._seed = seed
+
+        # Load the autoencoder model which will be used to decode the latents into image space.
+        self.vae = AutoencoderKL.from_pretrained(
+            "CompVis/stable-diffusion-v1-4", subfolder="vae")
+        
+        # Load the tokenizer and text encoder to tokenize and encode the text.
+        self.tokenizer = CLIPTokenizer.from_pretrained(
+            "openai/clip-vit-large-patch14")
+        self.text_encoder = CLIPTextModel.from_pretrained(
+            "openai/clip-vit-large-patch14")
+        
+        # The UNet model for generating the latents.
+        self.unet = UNet2DConditionModel.from_pretrained(
+            "CompVis/stable-diffusion-v1-4", subfolder="unet")
+        
+        self.scheduler = LMSDiscreteScheduler(
+            beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
+
+        self.generator = torch.Generator()
+
+        if self._seed is not None:
+
+            self.seed(seed)
+
+        self.eval()
+
+    def seed(self, seed):
+
+        self.generator = torch.manual_seed(seed)
+
+    def get_noise(self, batch_size, img_size):
+
+        param = list(self.parameters())[0]
+
+        return torch.randn(
+            (batch_size, self.unet.in_channels, img_size // 8, img_size // 8),
+            generator=self.generator).type(param.dtype).to(param.device)
+
+    def add_noise(self, latents, noise, step):
+
+        return self.scheduler.add_noise(latents, noise, torch.tensor([self.scheduler.timesteps[step]]))
+
+    def text_tokenize(self, prompts):
+
+        return self.tokenizer(prompts, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
+
+    def text_detokenize(self, tokens):
+
+        return [self.tokenizer.decode(token) for token in tokens if token != self.tokenizer.vocab_size - 1]
+
+    def text_encode(self, tokens):
+
+        return self.text_encoder(tokens.input_ids.to(self.unet.device))[0]
+
+    def decode(self, latents):
+
+        return self.vae.decode(1 / 0.18215 * latents).sample
+
+    def encode(self, tensors):
+
+        return self.vae.encode(tensors).latent_dist.mode() * 0.18215
+
+    def to_image(self, image):
+
+        image = (image / 2 + 0.5).clamp(0, 1)
+        image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
+        images = (image * 255).round().astype("uint8")
+        pil_images = [Image.fromarray(image) for image in images]
+
+        return pil_images
+
+    def set_scheduler_timesteps(self, n_steps):
+        self.scheduler.set_timesteps(n_steps, device=self.unet.device)
+
+    def get_initial_latents(self, n_imgs, img_size, n_prompts):
+
+        noise = self.get_noise(n_imgs, img_size).repeat(n_prompts, 1, 1, 1)
+
+        latents = noise * self.scheduler.init_noise_sigma
+
+        return latents
+
+    def get_text_embeddings(self, prompts, n_imgs):
+
+        text_tokens = self.text_tokenize(prompts)
+
+        text_embeddings = self.text_encode(text_tokens)
+
+        unconditional_tokens = self.text_tokenize([""] * len(prompts))
+
+        unconditional_embeddings = self.text_encode(unconditional_tokens)
+
+        text_embeddings = torch.cat([unconditional_embeddings, text_embeddings]).repeat_interleave(n_imgs, dim=0)
+
+        return text_embeddings
+
+    def predict_noise(self,
+             iteration,
+             latents,
+             text_embeddings,
+             guidance_scale=7.5
+             ):
+
+        # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
+        latents = torch.cat([latents] * 2)
+        latents = self.scheduler.scale_model_input(
+            latents, iteration)
+
+        # predict the noise residual
+        noise_prediction = self.unet(
+            latents, self.scheduler.timesteps[iteration], encoder_hidden_states=text_embeddings).sample
+
+        # perform guidance
+        noise_prediction_uncond, noise_prediction_text = noise_prediction.chunk(2)
+        noise_prediction = noise_prediction_uncond + guidance_scale * \
+            (noise_prediction_text - noise_prediction_uncond)
+
+        return noise_prediction
+
+    @torch.no_grad()
+    def diffusion(self,
+                  latents,
+                  text_embeddings,
+                  end_iteration=1000,
+                  start_iteration=0,
+                  return_steps=False,
+                  pred_x0=False,
+                  trace_args=None,                  
+                  show_progress=True,
+                  **kwargs):
+
+        latents_steps = []
+        trace_steps = []
+
+        trace = None
+
+        for iteration in tqdm(range(start_iteration, end_iteration), disable=not show_progress):
+
+            if trace_args:
+
+                trace = TraceDict(self, **trace_args)
+
+            noise_pred = self.predict_noise(
+                iteration, 
+                latents, 
+                text_embeddings,
+                **kwargs)
+
+            # compute the previous noisy sample x_t -> x_t-1
+            output = self.scheduler.step(noise_pred, iteration, latents)
+
+            if trace_args:
+
+                trace.close()
+
+                trace_steps.append(trace)
+
+            latents = output.prev_sample
+
+            if return_steps or iteration == end_iteration - 1:
+
+                output = output.pred_original_sample if pred_x0 else latents
+
+                if return_steps:
+                    latents_steps.append(output.cpu())
+                else:
+                    latents_steps.append(output)
+
+        return latents_steps, trace_steps
+
+    @torch.no_grad()
+    def __call__(self,
+                 prompts,
+                 img_size=512,
+                 n_steps=50,
+                 n_imgs=1,
+                 end_iteration=None,
+                 reseed=False,
+                 **kwargs
+                 ):
+
+        assert 0 <= n_steps <= 1000
+
+        if not isinstance(prompts, list):
+
+            prompts = [prompts]
+
+        self.set_scheduler_timesteps(n_steps)
+
+        if reseed:
+
+            self.seed(self._seed)
+
+        latents = self.get_initial_latents(n_imgs, img_size, len(prompts))
+
+        text_embeddings = self.get_text_embeddings(prompts,n_imgs=n_imgs)
+
+        end_iteration = end_iteration or n_steps
+
+        latents_steps, trace_steps = self.diffusion(
+            latents,
+            text_embeddings,
+            end_iteration=end_iteration,
+            **kwargs
+        )
+
+        latents_steps = [self.decode(latents.to(self.unet.device)) for latents in latents_steps]
+        images_steps = [self.to_image(latents) for latents in latents_steps]
+
+        images_steps = list(zip(*images_steps))
+
+        if trace_steps:
+
+            return images_steps, trace_steps
+
+        return images_steps
+
+
+if __name__ == '__main__':
+
+    parser = default_parser()
+
+    args = parser.parse_args()
+
+    diffuser = StableDiffuser(seed=args.seed).to(torch.device(args.device)).half()
+
+    images = diffuser(args.prompts,
+                      n_steps=args.nsteps,
+                      n_imgs=args.nimgs,
+                      start_iteration=args.start_itr,
+                      return_steps=args.return_steps,
+                      pred_x0=args.pred_x0
+                      )
+
+    util.image_grid(images, args.outpath)

app.py

@@ -2,73 +2,166 @@ import sys
 sys.path.insert(0,'stable_diffusion')
 import gradio as gr
 from train_esd import train_esd
-
+from convertModels import convert_ldm_unet_checkpoint, create_unet_diffusers_config
+from omegaconf import OmegaConf
+from StableDiffuser import StableDiffuser
+from diffusers import UNet2DConditionModel
 
 ckpt_path = "stable-diffusion/models/ldm/sd-v1-4-full-ema.ckpt"
 config_path = "stable-diffusion/configs/stable-diffusion/v1-inference.yaml"
 diffusers_config_path = "stable-diffusion/config.json"
 
-def train(prompt, train_method, neg_guidance, iterations, lr):
-
-    train_esd(prompt,
-              train_method,
-              3,
-              neg_guidance,
-              iterations,
-              lr,
-              config_path,
-              ckpt_path, 
-              diffusers_config_path,
-              ['cuda']
-              )
-
-
-with gr.Blocks() as demo:
-
-    prompt_input = gr.Text(
-        placeholder="Enter prompt...",
-        label="Prompt",
-        info="Prompt corresponding to concept to erase"
-    )
-    train_method_input = gr.Dropdown(
-        choices=['noxattn', 'selfattn', 'xattn', 'full'],
-        value='xattn',
-        label='Train Method',
-        info='Method of training'
-    )
-
-    neg_guidance_input = gr.Number(
-        value=1,
-        label="Negative Guidance",
-        info='Guidance of negative training used to train'
-    )
-
-    iterations_input = gr.Number(
-        value=1000,
-        precision=0,
-        label="Iterations",
-        info='iterations used to train'
-    )
-
-    lr_input = gr.Number(
-        value=1e-5,
-        label="Iterations",
-        info='Learning rate used to train'
-    )
-
-    train_button = gr.Button(
-        value="Train",
-    )
-    train_button.click(train, inputs = [
-            prompt_input,
-            train_method_input, 
-            neg_guidance_input,
-            iterations_input,
-            lr_input
-        ]
-    )
-
-
-
-
-demo.launch()
+
+class Demo:
+
+    def __init__(self) -> None:
+        demo = self.layout()
+        demo.launch()
+
+
+    def layout(self):
+
+        with gr.Blocks() as demo:
+
+            with gr.Row():
+                with gr.Column() as training_column:
+                    self.prompt_input = gr.Text(
+                        placeholder="Enter prompt...",
+                        label="Prompt",
+                        info="Prompt corresponding to concept to erase"
+                    )
+                    self.train_method_input = gr.Dropdown(
+                        choices=['noxattn', 'selfattn', 'xattn', 'full'],
+                        value='xattn',
+                        label='Train Method',
+                        info='Method of training'
+                    )
+
+                    self.neg_guidance_input = gr.Number(
+                        value=1,
+                        label="Negative Guidance",
+                        info='Guidance of negative training used to train'
+                    )
+
+                    self.iterations_input = gr.Number(
+                        value=1000,
+                        precision=0,
+                        label="Iterations",
+                        info='iterations used to train'
+                    )
+
+                    self.lr_input = gr.Number(
+                        value=1e-5,
+                        label="Learning Rate",
+                        info='Learning rate used to train'
+                    )
+
+                    self.train_button = gr.Button(
+                        value="Train",
+                    )
+                    self.train_button.click(self.train, inputs = [
+                            self.prompt_input,
+                            self.train_method_input, 
+                            self.neg_guidance_input,
+                            self.iterations_input,
+                            self.lr_input
+                        ]
+                    )
+                with gr.Column() as inference_column:
+
+                    with gr.Row():
+
+                        self.prompt_input_infr = gr.Text(
+                            placeholder="Enter prompt...",
+                            label="Prompt",
+                            info="Prompt corresponding to concept to erase"
+                        )
+
+                    with gr.Row():
+
+                        self.image_new = gr.Image(
+                            label="New Image",
+                            interactive=False
+                        )
+                        self.image_orig = gr.Image(
+                            label="Orig Image",
+                            interactive=False
+                        )
+
+                    with gr.Row():
+
+                        self.infr_button = gr.Button(
+                            value="Generate",
+                        )
+                        self.infr_button.click(self.inference, inputs = [
+                                self.prompt_input_infr,
+                            ],
+                            outputs=[
+                                self.image_new,
+                                self.image_orig
+                            ]
+                        )
+        return demo
+
+
+    def train(self, prompt, train_method, neg_guidance, iterations, lr):
+
+        model_orig, model_edited = train_esd(prompt,
+                train_method,
+                3,
+                neg_guidance,
+                iterations,
+                lr,
+                config_path,
+                ckpt_path, 
+                diffusers_config_path,
+                ['cuda', 'cuda'],
+                gr.Progress()
+                )
+        
+        original_config = OmegaConf.load(config_path)
+        original_config["model"]["params"]["unet_config"]["params"]["in_channels"] = 4
+        unet_config = create_unet_diffusers_config(original_config, image_size=512)
+        model_edited_sd = convert_ldm_unet_checkpoint(model_edited.state_dict(), unet_config)
+        model_orig_sd = convert_ldm_unet_checkpoint(model_orig.state_dict(), unet_config)
+
+        self.init_inference(model_edited_sd, model_orig_sd, unet_config)
+
+    def init_inference(self, model_edited_sd, model_orig_sd, unet_config):
+
+        self.model_edited_sd = model_edited_sd
+        self.model_orig_sd = model_orig_sd
+
+        self.diffuser = StableDiffuser(42)
+
+        self.diffuser.unet = UNet2DConditionModel(**unet_config)
+        self.diffuser.to('cuda')
+
+
+    def inference(self, prompt):
+
+        self.diffuser.unet.load_state_dict(self.model_orig_sd)
+
+        images = self.diffuser(
+            prompt,
+            n_steps=50,
+            reseed=True
+        )
+
+        orig_image = images[0][0]
+
+        self.diffuser.unet.load_state_dict(self.model_edited_sd)
+
+        images = self.diffuser(
+            prompt,
+            n_steps=50,
+            reseed=True
+        )
+
+        edited_image = images[0][0]
+
+        return edited_image, orig_image
+
+
+
+

requirements.txt

@@ -4,4 +4,8 @@ torchvision
 einops
 diffusers
 transformers
-pytorch_lightning
+pytorch_lightning==1.6.5
+taming-transformers
+kornia
+git+https://github.com/openai/CLIP.git@main#egg=clip
+git+https://github.com/davidbau/baukit.git

test.py

@@ -1,19 +1,32 @@
 import sys
 sys.path.insert(0,'stable_diffusion')
 from train_esd import train_esd
-
+import torch
 ckpt_path = "stable_diffusion/models/ldm/sd-v1-4-full-ema.ckpt"
 config_path = "stable_diffusion/configs/stable-diffusion/v1-inference.yaml"
 diffusers_config_path = "stable_diffusion/config.json"
 
-train_esd("England",
+orig, newm = train_esd("England",
               'xattn',
               3,
               1,
-              1000,
+              2,
               .003,
               config_path,
               ckpt_path, 
               diffusers_config_path,
-              ['cuda', 'cuda']
-              )
+              ['cuda', 'cuda'],
+              None
+              )
+
+
+from convertModels import convert_ldm_unet_checkpoint, create_unet_diffusers_config
+from diffusers import UNet2DConditionModel, AutoencoderKL, LMSDiscreteScheduler
+from omegaconf import OmegaConf
+from transformers import CLIPTextModel, CLIPTokenizer
+original_config = OmegaConf.load(config_path)
+original_config["model"]["params"]["unet_config"]["params"]["in_channels"] = 4
+unet_config = create_unet_diffusers_config(original_config, image_size=512)
+converted_unet_checkpoint = convert_ldm_unet_checkpoint(newm.state_dict(), unet_config)
+unet = UNet2DConditionModel(**unet_config)
+unet.load_state_dict(converted_unet_checkpoint)

train_esd.py

@@ -102,7 +102,7 @@ def get_models(config_path, ckpt_path, devices):
 
     return model_orig, sampler_orig, model, sampler
 
-def train_esd(prompt, train_method, start_guidance, negative_guidance, iterations, lr, config_path, ckpt_path, diffusers_config_path, devices, seperator=None, image_size=512, ddim_steps=50):
+def train_esd(prompt, train_method, start_guidance, negative_guidance, iterations, lr, config_path, ckpt_path, diffusers_config_path, devices, progress_bar, seperator=None, image_size=512, ddim_steps=50):
     '''
     Function to train diffusion models to erase concepts from model weights
 
@@ -251,17 +251,19 @@ def train_esd(prompt, train_method, start_guidance, negative_guidance, iteration
         pbar.set_postfix({"loss": loss.item()})
         history.append(loss.item())
         opt.step()
-        # save checkpoint and loss curve
-        if (i+1) % 500 == 0 and i+1 != iterations and i+1>= 500:
-            save_model(model, name, i-1, save_compvis=True, save_diffusers=False)
+        # # save checkpoint and loss curve
+        # if (i+1) % 500 == 0 and i+1 != iterations and i+1>= 500:
+        #     save_model(model, name, i-1, save_compvis=True, save_diffusers=False)
 
-        if i % 100 == 0:
-            save_history(losses, name, word_print)
+        # if i % 100 == 0:
+        #     save_history(losses, name, word_print)
 
     model.eval()
 
-    save_model(model, name, None, save_compvis=True, save_diffusers=True, compvis_config_file=config_path, diffusers_config_file=diffusers_config_path)
-    save_history(losses, name, word_print)
+    # save_model(model, name, None, save_compvis=True, save_diffusers=True, compvis_config_file=config_path, diffusers_config_file=diffusers_config_path)
+    # save_history(losses, name, word_print)
+
+    return model_orig, model
 
 def save_model(model, name, num, compvis_config_file=None, diffusers_config_file=None, device='cpu', save_compvis=True, save_diffusers=True):
     # SAVE MODEL

util.py

@@ -0,0 +1,107 @@
+from PIL import Image
+from matplotlib import pyplot as plt
+import textwrap
+
+
+def to_gif(images, path):
+
+    images[0].save(path, save_all=True,
+                   append_images=images[1:], loop=0, duration=len(images) * 20)
+
+
+def figure_to_image(figure):
+
+    figure.set_dpi(300)
+
+    figure.canvas.draw()
+
+    return Image.frombytes('RGB', figure.canvas.get_width_height(), figure.canvas.tostring_rgb())
+
+
+def image_grid(images, outpath=None, column_titles=None, row_titles=None):
+
+    n_rows = len(images)
+    n_cols = len(images[0])
+
+    fig, axs = plt.subplots(nrows=n_rows, ncols=n_cols,
+                            figsize=(n_cols, n_rows), squeeze=False)
+
+    for row, _images in enumerate(images):
+
+        for column, image in enumerate(_images):
+            ax = axs[row][column]
+            ax.imshow(image)
+            if column_titles and row == 0:
+                ax.set_title(textwrap.fill(
+                    column_titles[column], width=12), fontsize='x-small')
+            if row_titles and column == 0:
+                ax.set_ylabel(row_titles[row], rotation=0, fontsize='x-small', labelpad=1.6 * len(row_titles[row]))
+            ax.set_xticks([])
+            ax.set_yticks([])
+
+    plt.subplots_adjust(wspace=0, hspace=0)
+
+    if outpath is not None:
+        plt.savefig(outpath, bbox_inches='tight', dpi=300)
+        plt.close()
+    else:
+        plt.tight_layout(pad=0)
+        image = figure_to_image(plt.gcf())
+        plt.close()
+        return image
+
+
+
+
+    
+
+
+def get_module(module, module_name):
+
+    if isinstance(module_name, str):
+        module_name = module_name.split('.')
+
+    if len(module_name) == 0:
+        return module
+    else:
+        module = getattr(module, module_name[0])
+        return get_module(module, module_name[1:])
+
+
+def set_module(module, module_name, new_module):
+
+    if isinstance(module_name, str):
+        module_name = module_name.split('.')
+
+    if len(module_name) == 1:
+        return setattr(module, module_name[0], new_module)
+    else:
+        module = getattr(module, module_name[0])
+        return set_module(module, module_name[1:], new_module)
+
+
+def freeze(module):
+
+    for parameter in module.parameters():
+
+        parameter.requires_grad = False
+
+
+def unfreeze(module):
+
+    for parameter in module.parameters():
+
+        parameter.requires_grad = True
+
+
+def get_concat_h(im1, im2):
+    dst = Image.new('RGB', (im1.width + im2.width, im1.height))
+    dst.paste(im1, (0, 0))
+    dst.paste(im2, (im1.width, 0))
+    return dst
+
+def get_concat_v(im1, im2):
+    dst = Image.new('RGB', (im1.width, im1.height + im2.height))
+    dst.paste(im1, (0, 0))
+    dst.paste(im2, (0, im1.height))
+    return dst