Update utils.py

utils.py

@@ -81,4 +81,62 @@ def get_trainable_module(unet, trainable_module_name):
         raise ValueError(f"Unknown trainable_module_name: {trainable_module_name}")
 
                 
-    
+import torch
+import numpy as np
+from PIL import Image
+
+# =====================================================
+# Image and VAE utility functions used by CatVTONPipeline
+# =====================================================
+
+def compute_vae_encodings(image, vae):
+    """Encode an image tensor using the model's VAE encoder."""
+    if isinstance(image, list):
+        image = torch.cat(image, dim=0)
+    latents = vae.encode(image).latent_dist.sample()
+    latents = latents * vae.config.scaling_factor
+    return latents
+
+
+def numpy_to_pil(images):
+    """Convert numpy arrays to PIL Images."""
+    if images.ndim == 3:
+        images = images[None, ...]
+    images = (images * 255).round().astype("uint8")
+    return [Image.fromarray(image) for image in images]
+
+
+def prepare_image(image):
+    """Convert PIL image to normalized torch tensor."""
+    if isinstance(image, Image.Image):
+        image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
+        image = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0)
+    return image
+
+
+def prepare_mask_image(mask_image):
+    """Convert PIL mask to tensor in [0,1] range."""
+    if isinstance(mask_image, Image.Image):
+        mask_image = np.array(mask_image.convert("L")).astype(np.float32) / 255.0
+        mask_image = torch.from_numpy(mask_image).unsqueeze(0).unsqueeze(0)
+    return mask_image
+
+
+def resize_and_crop(image, size):
+    """Resize image keeping aspect ratio then center crop."""
+    if isinstance(image, Image.Image):
+        image = image.resize(size, Image.BICUBIC)
+    return image
+
+
+def resize_and_padding(image, size):
+    """Resize and pad to match target size."""
+    if isinstance(image, Image.Image):
+        image.thumbnail(size, Image.BICUBIC)
+        new_image = Image.new("RGB", size)
+        left = (size[0] - image.size[0]) // 2
+        top = (size[1] - image.size[1]) // 2
+        new_image.paste(image, (left, top))
+        image = new_image
+    return image
+