152 lines
5.1 KiB
Python
152 lines
5.1 KiB
Python
import os,sys
|
|
now_dir = os.path.dirname(os.path.abspath(__file__))
|
|
sys.path.append(now_dir)
|
|
|
|
import shutil
|
|
import math
|
|
import folder_paths
|
|
import numpy as np
|
|
from PIL import Image
|
|
from huggingface_hub import snapshot_download
|
|
from .infer import ViewCrafter
|
|
from viewcrafter.configs.infer_config import get_parser
|
|
|
|
output_dir = folder_paths.get_output_directory()
|
|
ckpt_dir = os.path.join(now_dir,"checkpoints")
|
|
default_traj_txt = "0 -40\n0 0\n0. -0.2"
|
|
|
|
|
|
|
|
class ViewCrafterTxTNode:
|
|
def __init__(self):
|
|
self.pvd = None
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required":{
|
|
"img":("IMAGE",),
|
|
"traj_txt":(["left","loop1","loop2","wave1","zoomin1"],),
|
|
"video_length":([25,16],),
|
|
"ddim_steps":("INT",{
|
|
"default":50,
|
|
}),
|
|
"center_scale":("FLOAT",{
|
|
"default":1.,
|
|
"min":0.0,
|
|
"max":2.0,
|
|
"display":"slider",
|
|
}),
|
|
"elevation":("FLOAT",{
|
|
"default":5.,
|
|
}),
|
|
"d_theta":("INT",{
|
|
"default":-30,
|
|
"min":-40,
|
|
"max":40,
|
|
"display":"slider",
|
|
}),
|
|
"d_phi":("INT",{
|
|
"default":45,
|
|
"min":-45,
|
|
"max":45,
|
|
"step":5,
|
|
"display":"slider",
|
|
}),
|
|
"d_r":("FLOAT",{
|
|
"default":-0.5,
|
|
"min":-0.5,
|
|
"max":0.5,
|
|
"display":"slider",
|
|
}),
|
|
"seed":("INT",{
|
|
"default":42,
|
|
}),
|
|
},
|
|
"optional":{
|
|
"custom_traj_txt":("STRING",{
|
|
"multiline": True,
|
|
"default":default_traj_txt
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VIDEO","VIDEO",)
|
|
RETURN_NAMES = ("result_video","traj_video",)
|
|
|
|
FUNCTION = "gen_video"
|
|
|
|
#OUTPUT_NODE = False
|
|
|
|
CATEGORY = "AIFSH_ViewCrafter"
|
|
|
|
def gen_video(self,img,traj_txt,video_length,ddim_steps,center_scale,elevation,
|
|
d_theta,d_phi,d_r,seed,custom_traj_txt=default_traj_txt):
|
|
parser = get_parser()
|
|
opts = parser.parse_args()
|
|
opts.exp_name = "AIFSH"
|
|
opts.save_dir = os.path.join(output_dir,"viewcrafter")
|
|
os.makedirs(opts.save_dir,exist_ok=True)
|
|
|
|
img = img.numpy()[0] * 255
|
|
img_np = img.astype(np.uint8)
|
|
img_pil = Image.fromarray(img_np)
|
|
|
|
org_h, org_w = img_pil.size
|
|
height,width = (1024,math.ceil(1024 * org_w/org_h/64)*64) if org_h > org_w else (math.ceil(1024 * org_h/org_w/64)*64,1024)
|
|
img_pil = img_pil.resize((height,width))
|
|
print(f"from {(org_h,org_w)} to {(height, width)}")
|
|
|
|
opts.height = 576
|
|
opts.width = 1024
|
|
tmp_img_path = os.path.join(opts.save_dir,"tmp.png")
|
|
img_pil.save(tmp_img_path)
|
|
|
|
opts.image_dir = tmp_img_path
|
|
if custom_traj_txt != default_traj_txt:
|
|
traj_path = os.path.join(opts.save_dir,"tmp.txt")
|
|
with open(traj_path,'w',encoding="utf-8") as w:
|
|
w.write(custom_traj_txt)
|
|
else:
|
|
traj_path = os.path.join(now_dir,"viewcrafter","trajs",f"{traj_txt}.txt")
|
|
opts.traj_txt = traj_path
|
|
opts.mode = "single_view_txt"
|
|
opts.center_scale = center_scale
|
|
opts.elevation = elevation
|
|
opts.seed = seed
|
|
opts.d_theta = d_theta
|
|
opts.d_phi = d_phi
|
|
opts.d_r = d_r
|
|
ckpt_name = f"ViewCrafter_{video_length}"
|
|
ckpt_path = os.path.join(ckpt_dir,ckpt_name,"model.ckpt")
|
|
if not os.path.exists(ckpt_path):
|
|
snapshot_download(repo_id=f"Drexubery/{ckpt_name}",local_dir=os.path.join(ckpt_dir,ckpt_name))
|
|
opts.ckpt_path = ckpt_path
|
|
opts.config = os.path.join(now_dir,"viewcrafter","configs","inference_pvd_1024.yaml")
|
|
opts.ddim_steps = ddim_steps
|
|
opts.video_length = video_length
|
|
opts.device = "cuda"
|
|
opts.model_path = os.path.join(ckpt_dir,"DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth")
|
|
if not os.path.exists(opts.model_path):
|
|
os.system(f"wget https://download.europe.naverlabs.com/ComputerVision/DUSt3R/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth -P {ckpt_dir}")
|
|
|
|
print(opts)
|
|
if self.pvd is None:
|
|
self.pvd = ViewCrafter(opts)
|
|
self.pvd.nvs_single_view()
|
|
|
|
res_video = os.path.join(output_dir, f'{traj_txt}_diffusion0.mp4')
|
|
shutil.copy(os.path.join(opts.save_dir, 'diffusion0.mp4'),res_video)
|
|
traj_video = os.path.join(output_dir,f'{traj_txt}_viz_traj.mp4')
|
|
shutil.copy(os.path.join(opts.save_dir,'viz_traj.mp4'),traj_video)
|
|
return (res_video, traj_video,)
|
|
|
|
|
|
WEB_DIRECTORY = "./js"
|
|
from .util_nodes import PreViewVideo,LoadVideo
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LoadVideo":LoadVideo,
|
|
"PreViewVideo":PreViewVideo,
|
|
"ViewCrafterTxTNode": ViewCrafterTxTNode
|
|
}
|