content: import os
  import json
  import torch
  import shutil
  from accelerate import init_empty_weights
  from accelerate.utils import set_module_tensor_to_device
  from safetensors.torch import load_file
  from diffusers.models import AutoencoderKL, UNet2DConditionModel
  from huggingface_hub import snapshot_download
  
  from .stablex.pipeline_yoso import YosoPipeline
  from .stablex.controlnetvae import ControlNetVAEModel
  from .trellis.pipelines.trellis_image_to_3d import TrellisImageTo3DPipeline as Stable3DGenTrellisImageTo3DPipeline
  from .trellis.backend_config import (
      set_attention_backend,
      set_sparse_backend,
      get_available_backends,
      get_available_sparse_backends
  )
  
  DEVICE_STR = "cuda" if torch.cuda.is_available() else "cpu"
  DEVICE = torch.device(DEVICE_STR)
  HF_DOWNLOAD_IGNORE = ["*.yaml", "*.json", "*.py", ".png", ".jpg", ".gif"]
  
  
  class StableGenPipelineBuilder:
      
      @staticmethod
      def get_available_backends():
          try:
              attn_backends = get_available_backends()
              sparse_backends = get_available_sparse_backends()
          except:
              attn_backends = {'flash_attn': True, 'sdpa': True}
              sparse_backends = {'spconv': True}
  
          available_attn = [k for k, v in attn_backends.items() if v]
          if not available_attn:
              available_attn = ['flash_attn']
  
          available_sparse = [k for k, v in sparse_backends.items() if v]
          if not available_sparse:
              available_sparse = ['spconv']
              
          return available_attn, available_sparse
      
      @staticmethod
      def create_trellis_pipeline_params(
          seed: int = 1234,
          ss_sampling_steps: int = 12,
          ss_guidance_strength: float = 7.5,
          slat_sampling_steps: int = 12,
          slat_guidance_strength: float = 3.0,
          formats: list = None
      ):
          if formats is None:
              formats = ["mesh"]
              
          return {
              "num_samples": 1,
              "seed": seed,
              "formats": formats,
              "preprocess_image": True,
              "sparse_structure_sampler_params": {
                  "steps": ss_sampling_steps,
                  "cfg_strength": ss_guidance_strength,
              },
              "slat_sampler_params": {
                  "steps": slat_sampling_steps,
                  "cfg_strength": slat_guidance_strength,
              }
          }
      
      @staticmethod
      def build_trellis_pipeline(
          repo: str, 
          dinov2_model: str, 
          use_fp16: bool, 
          attn_backend: str,
          sparse_backend: str, 
          spconv_algo: str, 
          smooth_k: bool,
          ckpt_path: str
      ):
          
          StableGenPipelineBuilder._ensure_trellis_weights(repo, ckpt_path)
          
          StableGenPipelineBuilder._setup_environment(attn_backend, sparse_backend, spconv_algo, smooth_k)
  
          model_dir = os.path.join(ckpt_path, "trellis", repo)
  
          try:
              pipe = Stable3DGenTrellisImageTo3DPipeline.from_pretrained(
                  model_dir,
                  dinov2_model=dinov2_model
              )
  
              if pipe is None:
                  raise Exception("Pipeline returned None from from_pretrained")
  
              pipe = StableGenPipelineBuilder._optimize_trellis_pipeline(pipe, use_fp16)
              return pipe
              
          except FileNotFoundError as e:
              if "Weights file not found" in str(e):
                  print(f"Trellis weights missing, re-downloading {repo}...")
                  StableGenPipelineBuilder._ensure_trellis_weights(repo, ckpt_path)
                  
                  try:
                      pipe = Stable3DGenTrellisImageTo3DPipeline.from_pretrained(
                          model_dir,
                          dinov2_model=dinov2_model
                      )
                      
                      if pipe is None:
                          raise Exception("Pipeline returned None from from_pretrained")
  
                      pipe = StableGenPipelineBuilder._optimize_trellis_pipeline(pipe, use_fp16)
                      return pipe
                      
                  except Exception as retry_e:
                      raise Exception(f"Failed to build Trellis pipeline after re-download: {retry_e}")
              else:
                  raise Exception(f"Failed to build Trellis pipeline: {e}")
          except Exception as e:
              raise Exception(f"Failed to build Trellis pipeline: {e}")
      
      @staticmethod
      def build_stablex_pipeline(repo: str, use_fp16: bool, ckpt_path: str):
          
          StableGenPipelineBuilder._ensure_stablex_weights(repo, ckpt_path)
  
          model_dir = os.path.join(ckpt_path, "stablex", repo)
          torch_dtype = torch.float16 if use_fp16 else torch.float32
  
          try:
              config_path = os.path.join(model_dir, 'unet', 'config.json')
              
              if not os.path.exists(config_path):
                  print(f"UNet config not found, re-downloading {repo}...")
                  StableGenPipelineBuilder._ensure_stablex_weights(repo, ckpt_path)
                  
                  if not os.path.exists(config_path):
                      raise FileNotFoundError(f"Config not found at {config_path}")
              
              with open(config_path, 'r', encoding='utf-8') as file:
                  config = json.load(file)
  
              with init_empty_weights():
                  unet = UNet2DConditionModel(**config)
  
              unet_file_candidates = [
                  'diffusion_pytorch_model.fp16.safetensors',
                  'diffusion_pytorch_model.safetensors', 
                  'pytorch_model.fp16.safetensors',
                  'pytorch_model.safetensors'
              ]
              
              unet_sd = None
              unet_ckpt_path = None
              
              for filename in unet_file_candidates:
                  candidate_path = os.path.join(model_dir, 'unet', filename)
                  if os.path.exists(candidate_path):
                      unet_ckpt_path = candidate_path
                      print(f"Found UNet checkpoint: {unet_ckpt_path}")
                      unet_sd = load_file(unet_ckpt_path)
                      break
              
              if unet_sd is None:
                  print(f"UNet checkpoint not found, re-downloading {repo}...")
                  StableGenPipelineBuilder._ensure_stablex_weights(repo, ckpt_path)
                  
                  for filename in unet_file_candidates:
                      candidate_path = os.path.join(model_dir, 'unet', filename)
                      if os.path.exists(candidate_path):
                          unet_ckpt_path = candidate_path
                          print(f"Found UNet checkpoint after re-download: {unet_ckpt_path}")
                          unet_sd = load_file(unet_ckpt_path)
                          break
                  
                  if unet_sd is None:
                      unet_dir = os.path.join(model_dir, 'unet')
                      if os.path.exists(unet_dir):
                          available_files = [f for f in os.listdir(unet_dir) if os.path.isfile(os.path.join(unet_dir, f))]
                          print(f"Available files in unet directory: {available_files}")
                      
                      raise FileNotFoundError(f"No UNet checkpoint found in {unet_dir}. Searched for: {unet_file_candidates}")
  
              for name, param in unet.named_parameters():
                  set_module_tensor_to_device(unet, name, device=DEVICE, dtype=torch_dtype, value=unet_sd[name])
  
              try:
                  vae = AutoencoderKL.from_pretrained(
                      model_dir, 
                      subfolder="vae", 
                      variant="fp16" if use_fp16 else None,
                      torch_dtype=torch_dtype
                  )
              except Exception as e:
                  print(f"Warning: Failed to load VAE with variant, trying without: {e}")
                  vae = AutoencoderKL.from_pretrained(
                      model_dir, 
                      subfolder="vae", 
                      torch_dtype=torch_dtype
                  )
              
              try:
                  controlnet = ControlNetVAEModel.from_pretrained(
                      model_dir, 
                      subfolder="controlnet", 
                      variant="fp16" if use_fp16 else None,
                      torch_dtype=torch_dtype
                  )
              except Exception as e:
                  print(f"Warning: Failed to load ControlNet with variant, trying without: {e}")
                  controlnet = ControlNetVAEModel.from_pretrained(
                      model_dir, 
                      subfolder="controlnet", 
                      torch_dtype=torch_dtype
                  )
  
              pipeline = YosoPipeline(
                  unet=unet,
                  vae=vae,
                  controlnet=controlnet,
              )
  
              pipeline.to(DEVICE)
              
              if torch_dtype == torch.float16 and DEVICE.type == "cuda":
                  pipeline.unet = pipeline.unet.half()
                  pipeline.vae = pipeline.vae.half()
                  pipeline.controlnet = pipeline.controlnet.half()
              
              return pipeline
              
          except Exception as e:
              raise Exception(f"Failed to build StableX pipeline: {e}")
      
      @staticmethod
      def _ensure_trellis_weights(repo: str, ckpt_path: str):
          base_dir = os.path.join(ckpt_path, "trellis", repo)
          
          def check_files_exist():
              required_files = [
                  "ckpts/ss_dec_conv3d_16l8_fp16.safetensors",
                  "ckpts/ss_flow_normal_dit_L_16l8_fp16.safetensors", 
                  "ckpts/slat_dec_mesh_swin8_B_64l8m256c_fp16.safetensors",
                  "ckpts/slat_flow_normal_dit_L_64l8p2_fp16.safetensors",
                  "ckpts/ss_dec_conv3d_16l8_fp16.json",
                  "ckpts/ss_flow_normal_dit_L_16l8_fp16.json",
                  "ckpts/slat_dec_mesh_swin8_B_64l8m256c_fp16.json", 
                  "ckpts/slat_flow_normal_dit_L_64l8p2_fp16.json"
              ]
              
              for req_file in required_files:
                  if not os.path.exists(os.path.join(base_dir, req_file)):
                      print(f"Missing required file: {req_file}")
                      return False
                      
              return True
          
          if not os.path.exists(base_dir) or not os.listdir(base_dir) or not check_files_exist():
              print(f"Downloading {repo} to {base_dir}")
              os.makedirs(base_dir, exist_ok=True)
              
              max_retries = 3
              for attempt in range(max_retries):
                  try:
                      print(f"Download attempt {attempt + 1}/{max_retries}")
                      snapshot_download(
                          repo_id=f"Stable-X/{repo}",
                          repo_type="model",
                          local_dir=base_dir,
                          resume_download=True,
                          ignore_patterns=HF_DOWNLOAD_IGNORE
                      )
                      
                      if check_files_exist():
                          print(f"Successfully downloaded {repo}")
                          break
                      else:
                          print(f"Download incomplete, retrying... ({attempt + 1}/{max_retries})")
                          
                  except Exception as e:
                      print(f"Download failed on attempt {attempt + 1}: {e}")
                      if attempt == max_retries - 1:
                          raise Exception(f"Failed to download {repo} after {max_retries} attempts")
                      
              if not check_files_exist():
                  raise Exception(f"Failed to download required files for {repo}")
      
      @staticmethod
      def _ensure_stablex_weights(repo: str, ckpt_path: str):
          base_dir = os.path.join(ckpt_path, "stablex", repo)
          
          def check_files_exist():
              required_files = [
                  "unet/config.json",
                  "vae/config.json", 
                  "controlnet/config.json",
                  "model_index.json"
              ]
              
              unet_files = [
                  "unet/diffusion_pytorch_model.fp16.safetensors",
                  "unet/diffusion_pytorch_model.safetensors"
              ]
              
              for req_file in required_files:
                  if not os.path.exists(os.path.join(base_dir, req_file)):
                      print(f"Missing required file: {req_file}")
                      return False
              
              unet_found = False
              for unet_file in unet_files:
                  if os.path.exists(os.path.join(base_dir, unet_file)):
                      unet_found = True
                      break
              
              if not unet_found:
                  print(f"No UNet model file found. Checked: {unet_files}")
                  return False
                  
              return True
          
          if not os.path.exists(base_dir) or not os.listdir(base_dir) or not check_files_exist():
              print(f"Downloading {repo} to {base_dir}")
              os.makedirs(base_dir, exist_ok=True)
              
              max_retries = 3
              for attempt in range(max_retries):
                  try:
                      print(f"Download attempt {attempt + 1}/{max_retries}")
                      snapshot_download(
                          repo_id=f"Stable-X/{repo}",
                          repo_type="model",
                          local_dir=base_dir,
                          resume_download=True,
                          ignore_patterns=["*text_encoder*", "tokenizer*", "*scheduler*"]
                      )
                      
                      if check_files_exist():
                          print(f"Successfully downloaded {repo}")
                          break
                      else:
                          print(f"Download incomplete, retrying... ({attempt + 1}/{max_retries})")
                          
                  except Exception as e:
                      print(f"Download failed on attempt {attempt + 1}: {e}")
                      if attempt == max_retries - 1:
                          raise Exception(f"Failed to download {repo} after {max_retries} attempts")
                      
              if not check_files_exist():
                  raise Exception(f"Failed to download required files for {repo}")
      
      @staticmethod
      def _setup_environment(attn_backend: str, sparse_backend: str, spconv_algo: str, smooth_k: bool):
          try:
              success = set_attention_backend(attn_backend)
              if not success:
                  print(f"Warning: Failed to set {attn_backend}, fallback to default")
  
              success2 = set_sparse_backend(sparse_backend, spconv_algo)
              if not success2:
                  print(f"Warning: Failed to set {sparse_backend}, fallback to default")
  
              os.environ['SAGEATTN_SMOOTH_K'] = '1' if smooth_k else '0'
          except Exception as e:
              print(f"Warning: Could not setup environment: {e}")
      
      @staticmethod
      def _optimize_trellis_pipeline(pipeline, use_fp16: bool = True):
          if DEVICE.type == "cuda":
              try:
                  if hasattr(pipeline, 'cuda'):
                      pipeline.cuda()
  
                  if use_fp16:
                      if hasattr(pipeline, 'enable_attention_slicing'):
                          pipeline.enable_attention_slicing(slice_size="auto")
                      if hasattr(pipeline, 'half'):
                          pipeline.half()
              except Exception as e:
                  print(f"Warning: Some pipeline optimizations failed: {str(e)}")
  
          return pipeline 