feat: add thread-safe model initialization, nest_asyncio support, and pre-loading on startup

app.py

@@ -14,11 +14,15 @@ from datetime import datetime
 from typing import *
 from PIL import Image
 
-# Re-install utils3d as in original app.py
-subprocess.run([
-    "pip", "install", "--force-reinstall", "--no-deps",
-    "https://github.com/LDYang694/Storages/releases/download/20260430/utils3d-0.0.2-py3-none-any.whl"
-], check=True)
+import threading
+try:
+    import nest_asyncio
+    nest_asyncio.apply()
+except ImportError:
+    pass
+
+# Lock for model initialization
+init_lock = threading.Lock()
 
 os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
@@ -119,36 +123,37 @@ envmap = None
 
 def init_models():
     global pipeline, moge_model, envmap
-    if pipeline is not None:
-        return
-
-    model_path = "TencentARC/Pixal3D-T"
-    print(f"[Pipeline] Loading from {model_path}...")
-    pipeline = Pixal3DImageTo3DPipeline.from_pretrained(model_path)
-    
-    print("[ImageCond] Building DinoV3ProjFeatureExtractor models...")
-    pipeline.image_cond_model_ss = build_image_cond_model(IMAGE_COND_CONFIGS["ss"])
-    pipeline.image_cond_model_shape_512 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_512"])
-    pipeline.image_cond_model_shape_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_1024"])
-    pipeline.image_cond_model_tex_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["tex_1024"])
-    
-    pipeline.cuda()
-    
-    print("[NAF] Pre-loading NAF upsampler model...")
-    for attr in ['image_cond_model_ss', 'image_cond_model_shape_512', 'image_cond_model_shape_1024', 'image_cond_model_tex_1024']:
-        model = getattr(pipeline, attr, None)
-        if model is not None and getattr(model, 'use_naf_upsample', False):
-            model._load_naf()
-            
-    print("[MoGe-2] Loading model for camera estimation...")
-    moge_model = load_moge_model(device="cuda")
-    
-    print("[EnvMap] Loading environment maps...")
-    envmap = {
-        'forest': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread('assets/hdri/forest.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device='cuda')),
-        'sunset': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread('assets/hdri/sunset.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device='cuda')),
-        'courtyard': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread('assets/hdri/courtyard.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device='cuda')),
-    }
+    with init_lock:
+        if pipeline is not None:
+            return
+
+        model_path = "TencentARC/Pixal3D-T"
+        print(f"[Pipeline] Loading from {model_path}...")
+        pipeline = Pixal3DImageTo3DPipeline.from_pretrained(model_path)
+        
+        print("[ImageCond] Building DinoV3ProjFeatureExtractor models...")
+        pipeline.image_cond_model_ss = build_image_cond_model(IMAGE_COND_CONFIGS["ss"])
+        pipeline.image_cond_model_shape_512 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_512"])
+        pipeline.image_cond_model_shape_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_1024"])
+        pipeline.image_cond_model_tex_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["tex_1024"])
+        
+        pipeline.cuda()
+        
+        print("[NAF] Pre-loading NAF upsampler model...")
+        for attr in ['image_cond_model_ss', 'image_cond_model_shape_512', 'image_cond_model_shape_1024', 'image_cond_model_tex_1024']:
+            model = getattr(pipeline, attr, None)
+            if model is not None and getattr(model, 'use_naf_upsample', False):
+                model._load_naf()
+                
+        print("[MoGe-2] Loading model for camera estimation...")
+        moge_model = load_moge_model(device="cuda")
+        
+        print("[EnvMap] Loading environment maps...")
+        envmap = {
+            'forest': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread('assets/hdri/forest.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device='cuda')),
+            'sunset': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread('assets/hdri/sunset.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device='cuda')),
+            'courtyard': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread('assets/hdri/courtyard.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device='cuda')),
+        }
 
 # ============================================================================
 # Utilities
@@ -341,4 +346,13 @@ def extract_glb_api(state_path: str, decimation_target: int, texture_size: int)
 app.mount("/assets", StaticFiles(directory="assets"), name="assets")
 
 if __name__ == "__main__":
-    app.launch(show_error=True)
+    # Re-install utils3d as in original app.py
+    subprocess.run([
+        "pip", "install", "--force-reinstall", "--no-deps",
+        "https://github.com/LDYang694/Storages/releases/download/20260430/utils3d-0.0.2-py3-none-any.whl"
+    ], check=True)
+    
+    # Pre-initialize models before launching the server
+    init_models()
+    
+    app.launch(show_error=True, share=True)