Upload 6 files

CogVideoX/README.md

@@ -0,0 +1,19 @@
+## Getting Started
+
+Download the corresponding pre-trained LoRA weights from the [model zoo](https://github.com/lwang592/TransPixar/README.md##model-zoo).
+
+## Inference
+
+To generate the RGBA video, run:
+
+```bash
+python cli.py \
+    --lora_path /path/to/lora \
+    --prompt "..." \
+```
+
+This command generates the RGB and Alpha videos simultaneously and saves them. Specifically, the RGB video is saved in its premultiplied form. To blend this video with any background image, you can simply use the following formula:
+
+```python
+com = rgb + (1 - alpha) * bgr
+```

CogVideoX/cli.py

@@ -0,0 +1,90 @@
+import os
+import torch
+from diffusers import CogVideoXDPMScheduler
+from pipeline_rgba import CogVideoXPipeline
+from diffusers.utils import export_to_video
+import argparse
+import numpy as np
+from rgba_utils import *
+
+def main(args):
+    # 1. load pipeline  
+    pipe = CogVideoXPipeline.from_pretrained(args.model_path, torch_dtype=torch.bfloat16)
+    pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
+    pipe.enable_sequential_cpu_offload()
+    pipe.vae.enable_slicing()
+    pipe.vae.enable_tiling()
+
+
+    # 2. define prompt and arguments
+    pipeline_args = {
+        "prompt": args.prompt,
+        "guidance_scale": args.guidance_scale,
+        "num_inference_steps": args.num_inference_steps,
+        "height": args.height,
+        "width": args.width,
+        "num_frames": args.num_frames,
+        "output_type": "latent",
+        "use_dynamic_cfg":True,
+    }
+
+    # 3. prepare rgbx utils
+    # breakpoint()
+    seq_length = 2 * (
+        (args.height // pipe.vae_scale_factor_spatial // 2)
+        * (args.width // pipe.vae_scale_factor_spatial // 2)
+        * ((args.num_frames - 1) // pipe.vae_scale_factor_temporal + 1)
+    )
+    # seq_length = 35100
+
+    prepare_for_rgba_inference(
+        pipe.transformer,
+        rgba_weights_path=args.lora_path,
+        device="cuda",
+        dtype=torch.bfloat16,
+        text_length=226,
+        seq_length=seq_length, # this is for the creation of attention mask.
+    )
+
+    # 4. run inference
+    generator = torch.manual_seed(args.seed) if args.seed else None
+    frames_latents = pipe(**pipeline_args, generator=generator).frames
+
+    frames_latents_rgb, frames_latents_alpha = frames_latents.chunk(2, dim=1)
+
+    frames_rgb = decode_latents(pipe, frames_latents_rgb)
+    frames_alpha = decode_latents(pipe, frames_latents_alpha)
+
+
+    pooled_alpha = np.max(frames_alpha, axis=-1, keepdims=True)
+    frames_alpha_pooled = np.repeat(pooled_alpha, 3, axis=-1)
+    premultiplied_rgb = frames_rgb * frames_alpha_pooled
+
+    if os.path.exists(args.output_path) == False:
+        os.makedirs(args.output_path)
+
+    export_to_video(premultiplied_rgb[0], os.path.join(args.output_path, "rgb.mp4"), fps=args.fps)
+    export_to_video(frames_alpha[0], os.path.join(args.output_path, "alpha.mp4"), fps=args.fps)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Generate a video from a text prompt")
+    parser.add_argument("--prompt", type=str, required=True, help="The description of the video to be generated")
+    parser.add_argument("--lora_path", type=str, default="/hpc2hdd/home/lwang592/projects/CogVideo/sat/outputs/training/ckpts-5b-attn_rebias-partial_lora-8gpu-wo_t2a/lora-rgba-12-21-19-11/5000/rgba_lora.safetensors", help="The path of the LoRA weights to be used")
+    
+    parser.add_argument(
+        "--model_path", type=str, default="THUDM/CogVideoX-5B", help="Path of the pre-trained model use"
+    )
+
+
+    parser.add_argument("--output_path", type=str, default="./output", help="The path save generated video")
+    parser.add_argument("--guidance_scale", type=float, default=6.0, help="The scale for classifier-free guidance")
+    parser.add_argument("--num_inference_steps", type=int, default=50, help="Inference steps")
+    parser.add_argument("--num_frames", type=int, default=49, help="Number of steps for the inference process")
+    parser.add_argument("--width", type=int, default=720, help="Number of steps for the inference process")
+    parser.add_argument("--height", type=int, default=480, help="Number of steps for the inference process")
+    parser.add_argument("--fps", type=int, default=8, help="Number of steps for the inference process")
+    parser.add_argument("--seed", type=int, default=None, help="The seed for reproducibility")
+    args = parser.parse_args()
+
+    main(args)

CogVideoX/pipeline_rgba.py

@@ -0,0 +1,744 @@
+# Copyright 2024 The CogVideoX team, Tsinghua University & ZhipuAI and The HuggingFace Team.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import math
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+from transformers import T5EncoderModel, T5Tokenizer
+
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.loaders import CogVideoXLoraLoaderMixin
+from diffusers.models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel
+from diffusers.models.embeddings import get_3d_rotary_pos_embed
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler
+from diffusers.utils import logging, replace_example_docstring
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.video_processor import VideoProcessor
+from diffusers.pipelines.cogvideo.pipeline_output import CogVideoXPipelineOutput
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```python
+        >>> import torch
+        >>> from diffusers import CogVideoXPipeline
+        >>> from diffusers.utils import export_to_video
+
+        >>> # Models: "THUDM/CogVideoX-2b" or "THUDM/CogVideoX-5b"
+        >>> pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-2b", torch_dtype=torch.float16).to("cuda")
+        >>> prompt = (
+        ...     "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. "
+        ...     "The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other "
+        ...     "pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, "
+        ...     "casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. "
+        ...     "The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical "
+        ...     "atmosphere of this unique musical performance."
+        ... )
+        >>> video = pipe(prompt=prompt, guidance_scale=6, num_inference_steps=50).frames[0]
+        >>> export_to_video(video, "output.mp4", fps=8)
+        ```
+"""
+
+
+# Similar to diffusers.pipelines.hunyuandit.pipeline_hunyuandit.get_resize_crop_region_for_grid
+def get_resize_crop_region_for_grid(src, tgt_width, tgt_height):
+    tw = tgt_width
+    th = tgt_height
+    h, w = src
+    r = h / w
+    if r > (th / tw):
+        resize_height = th
+        resize_width = int(round(th / h * w))
+    else:
+        resize_width = tw
+        resize_height = int(round(tw / w * h))
+
+    crop_top = int(round((th - resize_height) / 2.0))
+    crop_left = int(round((tw - resize_width) / 2.0))
+
+    return (crop_top, crop_left), (crop_top + resize_height, crop_left + resize_width)
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class CogVideoXPipeline(DiffusionPipeline, CogVideoXLoraLoaderMixin):
+    r"""
+    Pipeline for text-to-video generation using CogVideoX.
+
+    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 videos to and from latent representations.
+        text_encoder ([`T5EncoderModel`]):
+            Frozen text-encoder. CogVideoX 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 ([`CogVideoXTransformer3DModel`]):
+            A text conditioned `CogVideoXTransformer3DModel` to denoise the encoded video latents.
+        scheduler ([`SchedulerMixin`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded video latents.
+    """
+
+    _optional_components = []
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+
+    _callback_tensor_inputs = [
+        "latents",
+        "prompt_embeds",
+        "negative_prompt_embeds",
+    ]
+
+    def __init__(
+        self,
+        tokenizer: T5Tokenizer,
+        text_encoder: T5EncoderModel,
+        vae: AutoencoderKLCogVideoX,
+        transformer: CogVideoXTransformer3DModel,
+        scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler],
+    ):
+        super().__init__()
+
+        self.register_modules(
+            tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler
+        )
+        self.vae_scale_factor_spatial = (
+            2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8
+        )
+        self.vae_scale_factor_temporal = (
+            self.vae.config.temporal_compression_ratio if hasattr(self, "vae") and self.vae is not None else 4
+        )
+        self.vae_scaling_factor_image = (
+            self.vae.config.scaling_factor if hasattr(self, "vae") and self.vae is not None else 0.7
+        )
+
+        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
+
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]] = None,
+        num_videos_per_prompt: int = 1,
+        max_sequence_length: int = 226,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because `max_sequence_length` is set to "
+                f" {max_sequence_length} tokens: {removed_text}"
+            )
+
+        prompt_embeds = self.text_encoder(text_input_ids.to(device))[0]
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        _, 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)
+
+        return prompt_embeds
+
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        do_classifier_free_guidance: bool = True,
+        num_videos_per_prompt: int = 1,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        max_sequence_length: int = 226,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            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`).
+            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
+                Whether to use classifier free guidance or not.
+            num_videos_per_prompt (`int`, *optional*, defaults to 1):
+                Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
+            prompt_embeds (`torch.Tensor`, *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.Tensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            device: (`torch.device`, *optional*):
+                torch device
+            dtype: (`torch.dtype`, *optional*):
+                torch dtype
+        """
+        device = device or self._execution_device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_embeds = self._get_t5_prompt_embeds(
+                prompt=prompt,
+                num_videos_per_prompt=num_videos_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+                dtype=dtype,
+            )
+
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+            if prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+
+            negative_prompt_embeds = self._get_t5_prompt_embeds(
+                prompt=negative_prompt,
+                num_videos_per_prompt=num_videos_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+                dtype=dtype,
+            )
+
+        return prompt_embeds, negative_prompt_embeds
+
+    def prepare_latents(
+        self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None
+    ):
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        shape = (
+            batch_size,
+            (num_frames - 1) // self.vae_scale_factor_temporal + 1,
+            num_channels_latents,
+            height // self.vae_scale_factor_spatial,
+            width // self.vae_scale_factor_spatial,
+        )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    def decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        latents = latents.permute(0, 2, 1, 3, 4)  # [batch_size, num_channels, num_frames, height, width]
+        latents = 1 / self.vae_scaling_factor_image * latents
+
+        frames = self.vae.decode(latents).sample
+        return frames
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    # Copied from diffusers.pipelines.latte.pipeline_latte.LattePipeline.check_inputs
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        negative_prompt,
+        callback_on_step_end_tensor_inputs,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+    ):
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+    def fuse_qkv_projections(self) -> None:
+        r"""Enables fused QKV projections."""
+        self.fusing_transformer = True
+        self.transformer.fuse_qkv_projections()
+
+    def unfuse_qkv_projections(self) -> None:
+        r"""Disable QKV projection fusion if enabled."""
+        if not self.fusing_transformer:
+            logger.warning("The Transformer was not initially fused for QKV projections. Doing nothing.")
+        else:
+            self.transformer.unfuse_qkv_projections()
+            self.fusing_transformer = False
+
+    def _prepare_rotary_positional_embeddings(
+        self,
+        height: int,
+        width: int,
+        num_frames: int,
+        device: torch.device,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        grid_height = height // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
+        grid_width = width // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
+        base_size_width = 720 // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
+        base_size_height = 480 // (self.vae_scale_factor_spatial * self.transformer.config.patch_size)
+
+        grid_crops_coords = get_resize_crop_region_for_grid(
+            (grid_height, grid_width), base_size_width, base_size_height
+        )
+        freqs_cos, freqs_sin = get_3d_rotary_pos_embed(
+            embed_dim=self.transformer.config.attention_head_dim,
+            crops_coords=grid_crops_coords,
+            grid_size=(grid_height, grid_width),
+            temporal_size=num_frames,
+        )
+
+        freqs_cos = freqs_cos.to(device=device)
+        freqs_sin = freqs_sin.to(device=device)
+        return freqs_cos, freqs_sin
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def attention_kwargs(self):
+        return self._attention_kwargs
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        height: int = 480,
+        width: int = 720,
+        num_frames: int = 49,
+        num_inference_steps: int = 50,
+        timesteps: Optional[List[int]] = None,
+        guidance_scale: float = 6,
+        use_dynamic_cfg: bool = False,
+        num_videos_per_prompt: int = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[
+            Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+        ] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 226,
+    ) -> Union[CogVideoXPipelineOutput, 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.
+            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`).
+            height (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial):
+                The height in pixels of the generated image. This is set to 480 by default for the best results.
+            width (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial):
+                The width in pixels of the generated image. This is set to 720 by default for the best results.
+            num_frames (`int`, defaults to `48`):
+                Number of frames to generate. Must be divisible by self.vae_scale_factor_temporal. Generated video will
+                contain 1 extra frame because CogVideoX is conditioned with (num_seconds * fps + 1) frames where
+                num_seconds is 6 and fps is 8. However, since videos can be saved at any fps, the only condition that
+                needs to be satisfied is that of divisibility mentioned above.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                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 with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be 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_videos_per_prompt (`int`, *optional*, defaults to 1):
+                The number of videos to generate per prompt.
+            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. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. 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_xl.StableDiffusionXLPipelineOutput`] instead
+                of a plain tuple.
+            attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int`, defaults to `226`):
+                Maximum sequence length in encoded prompt. Must be consistent with
+                `self.transformer.config.max_text_seq_length` otherwise may lead to poor results.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipelineOutput`] or `tuple`:
+            [`~pipelines.cogvideo.pipeline_cogvideox.CogVideoXPipelineOutput`] if `return_dict` is True, otherwise a
+            `tuple`. When returning a tuple, the first element is a list with the generated images.
+        """
+
+        if num_frames > 49:
+            raise ValueError(
+                "The number of frames must be less than 49 for now due to static positional embeddings. This will be updated in the future to remove this limitation."
+            )
+
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        num_videos_per_prompt = 1
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            negative_prompt,
+            callback_on_step_end_tensor_inputs,
+            prompt_embeds,
+            negative_prompt_embeds,
+        )
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._interrupt = False
+
+        # 2. Default call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt,
+            negative_prompt,
+            do_classifier_free_guidance,
+            num_videos_per_prompt=num_videos_per_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            max_sequence_length=max_sequence_length,
+            device=device,
+        )
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+
+        # 4. Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+        self._num_timesteps = len(timesteps)
+
+        # 5. Prepare latents.
+        latent_channels = self.transformer.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_videos_per_prompt,
+            latent_channels,
+            num_frames,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        ).repeat(1,2,1,1,1) # Luozhou
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Create rotary embeds if required
+        image_rotary_emb = (
+            self._prepare_rotary_positional_embeddings(height, width, latents.size(1) // 2, device) # Luozhou
+            if self.transformer.config.use_rotary_positional_embeddings
+            else None
+        )
+
+        # 8. Denoising loop
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            # for DPM-solver++
+            old_pred_original_sample = None
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latent_model_input.shape[0])
+
+                # predict noise model_output
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    encoder_hidden_states=prompt_embeds,
+                    timestep=timestep,
+                    image_rotary_emb=image_rotary_emb,
+                    attention_kwargs=attention_kwargs,
+                    return_dict=False,
+                )[0]
+                noise_pred = noise_pred.float()
+
+                # perform guidance
+                if use_dynamic_cfg:
+                    self._guidance_scale = 1 + guidance_scale * (
+                        (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2
+                    )
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                if not isinstance(self.scheduler, CogVideoXDPMScheduler):
+                    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                else:
+                    latents, old_pred_original_sample = self.scheduler.step(
+                        noise_pred,
+                        old_pred_original_sample,
+                        t,
+                        timesteps[i - 1] if i > 0 else None,
+                        latents,
+                        **extra_step_kwargs,
+                        return_dict=False,
+                    )
+                latents = latents.to(prompt_embeds.dtype)
+
+                # call the callback, if provided
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+        if not output_type == "latent":
+            video = self.decode_latents(latents)
+            video = self.video_processor.postprocess_video(video=video, output_type=output_type)
+        else:
+            video = latents
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (video,)
+
+        return CogVideoXPipelineOutput(frames=video)

CogVideoX/rgba_utils.py

@@ -0,0 +1,313 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Any, Dict, Optional, Tuple, Union
+from diffusers.utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from safetensors.torch import load_file
+
+logger = logging.get_logger(__name__)
+
+@torch.no_grad()
+def decode_latents(pipe, latents):
+    video = pipe.decode_latents(latents)
+    video = pipe.video_processor.postprocess_video(video=video, output_type="np")
+    return video
+
+def create_attention_mask(text_length: int, seq_length: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    """
+    Create an attention mask to block text from attending to alpha.
+
+    Args:
+        text_length: Length of the text sequence.
+        seq_length: Length of the other sequence.
+        device: The device where the mask will be stored.
+        dtype: The data type of the mask tensor.
+
+    Returns:
+        An attention mask tensor.
+    """
+    total_length = text_length + seq_length
+    dense_mask = torch.ones((total_length, total_length), dtype=torch.bool)
+    dense_mask[:text_length, text_length + seq_length // 2:] = False
+    return dense_mask.to(device=device, dtype=dtype)
+
+class RGBALoRACogVideoXAttnProcessor:
+    r"""
+    Processor for implementing scaled dot-product attention for the CogVideoX model. 
+    It applies a rotary embedding on query and key vectors, but does not include spatial normalization.
+    """
+
+    def __init__(self, device, dtype, attention_mask, lora_rank=128, lora_alpha=1.0, latent_dim=3072):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("CogVideoXAttnProcessor requires PyTorch 2.0 or later.")
+
+        # Initialize LoRA layers
+        self.lora_alpha = lora_alpha
+        self.lora_rank = lora_rank
+
+        # Helper function to create LoRA layers
+        def create_lora_layer(in_dim, mid_dim, out_dim):
+            return nn.Sequential(
+                nn.Linear(in_dim, mid_dim, bias=False, device=device, dtype=dtype),
+                nn.Linear(mid_dim, out_dim, bias=False, device=device, dtype=dtype)
+            )
+
+        self.to_q_lora = create_lora_layer(latent_dim, lora_rank, latent_dim)
+        self.to_k_lora = create_lora_layer(latent_dim, lora_rank, latent_dim)
+        self.to_v_lora = create_lora_layer(latent_dim, lora_rank, latent_dim)
+        self.to_out_lora = create_lora_layer(latent_dim, lora_rank, latent_dim)
+
+        # Store attention mask
+        self.attention_mask = attention_mask
+
+    def _apply_lora(self, hidden_states, seq_len, query, key, value, scaling):
+        """Applies LoRA updates to query, key, and value tensors."""
+        query_delta = self.to_q_lora(hidden_states).to(query.device)
+        query[:, -seq_len // 2:, :] += query_delta[:, -seq_len // 2:, :] * scaling
+
+        key_delta = self.to_k_lora(hidden_states).to(key.device)
+        key[:, -seq_len // 2:, :] += key_delta[:, -seq_len // 2:, :] * scaling
+
+        value_delta = self.to_v_lora(hidden_states).to(value.device)
+        value[:, -seq_len // 2:, :] += value_delta[:, -seq_len // 2:, :] * scaling
+
+        return query, key, value
+
+    def _apply_rotary_embedding(self, query, key, image_rotary_emb, seq_len, text_seq_length, attn):
+        """Applies rotary embeddings to query and key tensors."""
+        from diffusers.models.embeddings import apply_rotary_emb
+
+        # Apply rotary embedding to RGB and alpha sections
+        query[:, :, text_seq_length:text_seq_length + seq_len // 2] = apply_rotary_emb(
+            query[:, :, text_seq_length:text_seq_length + seq_len // 2], image_rotary_emb)
+        query[:, :, text_seq_length + seq_len // 2:] = apply_rotary_emb(
+            query[:, :, text_seq_length + seq_len // 2:], image_rotary_emb)
+        
+        if not attn.is_cross_attention:
+            key[:, :, text_seq_length:text_seq_length + seq_len // 2] = apply_rotary_emb(
+                key[:, :, text_seq_length:text_seq_length + seq_len // 2], image_rotary_emb)
+            key[:, :, text_seq_length + seq_len // 2:] = apply_rotary_emb(
+                key[:, :, text_seq_length + seq_len // 2:], image_rotary_emb)
+
+        return query, key
+
+    def __call__(
+        self,
+        attn,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        # Concatenate encoder and decoder hidden states
+        text_seq_length = encoder_hidden_states.size(1)
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        batch_size, sequence_length, _ = hidden_states.shape
+        seq_len = hidden_states.shape[1] - text_seq_length
+        scaling = self.lora_alpha / self.lora_rank
+
+        # Apply LoRA to query, key, value
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        query, key, value = self._apply_lora(hidden_states, seq_len, query, key, value, scaling)
+
+        # Reshape query, key, value for multi-head attention
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # Normalize query and key if required
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # Apply rotary embeddings if provided
+        if image_rotary_emb is not None:
+            query, key = self._apply_rotary_embedding(query, key, image_rotary_emb, seq_len, text_seq_length, attn)
+
+        # Compute scaled dot-product attention
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=self.attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        # Reshape the output tensor back to the original shape
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+
+        # Apply linear projection and LoRA to the output
+        original_hidden_states = attn.to_out[0](hidden_states)
+        hidden_states_delta = self.to_out_lora(hidden_states).to(hidden_states.device)
+        original_hidden_states[:, -seq_len // 2:, :] += hidden_states_delta[:, -seq_len // 2:, :] * scaling
+
+        # Apply dropout
+        hidden_states = attn.to_out[1](original_hidden_states)
+
+        # Split back into encoder and decoder hidden states
+        encoder_hidden_states, hidden_states = hidden_states.split(
+            [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
+        )
+
+        return hidden_states, encoder_hidden_states
+
+def prepare_for_rgba_inference(
+    model, rgba_weights_path: str, device: torch.device, dtype: torch.dtype,
+    lora_rank: int = 128, lora_alpha: float = 1.0, text_length: int = 226, seq_length: int = 35100
+):
+    def load_lora_sequential_weights(lora_layer, lora_layers, prefix):
+        lora_layer[0].load_state_dict({'weight': lora_layers[f"{prefix}.lora_A.weight"]})
+        lora_layer[1].load_state_dict({'weight': lora_layers[f"{prefix}.lora_B.weight"]})
+
+
+    rgba_weights = load_file(rgba_weights_path)
+    aux_emb = rgba_weights['domain_emb']
+
+    attention_mask = create_attention_mask(text_length, seq_length, device, dtype)
+    attn_procs = {}
+
+    for name in model.attn_processors.keys():
+        attn_processor = RGBALoRACogVideoXAttnProcessor(
+            device=device, dtype=dtype, attention_mask=attention_mask,
+            lora_rank=lora_rank, lora_alpha=lora_alpha
+        )
+        
+        index = name.split('.')[1]
+        base_prefix = f'transformer.transformer_blocks.{index}.attn1'
+
+        for lora_layer, prefix in [
+            (attn_processor.to_q_lora, f'{base_prefix}.to_q'),
+            (attn_processor.to_k_lora, f'{base_prefix}.to_k'),
+            (attn_processor.to_v_lora, f'{base_prefix}.to_v'),
+            (attn_processor.to_out_lora, f'{base_prefix}.to_out.0'),
+        ]:
+            load_lora_sequential_weights(lora_layer, rgba_weights, prefix)
+
+        attn_procs[name] = attn_processor
+
+    model.set_attn_processor(attn_procs)
+
+    def custom_forward(self):
+        def forward(
+            hidden_states: torch.Tensor,
+            encoder_hidden_states: torch.Tensor,
+            timestep: Union[int, float, torch.LongTensor],
+            timestep_cond: Optional[torch.Tensor] = None,
+            ofs: Optional[Union[int, float, torch.LongTensor]] = None,
+            image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+            attention_kwargs: Optional[Dict[str, Any]] = None,
+            return_dict: bool = True,
+        ):
+            if attention_kwargs is not None:
+                attention_kwargs = attention_kwargs.copy()
+                lora_scale = attention_kwargs.pop("scale", 1.0)
+            else:
+                lora_scale = 1.0
+
+            if USE_PEFT_BACKEND:
+                # weight the lora layers by setting `lora_scale` for each PEFT layer
+                scale_lora_layers(self, lora_scale)
+            else:
+                if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
+                    logger.warning(
+                        "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
+                    )
+
+            batch_size, num_frames, channels, height, width = hidden_states.shape
+
+            # 1. Time embedding
+            timesteps = timestep
+            t_emb = self.time_proj(timesteps)
+
+            # timesteps does not contain any weights and will always return f32 tensors
+            # but time_embedding might actually be running in fp16. so we need to cast here.
+            # there might be better ways to encapsulate this.
+            t_emb = t_emb.to(dtype=hidden_states.dtype)
+            emb = self.time_embedding(t_emb, timestep_cond)
+
+            if self.ofs_embedding is not None:
+                ofs_emb = self.ofs_proj(ofs)
+                ofs_emb = ofs_emb.to(dtype=hidden_states.dtype)
+                ofs_emb = self.ofs_embedding(ofs_emb)
+                emb = emb + ofs_emb
+
+            # 2. Patch embedding
+            hidden_states = self.patch_embed(encoder_hidden_states, hidden_states)
+            hidden_states = self.embedding_dropout(hidden_states)
+
+            text_seq_length = encoder_hidden_states.shape[1]
+            encoder_hidden_states = hidden_states[:, :text_seq_length]
+            hidden_states = hidden_states[:, text_seq_length:]
+
+            hidden_states[:, hidden_states.size(1) // 2:, :] += aux_emb.expand(batch_size, -1, -1).to(hidden_states.device, dtype=hidden_states.dtype)
+
+            # 3. Transformer blocks
+            for i, block in enumerate(self.transformer_blocks):
+                if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                    def create_custom_forward(module):
+                        def custom_forward(*inputs):
+                            return module(*inputs)
+
+                        return custom_forward
+
+                    ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                    hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(block),
+                        hidden_states,
+                        encoder_hidden_states,
+                        emb,
+                        image_rotary_emb,
+                        **ckpt_kwargs,
+                    )
+                else:
+                    hidden_states, encoder_hidden_states = block(
+                        hidden_states=hidden_states,
+                        encoder_hidden_states=encoder_hidden_states,
+                        temb=emb,
+                        image_rotary_emb=image_rotary_emb,
+                    )
+
+            if not self.config.use_rotary_positional_embeddings:
+                # CogVideoX-2B
+                hidden_states = self.norm_final(hidden_states)
+            else:
+                # CogVideoX-5B
+                hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+                hidden_states = self.norm_final(hidden_states)
+                hidden_states = hidden_states[:, text_seq_length:]
+
+            # 4. Final block
+            hidden_states = self.norm_out(hidden_states, temb=emb)
+            hidden_states = self.proj_out(hidden_states)
+
+            # 5. Unpatchify
+            p = self.config.patch_size
+            p_t = self.config.patch_size_t
+
+            if p_t is None:
+                output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, -1, p, p)
+                output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4)
+            else:
+                output = hidden_states.reshape(
+                    batch_size, (num_frames + p_t - 1) // p_t, height // p, width // p, -1, p_t, p, p
+                )
+                output = output.permute(0, 1, 5, 4, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(1, 2)
+
+            if USE_PEFT_BACKEND:
+                # remove `lora_scale` from each PEFT layer
+                unscale_lora_layers(self, lora_scale)
+
+            if not return_dict:
+                return (output,)
+            return Transformer2DModelOutput(sample=output)
+
+
+        return forward
+
+    model.forward = custom_forward(model)
+

CogVideoX/test.py

@@ -0,0 +1,26 @@
+import torch
+from diffusers import CogVideoXImageToVideoPipeline
+from diffusers.utils import export_to_video, load_image
+
+prompt = "A dragon is flipping its wings"
+image = load_image(image="/hpc2hdd/home/lwang592/projects/CogVideo/sat/configs/i2v/Dragon.jpg")
+pipe = CogVideoXImageToVideoPipeline.from_pretrained(
+    "THUDM/CogVideoX-5b-I2V",
+    torch_dtype=torch.bfloat16
+)
+
+pipe.enable_sequential_cpu_offload()
+pipe.vae.enable_tiling()
+pipe.vae.enable_slicing()
+
+video = pipe(
+    prompt=prompt,
+    image=image,
+    num_videos_per_prompt=1,
+    num_inference_steps=50,
+    num_frames=13,
+    guidance_scale=6,
+    generator=torch.Generator(device="cuda").manual_seed(42),
+).frames[0]
+
+export_to_video(video, "output.mp4", fps=8)

app.py

@@ -0,0 +1,323 @@
+"""
+THis is the main file for the gradio web demo. It uses the CogVideoX-5B model to generate videos gradio web demo.
+set environment variable OPENAI_API_KEY to use the OpenAI API to enhance the prompt.
+Usage:
+    OpenAI_API_KEY=your_openai_api_key OPENAI_BASE_URL=https://api.openai.com/v1 python inference/gradio_web_demo.py
+"""
+
+import math
+import os
+import random
+import threading
+import time
+
+import cv2
+import tempfile
+import imageio_ffmpeg
+import gradio as gr
+import torch
+from PIL import Image
+# from diffusers import (
+#     CogVideoXPipeline,
+#     CogVideoXDPMScheduler,
+#     CogVideoXVideoToVideoPipeline,
+#     CogVideoXImageToVideoPipeline,
+#     CogVideoXTransformer3DModel,
+# )
+from typing import Union, List
+from CogVideoX.pipeline_rgba import CogVideoXPipeline
+from CogVideoX.rgba_utils import *
+from diffusers import CogVideoXDPMScheduler
+
+from diffusers.utils import load_video, load_image, export_to_video
+from datetime import datetime, timedelta
+
+from diffusers.image_processor import VaeImageProcessor
+import moviepy.editor as mp
+import numpy as np
+from huggingface_hub import hf_hub_download, snapshot_download
+import gc
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# hf_hub_download(repo_id="ai-forever/Real-ESRGAN", filename="RealESRGAN_x4.pth", local_dir="model_real_esran")
+hf_hub_download(repo_id="wileewang/TransPixar", filename="cogvideox_rgba_lora.safetensors", local_dir="model_cogvideox_rgba_lora")
+# snapshot_download(repo_id="AlexWortega/RIFE", local_dir="model_rife")
+
+pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5B", torch_dtype=torch.bfloat16)
+# pipe.enable_sequential_cpu_offload()
+pipe.vae.enable_slicing()
+pipe.vae.enable_tiling()
+pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
+seq_length = 2 * (
+    (480 // pipe.vae_scale_factor_spatial // 2)
+    * (720 // pipe.vae_scale_factor_spatial // 2)
+    * ((13 - 1) // pipe.vae_scale_factor_temporal + 1)
+)
+prepare_for_rgba_inference(
+    pipe.transformer,
+    rgba_weights_path="model_cogvideox_rgba_lora/cogvideox_rgba_lora.safetensors",
+    device="cuda",
+    dtype=torch.bfloat16,
+    text_length=226,
+    seq_length=seq_length, # this is for the creation of attention mask.
+)
+
+# pipe.transformer.to(memory_format=torch.channels_last)
+# pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune", fullgraph=True)
+# pipe_image.transformer.to(memory_format=torch.channels_last)
+# pipe_image.transformer = torch.compile(pipe_image.transformer, mode="max-autotune", fullgraph=True)
+
+os.makedirs("./output", exist_ok=True)
+os.makedirs("./gradio_tmp", exist_ok=True)
+
+# upscale_model = utils.load_sd_upscale("model_real_esran/RealESRGAN_x4.pth", device)
+# frame_interpolation_model = load_rife_model("model_rife")
+
+
+sys_prompt = """You are part of a team of bots that creates videos. You work with an assistant bot that will draw anything you say in square brackets.
+For example , outputting " a beautiful morning in the woods with the sun peaking through the trees " will trigger your partner bot to output an video of a forest morning , as described. You will be prompted by people looking to create detailed , amazing videos. The way to accomplish this is to take their short prompts and make them extremely detailed and descriptive.
+There are a few rules to follow:
+You will only ever output a single video description per user request.
+When modifications are requested , you should not simply make the description longer . You should refactor the entire description to integrate the suggestions.
+Other times the user will not want modifications , but instead want a new image . In this case , you should ignore your previous conversation with the user.
+Video descriptions must have the same num of words as examples below. Extra words will be ignored.
+"""
+def save_video(tensor: Union[List[np.ndarray], List[Image.Image]], fps: int = 8, prefix='rgb'):
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    video_path = f"./output/{prefix}_{timestamp}.mp4"
+    os.makedirs(os.path.dirname(video_path), exist_ok=True)
+    export_to_video(tensor, video_path, fps=fps)
+    return video_path
+
+def resize_if_unfit(input_video, progress=gr.Progress(track_tqdm=True)):
+    width, height = get_video_dimensions(input_video)
+
+    if width == 720 and height == 480:
+        processed_video = input_video
+    else:
+        processed_video = center_crop_resize(input_video)
+    return processed_video
+
+
+def get_video_dimensions(input_video_path):
+    reader = imageio_ffmpeg.read_frames(input_video_path)
+    metadata = next(reader)
+    return metadata["size"]
+
+
+def center_crop_resize(input_video_path, target_width=720, target_height=480):
+    cap = cv2.VideoCapture(input_video_path)
+
+    orig_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    orig_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    orig_fps = cap.get(cv2.CAP_PROP_FPS)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    width_factor = target_width / orig_width
+    height_factor = target_height / orig_height
+    resize_factor = max(width_factor, height_factor)
+
+    inter_width = int(orig_width * resize_factor)
+    inter_height = int(orig_height * resize_factor)
+
+    target_fps = 8
+    ideal_skip = max(0, math.ceil(orig_fps / target_fps) - 1)
+    skip = min(5, ideal_skip)  # Cap at 5
+
+    while (total_frames / (skip + 1)) < 49 and skip > 0:
+        skip -= 1
+
+    processed_frames = []
+    frame_count = 0
+    total_read = 0
+
+    while frame_count < 49 and total_read < total_frames:
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        if total_read % (skip + 1) == 0:
+            resized = cv2.resize(frame, (inter_width, inter_height), interpolation=cv2.INTER_AREA)
+
+            start_x = (inter_width - target_width) // 2
+            start_y = (inter_height - target_height) // 2
+            cropped = resized[start_y : start_y + target_height, start_x : start_x + target_width]
+
+            processed_frames.append(cropped)
+            frame_count += 1
+
+        total_read += 1
+
+    cap.release()
+
+    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file:
+        temp_video_path = temp_file.name
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+        out = cv2.VideoWriter(temp_video_path, fourcc, target_fps, (target_width, target_height))
+
+        for frame in processed_frames:
+            out.write(frame)
+
+        out.release()
+
+    return temp_video_path
+
+
+
+def infer(
+    prompt: str,
+    num_inference_steps: int,
+    guidance_scale: float,
+    seed: int = -1,
+    progress=gr.Progress(track_tqdm=True),
+):
+    if seed == -1:
+        seed = random.randint(0, 2**8 - 1)
+    pipe.to(device)
+    video_pt = pipe(
+        prompt=prompt + ", isolated background",
+        num_videos_per_prompt=1,
+        num_inference_steps=num_inference_steps,
+        num_frames=13,
+        use_dynamic_cfg=True,
+        output_type="latent",
+        guidance_scale=guidance_scale,
+        generator=torch.Generator(device=device).manual_seed(int(seed)),
+    ).frames
+    # pipe.to("cpu")
+    gc.collect()
+    return (video_pt, seed)
+
+
+def convert_to_gif(video_path):
+    clip = mp.VideoFileClip(video_path)
+    clip = clip.set_fps(8)
+    clip = clip.resize(height=240)
+    gif_path = video_path.replace(".mp4", ".gif")
+    clip.write_gif(gif_path, fps=8)
+    return gif_path
+
+
+def delete_old_files():
+    while True:
+        now = datetime.now()
+        cutoff = now - timedelta(minutes=10)
+        directories = ["./output", "./gradio_tmp"]
+
+        for directory in directories:
+            for filename in os.listdir(directory):
+                file_path = os.path.join(directory, filename)
+                if os.path.isfile(file_path):
+                    file_mtime = datetime.fromtimestamp(os.path.getmtime(file_path))
+                    if file_mtime < cutoff:
+                        os.remove(file_path)
+        time.sleep(600)
+
+
+threading.Thread(target=delete_old_files, daemon=True).start()
+# examples_videos = [["example_videos/horse.mp4"], ["example_videos/kitten.mp4"], ["example_videos/train_running.mp4"]]
+# examples_images = [["example_images/beach.png"], ["example_images/street.png"], ["example_images/camping.png"]]
+
+with gr.Blocks() as demo:
+    gr.Markdown("""
+           <div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
+               TransPixar + CogVideoX-5B Huggingface Space🤗
+           </div>
+           <div style="text-align: center;">
+               <a href="https://huggingface.co/wileewang/TransPixar">🤗 TransPixar LoRA Hub</a> |
+               <a href="https://github.com/wileewang/TransPixar">🌐 Github</a> |
+               <a href="https://arxiv.org/pdf/2408.06072">📜 arxiv </a>
+           </div>
+           <div style="text-align: center; font-size: 15px; font-weight: bold; color: red; margin-bottom: 20px;">
+            ⚠️ This demo is for academic research and experiential use only. 
+            </div>
+           """)
+    with gr.Row():
+        with gr.Column():
+            # with gr.Accordion("I2V: Image Input (cannot be used simultaneously with video input)", open=False):
+            #     image_input = gr.Image(label="Input Image (will be cropped to 720 * 480)")
+            #     examples_component_images = gr.Examples(examples_images, inputs=[image_input], cache_examples=False)
+            # with gr.Accordion("V2V: Video Input (cannot be used simultaneously with image input)", open=False):
+            #     video_input = gr.Video(label="Input Video (will be cropped to 49 frames, 6 seconds at 8fps)")
+            #     strength = gr.Slider(0.1, 1.0, value=0.8, step=0.01, label="Strength")
+            #     examples_component_videos = gr.Examples(examples_videos, inputs=[video_input], cache_examples=False)
+            prompt = gr.Textbox(label="Prompt (Less than 200 Words)", placeholder="Enter your prompt here", lines=5)
+            with gr.Group():
+                with gr.Column():
+                    with gr.Row():
+                        seed_param = gr.Number(
+                            label="Inference Seed (Enter a positive number, -1 for random)", value=-1
+                        )
+                    # with gr.Row():
+                    #     enable_scale = gr.Checkbox(label="Super-Resolution (720 × 480 -> 2880 × 1920)", value=False)
+                    #     enable_rife = gr.Checkbox(label="Frame Interpolation (8fps -> 16fps)", value=False)
+                    # gr.Markdown(
+                    #     "✨In this demo, we use [RIFE](https://github.com/hzwer/ECCV2022-RIFE) for frame interpolation and [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN) for upscaling(Super-Resolution).<br>&nbsp;&nbsp;&nbsp;&nbsp;The entire process is based on open-source solutions."
+                    # )
+
+            generate_button = gr.Button("🎬 Generate Video")
+
+            # Add the note at the bottom-left
+            with gr.Row():
+                gr.Markdown(
+                    """
+                    **Note:** The RGB is a premultiplied version to avoid the color decontamination problem.
+                    It can directly composite with a background using:
+                    ```
+                    composite = rgb + (1 - alpha) * background
+                    ```
+                    """
+                )
+
+        with gr.Column():
+            rgb_video_output = gr.Video(label="Generate RGB Video", width=720, height=480)
+            alpha_video_output = gr.Video(label="Generate Alpha Video", width=720, height=480)
+            with gr.Row():
+                download_rgb_video_button = gr.File(label="📥 Download RGB Video", visible=False)
+                download_alpha_video_button = gr.File(label="📥 Download Alpha Video", visible=False)
+                seed_text = gr.Number(label="Seed Used for Video Generation", visible=False)
+
+
+    def generate(
+        prompt,
+        seed_value,
+        progress=gr.Progress(track_tqdm=True)
+    ):
+        latents, seed = infer(
+            prompt,
+            num_inference_steps=25,  # NOT Changed
+            guidance_scale=7.0,  # NOT Changed
+            seed=seed_value,
+            progress=progress,
+        )
+
+        latents_rgb, latents_alpha = latents.chunk(2, dim=1)
+
+        frames_rgb = decode_latents(pipe, latents_rgb)
+        frames_alpha = decode_latents(pipe, latents_alpha)
+
+        pooled_alpha = np.max(frames_alpha, axis=-1, keepdims=True)
+        frames_alpha_pooled = np.repeat(pooled_alpha, 3, axis=-1)
+        premultiplied_rgb = frames_rgb * frames_alpha_pooled
+
+        rgb_video_path = save_video(premultiplied_rgb[0], fps=8, prefix='rgb')
+        rgb_video_update = gr.update(visible=True, value=rgb_video_path)
+
+        alpha_video_path = save_video(frames_alpha_pooled[0], fps=8, prefix='alpha')
+        alpha_video_update = gr.update(visible=True, value=alpha_video_path)
+        seed_update = gr.update(visible=True, value=seed)
+
+        return rgb_video_path, alpha_video_path, rgb_video_update, alpha_video_update, seed_update
+
+
+    generate_button.click(
+        generate,
+        inputs=[prompt, seed_param],
+        outputs=[rgb_video_output, alpha_video_output, download_rgb_video_button, download_alpha_video_button, seed_text],
+    )
+
+
+if __name__ == "__main__":
+    demo.queue(max_size=15)
+    demo.launch()