hipinis commited on
Commit
32d36ef
·
verified ·
1 Parent(s): 96d97a7

Upload 23 files

Browse files
.gitattributes CHANGED
@@ -483,3 +483,4 @@ custom_nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/usage_01.png filter=lfs diff=lfs
483
  custom_nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/usage_02.png filter=lfs diff=lfs merge=lfs -text
484
  custom_nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/example_inputs/Mustache_640x360.mp4 filter=lfs diff=lfs merge=lfs -text
485
  custom_nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/example_inputs/Sadhu_320x478.png filter=lfs diff=lfs merge=lfs -text
 
 
483
  custom_nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/usage_02.png filter=lfs diff=lfs merge=lfs -text
484
  custom_nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/example_inputs/Mustache_640x360.mp4 filter=lfs diff=lfs merge=lfs -text
485
  custom_nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/example_inputs/Sadhu_320x478.png filter=lfs diff=lfs merge=lfs -text
486
+ custom_nodes/ComfyUI-VFI/docs/image.png filter=lfs diff=lfs merge=lfs -text
custom_nodes/ComfyUI-VFI/.gitignore ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ *.pkl
2
+ *.safetensors
3
+ **/*.pkl
4
+ **/*.safetensors
5
+ **/__pycache__
6
+ .dev
7
+ .DS_Store
custom_nodes/ComfyUI-VFI/README.md ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # ComfyUI-VFI
2
+
3
+ Video Frame Interpolation nodes for ComfyUI using RIFE (Real-Time Intermediate Flow Estimation).
4
+
5
+ ![RIFE](./docs/image.png)
6
+
7
+ ## Features
8
+
9
+ - High-quality frame interpolation using RIFE
10
+ - Convert between different frame rates (e.g., 30fps to 60fps)
11
+ - Adjustable processing scale for performance/quality trade-off
12
+ - Model caching for efficient processing
13
+ - Progress tracking in ComfyUI
14
+
15
+ ## Installation
16
+
17
+ - Clone this repository into your ComfyUI custom_nodes directory:
18
+
19
+ ```bash
20
+ cd ComfyUI/custom_nodes
21
+ git clone https://github.com/your-username/ComfyUI-VFI.git
22
+ ```
23
+
24
+ - Install required dependencies:
25
+
26
+ ```bash
27
+ cd ComfyUI-VFI
28
+ pip install -r requirements.txt
29
+ ```
30
+
31
+ - The RIFE model will be automatically downloaded on first use
32
+ - Alternatively, you can manually place `flownet.pkl` in:
33
+ - `ComfyUI-VFI/rife/train_log/`
34
+ - Or `ComfyUI/models/rife/`
35
+
36
+ ## Usage
37
+
38
+ The node will appear in the "image/animation" category as "RIFE Frame Interpolation".
39
+
40
+ ### Inputs
41
+
42
+ - **images**: Image sequence tensor [N, H, W, C]
43
+ - **source_fps**: Original frame rate (default: 30.0)
44
+ - **target_fps**: Desired frame rate (default: 60.0)
45
+ - **scale**: Processing scale factor (default: 1.0)
46
+ - Lower values (0.25-0.5) for faster processing
47
+ - Higher values (1.0-4.0) for better quality
48
+
49
+ ### Output
50
+
51
+ - **images**: Interpolated image sequence tensor
52
+
53
+ ## Example Workflow
54
+
55
+ 1. Load video frames using a video loader node
56
+ 2. Connect to RIFE Frame Interpolation node
57
+ 3. Set source and target FPS
58
+ 4. Connect output to video encoder or preview
59
+
60
+ ## Model Download
61
+
62
+ The RIFE model (`flownet.pkl`) can be downloaded from the official RIFE repository.
custom_nodes/ComfyUI-VFI/__init__.py ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ """ComfyUI-VFI: Video Frame Interpolation nodes for ComfyUI"""
2
+
3
+ from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
4
+
5
+ __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
custom_nodes/ComfyUI-VFI/__pycache__/__init__.cpython-313.pyc ADDED
Binary file (375 Bytes). View file
 
custom_nodes/ComfyUI-VFI/__pycache__/nodes.cpython-313.pyc ADDED
Binary file (10.1 kB). View file
 
custom_nodes/ComfyUI-VFI/docs/image.png ADDED

Git LFS Details

  • SHA256: 78329fb3b6ca9bb0134ef1c5ddabaca2a78266979e6258b508aa838616466506
  • Pointer size: 131 Bytes
  • Size of remote file: 281 kB
custom_nodes/ComfyUI-VFI/examples/interp.json ADDED
@@ -0,0 +1,343 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "id": "68247ae1-7321-451f-9b71-b44476309c39",
3
+ "revision": 0,
4
+ "last_node_id": 19,
5
+ "last_link_id": 19,
6
+ "nodes": [
7
+ {
8
+ "id": 13,
9
+ "type": "VHS_VideoInfoLoaded",
10
+ "pos": [585.9952392578125, 234.33438110351562],
11
+ "size": [247.837890625, 106],
12
+ "flags": {},
13
+ "order": 3,
14
+ "mode": 0,
15
+ "inputs": [
16
+ {
17
+ "name": "video_info",
18
+ "type": "VHS_VIDEOINFO",
19
+ "link": 12
20
+ }
21
+ ],
22
+ "outputs": [
23
+ {
24
+ "name": "fps🟦",
25
+ "type": "FLOAT",
26
+ "links": [13, 15]
27
+ },
28
+ {
29
+ "name": "frame_count🟦",
30
+ "type": "INT",
31
+ "links": null
32
+ },
33
+ {
34
+ "name": "duration🟦",
35
+ "type": "FLOAT",
36
+ "links": null
37
+ },
38
+ {
39
+ "name": "width🟦",
40
+ "type": "INT",
41
+ "links": null
42
+ },
43
+ {
44
+ "name": "height🟦",
45
+ "type": "INT",
46
+ "links": null
47
+ }
48
+ ],
49
+ "properties": {
50
+ "Node name for S&R": "VHS_VideoInfoLoaded"
51
+ },
52
+ "widgets_values": {}
53
+ },
54
+ {
55
+ "id": 15,
56
+ "type": "easy showAnything",
57
+ "pos": [688.3895874023438, 443.7354736328125],
58
+ "size": [210, 88],
59
+ "flags": {},
60
+ "order": 5,
61
+ "mode": 0,
62
+ "inputs": [
63
+ {
64
+ "name": "anything",
65
+ "shape": 7,
66
+ "type": "*",
67
+ "link": 15
68
+ }
69
+ ],
70
+ "outputs": [
71
+ {
72
+ "name": "output",
73
+ "type": "*",
74
+ "links": null
75
+ }
76
+ ],
77
+ "properties": {
78
+ "Node name for S&R": "easy showAnything"
79
+ },
80
+ "widgets_values": ["32.0"]
81
+ },
82
+ {
83
+ "id": 16,
84
+ "type": "VHS_VideoCombine",
85
+ "pos": [1325.4146728515625, 65.38612365722656],
86
+ "size": [220.5830078125, 436.82794189453125],
87
+ "flags": {},
88
+ "order": 6,
89
+ "mode": 0,
90
+ "inputs": [
91
+ {
92
+ "name": "images",
93
+ "type": "IMAGE",
94
+ "link": 16
95
+ },
96
+ {
97
+ "name": "audio",
98
+ "shape": 7,
99
+ "type": "AUDIO",
100
+ "link": null
101
+ },
102
+ {
103
+ "name": "meta_batch",
104
+ "shape": 7,
105
+ "type": "VHS_BatchManager",
106
+ "link": null
107
+ },
108
+ {
109
+ "name": "vae",
110
+ "shape": 7,
111
+ "type": "VAE",
112
+ "link": null
113
+ },
114
+ {
115
+ "name": "frame_rate",
116
+ "type": "FLOAT",
117
+ "widget": {
118
+ "name": "frame_rate"
119
+ },
120
+ "link": 18
121
+ }
122
+ ],
123
+ "outputs": [
124
+ {
125
+ "name": "Filenames",
126
+ "type": "VHS_FILENAMES",
127
+ "links": null
128
+ }
129
+ ],
130
+ "properties": {
131
+ "Node name for S&R": "VHS_VideoCombine"
132
+ },
133
+ "widgets_values": {
134
+ "frame_rate": 8,
135
+ "loop_count": 0,
136
+ "filename_prefix": "AnimateDiff",
137
+ "format": "video/h264-mp4",
138
+ "pix_fmt": "yuv420p",
139
+ "crf": 19,
140
+ "save_metadata": true,
141
+ "pingpong": false,
142
+ "save_output": true,
143
+ "videopreview": {
144
+ "hidden": false,
145
+ "paused": false,
146
+ "params": {
147
+ "filename": "AnimateDiff_00594.mp4",
148
+ "subfolder": "",
149
+ "type": "output",
150
+ "format": "video/h264-mp4",
151
+ "frame_rate": 24
152
+ },
153
+ "muted": false
154
+ }
155
+ }
156
+ },
157
+ {
158
+ "id": 18,
159
+ "type": "PrimitiveFloat",
160
+ "pos": [930.571533203125, 273.9353942871094],
161
+ "size": [270, 58],
162
+ "flags": {},
163
+ "order": 0,
164
+ "mode": 0,
165
+ "inputs": [],
166
+ "outputs": [
167
+ {
168
+ "name": "FLOAT",
169
+ "type": "FLOAT",
170
+ "links": [17, 18, 19]
171
+ }
172
+ ],
173
+ "properties": {
174
+ "Node name for S&R": "PrimitiveFloat"
175
+ },
176
+ "widgets_values": [24]
177
+ },
178
+ {
179
+ "id": 10,
180
+ "type": "VHS_LoadVideo",
181
+ "pos": [278.06610107421875, 158.2479248046875],
182
+ "size": [253.279296875, 262],
183
+ "flags": {},
184
+ "order": 1,
185
+ "mode": 0,
186
+ "inputs": [
187
+ {
188
+ "name": "meta_batch",
189
+ "shape": 7,
190
+ "type": "VHS_BatchManager",
191
+ "link": null
192
+ },
193
+ {
194
+ "name": "vae",
195
+ "shape": 7,
196
+ "type": "VAE",
197
+ "link": null
198
+ }
199
+ ],
200
+ "outputs": [
201
+ {
202
+ "name": "IMAGE",
203
+ "type": "IMAGE",
204
+ "links": [14]
205
+ },
206
+ {
207
+ "name": "frame_count",
208
+ "type": "INT",
209
+ "links": null
210
+ },
211
+ {
212
+ "name": "audio",
213
+ "type": "AUDIO",
214
+ "links": null
215
+ },
216
+ {
217
+ "name": "video_info",
218
+ "type": "VHS_VIDEOINFO",
219
+ "links": [12]
220
+ }
221
+ ],
222
+ "properties": {
223
+ "Node name for S&R": "VHS_LoadVideo"
224
+ },
225
+ "widgets_values": {
226
+ "video": "TWYI-PKSR-A0TV-H5VK-5H7U_2X_32fps.mp4",
227
+ "force_rate": 0,
228
+ "force_size": "Disabled",
229
+ "custom_width": 512,
230
+ "custom_height": 512,
231
+ "frame_load_cap": 0,
232
+ "skip_first_frames": 0,
233
+ "select_every_nth": 1,
234
+ "choose video to upload": "image",
235
+ "videopreview": {
236
+ "hidden": false,
237
+ "paused": false,
238
+ "params": {
239
+ "frame_load_cap": 0,
240
+ "skip_first_frames": 0,
241
+ "force_rate": 0,
242
+ "select_every_nth": 1,
243
+ "filename": "TWYI-PKSR-A0TV-H5VK-5H7U_2X_32fps.mp4",
244
+ "type": "input",
245
+ "format": "video/mp4"
246
+ },
247
+ "muted": false
248
+ }
249
+ }
250
+ },
251
+ {
252
+ "id": 14,
253
+ "type": "RIFEInterpolation",
254
+ "pos": [943.4434204101562, 62.89570617675781],
255
+ "size": [270, 130],
256
+ "flags": {},
257
+ "order": 4,
258
+ "mode": 0,
259
+ "inputs": [
260
+ {
261
+ "name": "images",
262
+ "type": "IMAGE",
263
+ "link": 14
264
+ },
265
+ {
266
+ "name": "source_fps",
267
+ "type": "FLOAT",
268
+ "widget": {
269
+ "name": "source_fps"
270
+ },
271
+ "link": 13
272
+ },
273
+ {
274
+ "name": "target_fps",
275
+ "type": "FLOAT",
276
+ "widget": {
277
+ "name": "target_fps"
278
+ },
279
+ "link": 17
280
+ }
281
+ ],
282
+ "outputs": [
283
+ {
284
+ "name": "images",
285
+ "type": "IMAGE",
286
+ "links": [16]
287
+ }
288
+ ],
289
+ "properties": {
290
+ "Node name for S&R": "RIFEInterpolation"
291
+ },
292
+ "widgets_values": [30, 24, 1, "flownet.pkl"]
293
+ },
294
+ {
295
+ "id": 19,
296
+ "type": "easy showAnything",
297
+ "pos": [1063.455078125, 438.5444641113281],
298
+ "size": [140, 76],
299
+ "flags": {},
300
+ "order": 2,
301
+ "mode": 0,
302
+ "inputs": [
303
+ {
304
+ "name": "anything",
305
+ "shape": 7,
306
+ "type": "*",
307
+ "link": 19
308
+ }
309
+ ],
310
+ "outputs": [
311
+ {
312
+ "name": "output",
313
+ "type": "*",
314
+ "links": null
315
+ }
316
+ ],
317
+ "properties": {
318
+ "Node name for S&R": "easy showAnything"
319
+ },
320
+ "widgets_values": ["24.0"]
321
+ }
322
+ ],
323
+ "links": [
324
+ [12, 10, 3, 13, 0, "VHS_VIDEOINFO"],
325
+ [13, 13, 0, 14, 1, "FLOAT"],
326
+ [14, 10, 0, 14, 0, "IMAGE"],
327
+ [15, 13, 0, 15, 0, "*"],
328
+ [16, 14, 0, 16, 0, "IMAGE"],
329
+ [17, 18, 0, 14, 2, "FLOAT"],
330
+ [18, 18, 0, 16, 4, "FLOAT"],
331
+ [19, 18, 0, 19, 0, "*"]
332
+ ],
333
+ "groups": [],
334
+ "config": {},
335
+ "extra": {
336
+ "ds": {
337
+ "scale": 1.1000000000000005,
338
+ "offset": [-218.49763739436602, 25.55563144593943]
339
+ },
340
+ "frontendVersion": "1.19.9"
341
+ },
342
+ "version": 0.4
343
+ }
custom_nodes/ComfyUI-VFI/nodes.py ADDED
@@ -0,0 +1,297 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """ComfyUI nodes for Video Frame Interpolation using RIFE"""
2
+
3
+ import os
4
+ import subprocess
5
+ import sys
6
+
7
+ import torch
8
+
9
+ from .rife.rife_comfyui_wrapper import RIFEWrapper
10
+
11
+ try:
12
+ import comfy.utils
13
+ import folder_paths
14
+ except ImportError:
15
+ folder_paths = None
16
+ comfy = None
17
+
18
+
19
+ MODEL_CACHE = {}
20
+
21
+
22
+ class RIFEInterpolation:
23
+ """
24
+ ComfyUI node for RIFE (Real-Time Intermediate Flow Estimation) video frame interpolation.
25
+ Takes a sequence of images and interpolates frames to achieve a target frame rate.
26
+ """
27
+
28
+ @classmethod
29
+ def INPUT_TYPES(cls):
30
+ return {
31
+ "required": {
32
+ "images": ("IMAGE",),
33
+ "source_fps": (
34
+ "FLOAT",
35
+ {
36
+ "default": 30.0,
37
+ "min": 1.0,
38
+ "max": 120.0,
39
+ "step": 0.1,
40
+ "display": "number",
41
+ "tooltip": "Source video frame rate",
42
+ },
43
+ ),
44
+ "target_fps": (
45
+ "FLOAT",
46
+ {
47
+ "default": 60.0,
48
+ "min": 1.0,
49
+ "max": 240.0,
50
+ "step": 0.1,
51
+ "display": "number",
52
+ "tooltip": "Target frame rate after interpolation",
53
+ },
54
+ ),
55
+ "scale": (
56
+ "FLOAT",
57
+ {
58
+ "default": 1.0,
59
+ "min": 0.25,
60
+ "max": 4.0,
61
+ "step": 0.25,
62
+ "display": "number",
63
+ "tooltip": "Processing scale factor. Lower values process faster but may reduce quality",
64
+ },
65
+ ),
66
+ },
67
+ "optional": {
68
+ "model_name": (
69
+ ["flownet.pkl"],
70
+ {"default": "flownet.pkl", "tooltip": "RIFE model to use for interpolation"},
71
+ ),
72
+ "batch_size": (
73
+ "INT",
74
+ {
75
+ "default": 8,
76
+ "min": 1,
77
+ "max": 32,
78
+ "step": 1,
79
+ "display": "number",
80
+ "tooltip": "Number of frames to process in parallel. Higher values are faster but use more VRAM",
81
+ },
82
+ ),
83
+ "use_fp16": (
84
+ "BOOLEAN",
85
+ {
86
+ "default": True,
87
+ "tooltip": "Use half precision (FP16) for faster inference and lower VRAM usage. Requires CUDA GPU",
88
+ },
89
+ ),
90
+ },
91
+ }
92
+
93
+ RETURN_TYPES = ("IMAGE",)
94
+ RETURN_NAMES = ("images",)
95
+
96
+ FUNCTION = "interpolate"
97
+
98
+ CATEGORY = "image/animation"
99
+
100
+ DESCRIPTION = "Interpolate video frames using RIFE (Real-Time Intermediate Flow Estimation) to increase frame rate"
101
+
102
+ def interpolate(self, images, source_fps, target_fps, scale, model_name="flownet.pkl", batch_size=8, use_fp16=True):
103
+ # Validate inputs
104
+ if images is None or len(images) == 0:
105
+ raise ValueError("No images provided")
106
+
107
+ if len(images.shape) != 4 or images.shape[-1] != 3:
108
+ raise ValueError(f"Expected image tensor shape [N, H, W, 3], got {images.shape}")
109
+
110
+ if source_fps <= 0 or target_fps <= 0:
111
+ raise ValueError("Frame rates must be positive")
112
+
113
+ if scale <= 0:
114
+ raise ValueError("Scale must be positive")
115
+
116
+ # If source and target fps are the same, return original
117
+ if abs(source_fps - target_fps) < 0.01:
118
+ return (images,)
119
+
120
+ # Get or load model
121
+ model = self._get_or_load_model(model_name, use_fp16=use_fp16)
122
+
123
+ duration = len(images) / source_fps
124
+ total_target_frames = int(duration * target_fps)
125
+
126
+ pbar = None
127
+ if comfy and hasattr(comfy, "utils"):
128
+ pbar = comfy.utils.ProgressBar(total_target_frames)
129
+
130
+ def progress_callback(current, total):
131
+ if pbar:
132
+ pbar.update_absolute(current, total)
133
+
134
+ # Use autocast context for mixed precision
135
+ autocast_enabled = use_fp16 and torch.cuda.is_available()
136
+ autocast_context = torch.amp.autocast("cuda") if autocast_enabled else torch.nullcontext()
137
+
138
+ try:
139
+ with autocast_context:
140
+ interpolated_images = model.interpolate_frames(
141
+ images=images,
142
+ source_fps=source_fps,
143
+ target_fps=target_fps,
144
+ scale=scale,
145
+ progress_callback=progress_callback,
146
+ batch_size=batch_size,
147
+ )
148
+
149
+ return (interpolated_images,)
150
+
151
+ except Exception as e:
152
+ raise RuntimeError(f"Frame interpolation failed: {str(e)}")
153
+
154
+ def _get_or_load_model(self, model_name, use_fp16=False):
155
+ """Load model from cache or disk"""
156
+ global MODEL_CACHE
157
+
158
+ # Create cache key with fp16 flag
159
+ cache_key = f"{model_name}_fp16" if use_fp16 else model_name
160
+
161
+ if cache_key in MODEL_CACHE:
162
+ return MODEL_CACHE[cache_key]
163
+
164
+ # Look for model in multiple locations
165
+ model_paths = [
166
+ os.path.join(os.path.dirname(__file__), "rife", "train_log", model_name),
167
+ os.path.join(os.path.dirname(__file__), "models", model_name),
168
+ ]
169
+
170
+ # Add ComfyUI model directory if available
171
+ if folder_paths and hasattr(folder_paths, "models_dir"):
172
+ model_paths.insert(1, os.path.join(folder_paths.models_dir, "rife", model_name))
173
+
174
+ model_path = None
175
+ for path in model_paths:
176
+ if os.path.exists(path):
177
+ model_path = path
178
+ break
179
+
180
+ if model_path is None:
181
+ # Try to download the model automatically
182
+ print(f"RIFE model '{model_name}' not found. Attempting to download...")
183
+
184
+ # Default download location
185
+ download_target = os.path.join(os.path.dirname(__file__), "rife", "train_log")
186
+
187
+ try:
188
+ # Run the download script
189
+ download_script = os.path.join(os.path.dirname(__file__), "rife", "download_rife.py")
190
+
191
+ if os.path.exists(download_script):
192
+ result = subprocess.run(
193
+ [sys.executable, download_script, download_target], capture_output=True, text=True
194
+ )
195
+
196
+ if result.returncode == 0:
197
+ print("Model downloaded successfully!")
198
+ # Check if model now exists
199
+ model_path = os.path.join(download_target, model_name)
200
+ if not os.path.exists(model_path):
201
+ raise FileNotFoundError(
202
+ f"Model download completed but '{model_name}' not found at expected location."
203
+ )
204
+ else:
205
+ raise RuntimeError(f"Model download failed: {result.stderr}")
206
+ else:
207
+ raise FileNotFoundError(
208
+ f"Download script not found at {download_script}. "
209
+ f"Please manually download the model and place it in one of these locations:\n"
210
+ + "\n".join(f" - {p}" for p in model_paths)
211
+ )
212
+
213
+ except Exception as e:
214
+ raise RuntimeError(
215
+ f"Failed to automatically download RIFE model: {str(e)}\n"
216
+ f"Please manually download the model and place it in one of these locations:\n"
217
+ + "\n".join(f" - {p}" for p in model_paths)
218
+ )
219
+
220
+ # Load model
221
+ print(f"Loading RIFE model from: {model_path}")
222
+ model = RIFEWrapper(model_path, use_fp16=use_fp16)
223
+ MODEL_CACHE[cache_key] = model
224
+
225
+ return model
226
+
227
+ @classmethod
228
+ def IS_CHANGED(cls, **kwargs):
229
+ return float("NaN")
230
+
231
+
232
+ class CalculateLoadedFPS:
233
+ """
234
+ 计算加载后的FPS,根据原始FPS和每n帧选择一帧的参数
235
+ Calculate loaded FPS based on source FPS and select_every_nth parameter
236
+ """
237
+
238
+ @classmethod
239
+ def INPUT_TYPES(cls):
240
+ return {
241
+ "required": {
242
+ "source_fps": (
243
+ "FLOAT",
244
+ {
245
+ "default": 24,
246
+ "min": 0.1,
247
+ "max": 160.0,
248
+ "step": 0.1,
249
+ "display": "number",
250
+ "tooltip": "Source video frame rate",
251
+ },
252
+ ),
253
+ "select_every_nth": (
254
+ "INT",
255
+ {
256
+ "default": 1,
257
+ "min": 1,
258
+ "max": 100,
259
+ "step": 1,
260
+ "display": "number",
261
+ "tooltip": "Select every Nth frame (from VideoHelperSuite)",
262
+ },
263
+ ),
264
+ },
265
+ }
266
+
267
+ RETURN_TYPES = ("FLOAT",)
268
+ RETURN_NAMES = ("loaded_fps",)
269
+
270
+ FUNCTION = "calculate_fps"
271
+
272
+ CATEGORY = "image/animation"
273
+
274
+ DESCRIPTION = "Calculate loaded FPS after frame selection (source_fps / select_every_nth)"
275
+
276
+ def calculate_fps(self, source_fps, select_every_nth):
277
+ # 验证输入
278
+ if source_fps <= 0:
279
+ raise ValueError("source_fps must be positive")
280
+ if select_every_nth <= 0:
281
+ raise ValueError("select_every_nth must be positive")
282
+
283
+ # 计算加载后的FPS
284
+ loaded_fps = source_fps / select_every_nth
285
+ return (loaded_fps,)
286
+
287
+
288
+ # ComfyUI node mappings
289
+ NODE_CLASS_MAPPINGS = {
290
+ "RIFEInterpolation": RIFEInterpolation,
291
+ "CalculateLoadedFPS": CalculateLoadedFPS,
292
+ }
293
+
294
+ NODE_DISPLAY_NAME_MAPPINGS = {
295
+ "RIFEInterpolation": "RIFE Frame Interpolation",
296
+ "CalculateLoadedFPS": "Calculate Loaded FPS",
297
+ }
custom_nodes/ComfyUI-VFI/pyproject.toml ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [project]
2
+ name = "rife_comfyui_wrapper"
3
+ description = "ComfyUI-RIFE is an inference wrapper for RIFE designed for use with ComfyUI."
4
+ version = "1.0.0"
5
+ license = { file = "LICENSE" }
6
+
7
+ [project.urls]
8
+ Repository = "https://github.com/ModelTC/ComfyUI-VFI"
9
+ # Used by Comfy Registry https://comfyregistry.org
10
+
11
+ [tool.comfy]
12
+ PublisherId = "gaclove"
13
+ DisplayName = "ComfyUI-VFI"
14
+ Icon = ""
15
+
16
+ [tool.ruff]
17
+ line-length = 120
18
+
19
+ [tool.ruff.lint]
20
+ extend-select = ["I"]
21
+
22
+ [tool.ruff.lint.per-file-ignores]
23
+
24
+ "rife/train_log/RIFE_HDv3.py" = ["F"]
25
+ "rife/train_log/IFNet_HDv3.py" = ["F"]
custom_nodes/ComfyUI-VFI/requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ torch>=2.0.0
2
+ torchvision>=0.15.0
3
+ numpy>=1.21.0
4
+ requests>=2.25.0
custom_nodes/ComfyUI-VFI/rife/__pycache__/rife_comfyui_wrapper.cpython-313.pyc ADDED
Binary file (8.45 kB). View file
 
custom_nodes/ComfyUI-VFI/rife/download_rife.py ADDED
@@ -0,0 +1,136 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ # coding: utf-8
3
+
4
+ import argparse
5
+ import os
6
+ import shutil
7
+ import sys
8
+ import zipfile
9
+ from pathlib import Path
10
+
11
+ import requests
12
+
13
+
14
+ def is_huggingface_accessible():
15
+ try:
16
+ _response = requests.get("https://huggingface.co", timeout=3)
17
+ return True
18
+ except (requests.ConnectionError, requests.Timeout):
19
+ return False
20
+
21
+
22
+ def get_base_dir():
23
+ """Get project root directory"""
24
+ return Path(__file__).parent.parent
25
+
26
+
27
+ def download_file(url, save_path):
28
+ """Download file"""
29
+ print(f"Starting download: {url}")
30
+ response = requests.get(url, stream=True)
31
+ response.raise_for_status()
32
+
33
+ total_size = int(response.headers.get("content-length", 0))
34
+ downloaded_size = 0
35
+
36
+ with open(save_path, "wb") as f:
37
+ for chunk in response.iter_content(chunk_size=8192):
38
+ if chunk:
39
+ f.write(chunk)
40
+ downloaded_size += len(chunk)
41
+ if total_size > 0:
42
+ progress = (downloaded_size / total_size) * 100
43
+ print(f"\rDownload progress: {progress:.1f}%", end="", flush=True)
44
+
45
+ print(f"\nDownload completed: {save_path}")
46
+
47
+
48
+ def extract_zip(zip_path, extract_to):
49
+ """Extract zip file"""
50
+ print(f"Starting extraction: {zip_path}")
51
+ with zipfile.ZipFile(zip_path, "r") as zip_ref:
52
+ zip_ref.extractall(extract_to)
53
+ print(f"Extraction completed: {extract_to}")
54
+
55
+
56
+ def find_flownet_pkl(extract_dir):
57
+ """Find flownet.pkl file in extracted directory"""
58
+ for root, _dirs, files in os.walk(extract_dir):
59
+ for file in files:
60
+ if file == "flownet.pkl":
61
+ return os.path.join(root, file)
62
+ return None
63
+
64
+
65
+ def main():
66
+ parser = argparse.ArgumentParser(description="Download RIFE model to specified directory")
67
+ parser.add_argument("target_directory", help="Target directory path")
68
+
69
+ args = parser.parse_args()
70
+
71
+ target_dir = Path(args.target_directory)
72
+ if not target_dir.is_absolute():
73
+ target_dir = Path.cwd() / target_dir
74
+
75
+ base_dir = get_base_dir()
76
+ temp_dir = base_dir / "_temp"
77
+
78
+ # Create temporary directory
79
+ temp_dir.mkdir(exist_ok=True)
80
+
81
+ target_dir.mkdir(parents=True, exist_ok=True)
82
+
83
+ if not is_huggingface_accessible():
84
+ print("huggingface.co is not accessible, using hf-mirror.com")
85
+ zip_url = "https://hf-mirror.com/hzwer/RIFE/resolve/main/RIFEv4.26_0921.zip"
86
+ else:
87
+ zip_url = "https://huggingface.co/hzwer/RIFE/resolve/main/RIFEv4.26_0921.zip"
88
+
89
+ zip_path = temp_dir / "RIFEv4.26_0921.zip"
90
+
91
+ try:
92
+ download_file(zip_url, zip_path)
93
+ extract_zip(zip_path, temp_dir)
94
+ flownet_pkl = find_flownet_pkl(temp_dir)
95
+ if flownet_pkl:
96
+ target_file = target_dir / "flownet.pkl"
97
+ shutil.copy2(flownet_pkl, target_file)
98
+ print(f"flownet.pkl copied to: {target_file}")
99
+ else:
100
+ print("Error: flownet.pkl file not found")
101
+ return 1
102
+
103
+ print("RIFE model download and installation completed!")
104
+ return 0
105
+
106
+ except Exception as e:
107
+ print(f"Error: {e}")
108
+ return 1
109
+ finally:
110
+ print("Cleaning up temporary files...")
111
+
112
+ if zip_path.exists():
113
+ try:
114
+ zip_path.unlink()
115
+ print(f"Deleted: {zip_path}")
116
+ except Exception as e:
117
+ print(f"Error deleting zip file: {e}")
118
+
119
+ for item in temp_dir.iterdir():
120
+ if item.is_dir():
121
+ try:
122
+ shutil.rmtree(item)
123
+ print(f"Deleted directory: {item}")
124
+ except Exception as e:
125
+ print(f"Error deleting directory {item}: {e}")
126
+
127
+ if temp_dir.exists() and not any(temp_dir.iterdir()):
128
+ try:
129
+ temp_dir.rmdir()
130
+ print(f"Deleted temp directory: {temp_dir}")
131
+ except Exception as e:
132
+ print(f"Error deleting temp directory: {e}")
133
+
134
+
135
+ if __name__ == "__main__":
136
+ sys.exit(main())
custom_nodes/ComfyUI-VFI/rife/model/__pycache__/loss.cpython-313.pyc ADDED
Binary file (10.3 kB). View file
 
custom_nodes/ComfyUI-VFI/rife/model/__pycache__/warplayer.cpython-313.pyc ADDED
Binary file (2.25 kB). View file
 
custom_nodes/ComfyUI-VFI/rife/model/loss.py ADDED
@@ -0,0 +1,130 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import torch
3
+ import torch.nn as nn
4
+ import torch.nn.functional as F
5
+ import torchvision.models as models
6
+
7
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
8
+
9
+
10
+ class EPE(nn.Module):
11
+ def __init__(self):
12
+ super(EPE, self).__init__()
13
+
14
+ def forward(self, flow, gt, loss_mask):
15
+ loss_map = (flow - gt.detach()) ** 2
16
+ loss_map = (loss_map.sum(1, True) + 1e-6) ** 0.5
17
+ return loss_map * loss_mask
18
+
19
+
20
+ class Ternary(nn.Module):
21
+ def __init__(self):
22
+ super(Ternary, self).__init__()
23
+ patch_size = 7
24
+ out_channels = patch_size * patch_size
25
+ self.w = np.eye(out_channels).reshape((patch_size, patch_size, 1, out_channels))
26
+ self.w = np.transpose(self.w, (3, 2, 0, 1))
27
+ self.w = torch.tensor(self.w).float().to(device)
28
+
29
+ def transform(self, img):
30
+ patches = F.conv2d(img, self.w, padding=3, bias=None)
31
+ transf = patches - img
32
+ transf_norm = transf / torch.sqrt(0.81 + transf**2)
33
+ return transf_norm
34
+
35
+ def rgb2gray(self, rgb):
36
+ r, g, b = rgb[:, 0:1, :, :], rgb[:, 1:2, :, :], rgb[:, 2:3, :, :]
37
+ gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
38
+ return gray
39
+
40
+ def hamming(self, t1, t2):
41
+ dist = (t1 - t2) ** 2
42
+ dist_norm = torch.mean(dist / (0.1 + dist), 1, True)
43
+ return dist_norm
44
+
45
+ def valid_mask(self, t, padding):
46
+ n, _, h, w = t.size()
47
+ inner = torch.ones(n, 1, h - 2 * padding, w - 2 * padding).type_as(t)
48
+ mask = F.pad(inner, [padding] * 4)
49
+ return mask
50
+
51
+ def forward(self, img0, img1):
52
+ img0 = self.transform(self.rgb2gray(img0))
53
+ img1 = self.transform(self.rgb2gray(img1))
54
+ return self.hamming(img0, img1) * self.valid_mask(img0, 1)
55
+
56
+
57
+ class SOBEL(nn.Module):
58
+ def __init__(self):
59
+ super(SOBEL, self).__init__()
60
+ self.kernelX = torch.tensor(
61
+ [
62
+ [1, 0, -1],
63
+ [2, 0, -2],
64
+ [1, 0, -1],
65
+ ]
66
+ ).float()
67
+ self.kernelY = self.kernelX.clone().T
68
+ self.kernelX = self.kernelX.unsqueeze(0).unsqueeze(0).to(device)
69
+ self.kernelY = self.kernelY.unsqueeze(0).unsqueeze(0).to(device)
70
+
71
+ def forward(self, pred, gt):
72
+ N, C, H, W = pred.shape[0], pred.shape[1], pred.shape[2], pred.shape[3]
73
+ img_stack = torch.cat([pred.reshape(N * C, 1, H, W), gt.reshape(N * C, 1, H, W)], 0)
74
+ sobel_stack_x = F.conv2d(img_stack, self.kernelX, padding=1)
75
+ sobel_stack_y = F.conv2d(img_stack, self.kernelY, padding=1)
76
+ pred_X, gt_X = sobel_stack_x[: N * C], sobel_stack_x[N * C :]
77
+ pred_Y, gt_Y = sobel_stack_y[: N * C], sobel_stack_y[N * C :]
78
+
79
+ L1X, L1Y = torch.abs(pred_X - gt_X), torch.abs(pred_Y - gt_Y)
80
+ loss = L1X + L1Y
81
+ return loss
82
+
83
+
84
+ class MeanShift(nn.Conv2d):
85
+ def __init__(self, data_mean, data_std, data_range=1, norm=True):
86
+ c = len(data_mean)
87
+ super(MeanShift, self).__init__(c, c, kernel_size=1)
88
+ std = torch.Tensor(data_std)
89
+ self.weight.data = torch.eye(c).view(c, c, 1, 1)
90
+ if norm:
91
+ self.weight.data.div_(std.view(c, 1, 1, 1))
92
+ self.bias.data = -1 * data_range * torch.Tensor(data_mean) # type: ignore
93
+ self.bias.data.div_(std) # type: ignore
94
+ else:
95
+ self.weight.data.mul_(std.view(c, 1, 1, 1))
96
+ self.bias.data = data_range * torch.Tensor(data_mean) # type: ignore
97
+ self.requires_grad = False
98
+
99
+
100
+ class VGGPerceptualLoss(torch.nn.Module):
101
+ def __init__(self, rank=0):
102
+ super(VGGPerceptualLoss, self).__init__()
103
+ blocks = [] # noqa: F841
104
+ pretrained = True
105
+ self.vgg_pretrained_features = models.vgg19(pretrained=pretrained).features
106
+ self.normalize = MeanShift([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], norm=True).cuda()
107
+ for param in self.parameters():
108
+ param.requires_grad = False
109
+
110
+ def forward(self, X, Y, indices=None):
111
+ X = self.normalize(X)
112
+ Y = self.normalize(Y)
113
+ indices = [2, 7, 12, 21, 30]
114
+ weights = [1.0 / 2.6, 1.0 / 4.8, 1.0 / 3.7, 1.0 / 5.6, 10 / 1.5]
115
+ k = 0
116
+ loss = 0
117
+ for i in range(indices[-1]):
118
+ X = self.vgg_pretrained_features[i](X) # type: ignore
119
+ Y = self.vgg_pretrained_features[i](Y) # type: ignore
120
+ if (i + 1) in indices:
121
+ loss += weights[k] * (X - Y.detach()).abs().mean() * 0.1
122
+ k += 1
123
+ return loss
124
+
125
+
126
+ if __name__ == "__main__":
127
+ img0 = torch.zeros(3, 3, 256, 256).float().to(device)
128
+ img1 = torch.tensor(np.random.normal(0, 1, (3, 3, 256, 256))).float().to(device)
129
+ ternary_loss = Ternary()
130
+ print(ternary_loss(img0, img1).shape)
custom_nodes/ComfyUI-VFI/rife/model/pytorch_msssim/__init__.py ADDED
@@ -0,0 +1,203 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from math import exp
2
+
3
+ import torch
4
+ import torch.nn.functional as F
5
+
6
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7
+
8
+
9
+ def gaussian(window_size, sigma):
10
+ gauss = torch.Tensor([exp(-((x - window_size // 2) ** 2) / float(2 * sigma**2)) for x in range(window_size)])
11
+ return gauss / gauss.sum()
12
+
13
+
14
+ def create_window(window_size, channel=1):
15
+ _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
16
+ _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0).to(device)
17
+ window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
18
+ return window
19
+
20
+
21
+ def create_window_3d(window_size, channel=1):
22
+ _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
23
+ _2D_window = _1D_window.mm(_1D_window.t())
24
+ _3D_window = _2D_window.unsqueeze(2) @ (_1D_window.t())
25
+ window = _3D_window.expand(1, channel, window_size, window_size, window_size).contiguous().to(device)
26
+ return window
27
+
28
+
29
+ def ssim(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None):
30
+ # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh).
31
+ if val_range is None:
32
+ if torch.max(img1) > 128:
33
+ max_val = 255
34
+ else:
35
+ max_val = 1
36
+
37
+ if torch.min(img1) < -0.5:
38
+ min_val = -1
39
+ else:
40
+ min_val = 0
41
+ L = max_val - min_val
42
+ else:
43
+ L = val_range
44
+
45
+ padd = 0
46
+ (_, channel, height, width) = img1.size()
47
+ if window is None:
48
+ real_size = min(window_size, height, width)
49
+ window = create_window(real_size, channel=channel).to(img1.device)
50
+
51
+ # mu1 = F.conv2d(img1, window, padding=padd, groups=channel)
52
+ # mu2 = F.conv2d(img2, window, padding=padd, groups=channel)
53
+ mu1 = F.conv2d(F.pad(img1, (5, 5, 5, 5), mode="replicate"), window, padding=padd, groups=channel)
54
+ mu2 = F.conv2d(F.pad(img2, (5, 5, 5, 5), mode="replicate"), window, padding=padd, groups=channel)
55
+
56
+ mu1_sq = mu1.pow(2)
57
+ mu2_sq = mu2.pow(2)
58
+ mu1_mu2 = mu1 * mu2
59
+
60
+ sigma1_sq = F.conv2d(F.pad(img1 * img1, (5, 5, 5, 5), "replicate"), window, padding=padd, groups=channel) - mu1_sq
61
+ sigma2_sq = F.conv2d(F.pad(img2 * img2, (5, 5, 5, 5), "replicate"), window, padding=padd, groups=channel) - mu2_sq
62
+ sigma12 = F.conv2d(F.pad(img1 * img2, (5, 5, 5, 5), "replicate"), window, padding=padd, groups=channel) - mu1_mu2
63
+
64
+ C1 = (0.01 * L) ** 2
65
+ C2 = (0.03 * L) ** 2
66
+
67
+ v1 = 2.0 * sigma12 + C2
68
+ v2 = sigma1_sq + sigma2_sq + C2
69
+ cs = torch.mean(v1 / v2) # contrast sensitivity
70
+
71
+ ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)
72
+
73
+ if size_average:
74
+ ret = ssim_map.mean()
75
+ else:
76
+ ret = ssim_map.mean(1).mean(1).mean(1)
77
+
78
+ if full:
79
+ return ret, cs
80
+ return ret
81
+
82
+
83
+ def ssim_matlab(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None):
84
+ # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh).
85
+ if val_range is None:
86
+ if torch.max(img1) > 128:
87
+ max_val = 255
88
+ else:
89
+ max_val = 1
90
+
91
+ if torch.min(img1) < -0.5:
92
+ min_val = -1
93
+ else:
94
+ min_val = 0
95
+ L = max_val - min_val
96
+ else:
97
+ L = val_range
98
+
99
+ padd = 0
100
+ (_, _, height, width) = img1.size()
101
+ if window is None:
102
+ real_size = min(window_size, height, width)
103
+ window = create_window_3d(real_size, channel=1).to(img1.device)
104
+ # Channel is set to 1 since we consider color images as volumetric images
105
+
106
+ img1 = img1.unsqueeze(1)
107
+ img2 = img2.unsqueeze(1)
108
+
109
+ mu1 = F.conv3d(F.pad(img1, (5, 5, 5, 5, 5, 5), mode="replicate"), window, padding=padd, groups=1)
110
+ mu2 = F.conv3d(F.pad(img2, (5, 5, 5, 5, 5, 5), mode="replicate"), window, padding=padd, groups=1)
111
+
112
+ mu1_sq = mu1.pow(2)
113
+ mu2_sq = mu2.pow(2)
114
+ mu1_mu2 = mu1 * mu2
115
+
116
+ sigma1_sq = F.conv3d(F.pad(img1 * img1, (5, 5, 5, 5, 5, 5), "replicate"), window, padding=padd, groups=1) - mu1_sq
117
+ sigma2_sq = F.conv3d(F.pad(img2 * img2, (5, 5, 5, 5, 5, 5), "replicate"), window, padding=padd, groups=1) - mu2_sq
118
+ sigma12 = F.conv3d(F.pad(img1 * img2, (5, 5, 5, 5, 5, 5), "replicate"), window, padding=padd, groups=1) - mu1_mu2
119
+
120
+ C1 = (0.01 * L) ** 2
121
+ C2 = (0.03 * L) ** 2
122
+
123
+ v1 = 2.0 * sigma12 + C2
124
+ v2 = sigma1_sq + sigma2_sq + C2
125
+ cs = torch.mean(v1 / v2) # contrast sensitivity
126
+
127
+ ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)
128
+
129
+ if size_average:
130
+ ret = ssim_map.mean()
131
+ else:
132
+ ret = ssim_map.mean(1).mean(1).mean(1)
133
+
134
+ if full:
135
+ return ret, cs
136
+ return ret
137
+
138
+
139
+ def msssim(img1, img2, window_size=11, size_average=True, val_range=None, normalize=False):
140
+ device = img1.device
141
+ weights = torch.FloatTensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333]).to(device)
142
+ levels = weights.size()[0]
143
+ mssim = []
144
+ mcs = []
145
+ for _ in range(levels):
146
+ sim, cs = ssim(img1, img2, window_size=window_size, size_average=size_average, full=True, val_range=val_range)
147
+ mssim.append(sim)
148
+ mcs.append(cs)
149
+
150
+ img1 = F.avg_pool2d(img1, (2, 2))
151
+ img2 = F.avg_pool2d(img2, (2, 2))
152
+
153
+ mssim = torch.stack(mssim)
154
+ mcs = torch.stack(mcs)
155
+
156
+ # Normalize (to avoid NaNs during training unstable models, not compliant with original definition)
157
+ if normalize:
158
+ mssim = (mssim + 1) / 2
159
+ mcs = (mcs + 1) / 2
160
+
161
+ pow1 = mcs**weights
162
+ pow2 = mssim**weights
163
+ # From Matlab implementation https://ece.uwaterloo.ca/~z70wang/research/iwssim/
164
+ output = torch.prod(pow1[:-1] * pow2[-1])
165
+ return output
166
+
167
+
168
+ # Classes to re-use window
169
+ class SSIM(torch.nn.Module):
170
+ def __init__(self, window_size=11, size_average=True, val_range=None):
171
+ super(SSIM, self).__init__()
172
+ self.window_size = window_size
173
+ self.size_average = size_average
174
+ self.val_range = val_range
175
+
176
+ # Assume 3 channel for SSIM
177
+ self.channel = 3
178
+ self.window = create_window(window_size, channel=self.channel)
179
+
180
+ def forward(self, img1, img2):
181
+ (_, channel, _, _) = img1.size()
182
+
183
+ if channel == self.channel and self.window.dtype == img1.dtype:
184
+ window = self.window
185
+ else:
186
+ window = create_window(self.window_size, channel).to(img1.device).type(img1.dtype)
187
+ self.window = window
188
+ self.channel = channel
189
+
190
+ _ssim = ssim(img1, img2, window=window, window_size=self.window_size, size_average=self.size_average)
191
+ dssim = (1 - _ssim) / 2
192
+ return dssim
193
+
194
+
195
+ class MSSSIM(torch.nn.Module):
196
+ def __init__(self, window_size=11, size_average=True, channel=3):
197
+ super(MSSSIM, self).__init__()
198
+ self.window_size = window_size
199
+ self.size_average = size_average
200
+ self.channel = channel
201
+
202
+ def forward(self, img1, img2):
203
+ return msssim(img1, img2, window_size=self.window_size, size_average=self.size_average)
custom_nodes/ComfyUI-VFI/rife/model/warplayer.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
4
+ backwarp_tenGrid = {}
5
+
6
+
7
+ def warp(tenInput, tenFlow):
8
+ k = (str(tenFlow.device), str(tenFlow.size()))
9
+ if k not in backwarp_tenGrid:
10
+ tenHorizontal = (
11
+ torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=device)
12
+ .view(1, 1, 1, tenFlow.shape[3])
13
+ .expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1)
14
+ )
15
+ tenVertical = (
16
+ torch.linspace(-1.0, 1.0, tenFlow.shape[2], device=device)
17
+ .view(1, 1, tenFlow.shape[2], 1)
18
+ .expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3])
19
+ )
20
+ backwarp_tenGrid[k] = torch.cat([tenHorizontal, tenVertical], 1).to(device)
21
+
22
+ tenFlow = torch.cat(
23
+ [
24
+ tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0),
25
+ tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0),
26
+ ],
27
+ 1,
28
+ )
29
+
30
+ g = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1)
31
+ return torch.nn.functional.grid_sample(
32
+ input=tenInput, grid=g, mode="bilinear", padding_mode="border", align_corners=True
33
+ )
custom_nodes/ComfyUI-VFI/rife/rife_comfyui_wrapper.py ADDED
@@ -0,0 +1,191 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from typing import List, Optional, Tuple
3
+
4
+ import torch
5
+ from torch.nn import functional as F
6
+
7
+ from .train_log.RIFE_HDv3 import Model
8
+
9
+
10
+ class RIFEWrapper:
11
+ """Wrapper for RIFE model to work with ComfyUI Image tensors"""
12
+
13
+ BASE_DIR = os.path.dirname(os.path.abspath(__file__))
14
+
15
+ def __init__(self, model_path, device: Optional[torch.device] = None, use_fp16: bool = False):
16
+ self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
17
+ self.use_fp16 = use_fp16 and torch.cuda.is_available()
18
+
19
+ torch.set_grad_enabled(False)
20
+ if torch.cuda.is_available():
21
+ torch.backends.cudnn.enabled = True
22
+ torch.backends.cudnn.benchmark = True
23
+ if hasattr(torch.backends.cuda, "enable_mem_efficient_sdp"):
24
+ torch.backends.cuda.enable_mem_efficient_sdp(True)
25
+
26
+ self.model = Model()
27
+ self.model.load_model(model_path, -1)
28
+ self.model.eval()
29
+ self.model.device()
30
+
31
+ # Convert to fp16 if requested
32
+ if self.use_fp16:
33
+ self.model.flownet = self.model.flownet.half()
34
+
35
+ def interpolate_frames(
36
+ self,
37
+ images: torch.Tensor,
38
+ source_fps: float,
39
+ target_fps: float,
40
+ scale: float = 1.0,
41
+ progress_callback=None,
42
+ batch_size: int = 8,
43
+ ) -> torch.Tensor:
44
+ """
45
+ Interpolate frames from source FPS to target FPS
46
+
47
+ Args:
48
+ images: ComfyUI Image tensor [N, H, W, C] in range [0, 1]
49
+ source_fps: Source frame rate
50
+ target_fps: Target frame rate
51
+ scale: Scale factor for processing
52
+ progress_callback: Optional callback function that accepts (current, total) parameters
53
+ batch_size: Number of frames to process in parallel (default: 8)
54
+
55
+ Returns:
56
+ Interpolated ComfyUI Image tensor [M, H, W, C] in range [0, 1]
57
+ """
58
+
59
+ assert images.dim() == 4 and images.shape[-1] == 3, "Input must be [N, H, W, C] with C=3"
60
+
61
+ if source_fps == target_fps:
62
+ return images
63
+
64
+ total_source_frames = images.shape[0]
65
+ height, width = images.shape[1:3]
66
+
67
+ # Calculate padding
68
+ tmp = max(128, int(128 / scale))
69
+ ph = ((height - 1) // tmp + 1) * tmp
70
+ pw = ((width - 1) // tmp + 1) * tmp
71
+ padding = (0, pw - width, 0, ph - height)
72
+
73
+ # Calculate frame positions
74
+ frame_positions = self._calculate_target_frame_positions(source_fps, target_fps, total_source_frames)
75
+
76
+ # Pre-allocate output on CPU (NOT GPU to avoid OOM)
77
+ output_frames = []
78
+
79
+ # Build interpolation job list
80
+ interp_job_list = []
81
+ output_index_map = {} # Maps job_idx -> output position
82
+
83
+ for out_idx, (source_idx1, source_idx2, interp_factor) in enumerate(frame_positions):
84
+ if interp_factor == 0.0 or source_idx1 == source_idx2:
85
+ # Direct copy, no interpolation needed
86
+ output_frames.append(images[source_idx1])
87
+ else:
88
+ # Need interpolation - add placeholder
89
+ output_frames.append(None)
90
+ job_idx = len(interp_job_list)
91
+ interp_job_list.append((source_idx1, source_idx2, interp_factor))
92
+ output_index_map[job_idx] = out_idx
93
+
94
+ # Process interpolations in batches with streaming
95
+ num_jobs = len(interp_job_list)
96
+ gpu_dtype = torch.float16 if self.use_fp16 else torch.float32
97
+
98
+ with torch.inference_mode():
99
+ for batch_start in range(0, num_jobs, batch_size):
100
+ batch_end = min(batch_start + batch_size, num_jobs)
101
+ current_batch_size = batch_end - batch_start
102
+
103
+ # Collect unique source frames needed for this batch
104
+ source_frames_needed = set()
105
+ for job_idx in range(batch_start, batch_end):
106
+ source_idx1, source_idx2, _ = interp_job_list[job_idx]
107
+ source_frames_needed.add(source_idx1)
108
+ source_frames_needed.add(source_idx2)
109
+
110
+ # Load only required source frames to GPU
111
+ source_cache = {}
112
+ for src_idx in source_frames_needed:
113
+ source_cache[src_idx] = images[src_idx].to(device=self.device, dtype=gpu_dtype)
114
+
115
+ # Prepare batch tensors on GPU
116
+ batch_I0 = torch.empty((current_batch_size, 3, ph, pw), dtype=gpu_dtype, device=self.device)
117
+ batch_I1 = torch.empty((current_batch_size, 3, ph, pw), dtype=gpu_dtype, device=self.device)
118
+ timesteps = []
119
+
120
+ for i, job_idx in enumerate(range(batch_start, batch_end)):
121
+ source_idx1, source_idx2, interp_factor = interp_job_list[job_idx]
122
+
123
+ # Get frames from cache (already on GPU)
124
+ I0 = source_cache[source_idx1].permute(2, 0, 1).unsqueeze(0)
125
+ I1 = source_cache[source_idx2].permute(2, 0, 1).unsqueeze(0)
126
+
127
+ # Pad
128
+ batch_I0[i] = F.pad(I0, padding)[0]
129
+ batch_I1[i] = F.pad(I1, padding)[0]
130
+ timesteps.append(interp_factor)
131
+
132
+ # Batch inference
133
+ interpolated_batch = self.model.inference_batch(batch_I0, batch_I1, timesteps, scale=scale)
134
+
135
+ # Transfer results to CPU and store
136
+ for i, job_idx in enumerate(range(batch_start, batch_end)):
137
+ output_idx = output_index_map[job_idx]
138
+ result = interpolated_batch[i, :, :height, :width].permute(1, 2, 0).cpu().to(torch.float32)
139
+ output_frames[output_idx] = result
140
+
141
+ # Update progress
142
+ if progress_callback:
143
+ progress_callback(batch_end, num_jobs)
144
+
145
+ # Cleanup batch memory immediately
146
+ del batch_I0, batch_I1, interpolated_batch, source_cache
147
+ if torch.cuda.is_available():
148
+ torch.cuda.empty_cache()
149
+
150
+ # Stack all output frames
151
+ result = torch.stack(output_frames, dim=0)
152
+
153
+ return result
154
+
155
+ def _calculate_target_frame_positions(
156
+ self, source_fps: float, target_fps: float, total_source_frames: int
157
+ ) -> List[Tuple[int, int, float]]:
158
+ """
159
+ Calculate which frames need to be generated for the target frame rate.
160
+
161
+ Returns:
162
+ List of (source_frame_index1, source_frame_index2, interpolation_factor) tuples
163
+ """
164
+ frame_positions = []
165
+
166
+ # Calculate the time duration of the video
167
+ duration = total_source_frames / source_fps
168
+
169
+ # Calculate number of target frames
170
+ total_target_frames = int(duration * target_fps)
171
+
172
+ for target_idx in range(total_target_frames):
173
+ # Calculate the time position of this target frame
174
+ target_time = target_idx / target_fps
175
+
176
+ # Calculate the corresponding position in source frames
177
+ source_position = target_time * source_fps
178
+
179
+ # Find the two source frames to interpolate between
180
+ source_idx1 = int(source_position)
181
+ source_idx2 = min(source_idx1 + 1, total_source_frames - 1)
182
+
183
+ # Calculate interpolation factor (0 means use frame1, 1 means use frame2)
184
+ if source_idx1 == source_idx2:
185
+ interpolation_factor = 0.0
186
+ else:
187
+ interpolation_factor = source_position - source_idx1
188
+
189
+ frame_positions.append((source_idx1, source_idx2, interpolation_factor))
190
+
191
+ return frame_positions
custom_nodes/ComfyUI-VFI/rife/train_log/IFNet_HDv3.py ADDED
@@ -0,0 +1,214 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ import torch.nn.functional as F
4
+
5
+ from ..model.warplayer import warp
6
+
7
+ # from train_log.refine import *
8
+
9
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
10
+
11
+
12
+ def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
13
+ return nn.Sequential(
14
+ nn.Conv2d(
15
+ in_planes,
16
+ out_planes,
17
+ kernel_size=kernel_size,
18
+ stride=stride,
19
+ padding=padding,
20
+ dilation=dilation,
21
+ bias=True,
22
+ ),
23
+ nn.LeakyReLU(0.2, True),
24
+ )
25
+
26
+
27
+ def conv_bn(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
28
+ return nn.Sequential(
29
+ nn.Conv2d(
30
+ in_planes,
31
+ out_planes,
32
+ kernel_size=kernel_size,
33
+ stride=stride,
34
+ padding=padding,
35
+ dilation=dilation,
36
+ bias=False,
37
+ ),
38
+ nn.BatchNorm2d(out_planes),
39
+ nn.LeakyReLU(0.2, True),
40
+ )
41
+
42
+
43
+ class Head(nn.Module):
44
+ def __init__(self):
45
+ super(Head, self).__init__()
46
+ self.cnn0 = nn.Conv2d(3, 16, 3, 2, 1)
47
+ self.cnn1 = nn.Conv2d(16, 16, 3, 1, 1)
48
+ self.cnn2 = nn.Conv2d(16, 16, 3, 1, 1)
49
+ self.cnn3 = nn.ConvTranspose2d(16, 4, 4, 2, 1)
50
+ self.relu = nn.LeakyReLU(0.2, True)
51
+
52
+ def forward(self, x, feat=False):
53
+ x0 = self.cnn0(x)
54
+ x = self.relu(x0)
55
+ x1 = self.cnn1(x)
56
+ x = self.relu(x1)
57
+ x2 = self.cnn2(x)
58
+ x = self.relu(x2)
59
+ x3 = self.cnn3(x)
60
+ if feat:
61
+ return [x0, x1, x2, x3]
62
+ return x3
63
+
64
+
65
+ class ResConv(nn.Module):
66
+ def __init__(self, c, dilation=1):
67
+ super(ResConv, self).__init__()
68
+ self.conv = nn.Conv2d(c, c, 3, 1, dilation, dilation=dilation, groups=1)
69
+ self.beta = nn.Parameter(torch.ones((1, c, 1, 1)), requires_grad=True)
70
+ self.relu = nn.LeakyReLU(0.2, True)
71
+
72
+ def forward(self, x):
73
+ return self.relu(self.conv(x) * self.beta + x)
74
+
75
+
76
+ class IFBlock(nn.Module):
77
+ def __init__(self, in_planes, c=64):
78
+ super(IFBlock, self).__init__()
79
+ self.conv0 = nn.Sequential(
80
+ conv(in_planes, c // 2, 3, 2, 1),
81
+ conv(c // 2, c, 3, 2, 1),
82
+ )
83
+ self.convblock = nn.Sequential(
84
+ ResConv(c),
85
+ ResConv(c),
86
+ ResConv(c),
87
+ ResConv(c),
88
+ ResConv(c),
89
+ ResConv(c),
90
+ ResConv(c),
91
+ ResConv(c),
92
+ )
93
+ self.lastconv = nn.Sequential(nn.ConvTranspose2d(c, 4 * 13, 4, 2, 1), nn.PixelShuffle(2))
94
+
95
+ def forward(self, x, flow=None, scale=1):
96
+ x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False)
97
+ if flow is not None:
98
+ flow = F.interpolate(flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) * 1.0 / scale
99
+ x = torch.cat((x, flow), 1)
100
+ feat = self.conv0(x)
101
+ feat = self.convblock(feat)
102
+ tmp = self.lastconv(feat)
103
+ tmp = F.interpolate(tmp, scale_factor=scale, mode="bilinear", align_corners=False)
104
+ flow = tmp[:, :4] * scale
105
+ mask = tmp[:, 4:5]
106
+ feat = tmp[:, 5:]
107
+ return flow, mask, feat
108
+
109
+
110
+ class IFNet(nn.Module):
111
+ def __init__(self):
112
+ super(IFNet, self).__init__()
113
+ self.block0 = IFBlock(7 + 8, c=192)
114
+ self.block1 = IFBlock(8 + 4 + 8 + 8, c=128)
115
+ self.block2 = IFBlock(8 + 4 + 8 + 8, c=96)
116
+ self.block3 = IFBlock(8 + 4 + 8 + 8, c=64)
117
+ self.block4 = IFBlock(8 + 4 + 8 + 8, c=32)
118
+ self.encode = Head()
119
+
120
+ # not used during inference
121
+ """
122
+ self.teacher = IFBlock(8+4+8+3+8, c=64)
123
+ self.caltime = nn.Sequential(
124
+ nn.Conv2d(16+9, 8, 3, 2, 1),
125
+ nn.LeakyReLU(0.2, True),
126
+ nn.Conv2d(32, 64, 3, 2, 1),
127
+ nn.LeakyReLU(0.2, True),
128
+ nn.Conv2d(64, 64, 3, 1, 1),
129
+ nn.LeakyReLU(0.2, True),
130
+ nn.Conv2d(64, 64, 3, 1, 1),
131
+ nn.LeakyReLU(0.2, True),
132
+ nn.Conv2d(64, 1, 3, 1, 1),
133
+ nn.Sigmoid()
134
+ )
135
+ """
136
+
137
+ def forward(
138
+ self,
139
+ x,
140
+ timestep=0.5,
141
+ scale_list=[8, 4, 2, 1],
142
+ training=False,
143
+ fastmode=True,
144
+ ensemble=False,
145
+ ):
146
+ if not training:
147
+ channel = x.shape[1] // 2
148
+ img0 = x[:, :channel]
149
+ img1 = x[:, channel:]
150
+ if not torch.is_tensor(timestep):
151
+ timestep = (x[:, :1].clone() * 0 + 1) * timestep
152
+ else:
153
+ timestep = timestep.repeat(1, 1, img0.shape[2], img0.shape[3])
154
+ f0 = self.encode(img0[:, :3])
155
+ f1 = self.encode(img1[:, :3])
156
+ flow_list = []
157
+ merged = []
158
+ mask_list = []
159
+ warped_img0 = img0
160
+ warped_img1 = img1
161
+ flow = None
162
+ mask = None
163
+ loss_cons = 0
164
+ block = [self.block0, self.block1, self.block2, self.block3, self.block4]
165
+ for i in range(5):
166
+ if flow is None:
167
+ flow, mask, feat = block[i](
168
+ torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1),
169
+ None,
170
+ scale=scale_list[i],
171
+ )
172
+ if ensemble:
173
+ print("warning: ensemble is not supported since RIFEv4.21")
174
+ else:
175
+ wf0 = warp(f0, flow[:, :2])
176
+ wf1 = warp(f1, flow[:, 2:4])
177
+ fd, m0, feat = block[i](
178
+ torch.cat(
179
+ (
180
+ warped_img0[:, :3],
181
+ warped_img1[:, :3],
182
+ wf0,
183
+ wf1,
184
+ timestep,
185
+ mask,
186
+ feat,
187
+ ),
188
+ 1,
189
+ ),
190
+ flow,
191
+ scale=scale_list[i],
192
+ )
193
+ if ensemble:
194
+ print("warning: ensemble is not supported since RIFEv4.21")
195
+ else:
196
+ mask = m0
197
+ flow = flow + fd
198
+ mask_list.append(mask)
199
+ flow_list.append(flow)
200
+ warped_img0 = warp(img0, flow[:, :2])
201
+ warped_img1 = warp(img1, flow[:, 2:4])
202
+ merged.append((warped_img0, warped_img1))
203
+ mask = torch.sigmoid(mask)
204
+ merged[4] = warped_img0 * mask + warped_img1 * (1 - mask)
205
+ if not fastmode:
206
+ print("contextnet is removed")
207
+ """
208
+ c0 = self.contextnet(img0, flow[:, :2])
209
+ c1 = self.contextnet(img1, flow[:, 2:4])
210
+ tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
211
+ res = tmp[:, :3] * 2 - 1
212
+ merged[4] = torch.clamp(merged[4] + res, 0, 1)
213
+ """
214
+ return flow_list, mask_list[4], merged
custom_nodes/ComfyUI-VFI/rife/train_log/RIFE_HDv3.py ADDED
@@ -0,0 +1,130 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from torch.nn.parallel import DistributedDataParallel as DDP
3
+ from torch.optim import AdamW
4
+
5
+ from ..model.loss import *
6
+ from .IFNet_HDv3 import *
7
+
8
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
9
+
10
+
11
+ class Model:
12
+ def __init__(self, local_rank=-1):
13
+ self.flownet = IFNet()
14
+ self.device()
15
+ self.optimG = AdamW(self.flownet.parameters(), lr=1e-6, weight_decay=1e-4)
16
+ self.epe = EPE()
17
+ self.version = 4.25
18
+ # self.vgg = VGGPerceptualLoss().to(device)
19
+ self.sobel = SOBEL()
20
+ if local_rank != -1:
21
+ self.flownet = DDP(self.flownet, device_ids=[local_rank], output_device=local_rank)
22
+
23
+ def train(self):
24
+ self.flownet.train()
25
+
26
+ def eval(self):
27
+ self.flownet.eval()
28
+
29
+ def device(self):
30
+ self.flownet.to(device)
31
+
32
+ def load_model(self, path, rank=0):
33
+ def convert(param):
34
+ if rank == -1:
35
+ return {k.replace("module.", ""): v for k, v in param.items() if "module." in k}
36
+ else:
37
+ return param
38
+
39
+ if rank <= 0:
40
+ if torch.cuda.is_available():
41
+ try:
42
+ # Try with weights_only=True for PyTorch >= 1.13
43
+ self.flownet.load_state_dict(convert(torch.load(path, weights_only=True)), False)
44
+ except TypeError:
45
+ # Fallback for older PyTorch versions
46
+ self.flownet.load_state_dict(convert(torch.load(path)), False)
47
+ else:
48
+ try:
49
+ # Try with weights_only=True for PyTorch >= 1.13
50
+ self.flownet.load_state_dict(
51
+ convert(torch.load(path, map_location="cpu", weights_only=True)),
52
+ False,
53
+ )
54
+ except TypeError:
55
+ # Fallback for older PyTorch versions
56
+ self.flownet.load_state_dict(
57
+ convert(torch.load(path, map_location="cpu")),
58
+ False,
59
+ )
60
+
61
+ def save_model(self, path, rank=0):
62
+ if rank == 0:
63
+ torch.save(self.flownet.state_dict(), "{}/flownet.pkl".format(path))
64
+
65
+ def inference(self, img0, img1, timestep=0.5, scale=1.0):
66
+ imgs = torch.cat((img0, img1), 1)
67
+ scale_list = [16 / scale, 8 / scale, 4 / scale, 2 / scale, 1 / scale]
68
+ flow, mask, merged = self.flownet(imgs, timestep, scale_list)
69
+ # Return only the final result to save memory
70
+ result = merged[-1]
71
+ # Clear intermediate results
72
+ del flow, mask, merged
73
+ return result
74
+
75
+ def inference_batch(self, batch_img0, batch_img1, timesteps, scale=1.0):
76
+ """Batch inference for multiple frame pairs at once"""
77
+ batch_size = batch_img0.shape[0]
78
+
79
+ # Concatenate all pairs
80
+ imgs = torch.cat((batch_img0, batch_img1), 1)
81
+
82
+ # Pre-calculate scale list
83
+ scale_list = [16 / scale, 8 / scale, 4 / scale, 2 / scale, 1 / scale]
84
+
85
+ # Process all timesteps (convert list to tensor if needed)
86
+ if isinstance(timesteps, list):
87
+ timesteps = torch.tensor(timesteps, device=batch_img0.device, dtype=batch_img0.dtype)
88
+
89
+ # Batch process through network
90
+ results = []
91
+ for i in range(batch_size):
92
+ flow, mask, merged = self.flownet(
93
+ imgs[i : i + 1], timesteps[i] if timesteps.dim() > 0 else timesteps, scale_list
94
+ )
95
+ results.append(merged[-1])
96
+ # Clear intermediate results
97
+ del flow, mask, merged
98
+
99
+ # Stack results
100
+ result = torch.cat(results, dim=0)
101
+ return result
102
+
103
+ def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None):
104
+ for param_group in self.optimG.param_groups:
105
+ param_group["lr"] = learning_rate
106
+ img0 = imgs[:, :3]
107
+ img1 = imgs[:, 3:]
108
+ if training:
109
+ self.train()
110
+ else:
111
+ self.eval()
112
+ scale = [16, 8, 4, 2, 1]
113
+ flow, mask, merged = self.flownet(torch.cat((imgs, gt), 1), scale=scale, training=training)
114
+ loss_l1 = (merged[-1] - gt).abs().mean()
115
+ loss_smooth = self.sobel(flow[-1], flow[-1] * 0).mean()
116
+ # loss_vgg = self.vgg(merged[-1], gt)
117
+ if training:
118
+ self.optimG.zero_grad()
119
+ loss_G = loss_l1 + loss_cons + loss_smooth * 0.1 # noqa: F405
120
+ loss_G.backward()
121
+ self.optimG.step()
122
+ else:
123
+ flow_teacher = flow[2] # noqa: F841
124
+ return merged[-1], {
125
+ "mask": mask,
126
+ "flow": flow[-1][:, :2],
127
+ "loss_l1": loss_l1,
128
+ "loss_cons": loss_cons, # noqa
129
+ "loss_smooth": loss_smooth,
130
+ }
custom_nodes/ComfyUI-VFI/rife/train_log/__pycache__/IFNet_HDv3.cpython-313.pyc ADDED
Binary file (9.5 kB). View file
 
custom_nodes/ComfyUI-VFI/rife/train_log/__pycache__/RIFE_HDv3.cpython-313.pyc ADDED
Binary file (7.11 kB). View file
 
custom_nodes/ComfyUI-VFI/rife/train_log/refine.py ADDED
@@ -0,0 +1,113 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ import torch.nn.functional as F
4
+
5
+ from ..model.warplayer import warp
6
+
7
+
8
+ def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
9
+ return nn.Sequential(
10
+ nn.Conv2d(
11
+ in_planes,
12
+ out_planes,
13
+ kernel_size=kernel_size,
14
+ stride=stride,
15
+ padding=padding,
16
+ dilation=dilation,
17
+ bias=True,
18
+ ),
19
+ nn.LeakyReLU(0.2, True),
20
+ )
21
+
22
+
23
+ def conv_woact(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
24
+ return nn.Sequential(
25
+ nn.Conv2d(
26
+ in_planes,
27
+ out_planes,
28
+ kernel_size=kernel_size,
29
+ stride=stride,
30
+ padding=padding,
31
+ dilation=dilation,
32
+ bias=True,
33
+ ),
34
+ )
35
+
36
+
37
+ def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
38
+ return nn.Sequential(
39
+ torch.nn.ConvTranspose2d(
40
+ in_channels=in_planes,
41
+ out_channels=out_planes,
42
+ kernel_size=4,
43
+ stride=2,
44
+ padding=1,
45
+ bias=True,
46
+ ),
47
+ nn.LeakyReLU(0.2, True),
48
+ )
49
+
50
+
51
+ class Conv2(nn.Module):
52
+ def __init__(self, in_planes, out_planes, stride=2):
53
+ super(Conv2, self).__init__()
54
+ self.conv1 = conv(in_planes, out_planes, 3, stride, 1)
55
+ self.conv2 = conv(out_planes, out_planes, 3, 1, 1)
56
+
57
+ def forward(self, x):
58
+ x = self.conv1(x)
59
+ x = self.conv2(x)
60
+ return x
61
+
62
+
63
+ c = 16
64
+
65
+
66
+ class Contextnet(nn.Module):
67
+ def __init__(self):
68
+ super(Contextnet, self).__init__()
69
+ self.conv1 = Conv2(3, c)
70
+ self.conv2 = Conv2(c, 2 * c)
71
+ self.conv3 = Conv2(2 * c, 4 * c)
72
+ self.conv4 = Conv2(4 * c, 8 * c)
73
+
74
+ def forward(self, x, flow):
75
+ x = self.conv1(x)
76
+ flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5
77
+ f1 = warp(x, flow)
78
+ x = self.conv2(x)
79
+ flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5
80
+ f2 = warp(x, flow)
81
+ x = self.conv3(x)
82
+ flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5
83
+ f3 = warp(x, flow)
84
+ x = self.conv4(x)
85
+ flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5
86
+ f4 = warp(x, flow)
87
+ return [f1, f2, f3, f4]
88
+
89
+
90
+ class Unet(nn.Module):
91
+ def __init__(self):
92
+ super(Unet, self).__init__()
93
+ self.down0 = Conv2(17, 2 * c)
94
+ self.down1 = Conv2(4 * c, 4 * c)
95
+ self.down2 = Conv2(8 * c, 8 * c)
96
+ self.down3 = Conv2(16 * c, 16 * c)
97
+ self.up0 = deconv(32 * c, 8 * c)
98
+ self.up1 = deconv(16 * c, 4 * c)
99
+ self.up2 = deconv(8 * c, 2 * c)
100
+ self.up3 = deconv(4 * c, c)
101
+ self.conv = nn.Conv2d(c, 3, 3, 1, 1)
102
+
103
+ def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1):
104
+ s0 = self.down0(torch.cat((img0, img1, warped_img0, warped_img1, mask, flow), 1))
105
+ s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1))
106
+ s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1))
107
+ s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1))
108
+ x = self.up0(torch.cat((s3, c0[3], c1[3]), 1))
109
+ x = self.up1(torch.cat((x, s2), 1))
110
+ x = self.up2(torch.cat((x, s1), 1))
111
+ x = self.up3(torch.cat((x, s0), 1))
112
+ x = self.conv(x)
113
+ return torch.sigmoid(x)