Kleinhe/SemanticBoost
0
1import os, sys2import gradio as gr3from huggingface_hub import snapshot_download4css = """5.dfile {height: 85px}6.ov {height: 185px}7"""8 9 10from huggingface_hub import snapshot_download11from motion.visual_api import Visualize 12import torch13import json14from tqdm import tqdm15import imageio16 17with open("motion/path.json", "r") as f:18 json_dict = json.load(f)19 20def ref_video_fn(path_of_ref_video):21 if path_of_ref_video is not None:22 return gr.update(value=True)23 else:24 return gr.update(value=False)25 26def prepare():27 if not os.path.exists("body_models") or not os.path.exists("weights"):28 REPO_ID = 'Kleinhe/CAMD'29 snapshot_download(repo_id=REPO_ID, local_dir='./', local_dir_use_symlinks=False)30 31 if not os.path.exists("tada-extend"):32 import subprocess33 import platform34 command = "bash scripts/tada_goole.sh"35 subprocess.call(command, shell=platform.system() != 'Windows')36 37def demo(prompt, mode, condition, render_mode="joints", skip_steps=0, out_size=1024, tada_role=None):38 prompt = prompt39 if prompt is None:40 prompt = ""41 42 path = None43 out_paths = [None, None, None]44 joints_paths = [None, None, None]45 smpl_paths = [None, None, None]46 47 if tada_role == "None":48 tada_role = None49 50 for i in range(len(mode)):51 kargs = {52 "mode":mode[i],53 "device":"cuda" if torch.cuda.is_available() else "cpu",54 "condition":condition,55 "smpl_path":json_dict["smpl_path"],56 "skip_steps":skip_steps,57 "path":json_dict,58 "tada_base":json_dict["tada_base"],59 "tada_role":tada_role60 }61 visual = Visualize(**kargs)62 render_mode = render_mode63 64 joint_path = "results/joints/{}_joint.npy".format(mode[i])65 smpl_path = "results/smpls/{}_smpl.npy".format(mode[i])66 video_path = "results/motion/{}_video.gif".format(mode[i])67 68 output = visual.predict(prompt, path, render_mode, joint_path, smpl_path)69 70 if render_mode == "joints":71 pics = visual.joints_process(output, prompt)72 elif render_mode.startswith("pyrender"):73 meshes, _ = visual.get_mesh(output)74 pics = visual.pyrender_process(meshes, out_size, out_size)75 76 try:77 imageio.mimsave(video_path, pics, duration= 1000 / 20, loop=0)78 except:79 imageio.mimsave(video_path, pics, fps=20)80 81 if mode[i] == "cadm":82 out_paths[0] = video_path83 joints_paths[0] = joint_path84 smpl_paths[0] = smpl_path85 elif mode[i] == "cadm-augment":86 out_paths[1] = video_path87 joints_paths[1] = joint_path88 smpl_paths[1] = smpl_path89 elif mode[i] == "mdm":90 out_paths[2] = video_path91 joints_paths[2] = joint_path92 smpl_paths[2] = smpl_path93 94 return out_paths + joints_paths + smpl_paths95 96 97def t2m_demo():98 prepare()99 os.makedirs("results/motion", exist_ok=True)100 os.makedirs("results/joints", exist_ok=True)101 os.makedirs("results/smpls", exist_ok=True)102 103 tada_base = json_dict["tada_base"]104 files = os.listdir(os.path.join(tada_base, "MESH"))105 files = sorted(files)106 if files[0].startswith("."):107 files.pop(0)108 files = ["None"] + files109 110 with gr.Blocks(analytics_enabled=False, css=css) as t2m_interface:111 gr.Markdown("<div align='center'> <h2> 🤷♂️ SemanticBoost: Elevating Motion Generation with Augmented Textual Cues </span> </h2> \112 <a style='font-size:18px;' href='https://arxiv.org/abs/2310.20323'>Arxiv</a> \113 <a style='font-size:18px;' href='https://blackgold3.github.io/SemanticBoost/'>Homepage</a> \114 <a style='font-size:18px;' href='https://github.com/blackgold3/SemanticBoost'> Github </div>")115 116 with gr.Row().style(equal_height=True):117 with gr.Column(variant='panel'): 118 with gr.Tabs():119 with gr.TabItem('Settings'):120 with gr.Column(variant='panel'):121 with gr.Row():122 demo_mode = gr.CheckboxGroup(choices=['cadm', 'cadm-augment','mdm'], default=["cadm"], label='Mode', info="Choose models to run demos, more models cost more time.")123 skip_steps = gr.Number(value=0, label="Skip-Steps", info="The number of skip-steps during diffusion process (0 -> 999)", minimum=0, maximum=999, precision=0)124 125 with gr.Row():126 condition = gr.Radio(['text', 'uncond'], value='text', label='Condition', info="If sythesize motion with prompt?")127 out_size = gr.Number(value=256, label="Resolution", info="The resolution of output videos", minimum=128, maximum=2048, precision=0)128 129 with gr.Row():130 render_mode = gr.Radio(['joints','pyrender_fast', 'pyrender_slow'], value='joints', label='Render', info="If render results to 3D meshes? Pyrender need more time.")131 tada_role = gr.Dropdown(files, value="None", multiselect=False, label="TADA Role", info="Choose 3D role to render")132 133 with gr.Row():134 prompt = gr.Textbox(value=None, placeholder="120,A person walks forward and does a handstand.", label="Prompt for Model -> (Length,Text)")135 136 submit = gr.Button('Visualize', variant='primary')137 138 with gr.Column(variant='panel'): 139 with gr.Tabs():140 with gr.TabItem('Results'):141 with gr.Row():142 with gr.Column():143 gen_video = gr.Image(label="CADM", elem_classes="ov")144 with gr.Column():145 joint_file = gr.File(label="CADM-Joints", value=None, elem_classes="dfile")146 smpl_file = gr.File(label="CADM-SMPL", value=None, elem_classes="dfile")147 148 with gr.Row():149 with gr.Column():150 gen_video1 = gr.Image(label="CADM-Augment", elem_classes="ov")151 with gr.Column():152 joint_file1 = gr.File(label="CADM-Augment-Joints", value=None, elem_classes="dfile")153 smpl_file1 = gr.File(label="CADM-Augment-SMPL", value=None, elem_classes="dfile")154 155 with gr.Row():156 with gr.Column():157 gen_video2 = gr.Image(label="MDM", elem_classes="ov")158 with gr.Column():159 joint_file2 = gr.File(label="MDM-Joints", value=None, elem_classes="dfile")160 smpl_file2 = gr.File(label="MDM-SMPL", value=None, elem_classes="dfile")161 162 163 submit.click(164 fn=demo,165 inputs=[prompt,166 demo_mode,167 condition,168 render_mode, 169 skip_steps, 170 out_size,171 tada_role 172 ], 173 outputs=[gen_video, gen_video1, gen_video2, joint_file, joint_file1, joint_file2, smpl_file, smpl_file1, smpl_file2]174 )175 176 return t2m_interface177 178 179if __name__ == "__main__":180 demo = t2m_demo()181 demo.queue(max_size=10)182 demo.launch(debug=True)183 184 185 186