298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-04-20 03:13:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

save mp4 and other optimization

Browse Source
This commit is contained in:
drunkplato 2024-12-20 20:39:24 +00:00 committed by Ubuntu
parent 7c7d9f54da
commit 4e788eb2ea
6 changed files with 240 additions and 179 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
custom_nodes/MemedeckComfyNodes/__init__.py
View File

@ -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"
}

67
custom_nodes/MemedeckComfyNodes/nodes_input.py Normal file
View File

@ -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))

251
custom_nodes/MemedeckComfyNodes/nodes_output.py
View File

@ -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}}
# return {"ui": {"videos": results}}

67
custom_nodes/MemedeckComfyNodes/nodes_preprocessing.py
View File

@ -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)

6
custom_nodes/MemedeckComfyNodes/requirements.txt
View File

@ -3,4 +3,8 @@ numpy
torch
Pillow
opencv-python
torchvision
torchvision
lxml
cairosvg
ltx-video@git+https://github.com/Lightricks/LTX-Video@ltx-video-0.9.1

15
memedeck.py
View File

@ -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({