mirror of https://github.com/vladmandic/automatic
143 lines
7.8 KiB
Python
143 lines
7.8 KiB
Python
import json
|
|
import os
|
|
import shutil
|
|
|
|
import gradio as gr
|
|
from modules import scripts
|
|
from modules.images import save_image
|
|
from modules.sd_samplers import KDiffusionSampler, sample_to_image
|
|
|
|
# configurable section
|
|
video_rate = 30
|
|
author = 'https://github.com/vladmandic'
|
|
cli_template = "ffmpeg -hide_banner -loglevel {loglevel} -hwaccel auto -y -framerate {framerate} -i {inpath}/%5d.jpg -r {videorate} {preset} {minterpolate} {flags} -metadata title='{description}' -metadata description='{info}' -metadata author='stable-diffusion' -metadata album_artist='{author}' '{outfile}'" # note: <https://wiki.multimedia.cx/index.php/FFmpeg_Metadata>
|
|
presets = {
|
|
'x264': '-vcodec libx264 -preset medium -crf 23',
|
|
'x265': '-vcodec libx265 -preset faster -crf 28',
|
|
'vpx-vp9': '-vcodec libvpx-vp9 -crf 34 -b:v 0 -deadline realtime -cpu-used 4',
|
|
'aom-av1': '-vcodec libaom-av1 -crf 28 -b:v 0 -usage realtime -cpu-used 8 -pix_fmt yuv444p',
|
|
}
|
|
|
|
# internal state variables
|
|
current_step = 0
|
|
orig_callback_state = KDiffusionSampler.callback_state
|
|
|
|
|
|
class Script(scripts.Script):
|
|
# script title to show in ui
|
|
def title(self):
|
|
return "Save animation of intermediate steps"
|
|
|
|
|
|
# is ui visible: process/postprocess triggers for always-visible scripts otherwise use run as entry point
|
|
def show(self, is_img2img):
|
|
return scripts.AlwaysVisible
|
|
|
|
|
|
# ui components
|
|
def ui(self, is_visible):
|
|
with gr.Accordion("Save animation", open = False, elem_id="save-animation"):
|
|
gr.HTML("""
|
|
<a href="https://github.com/vladmandic/generative-art/tree/main/extensions">
|
|
Creates animation sequence from denoised intermediate steps with video frame interpolation to achieve desired animation duration</a><br>""")
|
|
with gr.Row():
|
|
is_enabled = gr.Checkbox(label = "Script Enabled", value = False)
|
|
codec = gr.Radio(label = 'Codec', choices = ['x264', 'x265', 'vpx-vp9', 'aom-av1'], value = 'x264')
|
|
interpolation = gr.Radio(label = 'Interpolation', choices = ['none', 'mci', 'blend'], value = 'mci')
|
|
with gr.Row():
|
|
duration = gr.Slider(label = "Duration", minimum = 0.5, maximum = 120, step = 0.1, value = 10)
|
|
skip_steps = gr.Slider(label = "Skip steps", minimum = 0, maximum = 100, step = 1, value = 5)
|
|
with gr.Row():
|
|
debug = gr.Checkbox(label = "Debug info", value = False)
|
|
run_incomplete = gr.Checkbox(label = "Run on incomplete", value = True)
|
|
tmp_delete = gr.Checkbox(label = "Delete intermediate", value = True)
|
|
out_create = gr.Checkbox(label = "Create animation", value = True)
|
|
with gr.Row():
|
|
tmp_path = gr.Textbox(label = "Path for intermediate files", lines = 1, value = "intermediate")
|
|
out_path = gr.Textbox(label = "Path for output animation file", lines = 1, value = "animation")
|
|
|
|
return [is_enabled, codec, interpolation, duration, skip_steps, debug, run_incomplete, tmp_delete, out_create, tmp_path, out_path]
|
|
|
|
|
|
# runs on each step for always-visible scripts
|
|
def process(self, p, is_enabled, codec, interpolation, duration, skip_steps, debug, run_incomplete, tmp_delete, out_create, tmp_path, out_path):
|
|
if is_enabled:
|
|
def callback_state(self, d):
|
|
global current_step
|
|
current_step = d["i"] + 1
|
|
if (skip_steps == 0) or (current_step > skip_steps):
|
|
image = sample_to_image(samples = d["denoised"], index = 0, approximation = None)
|
|
inpath = os.path.join(p.outpath_samples, tmp_path)
|
|
save_image(image, inpath, "", extension = 'jpg', short_filename = True, no_prompt = True) # filename using 00000 format so its easier for ffmpeg sequence parsing
|
|
return orig_callback_state(self, d)
|
|
|
|
setattr(KDiffusionSampler, "callback_state", callback_state)
|
|
|
|
|
|
# run at the end of sequence for always-visible scripts
|
|
def postprocess(self, p, processed, is_enabled, codec, interpolation, duration, skip_steps, debug, run_incomplete, tmp_delete, out_create, tmp_path, out_path):
|
|
global current_step
|
|
setattr(KDiffusionSampler, "callback_state", orig_callback_state)
|
|
if not is_enabled:
|
|
return
|
|
# callback happened too early, it happens with large number of steps and some samplers or if interrupted
|
|
if vars(processed)['steps'] != current_step:
|
|
print('Save animation warning: postprocess early call', { 'current': current_step, 'target': vars(processed)['steps'] })
|
|
if not run_incomplete:
|
|
return
|
|
# create dictionary with all input and output parameters
|
|
v = vars(processed)
|
|
params = {
|
|
'prompt': v['prompt'],
|
|
'negative': v['negative_prompt'],
|
|
'seed': v['seed'],
|
|
'sampler': v['sampler_name'],
|
|
'cfgscale': v['cfg_scale'],
|
|
'steps': v['steps'],
|
|
'current': current_step,
|
|
'skip': skip_steps,
|
|
'info': v['info'].replace('\n', ' '),
|
|
'model': v['info'].split('Model:')[1].split()[0] if ("Model:" in v['info']) else "unknown", # parse string if model info is present
|
|
'embedding': v['info'].split('Used embeddings:')[1].split()[0] if ("Used embeddings:" in v['info']) else "none", # parse string if embedding info is present
|
|
'faces': v['face_restoration_model'],
|
|
'timestamp': v['job_timestamp'],
|
|
'inpath': os.path.join(p.outpath_samples, tmp_path),
|
|
'outpath': os.path.join(p.outpath_samples, out_path),
|
|
'codec': 'lib' + codec,
|
|
'duration': duration,
|
|
'interpolation': interpolation,
|
|
'loglevel': 'error',
|
|
'cli': cli_template,
|
|
'framerate': 1.0 * (current_step - skip_steps) / duration,
|
|
'videorate': video_rate,
|
|
'author': author,
|
|
'preset': presets[codec],
|
|
'flags': "-movflags +faststart",
|
|
'ffmpeg': shutil.which("ffmpeg"), # detect if ffmpeg executable is present in path
|
|
}
|
|
if debug:
|
|
params['loglevel'] = 'info'
|
|
print("Save animation params:", json.dumps(params, indent = 2))
|
|
if out_create:
|
|
if not os.path.isdir(params['inpath']) or not os.path.isdir(params['outpath']):
|
|
print('Save animation error: folder not found', params['inpath'], params['outpath'])
|
|
return
|
|
if params['ffmpeg'] is None:
|
|
print("Save animation error: ffmpeg not found:")
|
|
return
|
|
# append conditionals to dictionary
|
|
params['minterpolate'] = "" if (params['interpolation'] == "none") else "-vf minterpolate=mi_mode={mi},fifo".format(mi = params['interpolation'])
|
|
params['outfile'] = os.path.join(params['outpath'], str(params['seed']) + "-" + str(params['prompt'])) + ('.webm' if (params['codec'] == 'libvpx-vp9') else '.mp4')
|
|
params['description'] = "{prompt} | negative {negative} | seed {seed} | sampler {sampler} | cfgscale {cfgscale} | steps {steps} | current {current} | model {model} | embedding {embedding} | faces {faces} | timestamp {timestamp} | interpolation {interpolation}".format(**params)
|
|
print("Save animation creating movie sequence:", params['outfile'])
|
|
cmd = params['cli'].format(**params)
|
|
# actual ffmpeg call
|
|
os.system(cmd)
|
|
if tmp_delete:
|
|
for root, _dirs, files in os.walk(params['inpath']):
|
|
print("Save animation removing {n} files from temp folder: {path}".format(path = root, n = len(files)))
|
|
for file in files:
|
|
f = os.path.join(root, file)
|
|
if os.path.isfile(f):
|
|
os.remove(f)
|