Update app.py

this should fix things

app.py

@@ -16,24 +16,25 @@ from trellis.pipelines import TrellisImageTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
 
-# Constants
+
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 os.makedirs(TMP_DIR, exist_ok=True)
 
-# Session Management Functions
+
 def start_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     print(f'Creating user directory: {user_dir}')
     os.makedirs(user_dir, exist_ok=True)
 
+
 def end_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     print(f'Removing user directory: {user_dir}')
     shutil.rmtree(user_dir)
 
-# Image Preprocessing Function
-def preprocess_image(image: Image.Image) -> Image.Image:
+
+def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
     """
     Preprocess the input image.
 
@@ -41,12 +42,13 @@ def preprocess_image(image: Image.Image) -> Image.Image:
         image (Image.Image): The input image.
 
     Returns:
+        str: uuid of the trial.
         Image.Image: The preprocessed image.
     """
     processed_image = pipeline.preprocess_image(image)
     return processed_image
 
-# State Packing and Unpacking Functions
+
 def pack_state(gs: Gaussian, mesh: MeshExtractResult, trial_id: str) -> dict:
     return {
         'gaussian': {
@@ -64,6 +66,7 @@ def pack_state(gs: Gaussian, mesh: MeshExtractResult, trial_id: str) -> dict:
         'trial_id': trial_id,
     }
 
+
 def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
     gs = Gaussian(
         aabb=state['gaussian']['aabb'],
@@ -86,21 +89,14 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
     
     return gs, mesh, state['trial_id']
 
-# Seed Management Function
+
 def get_seed(randomize_seed: bool, seed: int) -> int:
     """
     Get the random seed.
-
-    Args:
-        randomize_seed (bool): Whether to randomize the seed.
-        seed (int): The provided seed value.
-
-    Returns:
-        int: The final seed to use.
     """
     return np.random.randint(0, MAX_SEED) if randomize_seed else seed
 
-# Core 3D Generation Function
+
 @spaces.GPU
 def image_to_3d(
     image: Image.Image,
@@ -121,10 +117,10 @@ def image_to_3d(
         ss_sampling_steps (int): The number of sampling steps for sparse structure generation.
         slat_guidance_strength (float): The guidance strength for structured latent generation.
         slat_sampling_steps (int): The number of sampling steps for structured latent generation.
-        req (gr.Request): Gradio request object.
 
     Returns:
-        Tuple[dict, str]: The state dictionary and the path to the generated video.
+        dict: The information of the generated 3D model.
+        str: The path to the video of the 3D model.
     """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     outputs = pipeline.run(
@@ -151,7 +147,7 @@ def image_to_3d(
     torch.cuda.empty_cache()
     return state, video_path
 
-# Existing GLB Extraction Function
+
 @spaces.GPU
 def extract_glb(
     state: dict,
@@ -166,10 +162,9 @@ def extract_glb(
         state (dict): The state of the generated 3D model.
         mesh_simplify (float): The mesh simplification factor.
         texture_size (int): The texture resolution.
-        req (gr.Request): Gradio request object.
 
     Returns:
-        Tuple[str, str]: The path to the extracted GLB file.
+        str: The path to the extracted GLB file.
     """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     gs, mesh, trial_id = unpack_state(state)
@@ -179,7 +174,8 @@ def extract_glb(
     torch.cuda.empty_cache()
     return glb_path, glb_path
 
-# **New High-Quality GLB Extraction Function**
+
+# **Addition: High-Quality GLB Extraction Function**
 @spaces.GPU
 def extract_glb_high_quality(
     state: dict,
@@ -205,7 +201,7 @@ def extract_glb_high_quality(
     torch.cuda.empty_cache()
     return glb_path, glb_path
 
-# Gradio Interface Definition
+
 with gr.Blocks(delete_cache=(600, 600)) as demo:
     gr.Markdown("""
     ## Image to 3D Asset with [TRELLIS](https://trellis3d.github.io/)
@@ -213,121 +209,53 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     * If you find the generated 3D asset satisfactory, click "Extract GLB" to extract the GLB file and download it.
     * **New:** Click "Download High Quality GLB" to download the GLB file without any polygon reduction and with maximum texture quality.
     """)
-
+    
     with gr.Row():
         with gr.Column():
-            # Image Input
-            image_prompt = gr.Image(
-                label="Image Prompt",
-                format="png",
-                image_mode="RGBA",
-                type="pil",
-                height=300
-            )
+            image_prompt = gr.Image(label="Image Prompt", format="png", image_mode="RGBA", type="pil", height=300)
             
-            # Generation Settings Accordion
             with gr.Accordion(label="Generation Settings", open=False):
-                seed = gr.Slider(
-                    0,
-                    MAX_SEED,
-                    label="Seed",
-                    value=0,
-                    step=1
-                )
-                randomize_seed = gr.Checkbox(
-                    label="Randomize Seed",
-                    value=True
-                )
-                gr.Markdown("### Stage 1: Sparse Structure Generation")
+                seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
+                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+                gr.Markdown("Stage 1: Sparse Structure Generation")
                 with gr.Row():
-                    ss_guidance_strength = gr.Slider(
-                        0.0,
-                        10.0,
-                        label="Guidance Strength",
-                        value=7.5,
-                        step=0.1
-                    )
-                    ss_sampling_steps = gr.Slider(
-                        1,
-                        500,
-                        label="Sampling Steps",
-                        value=12,
-                        step=1
-                    )
-                gr.Markdown("### Stage 2: Structured Latent Generation")
+                    ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
+                    ss_sampling_steps = gr.Slider(1, 500, label="Sampling Steps", value=12, step=1)
+                gr.Markdown("Stage 2: Structured Latent Generation")
                 with gr.Row():
-                    slat_guidance_strength = gr.Slider(
-                        0.0,
-                        10.0,
-                        label="Guidance Strength",
-                        value=3.0,
-                        step=0.1
-                    )
-                    slat_sampling_steps = gr.Slider(
-                        1,
-                        500,
-                        label="Sampling Steps",
-                        value=12,
-                        step=1
-                    )
-
-            # Generate Button
+                    slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1)
+                    slat_sampling_steps = gr.Slider(1, 500, label="Sampling Steps", value=12, step=1)
+
             generate_btn = gr.Button("Generate")
             
-            # GLB Extraction Settings Accordion
             with gr.Accordion(label="GLB Extraction Settings", open=False):
-                mesh_simplify = gr.Slider(
-                    0.0,
-                    0.98,
-                    label="Simplify",
-                    value=0.95,
-                    step=0.01
-                )
-                texture_size = gr.Slider(
-                    512,
-                    2048,
-                    label="Texture Size",
-                    value=1024,
-                    step=512
-                )
+                mesh_simplify = gr.Slider(0.0, 0.98, label="Simplify", value=0.95, step=0.01)
+                texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
             
-            # Existing Extract GLB Button
-            extract_glb_btn = gr.Button("Extract GLB", interactive=True)
+            extract_glb_btn = gr.Button("Extract GLB", interactive=False)
             
-            # **New Download High Quality GLB Button**
-            download_high_quality_glb_btn = gr.Button("Download High Quality GLB", interactive=True)
+            # **Addition: Download High Quality GLB Button**
+            extract_glb_high_quality_btn = gr.Button("Download High Quality GLB", interactive=False)
 
         with gr.Column():
-            # Video Output
-            video_output = gr.Video(
-                label="Generated 3D Asset",
-                autoplay=True,
-                loop=True,
-                height=300
-            )
-            # 3D Model Display
-            model_output = LitModel3D(
-                label="Extracted GLB",
-                exposure=20.0,
-                height=300
-            )
-            # Existing Download GLB Button
+            video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
+            model_output = LitModel3D(label="Extracted GLB", exposure=20.0, height=300)
             download_glb = gr.DownloadButton(
                 label="Download GLB",
-                file_count="single",
+                # Removed 'file_count' to prevent runtime error
             )
-            # **New Download High Quality GLB Button**
+            
+            # **Addition: Download High Quality GLB DownloadButton**
             download_high_quality_glb = gr.DownloadButton(
                 label="Download High Quality GLB",
-                file_count="single",
+                # Removed 'file_count' to prevent runtime error
             )
 
-    # State Variables
     output_buf = gr.State()
     glb_path_state = gr.State()  # For standard GLB
     glb_high_quality_path_state = gr.State()  # For high-quality GLB
 
-    # Example Images
+    # Example images at the bottom of the page
     with gr.Row():
         examples = gr.Examples(
             examples=[
@@ -341,39 +269,35 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
             examples_per_page=64,
         )
 
-    # Event Handlers
+    # Handlers
     demo.load(start_session)
     demo.unload(end_session)
     
-    # Image Upload Handler
     image_prompt.upload(
         preprocess_image,
         inputs=[image_prompt],
         outputs=[image_prompt],
     )
 
-    # Generate Button Click Handler
     generate_btn.click(
         get_seed,
         inputs=[randomize_seed, seed],
         outputs=[seed],
     ).then(
         image_to_3d,
-        inputs=[
-            image_prompt,
-            seed,
-            ss_guidance_strength,
-            ss_sampling_steps,
-            slat_guidance_strength,
-            slat_sampling_steps
-        ],
+        inputs=[image_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
         outputs=[output_buf, video_output],
     ).then(
-        lambda: gr.Button.update(interactive=True),
-        outputs=[extract_glb_btn, download_high_quality_glb_btn],
+        # Enable the Extract GLB and Download High Quality GLB buttons after generation
+        lambda: (gr.Button.update(interactive=True), gr.Button.update(interactive=True)),
+        outputs=[extract_glb_btn, extract_glb_high_quality_btn],
+    )
+
+    video_output.clear(
+        lambda: (gr.Button.update(interactive=False), gr.Button.update(interactive=False)),
+        outputs=[extract_glb_btn, extract_glb_high_quality_btn],
     )
 
-    # Existing Extract GLB Button Click Handler
     extract_glb_btn.click(
         extract_glb,
         inputs=[output_buf, mesh_simplify, texture_size],
@@ -384,8 +308,8 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         outputs=[download_glb],
     )
 
-    # **New Download High Quality GLB Button Click Handler**
-    download_high_quality_glb_btn.click(
+    # **Addition: High-Quality GLB Extraction and Download**
+    extract_glb_high_quality_btn.click(
         extract_glb_high_quality,
         inputs=[output_buf],
         outputs=[model_output, glb_high_quality_path_state],
@@ -395,25 +319,19 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         outputs=[download_high_quality_glb],
     )
 
-    # Handle Clearing of Video Output
-    video_output.clear(
-        lambda: (gr.Button.update(interactive=True), gr.Button.update(interactive=True)),
-        outputs=[extract_glb_btn, download_high_quality_glb_btn],
-    )
-
-    # Handle Clearing of Model Output
     model_output.clear(
         lambda: (gr.File.update(value=None), gr.File.update(value=None)),
         outputs=[download_glb, download_high_quality_glb],
     )
 
-# Launch the Gradio app
+
+# **Addition: Configure Gradio's Queue to Handle Long GPU Operations**
 if __name__ == "__main__":
-    # Initialize the pipeline
     pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
     pipeline.cuda()
     try:
         pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg
     except:
         pass
-    demo.queue(concurrency_count=1, max_size=10).launch()
+    # Configure Gradio's queue with appropriate settings
+    demo.queue(concurrency_count=1, max_size=10, timeout=600).launch()