import re from typing import Optional, Tuple, Union import torch from diffusers.models import AutoencoderKL from videosys.core.pab_mgr import PABConfig, set_pab_manager from videosys.core.pipeline import VideoSysPipeline, VideoSysPipelineOutput from videosys.utils.utils import save_video from .datasets import get_image_size, get_num_frames from .inference_utils import ( append_generated, append_score_to_prompts, apply_mask_strategy, collect_references_batch, dframe_to_frame, extract_json_from_prompts, extract_prompts_loop, merge_prompt, prepare_multi_resolution_info, split_prompt, ) from .rflow import RFLOW from .stdit3 import STDiT3_XL_2 from .text_encoder import T5Encoder, text_preprocessing from .vae import OpenSoraVAE_V1_2 class OpenSoraPABConfig(PABConfig): def __init__( self, steps: int = 50, spatial_broadcast: bool = True, spatial_threshold: list = [450, 930], spatial_gap: int = 2, temporal_broadcast: bool = True, temporal_threshold: list = [450, 930], temporal_gap: int = 4, cross_broadcast: bool = True, cross_threshold: list = [450, 930], cross_gap: int = 6, diffusion_skip: bool = False, diffusion_timestep_respacing: list = None, diffusion_skip_timestep: list = None, mlp_skip: bool = True, mlp_spatial_skip_config: dict = { 676: {"block": [0, 1, 2, 3, 4], "skip_count": 2}, 788: {"block": [0, 1, 2, 3, 4], "skip_count": 2}, 864: {"block": [0, 1, 2, 3, 4], "skip_count": 2}, }, mlp_temporal_skip_config: dict = { 676: {"block": [0, 1, 2, 3, 4], "skip_count": 2}, 788: {"block": [0, 1, 2, 3, 4], "skip_count": 2}, 864: {"block": [0, 1, 2, 3, 4], "skip_count": 2}, }, ): super().__init__( steps=steps, spatial_broadcast=spatial_broadcast, spatial_threshold=spatial_threshold, spatial_gap=spatial_gap, temporal_broadcast=temporal_broadcast, temporal_threshold=temporal_threshold, temporal_gap=temporal_gap, cross_broadcast=cross_broadcast, cross_threshold=cross_threshold, cross_gap=cross_gap, diffusion_skip=diffusion_skip, diffusion_timestep_respacing=diffusion_timestep_respacing, diffusion_skip_timestep=diffusion_skip_timestep, mlp_skip=mlp_skip, mlp_spatial_skip_config=mlp_spatial_skip_config, mlp_temporal_skip_config=mlp_temporal_skip_config, ) class OpenSoraConfig: def __init__( self, world_size: int = 1, transformer: str = "hpcai-tech/OpenSora-STDiT-v3", vae: str = "hpcai-tech/OpenSora-VAE-v1.2", text_encoder: str = "DeepFloyd/t5-v1_1-xxl", # ======= scheduler ======= num_sampling_steps: int = 30, cfg_scale: float = 7.0, # ======= vae ======== tiling_size: int = 4, # ======= pab ======== enable_pab: bool = False, pab_config: PABConfig = OpenSoraPABConfig(), ): # ======= engine ======== self.world_size = world_size # ======= pipeline ======== self.pipeline_cls = OpenSoraPipeline self.transformer = transformer self.vae = vae self.text_encoder = text_encoder # ======= scheduler ======== self.num_sampling_steps = num_sampling_steps self.cfg_scale = cfg_scale # ======= vae ======== self.tiling_size = tiling_size # ======= pab ======== self.enable_pab = enable_pab self.pab_config = pab_config class OpenSoraPipeline(VideoSysPipeline): r""" Pipeline for text-to-image generation using PixArt-Alpha. This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.) Args: vae ([`AutoencoderKL`]): Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. text_encoder ([`T5EncoderModel`]): Frozen text-encoder. PixArt-Alpha uses [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the [t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant. tokenizer (`T5Tokenizer`): Tokenizer of class [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer). transformer ([`Transformer2DModel`]): A text conditioned `Transformer2DModel` to denoise the encoded image latents. scheduler ([`SchedulerMixin`]): A scheduler to be used in combination with `transformer` to denoise the encoded image latents. """ bad_punct_regex = re.compile( r"[" + "#®•©™&@·º½¾¿¡§~" + "\)" + "\(" + "\]" + "\[" + "\}" + "\{" + "\|" + "\\" + "\/" + "\*" + r"]{1,}" ) # noqa _optional_components = ["tokenizer", "text_encoder"] model_cpu_offload_seq = "text_encoder->transformer->vae" def __init__( self, config: OpenSoraConfig, text_encoder: Optional[T5Encoder] = None, vae: Optional[AutoencoderKL] = None, transformer: Optional[STDiT3_XL_2] = None, scheduler: Optional[RFLOW] = None, device: torch.device = torch.device("cuda"), dtype: torch.dtype = torch.bfloat16, ): super().__init__() self._config = config self._device = device self._dtype = dtype # initialize the model if not provided if text_encoder is None: text_encoder = T5Encoder( from_pretrained=config.text_encoder, model_max_length=300, device=device, dtype=dtype ) if vae is None: vae = OpenSoraVAE_V1_2( from_pretrained="hpcai-tech/OpenSora-VAE-v1.2", micro_frame_size=17, micro_batch_size=config.tiling_size, ).to(dtype) if transformer is None: transformer = STDiT3_XL_2( from_pretrained="hpcai-tech/OpenSora-STDiT-v3", qk_norm=True, enable_flash_attn=True, enable_layernorm_kernel=True, in_channels=vae.out_channels, caption_channels=text_encoder.output_dim, model_max_length=text_encoder.model_max_length, ).to(device, dtype) text_encoder.y_embedder = transformer.y_embedder if scheduler is None: scheduler = RFLOW( use_timestep_transform=True, num_sampling_steps=config.num_sampling_steps, cfg_scale=config.cfg_scale ) # pab if config.enable_pab: set_pab_manager(config.pab_config) # set eval and device self.set_eval_and_device(device, text_encoder, vae, transformer) self.register_modules(text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler) @torch.no_grad() def generate( self, prompt: str, resolution="480p", aspect_ratio="9:16", num_frames: int = 51, loop: int = 1, llm_refine: bool = False, negative_prompt: str = "", ms: Optional[str] = "", refs: Optional[str] = "", aes: float = 6.5, flow: Optional[float] = None, camera_motion: Optional[float] = None, condition_frame_length: int = 5, align: int = 5, condition_frame_edit: float = 0.0, return_dict: bool = True, verbose: bool = True, ) -> Union[VideoSysPipelineOutput, Tuple]: """ Function invoked when calling the pipeline for generation. Args: prompt (`str` or `List[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. resolution (`str`, *optional*, defaults to `"480p"`): The resolution of the generated video. aspect_ratio (`str`, *optional*, defaults to `"9:16"`): The aspect ratio of the generated video. num_frames (`int`, *optional*, defaults to 51): The number of frames to generate. negative_prompt (`str` or `List[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). num_inference_steps (`int`, *optional*, defaults to 100): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. timesteps (`List[int]`, *optional*): Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps` timesteps are used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to 7.0): Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). `guidance_scale` is defined as `w` of equation 2. of [Imagen Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. height (`int`, *optional*, defaults to self.unet.config.sample_size): The height in pixels of the generated image. width (`int`, *optional*, defaults to self.unet.config.sample_size): The width in pixels of the generated image. eta (`float`, *optional*, defaults to 0.0): Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to [`schedulers.DDIMScheduler`], will be ignored for others. generator (`torch.Generator` or `List[torch.Generator]`, *optional*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.FloatTensor`, *optional*): Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image generation. Can be used to tweak the same generation with different prompts. If not provided, a latents tensor will ge generated by sampling using the supplied random `generator`. prompt_embeds (`torch.FloatTensor`, *optional*): Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, text embeddings will be generated from `prompt` input argument. negative_prompt_embeds (`torch.FloatTensor`, *optional*): Pre-generated negative text embeddings. For PixArt-Alpha this negative prompt should be "". If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. output_type (`str`, *optional*, defaults to `"pil"`): The output format of the generate image. Choose between [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple. callback (`Callable`, *optional*): A function that will be called every `callback_steps` steps during inference. The function will be called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. callback_steps (`int`, *optional*, defaults to 1): The frequency at which the `callback` function will be called. If not specified, the callback will be called at every step. clean_caption (`bool`, *optional*, defaults to `True`): Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to be installed. If the dependencies are not installed, the embeddings will be created from the raw prompt. mask_feature (`bool` defaults to `True`): If set to `True`, the text embeddings will be masked. Examples: Returns: [`~pipelines.ImagePipelineOutput`] or `tuple`: If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is returned where the first element is a list with the generated images """ # == basic == fps = 24 image_size = get_image_size(resolution, aspect_ratio) num_frames = get_num_frames(num_frames) # == prepare batch prompts == batch_prompts = [prompt] ms = [ms] refs = [refs] # == get json from prompts == batch_prompts, refs, ms = extract_json_from_prompts(batch_prompts, refs, ms) # == get reference for condition == refs = collect_references_batch(refs, self.vae, image_size) # == multi-resolution info == model_args = prepare_multi_resolution_info( "OpenSora", len(batch_prompts), image_size, num_frames, fps, self._device, self._dtype ) # == process prompts step by step == # 0. split prompt # each element in the list is [prompt_segment_list, loop_idx_list] batched_prompt_segment_list = [] batched_loop_idx_list = [] for prompt in batch_prompts: prompt_segment_list, loop_idx_list = split_prompt(prompt) batched_prompt_segment_list.append(prompt_segment_list) batched_loop_idx_list.append(loop_idx_list) # 1. refine prompt by openai # if llm_refine: # only call openai API when # 1. seq parallel is not enabled # 2. seq parallel is enabled and the process is rank 0 # if not enable_sequence_parallelism or (enable_sequence_parallelism and coordinator.is_master()): # for idx, prompt_segment_list in enumerate(batched_prompt_segment_list): # batched_prompt_segment_list[idx] = refine_prompts_by_openai(prompt_segment_list) # # sync the prompt if using seq parallel # if enable_sequence_parallelism: # coordinator.block_all() # prompt_segment_length = [ # len(prompt_segment_list) for prompt_segment_list in batched_prompt_segment_list # ] # # flatten the prompt segment list # batched_prompt_segment_list = [ # prompt_segment # for prompt_segment_list in batched_prompt_segment_list # for prompt_segment in prompt_segment_list # ] # # create a list of size equal to world size # broadcast_obj_list = [batched_prompt_segment_list] * coordinator.world_size # dist.broadcast_object_list(broadcast_obj_list, 0) # # recover the prompt list # batched_prompt_segment_list = [] # segment_start_idx = 0 # all_prompts = broadcast_obj_list[0] # for num_segment in prompt_segment_length: # batched_prompt_segment_list.append( # all_prompts[segment_start_idx : segment_start_idx + num_segment] # ) # segment_start_idx += num_segment # 2. append score for idx, prompt_segment_list in enumerate(batched_prompt_segment_list): batched_prompt_segment_list[idx] = append_score_to_prompts( prompt_segment_list, aes=aes, flow=flow, camera_motion=camera_motion, ) # 3. clean prompt with T5 for idx, prompt_segment_list in enumerate(batched_prompt_segment_list): batched_prompt_segment_list[idx] = [text_preprocessing(prompt) for prompt in prompt_segment_list] # 4. merge to obtain the final prompt batch_prompts = [] for prompt_segment_list, loop_idx_list in zip(batched_prompt_segment_list, batched_loop_idx_list): batch_prompts.append(merge_prompt(prompt_segment_list, loop_idx_list)) # == Iter over loop generation == video_clips = [] for loop_i in range(loop): # == get prompt for loop i == batch_prompts_loop = extract_prompts_loop(batch_prompts, loop_i) # == add condition frames for loop == if loop_i > 0: refs, ms = append_generated( self.vae, video_clips[-1], refs, ms, loop_i, condition_frame_length, condition_frame_edit ) # == sampling == input_size = (num_frames, *image_size) latent_size = self.vae.get_latent_size(input_size) z = torch.randn( len(batch_prompts), self.vae.out_channels, *latent_size, device=self._device, dtype=self._dtype ) masks = apply_mask_strategy(z, refs, ms, loop_i, align=align) samples = self.scheduler.sample( self.transformer, self.text_encoder, z=z, prompts=batch_prompts_loop, device=self._device, additional_args=model_args, progress=verbose, mask=masks, ) samples = self.vae.decode(samples.to(self._dtype), num_frames=num_frames) video_clips.append(samples) for i in range(1, loop): video_clips[i] = video_clips[i][:, dframe_to_frame(condition_frame_length) :] video = torch.cat(video_clips, dim=1) low, high = -1, 1 video.clamp_(min=low, max=high) video.sub_(low).div_(max(high - low, 1e-5)) video = video.mul(255).add_(0.5).clamp_(0, 255).permute(0, 2, 3, 4, 1).to("cpu", torch.uint8) # Offload all models self.maybe_free_model_hooks() if not return_dict: return (video,) return VideoSysPipelineOutput(video=video) def save_video(self, video, output_path): save_video(video, output_path, fps=24)