content: #!/usr/bin/env python3
  
  import sys
  import os
  sys.path.append('.')
  
  import torch
  import numpy as np
  from partcrafter_src.models.autoencoders.autoencoder_kl_triposg import TripoSGVAEModel
  
  def test_both_sampling_versions():
      """Quick test of both sampling versions"""
      
      device = "cuda" if torch.cuda.is_available() else "cpu"
      print(f"🚀 Testing on device: {device}")
      
      # Create model
      model = TripoSGVAEModel(
          in_channels=3,
          latent_channels=64,
          num_attention_heads=8,
          width_encoder=512,
          width_decoder=1024,
          num_layers_encoder=8,
          num_layers_decoder=16,
      ).to(device)
      
      # Create realistic test data (simulating robot point cloud)
      batch_size = 2
      num_points = 8000
      
      # Create test data with different scales (like real scenarios)
      test_data = torch.zeros(batch_size, num_points, 6, device=device)
      
      # Batch 0: centered robot-like object
      test_data[0, :, :3] = torch.randn(num_points, 3, device=device) * 0.8
      test_data[0, :, 3:] = torch.randn(num_points, 3, device=device) * 0.1
      
      # Batch 1: offset robot with different scale  
      test_data[1, :, :3] = torch.randn(num_points, 3, device=device) * 1.5 + torch.tensor([2, 1, 0], device=device)
      test_data[1, :, 3:] = torch.randn(num_points, 3, device=device) * 0.2
      
      print(f"📊 Test data shape: {test_data.shape}")
      print(f"📊 Coordinate ranges:")
      for b in range(batch_size):
          xyz = test_data[b, :, :3]
          print(f"   Batch {b}: [{xyz.min():.2f}, {xyz.max():.2f}], std={xyz.std():.3f}")
      
      # Compare both versions
      sampled_v1, sampled_v2 = model.compare_sampling_versions(test_data, num_tokens=2048, seed=42)
      
      print(f"\n🎯 FINAL RECOMMENDATION:")
      v1_zeros = sum([(sampled_v1[b].abs().sum(-1) < 1e-6).sum().item() for b in range(batch_size)])
      v2_zeros = sum([(sampled_v2[b].abs().sum(-1) < 1e-6).sum().item() for b in range(batch_size)])
      
      if v1_zeros > v2_zeros:
          print(f"✅ Version 2 is BETTER (fewer zero artifacts: {v2_zeros} vs {v1_zeros})")
          recommended_version = 2
      elif v2_zeros > v1_zeros:
          print(f"❌ Version 1 is better (fewer zero artifacts: {v1_zeros} vs {v2_zeros})")
          recommended_version = 1  
      else:
          print(f"⚖️  Both versions have similar zero padding: {v1_zeros}")
          recommended_version = 2  # Default to fixed version
      
      return recommended_version
  
  if __name__ == "__main__":
      recommended = test_both_sampling_versions()
      print(f"\n🎯 Use version {recommended} for your robot inference!") 