292 lines
11 KiB
Python
292 lines
11 KiB
Python
# 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)}") |