omnigen2/utils/img_util.py

from typing import List

from PIL import Image

import torch
from torchvision.transforms.functional import to_pil_image

def resize_image(image, max_pixels, img_scale_num):
    width, height = image.size
    cur_pixels = height * width
    ratio = (max_pixels / cur_pixels) ** 0.5
    ratio = min(ratio, 1.0) # do not upscale input image

    new_height, new_width = int(height * ratio) // img_scale_num * img_scale_num, int(width * ratio) // img_scale_num * img_scale_num

    image = image.resize((new_width, new_height), resample=Image.BICUBIC)
    return image

def create_collage(images: List[torch.Tensor]) -> Image.Image:
    """Create a horizontal collage from a list of images."""
    max_height = max(img.shape[-2] for img in images)
    total_width = sum(img.shape[-1] for img in images)
    canvas = torch.zeros((3, max_height, total_width), device=images[0].device)
    
    current_x = 0
    for img in images:
        h, w = img.shape[-2:]
        canvas[:, :h, current_x:current_x+w] = img * 0.5 + 0.5
        current_x += w
    
    return to_pil_image(canvas)