add

app.py

@@ -231,25 +231,14 @@ class ShapERenderer:
     def __init__(self, device):
         print("Initializing Shap-E models...")
         self.device = device
-        self.xm = None
-        self.model = None
-        self.diffusion = None
+        torch.cuda.empty_cache()  # Clear GPU memory before loading
+        self.xm = load_model('transmitter', device=self.device)
+        self.model = load_model('text300M', device=self.device)
+        self.diffusion = diffusion_from_config(load_config('diffusion'))
         print("Shap-E models initialized!")
     
-    def ensure_models_loaded(self):
-        if self.xm is None:
-            try:
-                torch.cuda.empty_cache()  # Clear GPU memory before loading
-                self.xm = load_model('transmitter', device=self.device)
-                self.model = load_model('text300M', device=self.device)
-                self.diffusion = diffusion_from_config(load_config('diffusion'))
-            except Exception as e:
-                print(f"Error loading models: {e}")
-                raise
-
     def generate_views(self, prompt, guidance_scale=15.0, num_steps=64):
         try:
-            self.ensure_models_loaded()
             torch.cuda.empty_cache()  # Clear GPU memory before generation
             
             # Generate latents using the text-to-3D model
@@ -313,24 +302,13 @@ class ShapERenderer:
 class RefinerInterface:
     def __init__(self):
         print("Initializing InstantMesh models...")
-        self.pipeline = None
-        self.model = None
-        self.infer_config = None
+        torch.cuda.empty_cache()  # Clear GPU memory before loading
+        self.pipeline, self.model, self.infer_config = load_models()
         print("InstantMesh models initialized!")
     
-    def ensure_models_loaded(self):
-        if self.pipeline is None:
-            try:
-                torch.cuda.empty_cache()  # Clear GPU memory before loading
-                self.pipeline, self.model, self.infer_config = load_models()
-            except Exception as e:
-                print(f"Error loading models: {e}")
-                raise
-    
     def refine_model(self, input_image, prompt, steps=75, guidance_scale=7.5):
         """Main refinement function"""
         try:
-            self.ensure_models_loaded()
             torch.cuda.empty_cache()  # Clear GPU memory before processing
             
             # Process image and get refined output
@@ -404,9 +382,13 @@ class RefinerInterface:
             raise
 
 def create_demo():
+    print("Initializing models...")
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    # Initialize models at startup
     shap_e = ShapERenderer(device)
     refiner = RefinerInterface()
+    print("All models initialized!")
     
     with gr.Blocks() as demo:
         gr.Markdown("# Shap-E to InstantMesh Pipeline")