Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| import os | |
| import tempfile | |
| from pathlib import Path | |
| from typing import Callable | |
| from inference.config import TEMPORAL_DOWNSAMPLE, InferenceConfig | |
| from inference.loading import ComponentFactories, ModelBundle, ModelPaths, build_model_bundle | |
| from inference.output import ( | |
| IncrementalVideoWriter, | |
| mux_source_audio, | |
| save_tensor_video, | |
| ) | |
| from inference.preprocessing import ( | |
| expand_for_temporal_compression, | |
| load_lq_video, | |
| iter_lq_video, | |
| get_lq_video_fps, | |
| TemporalVideoStream, | |
| merge_temporal_output, | |
| ) | |
| from inference.postprocessing import WaveletColorCorrector | |
| from inference.request import InferenceRequest, InferenceResult | |
| from inference.runners import StreamRunner | |
| from inference.randomness import set_inference_seed | |
| class FlashVSRInferencePipeline: | |
| def __init__(self, config: InferenceConfig, bundle: ModelBundle): | |
| config.validate() | |
| self.config = config | |
| self.bundle = bundle | |
| color_config = config.color_correction | |
| self.color_corrector = ( | |
| WaveletColorCorrector( | |
| levels=color_config.levels, | |
| chunk_size=color_config.chunk_size, | |
| ) | |
| if color_config.enabled | |
| else None | |
| ) | |
| self.runner = StreamRunner(config, bundle) | |
| def infer( | |
| self, | |
| request: InferenceRequest | None = None, | |
| *, | |
| progress: Callable[[str], None] | None = None, | |
| video_preset: str = "veryfast", | |
| **kwargs, | |
| ) -> InferenceResult: | |
| request = request or InferenceRequest(**kwargs) | |
| request.validate() | |
| set_inference_seed(self.config.seed) | |
| input_stream = None | |
| expansion = None | |
| lq_video = None | |
| if self.config.decoder.decode_every != "auto": | |
| fps = get_lq_video_fps(request) | |
| chunk_frames = self.config.decoder.decode_every * self.config.temporal_compression | |
| input_stream = TemporalVideoStream( | |
| iter_lq_video( | |
| request, chunk_frames=chunk_frames, | |
| first_chunk_frames=chunk_frames + 1, | |
| ), | |
| self.config.temporal_compression, | |
| TEMPORAL_DOWNSAMPLE, | |
| ) | |
| model_input = input_stream | |
| else: | |
| lq_video, fps = load_lq_video(request, stream=True) | |
| expansion = expand_for_temporal_compression( | |
| lq_video, self.config.temporal_compression, | |
| model_compression=TEMPORAL_DOWNSAMPLE, | |
| ) | |
| model_input = expansion.video | |
| destination = Path(request.output_path) if request.output_path is not None else None | |
| audio_source = ( | |
| Path(request.input_path) | |
| if destination is not None and request.input_path is not None | |
| else None | |
| ) | |
| encoded_path = destination | |
| temporary_video = None | |
| if destination is not None and audio_source is not None: | |
| destination.parent.mkdir(parents=True, exist_ok=True) | |
| fd, temporary_name = tempfile.mkstemp( | |
| prefix=f".{destination.stem}.", | |
| suffix=destination.suffix or ".mp4", | |
| dir=str(destination.parent), | |
| ) | |
| os.close(fd) | |
| temporary_video = Path(temporary_name) | |
| temporary_video.unlink() | |
| encoded_path = temporary_video | |
| if progress is not None: | |
| progress("Iterative input ready" if input_stream is not None else "Input preprocessing complete") | |
| incremental = ( | |
| request.output_path is not None | |
| and self.config.decoder.decode_every != "auto" | |
| and self.config.temporal_compression in (1, 2, 4) | |
| ) | |
| writer = None | |
| if incremental: | |
| writer = IncrementalVideoWriter( | |
| encoded_path, fps, preset=video_preset | |
| ) | |
| try: | |
| output = self.runner.run( | |
| model_input, request.context, progress=progress, | |
| segment_writer=writer, | |
| output_frame_limit=(expansion.source_frames if writer is not None and expansion is not None else None), | |
| ) | |
| except Exception: | |
| if writer is not None: | |
| writer.close() | |
| if temporary_video is not None: | |
| temporary_video.unlink(missing_ok=True) | |
| raise | |
| finally: | |
| if input_stream is not None: | |
| input_stream.close() | |
| timeline = input_stream if input_stream is not None else expansion | |
| if writer is not None: | |
| writer.close() | |
| artifact = writer.verify() | |
| video = None | |
| else: | |
| if self.config.decoder.decode_every == "auto": | |
| video = merge_temporal_output( | |
| output.video, | |
| expansion, | |
| mode=self.config.temporal_merge, | |
| gaussian_sigma=self.config.temporal_gaussian_sigma, | |
| ) | |
| else: | |
| # Numeric decode intervals already merge each segment onto | |
| # the source timeline inside StreamRunner. | |
| video = output.video | |
| # The model input may contain repeated tail frames used only to | |
| # satisfy its temporal grid. Never expose those synthetic frames. | |
| video = video[:, :, : timeline.source_frames] | |
| color_config = self.config.color_correction | |
| if video is not None and self.color_corrector is not None and self.config.decoder.decode_every == "auto": | |
| reference = lq_video[:, :, : video.shape[2]] | |
| video = self.color_corrector(video, reference) | |
| if progress is not None: | |
| progress("Color correction complete") | |
| elif self.color_corrector is not None and progress is not None: | |
| progress("Color correction complete (per decode segment)") | |
| output_path = None | |
| artifact = locals().get("artifact", None) | |
| if request.output_path is not None and video is not None: | |
| artifact = save_tensor_video( | |
| video[0].permute(1, 0, 2, 3), | |
| encoded_path, | |
| fps, | |
| preset=video_preset, | |
| ) | |
| if progress is not None: | |
| progress("Video saved") | |
| if destination is not None and audio_source is not None and encoded_path is not None: | |
| mux_source_audio(encoded_path, audio_source, destination) | |
| if temporary_video is not None: | |
| temporary_video.unlink(missing_ok=True) | |
| if artifact is not None: | |
| artifact["path"] = str(destination) | |
| output_path = str(destination) | |
| elif destination is not None and artifact is not None: | |
| output_path = str(destination) | |
| runner_metadata = dict(output.metadata) | |
| if "input_frames" in runner_metadata: | |
| runner_metadata["model_input_frames"] = runner_metadata.pop("input_frames") | |
| if "output_frames" in runner_metadata: | |
| runner_metadata["model_output_frames"] = runner_metadata.pop( | |
| "output_frames" | |
| ) | |
| metadata = { | |
| **runner_metadata, | |
| "stream": True, | |
| "decoder": self.config.decoder.kind, | |
| "tiny_decoder_parallel_mode": ( | |
| self.config.decoder.tiny_decoder_parallel_mode | |
| if self.config.decoder.kind == "tiny" | |
| else None | |
| ), | |
| "tiny_decoder_channels_last": ( | |
| self.config.decoder.tiny_decoder_channels_last | |
| if self.config.decoder.kind == "tiny" | |
| else None | |
| ), | |
| "is_causal": self.config.is_causal, | |
| "is_sparse": self.config.is_sparse, | |
| "is_hybrid_sliding": self.config.is_hybrid_sliding, | |
| "attn_kernel": self.config.attn_kernel, | |
| "block_size": self.config.block_size, | |
| "sparse_local_range": self.config.sparse_local_range, | |
| "sliding_local_range": self.config.sliding_local_range, | |
| "sliding_temporal_block_size": ( | |
| self.config.sliding_temporal_block_size | |
| ), | |
| "sliding_temporal_mode": self.config.sliding_temporal_mode, | |
| "sliding_spatial_mode": self.config.sliding_spatial_mode, | |
| "sliding_cache_temporal_windows": ( | |
| self.config.sliding_cache_temporal_windows | |
| ), | |
| "num_inference_steps": self.config.num_inference_steps, | |
| "sigma_shift": self.config.sigma_shift, | |
| "replacement_layers": self.config.replacement_layers, | |
| "source_frames": timeline.source_frames, | |
| "effective_input_frames": timeline.effective_frames, | |
| "expanded_model_frames": timeline.expanded_frames, | |
| "source_fps": fps, | |
| "output_frames": (video.shape[2] if video is not None else int(artifact.get("frames", timeline.effective_frames) if artifact else timeline.effective_frames)), | |
| "temporal_compression": self.config.temporal_compression, | |
| "temporal_repeat_factor": timeline.repeat_factor, | |
| "temporal_merge": self.config.temporal_merge, | |
| "color_correction": { | |
| "enabled": color_config.enabled, | |
| "method": "wavelet", | |
| "levels": color_config.levels, | |
| "chunk_size": color_config.chunk_size, | |
| }, | |
| "artifact": artifact, | |
| } | |
| result_frames = video.shape[2] if video is not None else int(artifact.get("frames", 0) if artifact else 0) | |
| return InferenceResult(video, output_path or (str(request.output_path) if artifact else None), result_frames, fps, metadata) | |
| def build_inference_pipeline( | |
| config: InferenceConfig, | |
| model_paths: ModelPaths | None = None, | |
| *, | |
| bundle: ModelBundle | None = None, | |
| factories: ComponentFactories | None = None, | |
| ) -> FlashVSRInferencePipeline: | |
| config.validate() | |
| if bundle is None: | |
| if model_paths is None: | |
| raise ValueError("Provide model_paths or a prebuilt ModelBundle") | |
| bundle = build_model_bundle(config, model_paths, factories) | |
| return FlashVSRInferencePipeline(config, bundle.eval()) | |