content: # trellis_model_manager.py
  import os
  import logging
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from safetensors.torch import load_file
  import folder_paths
  from huggingface_hub import hf_hub_download, snapshot_download
  from typing import Dict, Union
  import json
  import importlib  # Import the importlib module
  from trellis.modules.utils import convert_module_to_f16, convert_module_to_f32
  
  logger = logging.getLogger('model_manager')
  
  __attributes = {
      'SparseStructureDecoder': 'trellis.models.sparse_structure_vae',
      'SparseStructureFlowModel': 'trellis.models.sparse_structure_flow',
      'SLatFlowModel': 'trellis.models.structured_latent_flow',
  }
  
  __all__ = list(__attributes.keys())
  
  def __getattr__(name):
      if name in __attributes:
          module_name = __attributes[name]
          module = importlib.import_module(module_name, package=None) # Import the module
          return getattr(module, name)
      raise AttributeError(f"module {__name__} has no attribute {name}")
  
  class TrellisModelManager:
      """
      Basic manager for Trellis models, using ComfyUI's new model path.
      """
      def __init__(self, model_dir: str, config=None, device: str = "cuda"):
          """
          Initialize the model manager with a specific model directory.
          
          Args:
              model_dir (str): Path to model directory (e.g. "models/checkpoints/TRELLIS-image-large")
              config (dict or object): Global configuration for Trellis
              device (str): Device to load models on (e.g. "cuda")
          """
          self.model_dir = model_dir
          if config is None:
              self.device = device
          elif isinstance(config, dict):
              self.device = config.get('device', device)
              self.config = config
          else:
              self.device = getattr(config, 'device', device)
              self.config = config
          self.model = None
          self.dinov2_model = None
          
      def load(self) -> None:
          """Load model configuration and checkpoints"""
          try:
              os.makedirs(self.model_dir, exist_ok=True)
              ckpts_folder = os.path.join(self.model_dir, "ckpts")
              os.makedirs(ckpts_folder, exist_ok=True)
              
              if not os.path.exists(os.path.join(self.model_dir, "pipeline.json")):
                  logger.info("Downloading TRELLIS models...")
                  try:
                      snapshot_download(
                          repo_id="Stable-X/trellis-normal-v0-1",
                          local_dir=self.model_dir,
                          local_dir_use_symlinks=False,
                          allow_patterns=["pipeline.json", "README.md"]
                      )
                      snapshot_download(
                          repo_id="Stable-X/trellis-normal-v0-1",
                          local_dir=ckpts_folder,
                          local_dir_use_symlinks=False,
                          allow_patterns=["*.safetensors", "*.json"],
                          cache_dir=os.path.join(self.model_dir, ".cache")
                      )
                      logger.info("Model files downloaded successfully")
                  except Exception as e:
                      logger.error(f"Error downloading model files: {str(e)}")
                      raise
              
              self.config = self._load_config()
              
          except Exception as e:
              logger.error(f"Error in load(): {str(e)}")
              raise
  
      def get_checkpoint_path(self, filename: str) -> str:
          """
          Returns the full path to a checkpoint file.
          """
          ckpts_folder = os.path.join(self.model_dir, "ckpts")
          if not filename.endswith('.safetensors'):
              filename = f"{filename}.safetensors"
          full_path = os.path.join(ckpts_folder, filename)
          if not os.path.exists(full_path):
              raise FileNotFoundError(f"Checkpoint file not found: {full_path}")
          return full_path
  
      def _load_config(self) -> Dict:
          """Load model configuration from pipeline.json"""
          try:
              config_path = os.path.join(self.model_dir, "pipeline.json")
              
              if os.path.exists(config_path):
                  logger.info(f"Loading config from {config_path}")
                  with open(config_path, 'r') as f:
                      config = json.load(f)
              else:
                  logger.info(f"Config not found locally, downloading from HuggingFace")
                  config_path = hf_hub_download(
                      repo_id=f"JeffreyXiang/{os.path.basename(self.model_dir)}", 
                      filename="pipeline.json",
                      cache_dir=os.path.join(self.model_dir, ".cache")
                  )
                  with open(config_path, 'r') as f:
                      config = json.load(f)
              
              logger.info("Raw config contents:")
              logger.info(json.dumps(config, indent=2))
              
              if not config:
                  raise ValueError(f"Could not load valid configuration from {self.model_dir}")
                  
              if 'name' not in config:
                  config['name'] = 'TrellisImageTo3DPipeline'
                  
              return config
              
          except Exception as e:
              logger.error(f"Error loading config from {self.model_dir}: {e}")
              return {
                  'name': 'TrellisImageTo3DPipeline',
                  'version': '1.0'
              }
  
      def load_models(self) -> Dict[str, nn.Module]:
          """Load all required models with current configuration"""
          return {
              'sparse_structure_flow_model': self.get_checkpoint_path("ss_flow_img_dit_L_16l8_fp16"),
              'slat_flow_model': self.get_checkpoint_path("slat_flow_img_dit_L_64l8p2_fp16")
          }
  
      def load_model_components(self) -> Dict[str, nn.Module]:
          """Loads individual model components."""
          models = {}
          model_paths = self.load_models()
          for name, path in model_paths.items():
              models[name] = models.from_pretrained(path, config=self.config)
              
              # Ensure each model is converted to the desired precision
              if self.config.get('use_fp16', True):
                  convert_module_to_f16(models[name])
              else:
                  convert_module_to_f32(models[name])
      
          # DINOv2 is handled separately
          # models['image_cond_model'] = self.load_dinov2(self.config.get("dinov2_model", "dinov2_vitl14"))
          
          return models
  
      def load_dinov2(self, model_name: str):
          """Load DINOv2 model from local repo and weights"""
          try:
              local_repo_path = self._ensure_dinov2_repo()
              weights_path = self._ensure_dinov2_weights(model_name)
              
              use_fp16 = (self.config.get('use_fp16', True) 
                      if isinstance(self.config, dict) 
                      else getattr(self.config, 'use_fp16', True))
              
              model = torch.hub.load(
                  local_repo_path,
                  model_name,
                  source='local',
                  pretrained=False,  
                  force_reload=False,
                  trust_repo=True
              )
              
              state_dict = torch.load(weights_path, map_location='cpu')
              
              if 'model' in state_dict:
                  state_dict = state_dict['model']
              elif 'state_dict' in state_dict:
                  state_dict = state_dict['state_dict']
              
              model.load_state_dict(state_dict, strict=False)
              
              model = model.to(self.device)
              if use_fp16:
                  model = model.half()
              
              model.eval()
              logger.info(f"DINOv2 model {model_name} loaded successfully")
              return model
                  
          except Exception as e:
              raise RuntimeError(f"Error loading DINOv2 model: {str(e)}")
  
      def _ensure_dinov2_repo(self):
          repo_path = os.path.join(os.path.dirname(__file__), 'dinov2')
          
          if not os.path.exists(repo_path):
              logger.error("DINOv2 repository not found. Please ensure dinov2/ folder exists with hubconf.py")
              raise FileNotFoundError(f"DINOv2 repo not found at {repo_path}")
          
          return repo_path
  
      def _ensure_dinov2_weights(self, model_name: str):
          weights_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 
                                    'Checkpoints', 'facebookresearch', 'dinov2')
          os.makedirs(weights_dir, exist_ok=True)
          
          weights_file = f"{model_name}.pth"
          weights_path = os.path.join(weights_dir, weights_file)
          
          if not os.path.exists(weights_path):
              logger.info(f"Downloading {model_name} weights...")
              
              success = self._download_from_facebook(model_name, weights_path)
              
              if not success:
                  logger.info("Trying HuggingFace...")
                  self._download_from_huggingface(model_name, weights_path)
                  logger.info("Weights downloaded via HuggingFace")
          
          return weights_path
      
      def _download_from_facebook(self, model_name: str, weights_path: str) -> bool:
          try:
              import urllib.request
              import shutil
              
              facebook_urls = {
                  'dinov2_vitl14_reg': 'https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_reg4_pretrain.pth'
              }
              
              if model_name not in facebook_urls:
                  return False
              
              url = facebook_urls[model_name]
              logger.info(f"Downloading from Facebook...")
              
              temp_path = weights_path + '.tmp'
              urllib.request.urlretrieve(url, temp_path)
              
              shutil.move(temp_path, weights_path)
              
              return True
              
          except Exception as e:
              logger.warning(f"Facebook download failed: {str(e)}")
              temp_path = weights_path + '.tmp'
              if os.path.exists(temp_path):
                  os.remove(temp_path)
              return False
      
      def _download_from_huggingface(self, model_name: str, weights_path: str):
          """Download DINOv2 weights from HuggingFace as fallback"""
          try:
              weights_dir = os.path.dirname(weights_path)
              
              model_mapping = {
                  'dinov2_vitl14_reg': ('DenisKochetov/dinov2_vitl14_reg', 'dinov2_vitl14_reg.pth')
              }
              
              if model_name not in model_mapping:
                  raise ValueError(f"Unknown model: {model_name}. Available: {list(model_mapping.keys())}")
              
              repo_id, filename = model_mapping[model_name]
              
              downloaded_path = hf_hub_download(
                  repo_id=repo_id,
                  filename=filename,
                  local_dir=weights_dir,
                  local_dir_use_symlinks=False,
                  resume_download=True
              )
              
              weights_file = os.path.basename(weights_path)
              if filename != weights_file:
                  final_path = weights_path
                  if os.path.exists(downloaded_path) and downloaded_path != final_path:
                      import shutil
                      shutil.move(downloaded_path, final_path)
              
          except Exception as e:
              raise RuntimeError(f"Failed to download weights for {model_name}: {str(e)}")