mirror of
https://github.com/vladmandic/automatic
synced 2026-09-14 18:48:43 +02:00
CalamitousFelicitousness
f01e752b06
Pipelines rarely report a sample rate, so the save path fell back to 24000. LTX-2.3 and 2.5 run at 48k, and muxing at half the rate drops the track an octave. The rate now comes from the vocoder, as the LTX tab already did.
387 lines
16 KiB
Python
387 lines
16 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_audio_rate(p=None, default: int = 24000) -> int:
|
|
# pipeline output wins when it reports a rate, else the loaded vocoder: LTX-2.0 runs at 24k,
|
|
# 2.3 and 2.5 at 48k, and muxing at the wrong rate shifts the pitch
|
|
rate = getattr(p, 'audio_sampling_rate', None) if p is not None else None
|
|
if not rate:
|
|
vocoder = getattr(shared.sd_model, 'vocoder', None)
|
|
rate = getattr(getattr(vocoder, 'config', None), 'output_sampling_rate', None)
|
|
return int(rate) if rate else default
|
|
|
|
|
|
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_packets(container, audio_stream, audio: dict):
|
|
if not audio or "frames" not in audio:
|
|
return
|
|
try:
|
|
av = check_av()
|
|
sr = audio.get("sr", 44100)
|
|
layout = audio.get("layout", "stereo")
|
|
resampler = av.AudioResampler(format="fltp", layout=layout, rate=sr)
|
|
fifo = av.AudioFifo()
|
|
for raw_frame in audio.get("frames", []):
|
|
for resampled in resampler.resample(raw_frame):
|
|
fifo.write(resampled)
|
|
for resampled in resampler.resample(None):
|
|
fifo.write(resampled)
|
|
pts_counter = 0
|
|
frame_size = audio_stream.codec_context.frame_size or 1024
|
|
while fifo.samples >= frame_size:
|
|
frame = fifo.read(frame_size)
|
|
frame.pts = pts_counter
|
|
pts_counter += frame.samples
|
|
for packet in audio_stream.encode(frame):
|
|
packet.stream = audio_stream
|
|
container.mux_one(packet)
|
|
if fifo.samples > 0:
|
|
frame = fifo.read(fifo.samples)
|
|
frame.pts = pts_counter
|
|
pts_counter += frame.samples
|
|
for packet in audio_stream.encode(frame):
|
|
packet.stream = audio_stream
|
|
container.mux_one(packet)
|
|
for packet in audio_stream.encode():
|
|
packet.stream = audio_stream
|
|
container.mux_one(packet)
|
|
except Exception as e:
|
|
log.error(f"Video audio encoding: type=packets {e}")
|
|
errors.display(e, "Audio")
|
|
|
|
|
|
def add_audio_tensor(container, audio_stream, audio: torch.Tensor, sample_rate: int):
|
|
av = check_av()
|
|
if torch.is_tensor(audio):
|
|
audio = audio.detach().float().cpu().numpy()
|
|
if audio.ndim > 2:
|
|
audio = np.squeeze(audio)
|
|
if audio.ndim == 1:
|
|
audio = audio[None, :]
|
|
elif audio.ndim == 2 and audio.shape[0] > audio.shape[1] and audio.shape[1] in (1, 2):
|
|
audio = audio.T
|
|
channels = audio.shape[0] if audio.shape[0] in (1, 2) else 1
|
|
layout = "stereo" if channels == 2 else "mono"
|
|
if audio.dtype != np.int16:
|
|
audio = np.clip(audio, -1.0, 1.0)
|
|
audio = (audio * 32767.0).astype(np.int16)
|
|
audio_frame = av.AudioFrame.from_ndarray(audio, format="s16p", layout=layout)
|
|
audio_frame.sample_rate = sample_rate
|
|
add_audio_packets(container, audio_stream, {"sr": sample_rate, "layout": layout, "frames": [audio_frame]})
|
|
|
|
|
|
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 = "",
|
|
sample_rate: 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)
|
|
parsed_options = {}
|
|
if isinstance(options, str):
|
|
for option in [opt.strip() for opt in options.split(',')]:
|
|
if '=' in option:
|
|
key, value = option.split('=', 1)
|
|
elif ':' in option:
|
|
key, value = option.split(':', 1)
|
|
else:
|
|
continue
|
|
parsed_options[key.strip()] = value.strip()
|
|
elif isinstance(options, dict):
|
|
parsed_options = options
|
|
log.info(f'Video: file="{filename}" codec={codec} frames={frames} width={width} height={height} fps={rate} audio={audio is not None} sample_rate={sample_rate} options={parsed_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=parsed_options)
|
|
stream.width = video_array.shape[2]
|
|
stream.height = video_array.shape[1]
|
|
stream.pix_fmt = pix_fmt
|
|
stream.time_base = Fraction(1, rate)
|
|
audio_stream = None
|
|
has_audio = (audio is not None) and ((torch.is_tensor(audio) or isinstance(audio, np.ndarray)) or (isinstance(audio, dict) and len(audio.get('frames', [])) > 0))
|
|
if has_audio:
|
|
sr = sample_rate if not isinstance(audio, dict) else audio.get("sr", sample_rate)
|
|
layout = "stereo" if not isinstance(audio, dict) else audio.get("layout", "stereo")
|
|
audio_stream = container.add_stream("aac", rate=sr)
|
|
audio_stream.layout = layout
|
|
audio_stream.time_base = Fraction(1, sr)
|
|
for i, img in enumerate(video_array):
|
|
frame = av.VideoFrame.from_ndarray(img, format="rgb24")
|
|
frame.pts = i
|
|
for packet in stream.encode_lazy(frame):
|
|
container.mux_one(packet)
|
|
if task is not None:
|
|
pbar.update(task, advance=1)
|
|
if (audio is not None) and (torch.is_tensor(audio) or isinstance(audio, np.ndarray)):
|
|
add_audio_tensor(container, audio_stream, audio, sample_rate)
|
|
elif (audio is not None) and isinstance(audio, dict) and len(audio.get('frames', [])) > 0:
|
|
add_audio_packets(container, audio_stream, audio)
|
|
for packet in stream.encode():
|
|
container.mux_one(packet)
|
|
|
|
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
|
|
reclamp: bool = True, # reclamp pixels to [-1, 1] range
|
|
):
|
|
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()
|
|
if pixels.ndim == 4:
|
|
pixels = pixels.unsqueeze(0)
|
|
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}')
|
|
if hasattr(audio, 'shape'):
|
|
audio_txt = f'audio={audio.shape} aac={aac_sample_rate}' if audio is not None else 'no audio'
|
|
elif isinstance(audio, dict):
|
|
audio_txt = f'audio={audio.get("format", None)} packets={len(audio.get("frames", []))} '
|
|
else:
|
|
audio_txt = None
|
|
log.debug(f'Video: encode={t} raw={size} latent={pixels.shape} {audio_txt} 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
|
|
|
|
if reclamp:
|
|
x = torch.clamp(pixels.float(), -1., 1.) * 127.5 + 127.5
|
|
else:
|
|
x = pixels.float() * 255.0
|
|
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, sample_rate=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
|