Update app.py

making the high quality source mesh available for download

app.py

@@ -17,24 +17,27 @@ 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)
 
 
-def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
+# Image Preprocessing Function
+def preprocess_image(image: Image.Image) -> Image.Image:
     """
     Preprocess the input image.
 
@@ -42,13 +45,13 @@ 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.
     """
     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': {
@@ -65,8 +68,8 @@ 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'],
@@ -90,13 +93,22 @@ 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,
@@ -117,10 +129,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:
-        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(
@@ -148,6 +160,7 @@ def image_to_3d(
     return state, video_path
 
 
+# Existing GLB Extraction Function
 @spaces.GPU
 def extract_glb(
     state: dict,
@@ -162,9 +175,10 @@ 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.
+        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)
@@ -175,45 +189,158 @@ def extract_glb(
     return glb_path, glb_path
 
 
+# New High-Quality GLB Extraction Function
+@spaces.GPU
+def extract_glb_high_quality(
+    state: dict,
+    req: gr.Request,
+) -> Tuple[str, str]:
+    """
+    Extract a high-quality GLB file from the 3D model without polygon reduction.
+
+    Args:
+        state (dict): The state of the generated 3D model.
+        req (gr.Request): Gradio request object.
+
+    Returns:
+        Tuple[str, str]: The path to the high-quality GLB file.
+    """
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    gs, mesh, trial_id = unpack_state(state)
+    # Set simplify to 0.0 to disable polygon reduction
+    # Set texture_size to 2048 for maximum texture quality
+    glb = postprocessing_utils.to_glb(gs, mesh, simplify=0.0, texture_size=2048, verbose=False)
+    glb_path = os.path.join(user_dir, f"{trial_id}_high_quality.glb")
+    glb.export(glb_path)
+    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/)
-    * Upload an image and click "Generate" to create a 3D asset. If the image has alpha channel, it be used as the mask. Otherwise, we use `rembg` to remove the background.
+    * Upload an image and click "Generate" to create a 3D asset. If the image has an alpha channel, it will be used as the mask. Otherwise, the background will be removed automatically.
     * 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)
+            # 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")
+                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)
-
+                    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)
+                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)
+            
+            # New Extract High Quality GLB Button
+            extract_glb_high_quality_btn = gr.Button("Download High Quality GLB", interactive=False)
 
         with gr.Column():
-            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", interactive=False)
-            
+            # 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",
+                file_name="model.glb",
+                interactive=False
+            )
+            # New Download High Quality GLB Button
+            download_high_quality_glb = gr.DownloadButton(
+                label="Download High Quality GLB",
+                file_name="model_high_quality.glb",
+                interactive=False
+            )
+
+    # 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=[
@@ -227,51 +354,78 @@ 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 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(interactive=True),
-        outputs=[extract_glb_btn],
-    )
-
-    video_output.clear(
-        lambda: gr.Button(interactive=False),
-        outputs=[extract_glb_btn],
+        lambda: gr.Button.update(interactive=True),
+        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],
         outputs=[model_output, download_glb],
     ).then(
-        lambda: gr.Button(interactive=True),
+        lambda: gr.Button.update(interactive=True),
         outputs=[download_glb],
     )
 
+    # New Extract High Quality GLB Button Click Handler
+    extract_glb_high_quality_btn.click(
+        extract_glb_high_quality,
+        inputs=[output_buf],
+        outputs=[model_output, glb_high_quality_path_state],
+    ).then(
+        lambda glb_path: {"value": glb_path} if glb_path else None,
+        inputs=[glb_high_quality_path_state],
+        outputs=[download_high_quality_glb],
+    ).then(
+        lambda: gr.Button.update(interactive=True),
+        outputs=[download_high_quality_glb],
+    )
+
+    # Handle Clearing of Video Output
+    video_output.clear(
+        lambda: (gr.Button.update(interactive=False), gr.Button.update(interactive=False)),
+        outputs=[extract_glb_btn, extract_glb_high_quality_btn],
+    )
+
+    # Handle Clearing of Model Output
     model_output.clear(
-        lambda: gr.Button(interactive=False),
-        outputs=[download_glb],
+        lambda: (gr.File.update(value=None), gr.File.update(value=None)),
+        outputs=[download_glb, download_high_quality_glb],
     )
-    
+
 
 # Launch the Gradio app
 if __name__ == "__main__":
+    # Initialize the pipeline
     pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
     pipeline.cuda()
     try: