Update code/demo.py: 主程序（支持Model Hub权重加载）

code/demo.py

@@ -319,33 +319,48 @@ def generate_ldr_gpu(ldr_content, ldr_path):
     bert_shift = 1
     shift = 0
 
-    # Prepare checkpoint paths (matching original demo.zip actual behavior)
+    # Prepare checkpoint paths (3 separate weight files as per original demo.zip design)
     config_path = os.path.join(os.path.dirname(__file__), 'cube3d/configs/open_model_v0.5.yaml')
 
     # Detect HuggingFace Spaces environment
     is_hf_space = os.getenv("SPACE_ID") is not None
 
     if is_hf_space:
-        # HF Spaces: Download checkpoint
+        # HF Spaces: Download all 3 checkpoints from Model Hub
         from huggingface_hub import hf_hub_download
-        checkpoint_path = hf_hub_download(
+        print("📥 Downloading model weights from HuggingFace Hub...")
+
+        # Download base GPT model (7.17 GB)
+        gpt_ckpt_path = hf_hub_download(
+            repo_id="0xZohar/object-assembler-models",
+            filename="shape_gpt.safetensors",
+            cache_dir=HF_CACHE_DIR,
+            local_files_only=False
+        )
+        print(f"   ✓ Base GPT model loaded")
+
+        # Download shape tokenizer (1.09 GB)
+        shape_ckpt_path = hf_hub_download(
+            repo_id="0xZohar/object-assembler-models",
+            filename="shape_tokenizer.safetensors",
+            cache_dir=HF_CACHE_DIR,
+            local_files_only=False
+        )
+        print(f"   ✓ Shape tokenizer loaded")
+
+        # Download fine-tuned adapter (1.68 GB)
+        save_gpt_ckpt_path = hf_hub_download(
             repo_id="0xZohar/object-assembler-models",
             filename="save_shape_cars_whole_p_rot_scratch_4mask_randp.safetensors",
             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
+        print(f"   ✓ Fine-tuned adapter loaded")
     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
+        # Local environment: Use local paths (matching original demo.zip structure)
+        gpt_ckpt_path = 'temp_weights/shape_gpt.safetensors'
+        shape_ckpt_path = 'temp_weights/shape_tokenizer.safetensors'
+        save_gpt_ckpt_path = '/private/tmp/demo_extracted/demo/code/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')