From 36a46d5e137b006c0959f5820ca4a5f623155a8e Mon Sep 17 00:00:00 2001 From: DenisKochetov Date: Thu, 5 Jun 2025 18:21:23 +0300 Subject: [PATCH] stablex model download fix --- .../Stable3DGen/pipeline_builders.py | 154 +++++++++++++++--- 1 file changed, 128 insertions(+), 26 deletions(-) diff --git a/Gen_3D_Modules/Stable3DGen/pipeline_builders.py b/Gen_3D_Modules/Stable3DGen/pipeline_builders.py index 6a26c8f..db39437 100644 --- a/Gen_3D_Modules/Stable3DGen/pipeline_builders.py +++ b/Gen_3D_Modules/Stable3DGen/pipeline_builders.py @@ -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):