73 lines
2.9 KiB
Python
73 lines
2.9 KiB
Python
# Copyright (C) 2023 Deforum LLC
|
|
#
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Affero General Public License as published by
|
|
# the Free Software Foundation, version 3 of the License.
|
|
#
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU Affero General Public License
|
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
|
|
# Contact the authors: https://deforum.github.io/
|
|
|
|
import os
|
|
import cv2
|
|
from modules.shared import opts
|
|
|
|
# Resume requires at least two actual frames in order to work
|
|
# 'Actual' frames are defined as frames that go through generation
|
|
# - Can't resume from a single frame.
|
|
# - If you have a cadence of 10, you need at least 10 frames in order to resume.
|
|
# - Resume grabs the last actual frame and the 2nd to last actual frame
|
|
# in order to work with cadence properly and feed it the prev_img/next_img
|
|
|
|
def get_resume_vars(folder, timestring, cadence):
|
|
DEBUG_MODE = opts.data.get("deforum_debug_mode_enabled", False)
|
|
# count previous frames
|
|
frame_count = 0
|
|
for item in os.listdir(folder):
|
|
# don't count txt files or mp4 files
|
|
if ".txt" in item or ".mp4" in item:
|
|
pass
|
|
else:
|
|
filename = item.split("_")
|
|
# other image file types may be supported in the future,
|
|
# so we just count files containing timestring
|
|
# that don't contain the depth keyword (depth maps are saved in same folder)
|
|
if timestring in filename and "depth" not in filename:
|
|
frame_count += 1
|
|
# add this to debugging var
|
|
if DEBUG_MODE:
|
|
print(f"\033[36mResuming:\033[0m File: {filename}")
|
|
|
|
print(f"\033[36mResuming:\033[0m Current frame count: {frame_count}")
|
|
|
|
# get last frame from frame count corrected for any trailing cadence frames
|
|
last_frame = frame_count - (frame_count % cadence)
|
|
|
|
# calculate previous actual frame
|
|
prev_frame = last_frame - cadence
|
|
|
|
# calculate next actual frame
|
|
next_frame = last_frame - 1
|
|
|
|
# get prev_img/next_img from prev/next frame index (files start at 0, so subtract 1 for index var)
|
|
path = os.path.join(folder, f"{timestring}_{prev_frame:09}.png")
|
|
prev_img = cv2.imread(path)
|
|
path = os.path.join(folder, f"{timestring}_{next_frame:09}.png")
|
|
next_img = cv2.imread(path)
|
|
|
|
# report resume last/next in console
|
|
print(f"\033[36mResuming:\033[0m Last frame: {prev_frame} - Next frame: {next_frame} ")
|
|
|
|
# returns:
|
|
# last frame count, accounting for cadence
|
|
# next frame count, accounting for cadence
|
|
# prev frame's image cv2 BGR
|
|
# next frame's image cv2 BGR
|
|
return prev_frame, next_frame, prev_img, next_img
|