update

example.py

@@ -1,7 +1,10 @@
 try:
     import spaces
-except:
-    pass
+except ImportError:
+    # Define a dummy decorator if spaces is not available
+    def GPU(func):
+        return func
+    spaces = type('spaces', (), {'GPU': GPU})
 
 import os
 import torch
@@ -39,66 +42,88 @@ if not os.path.exists(adapter_name_or_path):
 triplane_turbo_pipeline = TriplaneTurboTextTo3DPipeline.from_pretrained(adapter_name_or_path)
 triplane_turbo_pipeline.to(device)
 
+@spaces.GPU
+def generate_3d_model(prompt, num_results_per_prompt=1, seed=42, device="cuda"):
+    """
+    Generate 3D models using TriplaneTurbo pipeline.
+    
+    Args:
+        prompt (str): Text prompt for the 3D model
+        num_results_per_prompt (int): Number of results to generate
+        seed (int): Random seed for generation
+        device (str): Device to use for computation
+        
+    Returns:
+        dict: Output from the pipeline
+    """
+    output = triplane_turbo_pipeline(
+        prompt=prompt,
+        num_results_per_prompt=num_results_per_prompt,
+        generator=torch.Generator(device=device).manual_seed(seed),
+        device=device,
+    )
+    # Initialize a deque with maximum length of 100 to store obj file paths
+    obj_file_queue = deque(maxlen=max_obj_files)
+
+    # Save mesh
+    os.makedirs(output_dir, exist_ok=True)
+    for i, mesh in enumerate(output["mesh"]):
+        vertices = mesh.v_pos
+        
+        # 1. First rotate -90 degrees around X-axis to make the model face up
+        vertices = torch.stack([
+            vertices[:, 0],           # x remains unchanged
+            vertices[:, 2],           # y = z
+            -vertices[:, 1]           # z = -y
+        ], dim=1)
+        
+        # 2. Then rotate 90 degrees around Y-axis to make the model face the observer
+        vertices = torch.stack([
+            -vertices[:, 2],          # x = -z
+            vertices[:, 1],           # y remains unchanged
+            vertices[:, 0]            # z = x
+        ], dim=1)
+        
+        mesh.v_pos = vertices
+        
+        # If mesh has normals, they need to be rotated in the same way
+        if mesh.v_nrm is not None:
+            normals = mesh.v_nrm
+            # 1. Rotate -90 degrees around X-axis
+            normals = torch.stack([
+                normals[:, 0],
+                normals[:, 2],
+                -normals[:, 1]
+            ], dim=1)
+            # 2. Rotate 90 degrees around Y-axis
+            normals = torch.stack([
+                -normals[:, 2],
+                normals[:, 1],
+                normals[:, 0]
+            ], dim=1)
+            mesh._v_nrm = normals
+        
+        # Save obj file and add its path to the queue
+        name = f"{prompt.replace(' ', '_')}_{seed}_{i}"
+        save_paths = export_obj(mesh, f"{output_dir}/{name}.obj")
+        obj_file_queue.append(save_paths[0])
+        
+        # If an old file needs to be removed (queue is at max length)
+        # and the file exists, delete it
+        if len(obj_file_queue) == max_obj_files and os.path.exists(obj_file_queue[0]):
+            old_file = obj_file_queue[0]
+            try:
+                os.remove(old_file)
+            except OSError as e:
+                print(f"Error deleting file {old_file}: {e}")
+
+
+
 # Run the pipeline
-output = triplane_turbo_pipeline(
+output = generate_3d_model(
     prompt=prompt,
     num_results_per_prompt=num_results_per_prompt,
-    generator=torch.Generator(device=device).manual_seed(seed),
-    device=device,
+    seed=seed,
+    device=device
 )
 
-# Initialize a deque with maximum length of 100 to store obj file paths
-obj_file_queue = deque(maxlen=max_obj_files)
-
-# Save mesh
-os.makedirs(output_dir, exist_ok=True)
-for i, mesh in enumerate(output["mesh"]):
-    vertices = mesh.v_pos
-    
-    # 1. First rotate -90 degrees around X-axis to make the model face up
-    vertices = torch.stack([
-        vertices[:, 0],           # x remains unchanged
-        vertices[:, 2],           # y = z
-        -vertices[:, 1]           # z = -y
-    ], dim=1)
-    
-    # 2. Then rotate 90 degrees around Y-axis to make the model face the observer
-    vertices = torch.stack([
-        -vertices[:, 2],          # x = -z
-        vertices[:, 1],           # y remains unchanged
-        vertices[:, 0]            # z = x
-    ], dim=1)
-    
-    mesh.v_pos = vertices
-    
-    # If mesh has normals, they need to be rotated in the same way
-    if mesh.v_nrm is not None:
-        normals = mesh.v_nrm
-        # 1. Rotate -90 degrees around X-axis
-        normals = torch.stack([
-            normals[:, 0],
-            normals[:, 2],
-            -normals[:, 1]
-        ], dim=1)
-        # 2. Rotate 90 degrees around Y-axis
-        normals = torch.stack([
-            -normals[:, 2],
-            normals[:, 1],
-            normals[:, 0]
-        ], dim=1)
-        mesh._v_nrm = normals
-    
-    # Save obj file and add its path to the queue
-    name = f"{prompt.replace(' ', '_')}_{seed}_{i}"
-    save_paths = export_obj(mesh, f"{output_dir}/{name}.obj")
-    obj_file_queue.append(save_paths[0])
-    
-    # If an old file needs to be removed (queue is at max length)
-    # and the file exists, delete it
-    if len(obj_file_queue) == max_obj_files and os.path.exists(obj_file_queue[0]):
-        old_file = obj_file_queue[0]
-        try:
-            os.remove(old_file)
-        except OSError as e:
-            print(f"Error deleting file {old_file}: {e}")
-