update

app.py

@@ -1,3 +1,6 @@
+import os
+import subprocess
+import sys
 try:
     import spaces
 except ImportError:
@@ -5,68 +8,80 @@ except ImportError:
     def GPU(func):
         return func
     spaces = type('spaces', (), {'GPU': GPU})
+    
+# Check if setup has been run
+setup_marker = ".setup_complete"
+if not os.path.exists(setup_marker):
+    print("First run detected, installing dependencies...")
+    try:
+        subprocess.check_call(["bash", "setup.sh"])
+        # Create marker file to indicate setup is complete
+        with open(setup_marker, "w") as f:
+            f.write("Setup completed")
+        print("Setup completed successfully!")
+    except subprocess.CalledProcessError as e:
+        print(f"Setup failed with error: {e}")
+        sys.exit(1)
 
-import os
 import torch
-import argparse
+import gradio as gr
 from typing import *
-from diffusers import StableDiffusionPipeline
 from collections import deque
+from diffusers import StableDiffusionPipeline
 
+from triplaneturbo_executable import TriplaneTurboTextTo3DPipeline
 from triplaneturbo_executable.utils.mesh_exporter import export_obj
-from triplaneturbo_executable import TriplaneTurboTextTo3DPipeline, TriplaneTurboTextTo3DPipelineConfig
-
 
+# Initialize global variables
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+ADAPTER_PATH = "pretrained/triplane_turbo_sd_v1.pth" #"/home/user/app/pretrained/triplane_turbo_sd_v1.pth"
+PIPELINE = None  # Will hold our pipeline instance
+OBJ_FILE_QUEUE = deque(maxlen=100)  # Queue to store OBJ file paths
 
-# Initialize configuration and parameters
-prompt = "a beautiful girl"
-output_dir = "output"
-adapter_name_or_path = "pretrained/triplane_turbo_sd_v1.pth"
-num_results_per_prompt = 1
-seed = 42
-device = "cuda"
-max_obj_files = 100
-
-# download pretrained models if not exist
-if not os.path.exists(adapter_name_or_path):
-    print(f"Downloading pretrained models from huggingface")
-    os.system(
-        f"huggingface-cli download --resume-download ZhiyuanthePony/TriplaneTurbo \
-        --include \"triplane_turbo_sd_v1.pth\" \
-        --local-dir ./pretrained \
-        --local-dir-use-symlinks False"
+def download_model():
+    """Download the pretrained model if not exists"""
+    if not os.path.exists(ADAPTER_PATH):
+        print("Downloading pretrained models from huggingface")
+        os.system(
+            f"huggingface-cli download --resume-download ZhiyuanthePony/TriplaneTurbo \
+            --include \"triplane_turbo_sd_v1.pth\" \
+            --local-dir ./pretrained \
+            --local-dir-use-symlinks False"
         )
 
-
-# Initialize the TriplaneTurbo pipeline
-triplane_turbo_pipeline = TriplaneTurboTextTo3DPipeline.from_pretrained(adapter_name_or_path)
-triplane_turbo_pipeline.to(device)
+def initialize_pipeline():
+    """Initialize the pipeline once and keep it in memory"""
+    global PIPELINE
+    if PIPELINE is None:
+        print("Initializing pipeline...")
+        PIPELINE = TriplaneTurboTextTo3DPipeline.from_pretrained(ADAPTER_PATH)
+        PIPELINE.to(DEVICE)
+        print("Pipeline initialized!")
+    return PIPELINE
 
 @spaces.GPU
-def generate_3d_model(prompt, num_results_per_prompt=1, seed=42, device="cuda"):
-    """
-    Generate 3D models using TriplaneTurbo pipeline.
+def generate_3d_mesh(prompt: str) -> Tuple[str, str]:
+    """Generate 3D mesh from text prompt"""
+    global PIPELINE, OBJ_FILE_QUEUE
     
-    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(
+    # Use the global pipeline instance
+    pipeline = initialize_pipeline()
+    
+    # Use fixed seed value
+    seed = 42
+    
+    # Generate mesh
+    output = pipeline(
         prompt=prompt,
-        num_results_per_prompt=num_results_per_prompt,
-        generator=torch.Generator(device=device).manual_seed(seed),
-        device=device,
+        num_results_per_prompt=1,
+        generator=torch.Generator(device=DEVICE).manual_seed(seed),
     )
-    # Initialize a deque with maximum length of 100 to store obj file paths
-    obj_file_queue = deque(maxlen=max_obj_files)
-
+    
     # Save mesh
+    output_dir = "outputs"
     os.makedirs(output_dir, exist_ok=True)
+    
+    mesh_path = None
     for i, mesh in enumerate(output["mesh"]):
         vertices = mesh.v_pos
         
@@ -103,27 +118,96 @@ def generate_3d_model(prompt, num_results_per_prompt=1, seed=42, device="cuda"):
             ], dim=1)
             mesh._v_nrm = normals
         
-        # Save obj file and add its path to the queue
-        name = f"{prompt.replace(' ', '_')}_{seed}_{i}"
+        name = f"{prompt.replace(' ', '_')}"
         save_paths = export_obj(mesh, f"{output_dir}/{name}.obj")
-        obj_file_queue.append(save_paths[0])
+        mesh_path = 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}")
-
-
+        # Add new file path to queue
+        OBJ_FILE_QUEUE.append(mesh_path)
+        
+        # If queue is at max length, remove oldest file
+        if len(OBJ_FILE_QUEUE) == OBJ_FILE_QUEUE.maxlen:
+            old_file = OBJ_FILE_QUEUE[0]  # Get oldest file (will be automatically removed from queue)
+            if os.path.exists(old_file):
+                try:
+                    os.remove(old_file)
+                except OSError as e:
+                    print(f"Error deleting file {old_file}: {e}")
+        
+    return mesh_path, mesh_path  # Return the path twice - once for 3D preview, once for download
 
-# Run the pipeline
-output = generate_3d_model(
-    prompt=prompt,
-    num_results_per_prompt=num_results_per_prompt,
-    seed=seed,
-    device=device
-)
+with gr.Blocks(css=".output-image, .input-image, .image-preview {height: 512px !important}") as demo:
+    # Download model if needed
+    download_model()
+    
+    # Initialize pipeline at startup
+    initialize_pipeline()
+    
+    gr.Markdown(
+        """
+        # 🌟 Text to 3D Mesh Generation with TriplaneTurbo
+        
+        Demo of the paper "Progressive Rendering Distillation: Adapting Stable Diffusion for Instant Text-to-Mesh Generation beyond 3D Training Data" [CVPR 2025]
+        
+        [GitHub Repository](https://github.com/theEricMa/TriplaneTurbo)
+        
+        ## Instructions
+        1. Enter a text prompt describing what 3D object you want to generate
+        2. Click "Generate" and wait for the model to create your 3D mesh
+        3. View the result in the 3D viewer or download the OBJ file
+        """
+    )
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            prompt = gr.Textbox(
+                label="Text Prompt",
+                placeholder="Enter your text description...",
+                value="Armor dress style of outsiderzone fantasy helmet",
+                lines=2
+            )
+            
+            generate_btn = gr.Button("Generate", variant="primary")
+            
+            examples = gr.Examples(
+                examples=[
+                    ["Armor dress style of outsiderzone fantasy helmet"],
+                    ["Gandalf the grey riding a camel in a rock concert, victorian newspaper article, hyperrealistic"],
+                    ["A DSLR photo of a bald eagle"],
+                    ["A goblin riding a lawnmower in a hospital, victorian newspaper article, 4k hd"],
+                    ["An imperial stormtrooper, highly detailed"],
+                ],
+                inputs=[prompt],
+                label="Example Prompts"
+            )
+            
+        with gr.Column(scale=1):
+            output_model = gr.Model3D(
+                label="Generated 3D Mesh",
+                camera_position=(90, 90, 3),
+                clear_color=(0.5, 0.5, 0.5, 1),
+            )
+            output_file = gr.File(label="Download OBJ file")
+    
+    generate_btn.click(
+        fn=generate_3d_mesh,
+        inputs=[prompt],
+        outputs=[output_model, output_file]
+    )
+    
+    gr.Markdown(
+        """
+        ## About
+        
+        This demo uses TriplaneTurbo, which adapts Stable Diffusion for instant text-to-mesh generation.
+        The model can generate high-quality 3D meshes from text descriptions without requiring 3D training data.
+        
+        ### Limitations
+        - Generation is deterministic with a fixed seed
+        - Complex prompts may produce unpredictable results
+        - Generated meshes may require clean-up for professional use
+        """
+    )
 
+if __name__ == "__main__":
+    demo.launch()