Update README.md

README.md

@@ -35,17 +35,19 @@ def crop_image_to_nearest_divisible_by_8(img):
         new_width = img.shape[2] - (img.shape[2] % 8)
         
         # Use CenterCrop to crop the image
-        transform = transforms.Compose([transforms.CenterCrop((new_height, new_width)) , transforms.ToTensor])
+        transform = transforms.CenterCrop(new_height, new_width)
         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(image).unsqueeze(0)
+image = crop_image_to_nearest_divisible_by_8(to_tensor(image)).unsqueeze(0)
 
 upscaled_image = vae(image).sample
 # Save the reconstructed image