From b0c856424c4b9fce218947cbf21358b863ee053c Mon Sep 17 00:00:00 2001 From: LucipherDev Date: Fri, 10 Jan 2025 22:43:36 +0530 Subject: [PATCH] Bug fixes --- install.py | 45 ++++++++++++++++++++------------------- nodes.py | 58 +++++++++++++++++++++++++++++++++----------------- pyproject.toml | 2 +- 3 files changed, 62 insertions(+), 43 deletions(-) diff --git a/install.py b/install.py index 083f52d..213710f 100644 --- a/install.py +++ b/install.py @@ -31,47 +31,48 @@ except: COTRACKER = os.path.join(EXT_PATH, "cotracker") try: - log.info(f"Installing requirements") + log.info("Installing requirements") subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", f"{EXT_PATH}/requirements.txt", "--no-warn-script-location"]) if download_models: from huggingface_hub import snapshot_download + + log.info("Downloading Necessary models") try: - if not os.path.exists(ANIDOC_DIR): - log.info(f"Downloading AniDoc model to: {ANIDOC_DIR}") - snapshot_download( - repo_id="Yhmeng1106/anidoc", - local_dir=DIFFUSERS_DIR, - local_dir_use_symlinks=False, - ) + log.info(f"Downloading AniDoc model to: {ANIDOC_DIR}") + snapshot_download( + repo_id="Yhmeng1106/anidoc", + ignore_patterns=["*.md"], + local_dir=DIFFUSERS_DIR, + local_dir_use_symlinks=False, + ) except Exception: traceback.print_exc() - log.error(f"Failed to download AniDoc model") + log.error("Failed to download AniDoc model") try: - if not os.path.exists(SVD_I2V_DIR): - log.info(f"Downloading stable diffusion video img2vid to: {SVD_I2V_DIR}") - snapshot_download( - repo_id="vdo/stable-video-diffusion-img2vid-xt-1-1", - allow_patterns=[f"*.json", "*fp16*"], - ignore_patterns=["*unet*"], - local_dir=SVD_I2V_DIR, - local_dir_use_symlinks=False, - ) + log.info(f"Downloading stable diffusion video img2vid to: {SVD_I2V_DIR}") + snapshot_download( + repo_id="vdo/stable-video-diffusion-img2vid-xt-1-1", + allow_patterns=["*.json", "*fp16*"], + ignore_patterns=["*unet*"], + local_dir=SVD_I2V_DIR, + local_dir_use_symlinks=False, + ) except Exception: traceback.print_exc() - log.error(f"Failed to download stable diffusion video img2vid") + log.error("Failed to download stable diffusion video img2vid") try: log.info("Installing CoTracker") subprocess.check_call([sys.executable, "-m", "pip", "install", COTRACKER]) except Exception: traceback.print_exc() - log.error(f"Failed to install CoTracker") + log.error("Failed to install CoTracker") - log.info(f"AniDoc Installation completed") + log.info("AniDoc Installation completed") except Exception: traceback.print_exc() - log.error(f"AniDoc Installation failed") \ No newline at end of file + log.error("AniDoc Installation failed") \ No newline at end of file diff --git a/nodes.py b/nodes.py index 3a99b7c..e0c5fb3 100644 --- a/nodes.py +++ b/nodes.py @@ -14,10 +14,15 @@ import torchvision.transforms as T from LightGlue.lightglue import LightGlue, SuperPoint from LightGlue.lightglue.utils import rbd -from models_diffusers.unet_spatio_temporal_condition import UNetSpatioTemporalConditionModel +from models_diffusers.unet_spatio_temporal_condition import ( + UNetSpatioTemporalConditionModel, +) from models_diffusers.controlnet_svd import ControlNetSVDModel -from cotracker.predictor import CoTrackerPredictor, sample_trajectories_with_ref +try: + from cotracker.predictor import CoTrackerPredictor, sample_trajectories_with_ref +except: + pass from lineart_extractor.canny import CannyDetector from lineart_extractor.hed import HEDdetector @@ -50,6 +55,7 @@ folder_paths.folder_names_and_paths["cotracker"] = ( folder_paths.supported_pt_extensions, ) + class AniDocLoader: @classmethod def INPUT_TYPES(cls): @@ -80,17 +86,20 @@ class AniDocLoader: if self.pipeline is None: pbar = ProgressBar(5) - if not os.path.exists(anidoc_path): - log.info(f"Downloading AniDoc model to: {anidoc_path}") - snapshot_download( - repo_id="Yhmeng1106/anidoc", - local_dir=DIFFUSERS_DIR, - local_dir_use_symlinks=False, - ) - pbar.update(1) log.info(f"Loading model from: {anidoc_path}") + log.info("Missing models will be downloaded") + + try: + snapshot_download( + repo_id="Yhmeng1106/anidoc", + ignore_patterns=["*.md"], + local_dir=DIFFUSERS_DIR, + local_dir_use_symlinks=False, + ) + except: + log.info("Couldn't download models") unet = UNetSpatioTemporalConditionModel.from_pretrained( anidoc_path, @@ -111,23 +120,25 @@ class AniDocLoader: pbar.update(1) if XFORMERS_IS_AVAILABLE: - log.info(f"Enabling XFormers") + log.info("Enabling XFormers") unet.enable_xformers_memory_efficient_attention() - if not os.path.exists(svd_img2vid_path): - log.info(f"Downloading stable diffusion video img2vid to: {svd_img2vid_path}") + log.info(f"Loading model from: {svd_img2vid_path}") + log.info("Missing models will be downloaded") + + try: snapshot_download( repo_id="vdo/stable-video-diffusion-img2vid-xt-1-1", - allow_patterns=[f"*.json", "*fp16*"], + allow_patterns=["*.json", "*fp16*"], ignore_patterns=["*unet*"], local_dir=svd_img2vid_path, local_dir_use_symlinks=False, ) + except: + log.info("Couldn't download models") pbar.update(1) - log.info(f"Loading model from: {svd_img2vid_path}") - pipeline = AniDocPipeline.from_pretrained( svd_img2vid_path, unet=unet, @@ -190,6 +201,11 @@ class LoadAniDocCoTracker: device="cuda", dtype=torch.float32, ): + try: + import cotracker + except: + raise ImportError("Couldn't import cotracker module. Please install it to use this node") + if tracking: if self.tracker is None or self.tracker_shift_grid != tracker_shift_grid: cotracker_model_path = folder_paths.get_full_path( @@ -278,7 +294,7 @@ class GetAniDocControlnetImages: return cv2.bitwise_not(image) else: return 255 - image - + def get_controlnet_images( self, input_images, @@ -295,7 +311,7 @@ class GetAniDocControlnetImages: pbar.update(1) - log.info(f"Processing images with lineart detector") + log.info("Processing images with lineart detector") controlnet_images = [] @@ -307,7 +323,7 @@ class GetAniDocControlnetImages: sketch = detector(sketch, 100, 200) else: sketch = detector(sketch) - + if lineart_detector in ["canny", "hed"]: sketch = self.invert_images(sketch) @@ -689,7 +705,9 @@ class AniDocSampler: noise_aug_strength=noise_aug, decode_chunk_size=decode_chunk_size, generator=generator, - callback_on_step_end=lambda *args, **kwargs: (pbar.update(1), kwargs)[-1], + callback_on_step_end=lambda *args, **kwargs: (pbar.update(1), kwargs)[ + -1 + ], ).frames[0] tensor_frames = [T.ToTensor()(img) for img in video_frames] diff --git a/pyproject.toml b/pyproject.toml index 1be5208..f4a1e32 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-anidoc" 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." -version = "1.0.4" +version = "1.0.5" license = {file = "LICENSE"} dependencies = ["diffusers", "huggingface_hub", "Pillow", "accelerate", "omegaconf", "opencv-python", "einops", "kornia", "git+https://github.com/XPixelGroup/BasicSR.git"]