add Trellis pipeline integration for 3D model generation and improve error handling

app.py

@@ -50,6 +50,7 @@ os.makedirs(TMP_DIR, exist_ok=True)
 
 _text_gen_pipeline = None
 _image_gen_pipeline = None
+_trellis_pipeline = None
 
 
 def start_session(req: gr.Request):
@@ -107,6 +108,25 @@ def get_text_gen_pipeline():
             return None
     return _text_gen_pipeline
 
+@spaces.GPU()
+def get_trellis_pipeline():
+    global _trellis_pipeline
+    if _trellis_pipeline is None:
+        try:
+            print("Loading Trellis pipeline...")
+            _trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
+            _trellis_pipeline.cuda()
+            
+            # Preload rembg by processing a small test image
+            try:
+                _trellis_pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))
+            except Exception as e:
+                print(f"Warning when preloading rembg: {e}")
+        except Exception as e:
+            print(f"Error loading Trellis pipeline: {e}")
+            return None
+    return _trellis_pipeline
+
 @spaces.GPU()
 def refine_prompt(prompt, system_prompt=DEFAULT_SYSTEM_PROMPT, progress=gr.Progress()):
     text_gen = get_text_gen_pipeline()
@@ -217,11 +237,29 @@ def preload_models():
     print("Preloading models...")
     text_success = get_text_gen_pipeline() is not None
     image_success = get_image_gen_pipeline() is not None
-    success = text_success and image_success
+    trellis_success = get_trellis_pipeline() is not None
+    
+    success = text_success and image_success and trellis_success
     
-    status = "Models preloaded successfully!" if success else "Error preloading models"
+    status_parts = []
+    if text_success:
+        status_parts.append("Mistral ✓")
+    else:
+        status_parts.append("Mistral ✗")
+        
+    if image_success:
+        status_parts.append("Flux ✓")
+    else:
+        status_parts.append("Flux ✗")
+        
+    if trellis_success:
+        status_parts.append("Trellis ✓")
+    else:
+        status_parts.append("Trellis ✗")
+        
+    status = f"Models loaded: {', '.join(status_parts)}"
     print(status)
-    return success
+    return success, status
 
 
 def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
@@ -274,46 +312,40 @@ def image_to_3d(
     slat_sampling_steps: int,
     req: gr.Request,
 ) -> Tuple[dict, str]:
-    """
-    Convert an image to a 3D model.
-
-    Args:
-        image (Image.Image): The input image.
-        seed (int): The random seed.
-        ss_guidance_strength (float): The guidance strength for sparse structure generation.
-        ss_sampling_steps (int): The number of sampling steps for sparse structure generation.
-        slat_guidance_strength (float): The guidance strength for structured latent generation.
-        slat_sampling_steps (int): The number of sampling steps for structured latent generation.
-
-    Returns:
-        dict: The information of the generated 3D model.
-        str: The path to the video of the 3D model.
-    """
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    
-    outputs = pipeline.run(
-        image,
-        seed=seed,
-        formats=["gaussian", "mesh"],
-        preprocess_image=False,
-        sparse_structure_sampler_params={
-            "steps": ss_sampling_steps,
-            "cfg_strength": ss_guidance_strength,
-        },
-        slat_sampler_params={
-            "steps": slat_sampling_steps,
-            "cfg_strength": slat_guidance_strength,
-        },
-    )
+    try:
+        user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+        
+        # Get the pipeline using the getter function
+        pipeline = get_trellis_pipeline()
+        if pipeline is None:
+            return None, "Trellis pipeline is unavailable."
+        
+        outputs = pipeline.run(
+            image,
+            seed=seed,
+            formats=["gaussian", "mesh"],
+            preprocess_image=False,
+            sparse_structure_sampler_params={
+                "steps": ss_sampling_steps,
+                "cfg_strength": ss_guidance_strength,
+            },
+            slat_sampler_params={
+                "steps": slat_sampling_steps,
+                "cfg_strength": slat_guidance_strength,
+            },
+        )
 
-    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
-    video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
-    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
-    video_path = os.path.join(user_dir, 'sample.mp4')
-    imageio.mimsave(video_path, video, fps=15)
-    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
-    torch.cuda.empty_cache()
-    return state, video_path
+        video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
+        video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
+        video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+        video_path = os.path.join(user_dir, 'sample.mp4')
+        imageio.mimsave(video_path, video, fps=15)
+        state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
+        torch.cuda.empty_cache()
+        return state, video_path
+    except Exception as e:
+        print(f"Error in image_to_3d: {str(e)}")
+        return None, f"Error generating 3D model: {str(e)}"
 
 
 @spaces.GPU(duration=90)
@@ -382,8 +414,8 @@ def process_example_pipeline(example_prompt, system_prompt=DEFAULT_SYSTEM_PROMPT
 def create_interface():
     # Preload models if needed
     if PRELOAD_MODELS:
-        models_loaded = preload_models()
-        model_status = "✅ Models loaded successfully!" if models_loaded else "⚠️ Error loading models"
+        model_success, model_status_details = preload_models()
+        model_status = f"✅ {model_status_details}" if model_success else f"⚠️ {model_status_details}"
     else:
         model_status = "ℹ️ Models will be loaded on demand"
 
@@ -520,14 +552,10 @@ def create_interface():
 
 
 if __name__ == "__main__":
-    # Initialize the Trellis pipeline before creating the interface
-    pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
-    pipeline.cuda()
-    try:
-        # Preload rembg
-        pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))
-    except Exception as e:
-        print(f"Warning when preloading rembg: {e}")
+    # Initialize models if PRELOAD_MODELS is True
+    if PRELOAD_MODELS:
+        success, status = preload_models()
+        print(status)
         
     demo = create_interface()
     demo.launch()