diff --git a/.gitignore b/.gitignore
index ed8ebf5..1f17019 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,3 @@
-__pycache__
\ No newline at end of file
+__pycache__
+checkpoints
+!checkpoints/.gitkeep
\ No newline at end of file
diff --git a/README.md b/README.md
index 72ffa53..5395780 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,13 @@
# ViewCrafter-ComfyUI
a custom node for [ViewCrafter](https://github.com/Drexubery/ViewCrafter)
+
+## Example
+test on 4090,py310,torch==2.3.1
+|img|traj_text|output_traj_video|output_render_video|
+|--|--|--|--|
+||`left`| | |
+||`loop1`| | |
+
+
+## traj point
+you can refer [here](https://github.com/Drexubery/ViewCrafter/blob/main/docs/render_help.md)
\ No newline at end of file
diff --git a/__init__.py b/__init__.py
index 0fbd67a..f360da6 100644
--- a/__init__.py
+++ b/__init__.py
@@ -2,12 +2,13 @@ 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 viewcrafter.viewcrafter import ViewCrafter
+from .infer import ViewCrafter
from viewcrafter.configs.infer_config import get_parser
output_dir = folder_paths.get_output_directory()
@@ -88,14 +89,15 @@ class ViewCrafterTxTNode:
img = img.numpy()[0] * 255
img_np = img.astype(np.uint8)
- print(img_np.shape)
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))
+ #img_pil = img_pil.resize((height,width))
print(f"from {(org_h,org_w)} to {(height, width)}")
- opts.height = height
- opts.width = width
+ '''
+ opts.height = 576
+ opts.width = 1024
tmp_img_path = os.path.join(opts.save_dir,"tmp.png")
img_pil.save(tmp_img_path)
@@ -132,8 +134,10 @@ class ViewCrafterTxTNode:
self.pvd = ViewCrafter(opts)
self.pvd.nvs_single_view()
- res_video = os.path.join(opts.save_dir, 'diffusion0.mp4')
- traj_video = os.path.join(opts.save_dir,'viz_traj.mp4')
+ 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,)
diff --git a/doc/train.png b/doc/train.png
new file mode 100644
index 0000000..20827e9
Binary files /dev/null and b/doc/train.png differ
diff --git a/doc/web.png b/doc/web.png
new file mode 100644
index 0000000..e1d9508
Binary files /dev/null and b/doc/web.png differ
diff --git a/viewcrafter/viewcrafter.py b/infer.py
similarity index 97%
rename from viewcrafter/viewcrafter.py
rename to infer.py
index 84f42b4..b097da1 100644
--- a/viewcrafter/viewcrafter.py
+++ b/infer.py
@@ -1,6 +1,11 @@
import os,sys
now_dir = os.path.dirname(os.path.abspath(__file__))
-sys.path.append(os.path.join(now_dir,"extern"))
+sys.path.append(now_dir)
+sys.path.append(os.path.join(now_dir,"viewcrafter"))
+sys.path.append(os.path.join(now_dir,"viewcrafter/extern/dust3r"))
+import viewcrafter.utils as utils
+comfyui_utils = sys.modules['utils']
+sys.modules['utils'] = utils
from dust3r.inference import inference, load_model
from dust3r.utils.image import load_images
from dust3r.image_pairs import make_pairs
@@ -20,13 +25,13 @@ from torchvision.utils import save_image
import torch.nn.functional as F
import torchvision.transforms as transforms
from PIL import Image
-from .utils.pvd_utils import *
+from utils.pvd_utils import *
from omegaconf import OmegaConf
from pytorch_lightning import seed_everything
-from .utils.diffusion_utils import instantiate_from_config,load_model_checkpoint,image_guided_synthesis
+from utils.diffusion_utils import instantiate_from_config,load_model_checkpoint,image_guided_synthesis
from pathlib import Path
from torchvision.utils import save_image
-
+sys.modules['utils'] = comfyui_utils
class ViewCrafter:
def __init__(self, opts, gradio = False):
self.opts = opts
@@ -42,11 +47,11 @@ class ViewCrafter:
pairs = make_pairs(input_images, scene_graph='complete', prefilter=None, symmetrize=True)
output = inference(pairs, self.dust3r, self.device, batch_size=self.opts.batch_size)
- mode = GlobalAlignerMode.PointCloudOptimizer #if len(self.images) > 2 else GlobalAlignerMode.PairViewer
+ mode = GlobalAlignerMode.PointCloudOptimizer # if len(self.images) > 2 else GlobalAlignerMode.PairViewer
scene = global_aligner(output, device=self.device, mode=mode)
if mode == GlobalAlignerMode.PointCloudOptimizer:
loss = scene.compute_global_alignment(init='mst', niter=self.opts.niter, schedule=self.opts.schedule, lr=self.opts.lr)
-
+ # loss.requires_grad_(True)
if clean_pc:
self.scene = scene.clean_pointcloud()
else:
diff --git a/requirements.txt b/requirements.txt
index aedddd8..5f91e6f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,10 +3,8 @@ decord==0.6.0
einops==0.6.1
imageio==2.27.0
imageio-ffmpeg==0.4.8
-torch==1.13.1
-torchvision
kornia
-matplotlib==3.9.2
+matplotlib
moviepy==1.0.3
numpy==1.23.5
open-clip-torch==2.17.1
@@ -25,16 +23,14 @@ timm==0.6.13
tqdm==4.65.0
transformers==4.28.1
trimesh==4.4.3
-xformers
gradio==3.37.0
gradio_client==0.7.1
omegaconf==2.3.0
-triton
altair==5.4.0
certifi==2024.7.4
grpcio==1.66.0
httpx==0.27.0
Pygments==2.18.0
starlette==0.38.2
-tifffile==2024.8.24
+tifffile
yarl==1.9.4
\ No newline at end of file
diff --git a/viewcrafter/__init__.py b/viewcrafter/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/viewcrafter/extern/dust3r/dust3r/cloud_opt/base_opt.py b/viewcrafter/extern/dust3r/dust3r/cloud_opt/base_opt.py
index 18256de..72f7d78 100644
--- a/viewcrafter/extern/dust3r/dust3r/cloud_opt/base_opt.py
+++ b/viewcrafter/extern/dust3r/dust3r/cloud_opt/base_opt.py
@@ -368,6 +368,7 @@ def global_alignment_loop(net, lr=0.01, niter=300, schedule='cosine', lr_min=1e-
optimizer.zero_grad()
loss = net()
+ loss.requires_grad_(True)
loss.backward()
optimizer.step()
loss = float(loss)