Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import spaces
 from gradio_litmodel3d import LitModel3D
-
 import os
 import shutil
 os.environ['SPCONV_ALGO'] = 'native'
@@ -15,22 +14,18 @@ from trellis.pipelines import TrellisImageTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
 
-
 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)
 
-
 def start_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     os.makedirs(user_dir, exist_ok=True)
     
-    
 def end_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
     shutil.rmtree(user_dir)
 
-
 def preprocess_image(image: Image.Image) -> Image.Image:
     """
     Preprocess the input image.
@@ -42,7 +37,6 @@ def preprocess_image(image: Image.Image) -> Image.Image:
     processed_image = pipeline.preprocess_image(image)
     return processed_image
 
-
 def preprocess_images(images: List[Tuple[Image.Image, str]]) -> List[Image.Image]:
     """
     Preprocess a list of input images.
@@ -57,7 +51,6 @@ def preprocess_images(images: List[Tuple[Image.Image, str]]) -> List[Image.Image
     processed_images = [pipeline.preprocess_image(image) for image in images]
     return processed_images
 
-
 def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
     return {
         'gaussian': {
@@ -74,7 +67,6 @@ def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
         },
     }
     
-    
 def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
     gs = Gaussian(
         aabb=state['gaussian']['aabb'],
@@ -97,14 +89,12 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
     
     return gs, mesh
 
-
 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
 
-
 @spaces.GPU
 def image_to_3d(
     image: Image.Image,
@@ -175,7 +165,6 @@ def image_to_3d(
     torch.cuda.empty_cache()
     return state, video_path
 
-
 @spaces.GPU(duration=90)
 def extract_glb(
     state: dict,
@@ -200,7 +189,6 @@ def extract_glb(
     torch.cuda.empty_cache()
     return glb_path, glb_path
 
-
 @spaces.GPU
 def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
     """
@@ -217,39 +205,7 @@ def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
     torch.cuda.empty_cache()
     return gaussian_path, gaussian_path
 
-
-def prepare_multi_example() -> List[Image.Image]:
-    multi_case = list(set([i.split('_')[0] for i in os.listdir("assets/example_multi_image")]))
-    images = []
-    for case in multi_case:
-        _images = []
-        for i in range(1, 4):
-            img = Image.open(f'assets/example_multi_image/{case}_{i}.png')
-            W, H = img.size
-            img = img.resize((int(W / H * 512), 512))
-            _images.append(np.array(img))
-        images.append(Image.fromarray(np.concatenate(_images, axis=1)))
-    return images
-
-
-def split_image(image: Image.Image) -> List[Image.Image]:
-    """
-    Split an image into multiple views.
-    """
-    image = np.array(image)
-    alpha = image[..., 3]
-    alpha = np.any(alpha>0, axis=0)
-    start_pos = np.where(~alpha[:-1] & alpha[1:])[0].tolist()
-    end_pos = np.where(alpha[:-1] & ~alpha[1:])[0].tolist()
-    images = []
-    for s, e in zip(start_pos, end_pos):
-        images.append(Image.fromarray(image[:, s:e+1]))
-    return [preprocess_image(image) for image in images]
-
-
-with gr.Blocks(delete_cache=(600, 600)) as demo:
-    gr.Markdown("")
-    
+with gr.Blocks(theme=gr.themes.Default(), delete_cache=(600, 600)) as demo:
     with gr.Row():
         with gr.Column():
             with gr.Tabs() as input_tabs:
@@ -257,26 +213,19 @@ with gr.Blocks(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)
-                    gr.Markdown("""
-                        Input different views of the object in separate images. 
-                        
-                        *NOTE: this is an experimental algorithm without training a specialized model. It may not produce the best results for all images, especially those having different poses or inconsistent details.*
-                    """)
             
             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")
                 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)
                 multiimage_algo = gr.Radio(["stochastic", "multidiffusion"], label="Multi-image Algorithm", value="stochastic")
 
-            generate_btn = gr.Button("Generate")
+            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)
@@ -285,9 +234,6 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
             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.Column():
             video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
@@ -300,40 +246,17 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
     is_multiimage = gr.State(False)
     output_buf = gr.State()
 
-    # Example images at the bottom of the page
-    with gr.Row() as single_image_example:
-        examples = gr.Examples(
-            examples=[
-                f'assets/example_image/{image}'
-                for image in os.listdir("")
-            ],
-            inputs=[image_prompt],
-            fn=preprocess_image,
-            outputs=[image_prompt],
-            run_on_click=True,
-            examples_per_page=64,
-        )
-    with gr.Row(visible=False) as multiimage_example:
-        examples_multi = gr.Examples(
-            examples=prepare_multi_example(),
-            inputs=[image_prompt],
-            fn=split_image,
-            outputs=[multiimage_prompt],
-            run_on_click=True,
-            examples_per_page=8,
-        )
-
     # Handlers
     demo.load(start_session)
     demo.unload(end_session)
     
     single_image_input_tab.select(
-        lambda: tuple([False, gr.Row.update(visible=True), gr.Row.update(visible=False)]),
-        outputs=[is_multiimage, single_image_example, multiimage_example]
+        lambda: False,
+        outputs=[is_multiimage]
     )
     multiimage_input_tab.select(
-        lambda: tuple([True, gr.Row.update(visible=False), gr.Row.update(visible=True)]),
-        outputs=[is_multiimage, single_image_example, multiimage_example]
+        lambda: True,
+        outputs=[is_multiimage]
     )
     
     image_prompt.upload(
@@ -387,7 +310,6 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         lambda: gr.Button(interactive=False),
         outputs=[download_glb],
     )
-    
 
 # Launch the Gradio app
 if __name__ == "__main__":