Delete generation.py

generation.py

@@ -1,621 +0,0 @@
-import numpy as np
-import torch
-from PIL import Image, ImageDraw, ImageFont
-from tqdm import tqdm
-from typing import Union
-from IPython.display import display
-import p2p
-
-
-# Main function to run
-# ----------------------------------------------------------------------
-@torch.no_grad()
-def runner(
-        model,
-        prompt,
-        controller,
-        solver,
-        is_cons_forward=False,
-        num_inference_steps=50,
-        guidance_scale=7.5,
-        generator=None,
-        latent=None,
-        uncond_embeddings=None,
-        start_time=50,
-        return_type='image',
-        dynamic_guidance=False,
-        tau1=0.4,
-        tau2=0.6,
-        w_embed_dim=0,
-):
-    p2p.register_attention_control(model, controller)
-    height = width = 512
-    solver.init_prompt(prompt, None)
-    latent, latents = init_latent(latent, model, 512, 512, generator, len(prompt))
-    model.scheduler.set_timesteps(num_inference_steps)
-    dynamic_guidance = True if tau1 < 1.0 or tau1 < 1.0 else False
-
-    if not is_cons_forward:
-        latents = solver.ddim_loop(latents,
-                                   num_inference_steps,
-                                   is_forward=False,
-                                   guidance_scale=guidance_scale,
-                                   dynamic_guidance=dynamic_guidance,
-                                   tau1=tau1,
-                                   tau2=tau2,
-                                   w_embed_dim=w_embed_dim,
-                                   uncond_embeddings=uncond_embeddings if uncond_embeddings is not None else None,
-                                   controller=controller)
-        latents = latents[-1]
-    else:
-        latents = solver.cons_generation(
-            latents,
-            guidance_scale=guidance_scale,
-            w_embed_dim=w_embed_dim,
-            dynamic_guidance=dynamic_guidance,
-            tau1=tau1,
-            tau2=tau2,
-            controller=controller)
-        latents = latents[-1]
-
-    if return_type == 'image':
-        image = latent2image(model.vae, latents.to(model.vae.dtype))
-    else:
-        image = latents
-
-    return image, latent
-
-
-# ----------------------------------------------------------------------
-
-
-# Utils
-# ----------------------------------------------------------------------
-def linear_schedule_old(t, guidance_scale, tau1, tau2):
-    t = t / 1000
-    if t <= tau1:
-        gamma = 1.0
-    elif t >= tau2:
-        gamma = 0.0
-    else:
-        gamma = (tau2 - t) / (tau2 - tau1)
-    return gamma * guidance_scale
-
-
-def linear_schedule(t, guidance_scale, tau1=0.4, tau2=0.8):
-    t = t / 1000
-    if t <= tau1:
-        return guidance_scale
-    if t >= tau2:
-        return 1.0
-    gamma = (tau2 - t) / (tau2 - tau1) * (guidance_scale - 1.0) + 1.0
-
-    return gamma
-
-
-def guidance_scale_embedding(w, embedding_dim=512, dtype=torch.float32):
-    """
-    See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
-
-    Args:
-        timesteps (`torch.Tensor`):
-            generate embedding vectors at these timesteps
-        embedding_dim (`int`, *optional*, defaults to 512):
-            dimension of the embeddings to generate
-        dtype:
-            data type of the generated embeddings
-
-    Returns:
-        `torch.FloatTensor`: Embedding vectors with shape `(len(timesteps), embedding_dim)`
-    """
-    assert len(w.shape) == 1
-    w = w * 1000.0
-
-    half_dim = embedding_dim // 2
-    emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
-    emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
-    emb = w.to(dtype)[:, None] * emb[None, :]
-    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
-    if embedding_dim % 2 == 1:  # zero pad
-        emb = torch.nn.functional.pad(emb, (0, 1))
-    assert emb.shape == (w.shape[0], embedding_dim)
-    return emb
-
-
-# ----------------------------------------------------------------------
-
-
-# Diffusion step with scheduler from diffusers and controller for editing
-# ----------------------------------------------------------------------
-def extract_into_tensor(a, t, x_shape):
-    b, *_ = t.shape
-    out = a.gather(-1, t)
-    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
-
-
-def predicted_origin(model_output, timesteps, boundary_timesteps, sample, prediction_type, alphas, sigmas):
-    sigmas_s = extract_into_tensor(sigmas, boundary_timesteps, sample.shape)
-    alphas_s = extract_into_tensor(alphas, boundary_timesteps, sample.shape)
-
-    sigmas = extract_into_tensor(sigmas, timesteps, sample.shape)
-    alphas = extract_into_tensor(alphas, timesteps, sample.shape)
-
-    # Set hard boundaries to ensure equivalence with forward (direct) CD
-    alphas_s[boundary_timesteps == 0] = 1.0
-    sigmas_s[boundary_timesteps == 0] = 0.0
-
-    if prediction_type == "epsilon":
-        pred_x_0 = (sample - sigmas * model_output) / alphas  # x0 prediction
-        pred_x_0 = alphas_s * pred_x_0 + sigmas_s * model_output  # Euler step to the boundary step
-    elif prediction_type == "v_prediction":
-        assert boundary_timesteps == 0, "v_prediction does not support multiple endpoints at the moment"
-        pred_x_0 = alphas * sample - sigmas * model_output
-    else:
-        raise ValueError(f"Prediction type {prediction_type} currently not supported.")
-    return pred_x_0
-
-
-def guided_step(noise_prediction_text,
-                noise_pred_uncond,
-                t,
-                guidance_scale,
-                dynamic_guidance=False,
-                tau1=0.4,
-                tau2=0.6):
-    if dynamic_guidance:
-        if not isinstance(t, int):
-            t = t.item()
-        new_guidance_scale = linear_schedule(t, guidance_scale, tau1=tau1, tau2=tau2)
-    else:
-        new_guidance_scale = guidance_scale
-
-    noise_pred = noise_pred_uncond + new_guidance_scale * (noise_prediction_text - noise_pred_uncond)
-    return noise_pred
-
-
-# ----------------------------------------------------------------------
-
-
-# DDIM scheduler with inversion
-# ----------------------------------------------------------------------
-class Generator:
-
-    def prev_step(self, model_output: Union[torch.FloatTensor, np.ndarray], timestep: int,
-                  sample: Union[torch.FloatTensor, np.ndarray]):
-        prev_timestep = timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
-        alpha_prod_t = self.scheduler.alphas_cumprod[timestep]
-        alpha_prod_t_prev = self.scheduler.alphas_cumprod[
-            prev_timestep] if prev_timestep >= 0 else self.scheduler.final_alpha_cumprod
-        beta_prod_t = 1 - alpha_prod_t
-        pred_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
-        pred_sample_direction = (1 - alpha_prod_t_prev) ** 0.5 * model_output
-        prev_sample = alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction
-        return prev_sample
-
-    def next_step(self, model_output: Union[torch.FloatTensor, np.ndarray], timestep: int,
-                  sample: Union[torch.FloatTensor, np.ndarray]):
-        timestep, next_timestep = min(
-            timestep - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps, 999), timestep
-        alpha_prod_t = self.scheduler.alphas_cumprod[timestep] if timestep >= 0 else self.scheduler.final_alpha_cumprod
-        alpha_prod_t_next = self.scheduler.alphas_cumprod[next_timestep]
-        beta_prod_t = 1 - alpha_prod_t
-        next_original_sample = (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
-        next_sample_direction = (1 - alpha_prod_t_next) ** 0.5 * model_output
-        next_sample = alpha_prod_t_next ** 0.5 * next_original_sample + next_sample_direction
-        return next_sample
-
-    def get_noise_pred_single(self, latents, t, context):
-        noise_pred = self.model.unet(latents, t, encoder_hidden_states=context)["sample"]
-        return noise_pred
-
-    def get_noise_pred(self,
-                       model,
-                       latent,
-                       t,
-                       guidance_scale=1,
-                       context=None,
-                       w_embed_dim=0,
-                       dynamic_guidance=False,
-                       tau1=0.4,
-                       tau2=0.6):
-        latents_input = torch.cat([latent] * 2)
-        if context is None:
-            context = self.context
-
-        # w embed 
-        # --------------------------------------
-        if w_embed_dim > 0:
-            if dynamic_guidance:
-                if not isinstance(t, int):
-                    t_item = t.item()
-                guidance_scale = linear_schedule_old(t_item, guidance_scale, tau1=tau1, tau2=tau2)  # TODO UPDATE
-            if len(latents_input) == 4:
-                guidance_scale_tensor = torch.tensor([0.0, 0.0, 0.0, guidance_scale])
-            else:
-                guidance_scale_tensor = torch.tensor([guidance_scale] * len(latents_input))
-            w_embedding = guidance_scale_embedding(guidance_scale_tensor, embedding_dim=w_embed_dim)
-            w_embedding = w_embedding.to(device=latent.device, dtype=latent.dtype)
-        else:
-            w_embedding = None
-        # --------------------------------------
-        noise_pred = model.unet(latents_input.to(dtype=model.unet.dtype),
-                                t,
-                                timestep_cond=w_embedding.to(dtype=model.unet.dtype) if w_embed_dim > 0 else None,
-                                encoder_hidden_states=context)["sample"]
-        noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2)
-
-        if guidance_scale > 1 and w_embedding is None:
-            noise_pred = guided_step(noise_prediction_text, noise_pred_uncond, t, guidance_scale, dynamic_guidance,
-                                     tau1, tau2)
-        else:
-            noise_pred = noise_prediction_text
-
-        return noise_pred
-
-    @torch.no_grad()
-    def latent2image(self, latents, return_type='np'):
-        latents = 1 / 0.18215 * latents.detach()
-        image = self.model.vae.decode(latents.to(dtype=self.model.dtype))['sample']
-        if return_type == 'np':
-            image = (image / 2 + 0.5).clamp(0, 1)
-            image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
-            image = (image * 255).astype(np.uint8)
-        return image
-
-    @torch.no_grad()
-    def image2latent(self, image):
-        with torch.no_grad():
-            if type(image) is Image:
-                image = np.array(image)
-            if type(image) is torch.Tensor and image.dim() == 4:
-                latents = image
-            elif type(image) is list:
-                image = [np.array(i).reshape(1, 512, 512, 3) for i in image]
-                image = np.concatenate(image)
-                image = torch.from_numpy(image).float() / 127.5 - 1
-                image = image.permute(0, 3, 1, 2).to(self.model.device, dtype=self.model.vae.dtype)
-                latents = self.model.vae.encode(image)['latent_dist'].mean
-                latents = latents * 0.18215
-            else:
-                image = torch.from_numpy(image).float() / 127.5 - 1
-                image = image.permute(2, 0, 1).unsqueeze(0).to(self.model.device, dtype=self.model.dtype)
-                latents = self.model.vae.encode(image)['latent_dist'].mean
-                latents = latents * 0.18215
-        return latents
-
-    @torch.no_grad()
-    def init_prompt(self, prompt, uncond_embeddings=None):
-        if uncond_embeddings is None:
-            uncond_input = self.model.tokenizer(
-                [""], padding="max_length", max_length=self.model.tokenizer.model_max_length,
-                return_tensors="pt"
-            )
-            uncond_embeddings = self.model.text_encoder(uncond_input.input_ids.to(self.model.device))[0]
-        text_input = self.model.tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=self.model.tokenizer.model_max_length,
-            truncation=True,
-            return_tensors="pt",
-        )
-        text_embeddings = self.model.text_encoder(text_input.input_ids.to(self.model.device))[0]
-        self.context = torch.cat([uncond_embeddings.expand(*text_embeddings.shape), text_embeddings])
-        self.prompt = prompt
-
-    @torch.no_grad()
-    def ddim_loop(self,
-                  latent,
-                  n_steps,
-                  is_forward=True,
-                  guidance_scale=1,
-                  dynamic_guidance=False,
-                  tau1=0.4,
-                  tau2=0.6,
-                  w_embed_dim=0,
-                  uncond_embeddings=None,
-                  controller=None):
-        all_latent = [latent]
-        latent = latent.clone().detach()
-        for i in tqdm(range(n_steps)):
-            if uncond_embeddings is not None:
-                self.init_prompt(self.prompt, uncond_embeddings[i])
-            if is_forward:
-                t = self.model.scheduler.timesteps[len(self.model.scheduler.timesteps) - i - 1]
-            else:
-                t = self.model.scheduler.timesteps[i]
-            noise_pred = self.get_noise_pred(
-                model=self.model,
-                latent=latent,
-                t=t,
-                context=None,
-                guidance_scale=guidance_scale,
-                dynamic_guidance=dynamic_guidance,
-                w_embed_dim=w_embed_dim,
-                tau1=tau1,
-                tau2=tau2)
-            if is_forward:
-                latent = self.next_step(noise_pred, t, latent)
-            else:
-                latent = self.prev_step(noise_pred, t, latent)
-            if controller is not None:
-                latent = controller.step_callback(latent)
-            all_latent.append(latent)
-        return all_latent
-
-    @property
-    def scheduler(self):
-        return self.model.scheduler
-
-    @torch.no_grad()
-    def ddim_inversion(self,
-                       image,
-                       n_steps=None,
-                       guidance_scale=1,
-                       dynamic_guidance=False,
-                       tau1=0.4,
-                       tau2=0.6,
-                       w_embed_dim=0):
-
-        if n_steps is None:
-            n_steps = self.n_steps
-        latent = self.image2latent(image)
-        image_rec = self.latent2image(latent)
-        ddim_latents = self.ddim_loop(latent,
-                                      is_forward=True,
-                                      guidance_scale=guidance_scale,
-                                      n_steps=n_steps,
-                                      dynamic_guidance=dynamic_guidance,
-                                      tau1=tau1,
-                                      tau2=tau2,
-                                      w_embed_dim=w_embed_dim)
-        return image_rec, ddim_latents
-
-    @torch.no_grad()
-    def cons_generation(self,
-                        latent,
-                        guidance_scale=1,
-                        dynamic_guidance=False,
-                        tau1=0.4,
-                        tau2=0.6,
-                        w_embed_dim=0,
-                        controller=None, ):
-
-        all_latent = [latent]
-        latent = latent.clone().detach()
-        alpha_schedule = torch.sqrt(self.model.scheduler.alphas_cumprod).to(self.model.device)
-        sigma_schedule = torch.sqrt(1 - self.model.scheduler.alphas_cumprod).to(self.model.device)
-
-        for i, (t, s) in enumerate(tqdm(zip(self.reverse_timesteps, self.reverse_boundary_timesteps))):
-            noise_pred = self.get_noise_pred(
-                model=self.reverse_cons_model,
-                latent=latent,
-                t=t.to(self.model.device),
-                context=None,
-                tau1=tau1, tau2=tau2,
-                w_embed_dim=w_embed_dim,
-                guidance_scale=guidance_scale,
-                dynamic_guidance=dynamic_guidance)
-
-            latent = predicted_origin(
-                noise_pred,
-                torch.tensor([t] * len(latent), device=self.model.device),
-                torch.tensor([s] * len(latent), device=self.model.device),
-                latent,
-                self.model.scheduler.config.prediction_type,
-                alpha_schedule,
-                sigma_schedule,
-            )
-            if controller is not None:
-                latent = controller.step_callback(latent)
-            all_latent.append(latent)
-
-        return all_latent
-
-    @torch.no_grad()
-    def cons_inversion(self,
-                       image,
-                       guidance_scale=0.0,
-                       w_embed_dim=0,
-                       seed=0):
-        alpha_schedule = torch.sqrt(self.model.scheduler.alphas_cumprod).to(self.model.device)
-        sigma_schedule = torch.sqrt(1 - self.model.scheduler.alphas_cumprod).to(self.model.device)
-
-        # 5. Prepare latent variables
-        latent = self.image2latent(image)
-        generator = torch.Generator().manual_seed(seed)
-        noise = torch.randn(latent.shape, generator=generator).to(latent.device)
-        latent = self.noise_scheduler.add_noise(latent, noise, torch.tensor([self.start_timestep]))
-        image_rec = self.latent2image(latent)
-
-        for i, (t, s) in enumerate(tqdm(zip(self.forward_timesteps, self.forward_boundary_timesteps))):
-            # predict the noise residual
-            noise_pred = self.get_noise_pred(
-                model=self.forward_cons_model,
-                latent=latent,
-                t=t.to(self.model.device),
-                context=None,
-                guidance_scale=guidance_scale,
-                w_embed_dim=w_embed_dim,
-                dynamic_guidance=False)
-
-            latent = predicted_origin(
-                noise_pred,
-                torch.tensor([t] * len(latent), device=self.model.device),
-                torch.tensor([s] * len(latent), device=self.model.device),
-                latent,
-                self.model.scheduler.config.prediction_type,
-                alpha_schedule,
-                sigma_schedule,
-            )
-
-        return image_rec, [latent]
-
-    def _create_forward_inverse_timesteps(self,
-                                          num_endpoints,
-                                          n_steps,
-                                          max_inverse_timestep_index):
-        timestep_interval = n_steps // num_endpoints + int(n_steps % num_endpoints > 0)
-        endpoint_idxs = torch.arange(timestep_interval, n_steps, timestep_interval) - 1
-        inverse_endpoint_idxs = torch.arange(timestep_interval, n_steps, timestep_interval) - 1
-        inverse_endpoint_idxs = torch.tensor(inverse_endpoint_idxs.tolist() + [max_inverse_timestep_index])
-
-        endpoints = torch.tensor([0] + self.ddim_timesteps[endpoint_idxs].tolist())
-        inverse_endpoints = self.ddim_timesteps[inverse_endpoint_idxs]
-
-        return endpoints, inverse_endpoints
-
-    def __init__(self,
-                 model,
-                 n_steps,
-                 noise_scheduler,
-                 forward_cons_model=None,
-                 reverse_cons_model=None,
-                 num_endpoints=1,
-                 num_forward_endpoints=1,
-                 reverse_timesteps=None,
-                 forward_timesteps=None,
-                 max_forward_timestep_index=49,
-                 start_timestep=19):
-
-        self.model = model
-        self.forward_cons_model = forward_cons_model
-        self.reverse_cons_model = reverse_cons_model
-        self.noise_scheduler = noise_scheduler
-
-        self.n_steps = n_steps
-        self.tokenizer = self.model.tokenizer
-        self.model.scheduler.set_timesteps(n_steps)
-        self.prompt = None
-        self.context = None
-        step_ratio = 1000 // n_steps
-        self.ddim_timesteps = (np.arange(1, n_steps + 1) * step_ratio).round().astype(np.int64) - 1
-        self.ddim_timesteps = torch.from_numpy(self.ddim_timesteps).long()
-        self.start_timestep = start_timestep
-
-        # Set endpoints for direct CTM
-        if reverse_timesteps is None or forward_timesteps is None:
-            endpoints, inverse_endpoints = self._create_forward_inverse_timesteps(num_endpoints, n_steps,
-                                                                                  max_forward_timestep_index)
-            self.reverse_timesteps, self.reverse_boundary_timesteps = inverse_endpoints.flip(0), endpoints.flip(0)
-
-            # Set endpoints for forward CTM
-            endpoints, inverse_endpoints = self._create_forward_inverse_timesteps(num_forward_endpoints, n_steps,
-                                                                                  max_forward_timestep_index)
-            self.forward_timesteps, self.forward_boundary_timesteps = endpoints, inverse_endpoints
-            self.forward_timesteps[0] = self.start_timestep
-        else:
-            self.reverse_timesteps, self.reverse_boundary_timesteps = reverse_timesteps, reverse_timesteps
-            self.reverse_timesteps.reverse()
-            self.reverse_boundary_timesteps = self.reverse_boundary_timesteps[1:] + [self.reverse_boundary_timesteps[0]]
-            self.reverse_boundary_timesteps[-1] = 0
-            self.reverse_timesteps, self.reverse_boundary_timesteps = torch.tensor(reverse_timesteps), torch.tensor(
-                self.reverse_boundary_timesteps)
-
-            self.forward_timesteps, self.forward_boundary_timesteps = forward_timesteps, forward_timesteps
-            self.forward_boundary_timesteps = self.forward_boundary_timesteps[1:] + [self.forward_boundary_timesteps[0]]
-            self.forward_boundary_timesteps[-1] = 999
-            self.forward_timesteps, self.forward_boundary_timesteps = torch.tensor(
-                self.forward_timesteps), torch.tensor(self.forward_boundary_timesteps)
-
-        print(f"Endpoints reverse CTM: {self.reverse_timesteps}, {self.reverse_boundary_timesteps}")
-        print(f"Endpoints forward CTM: {self.forward_timesteps}, {self.forward_boundary_timesteps}")
-
-# ----------------------------------------------------------------------
-
-# 3rd party utils
-# ----------------------------------------------------------------------
-def latent2image(vae, latents):
-    latents = 1 / 0.18215 * latents
-    image = vae.decode(latents)['sample']
-    image = (image / 2 + 0.5).clamp(0, 1)
-    image = image.cpu().permute(0, 2, 3, 1).numpy()
-    image = (image * 255).astype(np.uint8)
-    return image
-
-
-def init_latent(latent, model, height, width, generator, batch_size):
-    if latent is None:
-        latent = torch.randn(
-            (1, model.unet.in_channels, height // 8, width // 8),
-            generator=generator,
-        )
-    latents = latent.expand(batch_size, model.unet.in_channels, height // 8, width // 8).to(model.device)
-    return latent, latents
-
-
-def load_512(image_path, left=0, right=0, top=0, bottom=0):
-    # if type(image_path) is str:
-    #    image = np.array(Image.open(image_path))[:, :, :3]
-    # else:
-    #    image = image_path
-    # h, w, c = image.shape
-    # left = min(left, w - 1)
-    # right = min(right, w - left - 1)
-    # top = min(top, h - left - 1)
-    # bottom = min(bottom, h - top - 1)
-    # image = image[top:h - bottom, left:w - right]
-    # h, w, c = image.shape
-    # if h < w:
-    #    offset = (w - h) // 2
-    #    image = image[:, offset:offset + h]
-    # elif w < h:
-    #    offset = (h - w) // 2
-    #    image = image[offset:offset + w]
-    image = np.array(Image.open(image_path).convert('RGB'))[:, :, :3]
-    image = np.array(Image.fromarray(image).resize((512, 512)))
-    return image
-
-
-def to_pil_images(images, num_rows=1, offset_ratio=0.02):
-    if type(images) is list:
-        num_empty = len(images) % num_rows
-    elif images.ndim == 4:
-        num_empty = images.shape[0] % num_rows
-    else:
-        images = [images]
-        num_empty = 0
-
-    empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255
-    images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty
-    num_items = len(images)
-
-    h, w, c = images[0].shape
-    offset = int(h * offset_ratio)
-    num_cols = num_items // num_rows
-    image_ = np.ones((h * num_rows + offset * (num_rows - 1),
-                      w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255
-    for i in range(num_rows):
-        for j in range(num_cols):
-            image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images[
-                i * num_cols + j]
-
-    pil_img = Image.fromarray(image_)
-    return pil_img
-
-
-def view_images(images, num_rows=1, offset_ratio=0.02):
-    if type(images) is list:
-        num_empty = len(images) % num_rows
-    elif images.ndim == 4:
-        num_empty = images.shape[0] % num_rows
-    else:
-        images = [images]
-        num_empty = 0
-
-    empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255
-    images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty
-    num_items = len(images)
-
-    h, w, c = images[0].shape
-    offset = int(h * offset_ratio)
-    num_cols = num_items // num_rows
-    image_ = np.ones((h * num_rows + offset * (num_rows - 1),
-                      w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255
-    for i in range(num_rows):
-        for j in range(num_cols):
-            image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images[
-                i * num_cols + j]
-
-    pil_img = Image.fromarray(image_)
-    display(pil_img)
-# ----------------------------------------------------------------------