modify the code

siliconflow · Dec 11, 2024 · 0540ef9 · 0540ef9
1 parent 9aab405
commit 0540ef9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/bizyair_extras/__init__.py b/bizyair_extras/__init__.py
@@ -2,6 +2,7 @@
 from .nodes_comfyui_detail_daemon import *
 from .nodes_comfyui_instantid import *
 from .nodes_comfyui_pulid_flux import *
+from .nodes_comfyui_rf_inversion import *
 from .nodes_controlnet import *
 from .nodes_custom_sampler import *
 from .nodes_differential_diffusion import *

diff --git a/bizyair_extras/nodes_comfyui_rf_inversion.py b/bizyair_extras/nodes_comfyui_rf_inversion.py
@@ -16,7 +16,7 @@ class OutFluxModelSamplingPredNode(BizyAirBaseNode):
     def INPUT_TYPES(s):
         return {
             "required": {
-                "model": ("MODEL",),
+                "model": (data_types.MODEL,),
                 "max_shift": (
                     "FLOAT",
                     {"default": 1.15, "min": 0.0, "max": 100.0, "step": 0.01},
@@ -57,7 +57,7 @@ class InFluxModelSamplingPredNode(BizyAirBaseNode):
     def INPUT_TYPES(s):
         return {
             "required": {
-                "model": ("MODEL",),
+                "model": (data_types.MODEL,),
                 "max_shift": (
                     "FLOAT",
                     {"default": 1.15, "min": 0.0, "max": 100.0, "step": 0.01},