()
| 123 | LOW_VRAM = os.environ.get("LOW_VRAM", "0") == "1" |
| 124 | |
| 125 | def init_models(): |
| 126 | global pipeline, moge_model, envmap |
| 127 | with init_lock: |
| 128 | if pipeline is not None: |
| 129 | return |
| 130 | |
| 131 | # GPU / CUDA Diagnostics (runs when GPU is allocated) |
| 132 | import subprocess as _sp |
| 133 | print("=" * 60) |
| 134 | print("[Diagnostics] PyTorch version:", torch.__version__) |
| 135 | print("[Diagnostics] CUDA available:", torch.cuda.is_available()) |
| 136 | if torch.cuda.is_available(): |
| 137 | print("[Diagnostics] CUDA version:", torch.version.cuda) |
| 138 | print("[Diagnostics] cuDNN version:", torch.backends.cudnn.version()) |
| 139 | for i in range(torch.cuda.device_count()): |
| 140 | name = torch.cuda.get_device_name(i) |
| 141 | cap = torch.cuda.get_device_capability(i) |
| 142 | mem = torch.cuda.get_device_properties(i).total_memory / 1024**3 |
| 143 | print(f"[Diagnostics] GPU {i}: {name}, sm_{cap[0]}{cap[1]}, {mem:.1f} GB") |
| 144 | try: |
| 145 | res = _sp.run(["nvidia-smi", "--query-gpu=name,compute_cap,memory.total", "--format=csv,noheader"], capture_output=True, text=True, timeout=10) |
| 146 | print("[Diagnostics] nvidia-smi:", res.stdout.strip()) |
| 147 | except Exception as e: |
| 148 | print(f"[Diagnostics] nvidia-smi failed: {e}") |
| 149 | print("=" * 60) |
| 150 | |
| 151 | model_path = "TencentARC/Pixal3D" |
| 152 | print(f"[Pipeline] Loading from {model_path}...") |
| 153 | pipeline = Pixal3DImageTo3DPipeline.from_pretrained(model_path) |
| 154 | |
| 155 | print("[ImageCond] Building DinoV3ProjFeatureExtractor models...") |
| 156 | pipeline.image_cond_model_ss = build_image_cond_model(IMAGE_COND_CONFIGS["ss"]) |
| 157 | pipeline.image_cond_model_shape_512 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_512"]) |
| 158 | pipeline.image_cond_model_shape_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_1024"]) |
| 159 | pipeline.image_cond_model_tex_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["tex_1024"]) |
| 160 | |
| 161 | if LOW_VRAM: |
| 162 | # Low-VRAM mode: models stay on CPU, loaded to GPU on-demand per stage. |
| 163 | print("[NAF] Pre-downloading NAF upsampler weights (CPU only)...") |
| 164 | for attr in ['image_cond_model_ss', 'image_cond_model_shape_512', |
| 165 | 'image_cond_model_shape_1024', 'image_cond_model_tex_1024']: |
| 166 | m = getattr(pipeline, attr, None) |
| 167 | if m is not None and getattr(m, 'use_naf_upsample', False): |
| 168 | m._load_naf() |
| 169 | pipeline._device = torch.device("cuda") |
| 170 | pipeline.low_vram = True |
| 171 | print("[Pipeline] Low-VRAM mode enabled.") |
| 172 | else: |
| 173 | # Standard mode: all models loaded to GPU at once. |
| 174 | pipeline.low_vram = False |
| 175 | pipeline.cuda() |
| 176 | pipeline.image_cond_model_ss.cuda() |
| 177 | pipeline.image_cond_model_shape_512.cuda() |
| 178 | pipeline.image_cond_model_shape_1024.cuda() |
| 179 | pipeline.image_cond_model_tex_1024.cuda() |
| 180 | print("[NAF] Pre-loading NAF upsampler model...") |
| 181 | for attr in ['image_cond_model_ss', 'image_cond_model_shape_512', |
| 182 | 'image_cond_model_shape_1024', 'image_cond_model_tex_1024']: |
no test coverage detected