content: from typing import Dict, List, Optional, Tuple, Union
  
  import numpy as np
  import torch
  import torch.nn as nn
  from diffusers.configuration_utils import ConfigMixin, register_to_config
  from diffusers.models.attention_processor import Attention, AttentionProcessor
  from diffusers.models.autoencoders.vae import DecoderOutput
  from diffusers.models.modeling_outputs import AutoencoderKLOutput
  from diffusers.models.modeling_utils import ModelMixin
  from diffusers.models.normalization import FP32LayerNorm, LayerNorm
  from diffusers.utils import logging
  from diffusers.utils.accelerate_utils import apply_forward_hook
  from einops import repeat
  from tqdm import tqdm
  
  # Check torch_cluster availability
  try:
      from torch_cluster import fps
      TORCH_CLUSTER_AVAILABLE = True
      print("torch_cluster available - using original FPS sampling method")
  except ImportError:
      TORCH_CLUSTER_AVAILABLE = False
      print("torch_cluster not available - using pytorch3d FPS sampling method")
      fps = None
  
  from ..attention_processor import FusedTripoSGAttnProcessor2_0, TripoSGAttnProcessor2_0, FlashTripo2AttnProcessor2_0
  from ..embeddings import FrequencyPositionalEmbedding
  from ..transformers.partcrafter_transformer import DiTBlock
  from .vae import DiagonalGaussianDistribution
  from pytorch3d.ops import sample_farthest_points
  
  
  logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
  
  
  class TripoSGEncoder(nn.Module):
      def __init__(
          self,
          in_channels: int = 3,
          dim: int = 512,
          num_attention_heads: int = 8,
          num_layers: int = 8,
      ):
          super().__init__()
  
          self.proj_in = nn.Linear(in_channels, dim, bias=True)
  
          self.blocks = nn.ModuleList(
              [
                  DiTBlock(
                      dim=dim,
                      num_attention_heads=num_attention_heads,
                      use_self_attention=False,
                      use_cross_attention=True,
                      cross_attention_dim=dim,
                      cross_attention_norm_type="layer_norm",
                      activation_fn="gelu",
                      norm_type="fp32_layer_norm",
                      norm_eps=1e-5,
                      qk_norm=False,
                      qkv_bias=False,
                  )  # cross attention
              ]
              + [
                  DiTBlock(
                      dim=dim,
                      num_attention_heads=num_attention_heads,
                      use_self_attention=True,
                      self_attention_norm_type="fp32_layer_norm",
                      use_cross_attention=False,
                      activation_fn="gelu",
                      norm_type="fp32_layer_norm",
                      norm_eps=1e-5,
                      qk_norm=False,
                      qkv_bias=False,
                  )
                  for _ in range(num_layers)  # self attention
              ]
          )
  
          self.norm_out = LayerNorm(dim)
      
      
  
  
      def forward(self, sample_1: torch.Tensor, sample_2: torch.Tensor):
          hidden_states = self.proj_in(sample_1)
          encoder_hidden_states = self.proj_in(sample_2)
  
          for layer, block in enumerate(self.blocks):
              if layer == 0:
                  hidden_states = block(
                      hidden_states, encoder_hidden_states=encoder_hidden_states
                  )
              else:
                  hidden_states = block(hidden_states)
  
          hidden_states = self.norm_out(hidden_states)
  
          return hidden_states
  
  
  class TripoSGDecoder(nn.Module):
      def __init__(
          self,
          in_channels: int = 3,
          out_channels: int = 1,
          dim: int = 512,
          num_attention_heads: int = 8,
          num_layers: int = 16,
          grad_type: str = "analytical",
          grad_interval: float = 0.001,
      ):
          super().__init__()
  
          if grad_type not in ["numerical", "analytical"]:
              raise ValueError(f"grad_type must be one of ['numerical', 'analytical']")
          self.grad_type = grad_type
          self.grad_interval = grad_interval
  
          self.blocks = nn.ModuleList(
              [
                  DiTBlock(
                      dim=dim,
                      num_attention_heads=num_attention_heads,
                      use_self_attention=True,
                      self_attention_norm_type="fp32_layer_norm",
                      use_cross_attention=False,
                      activation_fn="gelu",
                      norm_type="fp32_layer_norm",
                      norm_eps=1e-5,
                      qk_norm=False,
                      qkv_bias=False,
                  )
                  for _ in range(num_layers)  # self attention
              ]
              + [
                  DiTBlock(
                      dim=dim,
                      num_attention_heads=num_attention_heads,
                      use_self_attention=False,
                      use_cross_attention=True,
                      cross_attention_dim=dim,
                      cross_attention_norm_type="layer_norm",
                      activation_fn="gelu",
                      norm_type="fp32_layer_norm",
                      norm_eps=1e-5,
                      qk_norm=False,
                      qkv_bias=False,
                  )  # cross attention
              ]
          )
  
          self.proj_query = nn.Linear(in_channels, dim, bias=True)
  
          self.norm_out = LayerNorm(dim)
          self.proj_out = nn.Linear(dim, out_channels, bias=True)
  
      def set_topk(self, topk):
          self.blocks[-1].set_topk(topk)
  
      def set_flash_processor(self, processor):
          self.blocks[-1].set_flash_processor(processor)
  
      def query_geometry(
          self,
          model_fn: callable,
          queries: torch.Tensor,
          sample: torch.Tensor,
          grad: bool = False,
      ):
          logits = model_fn(queries, sample)
          if grad:
              with torch.autocast(device_type="cuda", dtype=torch.float32):
                  if self.grad_type == "numerical":
                      interval = self.grad_interval
                      grad_value = []
                      for offset in [
                          (interval, 0, 0),
                          (0, interval, 0),
                          (0, 0, interval),
                      ]:
                          offset_tensor = torch.tensor(offset, device=queries.device)[
                              None, :
                          ]
                          res_p = model_fn(queries + offset_tensor, sample)[..., 0]
                          res_n = model_fn(queries - offset_tensor, sample)[..., 0]
                          grad_value.append((res_p - res_n) / (2 * interval))
                      grad_value = torch.stack(grad_value, dim=-1)
                  else:
                      queries_d = torch.clone(queries)
                      queries_d.requires_grad = True
                      with torch.enable_grad():
                          res_d = model_fn(queries_d, sample)
                          grad_value = torch.autograd.grad(
                              res_d,
                              [queries_d],
                              grad_outputs=torch.ones_like(res_d),
                              create_graph=self.training,
                          )[0]
          else:
              grad_value = None
  
          return logits, grad_value
  
      def forward(
          self,
          sample: torch.Tensor,
          queries: torch.Tensor,
          kv_cache: Optional[torch.Tensor] = None,
      ):
          if kv_cache is None:
              hidden_states = sample
              for _, block in enumerate(self.blocks[:-1]):
                  hidden_states = block(hidden_states)
              kv_cache = hidden_states
  
          # query grid logits by cross attention
          def query_fn(q, kv):
              q = self.proj_query(q)
              l = self.blocks[-1](q, encoder_hidden_states=kv)
              return self.proj_out(self.norm_out(l))
  
          logits, grad = self.query_geometry(
              query_fn, queries, kv_cache, grad=self.training
          )
          logits = logits * -1 if not isinstance(logits, Tuple) else logits[0] * -1
  
          return logits, kv_cache
  
  
  class TripoSGVAEModel(ModelMixin, ConfigMixin):
      @register_to_config
      def __init__(
          self,
          in_channels: int = 3,  # NOTE xyz instead of feature dim
          latent_channels: int = 64,
          num_attention_heads: int = 8,
          width_encoder: int = 512,
          width_decoder: int = 1024,
          num_layers_encoder: int = 8,
          num_layers_decoder: int = 16,
          embedding_type: str = "frequency",
          embed_frequency: int = 8,
          embed_include_pi: bool = False,
      ):
          super().__init__()
  
          self.out_channels = 1
  
          if embedding_type == "frequency":
              self.embedder = FrequencyPositionalEmbedding(
                  num_freqs=embed_frequency,
                  logspace=True,
                  input_dim=in_channels,
                  include_pi=embed_include_pi,
              )
          else:
              raise NotImplementedError(
                  f"Embedding type {embedding_type} is not supported."
              )
  
          self.encoder = TripoSGEncoder(
              in_channels=in_channels + self.embedder.out_dim,
              dim=width_encoder,
              num_attention_heads=num_attention_heads,
              num_layers=num_layers_encoder,
          )
          self.decoder = TripoSGDecoder(
              in_channels=self.embedder.out_dim,
              out_channels=self.out_channels,
              dim=width_decoder,
              num_attention_heads=num_attention_heads,
              num_layers=num_layers_decoder,
          )
  
          self.quant = nn.Linear(width_encoder, latent_channels * 2, bias=True)
          self.post_quant = nn.Linear(latent_channels, width_decoder, bias=True)
  
          self.use_slicing = False
          self.slicing_length = 1
          
          # Auto-select sampling method based on torch_cluster availability
          if TORCH_CLUSTER_AVAILABLE:
              self.sampling_version = 1  # Original method with torch_cluster
              print("Using sampling version 1 (torch_cluster)")
          else:
              self.sampling_version = 2  # Fixed method with pytorch3d
              print("Using sampling version 2 (pytorch3d) - torch_cluster not available")
  
      def set_flash_decoder(self):
          self.decoder.set_flash_processor(FlashTripo2AttnProcessor2_0())
  
      # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedTripoSGAttnProcessor2_0
      def fuse_qkv_projections(self):
          """
          Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
          are fused. For cross-attention modules, key and value projection matrices are fused.
  
          <Tip warning={true}>
  
          This API is 🧪 experimental.
  
          </Tip>
          """
          self.original_attn_processors = None
  
          for _, attn_processor in self.attn_processors.items():
              if "Added" in str(attn_processor.__class__.__name__):
                  raise ValueError(
                      "`fuse_qkv_projections()` is not supported for models having added KV projections."
                  )
  
          self.original_attn_processors = self.attn_processors
  
          for module in self.modules():
              if isinstance(module, Attention):
                  module.fuse_projections(fuse=True)
  
          self.set_attn_processor(FusedTripoSGAttnProcessor2_0())
  
      # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections
      def unfuse_qkv_projections(self):
          """Disables the fused QKV projection if enabled.
  
          <Tip warning={true}>
  
          This API is 🧪 experimental.
  
          </Tip>
  
          """
          if self.original_attn_processors is not None:
              self.set_attn_processor(self.original_attn_processors)
  
      @property
      # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
      def attn_processors(self) -> Dict[str, AttentionProcessor]:
          r"""
          Returns:
              `dict` of attention processors: A dictionary containing all attention processors used in the model with
              indexed by its weight name.
          """
          # set recursively
          processors = {}
  
          def fn_recursive_add_processors(
              name: str,
              module: torch.nn.Module,
              processors: Dict[str, AttentionProcessor],
          ):
              if hasattr(module, "get_processor"):
                  processors[f"{name}.processor"] = module.get_processor()
  
              for sub_name, child in module.named_children():
                  fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
  
              return processors
  
          for name, module in self.named_children():
              fn_recursive_add_processors(name, module, processors)
  
          return processors
  
      # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
      def set_attn_processor(
          self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]
      ):
          r"""
          Sets the attention processor to use to compute attention.
  
          Parameters:
              processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
                  The instantiated processor class or a dictionary of processor classes that will be set as the processor
                  for **all** `Attention` layers.
  
                  If `processor` is a dict, the key needs to define the path to the corresponding cross attention
                  processor. This is strongly recommended when setting trainable attention processors.
  
          """
          count = len(self.attn_processors.keys())
  
          if isinstance(processor, dict) and len(processor) != count:
              raise ValueError(
                  f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
                  f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
              )
  
          def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
              if hasattr(module, "set_processor"):
                  if not isinstance(processor, dict):
                      module.set_processor(processor)
                  else:
                      module.set_processor(processor.pop(f"{name}.processor"))
  
              for sub_name, child in module.named_children():
                  fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
  
          for name, module in self.named_children():
              fn_recursive_attn_processor(name, module, processor)
  
      def set_default_attn_processor(self):
          """
          Disables custom attention processors and sets the default attention implementation.
          """
          self.set_attn_processor(TripoSGAttnProcessor2_0())
  
      def enable_slicing(self, slicing_length: int = 1) -> None:
          r"""
          Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
          compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
          """
          self.use_slicing = True
          self.slicing_length = slicing_length
  
      def disable_slicing(self) -> None:
          r"""
          Disable sliced VAE decoding. If `enable_slicing` was previously enabled, this method will go back to computing
          decoding in one step.
          """
          self.use_slicing = False
  
      def _sample_features(
          self, x: torch.Tensor, num_tokens: int = 2048, seed: Optional[int] = None
      ):
          """
          Sample points from features of the input point cloud.
  
          Args:
              x (torch.Tensor): The input point cloud. shape: (B, N, C)
              num_tokens (int, optional): The number of points to sample. Defaults to 2048.
              seed (Optional[int], optional): The random seed. Defaults to None.
          """
          rng = np.random.default_rng(seed)
          indices = rng.choice(
              x.shape[1], num_tokens * 4, replace=num_tokens * 4 > x.shape[1]
          )
          selected_points = x[:, indices]
  
          batch_size, num_points, num_channels = selected_points.shape
          flattened_points = selected_points.view(batch_size * num_points, num_channels)
          batch_indices = (
              torch.arange(batch_size).to(x.device).repeat_interleave(num_points)
          )
  
          # fps sampling
          sampling_ratio = 1.0 / 4
          
          if not TORCH_CLUSTER_AVAILABLE:
              raise RuntimeError("torch_cluster not available but trying to use original sampling method")
              
          sampled_indices = fps(
              flattened_points[:, :3],
              batch_indices,
              ratio=sampling_ratio,
              random_start=self.training,
          )
          sampled_points = flattened_points[sampled_indices].view(
              batch_size, -1, num_channels
          )
  
          return sampled_points
      
      def _sample_features2(
          self, x: torch.Tensor, num_tokens: int = 2048, seed: Optional[int] = None
      ):
          """
          Sample points from features of the input point cloud.
          FIXED VERSION: Per-batch FPS sampling with proper normalization and no zero padding.
  
          Args:
              x (torch.Tensor): The input point cloud. shape: (B, N, C)
              num_tokens (int, optional): The number of points to sample. Defaults to 2048.
              seed (Optional[int], optional): The random seed. Defaults to None.
          """
          rng = np.random.default_rng(seed)
          indices = rng.choice(
              x.shape[1], num_tokens * 4, replace=num_tokens * 4 > x.shape[1]
          )
          selected_points = x[:, indices]
  
          batch_size, num_points, num_channels = selected_points.shape
          sampling_ratio = 1.0 / 4
          target_points_per_batch = int(num_points * sampling_ratio)
          
          # FIXED: Apply FPS per-batch with coordinate normalization
          sampled_batches = []
          
          for b in range(batch_size):
              batch_points = selected_points[b]  # (num_points, num_channels)
              batch_xyz = batch_points[:, :3]   # (num_points, 3)
              
              # FIXED: Normalize coordinates within each batch to [0,1] for fair FPS
              xyz_min = batch_xyz.min(dim=0, keepdim=True)[0]
              xyz_max = batch_xyz.max(dim=0, keepdim=True)[0]
              xyz_range = xyz_max - xyz_min
              # Avoid division by zero
              xyz_range = torch.where(xyz_range < 1e-6, torch.ones_like(xyz_range), xyz_range)
              normalized_xyz = (batch_xyz - xyz_min) / xyz_range
              
              # Apply FPS on normalized coordinates
              normalized_xyz_batch = normalized_xyz.unsqueeze(0)  # (1, num_points, 3)
              _, sampled_indices = sample_farthest_points(normalized_xyz_batch, K=target_points_per_batch)
              sampled_indices = sampled_indices.squeeze(0)  # (target_points_per_batch,)
              
              # Get original (unnormalized) sampled points with all features
              sampled_batch = batch_points[sampled_indices]  # (target_points_per_batch, num_channels)
              sampled_batches.append(sampled_batch)
          
          # FIXED: Stack without zero padding - all batches have same number of points
          sampled_points = torch.stack(sampled_batches, dim=0)  # (batch_size, target_points_per_batch, num_channels)
  
          return sampled_points
  
      def set_sampling_version(self, version: int):
          """Switch between sampling versions for debugging"""
          if version not in [1, 2]:
              raise ValueError("Version must be 1 (original) or 2 (fixed)")
          
          # Check if version 1 is requested but torch_cluster is not available
          if version == 1 and not TORCH_CLUSTER_AVAILABLE:
              print("Warning: Cannot use version 1 - torch_cluster not available, using version 2")
              version = 2
              
          self.sampling_version = version
          print(f"Switched to sampling version {version}")
  
      def compare_sampling_versions(self, x: torch.Tensor, num_tokens: int = 2048, seed: Optional[int] = None):
          """Compare both sampling versions side by side"""
          print("COMPARING SAMPLING VERSIONS:")
          print("="*60)
          
          # Test version 1 (original) only if torch_cluster is available
          if TORCH_CLUSTER_AVAILABLE:
              print("Testing ORIGINAL version 1...")
              with torch.no_grad():
                  sampled_v1 = self._sample_features(x, num_tokens, seed)
          else:
              print("Skipping version 1 - torch_cluster not available")
              sampled_v1 = None
              
          # Test version 2 (fixed)  
          print("Testing FIXED version 2...")
          with torch.no_grad():
              sampled_v2 = self._sample_features2(x, num_tokens, seed)
              
          print(f"\nRESULTS COMPARISON:")
          if sampled_v1 is not None:
              print(f"  V1 shape: {sampled_v1.shape}")
          print(f"  V2 shape: {sampled_v2.shape}")
          
          # Check for zero padding
          versions_to_check = [(2, sampled_v2)]
          if sampled_v1 is not None:
              versions_to_check.insert(0, (1, sampled_v1))
              
          for version, sampled in versions_to_check:
              print(f"\n  Version {version} - Zero padding analysis:")
              for b in range(sampled.shape[0]):
                  real_points = (sampled[b].abs().sum(-1) > 1e-6).sum().item()
                  total_points = sampled[b].shape[0] 
                  zero_points = total_points - real_points
                  print(f"    Batch {b}: {real_points}/{total_points} real points ({zero_points} zeros)")
          
          # Coordinate range comparison
          print(f"\n  Coordinate ranges:")
          for version, sampled in versions_to_check:
              xyz = sampled[:, :, :3]
              real_mask = xyz.abs().sum(-1) > 1e-6
              if real_mask.any():
                  real_xyz = xyz[real_mask]
                  print(f"    V{version}: [{real_xyz.min():.3f}, {real_xyz.max():.3f}] range={real_xyz.max()-real_xyz.min():.3f}")
          
          print("="*60)
          return sampled_v1, sampled_v2
  
      def _encode(
          self, x: torch.Tensor, num_tokens: int = 2048, seed: Optional[int] = None
      ):
          position_channels = self.config.in_channels
          positions, features = x[..., :position_channels], x[..., position_channels:]
          x_kv = torch.cat([self.embedder(positions), features], dim=-1)
  
          # Switch between sampling versions
          if self.sampling_version == 2:
              print(f"Using FIXED sampling version 2")
              sampled_x = self._sample_features2(x, num_tokens, seed)
          else:
              print(f"Using ORIGINAL sampling version 1")
              sampled_x = self._sample_features(x, num_tokens, seed)
              
          positions, features = (
              sampled_x[..., :position_channels],
              sampled_x[..., position_channels:],
          )
          x_q = torch.cat([self.embedder(positions), features], dim=-1)
  
          x = self.encoder(x_q, x_kv)
  
          x = self.quant(x)
  
          return x
  
      @apply_forward_hook
      def encode(
          self, x: torch.Tensor, return_dict: bool = True, **kwargs
      ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
          """
          Encode a batch of point features into latents.
          """
          if self.use_slicing and x.shape[0] > 1:
              encoded_slices = [
                  self._encode(x_slice, **kwargs)
                  for x_slice in x.split(self.slicing_length)
              ]
              h = torch.cat(encoded_slices)
          else:
              h = self._encode(x, **kwargs)
  
          posterior = DiagonalGaussianDistribution(h, feature_dim=-1)
  
          if not return_dict:
              return (posterior,)
          return AutoencoderKLOutput(latent_dist=posterior)
  
      def _decode(
          self,
          z: torch.Tensor,
          sampled_points: torch.Tensor,
          num_chunks: int = 50000,
          to_cpu: bool = False,
          return_dict: bool = True,
      ) -> Union[DecoderOutput, torch.Tensor]:
          xyz_samples = sampled_points
  
          z = self.post_quant(z)
  
          num_points = xyz_samples.shape[1]
          kv_cache = None
          dec = []
  
          for i in range(0, num_points, num_chunks):
              queries = xyz_samples[:, i : i + num_chunks, :].to(z.device, dtype=z.dtype)
              queries = self.embedder(queries)
  
              z_, kv_cache = self.decoder(z, queries, kv_cache)
              dec.append(z_ if not to_cpu else z_.cpu())
  
          z = torch.cat(dec, dim=1)
  
          if not return_dict:
              return (z,)
  
          return DecoderOutput(sample=z)
  
      @apply_forward_hook
      def decode(
          self,
          z: torch.Tensor,
          sampled_points: torch.Tensor,
          return_dict: bool = True,
          **kwargs,
      ) -> Union[DecoderOutput, torch.Tensor]:
          if self.use_slicing and z.shape[0] > 1:
              decoded_slices = [
                  self._decode(z_slice, p_slice, **kwargs).sample
                  for z_slice, p_slice in zip(
                      z.split(self.slicing_length),
                      sampled_points.split(self.slicing_length),
                  )
              ]
              decoded = torch.cat(decoded_slices)
          else:
              decoded = self._decode(z, sampled_points, **kwargs).sample
  
          if not return_dict:
              return (decoded,)
          return DecoderOutput(sample=decoded)
  
      def forward(self, x: torch.Tensor):
          pass
  