import os
from PIL import Image
import numpy as np

batch_size = 4 # set a smaller batch size

images = ["/content/gdrive/MyDrive/castle/"+i for i in os.listdir("/content/gdrive/MyDrive/castle") if i.endswith(('jpeg', 'png', 'jpg','PNG','JPEG','JPG'))]

for i in range(0, len(images), batch_size):
  batch_paths = images[i:i+batch_size]
batch_images = []
for path in batch_paths:
  with Image.open(path) as img:
# Convert image to numpy array and normalize values
    x = np.array(img).astype(np.float32) / 255.0
# Add to batch
batch_images.append(x)
# Stack batch images into a single array
x_batch = np.stack(batch_images, axis=0)
# Call the predict function on the batch
y_batch = predict(x_batch)
# Loop over batch and save images
for j in range(len(batch_paths)):
# Convert output to PIL image and save to disk
    im = Image.fromarray(np.uint8(y_batch[j] * 255.0))
im.save("/content/gdrive/MyDrive/castle/processed/"+os.path.basename(batch_paths[j])).   ---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
<ipython-input-29-f687fc2e729c> in <cell line: 21>()
     19 x_batch = np.stack(batch_images, axis=0)
     20 # Call the predict function on the batch
---> 21 y_batch = predict(x_batch)
     22 # Loop over batch and save images
     23 for j in range(len(batch_paths)):

3 frames
<ipython-input-27-30b5223765f1> in predict(input_img)
    301 def predict(input_img):
    302   # handle multi-stage outputs, obtain the last scale output of last stage
--> 303   return model.apply({'params': flax.core.freeze(params)}, input_img)
    304 
    305 

    [... skipping hidden 6 frame]

/content/maxim/maxim/models/maxim.py in __call__(self, x, train)
    687       x_scales = []
    688       for i in range(self.num_supervision_scales):
--> 689         x_scale = Conv3x3(
    690             (2**i) * self.features,
    691             use_bias=self.use_bias,

    [... skipping hidden 2 frame]

/usr/local/lib/python3.9/dist-packages/flax/linen/linear.py in __call__(self, inputs)
    447     inputs, kernel, bias = promote_dtype(inputs, kernel, bias, dtype=self.dtype)
    448     if self.shared_weights:
--> 449       y = self.conv_general_dilated(
    450           inputs,
    451           kernel,

    [... skipping hidden 7 frame]

/usr/local/lib/python3.9/dist-packages/jax/_src/interpreters/pxla.py in __call__(self, *args)
   1914           results.consume_token())
   1915     else:
-> 1916       results = self.xla_executable.execute_sharded(input_bufs)
   1917     if dispatch.needs_check_special():
   1918       out_arrays = results.disassemble_into_single_device_arrays()

XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 493984768 bytes.
BufferAssignment OOM Debugging.
BufferAssignment stats:
             parameter allocation:   42.19MiB
              constant allocation:         0B
        maybe_live_out allocation:  450.00MiB
     preallocated temp allocation:  471.10MiB
  preallocated temp fragmentation:         0B (0.00%)
                 total allocation:  963.29MiB
              total fragmentation:   21.10MiB (2.19%)
Peak buffers:
	Buffer 1:
		Size: 450.00MiB
		Operator: op_name="jit(conv_general_dilated)/jit(main)/conv_general_dilated[window_strides=(1, 1) padding=((1, 1), (1, 1)) lhs_dilation=(1, 1) rhs_dilation=(1, 1) dimension_numbers=ConvDimensionNumbers(lhs_spec=(0, 3, 1, 2), rhs_spec=(3, 2, 0, 1), out_spec=(0, 3, 1, 2)) feature_group_count=1 batch_group_count=1 precision=None preferred_element_type=None]" source_file="/usr/local/lib/python3.9/dist-packages/flax/linen/linear.py" source_line=449
		XLA Label: custom-call
		Shape: f32[1,32,1920,1920]
		==========================

	Buffer 2:
		Size: 450.00MiB
		Operator: op_name="jit(conv_general_dilated)/jit(main)/conv_general_dilated[window_strides=(1, 1) padding=((1, 1), (1, 1)) lhs_dilation=(1, 1) rhs_dilation=(1, 1) dimension_numbers=ConvDimensionNumbers(lhs_spec=(0, 3, 1, 2), rhs_spec=(3, 2, 0, 1), out_spec=(0, 3, 1, 2)) feature_group_count=1 batch_group_count=1 precision=None preferred_element_type=None]" source_file="/usr/local/lib/python3.9/dist-packages/flax/linen/linear.py" source_line=449
		XLA Label: transpose
		Shape: f32[1,1920,1920,32]
		==========================

	Buffer 3:
		Size: 42.19MiB
		Entry Parameter Subshape: f32[1,1920,1920,3]
		==========================

	Buffer 4:
		Size: 21.10MiB
		Operator: op_name="jit(conv_general_dilated)/jit(main)/conv_general_dilated[window_strides=(1, 1) padding=((1, 1), (1, 1)) lhs_dilation=(1, 1) rhs_dilation=(1, 1) dimension_numbers=ConvDimensionNumbers(lhs_spec=(0, 3, 1, 2), rhs_spec=(3, 2, 0, 1), out_spec=(0, 3, 1, 2)) feature_group_count=1 batch_group_count=1 precision=None preferred_element_type=None]" source_file="/usr/local/lib/python3.9/dist-packages/flax/linen/linear.py" source_line=449
		XLA Label: custom-call
		Shape: u8[22121984]
		==========================

	Buffer 5:
		Size: 3.4KiB
		XLA Label: transpose
		Shape: f32[32,3,3,3]
		==========================

	Buffer 6:
		Size: 3.4KiB
		Entry Parameter Subshape: f32[3,3,3,32]
		==========================

	Buffer 7:
		Size: 16B
		Operator: op_name="jit(conv_general_dilated)/jit(main)/conv_general_dilated[window_strides=(1, 1) padding=((1, 1), (1, 1)) lhs_dilation=(1, 1) rhs_dilation=(1, 1) dimension_numbers=ConvDimensionNumbers(lhs_spec=(0, 3, 1, 2), rhs_spec=(3, 2, 0, 1), out_spec=(0, 3, 1, 2)) feature_group_count=1 batch_group_count=1 precision=None preferred_element_type=None]" source_file="/usr/local/lib/python3.9/dist-packages/flax/linen/linear.py" source_line=449
		XLA Label: custom-call
		Shape: (f32[1,32,1920,1920], u8[22121984])
		==========================