Initial commit

app.py

@@ -202,8 +202,14 @@ class FaceLiftPipeline:
             multiview_path = output_dir / "multiview.png"
             multiview_image.save(multiview_path)
             
-            # Clear CUDA cache after diffusion to free memory
+            # Move diffusion model to CPU to free GPU memory for GS-LRM
+            print("Moving diffusion model to CPU to free memory...")
+            self.mvdiffusion_pipeline.to("cpu")
+            
+            # Delete intermediate variables to free memory
+            del result, generator
             torch.cuda.empty_cache()
+            torch.cuda.synchronize()
             
             # Prepare 3D reconstruction input
             view_arrays = [np.array(view) for view in selected_views]
@@ -233,6 +239,21 @@ class FaceLiftPipeline:
                 "index": batch_indices,
             })
             
+            # Ensure GS-LRM model is on GPU
+            if next(self.gs_lrm_model.parameters()).device.type == "cpu":
+                print("Moving GS-LRM model to GPU...")
+                self.gs_lrm_model.to(self.device)
+                torch.cuda.empty_cache()
+            
+            # Log GPU memory status
+            if torch.cuda.is_available():
+                allocated = torch.cuda.memory_allocated() / 1024**3
+                reserved = torch.cuda.memory_reserved() / 1024**3
+                print(f"GPU memory before GS-LRM: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved")
+            
+            # Final memory cleanup before reconstruction
+            torch.cuda.empty_cache()
+            
             # Run 3D reconstruction
             with torch.no_grad(), torch.autocast(enabled=True, device_type="cuda", dtype=torch.float16):
                 result = self.gs_lrm_model.forward(batch, create_visual=False, split_data=True)