Update

README.md

@@ -1,11 +1,18 @@
 ---
 title: SceneDiffuserDemo
-emoji: 🌍
-colorFrom: green
-colorTo: indigo
+emoji: 🛋
+colorFrom: yellow
+colorTo: pink
 sdk: gradio
-sdk_version: 3.16.1
+sdk_version: 3.18.0
 app_file: app.py
+tags:
+  - 3D
+  - Scene Understanding
+  - Diffusion
+  - Generation
+  - Optimization
+  - Planning
 pinned: false
 ---
 

app.py

@@ -13,7 +13,7 @@ with gr.Blocks(css='style.css') as demo:
         gr.HTML(value="<p align='center' style='font-size: 1.2em; color: #485fc7;'><a href='https://arxiv.org/abs/2301.06015' target='_blank'>arXiv</a> | <a href='https://scenediffuser.github.io/' target='_blank'>Project Page</a> | <a href='https://github.com/scenediffuser/Scene-Diffuser' target='_blank'>Code</a></p>")
         gr.Markdown("<p align='center'><i>\"SceneDiffuser provides a unified model for solving scene-conditioned generation, optimization, and planning.\"</i></p>")
 
-        ## five task
+        ## five tasks
         ## pose generation
         with gr.Tab("Pose Generation"):
             with gr.Row():
@@ -32,33 +32,23 @@ with gr.Blocks(css='style.css') as demo:
         button1.click(IF.pose_generation, inputs=input1, outputs=[image1])
 
         ## motion generation
-        # with gr.Tab("Motion Generation"):
-        #     with gr.Row():
-        #         with gr.Column(scale=2):
-        #             selector2 = gr.Dropdown(choices=['MPH16', 'MPH1Library', 'N0SittingBooth', 'N3OpenArea'], label='Scenes', value='MPH16', interactive=True)
-        #             with gr.Row():
-        #                 sample2 = gr.Slider(minimum=1, maximum=8, step=1, label='Count', interactive=True, value=1)
-        #                 seed2 = gr.Slider(minimum=0, maximum=2 ** 16, step=1, label='Seed', interactive=True, value=2023)
-        #             with gr.Row():
-        #                 withstart = gr.Checkbox(label='With Start', interactive=True, value=False)
-        #             opt2 = gr.Checkbox(label='Optimizer Guidance', interactive=True, value=True)
-        #             scale_opt2 = gr.Slider(minimum=0.1, maximum=9.9, step=0.1, label='Scale', interactive=True, value=1.1)
-        #             button2 = gr.Button("Run")
-        #         with gr.Column(scale=3):
-        #             image2 = gr.Image(label="Result")
-        # input2 = [selector2, sample2, seed2, withstart, opt2, scale_opt2]
-        # button2.click(IF.motion_generation, inputs=input2, outputs=image2)
         with gr.Tab("Motion Generation"):
             with gr.Row():
                 with gr.Column(scale=2):
-                    input2 = [
-                        gr.Dropdown(choices=['MPH16', 'MPH1Library', 'N0SittingBooth', 'N3OpenArea'], label='Scenes')
-                    ]
-                    button2 = gr.Button("Generate")
-                    gr.HTML("<p style='font-size: 0.9em; color: #555555;'>Notes: the output results are pre-sampled results. We will deploy a real-time model for this task soon.</p>")
+                    selector2 = gr.Dropdown(choices=['MPH16', 'MPH1Library', 'N0SittingBooth', 'N3OpenArea'], label='Scenes', value='MPH16', interactive=True)
+                    with gr.Row():
+                        sample2 = gr.Slider(minimum=1, maximum=2, step=1, label='Count', interactive=True, value=1)
+                        seed2 = gr.Slider(minimum=0, maximum=2 ** 16, step=1, label='Seed', interactive=True, value=2023)
+                    with gr.Row():
+                        withstart = gr.Checkbox(label='With Start', interactive=True, value=False)
+                    opt2 = gr.Checkbox(label='Optimizer Guidance', interactive=True, value=False)
+                    scale_opt2 = gr.Slider(minimum=0.1, maximum=9.9, step=0.1, label='Scale', interactive=True, value=1.1)
+                    button2 = gr.Button("Run")
                 with gr.Column(scale=3):
-                    output2 = gr.Image(label="Result")
-        button2.click(IF.motion_generation, inputs=input2, outputs=output2)
+                    image2 = gr.Gallery(label="Image [Result]").style(grid=[1], height="50")
+                    gr.HTML("<p style='font-size: 0.9em; color: #555555;'>Notes: For motion generation, it will take a long time to do sampleing and rendering, especifically when you tick optimizer guidance.</p>")
+        input2 = [selector2, sample2, seed2, withstart, opt2, scale_opt2]
+        button2.click(IF.motion_generation, inputs=input2, outputs=image2)
         
         ## grasp generation
         with gr.Tab("Grasp Generation"):

interface.py

@@ -4,6 +4,8 @@ import torch
 import hydra
 import numpy as np
 import zipfile
+import time
+import uuid
 
 from typing import Any
 from hydra import compose, initialize
@@ -24,9 +26,9 @@ def model_weight_path(task, has_observation=False):
     if task == 'pose_gen':
         return hf_hub_download('SceneDiffuser/SceneDiffuser', 'weights/2022-11-09_11-22-52_PoseGen_ddm4_lr1e-4_ep100/ckpts/model.pth')
     elif task == 'motion_gen' and has_observation == True:
-        return hf_hub_download('SceneDiffuser/SceneDiffuser', 'weights//ckpts/model.pth')
+        return hf_hub_download('SceneDiffuser/SceneDiffuser', 'weights/2022-11-09_14-28-12_MotionGen_ddm_T200_lr1e-4_ep300_obser/ckpts/model.pth')
     elif task == 'motion_gen' and has_observation == False:
-        return hf_hub_download('SceneDiffuser/SceneDiffuser', 'weights//ckpts/model.pth')
+        return hf_hub_download('SceneDiffuser/SceneDiffuser', 'weights/2022-11-09_12-54-50_MotionGen_ddm_T200_lr1e-4_ep300/ckpts/model.pth')
     elif task == 'path_planning':
         return hf_hub_download('SceneDiffuser/SceneDiffuser', 'weights/2022-11-25_20-57-28_Path_ddm4_LR1e-4_E100_REL/ckpts/model.pth')
     else:
@@ -140,7 +142,10 @@ def _planning(cfg: DictConfig, scene: str) -> Any:
 
 
 ## interface for five task
-## real-time model: pose generation, path planning
+## real-time model:
+##   - pose generation
+##   - motion generation
+##   - path planning
 def pose_generation(scene, count, seed, opt, scale) -> Any:
     scene_model_weight_path = pretrain_pointtrans_weight_path()
     data_dir, smpl_dir, prox_dir, vposer_dir = pose_motion_data_path()
@@ -181,22 +186,63 @@ def pose_generation(scene, count, seed, opt, scale) -> Any:
     hydra.core.global_hydra.GlobalHydra.instance().clear()
     return res
 
-def motion_generation(scene):
-    assert isinstance(scene, str)
-    cnt = {
-        'MPH1Library': 3,
-        'MPH16': 6,
-        'N0SittingBooth': 7,
-        'N3OpenArea': 5
-    }[scene]
-
-    res = f"./results/motion_generation/results/{scene}/{random.randint(0, cnt-1)}.gif"
-    if not os.path.exists(res):
-        results_path = hf_hub_download('SceneDiffuser/SceneDiffuser', 'results/motion_generation/results.zip')
-        os.makedirs('./results/motion_generation/', exist_ok=True)
-        with zipfile.ZipFile(results_path, 'r') as zip_ref:
-            zip_ref.extractall('./results/motion_generation/')
+def motion_generation(scene, count, seed, withstart, opt, scale) -> Any:
+    scene_model_weight_path = pretrain_pointtrans_weight_path()
+    data_dir, smpl_dir, prox_dir, vposer_dir = pose_motion_data_path()
+    override_config = [
+        "diffuser=ddpm",
+        "diffuser.steps=200",
+        "model=unet",
+        "model.use_position_embedding=true",
+        f"model.scene_model.pretrained_weights={scene_model_weight_path}",
+        "task=motion_gen",
+        f"task.has_observation={withstart}",
+        "task.dataset.repr_type=absolute",
+        "task.dataset.frame_interval_test=20",
+        "task.visualizer.name=MotionGenVisualizerHF",
+        f"task.visualizer.ksample={count}",
+        f"task.dataset.data_dir={data_dir}",
+        f"task.dataset.smpl_dir={smpl_dir}",
+        f"task.dataset.prox_dir={prox_dir}",
+        f"task.dataset.vposer_dir={vposer_dir}",
+    ]
+    if opt == True:
+        override_config += [
+            "optimizer=motion_in_scene",
+            "optimizer.scale_type=div_var",
+            f"optimizer.scale={scale}",
+            "optimizer.vposer=false",
+            "optimizer.contact_weight=0.02",
+            "optimizer.collision_weight=1.0",
+            "optimizer.smoothness_weight=0.001",
+            "optimizer.frame_interval=1",
+        ]
     
+    initialize(config_path="./scenediffuser/configs", version_base=None)
+    config = compose(config_name="default", overrides=override_config)
+
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    res_gifs = _sampling(config, scene)
+
+    ## save sampled motion as .gif file
+    datestr = time.strftime("%Y-%m-%d", time.localtime(time.time()))
+    target_dir = os.path.join('./results/motion_generation/', f'd-{datestr}')
+    os.makedirs(target_dir, exist_ok=True)
+    res = []
+    uuid_str = uuid.uuid4()
+    for i, imgs in enumerate(res_gifs):
+        target_path = os.path.join(target_dir, f'{uuid_str}--{i}.gif')
+        imgs = [im.resize((720, 405)) for im in imgs] # resize image for low resolution to save space
+        img, *img_rest = imgs
+        img.save(fp=target_path, format='GIF', append_images=img_rest, save_all=True, duration=33.33, loop=0)
+        res.append(target_path)
+
+    hydra.core.global_hydra.GlobalHydra.instance().clear()
     return res
 
 def grasp_generation(case_id):

scenediffuser

@@ -1 +1 @@
-Subproject commit 2dfda611966515f5d2aff8ac45225f7f77bf9541
+Subproject commit ddcba15d05dcb52f3f3b576f7c60e5e255baa584