388 lines
15 KiB
Python
388 lines
15 KiB
Python
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 |