diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index dab545d81..ca095508c 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -435,14 +435,7 @@ class VideoGenerator: legacy_kwargs=kwargs, ) - fastvideo_args = self.fastvideo_args - pipeline_overrides = request_to_pipeline_overrides(request) - if pipeline_overrides: - fastvideo_args = deepcopy(self.fastvideo_args) - for key, value in pipeline_overrides.items(): - if not hasattr(fastvideo_args.pipeline_config, key): - raise ValueError(f"Request field {key!r} is not supported by pipeline config overrides") - setattr(fastvideo_args.pipeline_config, key, deepcopy(value)) + fastvideo_args = self._resolve_request_fastvideo_args(request) resolved_sampling_param = request_to_sampling_param( request, @@ -544,18 +537,23 @@ class VideoGenerator: return self._generate_single_request(request) + def _resolve_request_fastvideo_args(self, request: GenerationRequest) -> FastVideoArgs: + """Return ``self.fastvideo_args``, deep-copied with the request's pipeline overrides applied (if any).""" + pipeline_overrides = request_to_pipeline_overrides(request) + if not pipeline_overrides: + return self.fastvideo_args + fastvideo_args = deepcopy(self.fastvideo_args) + for key, value in pipeline_overrides.items(): + if not hasattr(fastvideo_args.pipeline_config, key): + raise ValueError(f"Request field {key!r} is not supported by pipeline config overrides") + setattr(fastvideo_args.pipeline_config, key, deepcopy(value)) + return fastvideo_args + def _generate_single_request( self, request: GenerationRequest, ) -> GenerationResult | list[GenerationResult]: - fastvideo_args = self.fastvideo_args - pipeline_overrides = request_to_pipeline_overrides(request) - if pipeline_overrides: - fastvideo_args = deepcopy(self.fastvideo_args) - for key, value in pipeline_overrides.items(): - if not hasattr(fastvideo_args.pipeline_config, key): - raise ValueError(f"Request field {key!r} is not supported by pipeline config overrides") - setattr(fastvideo_args.pipeline_config, key, deepcopy(value)) + fastvideo_args = self._resolve_request_fastvideo_args(request) sampling_param = request_to_sampling_param( request, @@ -832,6 +830,10 @@ class VideoGenerator: n_tokens=n_tokens, VSA_sparsity=fastvideo_args.VSA_sparsity, ) + # Allow precomputed prompt_embeds (e.g. from diffusers) to skip text encoding + prompt_embeds = kwargs.get("prompt_embeds") + if prompt_embeds is not None: + batch.prompt_embeds = (list(prompt_embeds) if isinstance(prompt_embeds, list | tuple) else [prompt_embeds]) extra_overrides = kwargs.get("_extra_overrides", {}) for _ek, _ev in extra_overrides.items(): @@ -1184,257 +1186,16 @@ class VideoGenerator: if fastvideo_args is None: fastvideo_args = self.fastvideo_args - # Validate inputs if not isinstance(prompt, str): raise TypeError(f"`prompt` must be a string, but got {type(prompt)}") - prompt = prompt.strip() - sampling_param = deepcopy(sampling_param) - output_path = kwargs["output_path"] - prompt_embeds = kwargs.get("prompt_embeds") - sampling_param.prompt = prompt - # Process negative prompt - if sampling_param.negative_prompt is not None: - sampling_param.negative_prompt = sampling_param.negative_prompt.strip() - # Validate dimensions - if (sampling_param.height <= 0 or sampling_param.width <= 0 or sampling_param.num_frames <= 0): - raise ValueError(f"Height, width, and num_frames must be positive integers, got " - f"height={sampling_param.height}, width={sampling_param.width}, " - f"num_frames={sampling_param.num_frames}") - - # Calculate sizes - target_height = align_to(sampling_param.height, 16) - target_width = align_to(sampling_param.width, 16) - - # Calculate latent sizes - latents_size = [(sampling_param.num_frames - 1) // 4 + 1, sampling_param.height // 8, sampling_param.width // 8] - n_tokens = latents_size[0] * latents_size[1] * latents_size[2] - - # Log parameters - debug_str = f""" - height: {target_height} - width: {target_width} - video_length: {sampling_param.num_frames} - prompt: {sampling_param.prompt} - image_path: {sampling_param.image_path} - neg_prompt: {sampling_param.negative_prompt} - seed: {sampling_param.seed} - infer_steps: {sampling_param.num_inference_steps} - num_videos_per_prompt: {sampling_param.num_videos_per_prompt} - guidance_scale: {sampling_param.guidance_scale} - n_tokens: {n_tokens} - flow_shift: {fastvideo_args.pipeline_config.flow_shift} - embedded_guidance_scale: {fastvideo_args.pipeline_config.embedded_cfg_scale} - save_video: {sampling_param.save_video} - output_path: {output_path} - """ # type: ignore[attr-defined] - logger.info(debug_str) - - # Prepare batch - batch = ForwardBatch( - **shallow_asdict(sampling_param), - eta=0.0, - n_tokens=n_tokens, - VSA_sparsity=fastvideo_args.VSA_sparsity, + work_item = self._prepare_generation_work_item( + prompt=prompt, + sampling_param=sampling_param, + fastvideo_args=fastvideo_args, + **kwargs, ) - # Allow precomputed prompt_embeds (e.g. from diffusers) to skip text encoding - if prompt_embeds is not None: - batch.prompt_embeds = (list(prompt_embeds) if isinstance(prompt_embeds, list | tuple) else [prompt_embeds]) - - extra_overrides = kwargs.pop("_extra_overrides", {}) - for _ek, _ev in extra_overrides.items(): - batch.extra[_ek] = _ev - - # Run inference - start_time = time.perf_counter() - - # Execute forward pass in a new thread for non-blocking tensor - # allocation. Capture thread exceptions so we can surface the true - # failure in the main thread instead of later hitting None outputs. - result_container = {"output_batch": ForwardBatch(data_type=batch.data_type)} - thread_error: dict[str, BaseException | None] = {"error": None} - thread_error_traceback: dict[str, str] = {"traceback": ""} - - def execute_forward_thread(): - import traceback - try: - result_container["output_batch"] = self.executor.execute_forward(batch, fastvideo_args) - except BaseException as error: # noqa: BLE001 - thread_error["error"] = error - thread_error_traceback["traceback"] = traceback.format_exc() - - thread = threading.Thread(target=execute_forward_thread) - thread.start() - latent_batch_size = _infer_latent_batch_size(batch) - # When ``output_type == "latent"`` the forward output has latent - # shape (e.g. ``[B, C_latent, T_latent, H_latent, W_latent]``) - # rather than the pre-allocation's pixel shape. Skip the pinned - # ~50 MB buffer entirely; we always fall through to the - # ``samples = output_batch.output.cpu()`` branch below in that - # mode. ``skip_pixel_prealloc`` also gates the slow-path warning. - skip_pixel_prealloc = fastvideo_args.output_type == "latent" - if skip_pixel_prealloc: - samples = torch.empty(0, device='cpu') - else: - samples = torch.empty( - (latent_batch_size, 3, sampling_param.num_frames, sampling_param.height, sampling_param.width), - device='cpu', - pin_memory=fastvideo_args.pin_cpu_memory) - thread.join() - - if thread_error["error"] is not None: - raise RuntimeError("Forward execution thread failed.\n" - f"{thread_error_traceback['traceback']}") from thread_error["error"] - - output_batch = result_container["output_batch"] - if output_batch.output is None: - raise RuntimeError("Forward execution returned no output tensor. " - "This usually means the executor/pipeline failed earlier.") - - if output_batch.output.shape == samples.shape: - samples.copy_(output_batch.output) - else: - if not skip_pixel_prealloc: - logger.warning("Output shape %s does not match expected shape %s; use slow path", - output_batch.output.shape, samples.shape) - samples = output_batch.output.cpu() - logging_info = output_batch.logging_info - - gen_time = time.perf_counter() - start_time - logger.info("Generated successfully in %.2f seconds", gen_time) - - # Three mutually-exclusive output modes determine whether (a) we - # build an RGB frame buffer and (b) what file we write to disk: - # - # 1. `output_type == "latent"` — VAE is bypassed in DecodingStage - # and `samples` holds raw latents (arbitrary channel count). - # The RGB grid / uint8 / mp4 / png pipeline below cannot - # consume those, so we skip it entirely and let callers work - # with the latent tensor directly via `result["samples"]`. - # 2. Audio-only workload — `samples` is a 1×3×1×8×8 placeholder - # no caller will use; skip the grid loop and save a `.wav`. - # 3. Pixel video / image — the historical happy path. - is_latent_output = fastvideo_args.output_type == "latent" - audio_only = bool(output_batch.extra.get("audio_only")) - - postprocess_start = time.perf_counter() - frames: list[np.ndarray] | None - if is_latent_output or audio_only: - frames = None if is_latent_output else [] - else: - videos = rearrange(samples, "b c t h w -> t b c h w") - frames = [] - for x in videos: - x = torchvision.utils.make_grid(x, nrow=6) - x = x.permute(1, 2, 0).squeeze(-1) - x = (x * 255).to(torch.uint8) - frames.append(x.contiguous().cpu().numpy()) - postprocess_time = time.perf_counter() - postprocess_start - logger.info("PostDecodeFrameProcessStage completed in %.3f s", postprocess_time) - if logging_info is not None: - logging_info.add_stage_execution_time("PostDecodeFrameProcessStage", postprocess_time) - - save_to_disk = batch.save_video and not is_latent_output - save_video_time = 0.0 - audio_mux_time = 0.0 - if save_to_disk: - if audio_only: - # Audio-only workload: write a standalone .wav rather than - # muxing the audio into a placeholder mp4 (which forces - # ffmpeg to round 8x8 placeholder frames up to 16x16). - output_path = self._rewrite_extension(output_path, ".wav") - save_start = time.perf_counter() - self._write_pcm_wav( - output_path, - output_batch.extra["audio"], - int(output_batch.extra["audio_sample_rate"]), - ) - save_video_time = time.perf_counter() - save_start - logger.info("Saved audio to %s", output_path) - elif self._is_image_workload(): - # Image workloads (t2i, i2i, …): save the first frame as PNG. - assert frames is not None # implied by save_to_disk and not audio_only - save_start = time.perf_counter() - imageio.imwrite(output_path, frames[0]) - save_video_time = time.perf_counter() - save_start - logger.info("Saved image to %s", output_path) - else: - assert frames is not None # implied by save_to_disk and not audio_only - audio = output_batch.extra.get("audio") - audio_sample_rate = output_batch.extra.get("audio_sample_rate") - if audio is not None and audio_sample_rate is not None: - # Single-pass save path: encode video+audio once, avoiding - # second-pass remux overhead. AudioMuxStage remains 0.0 on - # success because audio is already present in the saved MP4. - save_start = time.perf_counter() - save_ok = self._save_video_with_audio_ffmpeg_pipe( - output_path=output_path, - frames=frames, - fps=batch.fps, - audio=audio, - sample_rate=int(audio_sample_rate), - ) - if not save_ok: - logger.warning("ffmpeg pipe save failed; trying PyAV single-pass save.") - save_ok = self._save_video_with_audio_single_pass( - output_path=output_path, - frames=frames, - fps=batch.fps, - audio=audio, - sample_rate=int(audio_sample_rate), - ) - save_video_time = time.perf_counter() - save_start - if save_ok: - audio_mux_time = 0.0 - else: - logger.warning("Single-pass save failed; falling back to two-step save/mux.") - save_start = time.perf_counter() - imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4") - save_video_time = time.perf_counter() - save_start - mux_start = time.perf_counter() - mux_ok = self._mux_audio(output_path, audio, int(audio_sample_rate)) - audio_mux_time = time.perf_counter() - mux_start - if not mux_ok: - logger.warning("Audio mux failed; saved video without audio.") - else: - save_start = time.perf_counter() - imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4") - save_video_time = time.perf_counter() - save_start - audio_mux_time = 0.0 - logger.info("Saved video to %s", output_path) - - logger.info("VideoSaveStage completed in %.3f s", save_video_time) - if logging_info is not None: - logging_info.add_stage_execution_time("VideoSaveStage", save_video_time) - logger.info("AudioMuxStage completed in %.3f s", audio_mux_time) - if logging_info is not None: - logging_info.add_stage_execution_time("AudioMuxStage", audio_mux_time) - - e2e_time = time.perf_counter() - start_time - logger.info("End-to-end latency: %.2f seconds", e2e_time) - - result: dict[str, Any] = { - "prompts": prompt, - "samples": samples if batch.return_frames else None, - "frames": frames if batch.return_frames else None, - # Audio is the primary output for audio workloads — return it - # whenever the pipeline produced one, regardless of - # `return_frames` (which gates the video-shaped buffers). - "audio": output_batch.extra.get("audio"), - "audio_sample_rate": output_batch.extra.get("audio_sample_rate"), - "ltx2_audio_latents": output_batch.extra.get("ltx2_audio_latents"), - "size": (target_height, target_width, batch.num_frames), - "generation_time": gen_time, - "e2e_latency": e2e_time, - "logging_info": logging_info, - "trajectory": output_batch.trajectory_latents, - "trajectory_timesteps": output_batch.trajectory_timesteps, - "trajectory_decoded": output_batch.trajectory_decoded, - "video_path": output_path if save_to_disk else None, - "peak_memory_mb": output_batch.extra.get("peak_memory_mb"), - } - - return result + return self._execute_single_work_item(work_item) @staticmethod def _wrap_legacy_result(