cleaner code

configs/inpainting/lands_config_mountain.yaml

@@ -3,7 +3,7 @@ data:
   seq: {'half': [200, 300], 'box': [300, 350], 'random': [400,500]} #[400,500] #[350, 450], #, 'val': "random" : [350, 450], half : , val: [0,50]
   file_seq: None
   file_name: data/sflckr_all_images.npz 
-  channels: 3
+  channels: 182
   image_size: 512
   latent_size: 128
   latent_channels: 3

vipainting.py

@@ -94,6 +94,10 @@ def vipaint(num, mask_web, image_queue, sampling_queue):
         elif args.k_steps == 4: t_steps_hierarchy = inpaint_config[posterior]['t_steps_hierarchy'] # [550, 500, 450, 400]
         elif args.k_steps == 6: t_steps_hierarchy = [650, 600, 550, 500, 450, 400]
 
+    batch_size = inpaint_config[posterior]["batch_size"]
+    zero_tensor = torch.zeros(batch_size, 182, 512, 512, device=diff.device)
+    uc = diff.get_learned_conditioning({diff.cond_stage_key: zero_tensor}['segmentation']).detach()
+   
         
     # Prepare VI method
     print("=================== Prepare VI method")
@@ -121,45 +125,30 @@ def vipaint(num, mask_web, image_queue, sampling_queue):
             sub_dir = os.path.join(img_path, img_dir)
             os.makedirs(sub_dir, exist_ok=True)
 
-        bs = inpaint_config[posterior]["batch_size"]
-
-        batch_size = bs
-        channels = 182 
-        # For conditional models
-        segmentation = loader.dataset["segmentation"][random_num]
-        if inpaint_config["conditional_model"] : 
-            segment_c = torch.tensor(segmentation.transpose(2,0,1)[None]).to(dtype=torch.float32, device=diff.device)
-            segment_c = segment_c.repeat(batch_size, 1, 1, 1)  
-            uc = diff.get_learned_conditioning(
-                        {diff.cond_stage_key: segment_c.to(diff.device)}['segmentation']
-                        ).detach()
 
         #Get Image/Labels
-        print("==================== get image/labels")
-        #Get Image/Labels
+        print(f"==================== get image/labels")
         if len(loader.dataset) ==2: 
-            ref_img = loader.dataset["images"][random_num] #512, 512, 3
-            ref_img = torch.tensor(ref_img[None]).to(dtype=torch.float32, device=diff.device)
-            print(f"ref_img {ref_img.shape}") #1, 512, 512, 3
+            ref_img = torch.tensor(loader.dataset["images"][random_num][None], dtype=torch.float32, device=diff.device) #1, 512, 512, 3
             ref_img = ref_img/127.5 - 1
-
-            label = torch.tensor(segmentation.transpose(2,0,1)[None]).to(dtype=torch.float32, device=diff.device)
-            save_segmentation(label, img_path, 'input.png')
-            label = label.repeat(batch_size, 1, 1, 1)  # Now shape is [batch_size, 182, 128, 128]
-            xc = torch.tensor(label)
-            c = diff.get_learned_conditioning({diff.cond_stage_key: xc}['segmentation']).detach()
+            segmentation = torch.tensor(dataset["segmentation"][random_num].transpose(2,0,1)[None]).to(dtype=torch.float32, device=diff.device)
+            segmentation_repeated = segmentation.repeat(batch_size, 1, 1, 1)
+            save_segmentation(segmentation, img_path, 'input.png')
+            c = diff.get_learned_conditioning(
+                        {diff.cond_stage_key: segmentation_repeated.to(diff.device)}['segmentation']
+                        ).detach()
+            
         else:
-            ref_img = loader.dataset[random_num].reshape(1,x_size,x_size,channels)
+            ref_img = torch.tensor(loader.dataset[random_num].reshape(1, x_size, x_size, channels), dtype=torch.float32, device=diff.device)
             c = None
             uc = None
         
-        ref_img = torch.tensor(ref_img).to(device)
 
         # #Get mask
         mask_tensor = torch.tensor(mask_web).to(device)
         mask_tensor = mask_tensor.float() / 255.0  # Convert to float and normalize to [0, 1]
         ref_img = torch.permute(ref_img, (0,3,1,2)) 
-        y = torch.Tensor.repeat(mask_tensor*ref_img, [bs,1,1,1]).float()
+        y = torch.Tensor.repeat(mask_tensor*ref_img, [batch_size,1,1,1]).float()
         
         if inpaint_config[posterior]["first_stage"] == "kl":
             y_encoded = encoder_kl(diff, y)[0]
@@ -176,7 +165,7 @@ def vipaint(num, mask_web, image_queue, sampling_queue):
         # Fit posterior once
         print("============ fit posterior once")
         torch.cuda.empty_cache()  
-        h_inpainter.fit(lambda_ = lambda_, cond=c, shape = (bs, *y_encoded.shape[1:]),
+        h_inpainter.fit(lambda_ = lambda_, cond=c, shape = (batch_size, *y_encoded.shape[1:]),
                 quantize_denoised=False, mask_pixel = mask_tensor, y =y,
                 log_every_t=25, iterations = inpaint_config[posterior]['iterations'],
                 unconditional_guidance_scale= inpaint_config[posterior]["unconditional_guidance_scale"] ,
@@ -184,7 +173,7 @@ def vipaint(num, mask_web, image_queue, sampling_queue):
                  kl_weight_2 = inpaint_config[posterior]["beta_2"],
                 debug=True, wdb = False,
                 dir_name = img_path,
-                batch_size = bs, 
+                batch_size = batch_size, 
                 lr_init_gamma = inpaint_config[posterior]["lr_init_gamma"],
                 recon_weight = inpaint_config[posterior]["recon"], 
                 )  
@@ -197,7 +186,7 @@ def vipaint(num, mask_web, image_queue, sampling_queue):
         h_inpainter.sample(inpaint_config["sampling"]["scale"], inpaint_config[posterior]["eta"],
                             mu.cuda(), logvar.cuda(), gamma.cuda(), mask_tensor,  y,
                             n_samples=inpaint_config["sampling"]["n_samples"], 
-                            batch_size = bs, dir_name= img_path, cond=c, 
+                            batch_size = batch_size, dir_name= img_path, cond=c, 
                             unconditional_conditioning=uc, 
                             unconditional_guidance_scale=inpaint_config["sampling"]["unconditional_guidance_scale"], 
                             samples_iteration=inpaint_config[posterior]["iterations"])