Merge pull request #8 from FuouM/stab

StabStitch

Author: FuouM

README.md

@@ -1,11 +1,20 @@
-from .run import WFEN, RealViFormerSR, ProPIH_Harmonizer, CoLIE_LowLight_Enhance
-
+from .run import (
+    WFEN,
+    CoLIE_LowLight_Enhance,
+    ProPIH_Harmonizer,
+    RealViFormerSR,
+)
 from .run_convir import (
+    ConvIR_DefocusDeblur,
     ConvIR_DeHaze,
     ConvIR_DeRain,
     ConvIR_DeSnow,
     ConvIR_MotionDeBlur,
-    ConvIR_DefocusDeblur,
+)
+from .run_stabstitch import (
+    StabStitch_Crop_Resize,
+    StabStitch_Stabilize,
+    StabStitch_Stitch,
 )
 
 NODE_CLASS_MAPPINGS = {
@@ -18,6 +27,9 @@
     "ConvIR_DeSnow": ConvIR_DeSnow,
     "ConvIR_MotionDeBlur": ConvIR_MotionDeBlur,
     "ConvIR_DefocusDeblur": ConvIR_DefocusDeblur,
+    "StabStitch": StabStitch_Stitch,
+    "StabStitch_Stabilize": StabStitch_Stabilize,
+    "StabStitch_Crop_Resize": StabStitch_Crop_Resize,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -30,6 +42,9 @@
     "ConvIR_DeSnow": "ConvIR DeSnow",
     "ConvIR_MotionDeBlur": "ConvIR Motion DeBlur",
     "ConvIR_DefocusDeblur": "ConvIR Defocus Deblur",
+    "StabStitch": "StabStitch",
+    "StabStitch_Stabilize": "StabStitch Stabilize",
+    "StabStitch_Crop_Resize": "StabStitch Crop Resize",
 }
 
 try:

__init__.py

@@ -8,3 +8,9 @@
     "VFIMamba_S": "models/vfimamba/VFIMamba_S.pkl",
 }
 VFI_MAMBA_DEFAULT = "VFIMamba"
+
+STAB_TEMPORAL_PATH = "models/stabstitch/temporal_warp.pth"
+STAB_SPATIAL_PATH = "models/stabstitch/spatial_warp.pth"
+STAB_SMOOTH_PATH = "models/stabstitch/smooth_warp.pth"
+STAB_CONSTRUCT_MODES = ["FAST", "NORMAL"]
+STAB_CONSTRUCT_DEFAULT = "FAST"