s

dkt/pipelines/pipeline.py

@@ -991,7 +991,11 @@ class DKTPipeline:
 
             if iidx == 0:
                 # Read the input image and convert to tensor (3, H, W) with RGB values normalized to [0, 1]
-                input_image = torch.tensor(input_image_np / 255, dtype=torch.float32, device=moge_device).permute(2, 0, 1) 
+                if input_image_np.max() > 1:
+                    input_image = torch.tensor(input_image_np / 255, dtype=torch.float16, device=moge_device).permute(2, 0, 1) 
+                else:
+                    input_image = torch.tensor(input_image_np, dtype=torch.float16, device=moge_device).permute(2, 0, 1) 
+
                 print(f'moge devices: {moge_device}')
                 
                 output = self.moge_pipe.infer(input_image)