Update app.py

let's try a different approach

app.py

@@ -16,96 +16,26 @@ 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:
-    """
-    Preprocess the input image.
-
-    Args:
-        image (Image.Image): The input image.
-
-    Returns:
-        Image.Image: The preprocessed image.
-    """
-    # Validate image
-    if image is None:
-        raise ValueError("No image provided.")
-    if image.mode != "RGBA":
-        image = image.convert("RGBA")
+def preprocess_image(image: Image.Image) -> Tuple[str, Image.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': {
-            **gs.init_params,
-            '_xyz': gs._xyz.cpu().numpy(),
-            '_features_dc': gs._features_dc.cpu().numpy(),
-            '_scaling': gs._scaling.cpu().numpy(),
-            '_rotation': gs._rotation.cpu().numpy(),
-            '_opacity': gs._opacity.cpu().numpy(),
-        },
-        'mesh': {
-            'vertices': mesh.vertices.cpu().numpy(),
-            'faces': mesh.faces.cpu().numpy(),
-        },
-        'trial_id': trial_id,
-    }
-
-def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
-    gs = Gaussian(
-        aabb=state['gaussian']['aabb'],
-        sh_degree=state['gaussian']['sh_degree'],
-        mininum_kernel_size=state['gaussian']['mininum_kernel_size'],
-        scaling_bias=state['gaussian']['scaling_bias'],
-        opacity_bias=state['gaussian']['opacity_bias'],
-        scaling_activation=state['gaussian']['scaling_activation'],
-    )
-    gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
-    gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
-    gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
-    gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
-    gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
-    
-    mesh = edict(
-        vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
-        faces=torch.tensor(state['mesh']['faces'], device='cuda'),
-    )
-    
-    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,
@@ -115,21 +45,12 @@ def image_to_3d(
     slat_guidance_strength: float,
     slat_sampling_steps: int,
     req: gr.Request,
-) -> Tuple[dict, str]:
+) -> Tuple[str, str, str]:
     """
-    Convert an image to a 3D model.
-
-    Args:
-        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.
-
+    Convert an image to a 3D model and save both video preview and full-quality GLB.
+    
     Returns:
-        Tuple[dict, str]: The state dictionary and the path to the generated video.
+        Tuple[str, str, str]: (video_path, glb_path, download_path)
     """
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     outputs = pipeline.run(
@@ -146,134 +67,61 @@ def image_to_3d(
             "cfg_strength": slat_guidance_strength,
         },
     )
+    
+    # Generate and save video preview
     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()
+    trial_id = str(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,
-    mesh_simplify: float,
-    texture_size: int,
-    req: gr.Request,
-) -> Tuple[str, str]:
-    """
-    Extract a GLB file from the 3D model.
-
-    Args:
-        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.
-    """
-    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    gs, mesh, trial_id = unpack_state(state)
-    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
-    glb_path = os.path.join(user_dir, f"{trial_id}.glb")
-    glb.export(glb_path)
-    torch.cuda.empty_cache()
-    return glb_path, glb_path
-
-# **Addition: 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")
+    
+    # Save full-quality GLB directly from the generated mesh
+    glb = postprocessing_utils.to_glb(
+        outputs['gaussian'][0], 
+        outputs['mesh'][0],
+        simplify=0.0,  # No simplification
+        texture_size=2048,  # Maximum texture resolution
+        verbose=False
+    )
+    glb_path = os.path.join(user_dir, f"{trial_id}_full.glb")
     glb.export(glb_path)
+    
     torch.cuda.empty_cache()
-    return glb_path, glb_path
+    return video_path, 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 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.
+    * Upload an image and click "Generate" to create a high-quality 3D model
+    * Once generation is complete, you can download the full-quality GLB file
+    * The model will be in maximum quality with no reduction applied
     """)
     
     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
-            )
+            video_output = gr.Video(label="Generated 3D Asset Preview", autoplay=True, loop=True, height=300)
+            model_output = LitModel3D(label="3D Model Preview", exposure=20.0, height=300)
+            download_glb = gr.DownloadButton(label="Download Full-Quality GLB", interactive=False)
             
-            # **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
+    # Example images
     with gr.Row():
         examples = gr.Examples(
             examples=[
@@ -287,7 +135,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
             examples_per_page=64,
         )
 
-    # Event Handlers
+    # Event handlers
     demo.load(start_session)
     demo.unload(end_session)
     
@@ -301,59 +149,21 @@ 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=[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: (True, True),
-        outputs=[extract_glb_btn, extract_glb_high_quality_btn],
-    )
-
-    video_output.clear(
-        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, gr.Request()],
-        outputs=[model_output, glb_path_state],
-        concurrency_limit=1  # Set concurrency limit for extract_glb
+        inputs=[image_prompt, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
+        outputs=[video_output, model_output, download_glb],
     ).then(
-        lambda glb_path: glb_path if glb_path else "",
-        inputs=[glb_path_state],
+        lambda: gr.Button(interactive=True),
         outputs=[download_glb],
     )
 
-    # **Addition: High-Quality GLB Extraction and Download**
-    extract_glb_high_quality_btn.click(
-        extract_glb_high_quality,
-        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],
-        outputs=[download_high_quality_glb],
-    )
-
-    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
 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 Exception as e:
-        print(f"Preloading rembg failed: {e}")
-    
-    # Configure Gradio's queue without deprecated parameters
-    demo.queue().launch()
+    except:
+        pass
+    demo.launch()