from pathlib import Path import os import sys sys.path.append(str(Path(os.path.abspath('')))) import torch import numpy as np from tools.genrl_utils import ViCLIPGlobalInstance import time import torchvision from huggingface_hub import hf_hub_download import spaces # IMPORT HF_TOKEN hf_token = os.environ['HF_TOKEN'] def save_videos(batch_tensors, savedir, filenames, fps=10): # b,samples,c,t,h,w n_samples = batch_tensors.shape[1] for idx, vid_tensor in enumerate(batch_tensors): video = vid_tensor.detach().cpu() video = torch.clamp(video.float(), 0., 1.) video = video.permute(1, 0, 2, 3, 4) # t,n,c,h,w frame_grids = [torchvision.utils.make_grid(framesheet, nrow=int(n_samples)) for framesheet in video] #[3, 1*h, n*w] grid = torch.stack(frame_grids, dim=0) # stack in temporal dim [t, 3, n*h, w] grid = (grid * 255).to(torch.uint8).permute(0, 2, 3, 1) savepath = os.path.join(savedir, f"{filenames[idx]}.mp4") torchvision.io.write_video(savepath, grid, fps=fps, video_codec='h264', options={'crf': '10'}) class Text2Video(): def __init__(self,result_dir='./tmp/',gpu_num=1) -> None: model_folder = str(Path(os.path.abspath('')) / 'models') model_filename = 'genrl_stickman_500k_2.pt' if not os.path.isfile(os.path.join(model_folder, model_filename)): self.download_model(model_folder, model_filename) if not os.path.isfile(os.path.join(model_folder, 'InternVideo2-stage2_1b-224p-f4.pt')): self.download_internvideo2(model_folder) self.agent = torch.load(os.path.join(model_folder, model_filename),map_location='cpu') model_name = 'internvideo2' # Get ViCLIP viclip_global_instance = ViCLIPGlobalInstance(model_name) if not viclip_global_instance._instantiated: print("Instantiating InternVideo2") viclip_global_instance.instantiate(device='cpu') self.clip = viclip_global_instance.viclip self.tokenizer = viclip_global_instance.viclip_tokenizer self.result_dir = result_dir if not os.path.exists(self.result_dir): os.mkdir(self.result_dir) @spaces.GPU def get_prompt(self, prompt, duration): torch.cuda.empty_cache() self.agent.to('cuda') self.clip.to('cuda') print('start:', prompt, time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))) start = time.time() prompt_str = prompt.replace("/", "_slash_") if "/" in prompt else prompt prompt_str = prompt_str.replace(" ", "_") if " " in prompt else prompt_str labels_list = [prompt_str] with torch.no_grad(): wm = world_model = self.agent.wm connector = self.agent.wm.connector decoder = world_model.heads['decoder'] n_frames = connector.n_frames # Get text(video) embed text_feat = [] for text in labels_list: with torch.no_grad(): text_feat.append(self.clip.get_txt_feat(text,)) text_feat = torch.stack(text_feat, dim=0).to('cuda') video_embed = text_feat B = video_embed.shape[0] T = 1 # Get actions video_embed = video_embed.repeat(1, duration, 1) with torch.no_grad(): # Imagine prior = wm.connector.video_imagine(video_embed, None, sample=False, reset_every_n_frames=False, denoise=True) # Decode prior_recon = decoder(wm.decoder_input_fn(prior))['observation'].mean + 0.5 save_videos(prior_recon.unsqueeze(0), self.result_dir, filenames=[prompt_str], fps=15) print(f"Saved in {prompt_str}.mp4. Time used: {(time.time() - start):.2f} seconds") # Offload GPU self.agent.to('cpu') self.clip.to('cpu') return os.path.join(self.result_dir, f"{prompt_str}.mp4") def download_model(self, model_folder, model_filename): REPO_ID = 'mazpie/genrl_models' filename_list = [model_filename] if not os.path.exists(model_folder): os.makedirs(model_folder) for filename in filename_list: local_file = os.path.join(model_folder, filename) if not os.path.exists(local_file): hf_hub_download(repo_id=REPO_ID, filename=filename, local_dir=model_folder, local_dir_use_symlinks=False) def download_internvideo2(self, model_folder): REPO_ID = 'OpenGVLab/InternVideo2-Stage2_1B-224p-f4' filename_list = ['InternVideo2-stage2_1b-224p-f4.pt'] if not os.path.exists(model_folder): os.makedirs(model_folder) for filename in filename_list: local_file = os.path.join(model_folder, filename) if not os.path.exists(local_file): hf_hub_download(repo_id=REPO_ID, filename=filename, local_dir=model_folder, local_dir_use_symlinks=False, token=hf_token) if __name__ == '__main__': t2v = Text2Video() video_path = t2v.get_prompt('a black swan swims on the pond', 8) print('done', video_path)