Add model auto-update feature

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
 

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
 

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()

comfyui/workflows/workflow-ruyi-i2v-start-end-frames-80g.json

@@ -32,6 +32,7 @@
       },
       "widgets_values": [
         "Ruyi-Mini-7B",
+        "yes",
         "yes"
       ]
     },

comfyui/workflows/workflow-ruyi-i2v-start-end-frames.json

@@ -32,6 +32,7 @@
       },
       "widgets_values": [
         "Ruyi-Mini-7B",
+        "yes",
         "yes"
       ]
     },

comfyui/workflows/workflow-ruyi-i2v-start-frame-80g.json

@@ -91,6 +91,7 @@
       },
       "widgets_values": [
         "Ruyi-Mini-7B",
+        "yes",
         "yes"
       ]
     },

comfyui/workflows/workflow-ruyi-i2v-start-frame.json

@@ -32,6 +32,7 @@
       },
       "widgets_values": [
         "Ruyi-Mini-7B",
+        "yes",
         "yes"
       ]
     },

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} ...")

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} ...")

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"]