drizzymedia/StreamDiffusionV2-Realtime

1"""Readable staged video-to-video API for StreamDiffusionV2."""2 3from __future__ import annotations4 5from contextlib import ExitStack6from dataclasses import dataclass7from importlib.resources import as_file, files8from pathlib import Path9from typing import Literal10 11from diffusers.utils import export_to_video as diffusers_export_to_video12import numpy as np13import torch14 15from models.util import set_seed16from streamv2v.inference import (17    SingleGPUInferencePipeline as StreamBatchInferencePipeline,18    compute_noise_scale_and_step,19)20from streamv2v.inference_common import load_mp4_as_tensor, merge_cli_config, normalize_acceleration_flags21from streamv2v.inference_wo_batch import SingleGPUInferencePipeline as StreamNoBatchInferencePipeline22 23 24SingleMode = Literal["single", "single-wo"]25 26 27@dataclass28class VideoChunk:29    """One video chunk prepared for the encode -> denoise -> decode loop."""30 31    frames: torch.Tensor32    start_idx: int33    end_idx: int34    current_start: int35    current_end: int36 37 38@dataclass39class EncodedChunk:40    """Encoded latent chunk plus the schedule metadata needed for denoising."""41 42    noisy_latents: torch.Tensor43    current_start: int44    current_end: int45    noise_scale: float46    current_step: int | None = None47 48 49@dataclass50class DenoisedChunk:51    """Denoised latent chunk ready for VAE decoding."""52 53    denoised_pred: torch.Tensor54    last_frame_only: bool55 56 57def _resolve_default_config_path(resource_stack: ExitStack) -> str:58    resource = files("streamv2v.configs").joinpath("wan_causal_dmd_v2v.yaml")59    return str(resource_stack.enter_context(as_file(resource)))60 61 62def _resolve_device(device: str | torch.device | None) -> torch.device:63    cuda_available = torch.cuda.is_available()64    if device is None:65        return torch.device("cuda" if cuda_available else "cpu")66    resolved = torch.device(device)67    if resolved.type == "cuda" and not cuda_available:68        raise RuntimeError("CUDA is not available in the current Python environment")69    if resolved.type == "cuda" and resolved.index is not None:70        torch.cuda.set_device(resolved.index)71    return resolved72 73 74def _normalize_video_tensor(75    video: str | Path | torch.Tensor,76    *,77    height: int,78    width: int,79    device: torch.device,80) -> torch.Tensor:81    if isinstance(video, (str, Path)):82        tensor = load_mp4_as_tensor(str(video), resize_hw=(height, width)).unsqueeze(0)83    else:84        tensor = video85        if tensor.ndim == 4:86            tensor = tensor.unsqueeze(0)87        if tensor.ndim != 5:88            raise ValueError("video tensor must have shape [B, C, T, H, W] or [C, T, H, W]")89    if tensor.dtype != torch.bfloat16:90        tensor = tensor.to(dtype=torch.bfloat16)91    return tensor.to(device)92 93 94def load_video(video_path: str, *, height: int = 480, width: int = 832) -> torch.Tensor:95    """Load a video file as a normalized tensor with shape [C, T, H, W]."""96    return load_mp4_as_tensor(video_path, resize_hw=(height, width))97 98 99def export_video(video: np.ndarray, output_path: str, *, fps: int = 16) -> str:100    """Write a `[T, H, W, C]` float video array to an mp4 file."""101    output_file = Path(output_path)102    output_file.parent.mkdir(parents=True, exist_ok=True)103    diffusers_export_to_video(video, str(output_file), fps=fps)104    return str(output_file)105 106 107class StreamDiffusionV2Pipeline:108    """Readable staged single-GPU API that mirrors the offline inference flow."""109 110    def __init__(111        self,112        checkpoint_folder: str,113        *,114        mode: SingleMode = "single",115        config_path: str | None = None,116        device: str | torch.device | None = None,117        noise_scale: float = 0.8,118        height: int = 480,119        width: int = 832,120        fps: int = 16,121        step: int = 2,122        seed: int = 0,123        model_type: str = "T2V-1.3B",124        use_taehv: bool = False,125        use_tensorrt: bool = False,126        fast: bool = False,127        profile: bool = False,128    ) -> None:129        if mode not in {"single", "single-wo"}:130            raise ValueError("StreamDiffusionV2Pipeline only supports 'single' and 'single-wo'")131 132        self._resource_stack = ExitStack()133        self.mode = mode134        self.device = _resolve_device(device)135        self.checkpoint_folder = checkpoint_folder136        self.noise_scale = float(noise_scale)137        self.height = int(height)138        self.width = int(width)139        self.fps = int(fps)140        self.seed = int(seed)141        self.step = int(step)142        self.profile = bool(profile)143        self.model_type = model_type144        self.prompt: str | None = None145 146        resolved_config_path = config_path or _resolve_default_config_path(self._resource_stack)147        self.config_path = resolved_config_path148        flags = normalize_acceleration_flags(149            {150                "use_taehv": use_taehv,151                "use_tensorrt": use_tensorrt,152                "fast": fast,153            }154        )155        self.use_taehv = bool(flags["use_taehv"])156        self.use_tensorrt = bool(flags["use_tensorrt"])157        self.fast = bool(flags["fast"])158        config_args = {159            "config_path": resolved_config_path,160            "checkpoint_folder": checkpoint_folder,161            "noise_scale": noise_scale,162            "height": height,163            "width": width,164            "fps": fps,165            "step": step,166            "seed": seed,167            "model_type": model_type,168            "profile": profile,169            "use_taehv": self.use_taehv,170            "use_tensorrt": self.use_tensorrt,171            "fast": self.fast,172            "t2v": False,173            "target_fps": None,174            "fixed_noise_scale": False,175            "num_frames": 81,176        }177        self.config = merge_cli_config(resolved_config_path, config_args)178 179        manager_cls = (180            StreamBatchInferencePipeline if mode == "single" else StreamNoBatchInferencePipeline181        )182        torch.set_grad_enabled(False)183        set_seed(self.seed)184        self.pipeline_manager = manager_cls(self.config, self.device)185        self.pipeline_manager.load_model(checkpoint_folder)186        self.chunk_size = 4 * self.config.num_frame_per_block187        self.num_steps = len(self.pipeline_manager.pipeline.denoising_step_list)188        self._next_chunk_index = 0189 190    def close(self) -> None:191        self._resource_stack.close()192 193    def __enter__(self) -> "StreamDiffusionV2Pipeline":194        return self195 196    def __exit__(self, exc_type, exc_val, exc_tb) -> None:197        self.close()198 199    def enable_acceleration(200        self,201        *,202        use_taehv: bool = False,203        use_tensorrt: bool = False,204        fast: bool = False,205    ) -> "StreamDiffusionV2Pipeline":206        """Rebuild the pipeline with the requested acceleration flags."""207        replacement = StreamDiffusionV2Pipeline(208            checkpoint_folder=self.checkpoint_folder,209            mode=self.mode,210            config_path=self.config_path,211            device=self.device,212            noise_scale=self.noise_scale,213            height=self.height,214            width=self.width,215            fps=self.fps,216            step=self.step,217            seed=self.seed,218            model_type=self.model_type,219            use_taehv=use_taehv,220            use_tensorrt=use_tensorrt,221            fast=fast,222            profile=self.profile,223        )224        self.close()225        self.__dict__.update(replacement.__dict__)226        return self227 228    def prepare(self, prompt: str) -> None:229        """Reset the stream state and store the prompt for the next denoising pass."""230        self.prompt = prompt231        self.pipeline_manager.reset_stream_state(reset_vae_flags=True)232        self.pipeline_manager.processed = 0233        self._next_chunk_index = 0234 235    def chunk_video(self, video: str | Path | torch.Tensor) -> list[VideoChunk]:236        """Split a full input video into the same chunks used by the offline inference loop."""237        input_video = _normalize_video_tensor(238            video,239            height=self.height,240            width=self.width,241            device=self.device,242        )243        _, _, total_frames, _, _ = input_video.shape244        if total_frames < 1 + self.chunk_size:245            raise ValueError(f"video must contain at least {1 + self.chunk_size} frames")246 247        chunks: list[VideoChunk] = []248        start_idx = 0249        end_idx = 1 + self.chunk_size250        current_start = 0251        current_end = self.pipeline_manager.pipeline.frame_seq_length * (1 + self.chunk_size // 4)252 253        chunks.append(254            VideoChunk(255                frames=input_video[:, :, start_idx:end_idx],256                start_idx=start_idx,257                end_idx=end_idx,258                current_start=current_start,259                current_end=current_end,260            )261        )262 263        while True:264            start_idx = end_idx265            end_idx = end_idx + self.chunk_size266            if end_idx > total_frames:267                break268            current_start = current_end269            current_end = current_end + (self.chunk_size // 4) * self.pipeline_manager.pipeline.frame_seq_length270            chunks.append(271                VideoChunk(272                    frames=input_video[:, :, start_idx:end_idx],273                    start_idx=start_idx,274                    end_idx=end_idx,275                    current_start=current_start,276                    current_end=current_end,277                )278            )279        return chunks280 281    @torch.inference_mode()282    def encode_chunk(283        self,284        input_video: str | Path | torch.Tensor,285        chunk: VideoChunk,286        *,287        previous_noise_scale: float | None = None,288        initial_noise_scale: float | None = None,289    ) -> EncodedChunk:290        """Encode one chunk in the same style as the offline inference loop."""291        full_video = _normalize_video_tensor(292            input_video,293            height=self.height,294            width=self.width,295            device=self.device,296        )297        noise_scale = self.noise_scale if previous_noise_scale is None else float(previous_noise_scale)298        init_noise_scale = self.noise_scale if initial_noise_scale is None else float(initial_noise_scale)299        current_step = None300 301        if chunk.start_idx != 0:302            noise_scale, current_step = compute_noise_scale_and_step(303                full_video,304                chunk.end_idx,305                self.chunk_size,306                noise_scale,307                init_noise_scale,308            )309 310        latents = self.pipeline_manager._timed_stream_encode(chunk.frames)311        latents = latents.transpose(2, 1).contiguous().to(dtype=torch.bfloat16)312        noise = torch.randn_like(latents)313        return EncodedChunk(314            noisy_latents=noise * noise_scale + latents * (1 - noise_scale),315            current_start=chunk.current_start,316            current_end=chunk.current_end,317            noise_scale=float(noise_scale),318            current_step=current_step,319        )320 321    @torch.inference_mode()322    def encode_video(self, video: str | Path | torch.Tensor) -> list[EncodedChunk]:323        """Encode a full input video into noisy latent chunks."""324        chunks: list[EncodedChunk] = []325        noise_scale = float(self.noise_scale)326        init_noise_scale = noise_scale327        video_chunks = self.chunk_video(video)328        full_video = _normalize_video_tensor(329            video,330            height=self.height,331            width=self.width,332            device=self.device,333        )334        for chunk in video_chunks:335            encoded_chunk = self.encode_chunk(336                full_video,337                chunk,338                previous_noise_scale=noise_scale,339                initial_noise_scale=init_noise_scale,340            )341            noise_scale = encoded_chunk.noise_scale342            chunks.append(encoded_chunk)343        return chunks344 345    @torch.inference_mode()346    def denoise_chunks(self, chunks: list[EncodedChunk]) -> list[DenoisedChunk]:347        """Run DiT denoising over the encoded chunks."""348        if not chunks:349            raise ValueError("chunks must not be empty")350        if self.prompt is None:351            raise RuntimeError("Call prepare(prompt) before denoise_chunks(...)")352 353        self.prepare(self.prompt)354        outputs: list[DenoisedChunk] = []355        for chunk in chunks:356            denoised_chunk = self.denoise_chunk(chunk)357            if denoised_chunk is not None:358                outputs.append(denoised_chunk)359        return outputs360 361    @torch.inference_mode()362    def denoise_chunk(self, chunk: EncodedChunk) -> DenoisedChunk | None:363        """Run DiT on one encoded chunk and return a decodable latent when available."""364        if self.prompt is None:365            raise RuntimeError("Call prepare(prompt) before denoise_chunk(...)")366 367        if self._next_chunk_index == 0:368            if self.mode == "single":369                denoised_pred = self.pipeline_manager.prepare_pipeline(370                    text_prompts=[self.prompt],371                    noise=chunk.noisy_latents,372                    current_start=chunk.current_start,373                    current_end=chunk.current_end,374                )375            else:376                denoised_pred = self.pipeline_manager.prepare_pipeline(377                    text_prompts=[self.prompt],378                    noise=chunk.noisy_latents,379                    current_start=chunk.current_start,380                    current_end=chunk.current_end,381                    batch_denoise=False,382                )383            self._next_chunk_index += 1384            return DenoisedChunk(denoised_pred=denoised_pred, last_frame_only=False)385 386        current_start = chunk.current_start387        current_end = chunk.current_end388 389        if current_start // self.pipeline_manager.pipeline.frame_seq_length >= self.pipeline_manager.t_refresh:390            current_start = self.pipeline_manager.pipeline.kv_cache_length - self.pipeline_manager.pipeline.frame_seq_length391            current_end = current_start + (self.chunk_size // 4) * self.pipeline_manager.pipeline.frame_seq_length392 393        if self.mode == "single":394            denoised_pred = self.pipeline_manager.pipeline.inference_stream(395                noise=chunk.noisy_latents,396                current_start=current_start,397                current_end=current_end,398                current_step=chunk.current_step,399            )400            self.pipeline_manager.processed += 1401            self._next_chunk_index += 1402            if self.pipeline_manager.processed < self.num_steps:403                return None404            return DenoisedChunk(denoised_pred=denoised_pred, last_frame_only=True)405 406        denoised_pred = self.pipeline_manager.pipeline.inference_wo_batch(407            noise=chunk.noisy_latents,408            current_start=current_start,409            current_end=current_end,410            current_step=chunk.current_step,411        )412        self.pipeline_manager.processed += 1413        self._next_chunk_index += 1414        return DenoisedChunk(denoised_pred=denoised_pred, last_frame_only=True)415 416    @torch.inference_mode()417    def decode_chunks(self, chunks: list[DenoisedChunk]) -> np.ndarray:418        """Decode denoised latent chunks into a `[T, H, W, C]` video array."""419        if not chunks:420            raise ValueError("chunks must not be empty")421        decoded = [self.decode_chunk(chunk) for chunk in chunks]422        return np.concatenate(decoded, axis=0)423 424    @torch.inference_mode()425    def decode_chunk(self, chunk: DenoisedChunk) -> np.ndarray:426        """Decode one denoised latent chunk into `[T, H, W, C]` frames."""427        return self.pipeline_manager._decode_video_array(428            chunk.denoised_pred,429            last_frame_only=chunk.last_frame_only,430        )431 432    @torch.inference_mode()433    def __call__(self, video: str | Path | torch.Tensor) -> np.ndarray:434        """Run the full staged pipeline after `prepare(prompt)` has been called."""435        encoded = self.encode_video(video)436        denoised = self.denoise_chunks(encoded)437        return self.decode_chunks(denoised)438