diff --git a/depth_estimation_node.py b/depth_estimation_node.py
index bee5c8c..6e26f15 100644
--- a/depth_estimation_node.py
+++ b/depth_estimation_node.py
@@ -29,7 +29,7 @@ class DepthEstimationNode:
                 "image": ("IMAGE",),
                 "model_name": (list(DEPTH_MODELS.keys()),),
                 "blur_radius": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.1}),
-                "median_size": (cls.MEDIAN_SIZES, {"default": "5"}),  # Added default value
+                "median_size": (cls.MEDIAN_SIZES, {"default": "5"}),
                 "apply_auto_contrast": ("BOOLEAN", {"default": True}),
                 "apply_gamma": ("BOOLEAN", {"default": True})
             }
@@ -59,9 +59,6 @@ class DepthEstimationNode:
         Estimate depth from input image with optional post-processing.
         """
         try:
-            # Ensure median_size is a string
-            median_size = str(median_size)
-            
             # Validate median_size
             if median_size not in self.MEDIAN_SIZES:
                 raise ValueError(f"Invalid median_size. Must be one of {self.MEDIAN_SIZES}")
@@ -69,59 +66,60 @@ class DepthEstimationNode:
             median_size_int = int(median_size)
             self.ensure_model_loaded(model_name)
             
-            # Convert tensor to numpy if needed
+            # Handle tensor conversion
             if torch.is_tensor(image):
-                image = image.cpu().numpy()
+                # Convert tensor to numpy array
+                image_np = image.cpu().numpy()[0]  # Remove batch dimension
+                # Scale to 0-255 range if needed
+                if image_np.max() <= 1.0:
+                    image_np = (image_np * 255).astype(np.uint8)
+                else:
+                    image_np = image_np.astype(np.uint8)
+            else:
+                image_np = image
                 
-            # Ensure image is in range [0, 1]
-            if image.max() > 1.0:
-                image = image / 255.0
+            # Ensure RGB format
+            if image_np.shape[-1] == 4:  # RGBA to RGB
+                image_np = image_np[..., :3]
                 
-            # Convert to RGB if necessary
-            if image.shape[-1] == 4:  # RGBA to RGB
-                image = image[..., :3]
-                
-            # Convert to PIL Image for processing
-            pil_image = Image.fromarray((image[0] * 255).astype(np.uint8))
+            # Convert to PIL for processing
+            pil_image = Image.fromarray(image_np)
             
-            try:
-                # Generate depth map
-                depth_map = self.depth_estimator(pil_image)["predicted_depth"]
+            # Get depth map
+            depth_result = self.depth_estimator(pil_image)
+            depth_map = depth_result["predicted_depth"]
+            
+            # Convert tensor to numpy if necessary
+            if torch.is_tensor(depth_map):
+                depth_map = depth_map.cpu().numpy()
                 
-                # Convert depth map to PIL Image if it's not already
-                if not isinstance(depth_map, Image.Image):
-                    # Normalize depth values to 0-255 range
-                    depth_map = ((depth_map - depth_map.min()) * (255 / (depth_map.max() - depth_map.min()))).astype(np.uint8)
-                    depth_map = Image.fromarray(depth_map)
+            # Normalize depth values to 0-255 range
+            depth_map = ((depth_map - depth_map.min()) * (255 / (depth_map.max() - depth_map.min()))).astype(np.uint8)
+            depth_map = Image.fromarray(depth_map)
+            
+            # Apply post-processing
+            if blur_radius > 0:
+                depth_map = depth_map.filter(ImageFilter.GaussianBlur(radius=blur_radius))
+            
+            if median_size_int > 0:
+                depth_map = depth_map.filter(ImageFilter.MedianFilter(size=median_size_int))
                 
-                # Post-processing pipeline
-                if blur_radius > 0:
-                    depth_map = depth_map.filter(ImageFilter.GaussianBlur(radius=blur_radius))
-                
-                if median_size_int > 0:
-                    depth_map = depth_map.filter(ImageFilter.MedianFilter(size=median_size_int))
-                
-                if apply_auto_contrast:
-                    depth_map = ImageOps.autocontrast(depth_map)
-                
-                if apply_gamma:
-                    depth_array = np.array(depth_map).astype(np.float32) / 255.0
-                    mean_luminance = np.mean(depth_array)
-                    if mean_luminance > 0:
-                        gamma = np.log(0.5) / np.log(mean_luminance)
-                        depth_map = self.gamma_correction(depth_map, gamma)
-                
-                # Convert to numpy array and normalize
+            if apply_auto_contrast:
+                depth_map = ImageOps.autocontrast(depth_map)
+            
+            if apply_gamma:
                 depth_array = np.array(depth_map).astype(np.float32) / 255.0
-                
-                # Add batch and channel dimensions to match ComfyUI format (B,H,W,C)
-                depth_tensor = depth_array[None, ..., None]
-                
-                return (depth_tensor,)
-                
-            except Exception as e:
-                raise RuntimeError(f"Depth estimation failed: {str(e)}")
-
+                mean_luminance = np.mean(depth_array)
+                if mean_luminance > 0:
+                    gamma = np.log(0.5) / np.log(mean_luminance)
+                    depth_map = self.gamma_correction(depth_map, gamma)
+            
+            # Convert back to tensor format
+            depth_array = np.array(depth_map).astype(np.float32) / 255.0
+            depth_tensor = depth_array[None, ..., None]  # Add batch and channel dims
+            
+            return (depth_tensor,)
+            
         except Exception as e:
             raise RuntimeError(f"Depth estimation failed: {str(e)}")