Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,938 Bytes
4f7b5ea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
'''
nohup python preprocess/extract-vae1.py 0 2 >> vae1.log 2>&1 &
nohup python preprocess/extract-vae1.py 1 2 >> vae1.log 2>&1 &
0/1: What part does this process do
2: It consists of two parts in total.
'''
root_mp4s="test_data/mp4root"
h,w=480,832
num_frames = 49
opt_root="output_root/vae1"
checkpoint_dir="/DATA/bvac/personal/wan21/Wan2.1-I2V-14B-480P"
import os,sys,traceback
import pdb
os.environ["CUDA_VISIBLE_DEVICES"]=sys.argv[1]
all=int(sys.argv[2])
i_part=int(os.environ["CUDA_VISIBLE_DEVICES"])
import pdb,torch
from wan.modules.vae import WanVAE
device="cuda"
vae = WanVAE(vae_pth=os.path.join(checkpoint_dir, 'Wan2.1_VAE.pth'),device=device)
from decord import VideoReader
import torchvision.transforms.functional as TF
def read_img(path):
vr = VideoReader(uri=path, height=-1, width=-1)
temp_frms = vr.get_batch([2])
return (TF.to_tensor(temp_frms.asnumpy().astype("float32")[0])/255).sub_(0.5).div_(0.5).to(device)
os.makedirs(opt_root,exist_ok=True)
def go(todos):
for path in todos:
try:
name=os.path.basename(path).replace(".mp4",".pt")
if os.path.exists("%s/%s"%(opt_root,name)):continue
img = read_img(path)
tensorr = vae.encode([
torch.concat([
torch.nn.functional.interpolate(
img[None].cpu(), size=(h, w), mode='bicubic').transpose(
0, 1),
torch.zeros(3, num_frames-1, h, w)
], dim=1).to(device)
])[0].cpu() # torch.Size([16, 21, 90, 160])#21->13
save_path="%s/%s"%(opt_root,name)
torch.save(tensorr, save_path)
except:
print(path,traceback.format_exc())
todo=[]
for name in os.listdir(root_mp4s):
todo.append("%s/%s"%(root_mp4s,name))
todo=sorted(todo)
todo=todo[i_part::all]
go(todo)
|