Move logs

SUPIR/models/SUPIR_model.py

@@ -47,18 +47,29 @@ class SUPIRModel(DiffusionEngine):
 
     @torch.no_grad()
     def encode_first_stage_with_denoise(self, x, use_sample=True, is_stage1=False):
+        print('encode_first_stage_with_denoise 1')
         with torch.autocast("cuda", dtype=self.ae_dtype):
+            print('encode_first_stage_with_denoise 2')
             if is_stage1:
+                print('encode_first_stage_with_denoise 3')
                 h = self.first_stage_model.denoise_encoder_s1(x)
             else:
+                print('encode_first_stage_with_denoise 4')
                 h = self.first_stage_model.denoise_encoder(x)
+            print('encode_first_stage_with_denoise 5')
             moments = self.first_stage_model.quant_conv(h)
+            print('encode_first_stage_with_denoise 6')
             posterior = DiagonalGaussianDistribution(moments)
+            print('encode_first_stage_with_denoise 7')
             if use_sample:
+                print('encode_first_stage_with_denoise 8')
                 z = posterior.sample()
             else:
+                print('encode_first_stage_with_denoise 9')
                 z = posterior.mode()
+            print('encode_first_stage_with_denoise 10')
         z = self.scale_factor * z
+        print('encode_first_stage_with_denoise 11')
         return z
 
     @torch.no_grad()
@@ -73,9 +84,7 @@ class SUPIRModel(DiffusionEngine):
         '''
         [N, C, H, W], [-1, 1], RGB
         '''
-        print('Start batchify_denoise')
         x = self.encode_first_stage_with_denoise(x, use_sample=False, is_stage1=is_stage1)
-        print('End batchify_denoise')
         return self.decode_first_stage(x)
 
     @torch.no_grad()