From 5b0d3a114515b3752f4a7a42239ca872ae88a541 Mon Sep 17 00:00:00 2001 From: cellzero Date: Tue, 24 Dec 2024 14:22:03 +0800 Subject: [PATCH] Add model auto-update feature --- comfyui/README.md | 1 + comfyui/README_CN.md | 1 + comfyui/comfyui_nodes.py | 17 +++++++++++++++-- .../workflow-ruyi-i2v-start-end-frames-80g.json | 1 + .../workflow-ruyi-i2v-start-end-frames.json | 1 + .../workflow-ruyi-i2v-start-frame-80g.json | 1 + .../workflow-ruyi-i2v-start-frame.json | 1 + predict_i2v.py | 15 ++++++++++----- predict_i2v_80g.py | 15 ++++++++++----- pyproject.toml | 2 +- 10 files changed, 42 insertions(+), 13 deletions(-) diff --git a/comfyui/README.md b/comfyui/README.md index bf999e7..d68ecf6 100644 --- a/comfyui/README.md +++ b/comfyui/README.md @@ -20,6 +20,7 @@ The Load Model node is used to load the model from disk. It also provides the fu - **model**: Select which model to use. Currently, Ruyi-Mini-7B is the only option. - **auto_download**: Whether to automatically download. Defaults to yes. If the model is detected as missing or incomplete, it will automatically download the model to the *ComfyUI/models/Ruyi* path. +- **auto_update**: Whether to automatically check for and update the current model. Defaults to yes. When auto_download is enabled, the system will automatically check for updates to the model and download any updates to the *ComfyUI/models/Ruyi* directory. Please note that this feature relies on the caching mechanism of huggingface_hub, so do not delete the *.cache* folder in the model directory to ensure a smooth update process. ### Load LoRA diff --git a/comfyui/README_CN.md b/comfyui/README_CN.md index c645bf4..6df016c 100644 --- a/comfyui/README_CN.md +++ b/comfyui/README_CN.md @@ -20,6 +20,7 @@ - **model**: 选择使用哪个模型，目前只有 Ruyi-Mini-7B 一个选项。 - **auto_download**: 是否自动下载，默认为 yes，检测到模型不存在（或不完整）时，将自动下载模型到 *ComfyUI/models/Ruyi* 路径。 +- **auto_update**: 是否自动检查并更新当前模型，默认为 yes。当启用 auto_download 时，系统将自动检查模型是否有更新，并将更新内容下载到 *ComfyUI/models/Ruyi* 路径。请注意，此功能依赖于 huggingface_hub 的缓存机制，因此请勿删除模型路径中的 *.cache* 文件夹，以确保更新过程顺利进行。 ### Load LoRA diff --git a/comfyui/comfyui_nodes.py b/comfyui/comfyui_nodes.py index b6923c5..b2ac549 100644 --- a/comfyui/comfyui_nodes.py +++ b/comfyui/comfyui_nodes.py @@ -46,6 +46,12 @@ class Ruyi_LoadModel: "default": "yes" } ), + "auto_update": ( + ["yes", "no"], + { + "default": "yes" + } + ) }, } @@ -121,7 +127,7 @@ class Ruyi_LoadModel: print("[Ruyi] Setup pipeline failed:", e) return None - def load_model(self, model, auto_download): + def load_model(self, model, auto_download, auto_update): # Init weight_dtype and device device = mm.get_torch_device() offload_device = mm.unet_offload_device() @@ -135,9 +141,16 @@ class Ruyi_LoadModel: config_path = os.path.join(script_directory, "config", "default.yaml") config = OmegaConf.load(config_path) - # Init model + # Check for update repo_id = f"IamCreateAI/{model_name}" model_path = os.path.join(folder_paths.models_dir, "Ruyi", model_name) + if auto_download == "yes" and auto_update == "yes": + print(f"Checking for {model_name} updates ...") + + # Download the model + snapshot_download(repo_id=repo_id, local_dir=model_path) + + # Init model pipeline = self.try_setup_pipeline(model_path, weight_dtype, config) if pipeline is None and auto_download == "yes": mm.soft_empty_cache() diff --git a/comfyui/workflows/workflow-ruyi-i2v-start-end-frames-80g.json b/comfyui/workflows/workflow-ruyi-i2v-start-end-frames-80g.json index 4b16047..2ce29e6 100644 --- a/comfyui/workflows/workflow-ruyi-i2v-start-end-frames-80g.json +++ b/comfyui/workflows/workflow-ruyi-i2v-start-end-frames-80g.json @@ -32,6 +32,7 @@ }, "widgets_values": [ "Ruyi-Mini-7B", + "yes", "yes" ] }, diff --git a/comfyui/workflows/workflow-ruyi-i2v-start-end-frames.json b/comfyui/workflows/workflow-ruyi-i2v-start-end-frames.json index f604b26..540d0ac 100644 --- a/comfyui/workflows/workflow-ruyi-i2v-start-end-frames.json +++ b/comfyui/workflows/workflow-ruyi-i2v-start-end-frames.json @@ -32,6 +32,7 @@ }, "widgets_values": [ "Ruyi-Mini-7B", + "yes", "yes" ] }, diff --git a/comfyui/workflows/workflow-ruyi-i2v-start-frame-80g.json b/comfyui/workflows/workflow-ruyi-i2v-start-frame-80g.json index 9159b92..01d0100 100644 --- a/comfyui/workflows/workflow-ruyi-i2v-start-frame-80g.json +++ b/comfyui/workflows/workflow-ruyi-i2v-start-frame-80g.json @@ -91,6 +91,7 @@ }, "widgets_values": [ "Ruyi-Mini-7B", + "yes", "yes" ] }, diff --git a/comfyui/workflows/workflow-ruyi-i2v-start-frame.json b/comfyui/workflows/workflow-ruyi-i2v-start-frame.json index 1929a35..fdd4cb7 100644 --- a/comfyui/workflows/workflow-ruyi-i2v-start-frame.json +++ b/comfyui/workflows/workflow-ruyi-i2v-start-frame.json @@ -32,6 +32,7 @@ }, "widgets_values": [ "Ruyi-Mini-7B", + "yes", "yes" ] }, diff --git a/predict_i2v.py b/predict_i2v.py index 7c66838..78eb26b 100644 --- a/predict_i2v.py +++ b/predict_i2v.py @@ -46,7 +46,8 @@ config_path = "config/default.yaml" model_name = "Ruyi-Mini-7B" model_type = "Inpaint" model_path = f"models/{model_name}" # (Down)load mode in this path -auto_download = True +auto_download = True # Automatically download the model if the pipeline creation fails +auto_update = True # If auto_download is enabled, check for updates and update the model if necessary # LoRA settings lora_path = None @@ -145,9 +146,6 @@ def try_setup_pipeline(model_path, weight_dtype, config): return None -# Init weight_dtype and device -weight_dtype = torch.bfloat16 - # Load config config = OmegaConf.load(config_path) @@ -155,8 +153,15 @@ config = OmegaConf.load(config_path) start_img = [Image.open(start_image_path).convert("RGB")] end_img = [Image.open(end_image_path).convert("RGB")] if end_image_path is not None else None -# Init model +# Check for update repo_id = f"IamCreateAI/{model_name}" +if auto_download and auto_update: + print(f"Checking for {model_name} updates ...") + + # Download the model + snapshot_download(repo_id=repo_id, local_dir=model_path) + +# Init model pipeline = try_setup_pipeline(model_path, weight_dtype, config) if pipeline is None and auto_download: print(f"Downloading {model_name} ...") diff --git a/predict_i2v_80g.py b/predict_i2v_80g.py index 67455a9..84d0a33 100644 --- a/predict_i2v_80g.py +++ b/predict_i2v_80g.py @@ -46,7 +46,8 @@ config_path = "config/default.yaml" model_name = "Ruyi-Mini-7B" model_type = "Inpaint" model_path = f"models/{model_name}" # (Down)load mode in this path -auto_download = True +auto_download = True # Automatically download the model if the pipeline creation fails +auto_update = True # If auto_download is enabled, check for updates and update the model if necessary # LoRA settings lora_path = None @@ -145,9 +146,6 @@ def try_setup_pipeline(model_path, weight_dtype, config): return None -# Init weight_dtype and device -weight_dtype = torch.bfloat16 - # Load config config = OmegaConf.load(config_path) @@ -155,8 +153,15 @@ config = OmegaConf.load(config_path) start_img = [Image.open(start_image_path).convert("RGB")] end_img = [Image.open(end_image_path).convert("RGB")] if end_image_path is not None else None -# Init model +# Check for update repo_id = f"IamCreateAI/{model_name}" +if auto_download and auto_update: + print(f"Checking for {model_name} updates ...") + + # Download the model + snapshot_download(repo_id=repo_id, local_dir=model_path) + +# Init model pipeline = try_setup_pipeline(model_path, weight_dtype, config) if pipeline is None and auto_download: print(f"Downloading {model_name} ...") diff --git a/pyproject.toml b/pyproject.toml index df87b12..dea3e29 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "ruyi-models" description = "ComfyUI wrapper nodes for Ruyi, an image-to-video model by CreateAI." -version = "1.0.0" +version = "1.0.1" license = {file = "LICENSE"} dependencies = ["Pillow", "einops", "safetensors", "timm", "tomesd", "torch", "torchdiffeq", "torchsde", "decord", "datasets", "numpy", "scikit-image", "opencv-python", "omegaconf", "SentencePiece", "albumentations", "imageio[ffmpeg]", "imageio[pyav]", "tensorboard", "beautifulsoup4", "ftfy", "func_timeout", "huggingface_hub", "accelerate>=0.26.0", "diffusers>=0.28.2", "transformers>=4.37.2"]