# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from Helios (https://github.com/BestWishYsh/Helios) from __future__ import annotations import json import logging import math import os from collections.abc import Iterable, Sequence from typing import TYPE_CHECKING, Any, ClassVar import numpy as np import torch import torch.nn.functional as F from diffusers import AutoencoderKLWan from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel from vllm_omni.diffusion.models.helios.scheduling_helios import HeliosScheduler from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.prompt_update import PromptUpdateMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm_omni.diffusion.worker.input_batch import InputBatch from vllm_omni.diffusion.worker.utils import StepRequestState logger = logging.getLogger(__name__) def calculate_shift( image_seq_len, base_seq_len: int = 256, max_seq_len: int = 4096, base_shift: float = 0.5, max_shift: float = 1.15, ): m = (max_shift - base_shift) / (max_seq_len - base_seq_len) b = base_shift - m * base_seq_len mu = image_seq_len * m + b return mu def optimized_scale(positive_flat, negative_flat): positive_flat = positive_flat.float() negative_flat = negative_flat.float() dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8 st_star = dot_product / squared_norm return st_star def load_json_config(model_path: str, subfolder: str, filename: str, local_files_only: bool = True) -> dict: """Load a JSON config file from a local path or HuggingFace Hub repo.""" if local_files_only: config_path = os.path.join(model_path, subfolder, filename) if os.path.exists(config_path): with open(config_path) as f: return json.load(f) else: try: from huggingface_hub import hf_hub_download config_path = hf_hub_download( repo_id=model_path, filename=f"{subfolder}/{filename}", ) with open(config_path) as f: return json.load(f) except Exception: pass return {} def load_transformer_config(model_path: str, subfolder: str = "transformer", local_files_only: bool = True) -> dict: return load_json_config(model_path, subfolder, "config.json", local_files_only) def create_transformer_from_config( config: dict, quant_config: QuantizationConfig | None = None ) -> HeliosTransformer3DModel: kwargs = {} key_map = [ "patch_size", "num_attention_heads", "attention_head_dim", "in_channels", "out_channels", "text_dim", "freq_dim", "ffn_dim", "num_layers", "cross_attn_norm", "qk_norm", "eps", "added_kv_proj_dim", "rope_dim", "rope_theta", "guidance_cross_attn", "zero_history_timestep", "has_multi_term_memory_patch", "is_amplify_history", "history_scale_mode", ] for key in key_map: if key in config: val = config[key] if key in ("patch_size", "rope_dim") and isinstance(val, list): val = tuple(val) kwargs[key] = val return HeliosTransformer3DModel(quant_config=quant_config, **kwargs) def get_helios_post_process_func( od_config: OmniDiffusionConfig, ): from diffusers.video_processor import VideoProcessor video_processor = VideoProcessor(vae_scale_factor=8) def post_process_func( video: torch.Tensor, output_type: str = "np", ): if output_type == "latent": return video return video_processor.postprocess_video(video, output_type=output_type) return post_process_func def get_helios_pre_process_func( od_config: OmniDiffusionConfig, ): def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: return request return pre_process_func class HeliosPipeline( nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin, PromptUpdateMixin, SupportsComponentDiscovery, ): """Helios text-to-video / image-to-video / video-to-video pipeline for vllm-omni. Supports T2V, I2V (with image input), and V2V (with video input). Implements chunked video generation with multi-term memory history context. """ supports_step_execution: ClassVar[bool] = True _dit_modules: ClassVar[list[str]] = ["transformer"] _encoder_modules: ClassVar[list[str]] = ["text_encoder"] _vae_modules: ClassVar[list[str]] = ["vae"] def __init__( self, *, od_config: OmniDiffusionConfig, prefix: str = "", ): super().__init__() self.od_config = od_config self.device = get_local_device() dtype = getattr(od_config, "dtype", torch.bfloat16) model = od_config.model local_files_only = os.path.exists(model) self.weights_sources = [ DiffusersPipelineLoader.ComponentSource( model_or_path=od_config.model, subfolder="transformer", revision=None, prefix="transformer.", fall_back_to_pt=True, ) ] self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) # Helios checkpoints store embed_tokens under ``shared.weight`` only, # but the published config sets ``tie_word_embeddings=False``. When # transformers sees ``tie=False`` it creates a separate # ``encoder.embed_tokens.weight`` that is never loaded from the # checkpoint, leaving it as all-zeros. This silently destroys prompt # encoding and produces grey/meaningless video output. Force tying so # that ``embed_tokens`` shares ``shared.weight`` as intended. text_enc_cfg = AutoConfig.from_pretrained(model, subfolder="text_encoder", local_files_only=local_files_only) text_enc_cfg.tie_word_embeddings = True self.text_encoder = UMT5EncoderModel.from_pretrained( model, subfolder="text_encoder", config=text_enc_cfg, torch_dtype=dtype, local_files_only=local_files_only ).to(self.device) self.vae = AutoencoderKLWan.from_pretrained( model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only ).to(self.device) transformer_config = load_transformer_config(model, "transformer", local_files_only) self.transformer = create_transformer_from_config( transformer_config, quant_config=od_config.quantization_config ) # Read scheduler config to determine scheduler type sched_cfg = load_json_config(model, "scheduler", "scheduler_config.json", local_files_only) scheduler_kwargs = {} passthrough_keys = [ "num_train_timesteps", "shift", "stages", "stage_range", "gamma", "thresholding", "prediction_type", "solver_order", "predict_x0", "solver_type", "lower_order_final", "disable_corrector", "use_flow_sigmas", "scheduler_type", "use_dynamic_shifting", "time_shift_type", ] for key in passthrough_keys: if key in sched_cfg: scheduler_kwargs[key] = sched_cfg[key] self.scheduler = HeliosScheduler(**scheduler_kwargs) self.is_distilled = scheduler_kwargs.get("scheduler_type") == "dmd" self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4 self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8 self._guidance_scale = None self._num_timesteps = None self._current_timestep = None self.setup_diffusion_pipeline_profiler( enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler ) @property def guidance_scale(self): return self._guidance_scale @property def do_classifier_free_guidance(self): return self._guidance_scale is not None and self._guidance_scale > 1.0 @property def num_timesteps(self): return self._num_timesteps @property def current_timestep(self): return self._current_timestep def _stage1_sigmas(self, num_steps: int) -> np.ndarray: # DMD drops the last timestep in set_timesteps(). Only compensate for # the one-step dummy-run edge case; otherwise preserve the historical # stage-1 schedule used by request-level forward(). sigma_count = num_steps + 2 if self.is_distilled and num_steps == 1 else num_steps + 1 return np.linspace(0.999, 0.0, sigma_count)[:-1] def prepare_encode( self, state: StepRequestState, **kwargs: Any, ) -> StepRequestState: """Initialize Helios request state for chunk-wise step execution.""" del kwargs # Wrap the single request in a DiffusionRequestBatch so the batch # compatibility properties (`prompts`, etc.) used below are available; # OmniDiffusionRequest itself only exposes a singular `prompt`. req = DiffusionRequestBatch( requests=[ OmniDiffusionRequest( prompt=state.prompt, sampling_params=state.sampling, request_id=state.request_id, ) ] ) extra = getattr(state.sampling, "extra_args", {}) or {} history_sizes = sorted(extra.get("history_sizes", [16, 2, 1]), reverse=True) num_latent_frames_per_chunk = int(extra.get("num_latent_frames_per_chunk", 9)) keep_first_frame = bool(extra.get("keep_first_frame", True)) frame_num = int(extra.get("frame_num", 132)) height = (int(state.sampling.height or extra.get("height", 384)) // 16) * 16 width = (int(state.sampling.width or extra.get("width", 640)) // 16) * 16 num_frames = max(int(state.sampling.num_frames or frame_num), 1) num_steps = int(state.sampling.num_inference_steps or extra.get("num_inference_steps", 50)) output_type = extra.get("output_type", "np") image = extra.get("image") video = extra.get("video") if image is not None and video is not None: raise ValueError("image and video cannot be provided simultaneously") if len(req.prompts) > 1: raise ValueError("Helios step execution supports a single prompt, not a batched request.") prompt = None negative_prompt = None if len(req.prompts) == 1: prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt") negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt") if prompt is None: raise ValueError("Prompt is required for Helios generation.") guidance_scale = float(extra.get("guidance_scale", 5.0)) if state.sampling.guidance_scale_provided: guidance_scale = state.sampling.guidance_scale self._guidance_scale = guidance_scale device = self.device dtype = self.transformer.dtype generator = state.sampling.generator if generator is None and state.sampling.seed is not None: generator = torch.Generator(device=device).manual_seed(state.sampling.seed) state.sampling.generator = generator prompt_embeds, negative_prompt_embeds = self.encode_prompt( prompt=prompt, negative_prompt=negative_prompt, do_classifier_free_guidance=self.do_classifier_free_guidance, num_videos_per_prompt=state.sampling.num_outputs_per_prompt or 1, max_sequence_length=state.sampling.max_sequence_length or 226, device=device, dtype=dtype, ) batch_size = prompt_embeds.shape[0] latents_mean = ( torch.tensor(self.vae.config.latents_mean) .view(1, self.vae.config.z_dim, 1, 1, 1) .to(self.vae.device, self.vae.dtype) ) latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( self.vae.device, self.vae.dtype ) add_noise_to_image_latents = bool(extra.get("add_noise_to_image_latents", True)) image_noise_sigma_min = float(extra.get("image_noise_sigma_min", 0.111)) image_noise_sigma_max = float(extra.get("image_noise_sigma_max", 0.135)) add_noise_to_video_latents = bool(extra.get("add_noise_to_video_latents", True)) video_noise_sigma_min = float(extra.get("video_noise_sigma_min", 0.111)) video_noise_sigma_max = float(extra.get("video_noise_sigma_max", 0.135)) fake_image_latents = None image_latents = None if image is not None: image_latents, fake_image_latents = self.prepare_image_latents( image, latents_mean=latents_mean, latents_std=latents_std, num_latent_frames_per_chunk=num_latent_frames_per_chunk, dtype=torch.float32, device=device, generator=generator, ) if image_latents is not None and add_noise_to_image_latents: image_noise_sigma = ( torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min) + image_noise_sigma_min ) image_latents = ( image_noise_sigma * randn_tensor(image_latents.shape, generator=generator, device=device) + (1 - image_noise_sigma) * image_latents ) fake_image_noise_sigma = ( torch.rand(1, device=device, generator=generator) * (video_noise_sigma_max - video_noise_sigma_min) + video_noise_sigma_min ) fake_image_latents = ( fake_image_noise_sigma * randn_tensor(fake_image_latents.shape, generator=generator, device=device) + (1 - fake_image_noise_sigma) * fake_image_latents ) video_latents = None if video is not None: image_latents, video_latents = self.prepare_video_latents( video, latents_mean=latents_mean, latents_std=latents_std, num_latent_frames_per_chunk=num_latent_frames_per_chunk, dtype=torch.float32, device=device, generator=generator, ) if video_latents is not None and add_noise_to_video_latents: image_noise_sigma = ( torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min) + image_noise_sigma_min ) image_latents = ( image_noise_sigma * randn_tensor(image_latents.shape, generator=generator, device=device) + (1 - image_noise_sigma) * image_latents ) noisy_latents_chunks = [] num_latent_chunks = video_latents.shape[2] // num_latent_frames_per_chunk for i in range(num_latent_chunks): chunk_start = i * num_latent_frames_per_chunk chunk_end = chunk_start + num_latent_frames_per_chunk latent_chunk = video_latents[:, :, chunk_start:chunk_end, :, :] chunk_frames = latent_chunk.shape[2] frame_sigmas = ( torch.rand(chunk_frames, device=device, generator=generator) * (video_noise_sigma_max - video_noise_sigma_min) + video_noise_sigma_min ) frame_sigmas = frame_sigmas.view(1, 1, chunk_frames, 1, 1) noisy_chunk = ( frame_sigmas * randn_tensor(latent_chunk.shape, generator=generator, device=device) + (1 - frame_sigmas) * latent_chunk ) noisy_latents_chunks.append(noisy_chunk) video_latents = torch.cat(noisy_latents_chunks, dim=2) num_channels_latents = self.transformer.config.in_channels window_num_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1 num_latent_chunk = max(1, (num_frames + window_num_frames - 1) // window_num_frames) num_history_latent_frames = sum(history_sizes) total_generated_latent_frames = 0 if not keep_first_frame: history_sizes[-1] = history_sizes[-1] + 1 history_latents = torch.zeros( batch_size, num_channels_latents, num_history_latent_frames, height // self.vae_scale_factor_spatial, width // self.vae_scale_factor_spatial, device=device, dtype=torch.float32, ) if fake_image_latents is not None: history_latents = torch.cat([history_latents[:, :, :-1, :, :], fake_image_latents], dim=2) total_generated_latent_frames += 1 if video_latents is not None: history_frames = history_latents.shape[2] video_frames = video_latents.shape[2] if video_frames < history_frames: keep_frames = history_frames - video_frames history_latents = torch.cat([history_latents[:, :, :keep_frames, :, :], video_latents], dim=2) else: history_latents = video_latents total_generated_latent_frames += video_latents.shape[2] if keep_first_frame: indices = torch.arange(0, sum([1, *history_sizes, num_latent_frames_per_chunk])) ( indices_prefix, indices_latents_history_long, indices_latents_history_mid, indices_latents_history_1x, indices_hidden_states, ) = indices.split([1, *history_sizes, num_latent_frames_per_chunk], dim=0) indices_latents_history_short = torch.cat([indices_prefix, indices_latents_history_1x], dim=0) else: indices = torch.arange(0, sum([*history_sizes, num_latent_frames_per_chunk])) ( indices_latents_history_long, indices_latents_history_mid, indices_latents_history_short, indices_hidden_states, ) = indices.split([*history_sizes, num_latent_frames_per_chunk], dim=0) state.prompt_embeds = prompt_embeds state.negative_prompt_embeds = negative_prompt_embeds state.do_true_cfg = self.do_classifier_free_guidance state.chunk_index = 0 state.step_in_chunk = 0 state.total_chunks = num_latent_chunk state.extra.update( { "attention_kwargs": extra.get("attention_kwargs", {}) or {}, "batch_size": batch_size, "dtype": dtype, "generator": generator, "guidance_scale": guidance_scale, "height": height, "history_latents": history_latents, "history_sizes": history_sizes, "history_video": None, "image_latents": image_latents, "indices_hidden_states": indices_hidden_states.unsqueeze(0), "indices_latents_history_long": indices_latents_history_long.unsqueeze(0), "indices_latents_history_mid": indices_latents_history_mid.unsqueeze(0), "indices_latents_history_short": indices_latents_history_short.unsqueeze(0), "is_amplify_first_chunk": bool(extra.get("is_amplify_first_chunk", False)), "is_enable_stage2": bool(extra.get("is_enable_stage2", False)), "is_skip_first_chunk": bool(extra.get("is_skip_first_chunk", False)), "keep_first_frame": keep_first_frame, "latents_mean": latents_mean, "latents_std": latents_std, "num_channels_latents": num_channels_latents, "num_history_latent_frames": num_history_latent_frames, "num_latent_frames_per_chunk": num_latent_frames_per_chunk, "num_steps": num_steps, "output_type": output_type, "pyramid_num_inference_steps_list": extra.get("pyramid_num_inference_steps_list", [10, 10, 10]), "pyramid_num_stages": int(extra.get("pyramid_num_stages", 3)), "total_generated_latent_frames": total_generated_latent_frames, "use_cfg_zero_star": bool(extra.get("use_cfg_zero_star", False)), "use_zero_init": bool(extra.get("use_zero_init", True)), "vae_dtype": self.vae.dtype, "width": width, "window_num_frames": window_num_frames, "zero_steps": int(extra.get("zero_steps", 1)), } ) self._prepare_next_chunk(state) return state def _prepare_next_chunk(self, state: StepRequestState) -> None: extra = state.extra k = state.chunk_index is_first_chunk = k == 0 is_second_chunk = k == 1 history_sizes = extra["history_sizes"] history_latents = extra["history_latents"] keep_first_frame = extra["keep_first_frame"] image_latents = extra["image_latents"] num_history_latent_frames = extra["num_history_latent_frames"] if keep_first_frame: latents_history_long, latents_history_mid, latents_history_1x = history_latents[ :, :, -num_history_latent_frames: ].split(history_sizes, dim=2) if image_latents is None and is_first_chunk: latents_prefix = torch.zeros( ( extra["batch_size"], extra["num_channels_latents"], 1, latents_history_1x.shape[-2], latents_history_1x.shape[-1], ), device=latents_history_1x.device, dtype=latents_history_1x.dtype, ) else: latents_prefix = image_latents latents_history_short = torch.cat([latents_prefix, latents_history_1x], dim=2) else: latents_history_long, latents_history_mid, latents_history_short = history_latents[ :, :, -num_history_latent_frames: ].split(history_sizes, dim=2) state.latents = self.prepare_latents( extra["batch_size"], extra["num_channels_latents"], extra["height"], extra["width"], extra["window_num_frames"], dtype=torch.float32, device=self.device, generator=extra["generator"], latents=None, ) extra["stage1_start_latents"] = state.latents extra.update( { "is_first_chunk": is_first_chunk, "is_second_chunk": is_second_chunk, "latents_history_long": latents_history_long, "latents_history_mid": latents_history_mid, "latents_history_short": latents_history_short, } ) if not extra["is_enable_stage2"]: patch_size = self.transformer.config.patch_size image_seq_len = (state.latents.shape[-1] * state.latents.shape[-2] * state.latents.shape[-3]) // ( patch_size[0] * patch_size[1] * patch_size[2] ) sigmas = self._stage1_sigmas(extra["num_steps"]) mu = calculate_shift(image_seq_len) self.scheduler.set_timesteps(extra["num_steps"], device=self.device, sigmas=sigmas, mu=mu) state.timesteps = self.scheduler.timesteps state.chunk_num_steps = len(state.timesteps) else: self._prepare_stage2_chunk(state) state.step_in_chunk = 0 state.step_index = 0 self._num_timesteps = state.chunk_num_steps def _prepare_stage2_chunk(self, state: StepRequestState) -> None: extra = state.extra batch_size, num_channel, num_frames_lat, height, width = state.latents.shape latents_flat = state.latents.permute(0, 2, 1, 3, 4).reshape( batch_size * num_frames_lat, num_channel, height, width ) for _ in range(extra["pyramid_num_stages"] - 1): height //= 2 width //= 2 latents_flat = F.interpolate(latents_flat, size=(height, width), mode="bilinear") * 2 state.latents = latents_flat.reshape(batch_size, num_frames_lat, num_channel, height, width).permute( 0, 2, 1, 3, 4 ) extra["stage2_height"] = height extra["stage2_width"] = width extra["stage2_start_point_list"] = [state.latents] if self.is_distilled else None extra["stage_index"] = 0 extra["stage_step_index"] = 0 amplify_first_chunk = extra["is_amplify_first_chunk"] and extra["is_first_chunk"] state.chunk_num_steps = sum( self._stage2_effective_num_steps(int(num_steps), amplify_first_chunk) for num_steps in extra["pyramid_num_inference_steps_list"] ) self._set_stage2_timesteps(state) def _stage2_effective_num_steps(self, num_steps: int, is_amplify_first_chunk: bool) -> int: if self.is_distilled and is_amplify_first_chunk: return num_steps * 2 return num_steps def _set_stage2_timesteps(self, state: StepRequestState) -> None: extra = state.extra patch_size = self.transformer.config.patch_size image_seq_len = (state.latents.shape[-1] * state.latents.shape[-2] * state.latents.shape[-3]) // ( patch_size[0] * patch_size[1] * patch_size[2] ) mu = calculate_shift(image_seq_len) stage_index = extra["stage_index"] self.scheduler.set_timesteps( extra["pyramid_num_inference_steps_list"][stage_index], stage_index=stage_index, device=self.device, mu=mu, is_amplify_first_chunk=extra["is_amplify_first_chunk"] and extra["is_first_chunk"], ) state.timesteps = self.scheduler.timesteps state.step_index = extra["stage_step_index"] def denoise_step( self, input_batch: InputBatch, states: Sequence[StepRequestState], **kwargs: Any, ) -> torch.Tensor | None: del kwargs if len(states) != 1: raise ValueError("Helios step execution supports a single request, not a batched request.") state = states[0] if state.extra["is_enable_stage2"]: return self._denoise_stage2_step(state) return self._denoise_stage1_step(state, input_batch.latents, input_batch.timesteps) def _denoise_stage1_step( self, state: StepRequestState, latents: torch.Tensor, timesteps: torch.Tensor, ) -> torch.Tensor: extra = state.extra batch_size = latents.shape[0] t = timesteps self._current_timestep = t timestep = t.expand(batch_size) transformer_kwargs = { "hidden_states": latents.to(extra["dtype"]), "timestep": timestep, "indices_hidden_states": extra["indices_hidden_states"], "indices_latents_history_short": extra["indices_latents_history_short"], "indices_latents_history_mid": extra["indices_latents_history_mid"], "indices_latents_history_long": extra["indices_latents_history_long"], "latents_history_short": extra["latents_history_short"].to(extra["dtype"]), "latents_history_mid": extra["latents_history_mid"].to(extra["dtype"]), "latents_history_long": extra["latents_history_long"].to(extra["dtype"]), "attention_kwargs": extra["attention_kwargs"], "return_dict": False, } if extra["use_cfg_zero_star"] and state.do_true_cfg: noise_pred = self.transformer( encoder_hidden_states=state.prompt_embeds, **transformer_kwargs, )[0] noise_uncond = self.transformer( encoder_hidden_states=state.negative_prompt_embeds, **transformer_kwargs, )[0] positive_flat = noise_pred.view(batch_size, -1) negative_flat = noise_uncond.view(batch_size, -1) alpha_cfg = optimized_scale(positive_flat, negative_flat) alpha_cfg = alpha_cfg.view(batch_size, *([1] * (len(noise_pred.shape) - 1))).to(noise_pred.dtype) if (state.step_in_chunk <= extra["zero_steps"]) and extra["use_zero_init"]: return noise_pred * 0.0 return noise_uncond * alpha_cfg + extra["guidance_scale"] * (noise_pred - noise_uncond * alpha_cfg) positive_kwargs = { "encoder_hidden_states": state.prompt_embeds, **transformer_kwargs, } negative_kwargs = ( { "encoder_hidden_states": state.negative_prompt_embeds, **transformer_kwargs, } if state.do_true_cfg else None ) return self.predict_noise_maybe_with_cfg( do_true_cfg=state.do_true_cfg, true_cfg_scale=extra["guidance_scale"], positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs, cfg_normalize=False, ) def _denoise_stage2_step(self, state: StepRequestState) -> torch.Tensor: extra = state.extra latents = state.latents assert latents is not None batch_size = latents.shape[0] t = state.current_timestep assert t is not None self._current_timestep = t timestep = t.expand(batch_size).to(torch.int64) transformer_kwargs = { "hidden_states": latents.to(extra["dtype"]), "timestep": timestep, "indices_hidden_states": extra["indices_hidden_states"], "indices_latents_history_short": extra["indices_latents_history_short"], "indices_latents_history_mid": extra["indices_latents_history_mid"], "indices_latents_history_long": extra["indices_latents_history_long"], "latents_history_short": extra["latents_history_short"].to(extra["dtype"]), "latents_history_mid": extra["latents_history_mid"].to(extra["dtype"]), "latents_history_long": extra["latents_history_long"].to(extra["dtype"]), "attention_kwargs": extra["attention_kwargs"], "return_dict": False, } noise_pred = self.transformer(encoder_hidden_states=state.prompt_embeds, **transformer_kwargs)[0] if state.do_true_cfg: noise_uncond = self.transformer(encoder_hidden_states=state.negative_prompt_embeds, **transformer_kwargs)[0] if extra["use_cfg_zero_star"]: positive_flat = noise_pred.view(batch_size, -1) negative_flat = noise_uncond.view(batch_size, -1) alpha_cfg = optimized_scale(positive_flat, negative_flat) alpha_cfg = alpha_cfg.view(batch_size, *([1] * (len(noise_pred.shape) - 1))).to(noise_pred.dtype) if ( extra["stage_index"] == 0 and extra["stage_step_index"] <= extra["zero_steps"] and extra["use_zero_init"] ): noise_pred = noise_pred * 0.0 else: noise_pred = noise_uncond * alpha_cfg + extra["guidance_scale"] * ( noise_pred - noise_uncond * alpha_cfg ) else: noise_pred = noise_uncond + extra["guidance_scale"] * (noise_pred - noise_uncond) return noise_pred def step_scheduler( self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any, ) -> None: del kwargs if state.extra["is_enable_stage2"]: self._step_scheduler_stage2(state, noise_pred) else: t = state.current_timestep assert t is not None if self.is_distilled: state.latents = self.scheduler.step( noise_pred, t, state.latents, return_dict=False, cur_sampling_step=state.step_in_chunk, dmd_noisy_tensor=state.extra["stage1_start_latents"], dmd_sigmas=self.scheduler.sigmas, dmd_timesteps=self.scheduler.timesteps, all_timesteps=state.timesteps, )[0] else: state.latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, state.latents, state.do_true_cfg) state.step_in_chunk += 1 state.step_index = state.step_in_chunk def _step_scheduler_stage2(self, state: StepRequestState, noise_pred: torch.Tensor) -> None: extra = state.extra t = state.current_timestep assert t is not None and state.latents is not None start_points = extra["stage2_start_point_list"] state.latents = self.scheduler.step( noise_pred, t, state.latents, return_dict=False, cur_sampling_step=extra["stage_step_index"], dmd_noisy_tensor=start_points[extra["stage_index"]] if start_points is not None else None, dmd_sigmas=self.scheduler.sigmas, dmd_timesteps=self.scheduler.timesteps, all_timesteps=state.timesteps, )[0] extra["stage_step_index"] += 1 state.step_in_chunk += 1 state.step_index = extra["stage_step_index"] if extra["stage_step_index"] < len(state.timesteps): return if extra["stage_index"] >= extra["pyramid_num_stages"] - 1: return extra["stage_index"] += 1 extra["stage_step_index"] = 0 extra["stage2_height"] *= 2 extra["stage2_width"] *= 2 batch_size, num_channel, num_frames_cur = state.latents.shape[:3] latents_flat = state.latents.permute(0, 2, 1, 3, 4).reshape( batch_size * num_frames_cur, num_channel, extra["stage2_height"] // 2, extra["stage2_width"] // 2, ) latents_flat = F.interpolate(latents_flat, size=(extra["stage2_height"], extra["stage2_width"]), mode="nearest") state.latents = latents_flat.reshape( batch_size, num_frames_cur, num_channel, extra["stage2_height"], extra["stage2_width"], ).permute(0, 2, 1, 3, 4) ori_sigma = 1 - self.scheduler.ori_start_sigmas[extra["stage_index"]] gamma = self.scheduler.config.gamma alpha = 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma) beta = alpha * (1 - ori_sigma) / math.sqrt(gamma) patch_size = self.transformer.config.patch_size noise = self.sample_block_noise( batch_size, num_channel, state.latents.shape[2], extra["stage2_height"], extra["stage2_width"], patch_size, generator=extra["generator"], ).to(device=self.device, dtype=extra["dtype"]) state.latents = alpha * state.latents + beta * noise if self.is_distilled and extra["stage2_start_point_list"] is not None: extra["stage2_start_point_list"].append(state.latents) self._set_stage2_timesteps(state) def post_decode( self, state: StepRequestState, **kwargs: Any, ) -> DiffusionOutput: del kwargs extra = state.extra is_first_chunk = extra["is_first_chunk"] is_second_chunk = extra["is_second_chunk"] if extra["keep_first_frame"] and ( (is_first_chunk and extra["image_latents"] is None) or (extra["is_skip_first_chunk"] and is_second_chunk) ): extra["image_latents"] = state.latents[:, :, 0:1, :, :] extra["total_generated_latent_frames"] += state.latents.shape[2] extra["history_latents"] = torch.cat([extra["history_latents"], state.latents], dim=2) real_history_latents = extra["history_latents"][:, :, -extra["total_generated_latent_frames"] :] current_latents = ( real_history_latents[:, :, -extra["num_latent_frames_per_chunk"] :].to(extra["vae_dtype"]) / extra["latents_std"] + extra["latents_mean"] ) current_video = self.vae.decode(current_latents, return_dict=False)[0] if extra["history_video"] is None: extra["history_video"] = current_video else: extra["history_video"] = torch.cat([extra["history_video"], current_video], dim=2) output = current_latents if extra["output_type"] == "latent" else current_video completed_chunk_index = state.chunk_index prompt_update_metadata = state.extra.pop("prompt_update_chunk_metadata", {}) state.chunk_index += 1 finished = state.request_denoise_completed if not finished: # Apply queued/advancing prompt updates before preparing the next chunk. self._apply_prompt_update_at_chunk_boundary(state) self._prepare_next_chunk(state) else: self._current_timestep = None if current_omni_platform.is_available(): current_omni_platform.empty_cache() return DiffusionOutput( output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else {}, chunk_index=completed_chunk_index, total_chunks=state.total_chunks, started_event_ids=prompt_update_metadata.get("started_event_ids", []), active_event_ids=prompt_update_metadata.get("active_event_ids", []), completed_event_ids=prompt_update_metadata.get("completed_event_ids", []), finished=finished, ) def forward( self, req: DiffusionRequestBatch, prompt: str | None = None, negative_prompt: str | None = None, height: int = 384, width: int = 640, num_inference_steps: int = 50, guidance_scale: float = 5.0, frame_num: int = 132, output_type: str | None = "np", generator: torch.Generator | list[torch.Generator] | None = None, prompt_embeds: torch.Tensor | None = None, negative_prompt_embeds: torch.Tensor | None = None, attention_kwargs: dict | None = None, # Helios-specific history_sizes: list | None = None, num_latent_frames_per_chunk: int = 9, keep_first_frame: bool = True, # I2V image: torch.Tensor | None = None, add_noise_to_image_latents: bool = True, image_noise_sigma_min: float = 0.111, image_noise_sigma_max: float = 0.135, # V2V video: torch.Tensor | None = None, add_noise_to_video_latents: bool = True, video_noise_sigma_min: float = 0.111, video_noise_sigma_max: float = 0.135, # Stage 2 (pyramid multi-stage denoising) is_enable_stage2: bool = False, pyramid_num_stages: int = 3, pyramid_num_inference_steps_list: list | None = None, is_skip_first_chunk: bool = False, # DMD is_amplify_first_chunk: bool = False, # CFG Zero Star use_cfg_zero_star: bool = False, use_zero_init: bool = True, zero_steps: int = 1, **kwargs, ) -> DiffusionOutput: if pyramid_num_inference_steps_list is None: pyramid_num_inference_steps_list = [10, 10, 10] if history_sizes is None: history_sizes = [16, 2, 1] # Read Helios-specific params from extra_args extra = getattr(req.sampling_params, "extra_args", {}) or {} is_enable_stage2 = extra.get("is_enable_stage2", is_enable_stage2) pyramid_num_stages = extra.get("pyramid_num_stages", pyramid_num_stages) pyramid_num_inference_steps_list = extra.get( "pyramid_num_inference_steps_list", pyramid_num_inference_steps_list ) is_amplify_first_chunk = extra.get("is_amplify_first_chunk", is_amplify_first_chunk) use_cfg_zero_star = extra.get("use_cfg_zero_star", use_cfg_zero_star) use_zero_init = extra.get("use_zero_init", use_zero_init) zero_steps = extra.get("zero_steps", zero_steps) is_skip_first_chunk = extra.get("is_skip_first_chunk", is_skip_first_chunk) image = extra.get("image", image) video = extra.get("video", video) add_noise_to_image_latents = extra.get("add_noise_to_image_latents", add_noise_to_image_latents) image_noise_sigma_min = extra.get("image_noise_sigma_min", image_noise_sigma_min) image_noise_sigma_max = extra.get("image_noise_sigma_max", image_noise_sigma_max) add_noise_to_video_latents = extra.get("add_noise_to_video_latents", add_noise_to_video_latents) video_noise_sigma_min = extra.get("video_noise_sigma_min", video_noise_sigma_min) video_noise_sigma_max = extra.get("video_noise_sigma_max", video_noise_sigma_max) if image is not None and video is not None: raise ValueError("image and video cannot be provided simultaneously") if len(req.prompts) > 1: raise ValueError("This model only supports a single prompt, not a batched request.") if len(req.prompts) == 1: prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt") negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt") if prompt is None and prompt_embeds is None: raise ValueError("Prompt or prompt_embeds is required for Helios generation.") height = req.sampling_params.height or height width = req.sampling_params.width or width num_frames = req.sampling_params.num_frames if req.sampling_params.num_frames else frame_num num_steps = req.sampling_params.num_inference_steps or num_inference_steps if req.sampling_params.guidance_scale_provided: guidance_scale = req.sampling_params.guidance_scale self._guidance_scale = guidance_scale height = (height // 16) * 16 width = (width // 16) * 16 num_frames = max(num_frames, 1) device = self.device dtype = self.transformer.dtype if generator is None: generator = req.sampling_params.generator if generator is None and req.sampling_params.seed is not None: generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed) # Encode prompts if prompt_embeds is None: prompt_embeds, negative_prompt_embeds = self.encode_prompt( prompt=prompt, negative_prompt=negative_prompt, do_classifier_free_guidance=self.do_classifier_free_guidance, num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, max_sequence_length=req.sampling_params.max_sequence_length or 226, device=device, dtype=dtype, ) else: prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) if negative_prompt_embeds is not None: negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) batch_size = prompt_embeds.shape[0] history_sizes = sorted(history_sizes, reverse=True) latents_mean = ( torch.tensor(self.vae.config.latents_mean) .view(1, self.vae.config.z_dim, 1, 1, 1) .to(self.vae.device, self.vae.dtype) ) latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( self.vae.device, self.vae.dtype ) # Prepare I2V image latents fake_image_latents = None image_latents = None if image is not None: image_latents, fake_image_latents = self.prepare_image_latents( image, latents_mean=latents_mean, latents_std=latents_std, num_latent_frames_per_chunk=num_latent_frames_per_chunk, dtype=torch.float32, device=device, generator=generator, ) if image_latents is not None and add_noise_to_image_latents: image_noise_sigma = ( torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min) + image_noise_sigma_min ) image_latents = ( image_noise_sigma * randn_tensor(image_latents.shape, generator=generator, device=device) + (1 - image_noise_sigma) * image_latents ) fake_image_noise_sigma = ( torch.rand(1, device=device, generator=generator) * (video_noise_sigma_max - video_noise_sigma_min) + video_noise_sigma_min ) fake_image_latents = ( fake_image_noise_sigma * randn_tensor(fake_image_latents.shape, generator=generator, device=device) + (1 - fake_image_noise_sigma) * fake_image_latents ) # Prepare V2V video latents video_latents = None if video is not None: image_latents, video_latents = self.prepare_video_latents( video, latents_mean=latents_mean, latents_std=latents_std, num_latent_frames_per_chunk=num_latent_frames_per_chunk, dtype=torch.float32, device=device, generator=generator, ) if video_latents is not None and add_noise_to_video_latents: image_noise_sigma = ( torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min) + image_noise_sigma_min ) image_latents = ( image_noise_sigma * randn_tensor(image_latents.shape, generator=generator, device=device) + (1 - image_noise_sigma) * image_latents ) noisy_latents_chunks = [] num_latent_chunks = video_latents.shape[2] // num_latent_frames_per_chunk for i in range(num_latent_chunks): chunk_start = i * num_latent_frames_per_chunk chunk_end = chunk_start + num_latent_frames_per_chunk latent_chunk = video_latents[:, :, chunk_start:chunk_end, :, :] chunk_frames = latent_chunk.shape[2] frame_sigmas = ( torch.rand(chunk_frames, device=device, generator=generator) * (video_noise_sigma_max - video_noise_sigma_min) + video_noise_sigma_min ) frame_sigmas = frame_sigmas.view(1, 1, chunk_frames, 1, 1) noisy_chunk = ( frame_sigmas * randn_tensor(latent_chunk.shape, generator=generator, device=device) + (1 - frame_sigmas) * latent_chunk ) noisy_latents_chunks.append(noisy_chunk) video_latents = torch.cat(noisy_latents_chunks, dim=2) # Prepare latent variables num_channels_latents = self.transformer.config.in_channels window_num_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1 num_latent_chunk = max(1, (num_frames + window_num_frames - 1) // window_num_frames) num_history_latent_frames = sum(history_sizes) history_video = None total_generated_latent_frames = 0 if not keep_first_frame: history_sizes[-1] = history_sizes[-1] + 1 history_latents = torch.zeros( batch_size, num_channels_latents, num_history_latent_frames, height // self.vae_scale_factor_spatial, width // self.vae_scale_factor_spatial, device=device, dtype=torch.float32, ) if fake_image_latents is not None: history_latents = torch.cat([history_latents[:, :, :-1, :, :], fake_image_latents], dim=2) total_generated_latent_frames += 1 if video_latents is not None: history_frames = history_latents.shape[2] video_frames = video_latents.shape[2] if video_frames < history_frames: keep_frames = history_frames - video_frames history_latents = torch.cat([history_latents[:, :, :keep_frames, :, :], video_latents], dim=2) else: history_latents = video_latents total_generated_latent_frames += video_latents.shape[2] # Prepare frame indices if keep_first_frame: indices = torch.arange(0, sum([1, *history_sizes, num_latent_frames_per_chunk])) ( indices_prefix, indices_latents_history_long, indices_latents_history_mid, indices_latents_history_1x, indices_hidden_states, ) = indices.split([1, *history_sizes, num_latent_frames_per_chunk], dim=0) indices_latents_history_short = torch.cat([indices_prefix, indices_latents_history_1x], dim=0) else: indices = torch.arange(0, sum([*history_sizes, num_latent_frames_per_chunk])) ( indices_latents_history_long, indices_latents_history_mid, indices_latents_history_short, indices_hidden_states, ) = indices.split([*history_sizes, num_latent_frames_per_chunk], dim=0) indices_hidden_states = indices_hidden_states.unsqueeze(0) indices_latents_history_short = indices_latents_history_short.unsqueeze(0) indices_latents_history_mid = indices_latents_history_mid.unsqueeze(0) indices_latents_history_long = indices_latents_history_long.unsqueeze(0) if attention_kwargs is None: attention_kwargs = {} vae_dtype = self.vae.dtype # Chunked denoising loop for k in range(num_latent_chunk): is_first_chunk = k == 0 is_second_chunk = k == 1 if keep_first_frame: latents_history_long, latents_history_mid, latents_history_1x = history_latents[ :, :, -num_history_latent_frames: ].split(history_sizes, dim=2) if image_latents is None and is_first_chunk: latents_prefix = torch.zeros( ( batch_size, num_channels_latents, 1, latents_history_1x.shape[-2], latents_history_1x.shape[-1], ), device=latents_history_1x.device, dtype=latents_history_1x.dtype, ) else: latents_prefix = image_latents latents_history_short = torch.cat([latents_prefix, latents_history_1x], dim=2) else: latents_history_long, latents_history_mid, latents_history_short = history_latents[ :, :, -num_history_latent_frames: ].split(history_sizes, dim=2) # Prepare noise for this chunk latents = self.prepare_latents( batch_size, num_channels_latents, height, width, window_num_frames, dtype=torch.float32, device=device, generator=generator, latents=None, ) if not is_enable_stage2: # Stage 1 only: single-stage denoising patch_size = self.transformer.config.patch_size image_seq_len = (latents.shape[-1] * latents.shape[-2] * latents.shape[-3]) // ( patch_size[0] * patch_size[1] * patch_size[2] ) sigmas = np.linspace(0.999, 0.0, num_steps + 1)[:-1] mu = calculate_shift(image_seq_len) self.scheduler.set_timesteps(num_steps, device=device, sigmas=sigmas, mu=mu) timesteps = self.scheduler.timesteps self._num_timesteps = len(timesteps) latents = self._stage1_sample( latents=latents, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, timesteps=timesteps, guidance_scale=guidance_scale, indices_hidden_states=indices_hidden_states, indices_latents_history_short=indices_latents_history_short, indices_latents_history_mid=indices_latents_history_mid, indices_latents_history_long=indices_latents_history_long, latents_history_short=latents_history_short, latents_history_mid=latents_history_mid, latents_history_long=latents_history_long, attention_kwargs=attention_kwargs, transformer_dtype=dtype, use_cfg_zero_star=use_cfg_zero_star, use_zero_init=use_zero_init, zero_steps=zero_steps, ) else: # Stage 2: pyramid multi-stage denoising latents = self._stage2_sample( latents=latents, pyramid_num_stages=pyramid_num_stages, pyramid_num_inference_steps_list=pyramid_num_inference_steps_list, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, guidance_scale=guidance_scale, indices_hidden_states=indices_hidden_states, indices_latents_history_short=indices_latents_history_short, indices_latents_history_mid=indices_latents_history_mid, indices_latents_history_long=indices_latents_history_long, latents_history_short=latents_history_short, latents_history_mid=latents_history_mid, latents_history_long=latents_history_long, attention_kwargs=attention_kwargs, transformer_dtype=dtype, is_amplify_first_chunk=is_amplify_first_chunk and is_first_chunk, use_cfg_zero_star=use_cfg_zero_star, use_zero_init=use_zero_init, zero_steps=zero_steps, device=device, generator=generator, ) if keep_first_frame and ( (is_first_chunk and image_latents is None) or (is_skip_first_chunk and is_second_chunk) ): image_latents = latents[:, :, 0:1, :, :] total_generated_latent_frames += latents.shape[2] history_latents = torch.cat([history_latents, latents], dim=2) real_history_latents = history_latents[:, :, -total_generated_latent_frames:] index_slice = ( slice(None), slice(None), slice(-num_latent_frames_per_chunk, None), ) current_latents = real_history_latents[index_slice].to(vae_dtype) / latents_std + latents_mean current_video = self.vae.decode(current_latents, return_dict=False)[0] if history_video is None: history_video = current_video else: history_video = torch.cat([history_video, current_video], dim=2) if current_omni_platform.is_available(): current_omni_platform.empty_cache() self._current_timestep = None if output_type == "latent": output = real_history_latents else: output = history_video return DiffusionOutput( output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None ) def _stage1_sample( self, latents: torch.Tensor, prompt_embeds: torch.Tensor, negative_prompt_embeds: torch.Tensor | None, timesteps: torch.Tensor, guidance_scale: float, indices_hidden_states: torch.Tensor, indices_latents_history_short: torch.Tensor, indices_latents_history_mid: torch.Tensor, indices_latents_history_long: torch.Tensor, latents_history_short: torch.Tensor, latents_history_mid: torch.Tensor, latents_history_long: torch.Tensor, attention_kwargs: dict, transformer_dtype: torch.dtype, use_cfg_zero_star: bool = False, use_zero_init: bool = True, zero_steps: int = 1, ) -> torch.Tensor: """Single-stage denoising loop for one chunk.""" batch_size = latents.shape[0] do_true_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None with self.progress_bar(total=len(timesteps)) as pbar: for i, t in enumerate(timesteps): self._current_timestep = t timestep = t.expand(batch_size) transformer_kwargs = { "hidden_states": latents.to(transformer_dtype), "timestep": timestep, "indices_hidden_states": indices_hidden_states, "indices_latents_history_short": indices_latents_history_short, "indices_latents_history_mid": indices_latents_history_mid, "indices_latents_history_long": indices_latents_history_long, "latents_history_short": latents_history_short.to(transformer_dtype), "latents_history_mid": latents_history_mid.to(transformer_dtype), "latents_history_long": latents_history_long.to(transformer_dtype), "attention_kwargs": attention_kwargs, "return_dict": False, } if use_cfg_zero_star and do_true_cfg: noise_pred = self.transformer( encoder_hidden_states=prompt_embeds, **transformer_kwargs, )[0] noise_uncond = self.transformer( encoder_hidden_states=negative_prompt_embeds, **transformer_kwargs, )[0] positive_flat = noise_pred.view(batch_size, -1) negative_flat = noise_uncond.view(batch_size, -1) alpha_cfg = optimized_scale(positive_flat, negative_flat) alpha_cfg = alpha_cfg.view(batch_size, *([1] * (len(noise_pred.shape) - 1))) alpha_cfg = alpha_cfg.to(noise_pred.dtype) if (i <= zero_steps) and use_zero_init: noise_pred = noise_pred * 0.0 else: noise_pred = noise_uncond * alpha_cfg + guidance_scale * (noise_pred - noise_uncond * alpha_cfg) else: positive_kwargs = { "encoder_hidden_states": prompt_embeds, **transformer_kwargs, } if do_true_cfg: negative_kwargs = { "encoder_hidden_states": negative_prompt_embeds, **transformer_kwargs, } else: negative_kwargs = None noise_pred = self.predict_noise_maybe_with_cfg( do_true_cfg=do_true_cfg, true_cfg_scale=guidance_scale, positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs, cfg_normalize=False, ) latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) pbar.update() return latents def _stage2_sample( self, latents: torch.Tensor, pyramid_num_stages: int, pyramid_num_inference_steps_list: list[int], prompt_embeds: torch.Tensor, negative_prompt_embeds: torch.Tensor | None, guidance_scale: float, indices_hidden_states: torch.Tensor, indices_latents_history_short: torch.Tensor, indices_latents_history_mid: torch.Tensor, indices_latents_history_long: torch.Tensor, latents_history_short: torch.Tensor, latents_history_mid: torch.Tensor, latents_history_long: torch.Tensor, attention_kwargs: dict, transformer_dtype: torch.dtype, is_amplify_first_chunk: bool = False, use_cfg_zero_star: bool = False, use_zero_init: bool = True, zero_steps: int = 1, device: torch.device | None = None, generator: torch.Generator | list[torch.Generator] | None = None, ) -> torch.Tensor: """Pyramid multi-stage denoising for one chunk.""" batch_size, num_channel, num_frames_lat, height, width = latents.shape # Downsample latents to the smallest pyramid level latents_flat = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames_lat, num_channel, height, width) for _ in range(pyramid_num_stages - 1): height //= 2 width //= 2 latents_flat = F.interpolate(latents_flat, size=(height, width), mode="bilinear") * 2 latents = latents_flat.reshape(batch_size, num_frames_lat, num_channel, height, width).permute(0, 2, 1, 3, 4) start_point_list = None if self.is_distilled: start_point_list = [latents] do_true_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None for i_s in range(pyramid_num_stages): patch_size = self.transformer.config.patch_size image_seq_len = (latents.shape[-1] * latents.shape[-2] * latents.shape[-3]) // ( patch_size[0] * patch_size[1] * patch_size[2] ) mu = calculate_shift(image_seq_len) self.scheduler.set_timesteps( pyramid_num_inference_steps_list[i_s], stage_index=i_s, device=device, mu=mu, is_amplify_first_chunk=is_amplify_first_chunk, ) timesteps = self.scheduler.timesteps if i_s > 0: # Upsample latents to next pyramid level height *= 2 width *= 2 num_frames_cur = latents.shape[2] latents_flat = latents.permute(0, 2, 1, 3, 4).reshape( batch_size * num_frames_cur, num_channel, height // 2, width // 2 ) latents_flat = F.interpolate(latents_flat, size=(height, width), mode="nearest") latents = latents_flat.reshape(batch_size, num_frames_cur, num_channel, height, width).permute( 0, 2, 1, 3, 4 ) # Add block noise to fix artifacts between stages ori_sigma = 1 - self.scheduler.ori_start_sigmas[i_s] gamma = self.scheduler.config.gamma alpha = 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma) beta = alpha * (1 - ori_sigma) / math.sqrt(gamma) noise = self.sample_block_noise( batch_size, num_channel, latents.shape[2], height, width, patch_size, generator=generator ) noise = noise.to(device=device, dtype=transformer_dtype) latents = alpha * latents + beta * noise if self.is_distilled and start_point_list is not None: start_point_list.append(latents) with self.progress_bar(total=len(timesteps)) as pbar: for idx, t in enumerate(timesteps): self._current_timestep = t timestep = t.expand(latents.shape[0]).to(torch.int64) transformer_kwargs = { "hidden_states": latents.to(transformer_dtype), "timestep": timestep, "indices_hidden_states": indices_hidden_states, "indices_latents_history_short": indices_latents_history_short, "indices_latents_history_mid": indices_latents_history_mid, "indices_latents_history_long": indices_latents_history_long, "latents_history_short": latents_history_short.to(transformer_dtype), "latents_history_mid": latents_history_mid.to(transformer_dtype), "latents_history_long": latents_history_long.to(transformer_dtype), "attention_kwargs": attention_kwargs, "return_dict": False, } noise_pred = self.transformer( encoder_hidden_states=prompt_embeds, **transformer_kwargs, )[0] if do_true_cfg: noise_uncond = self.transformer( encoder_hidden_states=negative_prompt_embeds, **transformer_kwargs, )[0] if use_cfg_zero_star: positive_flat = noise_pred.view(batch_size, -1) negative_flat = noise_uncond.view(batch_size, -1) alpha_cfg = optimized_scale(positive_flat, negative_flat) alpha_cfg = alpha_cfg.view(batch_size, *([1] * (len(noise_pred.shape) - 1))) alpha_cfg = alpha_cfg.to(noise_pred.dtype) if (i_s == 0 and idx <= zero_steps) and use_zero_init: noise_pred = noise_pred * 0.0 else: noise_pred = noise_uncond * alpha_cfg + guidance_scale * ( noise_pred - noise_uncond * alpha_cfg ) else: noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond) latents = self.scheduler.step( noise_pred, t, latents, return_dict=False, cur_sampling_step=idx, dmd_noisy_tensor=start_point_list[i_s] if start_point_list is not None else None, dmd_sigmas=self.scheduler.sigmas, dmd_timesteps=self.scheduler.timesteps, all_timesteps=timesteps, )[0] pbar.update() return latents def sample_block_noise(self, batch_size, channel, num_frames, height, width, patch_size=(1, 2, 2), generator=None): gamma = self.scheduler.config.gamma _, ph, pw = patch_size block_size = ph * pw device = generator.device if generator is not None else self.device # Allocate directly on the execution device in float32 to use the device solver # and avoid fp16/bf16 Cholesky on the covariance matrix. eye = torch.eye(block_size, device=device, dtype=torch.float32) cov = eye * (1 + gamma) - torch.ones(block_size, block_size, device=device, dtype=torch.float32) * gamma cov += eye * 1e-8 L = torch.linalg.cholesky(cov) block_number = batch_size * channel * num_frames * (height // ph) * (width // pw) z = torch.randn(block_number, block_size, generator=generator, device=device) noise = z @ L.T noise = noise.view(batch_size, channel, num_frames, height // ph, width // pw, ph, pw) noise = noise.permute(0, 1, 2, 3, 5, 4, 6).reshape(batch_size, channel, num_frames, height, width) return noise def predict_noise(self, **kwargs: Any) -> torch.Tensor: return self.transformer(**kwargs)[0] def encode_prompt( self, prompt: str | list[str], negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, max_sequence_length: int | None = None, device: torch.device | None = None, dtype: torch.dtype | None = None, ): device = device or self.device dtype = dtype or self.text_encoder.dtype max_sequence_length = max_sequence_length or 226 prompt = [prompt] if isinstance(prompt, str) else prompt prompt_clean = [self._prompt_clean(p) for p in prompt] batch_size = len(prompt_clean) text_inputs = self.tokenizer( prompt_clean, padding="max_length", max_length=max_sequence_length, truncation=True, add_special_tokens=True, return_attention_mask=True, return_tensors="pt", ) ids, mask = text_inputs.input_ids, text_inputs.attention_mask seq_lens = mask.gt(0).sum(dim=1).long() prompt_embeds = self.text_encoder(ids.to(device), mask.to(device)).last_hidden_state prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0 ) _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) negative_prompt_embeds = None if do_classifier_free_guidance: negative_prompt = negative_prompt or "" negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt neg_text_inputs = self.tokenizer( [self._prompt_clean(p) for p in negative_prompt], padding="max_length", max_length=max_sequence_length, truncation=True, add_special_tokens=True, return_attention_mask=True, return_tensors="pt", ) ids_neg, mask_neg = neg_text_inputs.input_ids, neg_text_inputs.attention_mask seq_lens_neg = mask_neg.gt(0).sum(dim=1).long() negative_prompt_embeds = self.text_encoder(ids_neg.to(device), mask_neg.to(device)).last_hidden_state negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device) negative_prompt_embeds = [u[:v] for u, v in zip(negative_prompt_embeds, seq_lens_neg)] negative_prompt_embeds = torch.stack( [ torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in negative_prompt_embeds ], dim=0, ) negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1) negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) return prompt_embeds, negative_prompt_embeds @staticmethod def _prompt_clean(text: str) -> str: return " ".join(text.strip().split()) def prepare_latents( self, batch_size: int, num_channels_latents: int, height: int, width: int, num_frames: int, dtype: torch.dtype | None, device: torch.device | None, generator: torch.Generator | list[torch.Generator] | None, latents: torch.Tensor | None = None, ) -> torch.Tensor: if latents is not None: return latents.to(device=device, dtype=dtype) num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 shape = ( batch_size, num_channels_latents, num_latent_frames, int(height) // self.vae_scale_factor_spatial, int(width) // self.vae_scale_factor_spatial, ) if isinstance(generator, list) and len(generator) != batch_size: raise ValueError(f"Generator list length {len(generator)} does not match batch size {batch_size}.") latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) return latents def prepare_image_latents( self, image: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, num_latent_frames_per_chunk: int, dtype: torch.dtype | None = None, device: torch.device | None = None, generator: torch.Generator | list[torch.Generator] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Encode a single image into VAE latent space for I2V generation. Returns (image_latents, fake_image_latents) where fake_image_latents is the last-frame latent of a repeated-frame video, used to seed the history buffer for the first denoising chunk. """ device = device or self.device image = image.unsqueeze(2).to(device=device, dtype=self.vae.dtype) latents = self.vae.encode(image).latent_dist.sample(generator=generator) latents = (latents - latents_mean) * latents_std min_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1 fake_video = image.repeat(1, 1, min_frames, 1, 1) fake_latents_full = self.vae.encode(fake_video).latent_dist.sample(generator=generator) fake_latents_full = (fake_latents_full - latents_mean) * latents_std fake_latents = fake_latents_full[:, :, -1:, :, :] return latents.to(device=device, dtype=dtype), fake_latents.to(device=device, dtype=dtype) def prepare_video_latents( self, video: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, num_latent_frames_per_chunk: int, dtype: torch.dtype | None = None, device: torch.device | None = None, generator: torch.Generator | list[torch.Generator] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Encode a video into VAE latent space for V2V generation. Returns (first_frame_latent, video_latents) where first_frame_latent is used as the image prefix, and video_latents fills the history buffer. """ device = device or self.device video = video.to(device=device, dtype=self.vae.dtype) num_frames = video.shape[2] min_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1 num_chunks = num_frames // min_frames if num_chunks == 0: raise ValueError( f"Video must have at least {min_frames} frames (got {num_frames}). " f"Required: (num_latent_frames_per_chunk - 1) * {self.vae_scale_factor_temporal} + 1" ) total_valid_frames = num_chunks * min_frames start_frame = num_frames - total_valid_frames first_frame = video[:, :, 0:1, :, :] first_frame_latent = self.vae.encode(first_frame).latent_dist.sample(generator=generator) first_frame_latent = (first_frame_latent - latents_mean) * latents_std latents_chunks = [] for i in range(num_chunks): chunk_start = start_frame + i * min_frames chunk_end = chunk_start + min_frames video_chunk = video[:, :, chunk_start:chunk_end, :, :] chunk_latents = self.vae.encode(video_chunk).latent_dist.sample(generator=generator) chunk_latents = (chunk_latents - latents_mean) * latents_std latents_chunks.append(chunk_latents) video_latents = torch.cat(latents_chunks, dim=2) return first_frame_latent.to(device=device, dtype=dtype), video_latents.to(device=device, dtype=dtype) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) return loader.load_weights(weights)