update

app.py

@@ -23,6 +23,7 @@ from common import (
     MAX_SEED,
     VERSION,
     backproject_texture_v2,
+    custom_theme,
     end_session,
     generate_condition,
     generate_texture_mvimages,
@@ -33,8 +34,6 @@ from common import (
     render_result_video,
     start_session,
 )
-from gradio.themes import Default
-from gradio.themes.utils.colors import slate
 
 
 def active_btn_by_content(mesh_content: gr.Model3D, text_content: gr.Textbox):
@@ -50,14 +49,32 @@ def active_btn_by_content(mesh_content: gr.Model3D, text_content: gr.Textbox):
     return gr.Button(interactive=interactive)
 
 
-with gr.Blocks(
-    delete_cache=(43200, 43200), theme=Default(primary_hue=slate)
-) as demo:
+with gr.Blocks(delete_cache=(43200, 43200), theme=custom_theme) as demo:
     gr.Markdown(
-        f"""
-        ## ***EmbodiedGen***: Texture Generation \n
-        version: {VERSION} \n
-    """
+        """
+        ## ***EmbodiedGen***: Text-to-3D Asset
+        **🔖 Version**: {VERSION}
+        <p style="display: flex; gap: 10px; flex-wrap: nowrap;">
+            <a href="https://horizonrobotics.github.io/robot_lab/embodied_gen/index.html">
+                <img alt="🌐 Project Page" src="https://img.shields.io/badge/🌐-Project_Page-blue">
+            </a>
+            <a href="https://arxiv.org/abs/xxxx.xxxxx">
+                <img alt="📄 arXiv" src="https://img.shields.io/badge/📄-arXiv-b31b1b">
+            </a>
+            <a href="https://github.com/horizon-research/EmbodiedGen">
+                <img alt="💻 GitHub" src="https://img.shields.io/badge/GitHub-000000?logo=github">
+            </a>
+            <a href="https://www.youtube.com/watch?v=SnHhzHeb_aI">
+                <img alt="🎥 Video" src="https://img.shields.io/badge/🎥-Video-red">
+            </a>
+        </p>
+
+        🎨 Generate visually rich textures for 3D mesh.
+
+        """.format(
+            VERSION=VERSION
+        ),
+        elem_classes=["header"],
     )
     gr.HTML(image_css)
     gr.HTML(lighting_css)
@@ -123,7 +140,7 @@ with gr.Blocks(
                 )
 
             generate_mv_btn = gr.Button(
-                "Generate MV Images(~1min)",
+                "🎨 1. Generate MV Images(~1min)",
                 variant="primary",
                 interactive=False,
             )
@@ -206,12 +223,12 @@ with gr.Blocks(
 
                 with gr.Column(scale=1):
                     texture_bake_btn = gr.Button(
-                        "Texture Baking(~2min)",
+                        "🛠️ 2. Texture Baking(~2min)",
                         variant="primary",
                         interactive=False,
                     )
                     download_btn = gr.DownloadButton(
-                        label="Download Mesh",
+                        label="⬇️ 3. Download Mesh",
                         variant="primary",
                         interactive=False,
                     )

common.py

@@ -30,6 +30,8 @@ import torch
 import torch.nn.functional as F
 import trimesh
 from easydict import EasyDict as edict
+from gradio.themes import Soft
+from gradio.themes.utils.colors import gray, neutral, slate, stone, teal, zinc
 from PIL import Image
 from embodied_gen.data.backproject_v2 import entrypoint as backproject_api
 from embodied_gen.data.differentiable_render import entrypoint as render_api
@@ -233,6 +235,14 @@ height: 100% !important;
 </style>
 """
 
+custom_theme = Soft(
+    primary_hue=stone,
+    secondary_hue=gray,
+    radius_size="md",
+    text_size="sm",
+    spacing_size="sm",
+)
+
 
 def start_session(req: gr.Request) -> None:
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
@@ -670,8 +680,8 @@ def text2image_fn(
     ip_adapt_scale: float = 0.3,
     image_wh: int | tuple[int, int] = [1024, 1024],
     rmbg_tag: str = "rembg",
-    n_sample: int = 3,
     seed: int = None,
+    n_sample: int = 3,
     req: gr.Request = None,
 ):
     if isinstance(image_wh, int):

embodied_gen/models/text_model.py

@@ -16,10 +16,10 @@
 
 
 import logging
+import random
 
-import torch
 import numpy as np
-import random
+import torch
 from diffusers import (
     AutoencoderKL,
     EulerDiscreteScheduler,
@@ -143,9 +143,10 @@ def text2img_gen(
     seed: int = None,
 ) -> list[Image.Image]:
     prompt = "Single " + prompt + ", in the center of the image"
-    prompt += ", high quality, high resolution, best quality, white background, 3D style,"  # noqa
+    prompt += ", high quality, high resolution, best quality, white background, 3D style"  # noqa
     logger.info(f"Processing prompt: {prompt}")
 
+    generator = None
     if seed is not None:
         generator = torch.Generator(pipeline.device).manual_seed(seed)
         torch.manual_seed(seed)

embodied_gen/scripts/imageto3d.py

@@ -70,7 +70,9 @@ IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
 RBG_REMOVER = RembgRemover()
 RBG14_REMOVER = BMGG14Remover()
 SAM_PREDICTOR = SAMPredictor(model_type="vit_h", device="cpu")
-PIPELINE = TrellisImageTo3DPipeline.from_pretrained("microsoft/TRELLIS-image-large")
+PIPELINE = TrellisImageTo3DPipeline.from_pretrained(
+    "microsoft/TRELLIS-image-large"
+)
 PIPELINE.cuda()
 SEG_CHECKER = ImageSegChecker(GPT_CLIENT)
 GEO_CHECKER = MeshGeoChecker(GPT_CLIENT)