J

app.py

@@ -43,7 +43,7 @@ reg_model_bilin = register_model2((512,512), 'bilinear')
 def coords_grid_tensor(perturbed_img_shape):
     im_x, im_y = np.mgrid[0:perturbed_img_shape[0]-1:complex(perturbed_img_shape[0]), 0:perturbed_img_shape[1]-1:complex(perturbed_img_shape[1])]
     coords = np.stack((im_y,im_x), axis=2) # 先x后y，行序优先
-    coords = th.from_numpy(coords).float().permute(2,0,1).to(dist_util.dev())  # (2, 512, 512)
+    coords = th.from_numpy(coords).float().permute(2,0,1).to('cuda')  # (2, 512, 512)
     return coords.unsqueeze(0) # [2, 512, 512]
 
 def run_sample_lr_dewarping(
@@ -84,14 +84,7 @@ def run_sample_lr_dewarping(
     sample = th.clamp(sample, min=-1, max=1)
     return sample
 
-def visualize_dewarping(settings, sample, data, i, source_vis, data_path, ref_flow=None):
-    os.makedirs(f'vis_hp/{settings.env.eval_dataset_name}/{settings.name}/dewarped_pred', exist_ok=True) # pred dewarped
-    # warped_src = warp(source_vis.to(sample.device).float(), sample) # [1, 3, 1629, 981]
-    warped_src = reg_model_bilin([source_vis.to(sample.device).float(), sample])
-    warped_src = warped_src[0].permute(1, 2, 0).detach().cpu().numpy()#*255. # (1873, 1353, 3)
-    warped_src = Image.fromarray((warped_src).astype(np.uint8))
 
-    return warped_src
 
 def visualize_dewarping_single(settings, sample, source_vis):
     # os.makedirs(f'vis_hp/{settings.env.eval_dataset_name}/{settings.name}/dewarped_pred', exist_ok=True) # pred dewarped
@@ -105,77 +98,20 @@ def visualize_dewarping_single(settings, sample, source_vis):
 
 
 
-
-
-def prepare_data(settings, batch_preprocessing, SIZE, data):
-    if 'source_image_ori' in data:
-        source_vis = data['source_image_ori']  # B, C, 512, 512 torch.uint8 cpu
-    else:
-        source_vis = data['source_image']
-    if 'target_image' in data:
-        target_vis = data['target_image']
-    else:
-        target_vis = None
-
-    _, _, H_ori, W_ori = source_vis.shape
-
-    source = data['source_image'].to(dist_util.dev())  # [1, 3, 914, 1380]  torch.float32
-    if 'source_image_0' in data:
-        source_0 = data['source_image_0'].to(dist_util.dev())
-    else:
-        source_0 = None
-    if 'target_image' in data:  
-        target = data['target_image']  # [1, 3, 914, 1380]  torch.float32
-    else:
-        target = None 
-    if 'flow_map' in data:    
-        batch_ori = data['flow_map']   # [1, 2, 914, 1380]  torch.float32
-    else:
-        batch_ori = None    
-    if 'flow_map_inter' in data:    
-        batch_ori_inter = data['flow_map_inter']   # [1, 2, 914, 1380]  torch.float32
-    else:
-        batch_ori_inter = None    
-    if target is not None:
-        target = F.interpolate(target, size=512, mode='bilinear', align_corners=False) # [1, 3, 512, 512] 
-        target_256 = data['target_image_256'].to(dist_util.dev()) # [1, 3, 256, 256]
-    else:
-        target = None
-        target_256 = None
-    
-    if settings.env.eval_dataset == 'hp-240':# false
-        source_256 = source
-        target_256 = target
-
-    else: # true
-        data['source_image_256'] = torch.nn.functional.interpolate(input=source.float(), size=(256, 256), mode='area')
-        source_256 = data['source_image_256'].to(dist_util.dev())
-        
-        if 'target_image_256' in data:
-            target_256 = data['target_image_256']
-        else:
-            target_256 = None
-    if 'correspondence_mask' in data:
-        mask = data['correspondence_mask'] # torch.bool [1, 914, 1380]
-    else:
-        mask = torch.ones((1, 512, 512), dtype=torch.bool).to(dist_util.dev()) # None
-
-    return data, H_ori, W_ori, source, target, batch_ori, batch_ori_inter, source_256, target_256, source_vis, target_vis, mask, source_0
-
 def prepare_data_single(input_image, input_image_ori):
     source_vis = input_image_ori
     target_vis = None
     _, _, H_ori, W_ori = source_vis.shape
-    source = input_image.to(dist_util.dev())  # [1, 3, 914, 1380]  torch.float32
+    source = input_image.to('cuda')  # [1, 3, 914, 1380]  torch.float32
     source_0 = None
     target = None 
     batch_ori = None    
     batch_ori_inter = None 
     target = None
     target_256 = None
-    source_256 = torch.nn.functional.interpolate(input=source.float(), size=(256, 256), mode='area').to(dist_util.dev())
+    source_256 = torch.nn.functional.interpolate(input=source.float(), size=(256, 256), mode='area').to('cuda')
     target_256 = None
-    mask = torch.ones((1, 512, 512), dtype=torch.bool).to(dist_util.dev()) # None
+    mask = torch.ones((1, 512, 512), dtype=torch.bool).to('cuda') # None
 
     return input_image, H_ori, W_ori, source, target, batch_ori, batch_ori_inter, source_256, target_256, source_vis, target_vis, mask, source_0
 
@@ -201,16 +137,16 @@ def run_single_docunet(input_image_ori):
 
     os.makedirs(f'vis_hp/{settings.env.eval_dataset_name}/{settings.name}', exist_ok=True)
     batch_preprocessing = None
-    pyramid = VGGPyramid(train=False).to(dist_util.dev())
+    pyramid = VGGPyramid(train=False).to('cuda')
     SIZE = None
 
 
     radius = 4
     raw_corr = None
-    source_288 = F.interpolate(input_image, size=(288), mode='bilinear', align_corners=True).to(dist_util.dev())
+    source_288 = F.interpolate(input_image, size=(288), mode='bilinear', align_corners=True).to('cuda')
 
     if settings.env.time_variant == True:
-        init_feat = torch.zeros((input_image.shape[0], 256, 64, 64), dtype=torch.float32).to(dist_util.dev()) 
+        init_feat = torch.zeros((input_image.shape[0], 256, 64, 64), dtype=torch.float32).to('cuda') 
     else:
         init_feat = None
 
@@ -220,7 +156,7 @@ def run_single_docunet(input_image_ori):
     if settings.env.use_init_flow: 
         init_flow = F.interpolate(ref_flow, size=(64), mode='bilinear', align_corners=True) # [24, 2, 64, 64]                  
     else:
-        init_flow = torch.zeros((input_image.shape[0], 2, 64, 64), dtype=torch.float32).to(dist_util.dev()) 
+        init_flow = torch.zeros((input_image.shape[0], 2, 64, 64), dtype=torch.float32).to('cuda') 
 
     (
         data,
@@ -375,13 +311,13 @@ settings.severity = 0
 settings.corruption_number = 0
 
 
-dist_util.setup_dist()
+# dist_util.setup_dist()
 logger.configure(dir=f"SAMPLING_{settings.env.eval_dataset}_{settings.name}")
 logger.log(f"Corruption Disabled. Evaluating on Original {settings.env.eval_dataset}")
 logger.log("Loading model and diffusion...")
 
 model, diffusion = create_model_and_diffusion(
-    device=dist_util.dev(),
+    device='cuda',
     train_mode=settings.env.train_mode, # stage 1
     tv=settings.env.time_variant,
     **args_to_dict(settings, model_and_diffusion_defaults().keys()),
@@ -393,29 +329,30 @@ pretrained_dewarp_model = GeoTr_Seg_Inf()
 settings.env.seg_model_path = hf_hub_download(repo_id="hanquansanren/DvD", filename="seg.pth", token=token)
 reload_segmodel(pretrained_dewarp_model.msk, settings.env.seg_model_path)
 # reload_model(pretrained_dewarp_model.GeoTr, settings.env.dewarping_model_path)
-pretrained_dewarp_model.to(dist_util.dev())
+pretrained_dewarp_model.to('cuda')
 pretrained_dewarp_model.eval()
 
 if settings.env.use_line_mask:
     pretrained_line_seg_model = UNet(n_channels=3, n_classes=1)
     pretrained_seg_model = Seg()
     settings.env.line_seg_model_path = hf_hub_download(repo_id="hanquansanren/DvD", filename="line_model2.pth", token=token)
-    line_model_ckpt = dist_util.load_state_dict(settings.env.line_seg_model_path, map_location='cpu')['model']
+    line_model_ckpt = pretrained_line_seg_model.load_state_dict(settings.env.line_seg_model_path, map_location='cpu')['model']
     pretrained_line_seg_model.load_state_dict(line_model_ckpt, strict=True)    
-    pretrained_line_seg_model.to(dist_util.dev())
+    pretrained_line_seg_model.to('cuda')
     pretrained_line_seg_model.eval()
 
     settings.env.new_seg_model_path = hf_hub_download(repo_id="hanquansanren/DvD", filename="seg_model.pth", token=token)
-    seg_model_ckpt = dist_util.load_state_dict(settings.env.new_seg_model_path, map_location='cpu')['model']
+    seg_model_ckpt = pretrained_seg_model.load_state_dict(settings.env.new_seg_model_path, map_location='cpu')['model']
     pretrained_seg_model.load_state_dict(seg_model_ckpt, strict=True)    
-    pretrained_seg_model.to(dist_util.dev())
+    pretrained_seg_model.to('cuda')
     pretrained_seg_model.eval()
 
 settings.env.model_path = hf_hub_download(repo_id="hanquansanren/DvD", filename="model1852000.pt", token=token)
-model.cpu().load_state_dict(dist_util.load_state_dict(settings.env.model_path, map_location="cpu"), strict=False)
+# model.cpu().load_state_dict(dist_util.load_state_dict(settings.env.model_path, map_location="cpu"), strict=False)
+model.cpu().load_state_dict(settings.env.model_path, map_location="cpu", strict=False)
 logger.log(f"Model loaded with {settings.env.model_path}")
 
-model.to(dist_util.dev())
+model.to('cuda')
 model.eval()