from typing import Iterable, Tuple from .types import * from .categories import NodeCategories from .shared import DreamConfig, DreamImage import os, tempfile, subprocess, shutil, random from functools import lru_cache from PIL import Image CONFIG = DreamConfig() @lru_cache(5) def _load_image_cached(filename) -> Image: return Image.open(filename) class TempFileSet: def __init__(self): self._files = dict() def add(self, temppath, finalpath): self._files[temppath] = finalpath def remove(self): for f in self._files.keys(): os.unlink(f) def finalize(self): for a, b in self._files.items(): shutil.move(a, b) self._files = dict() class AnimationSeqProcessor: def __init__(self, sequence: AnimationSequence): self._sequence = sequence self._input_cache = {} self._inputs = {} self._output_dirs = {} for b in self._sequence.batches: self._inputs[b] = list(self._sequence.get_image_files_of_batch(b)) self._output_dirs[b] = os.path.dirname(os.path.abspath(self._inputs[b][0])) self._ext = os.path.splitext(self._inputs[0][0])[1].lower() self._length = len(self._inputs[0]) def _load_input(self, batch_id, index) -> DreamImage: files = self._inputs[batch_id] index = min(max(0, index), len(files) - 1) filename = files[index] return DreamImage(pil_image=_load_image_cached(filename)) def _process_single_batch(self, batch_id, indices, index_offsets: List[int], fun, output_dir) -> List[str]: all_indices = list(indices) last_index = max(all_indices) workset = TempFileSet() rnd = random.randint(0, 1000000) result_files = list() try: for index in all_indices: images = list(map(lambda offset: self._load_input(batch_id, index + offset), index_offsets)) result: Dict[int, DreamImage] = fun(index, last_index, images) for (result_index, img) in result.items(): filepath = os.path.join(output_dir, "tmp_" + str(rnd) + "_" + (str(result_index).zfill(8)) + self._ext) filepath_final = os.path.join(output_dir, "seq_" + (str(result_index).zfill(8)) + self._ext) if self._ext == ".png": img.save_png(filepath) else: img.save_jpg(filepath, quality=CONFIG.get("encoding.jpeg_quality", 98)) workset.add(filepath, filepath_final) result_files.append(filepath_final) # all done with batch - remove input files for oldfile in self._inputs[batch_id]: os.unlink(oldfile) workset.finalize() return result_files finally: workset.remove() def process(self, index_offsets: List[int], fun): results = dict() new_length = 0 for batch_id in self._sequence.batches: resulting_filenames = self._process_single_batch(batch_id, range(len(self._inputs[batch_id])), index_offsets, fun, self._output_dirs[batch_id]) for (index, filename) in enumerate(resulting_filenames): l = results.get(index, []) l.append(filename) results[index] = l new_length = len(resulting_filenames) new_fps = self._sequence.frame_counter.frames_per_second * (float(new_length) / self._length) counter = FrameCounter(new_length - 1, new_length, new_fps) return AnimationSequence(counter, results) def _ffmpeg(config, filenames, fps, output): fps = float(fps) duration = 1.0 / fps tmp = tempfile.NamedTemporaryFile(delete=False, mode="wb") tempfilepath = tmp.name try: for filename in filenames: filename = filename.replace("\\", "/") tmp.write(f"file '{filename}'\n".encode()) tmp.write(f"duration {duration}\n".encode()) finally: tmp.close() try: cmd = [config.get("ffmpeg.path", "ffmpeg")] cmd.extend(config.get("ffmpeg.arguments")) replacements = {"%FPS%": str(fps), "%FRAMES%": tempfilepath, "%OUTPUT%": output} for (key, value) in replacements.items(): cmd = list(map(lambda s: s.replace(key, value), cmd)) subprocess.run(cmd, shell=True) finally: os.unlink(tempfilepath) class DreamVideoEncoder: NODE_NAME = "FFMPEG Video Encoder" @classmethod def INPUT_TYPES(cls): return { "required": SharedTypes.sequence | { "filename": ("STRING", {"default": 'video.mp4', "multiline": False}), "framerate_factor": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 100.0}), "remove_images": (["yes", "no"],) }, } CATEGORY = NodeCategories.ANIMATION_POSTPROCESSING RETURN_TYPES = () RETURN_NAMES = () OUTPUT_NODE = True FUNCTION = "encode" @classmethod def IS_CHANGED(cls, sequence: AnimationSequence, **kwargs): return sequence.is_defined def _find_free_filename(self, filename, defaultdir): if os.path.basename(filename) == filename: filename = os.path.join(defaultdir, filename) n = 1 tested = filename while os.path.exists(tested): n += 1 (b, ext) = os.path.splitext(filename) tested = b + "_" + str(n) + ext return tested def generate_video(self, files, fps, filename, config): filename = self._find_free_filename(filename, os.path.dirname(files[0])) _ffmpeg(config, files, fps, filename) def encode(self, sequence: AnimationSequence, filename: str, remove_images, framerate_factor): if not sequence.is_defined: return () config = DreamConfig() for batch_num in sequence.batches: try: images = list(sequence.get_image_files_of_batch(batch_num)) self.generate_video(images, sequence.fps * framerate_factor, filename, config) if remove_images == "yes": for imagepath in images: if os.path.isfile(imagepath): os.unlink(imagepath) except Exception as e: print("Failed to encode files in dir {}!".format(os.path.dirname(images[0]))) print(str(e)) return () class DreamSequenceTweening: NODE_NAME = "Image Sequence Tweening" @classmethod def INPUT_TYPES(cls): return { "required": SharedTypes.sequence | { "multiplier": ("INT", {"default": 2, "min": 2, "max": 10}), }, } CATEGORY = NodeCategories.ANIMATION_POSTPROCESSING RETURN_TYPES = (AnimationSequence.ID,) RETURN_NAMES = ("sequence",) OUTPUT_NODE = False FUNCTION = "process" @classmethod def IS_CHANGED(cls, sequence: AnimationSequence, **kwargs): return sequence.is_defined def process(self, sequence: AnimationSequence, multiplier): if not sequence.is_defined: return (sequence,) def _generate_extra_frames(input_index, last_index, images): results = {} if input_index == last_index: # special case for i in range(multiplier): results[input_index * multiplier + i] = images[0] return results # normal case current_frame = images[0] next_frame = images[1] for i in range(multiplier): alpha = float(i + 1) / multiplier results[multiplier * input_index + i] = current_frame.blend(next_frame, 1.0 - alpha, alpha) return results proc = AnimationSeqProcessor(sequence) return (proc.process([0, 1], _generate_extra_frames),) class DreamSequenceBlend: NODE_NAME = "Image Sequence Blend" @classmethod def INPUT_TYPES(cls): return { "required": SharedTypes.sequence | { "fade_in": ("FLOAT", {"default": 0.1, "min": 0.01, "max": 0.5}), "fade_out": ("FLOAT", {"default": 0.1, "min": 0.01, "max": 0.5}), "iterations": ("INT", {"default": 1, "min": 1, "max": 10}), }, } CATEGORY = NodeCategories.ANIMATION_POSTPROCESSING RETURN_TYPES = (AnimationSequence.ID,) RETURN_NAMES = ("sequence",) OUTPUT_NODE = False FUNCTION = "process" @classmethod def IS_CHANGED(cls, sequence: AnimationSequence, **kwargs): return sequence.is_defined def process(self, sequence: AnimationSequence, fade_in, fade_out, iterations): if not sequence.is_defined: return (sequence,) current_sequence = sequence for i in range(iterations): proc = AnimationSeqProcessor(current_sequence) def _blur(index: int, last_index: int, images: List[DreamImage]): pre_frame = images[0].blend(images[1], fade_in, 1.0) post_frame = images[2].blend(images[1], fade_out, 1.0) return {index: pre_frame.blend(post_frame)} current_sequence = proc.process([-1, 0, 1], _blur) return (current_sequence,)