Austin Mroz
414c5d3bb8
This ports the video format code written for the upstream changes to the ffmpeg implementation. It improves the quality of webm outputs and adds support for additional codecs (h264, h265, av1) It also improves the logging by passing errors and more selectively blocking the logging of encoders. While h265 has been included, most browsers will be unable to display the resulting video.
463 lines
15 KiB
Python
463 lines
15 KiB
Python
import os
|
|
import json
|
|
import torch
|
|
import numpy as np
|
|
from typing import Dict, List
|
|
from torch import Tensor
|
|
from torch.nn.functional import group_norm
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
from einops import rearrange
|
|
|
|
import folder_paths
|
|
import comfy.ldm.modules.diffusionmodules.openaimodel as openaimodel
|
|
import comfy.model_management as model_management
|
|
from comfy.model_base import BaseModel
|
|
from comfy.ldm.modules.attention import SpatialTransformer
|
|
from comfy.utils import load_torch_file, calculate_parameters
|
|
from nodes import KSampler
|
|
|
|
from .logger import logger
|
|
from .motion_module import MotionWrapper, VanillaTemporalModule
|
|
from .model_utils import get_available_models, get_model_path, get_model_hash
|
|
|
|
|
|
def forward_timestep_embed(
|
|
ts, x, emb, context=None, transformer_options={}, output_shape=None
|
|
):
|
|
for layer in ts:
|
|
if isinstance(layer, openaimodel.TimestepBlock):
|
|
x = layer(x, emb)
|
|
elif isinstance(layer, VanillaTemporalModule):
|
|
x = layer(x, context)
|
|
elif isinstance(layer, SpatialTransformer):
|
|
x = layer(x, context, transformer_options)
|
|
transformer_options["current_index"] += 1
|
|
elif isinstance(layer, openaimodel.Upsample):
|
|
x = layer(x, output_shape=output_shape)
|
|
else:
|
|
x = layer(x)
|
|
return x
|
|
|
|
|
|
def groupnorm_mm_factory(video_length: int):
|
|
def groupnorm_mm_forward(self, input: Tensor) -> Tensor:
|
|
# axes_factor normalizes batch based on total conds and unconds passed in batch;
|
|
# the conds and unconds per batch can change based on VRAM optimizations that may kick in
|
|
axes_factor = input.size(0) // video_length
|
|
|
|
input = rearrange(input, "(b f) c h w -> b c f h w", b=axes_factor)
|
|
input = group_norm(input, self.num_groups, self.weight, self.bias, self.eps)
|
|
input = rearrange(input, "b c f h w -> (b f) c h w", b=axes_factor)
|
|
return input
|
|
|
|
return groupnorm_mm_forward
|
|
|
|
|
|
orig_forward_timestep_embed = openaimodel.forward_timestep_embed
|
|
orig_maximum_batch_area = model_management.maximum_batch_area
|
|
orig_groupnorm_forward = torch.nn.GroupNorm.forward
|
|
openaimodel.forward_timestep_embed = forward_timestep_embed
|
|
|
|
motion_modules: Dict[str, MotionWrapper] = {}
|
|
|
|
|
|
def load_motion_module(model_name: str):
|
|
model_path = get_model_path(model_name)
|
|
model_hash = get_model_hash(model_path)
|
|
if model_hash not in motion_modules:
|
|
logger.info(f"Loading motion module {model_name}")
|
|
mm_state_dict = load_torch_file(model_path)
|
|
motion_module = MotionWrapper.from_pretrained(mm_state_dict, model_name)
|
|
|
|
params = calculate_parameters(mm_state_dict, "")
|
|
if model_management.should_use_fp16(model_params=params):
|
|
logger.info(f"Converting motion module to fp16.")
|
|
motion_module.half()
|
|
offload_device = model_management.unet_offload_device()
|
|
motion_module = motion_module.to(offload_device)
|
|
|
|
motion_modules[model_hash] = motion_module
|
|
|
|
return motion_modules[model_hash]
|
|
|
|
|
|
def inject_motion_module_to_unet_legacy(unet, motion_module: MotionWrapper):
|
|
for mm_idx, unet_idx in enumerate([1, 2, 4, 5, 7, 8, 10, 11]):
|
|
mm_idx0, mm_idx1 = mm_idx // 2, mm_idx % 2
|
|
unet.input_blocks[unet_idx].append(
|
|
motion_module.down_blocks[mm_idx0].motion_modules[mm_idx1]
|
|
)
|
|
|
|
for unet_idx in range(12):
|
|
mm_idx0, mm_idx1 = unet_idx // 3, unet_idx % 3
|
|
if unet_idx % 2 == 2:
|
|
unet.output_blocks[unet_idx].insert(
|
|
-1, motion_module.up_blocks[mm_idx0].motion_modules[mm_idx1]
|
|
)
|
|
else:
|
|
unet.output_blocks[unet_idx].append(
|
|
motion_module.up_blocks[mm_idx0].motion_modules[mm_idx1]
|
|
)
|
|
if motion_module.is_v2:
|
|
unet.middle_block.insert(-1, motion_module.mid_block.motion_modules[0])
|
|
|
|
unet.motion_module = motion_module
|
|
|
|
|
|
def eject_motion_module_from_unet_legacy(unet):
|
|
for unet_idx in [1, 2, 4, 5, 7, 8, 10, 11]:
|
|
unet.input_blocks[unet_idx].pop(-1)
|
|
|
|
for unet_idx in range(12):
|
|
if unet_idx % 2 == 2:
|
|
unet.output_blocks[unet_idx].pop(-2)
|
|
else:
|
|
unet.output_blocks[unet_idx].pop(-1)
|
|
|
|
if unet.motion_module.is_v2:
|
|
unet.middle_block.pop(-2)
|
|
|
|
del unet.motion_module
|
|
|
|
|
|
def inject_motion_module_to_unet(unet, motion_module: MotionWrapper):
|
|
for mm_idx, unet_idx in enumerate([1, 2, 4, 5, 7, 8, 10, 11]):
|
|
mm_idx0, mm_idx1 = mm_idx // 2, mm_idx % 2
|
|
unet.input_blocks[unet_idx].append(
|
|
motion_module.down_blocks[mm_idx0].motion_modules[mm_idx1]
|
|
)
|
|
|
|
for unet_idx in range(12):
|
|
mm_idx0, mm_idx1 = unet_idx // 3, unet_idx % 3
|
|
if unet_idx % 3 == 2 and unet_idx != 11:
|
|
unet.output_blocks[unet_idx].insert(
|
|
-1, motion_module.up_blocks[mm_idx0].motion_modules[mm_idx1]
|
|
)
|
|
else:
|
|
unet.output_blocks[unet_idx].append(
|
|
motion_module.up_blocks[mm_idx0].motion_modules[mm_idx1]
|
|
)
|
|
if motion_module.is_v2:
|
|
unet.middle_block.insert(-1, motion_module.mid_block.motion_modules[0])
|
|
|
|
unet.motion_module = motion_module
|
|
|
|
|
|
def eject_motion_module_from_unet(unet):
|
|
for unet_idx in [1, 2, 4, 5, 7, 8, 10, 11]:
|
|
unet.input_blocks[unet_idx].pop(-1)
|
|
|
|
for unet_idx in range(12):
|
|
if unet_idx % 3 == 2 and unet_idx != 11:
|
|
unet.output_blocks[unet_idx].pop(-2)
|
|
else:
|
|
unet.output_blocks[unet_idx].pop(-1)
|
|
|
|
if unet.motion_module.is_v2:
|
|
unet.middle_block.pop(-2)
|
|
|
|
del unet.motion_module
|
|
|
|
|
|
injectors = {
|
|
"legacy": inject_motion_module_to_unet_legacy,
|
|
"default": inject_motion_module_to_unet,
|
|
}
|
|
|
|
ejectors = {
|
|
"legacy": eject_motion_module_from_unet_legacy,
|
|
"default": eject_motion_module_from_unet,
|
|
}
|
|
|
|
|
|
class AnimateDiffModuleLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_name": (get_available_models(),),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MOTION_MODULE",)
|
|
CATEGORY = "Animate Diff"
|
|
FUNCTION = "load_motion_module"
|
|
|
|
def load_motion_module(
|
|
self,
|
|
model_name: str,
|
|
):
|
|
motion_module = load_motion_module(model_name)
|
|
|
|
return (motion_module,)
|
|
|
|
|
|
class AnimateDiffSampler(KSampler):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
inputs = {
|
|
"required": {
|
|
"motion_module": ("MOTION_MODULE",),
|
|
"inject_method": (["default", "legacy"],),
|
|
"frame_number": (
|
|
"INT",
|
|
{"default": 16, "min": 2, "max": 32, "step": 1},
|
|
),
|
|
}
|
|
}
|
|
inputs["required"].update(KSampler.INPUT_TYPES()["required"])
|
|
return inputs
|
|
|
|
FUNCTION = "animatediff_sample"
|
|
CATEGORY = "Animate Diff"
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.prev_beta = None
|
|
self.prev_linear_start = None
|
|
self.prev_linear_end = None
|
|
|
|
def override_beta_schedule(self, model: BaseModel):
|
|
logger.info(f"Override beta schedule.")
|
|
self.prev_beta = model.get_buffer("betas").cpu().clone()
|
|
self.prev_linear_start = model.linear_start
|
|
self.prev_linear_end = model.linear_end
|
|
model.register_schedule(
|
|
given_betas=None,
|
|
beta_schedule="sqrt_linear",
|
|
timesteps=1000,
|
|
linear_start=0.00085,
|
|
linear_end=0.012,
|
|
cosine_s=8e-3,
|
|
)
|
|
|
|
def restore_beta_schedule(self, model: BaseModel):
|
|
logger.info(f"Restoring beta schedule.")
|
|
model.register_schedule(
|
|
given_betas=self.prev_beta,
|
|
linear_start=self.prev_linear_start,
|
|
linear_end=self.prev_linear_end,
|
|
)
|
|
self.prev_beta = None
|
|
self.prev_linear_start = None
|
|
self.prev_linear_end = None
|
|
|
|
def inject_motion_module(
|
|
self, model, motion_module: MotionWrapper, inject_method: str, frame_number: int
|
|
):
|
|
model = model.clone()
|
|
unet = model.model.diffusion_model
|
|
|
|
logger.info(f"Injecting motion module with method {inject_method}.")
|
|
motion_module.set_video_length(frame_number)
|
|
injectors[inject_method](unet, motion_module)
|
|
self.override_beta_schedule(model.model)
|
|
if not motion_module.is_v2:
|
|
logger.info(f"Hacking GroupNorm.forward function.")
|
|
torch.nn.GroupNorm.forward = groupnorm_mm_factory(frame_number)
|
|
|
|
return model
|
|
|
|
def eject_motion_module(self, model, inject_method):
|
|
unet = model.model.diffusion_model
|
|
|
|
self.restore_beta_schedule(model.model)
|
|
if not unet.motion_module.is_v2:
|
|
logger.info(f"Restore GroupNorm.forward function.")
|
|
torch.nn.GroupNorm.forward = orig_groupnorm_forward
|
|
|
|
logger.info(f"Ejecting motion module with method {inject_method}.")
|
|
ejectors[inject_method](unet)
|
|
|
|
def animatediff_sample(
|
|
self,
|
|
motion_module,
|
|
inject_method,
|
|
frame_number,
|
|
model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
denoise=1.0,
|
|
):
|
|
model = self.inject_motion_module(
|
|
model, motion_module, inject_method, frame_number
|
|
)
|
|
|
|
init_frames = len(latent_image["samples"])
|
|
samples = latent_image["samples"][:init_frames, :, :, :].clone().cpu()
|
|
|
|
if init_frames < frame_number:
|
|
last_frame = samples[-1].unsqueeze(0)
|
|
repeated_last_frames = last_frame.repeat(
|
|
frame_number - init_frames, 1, 1, 1
|
|
)
|
|
samples = torch.cat((samples, repeated_last_frames), dim=0)
|
|
|
|
latent_image = {"samples": samples}
|
|
|
|
try:
|
|
return super().sample(
|
|
model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
denoise=denoise,
|
|
)
|
|
except:
|
|
raise
|
|
finally:
|
|
self.eject_motion_module(model, inject_method)
|
|
|
|
|
|
class AnimateDiffCombine:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"frame_rate": (
|
|
"INT",
|
|
{"default": 8, "min": 1, "max": 24, "step": 1},
|
|
),
|
|
"loop_count": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
|
|
"save_image": ([True, False],),
|
|
"filename_prefix": ("STRING", {"default": "animate_diff"}),
|
|
"format": (["image/gif", "image/webp"] +
|
|
["video/"+x[:-5] for x in folder_paths.get_filename_list("video_formats")],),
|
|
"pingpong": ([False, True],),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("GIF",)
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "Animate Diff"
|
|
FUNCTION = "generate_gif"
|
|
|
|
def generate_gif(
|
|
self,
|
|
images,
|
|
frame_rate: int,
|
|
loop_count: int,
|
|
save_image=True,
|
|
filename_prefix="AnimateDiff",
|
|
format="image/gif",
|
|
pingpong=False,
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
):
|
|
# convert images to numpy
|
|
frames: List[Image.Image] = []
|
|
for image in images:
|
|
img = 255.0 * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))
|
|
frames.append(img)
|
|
|
|
# save image
|
|
output_dir = (
|
|
folder_paths.get_output_directory()
|
|
if save_image
|
|
else folder_paths.get_temp_directory()
|
|
)
|
|
(
|
|
full_output_folder,
|
|
filename,
|
|
counter,
|
|
subfolder,
|
|
_,
|
|
) = folder_paths.get_save_image_path(filename_prefix, output_dir)
|
|
|
|
metadata = PngInfo()
|
|
if prompt is not None:
|
|
metadata.add_text("prompt", json.dumps(prompt))
|
|
if extra_pnginfo is not None:
|
|
for x in extra_pnginfo:
|
|
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
|
|
|
|
# save first frame as png to keep metadata
|
|
file = f"{filename}_{counter:05}_.png"
|
|
file_path = os.path.join(full_output_folder, file)
|
|
frames[0].save(
|
|
file_path,
|
|
pnginfo=metadata,
|
|
compress_level=4,
|
|
)
|
|
if pingpong:
|
|
frames = frames + frames[-2:0:-1]
|
|
|
|
format_type, format_ext = format.split("/")
|
|
|
|
if format_type == "image":
|
|
file = f"{filename}_{counter:05}_.{format_ext}"
|
|
file_path = os.path.join(full_output_folder, file)
|
|
frames[0].save(
|
|
file_path,
|
|
format=format_ext.upper(),
|
|
save_all=True,
|
|
append_images=frames[1:],
|
|
duration=round(1000 / frame_rate),
|
|
loop=loop_count,
|
|
compress_level=4,
|
|
)
|
|
else:
|
|
# save webm
|
|
import shutil
|
|
import subprocess
|
|
|
|
ffmpeg_path = shutil.which("ffmpeg")
|
|
if ffmpeg_path is None:
|
|
raise ProcessLookupError("Could not find ffmpeg")
|
|
video_format_path = folder_paths.get_full_path("video_formats", format_ext + ".json")
|
|
with open(video_format_path, 'r') as stream:
|
|
video_format = json.load(stream)
|
|
file = f"{filename}_{counter:05}_.{video_format['extension']}"
|
|
file_path = os.path.join(full_output_folder, file)
|
|
dimensions = f"{frames[0].width}x{frames[0].height}"
|
|
args = [ffmpeg_path, "-v", "error", "-f", "rawvideo", "-pix_fmt", "rgb24",
|
|
"-s", dimensions, "-r", str(frame_rate), "-i", "-"] \
|
|
+ video_format['main_pass'] + [file_path]
|
|
|
|
env=os.environ
|
|
if "environment" in video_format:
|
|
env.update(video_format["environment"])
|
|
with subprocess.Popen(args, stdin=subprocess.PIPE, env=env) as proc:
|
|
for frame in frames:
|
|
proc.stdin.write(frame.tobytes())
|
|
|
|
previews = [
|
|
{
|
|
"filename": file,
|
|
"subfolder": subfolder,
|
|
"type": "output" if save_image else "temp",
|
|
"format": format,
|
|
}
|
|
]
|
|
return {"ui": {"gifs": previews}}
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AnimateDiffModuleLoader": AnimateDiffModuleLoader,
|
|
"AnimateDiffCombine": AnimateDiffCombine,
|
|
"AnimateDiffSampler": AnimateDiffSampler,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AnimateDiffModuleLoader": "Animate Diff Module Loader",
|
|
"AnimateDiffSampler": "Animate Diff Sampler",
|
|
"AnimateDiffCombine": "Animate Diff Combine",
|
|
}
|