update

app.py

@@ -1,83 +1,72 @@
-import os
-import subprocess
-import sys
 try:
     import spaces
-except:
-    pass
-
-# 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)
+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
-import gradio as gr
-from collections import deque
+import argparse
+from typing import *
 from diffusers import StableDiffusionPipeline
+from collections import deque
 
-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
 
-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 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 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
+
+# Initialize the TriplaneTurbo pipeline
+triplane_turbo_pipeline = TriplaneTurboTextTo3DPipeline.from_pretrained(adapter_name_or_path)
+triplane_turbo_pipeline.to(device)
 
 @spaces.GPU
-def generate_3d_mesh(prompt):
-    """Generate 3D mesh from text prompt"""
-    global PIPELINE, OBJ_FILE_QUEUE
+def generate_3d_model(prompt, num_results_per_prompt=1, seed=42, device="cuda"):
+    """
+    Generate 3D models using TriplaneTurbo pipeline.
     
-    # Use the global pipeline instance
-    pipeline = initialize_pipeline()
-    
-    # Use fixed seed value
-    seed = 42
-    
-    # Generate mesh
-    output = 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=1,
-        generator=torch.Generator(device=DEVICE).manual_seed(seed),
+        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
-    output_dir = "outputs"
     os.makedirs(output_dir, exist_ok=True)
-    
-    mesh_path = None
     for i, mesh in enumerate(output["mesh"]):
         vertices = mesh.v_pos
         
@@ -114,96 +103,27 @@ def generate_3d_mesh(prompt):
             ], dim=1)
             mesh._v_nrm = normals
         
-        name = f"{prompt.replace(' ', '_')}"
+        # 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")
-        mesh_path = save_paths[0]
-        
-        # 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}")
+        obj_file_queue.append(save_paths[0])
         
-    return mesh_path, mesh_path  # Return the path twice - once for 3D preview, once for download
+        # 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}")
 
-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()
+
+# Run the pipeline
+output = generate_3d_model(
+    prompt=prompt,
+    num_results_per_prompt=num_results_per_prompt,
+    seed=seed,
+    device=device
+)
+