nodes/MrForExample-ComfyUI-3D-Pack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

stablex model download fix

Browse Source
This commit is contained in:
DenisKochetov
2025-06-05 18:21:23 +03:00
parent 56640bc297
commit 36a46d5e13
1 changed files with 128 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Gen_3D_Modules/Stable3DGen/pipeline_builders.py
+128 -26
View File
@@ -1,6 +1,7 @@
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
@@ -113,10 +114,13 @@ class StableGenPipelineBuilder:
try:
config_path = os.path.join(model_dir, 'unet', 'config.json')
unet_ckpt_path = os.path.join(model_dir, 'unet', 'diffusion_pytorch_model.fp16.safetensors')
if not os.path.exists(config_path):
raise FileNotFoundError(f"Config not found at {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)
@@ -124,27 +128,76 @@ class StableGenPipelineBuilder:
with init_empty_weights():
unet = UNet2DConditionModel(**config)
if os.path.exists(unet_ckpt_path):
unet_sd = load_file(unet_ckpt_path)
else:
raise FileNotFoundError(f"No checkpoint found at {unet_ckpt_path}")
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])
vae = AutoencoderKL.from_pretrained(
model_dir,
subfolder="vae",
variant="fp16" if use_fp16 else None,
torch_dtype=torch_dtype
)
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
)
controlnet = ControlNetVAEModel.from_pretrained(
model_dir,
subfolder="controlnet",
variant="fp16" if use_fp16 else None,
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,
@@ -183,16 +236,65 @@ class StableGenPipelineBuilder:
def _ensure_stablex_weights(repo: str, ckpt_path: str):
base_dir = os.path.join(ckpt_path, "stablex", repo)
if not os.path.exists(base_dir) or not os.listdir(base_dir):
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)
snapshot_download(
repo_id=f"Stable-X/{repo}",
repo_type="model",
local_dir=base_dir,
resume_download=True,
ignore_patterns=["*text_encoder*", "tokenizer*", "*scheduler*"]
)
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):