Update app.py

working on timeouts

app.py

@@ -33,7 +33,7 @@ def end_session(req: gr.Request):
     shutil.rmtree(user_dir)
 
 # Image Preprocessing Function
-def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
+def preprocess_image(image: Image.Image) -> Image.Image:
     """
     Preprocess the input image.
 
@@ -41,12 +41,15 @@ def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
         image (Image.Image): The input image.
 
     Returns:
-        str: uuid of the trial.
         Image.Image: The preprocessed image.
     """
+    # Validate image
+    if image is None:
+        raise ValueError("No image provided.")
+    if image.mode != "RGBA":
+        image = image.convert("RGBA")
     processed_image = pipeline.preprocess_image(image)
-    trial_id = str(uuid.uuid4())
-    return trial_id, processed_image
+    return processed_image
 
 # State Packing and Unpacking Functions
 def pack_state(gs: Gaussian, mesh: MeshExtractResult, trial_id: str) -> dict:
@@ -92,13 +95,19 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
 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(
-    trial_id: str,
     image: Image.Image,
     seed: int,
     ss_guidance_strength: float,
@@ -111,7 +120,6 @@ def image_to_3d(
     Convert an image to a 3D model.
 
     Args:
-        trial_id (str): The UUID of the trial.
         image (Image.Image): The input image.
         seed (int): The random seed.
         ss_guidance_strength (float): The guidance strength for sparse structure generation.
@@ -121,8 +129,7 @@ def image_to_3d(
         req (gr.Request): Gradio request object.
 
     Returns:
-        dict: The information of the generated 3D model.
-        str: The path to the video of the 3D model.
+        Tuple[dict, str]: The state dictionary and the path to the generated video.
     """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     outputs = pipeline.run(
@@ -142,6 +149,7 @@ def image_to_3d(
     video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
     video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
     video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+    trial_id = uuid.uuid4()
     video_path = os.path.join(user_dir, f"{trial_id}.mp4")
     imageio.mimsave(video_path, video, fps=15)
     state = pack_state(outputs['gaussian'][0], outputs['mesh'][0], trial_id)
@@ -166,7 +174,7 @@ def extract_glb(
         req (gr.Request): Gradio request object.
 
     Returns:
-        str: The path to the extracted GLB file.
+        Tuple[str, 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)
@@ -210,53 +218,62 @@ 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)
-
+            
+            # 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")
+                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")
+                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
             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)
-
+            
+            # Existing Extract GLB Button
             extract_glb_btn = gr.Button("Extract GLB", interactive=False)
-
+            
             # **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
             download_glb = gr.DownloadButton(
                 label="Download GLB",
                 # Removed 'file_count' to prevent runtime error
             )
-
+            
             # **Addition: Download High Quality GLB DownloadButton**
             download_high_quality_glb = gr.DownloadButton(
                 label="Download High Quality GLB",
                 # 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 at the bottom of the page
+    # Example Images
     with gr.Row():
         examples = gr.Examples(
             examples=[
@@ -270,10 +287,10 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
             examples_per_page=64,
         )
 
-    # Handlers
+    # Event Handlers
     demo.load(start_session)
     demo.unload(end_session)
-
+    
     image_prompt.upload(
         preprocess_image,
         inputs=[image_prompt],
@@ -284,25 +301,28 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         get_seed,
         inputs=[randomize_seed, seed],
         outputs=[seed],
+        concurrency_limit=1  # Set concurrency limit for Generate
     ).then(
         image_to_3d,
-        inputs=[output_buf, 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, gr.Request()],
         outputs=[output_buf, video_output],
+        concurrency_limit=1  # Set concurrency limit for image_to_3d
     ).then(
         # Enable the Extract GLB and Download High Quality GLB buttons after generation
-        lambda: (gr.Button.update(interactive=True), gr.Button.update(interactive=True)),
+        lambda: (True, True),
         outputs=[extract_glb_btn, extract_glb_high_quality_btn],
     )
 
     video_output.clear(
-        lambda: (gr.Button.update(interactive=False), gr.Button.update(interactive=False)),
+        lambda: (False, False),
         outputs=[extract_glb_btn, extract_glb_high_quality_btn],
     )
 
     extract_glb_btn.click(
         extract_glb,
-        inputs=[output_buf, mesh_simplify, texture_size],
+        inputs=[output_buf, mesh_simplify, texture_size, gr.Request()],
         outputs=[model_output, glb_path_state],
+        concurrency_limit=1  # Set concurrency limit for extract_glb
     ).then(
         lambda glb_path: glb_path if glb_path else "",
         inputs=[glb_path_state],
@@ -312,8 +332,9 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     # **Addition: High-Quality GLB Extraction and Download**
     extract_glb_high_quality_btn.click(
         extract_glb_high_quality,
-        inputs=[output_buf],
+        inputs=[output_buf, gr.Request()],
         outputs=[model_output, glb_high_quality_path_state],
+        concurrency_limit=1  # Set concurrency limit for extract_glb_high_quality
     ).then(
         lambda glb_path: glb_path if glb_path else "",
         inputs=[glb_high_quality_path_state],
@@ -325,14 +346,14 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         outputs=[download_glb, download_high_quality_glb],
     )
 
-# **Addition: Configure Gradio's Queue to Handle Long GPU Operations**
+# Launch the Gradio app
 if __name__ == "__main__":
     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
-    # Configure Gradio's queue with appropriate settings
-    # Removed 'concurrency_count' and 'timeout' as they are deprecated
+    except Exception as e:
+        print(f"Preloading rembg failed: {e}")
+    
+    # Configure Gradio's queue without deprecated parameters
     demo.queue().launch()