Update main.py

main.py

@@ -8,12 +8,24 @@ from pipeline_dedit_sd import DEditSDPipeline
 from utils import load_image, load_mask, load_mask_edit
 from utils_mask import process_mask_move_torch, process_mask_remove_torch, mask_union_torch, mask_substract_torch, create_outer_edge_mask_torch
 from utils_mask import check_mask_overlap_torch, check_cover_all_torch, visualize_mask_list, get_mask_difference_torch, save_mask_list_to_npys
-print(torch.__version__)
-if torch.cuda.is_available():
-    print("CUDA version:", torch.version.cuda)
-    print("GPU model name:", torch.cuda.get_device_name(0))
-else:
-    print('no cuda.')
+
+print("PyTorch 版本:", torch.__version__)
+print("CUDA 版本:", torch.version.cuda)
+print("CUDA 是否可用:", torch.cuda.is_available())
+import subprocess
+
+# 检查 nvidia-smi 命令的输出
+try:
+    result = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    print(result.stdout.decode())
+except FileNotFoundError:
+    print("nvidia-smi 命令不可用，可能未安装 NVIDIA 驱动。")
+
+device_count = torch.cuda.device_count()
+print("可用的 CUDA 设备数量:", device_count)
+for i in range(device_count):
+    print(f"设备 {i} 名称:", torch.cuda.get_device_name(i))
+
 def run_main(
     name="example_tmp",
     name_2=None,