Heasterian/

AsymmetricAutoencoderKLUpscaler

import torch

from torchvision import transforms, utils

import diffusers
from diffusers import AsymmetricAutoencoderKL

from diffusers.utils import load_image

def crop_image_to_nearest_divisible_by_8(img):
    # Check if the image height and width are divisible by 8
    if img.shape[1] % 8 == 0 and img.shape[2] % 8 == 0:
        return img
    else:
        # Calculate the closest lower resolution divisible by 8
        new_height = img.shape[1] - (img.shape[1] % 8)
        new_width = img.shape[2] - (img.shape[2] % 8)
        
        # Use CenterCrop to crop the image
        transform = transforms.CenterCrop((new_height, new_width), interpolation=transforms.InterpolationMode.BILINEAR)
        img = transform(img).to(torch.float32).clamp(-1, 1)
        
        return img
        
to_tensor = transforms.ToTensor()

vae = AsymmetricAutoencoderKL.from_pretrained("Heasterian/AsymmetricAutoencoderKLUpscaler", weight_dtype=torch.float32)
vae.requires_grad_(False)

image = load_image(r"/home/heasterian/test/a/F8VlGmCWEAAUVpc (copy).jpeg")

image = crop_image_to_nearest_divisible_by_8(to_tensor(image)).unsqueeze(0)

upscaled_image = vae(image).sample
# Save the reconstructed image
utils.save_image(upscaled_image, "test.png")