Image to tensor proper format

app.py

@@ -360,6 +360,10 @@ def invert(image):
     s = 1.0 - image
     return s
 
+def image2tensor(image) -> torch.Tensor:
+    image_tensor = torch.from_numpy(np.array(image)).permute(2, 0, 1).float() / 255.0
+    return image_tensor
+
 def tensor2pil(image: torch.Tensor, mode=None):
     if image.ndim == 2:  # Grayscale image
         image = image.unsqueeze(0)  # Add channel dimension
@@ -440,9 +444,9 @@ def generate_ai_bg(input_img, prompt):
 def blend_details(input_image, relit_image, masked_image):
     with torch.inference_mode():
         # Convert images to tensors
-        input_image = pil2tensor(input_image)
-        relit_image = pil2tensor(relit_image)
-        masked_image = pil2tensor(masked_image)
+        input_image = image2tensor(input_image)
+        relit_image = image2tensor(relit_image)
+        masked_image = image2tensor(masked_image)
 
         # Resize input image
         input_image = torch.nn.functional.interpolate(