mirror of
https://github.com/vladmandic/automatic
synced 2026-08-29 16:41:01 +02:00
Vladimir Mandic
352 lines
14 KiB
Python
352 lines
14 KiB
Python
from fractions import Fraction
|
|
import os
|
|
import time
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
import einops
|
|
from PIL import Image
|
|
from modules import shared, errors ,timer, rife, processing
|
|
from modules.logger import log
|
|
from modules.video_models.video_utils import check_av
|
|
|
|
|
|
def get_video_filename(p:processing.StableDiffusionProcessingVideo):
|
|
from modules.image.namegen import FilenameGenerator
|
|
from modules.paths import resolve_output_path
|
|
namegen = FilenameGenerator(p, seed=p.seed if p is not None else 0, prompt=p.prompt if p is not None else '')
|
|
filename = namegen.apply(shared.opts.samples_filename_pattern if shared.opts.samples_filename_pattern and len(shared.opts.samples_filename_pattern) > 0 else "[seq]-[prompt_words]")
|
|
base_path = resolve_output_path(shared.opts.outdir_samples, shared.opts.outdir_video)
|
|
if shared.opts.save_to_dirs:
|
|
dirname = namegen.apply(shared.opts.directories_filename_pattern or "[prompt_words]")
|
|
dirfile = os.path.dirname(filename)
|
|
dirname = os.path.join(base_path, dirname, dirfile)
|
|
else:
|
|
dirname = base_path
|
|
if not os.path.exists(dirname):
|
|
os.makedirs(dirname, exist_ok=True)
|
|
filename = os.path.join(dirname, filename)
|
|
filename = namegen.sequence(filename)
|
|
filename = namegen.sanitize(filename)
|
|
return filename
|
|
|
|
|
|
def save_params(p, filename: str | None = None):
|
|
from modules.paths import params_path
|
|
if p is None:
|
|
dct = {}
|
|
else:
|
|
# sampler_index, sampler_shift, dynamic_shift, guidance_scale, guidance_true, init_image, init_strength, last_image, vae_type, vae_tile_frames, mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb, vlm_enhance, vlm_model, vlm_system_prompt, override_settings = args
|
|
dct = {
|
|
"Prompt": p.prompt,
|
|
"Negative prompt": p.negative_prompt,
|
|
"Steps": p.steps,
|
|
"Sampler": p.sampler_name,
|
|
"Seed": p.seed,
|
|
"Engine": p.video_engine,
|
|
"Model": p.video_model,
|
|
"Frames": p.frames,
|
|
"Size": f"{p.width}x{p.height}",
|
|
"Styles": ','.join(p.styles) if isinstance(p.styles, list) else p.styles,
|
|
}
|
|
params = ', '.join([f'{k}: {v}' for k, v in dct.items() if v is not None and v != ''])
|
|
fn = filename if filename is not None else params_path
|
|
with open(fn, "w", encoding="utf8") as file:
|
|
file.write(params)
|
|
|
|
|
|
def create_video_metadata(p: processing.StableDiffusionProcessingVideo | None, metadata: dict | None = None, filename: str | None = None):
|
|
metadata = metadata.copy() if metadata is not None else {}
|
|
if not shared.opts.image_metadata:
|
|
return metadata
|
|
if p is None:
|
|
return metadata
|
|
try:
|
|
info = processing.create_infotext(p)
|
|
except Exception as e:
|
|
log.debug(f'Video metadata: infotext failed: {e}')
|
|
info = ''
|
|
if len(info) == 0:
|
|
info = getattr(p, 'prompt', '') or ''
|
|
if len(info) == 0:
|
|
return metadata
|
|
title = os.path.basename(filename) if filename else 'SD.Next video'
|
|
metadata.setdefault('title', title)
|
|
metadata.setdefault('encoder', 'SD.Next')
|
|
metadata.setdefault('comment', info)
|
|
metadata.setdefault('description', info)
|
|
return metadata
|
|
|
|
|
|
def images_to_tensor(images):
|
|
if images is None or len(images) == 0:
|
|
return None
|
|
array = [torch.from_numpy(np.array(image)) for image in images]
|
|
tensor = torch.stack(array, dim=0) # n h w c
|
|
tensor = tensor.unsqueeze(0) # 1, n, h, w, c
|
|
tensor = tensor.permute(0, 4, 1, 2, 3).contiguous() # 1, c, n, h, w
|
|
tensor = (tensor.float() / 127.5) - 1.0 # from [0,255] to [-1,1]
|
|
# log.debug(f'Video output: images={len(images)} tensor={tensor.shape}')
|
|
return tensor
|
|
|
|
|
|
def numpy_to_tensor(images):
|
|
if images is None or len(images) == 0:
|
|
return None
|
|
images = (2.0 * images) - 1.0 # from [0,1] to [-1,1]
|
|
array = [torch.from_numpy(images[i]) for i in range(images.shape[0])]
|
|
tensor = torch.stack(array, dim=0) # n h w c
|
|
tensor = tensor.unsqueeze(0) # 1, n, h, w, c
|
|
tensor = tensor.permute(0, 4, 1, 2, 3).contiguous() # 1, c, n, h, w
|
|
# tensor = (tensor.float() / 127.5) - 1.0 # from [0,255] to [-1,1]
|
|
# log.debug(f'Video output: images={len(images)} tensor={tensor.shape}')
|
|
return tensor
|
|
|
|
|
|
def add_audio_stream(container, audio_sample_rate: int):
|
|
# Must be registered before the first container.mux(); avformat_write_header runs there
|
|
# and freezes the stream set, after which new streams have time_base=0/0.
|
|
audio_stream = container.add_stream("aac", rate=audio_sample_rate)
|
|
audio_stream.codec_context.sample_rate = audio_sample_rate
|
|
audio_stream.codec_context.layout = "stereo"
|
|
audio_stream.codec_context.time_base = Fraction(1, audio_sample_rate)
|
|
log.debug(f'Audio: codec={audio_stream.codec_context.name} rate={audio_stream.codec_context.sample_rate} layout={audio_stream.codec_context.layout} base={audio_stream.codec_context.time_base}')
|
|
return audio_stream
|
|
|
|
|
|
def write_audio(
|
|
container,
|
|
audio_stream,
|
|
samples: torch.Tensor,
|
|
audio_sample_rate: int,
|
|
) -> None:
|
|
av = check_av()
|
|
audio_stream.codec_context.format = "fltp"
|
|
if samples.ndim == 1:
|
|
samples = samples[:, None]
|
|
if samples.shape[1] != 2 and samples.shape[0] == 2:
|
|
samples = samples.T
|
|
if samples.shape[1] != 2:
|
|
raise ValueError(f"Expected samples with 2 channels; got shape {samples.shape}.")
|
|
if samples.dtype != torch.int16:
|
|
samples = torch.clip(samples, -1.0, 1.0)
|
|
samples = (samples * 32767.0).to(torch.int16)
|
|
audio_frames = av.AudioFrame.from_ndarray(
|
|
samples.contiguous().reshape(1, -1).cpu().numpy(),
|
|
format="s16",
|
|
layout="stereo",
|
|
)
|
|
audio_frames.sample_rate = audio_sample_rate
|
|
audio_resampler = av.audio.resampler.AudioResampler(
|
|
format=audio_stream.codec_context.format,
|
|
layout=audio_stream.codec_context.layout,
|
|
rate=audio_stream.codec_context.sample_rate,
|
|
)
|
|
pts = 0
|
|
for resampled in audio_resampler.resample(audio_frames):
|
|
resampled.pts = resampled.pts or 0
|
|
resampled.sample_rate = audio_frames.sample_rate
|
|
packets = audio_stream.encode(resampled)
|
|
for packet in packets:
|
|
container.mux(packet)
|
|
pts += resampled.samples
|
|
for packet in audio_stream.encode():
|
|
container.mux(packet)
|
|
|
|
|
|
def atomic_save_video(
|
|
filename: str,
|
|
tensor: torch.Tensor,
|
|
audio: torch.Tensor | None = None,
|
|
fps: float = 24,
|
|
codec: str = "libx264",
|
|
pix_fmt: str = "yuv420p",
|
|
options: str = "",
|
|
aac: int = 24000,
|
|
metadata: dict | None = None,
|
|
pbar=None,
|
|
):
|
|
if metadata is None:
|
|
metadata = {}
|
|
av = check_av()
|
|
if av is None or av is False:
|
|
log.error('Video: ffmpeg/av not available')
|
|
return
|
|
|
|
savejob = shared.state.begin('Save video')
|
|
frames, height, width, _channels = tensor.shape
|
|
rate = round(fps)
|
|
options_str = options
|
|
options = {}
|
|
for option in [option.strip() for option in options_str.split(',')]:
|
|
if '=' in option:
|
|
key, value = option.split('=', 1)
|
|
elif ':' in option:
|
|
key, value = option.split(':', 1)
|
|
else:
|
|
continue
|
|
options[key.strip()] = value.strip()
|
|
log.info(f'Video: file="{filename}" codec={codec} frames={frames} width={width} height={height} fps={rate} audio={audio is not None} aac={aac} options={options}')
|
|
video_array = torch.as_tensor(tensor, dtype=torch.uint8).numpy(force=True)
|
|
|
|
task = pbar.add_task('encoding', total=frames) if pbar is not None else None
|
|
if task is not None:
|
|
pbar.update(task, description='video encoding')
|
|
|
|
with av.open(filename, mode="w") as container:
|
|
for k, v in metadata.items():
|
|
container.metadata[k] = v
|
|
stream: av.VideoStream = container.add_stream(codec, rate=rate, options=options)
|
|
stream.width = video_array.shape[2]
|
|
stream.height = video_array.shape[1]
|
|
stream.pix_fmt = pix_fmt
|
|
audio_stream = add_audio_stream(container, aac) if audio is not None else None
|
|
for img in video_array:
|
|
frame = av.VideoFrame.from_ndarray(img, format="rgb24")
|
|
for packet in stream.encode_lazy(frame):
|
|
container.mux(packet)
|
|
if task is not None:
|
|
pbar.update(task, advance=1)
|
|
for packet in stream.encode(): # flush
|
|
container.mux(packet)
|
|
if audio_stream is not None:
|
|
try:
|
|
write_audio(container, audio_stream, audio, aac)
|
|
except Exception as e:
|
|
log.error(f'Video audio encoding: {e}')
|
|
errors.display(e, 'Audio')
|
|
|
|
shared.state.outputs(filename)
|
|
shared.state.end(savejob)
|
|
|
|
|
|
def save_thumbnail(video_path, tensor=None):
|
|
try:
|
|
base = os.path.splitext(video_path)[0]
|
|
thumb_path = f'{base}.thumb.jpg'
|
|
if tensor is not None and len(tensor) > 0:
|
|
frame = Image.fromarray(tensor[0].numpy())
|
|
else:
|
|
from modules.video import get_video_params
|
|
_frames, _fps, _dur, _w, _h, _codec, frame = get_video_params(video_path, capture=True)
|
|
if frame is not None:
|
|
frame.thumbnail((512, 512), Image.Resampling.LANCZOS)
|
|
frame.save(thumb_path, quality=80)
|
|
return thumb_path
|
|
except Exception as e:
|
|
log.debug(f'Video thumbnail: {e}')
|
|
return None
|
|
|
|
|
|
def save_video(
|
|
p: processing.StableDiffusionProcessingVideo,
|
|
pixels: torch.Tensor | None = None,
|
|
audio: torch.Tensor | None = None,
|
|
binary: bytes | None = None,
|
|
mp4_fps: int = 24,
|
|
mp4_codec: str = "libx264",
|
|
mp4_opt: str = "",
|
|
mp4_ext: str = "mp4",
|
|
mp4_sf: bool = False, # save safetensors
|
|
mp4_video: bool = True, # save video
|
|
mp4_frames: bool = False, # save frames
|
|
mp4_thumb: bool = True, # save thumbnail
|
|
mp4_interpolate: int = 0, # rife interpolation
|
|
aac_sample_rate: int = 24000, # audio sample rate
|
|
stream=None, # async progress reporting stream
|
|
metadata: dict | None = None, # metadata for video
|
|
pbar=None, # progress bar for video
|
|
):
|
|
if metadata is None:
|
|
metadata = {}
|
|
output_video = None
|
|
|
|
if binary is not None:
|
|
output_filename = get_video_filename(p)
|
|
output_video = f'{output_filename}.{mp4_ext}'
|
|
try:
|
|
try:
|
|
with open(output_video, 'wb') as f:
|
|
f.write(binary)
|
|
log.info(f'Video output: file="{output_video}" size={len(binary)}')
|
|
shared.state.outputs(output_video)
|
|
except Exception as e:
|
|
log.error(f'Video output: file="{output_video}" {e}')
|
|
except Exception as e:
|
|
log.error(f'Video output: file="{output_video}" write error {e}')
|
|
errors.display(e, 'video')
|
|
thumb = save_thumbnail(output_video) if mp4_thumb else None
|
|
return 0, output_video, thumb
|
|
|
|
if pixels is None:
|
|
return 0, output_video, None
|
|
if isinstance(pixels, np.ndarray):
|
|
pixels = numpy_to_tensor(pixels)
|
|
if isinstance(pixels, list) and isinstance(pixels[0], Image.Image):
|
|
pixels = images_to_tensor(pixels)
|
|
if not torch.is_tensor(pixels):
|
|
log.error(f'Video: type={type(pixels)} not a tensor')
|
|
return 0, output_video, None
|
|
t_save = time.time()
|
|
n, _c, t, h, w = pixels.shape
|
|
size = pixels.element_size() * pixels.numel()
|
|
log.debug(f'Video: video={mp4_video} export={mp4_frames} safetensors={mp4_sf} interpolate={mp4_interpolate}')
|
|
log.debug(f'Video: encode={t} raw={size} latent={pixels.shape} audio={audio.shape if audio is not None else None} fps={mp4_fps} codec={mp4_codec} ext={mp4_ext} options="{mp4_opt}"')
|
|
try:
|
|
preparejob = shared.state.begin('Prepare video')
|
|
if stream is not None:
|
|
stream.output_queue.push(('progress', (None, 'Saving video...')))
|
|
if mp4_interpolate > 0 and not getattr(p, 'video_interpolated', False):
|
|
x = pixels.squeeze(0).permute(1, 0, 2, 3)
|
|
x = (x.clamp(-1., 1.) + 1.0) * 0.5 # RIFE expects [0, 1]; video pixels are [-1, 1]
|
|
interpolated = rife.interpolate_nchw(x, count=mp4_interpolate+1)
|
|
pixels = torch.stack(interpolated, dim=0)
|
|
pixels = pixels.permute(1, 2, 0, 3, 4)
|
|
pixels = pixels * 2.0 - 1.0
|
|
|
|
n, _c, t, h, w = pixels.shape
|
|
x = torch.clamp(pixels.float(), -1., 1.) * 127.5 + 127.5
|
|
x = x.detach().cpu().to(torch.uint8)
|
|
x = einops.rearrange(x, '(m n) c t h w -> t (m h) (n w) c', n=n)
|
|
x = x.contiguous()
|
|
|
|
output_filename = get_video_filename(p)
|
|
if shared.opts.save_txt:
|
|
save_params(p, f'{output_filename}.txt')
|
|
save_params(p)
|
|
|
|
if mp4_sf:
|
|
fn = f'{output_filename}.safetensors'
|
|
log.info(f'Video export: file="{fn}" type=savetensors shape={x.shape}')
|
|
from safetensors.torch import save_file
|
|
shared.state.outputs(fn)
|
|
save_file({ 'frames': x }, fn, metadata={'format': 'video', 'frames': str(t), 'width': str(w), 'height': str(h), 'fps': str(mp4_fps), 'codec': mp4_codec, 'options': mp4_opt, 'ext': mp4_ext, 'interpolate': str(mp4_interpolate)})
|
|
|
|
if mp4_frames:
|
|
log.info(f'Video frames: files="{output_filename}-00000.jpg" frames={t} width={w} height={h}')
|
|
for i in range(t):
|
|
image = cv2.cvtColor(x[i].numpy(), cv2.COLOR_RGB2BGR)
|
|
fn = f'{output_filename}-{i:05d}.jpg'
|
|
shared.state.outputs(fn)
|
|
cv2.imwrite(fn, image)
|
|
|
|
shared.state.end(preparejob)
|
|
|
|
if mp4_video and (mp4_codec != 'none'):
|
|
output_video = f'{output_filename}.{mp4_ext}'
|
|
metadata = create_video_metadata(p, metadata, output_filename)
|
|
atomic_save_video(output_video, tensor=x, audio=audio, fps=mp4_fps, codec=mp4_codec, options=mp4_opt, aac=aac_sample_rate, metadata=metadata, pbar=pbar)
|
|
if stream is not None:
|
|
stream.output_queue.push(('progress', (None, f'Video {os.path.basename(output_video)} | Codec {mp4_codec} | Size {w}x{h}x{t} | FPS {mp4_fps}')))
|
|
stream.output_queue.push(('file', output_video))
|
|
else:
|
|
if stream is not None:
|
|
stream.output_queue.push(('progress', (None, '')))
|
|
|
|
except Exception as e:
|
|
log.error(f'Video save: raw={size} {e}')
|
|
errors.display(e, 'video')
|
|
timer.process.add('save', time.time()-t_save)
|
|
thumb = save_thumbnail(output_video) if mp4_thumb and output_video is not None else None
|
|
return t, output_video, thumb
|