Spaces:
Running
Running
| from calendar import EPOCH | |
| from geometry_utils import diffout2motion | |
| import gradio as gr | |
| import spaces | |
| import torch | |
| import random | |
| import os | |
| from pathlib import Path | |
| import smplx | |
| from body_renderer import get_render | |
| import joblib | |
| # import cv2 | |
| # import moderngl | |
| # ctx = moderngl.create_context(standalone=True) | |
| # print(ctx) | |
| access_token_smpl = os.environ.get('HF_SMPL_TOKEN') | |
| os.environ["PYOPENGL_PLATFORM"] = "egl" | |
| zero = torch.Tensor([0]).cuda() | |
| print(zero.device) # <-- 'cuda:0' 🤗 | |
| DEFAULT_TEXT = "do it slower " | |
| def get_smpl_models(): | |
| REPO_ID = 'athn-nik/smpl_models' | |
| from huggingface_hub import snapshot_download | |
| return snapshot_download(repo_id=REPO_ID, allow_patterns="smplh*", | |
| token=access_token_smpl) | |
| WEBSITE = ("""<div class="embed_hidden" style="text-align: center;"> | |
| <h1>MotionFix: Text-Driven 3D Human Motion Editing</h1> | |
| <h3> | |
| <a href="https://is.mpg.de/person/~nathanasiou" target="_blank" rel="noopener noreferrer">Nikos Athanasiou</a><sup>1</sup>, | |
| <a href="https://is.mpg.de/person/acseke" target="_blank" rel="noopener noreferrer">Alpar Cseke</a><sup>1</sup>, | |
| <br> | |
| <a href="https://ps.is.mpg.de/person/mdiomataris" target="_blank" rel="noopener noreferrer">Markos Diomataris</a><sup>1, 3</sup>, | |
| <a href="https://is.mpg.de/person/black" target="_blank" rel="noopener noreferrer">Michael J. Black</a><sup>1</sup>, | |
| <a href="https://imagine.enpc.fr/~varolg/" target="_blank" rel="noopener noreferrer">Gül Varol</a><sup>2</sup> | |
| </h3> | |
| <h3> | |
| <sup>1</sup>Max Planck Institute for Intelligent Systems, Tübingen, Germany; | |
| <sup>2</sup>LIGM, École des Ponts, Univ Gustave Eiffel, CNRS, France, | |
| <sup>3</sup>ETH Zürich, Switzerland | |
| </h3> | |
| </div> | |
| <div style="display:flex; gap: 0.3rem; justify-content: center; align-items: center;" align="center"> | |
| <a href='https://arxiv.org/pdf/2408.00712'><img src='https://img.shields.io/badge/Arxiv-2405.20340-A42C25?style=flat&logo=arXiv&logoColor=A42C25'></a> | |
| <a href='https://motionfix.is.tue.mpg.de'><img src='https://img.shields.io/badge/Project-Page-%23df5b46?style=flat&logo=Google%20chrome&logoColor=%23df5b46'></a> | |
| <a href='https://www.youtube.com/watch?v=cFa6V6Ua-TY'><img src='https://img.shields.io/badge/YouTube-red?style=flat&logo=youtube&logoColor=white'></a> | |
| </div> | |
| """) | |
| CREDITS=("""<div class="embed_hidden" style="text-align: center;"> | |
| <h3> | |
| The renderer of this demo is adapted from the render of | |
| <a href="https://geometry.stanford.edu/projects/humor/" target="_blank" rel="noopener noreferrer">HuMoR</a> | |
| with the help of <a href="https://ps.is.mpg.de/person/trakshit" target="_blank" rel="noopener noreferrer">Tithi Rakshit</a> :) | |
| </h3> | |
| """) | |
| WEB_source = ("""<div class="embed_hidden" style="text-align: center;"> | |
| <h1>Pick a motion to edit!</h1> | |
| <h3> | |
| Here you should pick a source motion | |
| <hr class="double"> | |
| </h3> | |
| </div> | |
| """) | |
| WEB_target = ("""<div class="embed_hidden" style="text-align: center;"> | |
| <h1>Now type the text to edit that motion!</h1> | |
| <h3> | |
| Here you should get the generated motion! | |
| <hr class="double"> | |
| </h3> | |
| </div> | |
| """) | |
| def greet(n): | |
| print(zero.device) # <-- 'cuda:0' 🤗 | |
| try: | |
| number = float(n) | |
| except ValueError: | |
| return "Invalid input. Please enter a number." | |
| return f"Hello {zero + number} Tensor" | |
| def clear(): | |
| return "" | |
| def show_video(input_text, key_to_use): | |
| from normalization import Normalizer | |
| normalizer = Normalizer() | |
| from diffusion import create_diffusion | |
| from text_encoder import ClipTextEncoder | |
| from tmed_denoiser import TMED_denoiser | |
| model_ckpt = download_models() | |
| checkpoint = torch.load(model_ckpt) | |
| motion_to_edit = download_motion_from_dataset(key_to_use) | |
| ds_sample = joblib.load(motion_to_edit) | |
| source_motion_norm = ds_sample['source_feats_norm'].to('cuda') | |
| seqlen_tgt = ds_sample['target_feats_norm'].shape[0] | |
| seqlen_src = ds_sample['source_feats_norm'].shape[0] | |
| # import ipdb; ipdb.set_trace() | |
| checkpoint = {k.replace('denoiser.', ''): v for k, v in checkpoint.items()} | |
| tmed_denoiser = TMED_denoiser().to('cuda') | |
| tmed_denoiser.load_state_dict(checkpoint, strict=False) | |
| tmed_denoiser.eval() | |
| text_encoder = ClipTextEncoder() | |
| texts_cond = [input_text] | |
| diffusion_process = create_diffusion(timestep_respacing=None, | |
| learn_sigma=False, sigma_small=True, | |
| diffusion_steps=300, | |
| noise_schedule='squaredcos_cap_v2', | |
| predict_xstart=True) | |
| bsz = 1 | |
| no_of_texts = len(texts_cond) | |
| texts_cond = ['']*no_of_texts + texts_cond | |
| texts_cond = ['']*no_of_texts + texts_cond | |
| text_emb, text_mask = text_encoder(texts_cond) | |
| cond_emb_motion = source_motion_norm.unsqueeze(0).permute(1, 0, 2) | |
| cond_motion_mask = torch.ones((bsz, seqlen_src), | |
| dtype=bool, device='cuda') | |
| mask_target = torch.ones((bsz, seqlen_tgt), | |
| dtype=bool, device='cuda') | |
| diff_out = tmed_denoiser._diffusion_reverse(text_emb.to(cond_emb_motion.device), | |
| text_mask.to(cond_emb_motion.device), | |
| cond_emb_motion, | |
| cond_motion_mask, | |
| mask_target, | |
| diffusion_process, | |
| init_vec=None, | |
| init_from='noise', | |
| gd_text=2.0, | |
| gd_motion=2.0, | |
| steps_num=300) | |
| edited_motion = diffout2motion(diff_out, normalizer).squeeze() | |
| # import ipdb; ipdb.set_trace() | |
| # aitrenderer = get_renderer() | |
| # SMPL_LAYER = SMPLLayer(model_type='smplh', ext='npz', gender='neutral') | |
| # edited_mot_to_render = pack_to_render(rots=edited_motion[..., 3:], | |
| # trans=edited_motion[..., :3]) | |
| SMPL_MODELS_PATH = str(Path(get_smpl_models())) | |
| body_model=smplx.SMPLHLayer(f"{SMPL_MODELS_PATH}/smplh", | |
| model_type='smplh', | |
| gender='neutral',ext='npz') | |
| # run_smpl_fwd_verticesbody_model, body_transl, body_orient, body_pose, | |
| import random | |
| xx = random.randint(1, 1000) | |
| # edited_mot_to_render | |
| from body_renderer import get_render | |
| from transform3d import transform_body_pose | |
| edited_motion_aa = transform_body_pose(edited_motion[:, 3:], | |
| '6d->aa') | |
| if os.path.exists('./output_movie.mp4'): | |
| os.remove('./output_movie.mp4') | |
| fname = get_render(body_model, | |
| [edited_motion[..., :3].detach().cpu()], | |
| [edited_motion_aa[..., :3].detach().cpu()], | |
| [edited_motion_aa[..., 3:].detach().cpu()], | |
| output_path='./output_movie.mp4', | |
| text='', colors=['sky blue']) | |
| # import ipdb; ipdb.set_trace() | |
| # fname = render_motion(AIT_RENDERER, [edited_mot_to_render], | |
| # f"movie_example--{str(xx)}", | |
| # pose_repr='aa', | |
| # color=[color_map['generated']], | |
| # smpl_layer=SMPL_LAYER) | |
| print(fname) | |
| print(os.path.abspath(fname)) | |
| return fname | |
| from huggingface_hub import hf_hub_download | |
| def download_models(): | |
| REPO_ID = 'athn-nik/example-model' | |
| return hf_hub_download(REPO_ID, filename="tmed_compressed.ckpt") | |
| def download_motion_from_dataset(key_to_dl): | |
| REPO_ID = 'athn-nik/example-model' | |
| from huggingface_hub import snapshot_download | |
| keytodl = key_to_dl | |
| keytodl = '000008' | |
| path_for_ds = snapshot_download(repo_id=REPO_ID, | |
| allow_patterns=f"dataset_inputs/{keytodl}", | |
| token=access_token_smpl) | |
| path_for_ds_sample = path_for_ds + f'/dataset_inputs/{keytodl}.pth.tar' | |
| return path_for_ds_sample | |
| def download_tmr(): | |
| REPO_ID = 'athn-nik/example-model' | |
| # return hf_hub_download(REPO_ID, filename="min_checkpoint.ckpt") | |
| from huggingface_hub import snapshot_download | |
| return snapshot_download(repo_id=REPO_ID, allow_patterns="tmr*", | |
| token=access_token_smpl) | |
| def download_motionfix(): | |
| REPO_ID = 'athn-nik/example-model' | |
| # return hf_hub_download(REPO_ID, filename="min_checkpoint.ckpt") | |
| from huggingface_hub import snapshot_download | |
| return snapshot_download(repo_id=REPO_ID, allow_patterns="motionfix*", | |
| token=access_token_smpl) | |
| def download_motionfix_dataset(): | |
| REPO_ID = 'athn-nik/example-model' | |
| dataset_downloaded_path = hf_hub_download(REPO_ID, filename="tmed_compressed.ckpt") | |
| dataset_dict = joblib.load(dataset_downloaded_path) | |
| return dataset_dict | |
| def download_embeddings(): | |
| REPO_ID = 'athn-nik/example-model' | |
| # return hf_hub_download(REPO_ID, filename="min_checkpoint.ckpt") | |
| from huggingface_hub import snapshot_download | |
| return snapshot_download(repo_id=REPO_ID, allow_patterns="embeddings*", | |
| token=access_token_smpl) | |
| MFIX_p = download_motionfix() + '/motionfix' | |
| SOURCE_MOTS_p = download_embeddings() + '/embeddings' | |
| MFIX_DATASET_DICT = download_motionfix_dataset() | |
| import gradio as gr | |
| def clear(): | |
| return "" | |
| def random_source_motion(set_to_pick): | |
| # import ipdb;ipdb.set_trace() | |
| mfix_train, mfix_test = load_motionfix(MFIX_p) | |
| if set_to_pick == 'all': | |
| current_set = mfix_test | mfix_train | |
| elif set_to_pick == 'train': | |
| current_set = mfix_train | |
| elif set_to_pick == 'test': | |
| current_set = mfix_test | |
| import random | |
| random_key = random.choice(list(current_set.keys())) | |
| curvid = current_set[random_key]['motion_a'] | |
| text_annot = current_set[random_key]['annotation'] | |
| return curvid, text_annot, random_key, text_annot | |
| def retrieve_video(retrieve_text): | |
| tmr_text_encoder = get_tmr_model(download_tmr()) | |
| # import ipdb;ipdb.set_trace() | |
| # text_encoded = tmr_text_encoder([retrieve_text]) | |
| motion_embeds = None | |
| from gen_utils import read_json | |
| import numpy as np | |
| motion_embeds = torch.load(SOURCE_MOTS_p+'/source_motions_embeddings.pt') | |
| motion_keyids =np.array(read_json(SOURCE_MOTS_p+'/keyids_embeddings.json')) | |
| mfix_train, mfix_test = load_motionfix(MFIX_p) | |
| all_mots = mfix_test | mfix_train | |
| scores = tmr_text_encoder.compute_scores(retrieve_text, embs=motion_embeds) | |
| sorted_idxs = np.argsort(-scores) | |
| best_keyids = motion_keyids[sorted_idxs] | |
| # best_scores = scores[sorted_idxs] | |
| top_mot = best_keyids[0] | |
| curvid = all_mots[top_mot]['motion_a'] | |
| text_annot = all_mots[top_mot]['annotation'] | |
| return curvid, text_annot | |
| with gr.Blocks(css=""" | |
| .gradio-row { | |
| display: flex; | |
| gap: 20px; | |
| } | |
| .gradio-column { | |
| flex: 1; | |
| } | |
| .gradio-container { | |
| display: flex; | |
| flex-direction: column; | |
| gap: 10px; | |
| } | |
| .gradio-button-row { | |
| display: flex; | |
| gap: 10px; | |
| } | |
| .gradio-textbox-row { | |
| display: flex; | |
| gap: 10px; | |
| align-items: center; | |
| } | |
| .gradio-edit-row { | |
| gap: 10px; | |
| align-items: center; | |
| } | |
| .gradio-textbox-with-button { | |
| display: flex; | |
| align-items: center; | |
| } | |
| .gradio-textbox-with-button input { | |
| flex-grow: 1; | |
| } | |
| """) as demo: | |
| gr.Markdown(WEBSITE) | |
| random_key_state = gr.State() | |
| with gr.Row(elem_id="gradio-row"): | |
| with gr.Column(scale=5, elem_id="gradio-column"): | |
| gr.Markdown(WEB_source) | |
| with gr.Row(elem_id="gradio-button-row"): | |
| # iterative_button = gr.Button("Iterative") | |
| # retrieve_button = gr.Button("TMRetrieve") | |
| random_button = gr.Button("Random") | |
| with gr.Row(elem_id="gradio-textbox-row"): | |
| with gr.Column(scale=5, elem_id="gradio-textbox-with-button"): | |
| # retrieve_text = gr.Textbox(placeholder="Type the text for the motion you want to Retrieve:", | |
| # show_label=True, label="Retrieval Text", | |
| # value=DEFAULT_TEXT) | |
| clear_button_retrieval = gr.Button("Clear", scale=0) | |
| with gr.Row(elem_id="gradio-textbox-row"): | |
| suggested_edit_text = gr.Textbox(placeholder="Texts likely to edit the motion:", | |
| show_label=True, label="Suggested Edit Text", | |
| value='') | |
| xxx = 'https://motion-editing.s3.eu-central-1.amazonaws.com/collection_wo_walks_runs/rendered_pairs/011327_120_240-002682_120_240.mp4' | |
| set_to_pick = gr.Radio(['all', 'train', 'test'], | |
| value='all', | |
| label="Set to pick from", | |
| info="Motion will be picked from whole dataset or test or train data.") | |
| # import ipdb; ipdb.set_trace() | |
| retrieved_video_output = gr.Video(label="Retrieved Motion", | |
| # value=xxx, | |
| height=360, width=480) | |
| with gr.Column(scale=5, elem_id="gradio-column"): | |
| gr.Markdown(WEB_target) | |
| with gr.Row(elem_id="gradio-edit-row"): | |
| clear_button_edit = gr.Button("Clear", scale=0) | |
| edit_button = gr.Button("Edit", scale=0) | |
| with gr.Row(elem_id="gradio-textbox-row"): | |
| input_text = gr.Textbox(placeholder="Type the edit text you want:", | |
| show_label=False, label="Input Text", | |
| value=DEFAULT_TEXT) | |
| video_output = gr.Video(label="Generated Video", height=360, | |
| width=480) | |
| def process_and_show_video(input_text, random_key_state): | |
| fname = show_video(input_text, random_key_state) | |
| return fname | |
| def process_and_retrieve_video(input_text): | |
| fname = retrieve_video(input_text) | |
| return fname | |
| from retrieval_loader import get_tmr_model | |
| from dataset_utils import load_motionfix | |
| edit_button.click(process_and_show_video, inputs=[input_text, random_key_state], outputs=video_output) | |
| # retrieve_button.click(process_and_retrieve_video, inputs=retrieve_text, outputs=[retrieved_video_output, suggested_edit_text]) | |
| random_button.click(random_source_motion, inputs=set_to_pick, | |
| outputs=[retrieved_video_output, | |
| suggested_edit_text, | |
| random_key_state, | |
| input_text]) | |
| clear_button_edit.click(clear, outputs=input_text) | |
| # clear_button_retrieval.click(clear, outputs=retrieve_text) | |
| gr.Markdown(CREDITS) | |
| demo.launch(share=True) | |