Update app.py

fixing timeout issues

app.py

@@ -16,24 +16,23 @@ 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) -> Tuple[str, Image.Image]:
     """
     Preprocess the input image.
@@ -46,9 +45,10 @@ def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
         Image.Image: The preprocessed image.
     """
     processed_image = pipeline.preprocess_image(image)
-    return processed_image
-
+    trial_id = str(uuid.uuid4())
+    return trial_id, processed_image
 
+# State Packing and Unpacking Functions
 def pack_state(gs: Gaussian, mesh: MeshExtractResult, trial_id: str) -> dict:
     return {
         'gaussian': {
@@ -66,7 +66,6 @@ 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'],
@@ -89,16 +88,17 @@ 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.
     """
     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,12 +111,14 @@ 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.
         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:
         dict: The information of the generated 3D model.
@@ -140,14 +142,13 @@ 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)
     torch.cuda.empty_cache()
     return state, video_path
 
-
+# Existing GLB Extraction Function
 @spaces.GPU
 def extract_glb(
     state: dict,
@@ -162,6 +163,7 @@ 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:
         str: The path to the extracted GLB file.
@@ -174,7 +176,6 @@ def extract_glb(
     torch.cuda.empty_cache()
     return glb_path, glb_path
 
-
 # **Addition: High-Quality GLB Extraction Function**
 @spaces.GPU
 def extract_glb_high_quality(
@@ -201,7 +202,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/)
@@ -209,11 +210,11 @@ 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_prompt = gr.Image(label="Image Prompt", format="png", image_mode="RGBA", type="pil", height=300)
-            
+
             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)
@@ -227,13 +228,13 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
                     slat_sampling_steps = gr.Slider(1, 500, label="Sampling Steps", value=12, step=1)
 
             generate_btn = gr.Button("Generate")
-            
+
             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)
-            
+
             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)
 
@@ -244,7 +245,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
                 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",
@@ -272,7 +273,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     # Handlers
     demo.load(start_session)
     demo.unload(end_session)
-    
+
     image_prompt.upload(
         preprocess_image,
         inputs=[image_prompt],
@@ -285,7 +286,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         outputs=[seed],
     ).then(
         image_to_3d,
-        inputs=[image_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
+        inputs=[output_buf, image_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
         outputs=[output_buf, video_output],
     ).then(
         # Enable the Extract GLB and Download High Quality GLB buttons after generation
@@ -324,7 +325,6 @@ 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**
 if __name__ == "__main__":
     pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
@@ -334,4 +334,5 @@ if __name__ == "__main__":
     except:
         pass
     # Configure Gradio's queue with appropriate settings
-    demo.queue(concurrency_count=1, max_size=10).launch()
+    # Removed 'concurrency_count' and 'timeout' as they are deprecated
+    demo.queue().launch()