add glb vis

app.py

@@ -12,6 +12,7 @@ import torch.nn.functional as F
 import torchvision.transforms.functional as TF
 import rembg
 import gradio as gr
+from gradio_litmodel3d import LitModel3D
 from dva.io import load_from_config
 from dva.ray_marcher import RayMarcher
 from dva.visualize import visualize_primvolume, visualize_video_primvolume
@@ -75,7 +76,7 @@ model_primx = load_from_config(config.model)
 rembg_session = rembg.new_session()
 
 # process function
-def process(input_image, input_num_steps=25, input_seed=42, input_cfg=6.0):
+def process(input_image, input_num_steps, input_seed=42, input_cfg=6.0):
     # seed
     torch.manual_seed(input_seed)
 
@@ -175,23 +176,31 @@ with block:
             # input image
             input_image = gr.Image(label="image", type='pil')
             # inference steps
-            input_num_steps = gr.Slider(label="inference steps", minimum=1, maximum=100, step=1, value=25)
+            input_num_steps = gr.Radio(choices=[25, 50, 100], label="DDIM steps")
             # random seed
-            input_cfg = gr.Slider(label="CFG scale", minimum=0, maximum=15, step=1, value=6)
+            input_cfg = gr.Slider(label="CFG scale", minimum=0, maximum=15, step=0.5, value=6)
             # random seed
-            input_seed = gr.Slider(label="random seed", minimum=0, maximum=100000, step=1, value=42)
+            input_seed = gr.Slider(label="random seed", minimum=0, maximum=10000, step=1, value=42)
             # gen button
             button_gen = gr.Button("Generate")
 
         with gr.Column(scale=1):
             with gr.Tab("Video"):
-                # final video results
-                output_rgb_video = gr.Video(label="video")
-                output_prim_video = gr.Video(label="video")
-                output_mat_video = gr.Video(label="video")
+                with gr.Row():
+                    # final video results
+                    output_rgb_video = gr.Video(label="video")
+                    output_prim_video = gr.Video(label="video")
+                    output_mat_video = gr.Video(label="video")
             with gr.Tab("GLB"):
                 # glb file
-                output_glb = gr.File(label="glb")
+                output_glb = LitModel3D(
+                    label="3D GLB Model",
+                    visible=True,
+                    clear_color=[0.0, 0.0, 0.0, 0.0],
+                    tonemapping="aces",
+                    contrast=1.0,
+                    scale=1.0,
+                )
 
         button_gen.click(process, inputs=[input_image, input_num_steps, input_seed, input_cfg], outputs=[output_rgb_video, output_prim_video, output_mat_video, output_glb])
     

requirements.txt

@@ -13,11 +13,12 @@ git+https://github.com/NVlabs/nvdiffrast/
 scikit-learn
 open_clip_torch
 triton==2.1.0
-rembg
+rembg[gpu]
 gradio
 tqdm
 transformers==4.40.1
 diffusers==0.19.3
 ninja
 imageio
-imageio-ffmpeg
+imageio-ffmpeg
+gradio-litmodel3d==0.0.1