fixing bugs

requirements.txt

@@ -25,7 +25,7 @@ opt-einsum==3.3.0
 optax==0.2.0
 orbax-checkpoint==0.2.4
 peft
-pillow
+Pillow==10.0.0
 scipy==1.10.1
 timm==0.9.6
 torch

utils.py

@@ -4,7 +4,9 @@ import jax
 import numpy as np
 
 
-def repeat_vmap(fun, in_axes=[0]):
+def repeat_vmap(fun, in_axes=None):
+    if in_axes is None:
+        in_axes = [0]
     for axes in in_axes:
         fun = jax.vmap(fun, in_axes=axes)
     return fun
@@ -21,16 +23,37 @@ def make_grid(patch_size: int | tuple[int, int]):
 
 def interpolate_grid(coords, grid, order=0):
     """
-    args:
-        coords: Tensor of shape (B, H, W, 2) with coordinates in [-0.5, 0.5]
-        grid: Tensor of shape (B, H', W', C)
-    returns:
-        Tensor of shape (B, H, W, C) with interpolated values
+    Args:
+        coords: Tensor de shape (B, H, W, 2) ou (H, W, 2)
+        grid: Tensor de shape (B, H', W', C)
     """
-    # convert [-0.5, 0.5] -> [0, size], where pixel centers are expected at
-    # [-0.5 + 1 / (2*size), ..., 0.5 - 1 / (2*size)]
-    coords = coords.transpose((0, 3, 1, 2))
+    # Adicionar dimensão de batch se necessário
+    if coords.ndim == 3:
+        coords = coords[np.newaxis, ...]
+
+    # Verificar dimensões
+    assert coords.ndim == 4, f"Dimensões inválidas para coords: {coords.shape}"
+    assert grid.ndim == 4, f"Dimensões inválidas para grid: {grid.shape}"
+
+    # Ajustar transposição de forma segura
+    try:
+        coords = coords.transpose((0, 3, 1, 2))
+    except ValueError as e:
+        raise ValueError(f"Falha na transposição: {coords.shape} → (0,3,1,2)") from e
+
+    # Conversão de coordenadas
     coords = coords.at[:, 0].set(coords[:, 0] * grid.shape[-3] + (grid.shape[-3] - 1) / 2)
     coords = coords.at[:, 1].set(coords[:, 1] * grid.shape[-2] + (grid.shape[-2] - 1) / 2)
-    map_coordinates = partial(jax.scipy.ndimage.map_coordinates, order=order, mode='nearest')
-    return jax.vmap(jax.vmap(map_coordinates, in_axes=(2, None), out_axes=2))(grid, coords)
+
+    # Interpolação com JAX
+    map_coordinates = partial(jax.scipy.ndimage.map_coordinates,
+                              order=order,
+                              mode='nearest')
+
+    return jax.vmap(  # Sobre batches
+        jax.vmap(  # Sobre canais
+            map_coordinates,
+            in_axes=(2, None),  # (C, H', W'), (B, 2, H, W)
+            out_axes=2
+        )
+    )(grid, coords)