nodes/LonicaMewinsky-ComfyUI-MakeFrame
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

enabled cuda

Browse Source
This commit is contained in:
LonicaMewinsky
2023-11-10 16:15:18 -06:00
committed by GitHub
parent e400a8c7df
commit 860adfee22
1 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
LoadAnim.py
+5 -3
View File
@@ -4,6 +4,8 @@ import torchvision.transforms as transforms
import numpy as np import numpy as np
import os import os
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class BreakFrames: class BreakFrames:
def __init__(self): def __init__(self):
pass pass
@@ -51,13 +53,13 @@ class BreakFrames:
ret, frame = video_capture.read() ret, frame = video_capture.read()
if ret: if ret:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
tensors.append(transformer(frame).unsqueeze(0)) tensors.append(transformer(frame).to(device).unsqueeze(0))
frame_count += 1 frame_count += 1
else: else:
break break
video_capture.release() video_capture.release()
if num_keyframes > 0: if num_keyframes > 0:
N = np.clip(num_keyframes, 2, len(tensors)-1) # Save a spot for first frame N = np.clip(num_keyframes, 2, len(tensors)-1)
differences = [torch.norm(tensors[i+1] - tensors[i], p=2) for i in range(len(tensors)-1)] differences = [torch.norm(tensors[i+1] - tensors[i], p=2) for i in range(len(tensors)-1)]
_, top_indices = torch.topk(torch.tensor(differences), k=N, largest=True) _, top_indices = torch.topk(torch.tensor(differences), k=N, largest=True)
keyframe_indices = sorted([index.item() + 1 for index in top_indices]) keyframe_indices = sorted([index.item() + 1 for index in top_indices])
@@ -75,4 +77,4 @@ NODE_CLASS_MAPPINGS = {
NODE_DISPLAY_NAME_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = {
"BreakFrames": "BreakFrames" "BreakFrames": "BreakFrames"
} }