add

app.py

@@ -15,6 +15,7 @@ from shap_e.models.download import load_model, load_config
 from shap_e.util.notebooks import create_pan_cameras, decode_latent_images
 import spaces
 from shap_e.models.nn.camera import DifferentiableCameraBatch, DifferentiableProjectiveCamera
+import math
 
 from src.utils.train_util import instantiate_from_config
 from src.utils.camera_util import (
@@ -27,15 +28,14 @@ from src.utils.mesh_util import save_obj, save_glb
 from src.utils.infer_util import remove_background, resize_foreground
 
 def create_custom_cameras(size: int, device: torch.device, azimuths: list, elevations: list, 
-                          fov_degrees: float,distance) -> DifferentiableCameraBatch:
+                          fov_degrees: float, distance: float) -> DifferentiableCameraBatch:
     # Object is in a 2x2x2 bounding box (-1 to 1 in each dimension)
-    object_diagonal =  distance # Correct diagonal calculation for the cube
+    object_diagonal = distance # Correct diagonal calculation for the cube
     
     # Calculate radius based on object size and FOV
     fov_radians = math.radians(fov_degrees)
     radius = (object_diagonal / 2) / math.tan(fov_radians / 2)  # Correct radius calculation
-    # print(radius)
-    # exit(0)
+    
     origins = []
     xs = []
     ys = []
@@ -75,8 +75,6 @@ def create_custom_cameras(size: int, device: torch.device, azimuths: list, eleva
         ),
     )
 
-
-@spaces.GPU(duration=60)
 def load_models():
     """Initialize and load all required models"""
     config = OmegaConf.load('configs/instant-nerf-large-best.yaml')
@@ -231,15 +229,22 @@ def create_mesh(refined_image, model, infer_config):
 
 class ShapERenderer:
     def __init__(self, device):
-        print("Loading Shap-E models...")
+        print("Initializing Shap-E models...")
         self.device = device
-        self.xm = load_model('transmitter', device=device)
-        self.model = load_model('text300M', device=device)
-        self.diffusion = diffusion_from_config(load_config('diffusion'))
-        print("Shap-E models loaded!")
+        self.xm = None
+        self.model = None
+        self.diffusion = None
+        print("Shap-E models initialized!")
+    
+    def ensure_models_loaded(self):
+        if self.model is None:
+            self.xm = load_model('transmitter', device=self.device)
+            self.model = load_model('text300M', device=self.device)
+            self.diffusion = diffusion_from_config(load_config('diffusion'))
 
-    @spaces.GPU(duration=60)
     def generate_views(self, prompt, guidance_scale=15.0, num_steps=64):
+        self.ensure_models_loaded()
+        
         # Generate latents using the text-to-3D model
         batch_size = 1
         guidance_scale = float(guidance_scale)
@@ -272,13 +277,13 @@ class ShapERenderer:
             rendered_image = decode_latent_images(
                 self.xm,
                 latents[0],
-                rendering_mode='stf',
-                cameras=cameras
+                cameras=cameras,
+                rendering_mode='stf'
             )
-            images.append(rendered_image.detach().cpu().numpy())
+            images.append(rendered_image[0])
         
         # Convert images to uint8
-        images = [(image).astype(np.uint8) for image in images]
+        images = [np.array(image) for image in images]
         
         # Create 2x3 grid layout (640x960) instead of 3x2 (960x640)
         layout = np.zeros((960, 640, 3), dtype=np.uint8)
@@ -292,12 +297,19 @@ class ShapERenderer:
 class RefinerInterface:
     def __init__(self):
         print("Initializing InstantMesh models...")
-        self.pipeline, self.model, self.infer_config = load_models()
-        print("InstantMesh models loaded!")
+        self.pipeline = None
+        self.model = None
+        self.infer_config = None
+        print("InstantMesh models initialized!")
+    
+    def ensure_models_loaded(self):
+        if self.pipeline is None:
+            self.pipeline, self.model, self.infer_config = load_models()
     
-    @spaces.GPU(duration=65)
     def refine_model(self, input_image, prompt, steps=75, guidance_scale=7.5):
         """Main refinement function"""
+        self.ensure_models_loaded()
+        
         # Process image and get refined output
         input_image = Image.fromarray(input_image)
         
@@ -434,11 +446,13 @@ def create_demo():
             )
 
         # Set up event handlers
+        @spaces.GPU(duration=60)
         def generate(prompt, guidance_scale, num_steps):
             with torch.no_grad():
                 layout, _ = shap_e.generate_views(prompt, guidance_scale, num_steps)
             return layout
 
+        @spaces.GPU(duration=60)
         def refine(input_image, prompt, steps, guidance_scale):
             refined_img, mesh_path = refiner.refine_model(
                 input_image, 

app2.py

@@ -12,7 +12,8 @@ from einops import rearrange
 from shap_e.diffusion.sample import sample_latents
 from shap_e.diffusion.gaussian_diffusion import diffusion_from_config
 from shap_e.models.download import load_model, load_config
-from shap_e.util.notebooks import create_pan_cameras, decode_latent_images, create_custom_cameras
+from shap_e.util.notebooks import create_pan_cameras, decode_latent_images
+from util import create_custom_cameras
 
 from src.utils.train_util import instantiate_from_config
 from src.utils.camera_util import (