niulx commited on
Commit
6d51956
1 Parent(s): 99149fd

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +6 -1
main.py CHANGED
@@ -8,7 +8,12 @@ from pipeline_dedit_sd import DEditSDPipeline
8
  from utils import load_image, load_mask, load_mask_edit
9
  from utils_mask import process_mask_move_torch, process_mask_remove_torch, mask_union_torch, mask_substract_torch, create_outer_edge_mask_torch
10
  from utils_mask import check_mask_overlap_torch, check_cover_all_torch, visualize_mask_list, get_mask_difference_torch, save_mask_list_to_npys
11
-
 
 
 
 
 
12
  def run_main(
13
  name="example_tmp",
14
  name_2=None,
 
8
  from utils import load_image, load_mask, load_mask_edit
9
  from utils_mask import process_mask_move_torch, process_mask_remove_torch, mask_union_torch, mask_substract_torch, create_outer_edge_mask_torch
10
  from utils_mask import check_mask_overlap_torch, check_cover_all_torch, visualize_mask_list, get_mask_difference_torch, save_mask_list_to_npys
11
+ print(torch.__version__)
12
+ if torch.cuda.is_available():
13
+ print("CUDA version:", torch.version.cuda)
14
+ print("GPU model name:", torch.cuda.get_device_name(0))
15
+ else:
16
+ print('no cuda.')
17
  def run_main(
18
  name="example_tmp",
19
  name_2=None,