Update app.py

app.py

@@ -84,15 +84,31 @@ def preprocess_image(image, target_width, target_height, resize_to_224=False):
     image = torch.from_numpy(image).to(device)
     return image
 
+# def get_depth_map(image, depth_estimator):
+#     image = depth_estimator(image)["depth"]
+#     image = np.array(image)
+#     image = image[:, :, None]
+#     image = np.concatenate([image, image, image], axis=2)
+#     detected_map = torch.from_numpy(image).float() / 255.0
+#     depth_map = detected_map.permute(2, 0, 1)
+#     return depth_map
+
 def get_depth_map(image, depth_estimator):
-    image = depth_estimator(image)["depth"]
-    image = np.array(image)
-    image = image[:, :, None]
-    image = np.concatenate([image, image, image], axis=2)
-    detected_map = torch.from_numpy(image).float() / 255.0
-    depth_map = detected_map.permute(2, 0, 1)
+    # Преобразуем изображение в PIL, если это необходимо
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    elif isinstance(image, torch.Tensor):
+        image = Image.fromarray(image.cpu().numpy())
+    
+    # Получаем карту глубины
+    depth_map = depth_estimator(image)["depth"]
+    depth_map = np.array(depth_map)
+    depth_map = depth_map[:, :, None]  # Добавляем третье измерение
+    depth_map = np.concatenate([depth_map, depth_map, depth_map], axis=2)  # Преобразуем в 3 канала
+    depth_map = torch.from_numpy(depth_map).float() / 255.0  # Нормализация [0, 1]
+    depth_map = depth_map.permute(2, 0, 1)  # Меняем порядок осей (C, H, W)
     return depth_map
-
+  
 pipe_default = get_lora_sd_pipeline(lora_dir='lora_man_animestyle', base_model_name_or_path=model_default, dtype=torch_dtype).to(device)
 
 # ----------------------------------------------------------------------------------------------------------------------------------------------------
@@ -424,9 +440,9 @@ def infer(
 
                         print('control_mode = ', control_mode)
 
-                        # Преобразуем изображения
-                        cn_source_image = preprocess_image(cn_source_image, width, height)
-                        control_image = preprocess_image(control_image, width, height)
+                        # # Преобразуем изображения
+                        # cn_source_image = preprocess_image(cn_source_image, width, height)
+                        # control_image = preprocess_image(control_image, width, height)
 
                         depth_estimator = pipeline("depth-estimation")
                         depth_map = get_depth_map(control_image, depth_estimator).unsqueeze(0).half().to(device)