diff --git a/CHANGELOG.md b/CHANGELOG.md index 9f2dca43d..0128ba79e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,13 @@ # Change Log for SD.Next +## Update for 2026-07-27 + +- **Features** + - optimized server startup + - video processing preserve audio +- **Fixes** + - seedvr quality + ## Update for 2026-07-23 Primarily a service release with updates to compute packages: torch, CUDA, ROCm, etc. diff --git a/modules/postprocess/seedvr_model.py b/modules/postprocess/seedvr_model.py index 5da179635..27f8e775b 100644 --- a/modules/postprocess/seedvr_model.py +++ b/modules/postprocess/seedvr_model.py @@ -170,7 +170,7 @@ class UpscalerSeedVR(Upscaler): devices.torch_gc(fast=True) t0 = time.time() with devices.inference_context(): - self.pbar.update(self.task, description=f'inference: batch={self.step}') + self.pbar.update(self.task, description='inference') result = generation.generation_step_original(*args, **kwargs) self.pbar.update(self.task, advance=self.step) self.timer.ts('step', t0) @@ -187,13 +187,38 @@ class UpscalerSeedVR(Upscaler): image = Image.open(image) image = image.convert("RGB") width = image.width + height = image.height tensor = np.array(image) tensor = torch.from_numpy(tensor).to(device=devices.device, dtype=devices.dtype).unsqueeze(0) / 255.0 self.frames = 1 - return tensor, width + return tensor, width, height except Exception as e: log.error(f'Upscaler: name="SeedVR2" image="{image}" {e}') - return None, None + return None, None, None + + def read_audio(self, video_path: str): + audio_frames = [] + audio_meta = None + try: + from modules.video_models.video_utils import check_av + av = check_av() + container = av.open(video_path) + if container.streams.audio: + audio_stream = container.streams.audio[0] + audio_meta = { + "sr": audio_stream.codec_context.sample_rate, + "channels": audio_stream.codec_context.channels, + "layout": audio_stream.layout.name if audio_stream.layout else "stereo", + "format": audio_stream.codec_context.format.name, + } + for frame in container.decode(audio_stream): + audio_frames.append(frame) + container.close() + except Exception as e: + log.error(f'Upscaler: name="SeedVR2" video="{video_path}" {e}') + if audio_meta and len(audio_frames) > 0: + return {"frames": audio_frames, **audio_meta} + return None def read_video(self, video_path: str): try: @@ -201,9 +226,10 @@ class UpscalerSeedVR(Upscaler): cap = cv2.VideoCapture(video_path) if not cap.isOpened(): log.error(f'Upscaler: name="SeedVR2" video="{video_path}" failed to open') - return None, None + return None, None, None frames = [] width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) self.fps = int(cap.get(cv2.CAP_PROP_FPS)) while True: ret, frame = cap.read() @@ -214,20 +240,21 @@ class UpscalerSeedVR(Upscaler): cap.release() if len(frames) == 0: log.error(f'Upscaler: name="SeedVR2" video="{video_path}" no frames read') - return None, None + return None, None, None tensor = torch.from_numpy(np.array(frames)).to(device=devices.device, dtype=devices.dtype) / 255.0 self.frames = tensor.shape[0] - return tensor, width + return tensor, width, height except Exception as e: log.error(f'Upscaler: name="SeedVR2" video="{video_path}" {e}') - return None, None + return None, None, None - def create_video(self, tensor: torch.Tensor, codec: str = 'libx264', codec_opt: str = 'crf:16', interpolate: int = 0): + def create_video(self, tensor: torch.Tensor, audio, codec: str = 'libx264', codec_opt: str = 'crf:16', interpolate: int = 0): t0 = time.time() from modules.video_models.video_save import save_video pixels = tensor.permute(3, 0, 1, 2).unsqueeze(0) # from (t, h, w, c) to (n, c, t, h, w) _frames, filename, _thumb = save_video(p=None, pixels=pixels, + audio=audio, mp4_fps=self.fps, mp4_thumb=False, mp4_frames=False, @@ -255,7 +282,7 @@ class UpscalerSeedVR(Upscaler): interpolate: int = 1, codec: str = 'libx264', codec_opt: str = 'crf:16', - vae_memory: float = 0.2, + vae_memory: float = 0.5, vae_tile_encode: bool = True, vae_tile_decode: bool = True, ): @@ -273,11 +300,13 @@ class UpscalerSeedVR(Upscaler): from modules.seedvr.src.core import generation + audio = None self.scale = self.scale if scale is None else scale if isinstance(img, Image.Image): - tensor, width = self.read_image(img) + tensor, width, height = self.read_image(img) elif isinstance(img, str): - tensor, width = self.read_video(img) + tensor, width, height = self.read_video(img) + audio = self.read_audio(img) else: log.error(f'Upscaler: name="SeedVR2" image="{img}" unsupported type {type(img)}') return img @@ -287,6 +316,7 @@ class UpscalerSeedVR(Upscaler): log.error(f'Upscaler: name="SeedVR2" image="{img}" failed to read') return img width = int(self.scale * width) // 8 * 8 + height = int(self.scale * height) // 8 * 8 random.seed() seed = int(random.randrange(4294967294)) if seed == -1 else int(seed) self.step = 1 if self.frames == 1 else batch_size - batch_overlap @@ -294,7 +324,7 @@ class UpscalerSeedVR(Upscaler): mode = "mode=image" if self.frames == 1 else f"mode=video frames={self.frames}" batch_info = f'batch=(size={batch_size} overlap={batch_overlap})' vae_info = f'vae=(tiled={vae_tile_encode}/{vae_tile_decode} memory={vae_memory} size={tile_size} overlap={tile_overlap})' - log.info(f'Upscaler: type="{self.name}" model="{selected_file}" {mode} scale={self.scale} cfg={cfg_scale}:{cfg_rescale} seed={seed} steps={steps} offload={self.offload} {batch_info} {vae_info}') + log.info(f'Upscaler: type="{self.name}" model="{selected_file}" {mode} scale={self.scale} width={width} height={height} cfg={cfg_scale}:{cfg_rescale} seed={seed} steps={steps} offload={self.offload} {batch_info} {vae_info}') import rich.progress as rp self.pbar = rp.Progress(rp.TextColumn('[cyan]SeedVR:'), rp.BarColumn(), rp.MofNCompleteColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=console) @@ -342,7 +372,7 @@ class UpscalerSeedVR(Upscaler): if self.frames == 1: result = convert.to_pil(result_tensor.squeeze()) elif self.frames > 1: - result = self.create_video(result_tensor, codec=codec, codec_opt=codec_opt, interpolate=interpolate) + result = self.create_video(result_tensor, audio, codec=codec, codec_opt=codec_opt, interpolate=interpolate) else: log.error(f'Upscaler: name="SeedVR2" model="{selected_file}" no frames generated') result = img diff --git a/modules/seedvr/src/common/decorators.py b/modules/seedvr/src/common/decorators.py deleted file mode 100644 index 52ab8ac03..000000000 --- a/modules/seedvr/src/common/decorators.py +++ /dev/null @@ -1,126 +0,0 @@ -import functools -import threading -from typing import Callable -import torch -from .distributed import barrier_if_distributed, get_global_rank, get_local_rank -from .logger import get_logger - - -logger = get_logger(__name__) - - -def log_on_entry(func: Callable) -> Callable: - """ - Functions with this decorator will log the function name at entry. - When using multiple decorators, this must be applied innermost to properly capture the name. - """ - - def log_on_entry_wrapper(*args, **kwargs): - logger.info(f"Entering {func.__name__}") - return func(*args, **kwargs) - - return log_on_entry_wrapper - - -def barrier_on_entry(func: Callable) -> Callable: - """ - Functions with this decorator will start executing when all ranks are ready to enter. - """ - - def barrier_on_entry_wrapper(*args, **kwargs): - barrier_if_distributed() - return func(*args, **kwargs) - - return barrier_on_entry_wrapper - - -def _conditional_execute_wrapper_factory(execute: bool, func: Callable) -> Callable: - """ - Helper function for local_rank_zero_only and global_rank_zero_only. - """ - - def conditional_execute_wrapper(*args, **kwargs): - # Only execute if needed. - result = func(*args, **kwargs) if execute else None - # All GPUs must wait. - barrier_if_distributed() - # Return results. - return result - - return conditional_execute_wrapper - - -def _asserted_wrapper_factory(condition: bool, func: Callable, err_msg: str = "") -> Callable: - """ - Helper function for some functions with special constraints, - especially functions called by other global_rank_zero_only / local_rank_zero_only ones, - in case they are wrongly invoked in other scenarios. - """ - - def asserted_execute_wrapper(*args, **kwargs): - assert condition, err_msg - result = func(*args, **kwargs) - return result - - return asserted_execute_wrapper - - -def local_rank_zero_only(func: Callable) -> Callable: - """ - Functions with this decorator will only execute on local rank zero. - """ - return _conditional_execute_wrapper_factory(get_local_rank() == 0, func) - - -def global_rank_zero_only(func: Callable) -> Callable: - """ - Functions with this decorator will only execute on global rank zero. - """ - return _conditional_execute_wrapper_factory(get_global_rank() == 0, func) - - -def assert_only_global_rank_zero(func: Callable) -> Callable: - """ - Functions with this decorator are only accessible to processes with global rank zero. - """ - return _asserted_wrapper_factory( - get_global_rank() == 0, func, err_msg="Not accessible to processes with global_rank != 0" - ) - - -def assert_only_local_rank_zero(func: Callable) -> Callable: - """ - Functions with this decorator are only accessible to processes with local rank zero. - """ - return _asserted_wrapper_factory( - get_local_rank() == 0, func, err_msg="Not accessible to processes with local_rank != 0" - ) - - -def new_thread(func: Callable) -> Callable: - """ - Functions with this decorator will run in a new thread. - The function will return the thread, which can be joined to wait for completion. - """ - - def new_thread_wrapper(*args, **kwargs): - thread = threading.Thread(target=func, args=args, kwargs=kwargs) - thread.start() - return thread - - return new_thread_wrapper - - -def log_runtime(func: Callable) -> Callable: - """ - Functions with this decorator will logging the runtime. - """ - - @functools.wraps(func) - def wrapped(*args, **kwargs): - barrier_if_distributed() - result = func(*args, **kwargs) - barrier_if_distributed() - return result - - return wrapped diff --git a/modules/seedvr/src/common/seed.py b/modules/seedvr/src/common/seed.py index 2469ad944..1129ebd5f 100644 --- a/modules/seedvr/src/common/seed.py +++ b/modules/seedvr/src/common/seed.py @@ -16,13 +16,14 @@ import random from typing import Optional import numpy as np import torch -from .distributed import get_global_rank -def set_seed(seed: Optional[int], same_across_ranks: bool = False): +def set_seed(seed: Optional[int]): """Function that sets the seed for pseudo-random number generators.""" + if (seed is None) or (seed == '') or (seed == -1): + random.seed() + seed = int(random.randrange(4294967294)) if seed is not None: - seed += get_global_rank() if not same_across_ranks else 0 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) diff --git a/modules/seedvr/src/core/generation.py b/modules/seedvr/src/core/generation.py index 0b1d42f0f..811bdac88 100644 --- a/modules/seedvr/src/core/generation.py +++ b/modules/seedvr/src/core/generation.py @@ -107,7 +107,19 @@ def cut_videos(videos): return result -def generation_loop(runner, images, cfg_scale=1.0, cfg_rescale=0.0, steps=1, seed=666, res_w=720, batch_size=90, temporal_overlap=0, progress_callback=None, device:str='cpu', color_reconstruct=True): +def generation_loop(runner, + images, + cfg_scale=1.5, + cfg_rescale=0.0, + steps=1, + seed=-1, + res_w=720, + batch_size=1, + temporal_overlap=0, + progress_callback=None, + device:str='cpu', + color_reconstruct=True, + ): """ Main generation loop with context-aware temporal processing diff --git a/modules/upscaler.py b/modules/upscaler.py index 6bafa3db0..b02633f58 100644 --- a/modules/upscaler.py +++ b/modules/upscaler.py @@ -118,7 +118,7 @@ class Upscaler: img = self.do_upscale(img, selected_model) if shape == (img.width, img.height): break - if img.width >= dest_w and img.height >= dest_h: + if img.width >= (dest_w - 8) and img.height >= (dest_h - 8): break if img.width != dest_w or img.height != dest_h: from modules.image import sharpfin diff --git a/modules/video_models/video_save.py b/modules/video_models/video_save.py index 49166e237..2ce0fc6ad 100644 --- a/modules/video_models/video_save.py +++ b/modules/video_models/video_save.py @@ -103,55 +103,62 @@ def numpy_to_tensor(images): 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 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 write_audio( - container, - audio_stream, - samples: torch.Tensor, - audio_sample_rate: int, -) -> None: +def add_audio_tensor(container, audio_stream, audio: torch.Tensor, sample_rate: int): 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) + if torch.is_tensor(audio): + audio = audio.detach().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(np.ascontiguousarray(audio.T), format="s16", 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( @@ -162,7 +169,7 @@ def atomic_save_video( codec: str = "libx264", pix_fmt: str = "yuv420p", options: str = "", - aac: int = 24000, + sample_rate: int = 24000, metadata: dict | None = None, pbar=None, ): @@ -172,23 +179,23 @@ def atomic_save_video( 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}') + 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') @@ -196,25 +203,32 @@ def atomic_save_video( 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: 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 - audio_stream = add_audio_stream(container, aac) if audio is not None else None - for img in video_array: + 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(packet) + container.mux_one(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') + 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) @@ -294,7 +308,13 @@ def save_video( 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}"') + 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: @@ -340,7 +360,7 @@ def save_video( 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) + 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)) diff --git a/scripts/daam/utils.py b/scripts/daam/utils.py index 58402c9ed..410de9b30 100644 --- a/scripts/daam/utils.py +++ b/scripts/daam/utils.py @@ -94,8 +94,8 @@ nlp = None @lru_cache(maxsize=100000) -def cached_nlp(prompt: str, type='en_core_web_md'): - global nlp +def cached_nlp(prompt: str, type='en_core_web_md'): # pylint: disable=redefined-builtin + global nlp # pylint: disable=global-statement if nlp is None: try: diff --git a/scripts/postprocessing_upscale.py b/scripts/postprocessing_upscale.py index 92e6da29a..9726b23b3 100644 --- a/scripts/postprocessing_upscale.py +++ b/scripts/postprocessing_upscale.py @@ -68,7 +68,6 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing): return image def process(self, pp: scripts_postprocessing.PostprocessedImage, upscale_mode=1, upscale_by=2.0, upscale_to_width=None, upscale_to_height=None, upscale_crop=False, upscaler_1_name=None, upscaler_2_name=None, upscaler_2_visibility=0.0): # pylint: disable=arguments-differ - if upscaler_1_name == "None": upscaler_1_name = None upscaler1 = next(iter([x for x in shared.sd_upscalers if x.name == upscaler_1_name]), None)