Fix: Match demo.zip checkpoint loading logic exactly - all 3 paths point to same existing file

code/demo.py

@@ -7,6 +7,11 @@ from PIL import Image, ImageDraw, ImageFont
 import numpy as np
 import torch
 
+# Normalize OMP threads for libgomp (HF Spaces sometimes inject invalid values)
+_omp_env = os.getenv("OMP_NUM_THREADS", "")
+if _omp_env and not _omp_env.isdigit():
+    os.environ["OMP_NUM_THREADS"] = "4"
+
 # ZeroGPU Support - CRITICAL for HuggingFace Spaces
 try:
     import spaces
@@ -79,17 +84,28 @@ def model_predict(ldr_content):
 
 DEFAULT_PART_RENDER_PATH = "../data/car_1k/demos/example/part_ldr_1k_render/"
 os.makedirs(DEFAULT_PART_RENDER_PATH, exist_ok=True)
+# Ensure a visual placeholder exists to avoid broken images in the gallery
+UNKNOWN_PART_IMG = os.path.join(DEFAULT_PART_RENDER_PATH, "unknown_part.png")
+if not os.path.exists(UNKNOWN_PART_IMG):
+    os.makedirs(os.path.dirname(UNKNOWN_PART_IMG), exist_ok=True)
+    img = Image.new("RGB", (256, 256), color=(240, 240, 240))
+    draw = ImageDraw.Draw(img)
+    text = "No preview"
+    draw.text((70, 120), text, fill=(120, 120, 120))
+    img.save(UNKNOWN_PART_IMG)
+
+
 def get_part_renderings(part_names):
     renderings = []
     for part in part_names:
         # 拼接零件对应的渲染图路径（假设文件名与part_name一致，后缀为.png）
         # 例如：part为"3001.dat" → 对应路径为 "./part_renders/3001.dat.png"
-        part_base = part.replace(".dat", "")  # 统一转为小写并移除.dat
+        part_base = part.replace(".dat", "").replace("/", "_")  # 统一转为小写并移除非法路径分隔符
         part_render_path = os.path.join(DEFAULT_PART_RENDER_PATH, f"{part_base}.png")
         # 检查文件是否存在，不存在则使用默认缺失图（可选逻辑）
         if not os.path.exists(part_render_path):
             # 若需要，可指定一张"未知零件"的默认图路径
-            part_render_path = os.path.join(DEFAULT_PART_RENDER_PATH, "unknown_part.png")
+            part_render_path = UNKNOWN_PART_IMG
         
         renderings.append((part_render_path, part))  # (图片路径, 零件名)
     return renderings
@@ -303,14 +319,14 @@ def generate_ldr_gpu(ldr_content, ldr_path):
     bert_shift = 1
     shift = 0
 
-    # Prepare checkpoint paths
+    # Prepare checkpoint paths (matching original demo.zip actual behavior)
     config_path = os.path.join(os.path.dirname(__file__), 'cube3d/configs/open_model_v0.5.yaml')
-    gpt_ckpt_path = None  # test mode doesn't use this
 
     # Detect HuggingFace Spaces environment
     is_hf_space = os.getenv("SPACE_ID") is not None
 
     if is_hf_space:
+        # HF Spaces: Download checkpoint
         from huggingface_hub import hf_hub_download
         checkpoint_path = hf_hub_download(
             repo_id="0xZohar/object-assembler-models",
@@ -318,11 +334,18 @@ def generate_ldr_gpu(ldr_content, ldr_path):
             cache_dir=HF_CACHE_DIR,
             local_files_only=False
         )
+        # Original demo.zip references 3 files, but only 1 exists
+        # In test mode, engine.py ignores gpt_ckpt_path, uses save_gpt_ckpt_path for gpt_model
+        # and uses shape_ckpt_path for shape_model. Both must point to existing file.
+        gpt_ckpt_path = checkpoint_path  # Unused in test mode, but must be valid path
+        shape_ckpt_path = checkpoint_path  # Used for shape_model
+        save_gpt_ckpt_path = checkpoint_path  # Used for gpt_model
+    else:
+        # Local environment: All point to the same existing file (matching demo.zip reality)
+        checkpoint_path = 'model_weights/save_shape_cars_whole_p_rot_scratch_4mask_randp.safetensors'
+        gpt_ckpt_path = checkpoint_path
         shape_ckpt_path = checkpoint_path
         save_gpt_ckpt_path = checkpoint_path
-    else:
-        shape_ckpt_path = 'model_weights/save_shape_cars_whole_p_rot_scratch_4mask_randp.safetensors'
-        save_gpt_ckpt_path = 'model_weights/save_shape_cars_whole_p_rot_scratch_4mask_randp.safetensors'
 
     # Create fresh engine instance (fixes state corruption from caching)
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')