deactivate 3D

app.py

@@ -163,7 +163,7 @@ def validate_dimensions(width, height):
     return True, None
 
 @spaces.GPU()
-def infer(prompt, seed=DEFAULT_SEED, 
+def generate_image(prompt, seed=DEFAULT_SEED, 
           randomize_seed=DEFAULT_RANDOMIZE_SEED, 
           width=DEFAULT_WIDTH, 
           height=DEFAULT_HEIGHT, 
@@ -251,136 +251,136 @@ def preload_models():
     return success, status
 
 
-def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> 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(),
-        },
-    }
+# def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> 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(),
+#         },
+#     }
     
     
-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')
+# 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'),
-    )
+#     mesh = edict(
+#         vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
+#         faces=torch.tensor(state['mesh']['faces'], device='cuda'),
+#     )
     
-    return gs, mesh
-
-
-@spaces.GPU
-def image_to_3d(
-    image: Image.Image,
-    seed: int,
-    ss_guidance_strength: float,
-    ss_sampling_steps: int,
-    slat_guidance_strength: float,
-    slat_sampling_steps: int,
-) -> Tuple[dict, str]:
-    try:
-        # Use a fixed temp directory instead of user-specific
-        temp_dir = os.path.join(TMP_DIR, "temp_output")
-        os.makedirs(temp_dir, exist_ok=True)
+#     return gs, mesh
+
+
+# @spaces.GPU
+# def image_to_3d(
+#     image: Image.Image,
+#     seed: int,
+#     ss_guidance_strength: float,
+#     ss_sampling_steps: int,
+#     slat_guidance_strength: float,
+#     slat_sampling_steps: int,
+# ) -> Tuple[dict, str]:
+#     try:
+#         # Use a fixed temp directory instead of user-specific
+#         temp_dir = os.path.join(TMP_DIR, "temp_output")
+#         os.makedirs(temp_dir, exist_ok=True)
         
-        # Get the pipeline using the getter function
-        pipeline = get_trellis_pipeline()
-        if pipeline is None:
-            return None, "Trellis pipeline is unavailable."
+#         # Get the pipeline using the getter function
+#         pipeline = get_trellis_pipeline()
+#         if pipeline is None:
+#             return None, "Trellis pipeline is unavailable."
         
-        outputs = pipeline.run(
-            image,
-            seed=seed,
-            formats=["gaussian", "mesh"],
-            preprocess_image=False,
-            sparse_structure_sampler_params={
-                "steps": ss_sampling_steps,
-                "cfg_strength": ss_guidance_strength,
-            },
-            slat_sampler_params={
-                "steps": slat_sampling_steps,
-                "cfg_strength": slat_guidance_strength,
-            },
-        )
-
-        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))]
-        video_path = os.path.join(temp_dir, 'sample.mp4')
-        imageio.mimsave(video_path, video, fps=15)
-        state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
-        torch.cuda.empty_cache()
-        return state, video_path
-    except Exception as e:
-        print(f"Error in image_to_3d: {str(e)}")
-        return None, f"Error generating 3D model: {str(e)}"
-
-
-@spaces.GPU(duration=90)
-def extract_glb(
-    state: dict,
-    mesh_simplify: float,
-    texture_size: int,
-) -> 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.
-
-    Returns:
-        str: The path to the extracted GLB file.
-    """
-    temp_dir = os.path.join(TMP_DIR, "temp_output")
-    gs, mesh = unpack_state(state)
-    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
-    glb_path = os.path.join(temp_dir, 'sample.glb')
-    glb.export(glb_path)
-    torch.cuda.empty_cache()
-    return glb_path, glb_path
-
-
-@spaces.GPU
-def extract_gaussian(state: dict) -> Tuple[str, str]:
-    """
-    Extract a Gaussian file from the 3D model.
-
-    Args:
-        state (dict): The state of the generated 3D model.
-
-    Returns:
-        str: The path to the extracted Gaussian file.
-    """
-    temp_dir = os.path.join(TMP_DIR, "temp_output")
-    gs, _ = unpack_state(state)
-    gaussian_path = os.path.join(temp_dir, 'sample.ply')
-    gs.save_ply(gaussian_path)
-    torch.cuda.empty_cache()
-    return gaussian_path, gaussian_path
+#         outputs = pipeline.run(
+#             image,
+#             seed=seed,
+#             formats=["gaussian", "mesh"],
+#             preprocess_image=False,
+#             sparse_structure_sampler_params={
+#                 "steps": ss_sampling_steps,
+#                 "cfg_strength": ss_guidance_strength,
+#             },
+#             slat_sampler_params={
+#                 "steps": slat_sampling_steps,
+#                 "cfg_strength": slat_guidance_strength,
+#             },
+#         )
+
+#         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))]
+#         video_path = os.path.join(temp_dir, 'sample.mp4')
+#         imageio.mimsave(video_path, video, fps=15)
+#         state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
+#         torch.cuda.empty_cache()
+#         return state, video_path
+#     except Exception as e:
+#         print(f"Error in image_to_3d: {str(e)}")
+#         return None, f"Error generating 3D model: {str(e)}"
+
+
+# @spaces.GPU(duration=90)
+# def extract_glb(
+#     state: dict,
+#     mesh_simplify: float,
+#     texture_size: int,
+# ) -> 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.
+
+#     Returns:
+#         str: The path to the extracted GLB file.
+#     """
+#     temp_dir = os.path.join(TMP_DIR, "temp_output")
+#     gs, mesh = unpack_state(state)
+#     glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
+#     glb_path = os.path.join(temp_dir, 'sample.glb')
+#     glb.export(glb_path)
+#     torch.cuda.empty_cache()
+#     return glb_path, glb_path
+
+
+# @spaces.GPU
+# def extract_gaussian(state: dict) -> Tuple[str, str]:
+#     """
+#     Extract a Gaussian file from the 3D model.
+
+#     Args:
+#         state (dict): The state of the generated 3D model.
+
+#     Returns:
+#         str: The path to the extracted Gaussian file.
+#     """
+#     temp_dir = os.path.join(TMP_DIR, "temp_output")
+#     gs, _ = unpack_state(state)
+#     gaussian_path = os.path.join(temp_dir, 'sample.ply')
+#     gs.save_ply(gaussian_path)
+#     torch.cuda.empty_cache()
+#     return gaussian_path, gaussian_path
 
 
 # Create a combined function that handles the whole pipeline from example to image
@@ -435,14 +435,14 @@ def create_interface():
             visual_button = gr.Button("Create visual with Flux")
 
             generated_image = gr.Image(show_label=False)
-            gen3d_button = gr.Button("Create 3D visual with Trellis")
+            # gen3d_button = gr.Button("Create 3D visual with Trellis")
 
-            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)
+            # 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)  
+            # with gr.Row():
+            #     download_glb = gr.DownloadButton(label="Download GLB", interactive=False)
+            #     download_gs = gr.DownloadButton(label="Download Gaussian", interactive=False)  
     
             message_box = gr.Textbox(
                 label="Status Messages", 
@@ -487,28 +487,28 @@ def create_interface():
                         value=DEFAULT_NUM_INFERENCE_STEPS,
                     )
                 
-                with gr.Tab("3D Generation Settings"):
-                    trellis_seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
-                    trellis_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, 50, 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, 50, label="Sampling Steps", value=12, step=1)
-
-                with gr.Tab("GLB Extraction Settings"):
-                    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.Tab("3D Generation Settings"):
+                #     trellis_seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
+                #     trellis_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, 50, 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, 50, label="Sampling Steps", value=12, step=1)
+
+                # with gr.Tab("GLB Extraction Settings"):
+                #     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)
-            gr.Markdown("""
-                        *NOTE: Gaussian file can be very large (~50MB), it will take a while to display and download.*
-                        """)
+            # with gr.Row():
+            #     extract_glb_btn = gr.Button("Extract GLB", interactive=False)
+            #     extract_gs_btn = gr.Button("Extract Gaussian", interactive=False)
+            # gr.Markdown("""
+            #             *NOTE: Gaussian file can be very large (~50MB), it will take a while to display and download.*
+            #             """)
 
             output_buf = gr.State()
 
@@ -531,44 +531,44 @@ def create_interface():
 
         gr.on(
             triggers=[visual_button.click],
-            fn=infer,
+            fn=generate_image,
             inputs=[refined_prompt, flux_seed, flux_randomize_seed, width, height, num_inference_steps],
             outputs=[generated_image, message_box]
         )
 
-        gr.on(
-            triggers=[gen3d_button.click],
-            fn=image_to_3d,
-            inputs=[generated_image, trellis_seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
-            outputs=[output_state, video_output],
-        ).then(
-            # Update button states after successful 3D generation
-            lambda: (gr.Button.update(interactive=True), gr.Button.update(interactive=True), "3D model generated successfully"),
-            outputs=[extract_glb_btn, extract_gs_btn, message_box]
-        )
+        # gr.on(
+        #     triggers=[gen3d_button.click],
+        #     fn=image_to_3d,
+        #     inputs=[generated_image, trellis_seed, ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps],
+        #     outputs=[output_state, video_output],
+        # ).then(
+        #     # Update button states after successful 3D generation
+        #     lambda: (gr.Button.update(interactive=True), gr.Button.update(interactive=True), "3D model generated successfully"),
+        #     outputs=[extract_glb_btn, extract_gs_btn, message_box]
+        # )
         
-        # Add handlers for GLB and Gaussian extraction
-        gr.on(
-            triggers=[extract_glb_btn.click],
-            fn=extract_glb,
-            inputs=[output_state, mesh_simplify, texture_size],
-            outputs=[model_output, download_glb]
-        ).then(
-            lambda path: (gr.DownloadButton.update(interactive=True, value=path), "GLB extraction completed"),
-            inputs=[model_output],
-            outputs=[download_glb, message_box]
-        )
-
-        gr.on(
-            triggers=[extract_gs_btn.click],
-            fn=extract_gaussian,
-            inputs=[output_state],
-            outputs=[model_output, download_gs]
-        ).then(
-            lambda path: (gr.DownloadButton.update(interactive=True, value=path), "Gaussian extraction completed"),
-            inputs=[model_output],
-            outputs=[download_gs, message_box]
-        )
+        # # Add handlers for GLB and Gaussian extraction
+        # gr.on(
+        #     triggers=[extract_glb_btn.click],
+        #     fn=extract_glb,
+        #     inputs=[output_state, mesh_simplify, texture_size],
+        #     outputs=[model_output, download_glb]
+        # ).then(
+        #     lambda path: (gr.DownloadButton.update(interactive=True, value=path), "GLB extraction completed"),
+        #     inputs=[model_output],
+        #     outputs=[download_glb, message_box]
+        # )
+
+        # gr.on(
+        #     triggers=[extract_gs_btn.click],
+        #     fn=extract_gaussian,
+        #     inputs=[output_state],
+        #     outputs=[model_output, download_gs]
+        # ).then(
+        #     lambda path: (gr.DownloadButton.update(interactive=True, value=path), "Gaussian extraction completed"),
+        #     inputs=[model_output],
+        #     outputs=[download_gs, message_box]
+        # )
 
     return demo