Skip to content

Commit b0c8564

Browse files
LucipherDevLucipherDev
committed
Bug fixes
1 parent 3185eb2 commit b0c8564

File tree

3 files changed

+60
-41
lines changed

3 files changed

+60
-41
lines changed

‎install.py

Lines changed: 23 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -31,47 +31,48 @@
3131
COTRACKER = os.path.join(EXT_PATH, "cotracker")
3232

3333
try:
34-
log.info(f"Installing requirements")
34+
log.info("Installing requirements")
3535
subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", f"{EXT_PATH}/requirements.txt", "--no-warn-script-location"])
3636

3737
if download_models:
3838
from huggingface_hub import snapshot_download
39+
40+
log.info("Downloading Necessary models")
3941

4042
try:
41-
if not os.path.exists(ANIDOC_DIR):
42-
log.info(f"Downloading AniDoc model to: {ANIDOC_DIR}")
43-
snapshot_download(
44-
repo_id="Yhmeng1106/anidoc",
45-
local_dir=DIFFUSERS_DIR,
46-
local_dir_use_symlinks=False,
47-
)
43+
log.info(f"Downloading AniDoc model to: {ANIDOC_DIR}")
44+
snapshot_download(
45+
repo_id="Yhmeng1106/anidoc",
46+
ignore_patterns=["*.md"],
47+
local_dir=DIFFUSERS_DIR,
48+
local_dir_use_symlinks=False,
49+
)
4850
except Exception:
4951
traceback.print_exc()
50-
log.error(f"Failed to download AniDoc model")
52+
log.error("Failed to download AniDoc model")
5153

5254
try:
53-
if not os.path.exists(SVD_I2V_DIR):
54-
log.info(f"Downloading stable diffusion video img2vid to: {SVD_I2V_DIR}")
55-
snapshot_download(
56-
repo_id="vdo/stable-video-diffusion-img2vid-xt-1-1",
57-
allow_patterns=[f"*.json", "*fp16*"],
58-
ignore_patterns=["*unet*"],
59-
local_dir=SVD_I2V_DIR,
60-
local_dir_use_symlinks=False,
61-
)
55+
log.info(f"Downloading stable diffusion video img2vid to: {SVD_I2V_DIR}")
56+
snapshot_download(
57+
repo_id="vdo/stable-video-diffusion-img2vid-xt-1-1",
58+
allow_patterns=["*.json", "*fp16*"],
59+
ignore_patterns=["*unet*"],
60+
local_dir=SVD_I2V_DIR,
61+
local_dir_use_symlinks=False,
62+
)
6263
except Exception:
6364
traceback.print_exc()
64-
log.error(f"Failed to download stable diffusion video img2vid")
65+
log.error("Failed to download stable diffusion video img2vid")
6566

6667
try:
6768
log.info("Installing CoTracker")
6869
subprocess.check_call([sys.executable, "-m", "pip", "install", COTRACKER])
6970
except Exception:
7071
traceback.print_exc()
71-
log.error(f"Failed to install CoTracker")
72+
log.error("Failed to install CoTracker")
7273

73-
log.info(f"AniDoc Installation completed")
74+
log.info("AniDoc Installation completed")
7475

7576
except Exception:
7677
traceback.print_exc()
77-
log.error(f"AniDoc Installation failed")
78+
log.error("AniDoc Installation failed")

‎nodes.py

Lines changed: 36 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,15 @@
1414
from LightGlue.lightglue import LightGlue, SuperPoint
1515
from LightGlue.lightglue.utils import rbd
1616

17-
from models_diffusers.unet_spatio_temporal_condition import UNetSpatioTemporalConditionModel
17+
from models_diffusers.unet_spatio_temporal_condition import (
18+
UNetSpatioTemporalConditionModel,
19+
)
1820
from models_diffusers.controlnet_svd import ControlNetSVDModel
1921

20-
from cotracker.predictor import CoTrackerPredictor, sample_trajectories_with_ref
22+
try:
23+
from cotracker.predictor import CoTrackerPredictor, sample_trajectories_with_ref
24+
except:
25+
pass
2126

2227
from lineart_extractor.canny import CannyDetector
2328
from lineart_extractor.hed import HEDdetector
@@ -50,6 +55,7 @@
5055
folder_paths.supported_pt_extensions,
5156
)
5257

58+
5359
class AniDocLoader:
5460
@classmethod
5561
def INPUT_TYPES(cls):
@@ -80,17 +86,20 @@ def load_anidoc(
8086
if self.pipeline is None:
8187
pbar = ProgressBar(5)
8288

83-
if not os.path.exists(anidoc_path):
84-
log.info(f"Downloading AniDoc model to: {anidoc_path}")
89+
pbar.update(1)
90+
91+
log.info(f"Loading model from: {anidoc_path}")
92+
log.info("Missing models will be downloaded")
93+
94+
try:
8595
snapshot_download(
8696
repo_id="Yhmeng1106/anidoc",
97+
ignore_patterns=["*.md"],
8798
local_dir=DIFFUSERS_DIR,
8899
local_dir_use_symlinks=False,
89100
)
90-
91-
pbar.update(1)
92-
93-
log.info(f"Loading model from: {anidoc_path}")
101+
except:
102+
log.info("Couldn't download models")
94103

95104
unet = UNetSpatioTemporalConditionModel.from_pretrained(
96105
anidoc_path,
@@ -111,23 +120,25 @@ def load_anidoc(
111120
pbar.update(1)
112121

113122
if XFORMERS_IS_AVAILABLE:
114-
log.info(f"Enabling XFormers")
123+
log.info("Enabling XFormers")
115124
unet.enable_xformers_memory_efficient_attention()
116125

117-
if not os.path.exists(svd_img2vid_path):
118-
log.info(f"Downloading stable diffusion video img2vid to: {svd_img2vid_path}")
126+
log.info(f"Loading model from: {svd_img2vid_path}")
127+
log.info("Missing models will be downloaded")
128+
129+
try:
119130
snapshot_download(
120131
repo_id="vdo/stable-video-diffusion-img2vid-xt-1-1",
121-
allow_patterns=[f"*.json", "*fp16*"],
132+
allow_patterns=["*.json", "*fp16*"],
122133
ignore_patterns=["*unet*"],
123134
local_dir=svd_img2vid_path,
124135
local_dir_use_symlinks=False,
125136
)
137+
except:
138+
log.info("Couldn't download models")
126139

127140
pbar.update(1)
128141

129-
log.info(f"Loading model from: {svd_img2vid_path}")
130-
131142
pipeline = AniDocPipeline.from_pretrained(
132143
svd_img2vid_path,
133144
unet=unet,
@@ -190,6 +201,11 @@ def load_tracker(
190201
device="cuda",
191202
dtype=torch.float32,
192203
):
204+
try:
205+
import cotracker
206+
except:
207+
raise ImportError("Couldn't import cotracker module. Please install it to use this node")
208+
193209
if tracking:
194210
if self.tracker is None or self.tracker_shift_grid != tracker_shift_grid:
195211
cotracker_model_path = folder_paths.get_full_path(
@@ -278,7 +294,7 @@ def invert_images(self, image):
278294
return cv2.bitwise_not(image)
279295
else:
280296
return 255 - image
281-
297+
282298
def get_controlnet_images(
283299
self,
284300
input_images,
@@ -295,7 +311,7 @@ def get_controlnet_images(
295311

296312
pbar.update(1)
297313

298-
log.info(f"Processing images with lineart detector")
314+
log.info("Processing images with lineart detector")
299315

300316
controlnet_images = []
301317

@@ -307,7 +323,7 @@ def get_controlnet_images(
307323
sketch = detector(sketch, 100, 200)
308324
else:
309325
sketch = detector(sketch)
310-
326+
311327
if lineart_detector in ["canny", "hed"]:
312328
sketch = self.invert_images(sketch)
313329

@@ -689,7 +705,9 @@ def sample(
689705
noise_aug_strength=noise_aug,
690706
decode_chunk_size=decode_chunk_size,
691707
generator=generator,
692-
callback_on_step_end=lambda *args, **kwargs: (pbar.update(1), kwargs)[-1],
708+
callback_on_step_end=lambda *args, **kwargs: (pbar.update(1), kwargs)[
709+
-1
710+
],
693711
).frames[0]
694712

695713
tensor_frames = [T.ToTensor()(img) for img in video_frames]

‎pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
[project]
22
name = "comfyui-anidoc"
33
description = "ComfyUI Custom Nodes for 'AniDoc: Animation Creation Made Easier'. This approach automates line art video colorization using a novel model that aligns color information from references, ensures temporal consistency, and reduces manual effort in animation production."
4-
version = "1.0.4"
4+
version = "1.0.5"
55
license = {file = "LICENSE"}
66
dependencies = ["diffusers", "huggingface_hub", "Pillow", "accelerate", "omegaconf", "opencv-python", "einops", "kornia", "git+https://github.com/XPixelGroup/BasicSR.git"]
77

0 commit comments

Comments
 (0)