mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-09-09 17:27:05 +08:00
100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
import io
|
|
|
|
import av
|
|
import comfy.latent_formats
|
|
import comfy.model_base
|
|
import comfy.model_management
|
|
import comfy.model_patcher
|
|
import comfy.sd
|
|
import comfy.supported_models_base
|
|
import comfy.utils
|
|
import numpy as np
|
|
import torch
|
|
from ltx_video.models.autoencoders.vae_encode import get_vae_size_scale_factor
|
|
|
|
|
|
def encode_single_frame(output_file, image_array: np.ndarray, crf):
|
|
container = av.open(output_file, "w", format="mp4")
|
|
try:
|
|
stream = container.add_stream(
|
|
"h264", rate=1, options={"crf": str(crf), "preset": "veryfast"}
|
|
)
|
|
stream.height = image_array.shape[0]
|
|
stream.width = image_array.shape[1]
|
|
av_frame = av.VideoFrame.from_ndarray(image_array, format="rgb24").reformat(
|
|
format="yuv420p"
|
|
)
|
|
container.mux(stream.encode(av_frame))
|
|
container.mux(stream.encode())
|
|
finally:
|
|
container.close()
|
|
|
|
|
|
def decode_single_frame(video_file):
|
|
container = av.open(video_file)
|
|
try:
|
|
stream = next(s for s in container.streams if s.type == "video")
|
|
frame = next(container.decode(stream))
|
|
finally:
|
|
container.close()
|
|
return frame.to_ndarray(format="rgb24")
|
|
|
|
|
|
def videofy(image: torch.Tensor, crf=29):
|
|
if crf == 0:
|
|
return image
|
|
|
|
image_array = (image * 255.0).byte().cpu().numpy()
|
|
with io.BytesIO() as output_file:
|
|
encode_single_frame(output_file, image_array, crf)
|
|
video_bytes = output_file.getvalue()
|
|
with io.BytesIO(video_bytes) as video_file:
|
|
image_array = decode_single_frame(video_file)
|
|
tensor = torch.tensor(image_array, dtype=image.dtype, device=image.device) / 255.0
|
|
return tensor
|
|
|
|
|
|
def pad_tensor(tensor, target_len):
|
|
dim = 2
|
|
repeat_factor = target_len - tensor.shape[dim] # Ceiling division
|
|
last_element = tensor.select(dim, -1).unsqueeze(dim)
|
|
padding = last_element.repeat(1, 1, repeat_factor, 1, 1)
|
|
return torch.cat([tensor, padding], dim=dim)
|
|
|
|
|
|
def encode_media_conditioning(
|
|
init_media, vae, width, height, frames_number, image_compression, initial_latent
|
|
):
|
|
pixels = comfy.utils.common_upscale(
|
|
init_media.movedim(-1, 1), width, height, "bilinear", ""
|
|
).movedim(1, -1)
|
|
encode_pixels = pixels[:, :, :, :3]
|
|
if image_compression > 0:
|
|
for i in range(encode_pixels.shape[0]):
|
|
image = videofy(encode_pixels[i], image_compression)
|
|
encode_pixels[i] = image
|
|
|
|
encoded_latents = vae.encode(encode_pixels).float()
|
|
|
|
video_scale_factor, _, _ = get_vae_size_scale_factor(vae.first_stage_model)
|
|
video_scale_factor = video_scale_factor if frames_number > 1 else 1
|
|
target_len = (frames_number // video_scale_factor) + 1
|
|
encoded_latents = encoded_latents[:, :, :target_len]
|
|
|
|
if initial_latent is None:
|
|
initial_latent = encoded_latents
|
|
else:
|
|
if encoded_latents.shape[2] > initial_latent.shape[2]:
|
|
initial_latent = pad_tensor(initial_latent, encoded_latents.shape[2])
|
|
initial_latent[:, :, : encoded_latents.shape[2], ...] = encoded_latents
|
|
|
|
init_image_frame_number = init_media.shape[0]
|
|
if init_image_frame_number == 1:
|
|
result = pad_tensor(initial_latent, target_len)
|
|
elif init_image_frame_number % 8 != 1:
|
|
result = pad_tensor(initial_latent, target_len)
|
|
else:
|
|
result = initial_latent
|
|
|
|
return result
|