s

dkt/pipelines/pipeline.py

@@ -983,7 +983,17 @@ class DKTPipeline:
 
 
         # 5. 放到 GPU
-        input_image = input_image.to(device=device, dtype=torch.bfloat16)
+        input_image = input_image.to(device=device, dtype=torch.float32)
+
+
+
+
+        # 🔴 必须补 batch 维度
+        if input_image.dim() == 3:
+            input_image = input_image.unsqueeze(0)  # (1, 3, H, W)
+        
+
+    
 
 
         model_device = next(self.moge_pipe.parameters()).device
@@ -1013,9 +1023,9 @@ class DKTPipeline:
             if iidx == 0:
                 # Read the input image and convert to tensor (3, H, W) with RGB values normalized to [0, 1]
                 if input_image_np.max() > 1:
-                    input_image = torch.tensor(input_image_np / 255, dtype=torch.bfloat16, device=moge_device).permute(2, 0, 1) 
+                    input_image = torch.tensor(input_image_np / 255, dtype=torch.float32, device=moge_device).permute(2, 0, 1) 
                 else:
-                    input_image = torch.tensor(input_image_np, dtype=torch.bfloat16, device=moge_device).permute(2, 0, 1) 
+                    input_image = torch.tensor(input_image_np, dtype=torch.float32, device=moge_device).permute(2, 0, 1) 
 
                 print(f'moge devices: {moge_device}') #* why cpu?