Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import spaces
 from gradio_litmodel3d import LitModel3D
 from transformers import pipeline
+
 import os
 import shutil
 os.environ['SPCONV_ALGO'] = 'native'
@@ -16,6 +17,7 @@ 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)
@@ -30,21 +32,26 @@ try:
 except:
     pass
 
+
 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]:
     """
     Preprocess the input image.
+
     Args:
         image (Image.Image): The input image.
+
     Returns:
         str: uuid of the trial.
         Image.Image: The preprocessed image.
@@ -52,6 +59,7 @@ def preprocess_image(image: Image.Image) -> Tuple[str, Image.Image]:
     processed_image = pipeline.preprocess_image(image)
     return processed_image
 
+
 def pack_state(gs: Gaussian, mesh: MeshExtractResult, trial_id: str) -> dict:
     return {
         'gaussian': {
@@ -68,7 +76,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'],
@@ -91,12 +100,14 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
     
     return gs, mesh, state['trial_id']
 
+
 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,
@@ -109,6 +120,7 @@ def image_to_3d(
 ) -> Tuple[dict, str]:
     """
     Convert an image to a 3D model.
+
     Args:
         image (Image.Image): The input image.
         seed (int): The random seed.
@@ -116,6 +128,7 @@ 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.
+
     Returns:
         dict: The information of the generated 3D model.
         str: The path to the video of the 3D model.
@@ -145,6 +158,7 @@ def image_to_3d(
     torch.cuda.empty_cache()
     return state, video_path
 
+
 @spaces.GPU
 def extract_glb(
     state: dict,
@@ -154,10 +168,12 @@ def extract_glb(
 ) -> 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.
     """
@@ -169,55 +185,123 @@ def extract_glb(
     torch.cuda.empty_cache()
     return glb_path, glb_path
 
-with gr.Blocks(delete_cache=(600, 600)) as demo:
-    gr.Markdown("""<h1 style='color: #4CAF50; text-align: center;'>Image to 3D Model Generator</h1>""")
 
+with open("styles.css", "w") as f:
+    f.write("""
+    .gradio-container {
+        background-color: #f0f0f0;
+        font-family: sans-serif;
+    }
+    .my-button {
+        background-color: #4CAF50;
+        color: white;
+        padding: 10px 20px;
+        border: none;
+        border-radius: 5px;
+        cursor: pointer;
+    }
+    .my-button:hover {
+        background-color: #3e8e41;
+    }
+    .gradio-container h1{
+        text-align: center;
+        color: #333;
+        margin-bottom: 20px;
+    }
+    """)
+
+
+with gr.Blocks(css=open("styles.css").read(), delete_cache=(600, 600)) as demo:
+    gr.Markdown("# 3D Model Generator")
+    
     with gr.Row():
-        with gr.Column():
-            image_prompt = gr.Image(
-                label="Image Prompt", 
-                format="png", 
-                image_mode="RGBA", 
-                type="pil", 
-                height=300, 
-                style={"border": "2px solid #4CAF50", "border-radius": "10px", "background-color": "#F9F9F9"}
-            )
-
-            with gr.Accordion(label="Generation Settings", open=False, style={"border": "1px solid #2196F3", "border-radius": "8px"}):
-                seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1, style={"color": "#2196F3"})
-                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True, style={"color": "#2196F3"})
-                gr.Markdown("<strong>Stage 1: Sparse Structure Generation</strong>", style={"color": "#2196F3"})
+        with gr.Column(style={"padding": "20px"}):
+            image_prompt = gr.Image(label="Image Prompt", format="png", image_mode="RGBA", type="pil", height=300)
+            
+            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("<strong>Stage 2: Structured Latent Generation</strong>", style={"color": "#2196F3"})
+                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)
 
-            generate_btn = gr.Button("Generate", style={"background-color": "#4CAF50", "color": "white", "border": "none", "border-radius": "8px"})
-
-            with gr.Accordion(label="GLB Extraction Settings", open=False, style={"border": "1px solid #2196F3", "border-radius": "8px"}):
+            generate_btn = gr.Button("Generate", style={"classes": ["my-button"]})
+            
+            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)
+            
+            extract_glb_btn = gr.Button("Extract GLB", interactive=False)
 
-            extract_glb_btn = gr.Button("Extract GLB", interactive=False, style={"background-color": "#FF9800", "color": "white", "border-radius": "8px"})
+        with gr.Column(style={"padding": "20px"}):
+            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)
+            
+    output_buf = gr.State()
 
-        with gr.Column():
-            video_output = gr.Video(label="Generated 3D Model", interactive=True, style={"border-radius": "10px"})
-            glb_output = gr.File(label="Download GLB", interactive=True)
+    # Example images at the bottom of the page
+    with gr.Row():
+        examples = gr.Examples(
+            examples=[
+                f'assets/example_image/{image}'
+                for image in os.listdir("assets/example_image")
+            ],
+            inputs=[image_prompt],
+            fn=preprocess_image,
+            outputs=[image_prompt],
+            run_on_click=True,
+            examples_per_page=64,
+        )
+
+    # Handlers
+    demo.load(start_session)
+    demo.unload(end_session)
+    
+    image_prompt.upload(
+        preprocess_image,
+        inputs=[image_prompt],
+        outputs=[image_prompt],
+    )
 
-    # Trigger functions for generating 3D model and extracting GLB
     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, gr.Request()],
-        outputs=[glb_output, video_output]
+        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],
     )
 
     extract_glb_btn.click(
         extract_glb,
-        inputs=[state, mesh_simplify, texture_size, gr.Request()],
-        outputs=[glb_output, glb_output]
+        inputs=[output_buf, mesh_simplify, texture_size],
+        outputs=[model_output, download_glb],
+    ).then(
+        lambda: gr.Button(interactive=True),
+        outputs=[download_glb],
+    )
+
+    model_output.clear(
+        lambda: gr.Button(interactive=False),
+        outputs=[download_glb],
     )
+    
 
-    demo.launch()  # Start the demo
+# Launch the Gradio app
+if __name__ == "__main__":
+    demo.launch()