169 lines
6.9 KiB
Python
169 lines
6.9 KiB
Python
import math
|
|
import os
|
|
|
|
import cv2
|
|
from PIL import Image
|
|
|
|
from scripts.functions import keyframe_functions, postprocessing
|
|
from modules import processing, shared, sd_models
|
|
from modules.processing import Processed
|
|
from modules.shared import state
|
|
|
|
|
|
def main_process(myset: dict,
|
|
ptxt: processing.StableDiffusionProcessingTxt2Img) -> any:
|
|
|
|
frame_count = math.ceil(myset['fps'] * myset['total_time'])
|
|
shared.state.job_count = frame_count
|
|
|
|
df = keyframe_functions.process_keyframes(myset)
|
|
|
|
all_images = []
|
|
|
|
state.job_count = frame_count
|
|
|
|
# Need to check input source and load video if required.
|
|
if myset['source'] == 'video':
|
|
source_cap = cv2.VideoCapture(myset['source_file'])
|
|
else:
|
|
source_cap = None
|
|
|
|
# Post Processing object dicts
|
|
text_blocks = {}
|
|
props = {}
|
|
stamps = {}
|
|
|
|
last_frame = None
|
|
frame_save = 0
|
|
|
|
# Main loop through frames
|
|
for frame_no in range(frame_count):
|
|
|
|
if state.interrupted:
|
|
# Interrupt button pressed in WebUI
|
|
break
|
|
|
|
#
|
|
# Process Keyframes
|
|
#
|
|
# Check if keyframes exists for this frame
|
|
# print("process keyframes")
|
|
if frame_no in myset['keyframes']:
|
|
# Keyframes exist for this frame.
|
|
# print(f"\r\nKeyframe at {frame_no}: {myset['keyframes'][frame_no]}\r\n")
|
|
|
|
for keyframe in myset['keyframes'][frame_no]:
|
|
keyframe_command = keyframe[0].lower().strip()
|
|
# Check the command, should be first item.
|
|
if keyframe_command == "model" and len(keyframe) == 2:
|
|
# Time (s) | model | model name
|
|
info = sd_models.get_closet_checkpoint_match(keyframe[1].strip() + ".ckpt")
|
|
if info is None:
|
|
raise RuntimeError(f"Unknown checkpoint: {keyframe[1]}")
|
|
sd_models.reload_model_weights(shared.sd_model, info)
|
|
|
|
elif keyframe_command == "prop" and len(keyframe) == 6:
|
|
# Time (s) | prop | prop_filename | x pos | y pos | scale | rotation
|
|
# bit of a hack, no prop name is supplied, but same function is used to draw.
|
|
# so the command is passed in place of prop name, which will be ignored anyway.
|
|
props[len(props)] = keyframe
|
|
elif keyframe_command == "set_stamp" and len(keyframe) == 7:
|
|
# Time (s) | set_stamp | stamp_name | stamp_filename | x pos | y pos | scale | rotation
|
|
stamps[keyframe[1].strip()] = keyframe[1:]
|
|
elif keyframe_command == "clear_stamp" and len(keyframe) == 2:
|
|
# Time (s) | clear_stamp | stamp_name
|
|
if keyframe[1].strip() in stamps:
|
|
stamps.pop(keyframe[1].strip())
|
|
|
|
elif keyframe_command == "set_text" and len(keyframe) == 10:
|
|
# time_s | set_text | name | text_prompt | x | y | w | h | fore_color | back_color | font_name
|
|
text_blocks[keyframe[1].strip()] = keyframe[1:]
|
|
elif keyframe_command == "clear_text" and len(keyframe) == 2:
|
|
# Time (s) | clear_text | textblock_name
|
|
if keyframe[1].strip() in text_blocks:
|
|
text_blocks.pop(keyframe[1].strip())
|
|
|
|
#
|
|
# Get source frame
|
|
#
|
|
# print("Animator: Get/Generate Source Image.")
|
|
|
|
|
|
#
|
|
# Pre-process source frame
|
|
#
|
|
# print("Animator: Pre-process Source Frame.")
|
|
|
|
#
|
|
# Process source frame into destination frame
|
|
#
|
|
# print("Animator: Process Source Frame.")
|
|
|
|
# Set prompts
|
|
ptxt.prompt = str(df.loc[frame_no, ['pos_prompt']][0])
|
|
ptxt.negative_prompt = str(df.loc[frame_no, ['neg_prompt']][0])
|
|
|
|
ptxt.seed = int(df.loc[frame_no, ['seed_start']][0])
|
|
ptxt.subseed = None \
|
|
if df.loc[frame_no, ['seed_end']][0] is None else int(df.loc[frame_no, ['seed_end']][0])
|
|
ptxt.subseed_strength = None \
|
|
if df.loc[frame_no, ['seed_str']][0] is None else float(df.loc[frame_no, ['seed_str']][0])
|
|
# print(f"Frame:{frame_no} Seed:{ptxt.seed} Sub:{ptxt.subseed} Str:{ptxt.subseed_strength}")
|
|
|
|
# Check if a source is set, and grab frame from there. If not, process.
|
|
# TODO: Maybe figure out blending options for source frame and generated frame.
|
|
if myset['source'] == 'video':
|
|
source_cap.set(1, frame_no)
|
|
ret, tmp_array = source_cap.read()
|
|
post_processed_image = Image.fromarray(cv2.cvtColor(tmp_array, cv2.COLOR_BGR2RGB).astype('uint8'), 'RGB')
|
|
elif myset['source'] == 'images':
|
|
if frame_no >= len(source_cap):
|
|
post_processed_image = Image.open(source_cap[-1])
|
|
print('Out of frames, reverting to last frame!')
|
|
else:
|
|
post_processed_image = Image.open(source_cap[frame_no])
|
|
|
|
else:
|
|
processed = processing.process_images(ptxt)
|
|
post_processed_image = processed.images[0].copy()
|
|
|
|
if post_processed_image.mode != 'RGBA':
|
|
post_processed_image = post_processed_image.convert('RGBA')
|
|
#
|
|
# Post-process destination frame
|
|
#
|
|
# print("Animator: Post-Process Source Frame.")
|
|
if len(stamps) > 0:
|
|
post_processed_image = postprocessing.paste_prop(post_processed_image,
|
|
stamps,
|
|
shared.opts.animatoranon_prop_folder)
|
|
if len(text_blocks) > 0:
|
|
post_processed_image = postprocessing.render_text_block(post_processed_image, text_blocks)
|
|
|
|
#
|
|
# Save frame
|
|
#
|
|
# Create and save smoothed intermediate frames
|
|
if frame_no > 0 and myset['smoothing'] > 0 and not myset['film_interpolation']:
|
|
# working a frame behind, smooth from last_frame -> post_processed_image
|
|
for idx, img in enumerate(postprocessing.morph(last_frame, post_processed_image, myset['smoothing'])):
|
|
img.save(os.path.join(myset['output_path'], f"frame_{frame_save:05}.png"))
|
|
print(f"{frame_save:03}: {frame_no:03} > {idx} smooth frame")
|
|
frame_save += 1
|
|
|
|
# print("Animator: Save Frame")
|
|
if frame_no % int(myset['fps']) == 0:
|
|
all_images.append(post_processed_image)
|
|
|
|
post_processed_image.save(os.path.join(myset['output_path'], f"frame_{frame_save:05}.png"))
|
|
frame_save += 1
|
|
|
|
last_frame = post_processed_image.copy()
|
|
|
|
shared.state.current_image = post_processed_image
|
|
|
|
Processed(ptxt, all_images, 0, "")
|
|
print("Done.")
|
|
|
|
return all_images
|