diff --git a/custom_nodes/MemedeckComfyNodes/__init__.py b/custom_nodes/MemedeckComfyNodes/__init__.py index 25cdb214..00e7433d 100644 --- a/custom_nodes/MemedeckComfyNodes/__init__.py +++ b/custom_nodes/MemedeckComfyNodes/__init__.py @@ -1,6 +1,9 @@ from .nodes_preprocessing import MD_LoadImageFromUrl, MD_CompressAdjustNode, MD_ImageToMotionPrompt from .nodes_model import MD_LoadVideoModel, MD_ImgToVideo, MD_VideoSampler -from .nodes_output import MD_SaveAnimatedWEBP, MD_SaveMP4 +from .nodes_output import MD_SaveMP4 + +from .nodes_input import MD_VideoInputs + NODE_CLASS_MAPPINGS = { # PREPROCESSING "Memedeck_ImageToMotionPrompt": MD_ImageToMotionPrompt, @@ -12,7 +15,9 @@ NODE_CLASS_MAPPINGS = { "Memedeck_VideoSampler": MD_VideoSampler, # POSTPROCESSING "Memedeck_SaveMP4": MD_SaveMP4, - "Memedeck_SaveAnimatedWEBP": MD_SaveAnimatedWEBP + # "Memedeck_SaveAnimatedWEBP": MD_SaveAnimatedWEBP, + # INPUT NODES + "Memedeck_VideoInputs": MD_VideoInputs # "Memedeck_SaveAnimatedGIF": MD_SaveAnimatedGIF } @@ -27,6 +32,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "Memedeck_ImgToVideo": "MemeDeck: Image To Video", "Memedeck_VideoSampler": "MemeDeck: Video Sampler", # POSTPROCESSING - "Memedeck_SaveMP4": "MemeDeck: Save MP4" + "Memedeck_SaveMP4": "MemeDeck: Save MP4", + # INPUT NODES + "Memedeck_VideoInputs": "MemeDeck: Video Inputs" # "Memedeck_SaveAnimatedGIF": "MemeDeck: Save Animated GIF" } diff --git a/custom_nodes/MemedeckComfyNodes/nodes_input.py b/custom_nodes/MemedeckComfyNodes/nodes_input.py new file mode 100644 index 00000000..efaa5a2f --- /dev/null +++ b/custom_nodes/MemedeckComfyNodes/nodes_input.py @@ -0,0 +1,67 @@ +from comfy_extras.nodes_custom_sampler import Noise_RandomNoise + + +class MD_VideoInputs: + """One node to load all input parameters for video generation""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image_url": ( + "STRING", + ), + "length": ("INT", { + "default": 121, + "description": "The length of the video." + }), + "steps": ("INT", { + "default": 25, + "description": "Number of steps to generate the video." + }), + "width": ("INT", { + "default": 768, + "description": "The width of the video." + }), + "height": ("INT", { + "default": 768, + "description": "The height of the video." + }), + "crf": ("INT", { + "default": 28, + "min": 0, + "max": 51, + "step": 1 + }), + "terminal": ("FLOAT", { + "default": 0.1, + "step": 0.01, + "description": "The terminal values of the sigmas after stretching." + }), + }, + "optional": { + "seed": ( + "INT", + ), + "user_prompt": ( + "STRING", + ), + "pre_prompt": ( + "STRING", + ), + "post_prompt": ( + "STRING", + ), + "negative_prompt": ( + "STRING", + ), + } + } + + RETURN_TYPES = ("STRING", "INT", "INT", "INT", "INT", "INT", "FLOAT", "STRING", "STRING", "STRING", "STRING", "NOISE",) + RETURN_NAMES = ("image_url", "length", "steps", "width", "height", "crf", "terminal", "user_prompt", "pre_prompt", "post_prompt", "negative_prompt", "seed") + FUNCTION = "load_inputs" + CATEGORY = "MemeDeck" + + def load_inputs(self, image_url, length=121, steps=25, width=768, height=768, crf=28, terminal=0.1, user_prompt="", pre_prompt="", post_prompt="", negative_prompt="", seed=None): + return (image_url, length, steps, width, height, crf, terminal, user_prompt, pre_prompt, post_prompt, negative_prompt, Noise_RandomNoise(seed)) \ No newline at end of file diff --git a/custom_nodes/MemedeckComfyNodes/nodes_output.py b/custom_nodes/MemedeckComfyNodes/nodes_output.py index 75dc1553..def6173d 100644 --- a/custom_nodes/MemedeckComfyNodes/nodes_output.py +++ b/custom_nodes/MemedeckComfyNodes/nodes_output.py @@ -1,5 +1,8 @@ from io import BytesIO +import subprocess import time +import uuid +from custom_nodes.MemedeckComfyNodes.nodes_preprocessing import ffmpeg_process import folder_paths from comfy.cli_args import args import torch @@ -11,6 +14,7 @@ import numpy as np import json import os import logging +from .lib import utils # setup logger logger = logging.getLogger(__name__) @@ -23,7 +27,7 @@ WATERMARK = """ """ WATERMARK_SIZE = 28 -class MD_SaveAnimatedWEBP: +class MD_SaveMP4: def __init__(self): self.output_dir = folder_paths.get_output_directory() self.type = "output" @@ -37,9 +41,9 @@ class MD_SaveAnimatedWEBP: "images": ("IMAGE", ), "filename_prefix": ("STRING", {"default": "memedeck_video"}), "fps": ("FLOAT", {"default": 24.0, "min": 0.01, "max": 1000.0, "step": 0.01}), - "lossless": ("BOOLEAN", {"default": False}), - "quality": ("INT", {"default": 90, "min": 0, "max": 100}), - "method": (list(s.methods.keys()),), + # "lossless": ("BOOLEAN", {"default": False}), + # "quality": ("INT", {"default": 90, "min": 0, "max": 100}), + # "method": (list(s.methods.keys()),), "crf": ("FLOAT",), "motion_prompt": ("STRING", ), "negative_prompt": ("STRING", ), @@ -56,9 +60,8 @@ class MD_SaveAnimatedWEBP: CATEGORY = "MemeDeck" - def save_images(self, images, fps, filename_prefix, lossless, quality, method, crf=None, motion_prompt=None, negative_prompt=None, img2vid_metadata=None, sampler_metadata=None): + def save_images(self, images, fps, filename_prefix, crf=None, motion_prompt=None, negative_prompt=None, img2vid_metadata=None, sampler_metadata=None): start_time = time.time() - method = self.methods.get(method) filename_prefix += self.prefix_append full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) results = [] @@ -67,15 +70,18 @@ class MD_SaveAnimatedWEBP: pil_images = [Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8)) for image in images] first_image = pil_images[0] + width = first_image.width + height = first_image.height + padding = 8 - x = first_image.width - WATERMARK_SIZE - padding - y = first_image.height - WATERMARK_SIZE - padding + x = width - WATERMARK_SIZE - padding + y = height - WATERMARK_SIZE - padding first_image_background_brightness = self.analyze_background_brightness(first_image, x, y, WATERMARK_SIZE) watermarked_images = [self.add_watermark_to_image(img, first_image_background_brightness) for img in pil_images] - metadata = pil_images[0].getexif() - num_frames = len(pil_images) + # metadata = pil_images[0].getexif() + # num_frames = len(pil_images) json_metadata = { "crf": crf, @@ -83,99 +89,64 @@ class MD_SaveAnimatedWEBP: "negative_prompt": negative_prompt, "img2vid_metadata": json.loads(img2vid_metadata), "sampler_metadata": json.loads(sampler_metadata), - } + } - # Optimized saving logic - if num_frames == 1: # Single image, save once - file = f"{filename}_{counter:05}_.webp" - watermarked_images[0].save(os.path.join(full_output_folder, file), exif=metadata, lossless=lossless, quality=quality, method=method) - results.append({ - "filename": file, - "subfolder": subfolder, - "type": self.type, - }) - else: # multiple images, save as animation - file = f"{filename}_{counter:05}_.webp" - watermarked_images[0].save(os.path.join(full_output_folder, file), save_all=True, duration=int(1000.0 / fps), append_images=watermarked_images[1:], exif=metadata, lossless=lossless, quality=quality, method=method) - results.append({ - "filename": file, - "subfolder": subfolder, - "type": self.type, - }) + # Use ffmpeg to create MP4 with watermark + output_file = f"{filename}_{counter:05}_.mp4" + output_path = os.path.join(full_output_folder, output_file) + ffmpeg_cmd = [ + utils.ffmpeg_path, + "-v", "error", + '-f', 'rawvideo', + '-pix_fmt', 'rgb24', + "-r", str(fps), # Set frame rate + "-s", f"{width}x{height}", + "-i", "-", + "-y", # Overwrite output file if it exists + "-c:v", "libx264", # Use x264 encoder + "-crf", "19", # Set CRF (quality) + "-pix_fmt", "yuv420p", # Set pixel format + ] - animated = num_frames != 1 + env = os.environ.copy() + output_process = ffmpeg_process(ffmpeg_cmd, output_path, env) + + # Proceed to first yield + output_process.send(None) + output_process.send(first_image.tobytes()) + for image in watermarked_images: + output_process.send(image.tobytes()) + try: + output_process.send(None) # Signal end of input + next(output_process) # Get the final yield + except StopIteration: + pass + + results.append({ + "filename": output_file, + "subfolder": subfolder, + "type": self.type, + }) end_time = time.time() logger.info(f"Save images took: {end_time - start_time} seconds") + + preview = { + "filename": output_path, + "subfolder": subfolder, + "type": "output", + "format": "video/mp4", + "frame_rate": fps, + } return { "ui": { - "images": results, - "animated": (animated,), + "gifs": [preview], "metadata": (json.dumps(json_metadata),) }, + "result": ((True, [output_path]),) } - # def save_images(self, images, fps, filename_prefix, lossless, quality, method, crf=None, motion_prompt=None, negative_prompt=None, img2vid_metadata=None, sampler_metadata=None): - # start_time = time.time() - # method = self.methods.get(method) - # filename_prefix += self.prefix_append - # full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) - # results = [] - - # # Vectorized conversion to PIL images - # pil_images = [Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8)) for image in images] - - # first_image = pil_images[0] - # padding = 12 - # x = first_image.width - WATERMARK_SIZE - padding - # y = first_image.height - WATERMARK_SIZE - padding - # first_image_background_brightness = self.analyze_background_brightness(first_image, x, y, WATERMARK_SIZE) - - # watermarked_images = [self.add_watermark_to_image(img, first_image_background_brightness) for img in pil_images] - - # metadata = pil_images[0].getexif() - # num_frames = len(pil_images) - - # json_metadata = { - # "crf": crf, - # "motion_prompt": motion_prompt, - # "negative_prompt": negative_prompt, - # "img2vid_metadata": json.loads(img2vid_metadata), - # "sampler_metadata": json.loads(sampler_metadata), - # } - - # # Optimized saving logic - # if num_frames == 1: # Single image, save once - # file = f"{filename}_{counter:05}_.webp" - # watermarked_images[0].save(os.path.join(full_output_folder, file), exif=metadata, lossless=lossless, quality=quality, method=method) - # results.append({ - # "filename": file, - # "subfolder": subfolder, - # "type": self.type, - # }) - # else: # multiple images, save as animation - # file = f"{filename}_{counter:05}_.webp" - # watermarked_images[0].save(os.path.join(full_output_folder, file), save_all=True, duration=int(1000.0 / fps), append_images=watermarked_images[1:], exif=metadata, lossless=lossless, quality=quality, method=method) - # results.append({ - # "filename": file, - # "subfolder": subfolder, - # "type": self.type, - # }) - - # animated = num_frames != 1 - - # end_time = time.time() - # logger.info(f"Save images took: {end_time - start_time} seconds") - - # return { - # "ui": { - # "images": results, - # "animated": (animated,), - # "metadata": (json.dumps(json_metadata),) - # }, - # } - def add_watermark_to_image(self, img, background_brightness=None): """ Adds a watermark to a single PIL Image. @@ -402,63 +373,63 @@ class MD_VAEDecode: return (images,) -class MD_SaveMP4: - def __init__(self): - # Get absolute path of the output directory - self.output_dir = os.path.abspath("output/video_gen") - self.type = "output" - self.prefix_append = "" +# class MD_SaveMP4: +# def __init__(self): +# # Get absolute path of the output directory +# self.output_dir = os.path.abspath("output/video_gen") +# self.type = "output" +# self.prefix_append = "" - methods = {"default": 4, "fastest": 0, "slowest": 6} +# methods = {"default": 4, "fastest": 0, "slowest": 6} - @classmethod - def INPUT_TYPES(s): - return {"required": - {"images": ("IMAGE", ), - "filename_prefix": ("STRING", {"default": "ComfyUI"}), - "fps": ("FLOAT", {"default": 24.0, "min": 0.01, "max": 1000.0, "step": 0.01}), - "quality": ("INT", {"default": 80, "min": 0, "max": 100}), - }, - "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, - } +# @classmethod +# def INPUT_TYPES(s): +# return {"required": +# {"images": ("IMAGE", ), +# "filename_prefix": ("STRING", {"default": "ComfyUI"}), +# "fps": ("FLOAT", {"default": 24.0, "min": 0.01, "max": 1000.0, "step": 0.01}), +# "quality": ("INT", {"default": 80, "min": 0, "max": 100}), +# }, +# "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, +# } - RETURN_TYPES = () - FUNCTION = "save_video" +# RETURN_TYPES = () +# FUNCTION = "save_video" - OUTPUT_NODE = True +# OUTPUT_NODE = True - CATEGORY = "MemeDeck" +# CATEGORY = "MemeDeck" - def save_video(self, images, fps, filename_prefix, quality, prompt=None, extra_pnginfo=None): - filename_prefix += self.prefix_append - full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path( - filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0] - ) - results = list() - video_path = os.path.join(full_output_folder, f"{filename}_{counter:05}.mp4") +# def save_video(self, images, fps, filename_prefix, quality, prompt=None, extra_pnginfo=None): +# filename_prefix += self.prefix_append +# full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path( +# filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0] +# ) +# results = list() +# video_path = os.path.join(full_output_folder, f"{filename}_{counter:05}.mp4") - # Determine video resolution - height, width = images[0].shape[1], images[0].shape[2] - video_writer = cv2.VideoWriter( - video_path, - cv2.VideoWriter_fourcc(*'mp4v'), - fps, - (width, height) - ) +# # Determine video resolution +# height, width = images[0].shape[1], images[0].shape[2] +# video_writer = cv2.VideoWriter( +# video_path, +# cv2.VideoWriter_fourcc(*'mp4v'), +# fps, +# (width, height) +# ) - # Write each frame to the video - for image in images: - i = 255. * image.cpu().numpy() - frame = np.clip(i, 0, 255).astype(np.uint8) - frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) # Convert RGB to BGR for OpenCV - video_writer.write(frame) +# # Write each frame to the video +# for image in images: +# i = 255. * image.cpu().numpy() +# frame = np.clip(i, 0, 255).astype(np.uint8) +# frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) # Convert RGB to BGR for OpenCV +# video_writer.write(frame) - video_writer.release() +# video_writer.release() - results.append({ - "filename": os.path.basename(video_path), - "subfolder": subfolder, - "type": self.type - }) +# results.append({ +# "filename": os.path.basename(video_path), +# "subfolder": subfolder, +# "type": self.type +# }) - return {"ui": {"videos": results}} \ No newline at end of file +# return {"ui": {"videos": results}} \ No newline at end of file diff --git a/custom_nodes/MemedeckComfyNodes/nodes_preprocessing.py b/custom_nodes/MemedeckComfyNodes/nodes_preprocessing.py index 5e56c5aa..142d4ff3 100644 --- a/custom_nodes/MemedeckComfyNodes/nodes_preprocessing.py +++ b/custom_nodes/MemedeckComfyNodes/nodes_preprocessing.py @@ -23,6 +23,29 @@ import logging logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) +def ffmpeg_process(args, file_path, env): + res = None + frame_data = yield + total_frames_output = 0 + if res != b'': + with subprocess.Popen(args + [file_path], stderr=subprocess.PIPE, + stdin=subprocess.PIPE, env=env) as proc: + try: + while frame_data is not None: + proc.stdin.write(frame_data) + frame_data = yield + total_frames_output+=1 + proc.stdin.flush() + proc.stdin.close() + res = proc.stderr.read() + except BrokenPipeError as e: + res = proc.stderr.read() + raise Exception("An error occurred in the ffmpeg subprocess:\n" \ + + res.decode("utf-8")) + yield total_frames_output + if len(res) > 0: + print(res.decode("utf-8"), end="", file=sys.stderr) + class MD_LoadImageFromUrl: """Load an image from the given URL""" @@ -66,6 +89,13 @@ class MD_ImageToMotionPrompt: "default": "masterpiece, 4k, HDR, cinematic,", }, ), + "post_prompt": ( + "STRING", + { + "multiline": False, + "default": "The scene appears to be from a movie or TV show.", + }, + ), "prompt": ( "STRING", { @@ -91,7 +121,7 @@ class MD_ImageToMotionPrompt: CATEGORY = "MemeDeck" def generate_completion( - self, pre_prompt: str, Image: torch.Tensor, clip, prompt: str, negative_prompt: str, max_tokens: int + self, pre_prompt: str, post_prompt: str, Image: torch.Tensor, clip, prompt: str, negative_prompt: str, max_tokens: int ) -> Tuple[str]: # start a timer start_time = time.time() @@ -104,7 +134,7 @@ class MD_ImageToMotionPrompt: end_time = time.time() logger.info(f"Motion prompt took: {end_time - start_time} seconds") - full_prompt = f"{pre_prompt}\n{response.json()['result']}" + full_prompt = f"{pre_prompt}\n{response.json()['result']} {post_prompt}" pos_tokens = clip.tokenize(full_prompt) pos_output = clip.encode_from_tokens(pos_tokens, return_pooled=True, return_dict=True) @@ -165,29 +195,6 @@ class MD_CompressAdjustNode: def tensor_to_bytes(self, tensor): return self.tensor_to_int(tensor, 8).astype(np.uint8) - - def ffmpeg_process(self, args, file_path, env): - res = None - frame_data = yield - total_frames_output = 0 - if res != b'': - with subprocess.Popen(args + [file_path], stderr=subprocess.PIPE, - stdin=subprocess.PIPE, env=env) as proc: - try: - while frame_data is not None: - proc.stdin.write(frame_data) - frame_data = yield - total_frames_output+=1 - proc.stdin.flush() - proc.stdin.close() - res = proc.stderr.read() - except BrokenPipeError as e: - res = proc.stderr.read() - raise Exception("An error occurred in the ffmpeg subprocess:\n" \ - + res.decode("utf-8")) - yield total_frames_output - if len(res) > 0: - print(res.decode("utf-8"), end="", file=sys.stderr) def detect_image_clarity(self, image): # detect the clarity of the image @@ -286,7 +293,7 @@ class MD_CompressAdjustNode: i_pix_fmt = 'rgb24' # default bitrate and frame rate - frame_rate = 25 + frame_rate = 24 image_cv2 = cv2.cvtColor(np.array(tensor2pil(image)), cv2.COLOR_RGB2BGR) # calculate the crf based on the image @@ -295,7 +302,11 @@ class MD_CompressAdjustNode: self.ideal_color_variation, self.blockiness_weight, self.edge_density_weight, self.color_variation_weight) - logger.info(f"detected crf: {calculated_crf}") + if desired_crf is 0: + desired_crf = calculated_crf + + logger.info(f"calculated_crf: {calculated_crf}") + logger.info(f"desired_crf: {desired_crf}") args = [ utils.ffmpeg_path, "-v", "error", @@ -312,7 +323,7 @@ class MD_CompressAdjustNode: video_path = os.path.abspath(str(Path(temp_dir) / f"{filename}.mp4")) env = os.environ.copy() - output_process = self.ffmpeg_process(args, video_path, env) + output_process = ffmpeg_process(args, video_path, env) # Proceed to first yield output_process.send(None) diff --git a/custom_nodes/MemedeckComfyNodes/requirements.txt b/custom_nodes/MemedeckComfyNodes/requirements.txt index 29a06726..d47723d8 100644 --- a/custom_nodes/MemedeckComfyNodes/requirements.txt +++ b/custom_nodes/MemedeckComfyNodes/requirements.txt @@ -3,4 +3,8 @@ numpy torch Pillow opencv-python -torchvision \ No newline at end of file +torchvision +lxml +cairosvg + +ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9.1 diff --git a/memedeck.py b/memedeck.py index 0c33fb52..a4c63a5f 100644 --- a/memedeck.py +++ b/memedeck.py @@ -414,23 +414,24 @@ class MemedeckWorker: if event == "executed": if data['node'] == task['end_node_id']: - filename = data['output']['images'][0]['filename'] + # self.logger.info(f"[memedeck]: video gen completed {data}") + file_path = data['output']['gifs'][0]['filename'] metadata = json.loads(data['output']['metadata'][0]) self.logger.info(f"[memedeck]: video gen completed {metadata}") - current_dir = os.path.dirname(os.path.abspath(__file__)) - file_path = os.path.join(current_dir, "output", filename) - blob_name = f"{task['user_id']}/video_gen/video_{task['image_id'].replace('image:', '')}_{task['prompt_id']}.webp" + # current_dir = os.path.dirname(os.path.abspath(__file__)) + # file_path = os.path.join(current_dir, "output", filename) + blob_name = f"{task['user_id']}/video_gen/video_{task['image_id'].replace('image:', '')}_{task['prompt_id']}.mp4" # TODO: take the file path and upload to azure blob storage # load image bytes - with open(file_path, "rb") as image_file: - image_bytes = image_file.read() + with open(file_path, "rb") as video_file: + video_bytes = video_file.read() self.logger.info(f"[memedeck]: video gen completed for {sid}, file={file_path}, blob={blob_name}") - url = await self.azure_storage.save_image(blob_name, "image/webp", image_bytes) + url = await self.azure_storage.save_image(blob_name, "video/mp4", video_bytes) self.logger.info(f"[memedeck]: video gen completed for {sid}, {url}") await self.send_to_api({