Update app.py

app.py

@@ -17,9 +17,6 @@ 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)
 
-# Initialize pipeline
-pipeline = TrellisImageTo3DPipeline()
-
 # Session management
 def start_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
@@ -68,10 +65,12 @@ def unpack_state(state: dict) -> tuple:
     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
 
 def get_seed(randomize_seed: bool, seed: int) -> int:
@@ -123,6 +122,7 @@ def image_to_3d(
             },
             mode=multiimage_algo,
         )
+
     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))]
@@ -141,12 +141,14 @@ def extract_glb(
 ) -> tuple:
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     gs, mesh = unpack_state(state)
+
     # Convert the mesh to polygonal surfaces (quads)
     mesh.vertices, mesh.faces = postprocessing_utils.remesh_to_quads(
         vertices=mesh.vertices.cpu().numpy(),
         faces=mesh.faces.cpu().numpy(),
         simplify=mesh_simplify
     )
+
     glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
     glb_path = os.path.join(user_dir, 'sample.glb')
     glb.export(glb_path)
@@ -156,7 +158,7 @@ def extract_glb(
 @spaces.GPU
 def extract_gaussian(state: dict, req: gr.Request) -> tuple:
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-    gs, * = unpack_state(state)
+    gs, _ = unpack_state(state)
     gaussian_path = os.path.join(user_dir, 'sample.ply')
     gs.save_ply(gaussian_path)
     torch.cuda.empty_cache()
@@ -171,6 +173,7 @@ with gr.Blocks(theme=gr.themes.Default(), delete_cache=(600, 600)) as demo:
                     image_prompt = gr.Image(label="Image Prompt", format="png", image_mode="RGBA", type="pil", height=300)
                 with gr.Tab(label="Multiple Images", id=1) as multiimage_input_tab:
                     multiimage_prompt = gr.Gallery(label="Image Prompt", format="png", type="pil", height=300, columns=3)
+
             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)
@@ -181,34 +184,38 @@ with gr.Blocks(theme=gr.themes.Default(), delete_cache=(600, 600)) as demo:
                     slat_guidance_strength = gr.Slider(0.0, 10.0, label="Latent Guidance Strength", value=3.0, step=0.1)
                     slat_sampling_steps = gr.Slider(1, 50, label="Latent Sampling Steps", value=12, step=1)
                 multiimage_algo = gr.Radio(["stochastic", "multidiffusion"], label="Multi-image Algorithm", value="stochastic")
+
             generate_btn = gr.Button("Generate", variant="primary")
+
             with gr.Accordion(label="GLB Extraction Settings", open=False):
                 mesh_simplify = gr.Slider(0.9, 0.98, label="Simplify", value=0.95, step=0.01)
                 texture_size = gr.Slider(512, 2048, label="Texture Size", value=1024, step=512)
+
             with gr.Row():
                 extract_glb_btn = gr.Button("Extract GLB", interactive=False)
                 extract_gs_btn = gr.Button("Extract Gaussian", 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/Gaussian", exposure=10.0, height=300)
+
             with gr.Row():
                 download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
                 download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False)
-    
-    # State Management
+
     is_multiimage = gr.State(False)
     output_buf = gr.State()
-    
+
     # Handlers
     demo.load(start_session)
     demo.unload(end_session)
-    
+
     single_image_input_tab.select(lambda: False, outputs=[is_multiimage])
     multiimage_input_tab.select(lambda: True, outputs=[is_multiimage])
-    
+
     image_prompt.upload(preprocess_image, inputs=[image_prompt], outputs=[image_prompt])
     multiimage_prompt.upload(preprocess_images, inputs=[multiimage_prompt], outputs=[multiimage_prompt])
-    
+
     generate_btn.click(get_seed, inputs=[randomize_seed, seed], outputs=[seed]).then(
         image_to_3d,
         inputs=[image_prompt, multiimage_prompt, is_multiimage, seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps, multiimage_algo],
@@ -217,18 +224,18 @@ with gr.Blocks(theme=gr.themes.Default(), delete_cache=(600, 600)) as demo:
         lambda: tuple([gr.Button(interactive=True), gr.Button(interactive=True)]),
         outputs=[extract_glb_btn, extract_gs_btn],
     )
-    
+
     extract_glb_btn.click(
         extract_glb,
         inputs=[output_buf, mesh_simplify, texture_size],
         outputs=[model_output, download_glb],
     )
-    
+
     extract_gs_btn.click(
         extract_gaussian,
         inputs=[output_buf],
-        outputs=[model_output, download_gs]
+        outputs=[model_output, download_gs],
     )
 
-# Launch the app
-demo.launch()
+# Launch the Gradio demo for Hugging Face Spaces
+demo.launch()