diff --git a/README.md b/README.md new file mode 100644 index 0000000..152abf8 --- /dev/null +++ b/README.md @@ -0,0 +1,219 @@ +# Ruyi-Models + +English | [简体中文](./README_CN.md) + +Welcome to Ruyi-Models! + +Ruyi is an image-to-video model capable of generating cinematic-quality videos at a **resolution of 768**, with a frame rate of **24 frames per second**, totaling **5 seconds and 120 frames**. It supports **lens control** and **motion amplitude control**. Using a **RTX 3090 or RTX 4090**, you can generate 512 resolution, 120 frames (or 768 resolution, ~72 frames) videos **without any loss of quality**. + +## Table of Contents + +- [Installation Instructions](#installation-instructions) +- [Download Model (Optional)](#download-model-optional) +- [How to Use](#how-to-use) +- [Showcase](#showcase) +- [GPU Memory Optimization](#gpu-memory-optimization) +- [License](#license) + +## Installation Instructions + +The installation instructions are simple. Just clone the repo and install the requirements. + +```shell +git clone https://github.com/IamCreateAI/Ruyi-Models +cd Ruyi-Models +pip install -r requirements.txt +``` + +### For ComfyUI Users + +#### Method (1): Installation via ComfyUI Manager + +TBD, not supported yet. + +#### Method (2): Manual Installation + +Download and save this repository to the path *ComfyUI/custom_nodes/Ruyi-Models*. + +```shell +# download the repo +cd ComfyUI/custom_nodes/ +git clone https://github.com/IamCreateAI/Ruyi-Models.git + +# install requirements +pip install -r Ruyi-Models/requirements.txt +``` + +Install the dependency ComfyUI-VideoHelperSuite to display video output (skip this step if already installed). + +```shell +# download ComfyUI-VideoHelperSuite +cd ComfyUI/custom_nodes/ +git clone https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git + +# install requirements +pip install -r ComfyUI-VideoHelperSuite/requirements.txt +``` + +##### For Windows Users + +When using the Windows operating system, a common distribution is [ComfyUI_windows_portable_nvidia](https://github.com/comfyanonymous/ComfyUI/releases). When launched with `run_nvidia_gpu.bat`, it utilizes the embedded Python interpreter included with the package. Therefore, the environment needs to be set up within this built-in Python. + +For example, if the extracted directory of the distribution is ComfyUI_windows_portable, you can typically use the following command to download the repository and install the runtime environment: + +```shell +# download the repo +cd ComfyUI_windows_portable\ComfyUI\custom_nodes +git clone https://github.com/IamCreateAI/Ruyi-Models.git + +# install requirements using embedded Python interpreter +..\..\python_embeded\python.exe -m pip install -r Ruyi-Models\requirements.txt +``` + +## Download Model (Optional) + +Download the model and save it to certain path. To directly run our model, it is recommand to save the models into _Ruyi-Models/models_ folder. For ComfyUI users, the path should be _ComfyUI/models/Ruyi_. + +| Model Name | Type | Resolution | Max Frames | Frames per Second | Storage Space | Download | +| :---: | :---: | :---: | :---: | :---: | :---: | :---: | +| Ruyi-Mini-7B | Image to Video | 512 & 768 | 120 | 24 | 17 GB | [🤗](https://huggingface.co/IamCreateAI/Ruyi-Mini-7B) | + +For example, after downloading Ruyi-Mini-7B, the file path structure should be: + +``` +📦 Ruyi-Models/models/ or ComfyUI/models/Ruyi/ +├── 📂 Ruyi-Mini-7B/ +│ ├── 📂 transformers/ +│ ├── 📂 vae/ +│ └── 📂 ... +``` + +> This repository **supports automatic model downloading**, but manual downloading provides more control. For instance, you can download the model to another location and then link it to the *ComfyUI/models/Ruyi* path using symbolic links or similar methods. + +## How to Use + +We provide two ways to run our model. The first is directly using python code. + +``` +python3 predict_i2v.py +``` + +Specifically, the script downloads the model to the _Ruyi-Models/models_ folder and uses images from the [_assets_](./assets/) folder as the start and end frames for video inference. You can modify the variables in the script to replace the input images and set parameters such as video length and resolution. + +For users with more than 24GB of GPU memory, you can use predict_i2v_80g.py to enhance generation speed. For those with less GPU memory, we offer parameters to optimize memory usage, enabling the generation of higher resolution and longer videos by extending the inference time. The effects of these parameters can be found in the [GPU memory optimization section](#gpu-memory-optimization) section below. + +Or use ComfyUI wrapper in our github repo, the detail of ComfyUI nodes is described in [_comfyui/README.md_](./comfyui/README.md). + +## Showcase + +### Image to Video Effects + + + + + + + + +
+ +### Camera Control + + + + + + + + + + + + +
inputleftright
staticupdown
+ +### Motion Amplitude Control + + + + + + + + +
motion 1motion 2motion 3motion 4
+ +## GPU Memory Optimization + +We provide the options **`GPU_memory_mode` and `GPU_offload_steps` to reduce GPU memory usage**, catering to different user needs. + +Generally speaking, using **less GPU memory requires more RAM and results in longer generation times**. Below is a reference table of expected GPU memory usage and generation times. Note that, the GPU memory reported below is the `max_memory_allocated()` value. The values read from nvidia-smi may be higher than the reported values because CUDA occupies some GPU memory (usually between 500 - 800 MiB), and PyTorch's caching mechanism also requests additional GPU memory. + +### A100 Results + +- Resolution of 512 + +| Num frames | normal_mode + 0 steps | normal_mode + 10 steps | normal_mode + 7 steps | normal_mode + 5 steps | normal_mode + 1 steps | low_gpu_mode + 0 steps | +| :---: | :---: | :---: | :---: | :---: | :---: | :---: | +| 24 frames | 16119MiB
_01:01s_ | 15535MiB
_01:07s_ | 15340MiB
_01:13s_ | 15210MiB
_01:20s_ | 14950MiB
_01:32s_ | 4216MiB
_05:14s_ | +| 48 frames | 18398MiB
_01:53s_ | 17230MiB
_02:15s_ | 16840MiB
_02:29s_ | 16580MiB
_02:32s_ | 16060MiB
_02:54s_ | 4590MiB
_09:59s_ | +| 72 frames | 20678MiB
_03:00s_ | 18925MiB
_03:31s_ | 18340MiB
_03:53s_ | 17951MiB
_03:57s_ | 17171MiB
_04:25s_ | 6870MiB
_14:42s_ | +| 96 frames | 22958MiB
_04:11s_ | 20620MiB
_04:54s_ | 19841MiB
_05:10s_ | 19321MiB
_05:14s_ | 18281MiB
_05:47s_ | 9150MiB
_19:17s_ | +| 120 frames | 25238MiB
_05:42s_ | 22315MiB
_06:34s_ | 21341MiB
_06:59s_ | 20691MiB
_07:07s_ | 19392MiB
_07:41s_ | 11430MiB
_24:08s_ | + +- Resolution of 768 + +| Num frames | normal_mode + 0 steps | normal_mode + 10 steps | normal_mode + 7 steps | normal_mode + 5 steps | normal_mode + 1 steps | low_gpu_mode + 0 steps | +| :---: | :---: | :---: | :---: | :---: | :---: | :---: | +| 24 frames | 18971MiB
_02:06s_ | 17655MiB
_02:40s_ | 17217MiB
_02:39s_ | 16925MiB
_02:41s_ | 16339MiB
_03:13s_ | 5162MiB
_13:42s_ | +| 48 frames | 24101MiB
_04:52s_ | 21469MiB
_05:44s_ | 20592MiB
_05:51s_ | 20008MiB
_06:00s_ | 18837MiB
_06:49s_ | 10292MiB
_20:58s_ | +| 72 frames | 29230MiB
_08:24s_ | 25283MiB
_09:45s_ | 25283MiB
_09:45s_ | 23091MiB
_10:10s_ | 21335MiB
_11:10s_ | 15421MiB
_39:12s_ | +| 96 frames | 34360MiB
_12:49s_ | 29097MiB
_14:41s_ | 27343MiB
_15:33s_ | 26174MiB
_15:44s_ | 23834MiB
_16:33s_ | 20550MiB
_43:47s_ | +| 120 frames | 39489MiB
_18:21s_ | 32911MiB
_20:39s_ | 30719MiB
_21:34s_ | 29257MiB
_21:48s_ | 26332MiB
_23:02s_ | 25679MiB
_63:01s_ | + +### RTX 4090 Results + +The values marked with `---` in the table indicate that an out-of-memory (OOM) error occurred, preventing generation. + +- Resolution of 512 + +| Num frames | normal_mode + 0 steps | normal_mode + 10 steps | normal_mode + 7 steps | normal_mode + 5 steps | normal_mode + 1 steps | low_gpu_mode + 0 steps | +| :---: | :---: | :---: | :---: | :---: | :---: | :---: | +| 24 frames | 16366MiB
_01:18s_ | 15805MiB
_01:26s_ | 15607MiB
_01:37s_ | 15475MiB
_01:36s_ | 15211MiB
_01:39s_ | 4211MiB
_03:57s_ | +| 48 frames | 18720MiB
_02:21s_ | 17532MiB
_02:49s_ | 17136MiB
_02:55s_ | 16872MiB
_02:58s_ | 16344MiB
_03:01s_ | 4666MiB
_05:01s_ | +| 72 frames | 21036MiB
_03:41s_ | 19254MiB
_04:25s_ | 18660MiB
_04:34s_ | 18264MiB
_04:36s_ | 17472MiB
_04:51s_ | 6981MiB
_06:36s_ | +| 96 frames | -----MiB
_--:--s_ | 20972MiB
_06:18s_ | 20180MiB
_06:24s_ | 19652MiB
_06:36s_ | 18596MiB
_06:56s_ | 9298MiB
_10:03s_ | +| 120 frames | -----MiB
_--:--s_ | -----MiB
_--:--s_ | 21704MiB
_08:50s_ | 21044MiB
_08:53s_ | 19724MiB
_09:08s_ | 11613MiB
_13:57s_ | + +- Resolution of 768 + +| Num frames | normal_mode + 0 steps | normal_mode + 10 steps | normal_mode + 7 steps | normal_mode + 5 steps | normal_mode + 1 steps | low_gpu_mode + 0 steps | +| :---: | :---: | :---: | :---: | :---: | :---: | :---: | +| 24 frames | 19223MiB
_02:38s_ | 17900MiB
_03:06s_ | 17448MiB
_03:18s_ | 17153MiB
_03:23s_ | 16624MiB
_03:34s_ | 5251MiB
_05:54s_ | +| 48 frames | -----MiB
_--:--s_ | -----MiB
_--:--s_ | 20946MiB
_07:28s_ | 20352MiB
_07:35s_ | 19164MiB
_08:04s_ | 10457MiB
_10:55s_ | +| 72 frames | -----MiB
_--:--s_ | -----MiB
_--:--s_ | -----MiB
_--:--s_ | -----MiB
_--:--s_ | -----MiB
_--:--s_ | 15671MiB
_18:52s_ | + +## License + +We’re releasing the model under a permissive **Apache 2.0 license**. + +## BibTeX + +``` +@misc{createai2024ruyi, + title={Ruyi-Mini-7B}, + author={CreateAI Team}, + year={2024}, + publisher = {GitHub}, + journal = {GitHub repository}, + howpublished={\url{https://github.com/IamCreateAI/Ruyi-Models}} +} +``` + +## Welcome Feedback and Collaborative Optimization + +We sincerely welcome everyone to actively provide valuable feedback and suggestions, and we hope to work together to optimize our services and products. Your words will help us better understand user needs, allowing us to continuously enhance the user experience. Thank you for your support and attention to our work! + +You are welcomed to join our [Discord](https://discord.com/invite/nueQFQwwGw) or Wechat Group (Scan QR code) for further discussion! + +![wechat](assets/wechat_group.png) diff --git a/README_CN.md b/README_CN.md new file mode 100644 index 0000000..875d9e3 --- /dev/null +++ b/README_CN.md @@ -0,0 +1,206 @@ +# Ruyi-Models + +欢迎使用 Ruyi-Models！ + +Ruyi 是一款图生视频模型，能够生成 **768 分辨率、每秒 24 帧总计 5 秒 120 帧的影视级视频**，支持**镜头控制**与**运动幅度控制**，使用 **RTX 3090** 或 **RTX 4090** 可**无精度损失**地生成 512 分辨率、120 帧（768分辨率、~72帧）的视频。 + +## 安装方法 + +克隆本仓库并安装所需的依赖。 + +```shell +git clone https://github.com/IamCreateAI/Ruyi-Models +cd Ruyi-Models +pip install -r requirements.txt +``` + +### ComfyUI 的安装方法 + +#### 方法（1）：通过 ComfyUI Manager 安装 + +暂不支持，有待完善。 + +#### 方法（2）：手动安装 + +- 下载并保存本仓库到 *ComfyUI/custom_nodes/Ruyi-Models* 路径。 + +```shell +# download the repo +cd ComfyUI/custom_nodes/ +git clone https://github.com/IamCreateAI/Ruyi-Models.git + +# install requirements +pip install -r Ruyi-Models/requirements.txt +``` + +- 安装依赖项 [ComfyUI-VideoHelperSuite](https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite) 以显示视频输出（如已安装，请跳过本步骤）。 + +```shell +# download ComfyUI-VideoHelperSuite +cd ComfyUI/custom_nodes/ +git clone https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git + +# install requirements +pip install -r ComfyUI-VideoHelperSuite/requirements.txt +``` + +##### Windows 操作系统下的安装方法 + +在使用Windows操作系统时，[ComfyUI_windows_portable_nvidia](https://github.com/comfyanonymous/ComfyUI/releases) 是一种常见的发行版。当通过 `run_nvidia_gpu.bat` 启动时，会使用其中嵌入的 Python 解释器。因此，需要在这个内置的 Python 环境中安装运行环境。 + +例如，如果提取后的发行版目录是 _ComfyUI_windows_portable_，通常可使用以下命令下载仓库并安装运行时环境： + +```shell +# download the repo +cd ComfyUI_windows_portable\ComfyUI\custom_nodes +git clone https://github.com/IamCreateAI/Ruyi-Models.git + +# install requirements using embedded Python interpreter +..\..\python_embeded\python.exe -m pip install -r Ruyi-Models\requirements.txt +``` + +## 下载模型（可选） + +下载模型并将其保存到指定路径。为了直接运行 Ruyi 模型，建议将模型保存到 _Ruyi-Models/models_ 文件夹中。对于 ComfyUI 用户，路径应为 _ComfyUI/models/Ruyi_。 + +| 名称 | 类型 | 分辨率 | 最大帧数 | 每秒帧数 | 存储空间 | 下载地址 | +| :---: | :---: | :---: | :---: | :---: | :---: | :---: | +| Ruyi-Mini-7B | 图生视频 | 512 & 768 | 120 | 24 | 17 GB | [🤗](https://huggingface.co/IamCreateAI/Ruyi-Mini-7B) | + +例如，下载 Ruyi-Mini-7B 后，文件的路径结构应该为： + +``` +📦 Ruyi-Models/models/ or ComfyUI/models/Ruyi/ +├── 📂 Ruyi-Mini-7B/ +│ ├── 📂 transformers/ +│ ├── 📂 vae/ +│ └── 📂 ... +``` + +> 本仓库**支持自动下载模型功能**，但手动下载提供了更多的可控性。例如，可以下载模型到其他位置，再通过软链接等方式链接到 *ComfyUI/models/Ruyi* 路径。 + +## 使用方法 + +我们提供两种运行模型的方法。第一种是直接使用Python代码。 + +```shell +python3 predict_i2v.py +``` + +具体来说，该脚本将模型下载到 _Ruyi-Models/models_ 文件夹，并使用 _[assets](./assets)_ 文件夹中的图像作为视频推理的起始帧和结束帧。您可以修改脚本中的变量来替换输入图像，并设置视频长度和分辨率等参数。 + +对于显存超过 24GB 的用户，可以使用 predict_i2v_80g.py 来提高生成速度。对于显存较少的用户，提供了优化显存使用的参数，这些参数可以通过延长推理时间来生成更高分辨率和更长时长的视频。这些参数的影响可以在下面的显存优化选项小节找到。 + +或者，您可以使用我们 GitHub 仓库中的 ComfyUI 封装，ComfyUI 节点的详细信息在 _[comfyui/README_CN.md](comfyui/README_CN.md)_ 中描述。 + +## 效果展示 + +### 图生视频效果 + + + + + + + + +
+ +### 镜头控制 + + + + + + + + + + + + +
inputleftright
staticupdown
+ +### 运动幅度控制 + + + + + + + + +
motion 1motion 2motion 3motion 4
+ +## 显存优化选项 + +提供了 **`GPU_memory_mode` 和 `GPU_offload_steps`** 选项以**降低显存占用**，满足不同用户的需求。 + +通常来说，**使用更少的显存，需要更多的内存和更长的生成时间**。以下列出预期显存使用和生成时间的参考表格。请注意，下面报告的显存是 `max_memory_allocated()` 的返回值，而 nvidia-smi 的显存数值通常会高于报告的数值。因为 CUDA 会占用一些显存（通常在500 - 800 MiB之间），而 PyTorch 的缓存机制也会请求额外的显存。 + +### A100 的显存占用与运行时间 + +- Resolution of 512 + +| 帧数 | normal_mode + 0 steps | normal_mode + 10 steps | normal_mode + 7 steps | normal_mode + 5 steps | normal_mode + 1 steps | low_gpu_mode + 0 steps | +| :---: | :---: | :---: | :---: | :---: | :---: | :---: | +| 24 frames | 16119MiB
_01:01s_ | 15535MiB
_01:07s_ | 15340MiB
_01:13s_ | 15210MiB
_01:20s_ | 14950MiB
_01:32s_ | 4216MiB
_05:14s_ | +| 48 frames | 18398MiB
_01:53s_ | 17230MiB
_02:15s_ | 16840MiB
_02:29s_ | 16580MiB
_02:32s_ | 16060MiB
_02:54s_ | 4590MiB
_09:59s_ | +| 72 frames | 20678MiB
_03:00s_ | 18925MiB
_03:31s_ | 18340MiB
_03:53s_ | 17951MiB
_03:57s_ | 17171MiB
_04:25s_ | 6870MiB
_14:42s_ | +| 96 frames | 22958MiB
_04:11s_ | 20620MiB
_04:54s_ | 19841MiB
_05:10s_ | 19321MiB
_05:14s_ | 18281MiB
_05:47s_ | 9150MiB
_19:17s_ | +| 120 frames | 25238MiB
_05:42s_ | 22315MiB
_06:34s_ | 21341MiB
_06:59s_ | 20691MiB
_07:07s_ | 19392MiB
_07:41s_ | 11430MiB
_24:08s_ | + +- Resolution of 768 + +| 帧数 | normal_mode + 0 steps | normal_mode + 10 steps | normal_mode + 7 steps | normal_mode + 5 steps | normal_mode + 1 steps | low_gpu_mode + 0 steps | +| :---: | :---: | :---: | :---: | :---: | :---: | :---: | +| 24 frames | 18971MiB
_02:06s_ | 17655MiB
_02:40s_ | 17217MiB
_02:39s_ | 16925MiB
_02:41s_ | 16339MiB
_03:13s_ | 5162MiB
_13:42s_ | +| 48 frames | 24101MiB
_04:52s_ | 21469MiB
_05:44s_ | 20592MiB
_05:51s_ | 20008MiB
_06:00s_ | 18837MiB
_06:49s_ | 10292MiB
_20:58s_ | +| 72 frames | 29230MiB
_08:24s_ | 25283MiB
_09:45s_ | 25283MiB
_09:45s_ | 23091MiB
_10:10s_ | 21335MiB
_11:10s_ | 15421MiB
_39:12s_ | +| 96 frames | 34360MiB
_12:49s_ | 29097MiB
_14:41s_ | 27343MiB
_15:33s_ | 26174MiB
_15:44s_ | 23834MiB
_16:33s_ | 20550MiB
_43:47s_ | +| 120 frames | 39489MiB
_18:21s_ | 32911MiB
_20:39s_ | 30719MiB
_21:34s_ | 29257MiB
_21:48s_ | 26332MiB
_23:02s_ | 25679MiB
_63:01s_ | + +### RTX 4090 的显存占用与运行时间 + +表格中以 `---` 显示的值表示触发了显存溢出（OOM），无法生成视频。 + +- Resolution of 512 + +| 帧数 | normal_mode + 0 steps | normal_mode + 10 steps | normal_mode + 7 steps | normal_mode + 5 steps | normal_mode + 1 steps | low_gpu_mode + 0 steps | +| :---: | :---: | :---: | :---: | :---: | :---: | :---: | +| 24 frames | 16366MiB
_01:18s_ | 15805MiB
_01:26s_ | 15607MiB
_01:37s_ | 15475MiB
_01:36s_ | 15211MiB
_01:39s_ | 4211MiB
_03:57s_ | +| 48 frames | 18720MiB
_02:21s_ | 17532MiB
_02:49s_ | 17136MiB
_02:55s_ | 16872MiB
_02:58s_ | 16344MiB
_03:01s_ | 4666MiB
_05:01s_ | +| 72 frames | 21036MiB
_03:41s_ | 19254MiB
_04:25s_ | 18660MiB
_04:34s_ | 18264MiB
_04:36s_ | 17472MiB
_04:51s_ | 6981MiB
_06:36s_ | +| 96 frames | -----MiB
_--:--s_ | 20972MiB
_06:18s_ | 20180MiB
_06:24s_ | 19652MiB
_06:36s_ | 18596MiB
_06:56s_ | 9298MiB
_10:03s_ | +| 120 frames | -----MiB
_--:--s_ | -----MiB
_--:--s_ | 21704MiB
_08:50s_ | 21044MiB
_08:53s_ | 19724MiB
_09:08s_ | 11613MiB
_13:57s_ | + +- Resolution of 768 + +| 帧数 | normal_mode + 0 steps | normal_mode + 10 steps | normal_mode + 7 steps | normal_mode + 5 steps | normal_mode + 1 steps | low_gpu_mode + 0 steps | +| :---: | :---: | :---: | :---: | :---: | :---: | :---: | +| 24 frames | 19223MiB
_02:38s_ | 17900MiB
_03:06s_ | 17448MiB
_03:18s_ | 17153MiB
_03:23s_ | 16624MiB
_03:34s_ | 5251MiB
_05:54s_ | +| 48 frames | -----MiB
_--:--s_ | -----MiB
_--:--s_ | 20946MiB
_07:28s_ | 20352MiB
_07:35s_ | 19164MiB
_08:04s_ | 10457MiB
_10:55s_ | +| 72 frames | -----MiB
_--:--s_ | -----MiB
_--:--s_ | -----MiB
_--:--s_ | -----MiB
_--:--s_ | -----MiB
_--:--s_ | 15671MiB
_18:52s_ | + +## 许可证 + +我们将以宽松的 **Apache 2.0** 许可证发布该模型。 + +## 引用 + +``` +@misc{createai2024ruyi, + title={Ruyi-Mini-7B}, + author={CreateAI Team}, + year={2024}, + publisher = {GitHub}, + journal = {GitHub repository}, + howpublished={\url{https://github.com/IamCreateAI/Ruyi-Models}} +} +``` + +## 欢迎建议反馈与协同优化 + +我们真诚欢迎大家积极提供宝贵的反馈和建议，希望能够共同努力优化我们的服务和产品。您的意见将帮助我们更好地理解用户需求，从而不断提升用户体验。感谢您对我们工作的支持和关注！欢迎加入我们的[Discord](https://discord.com/invite/nueQFQwwGw) 或者微信群（扫描下方二维码）！ + +![wechat](assets/wechat_group.png) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..273bba3 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .comfyui.comfyui_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/assets/example_01.png b/assets/example_01.png new file mode 100644 index 0000000..f18c97e Binary files /dev/null and b/assets/example_01.png differ diff --git a/assets/example_02.jpg b/assets/example_02.jpg new file mode 100644 index 0000000..7dd5375 Binary files /dev/null and b/assets/example_02.jpg differ diff --git a/assets/example_03.jpg b/assets/example_03.jpg new file mode 100644 index 0000000..068ccbb Binary files /dev/null and b/assets/example_03.jpg differ diff --git a/assets/girl_01.jpg b/assets/girl_01.jpg new file mode 100644 index 0000000..243f12b Binary files /dev/null and b/assets/girl_01.jpg differ diff --git a/assets/girl_02.jpg b/assets/girl_02.jpg new file mode 100644 index 0000000..d5c26bd Binary files /dev/null and b/assets/girl_02.jpg differ diff --git a/assets/wechat_group.png b/assets/wechat_group.png new file mode 100644 index 0000000..76384c4 Binary files /dev/null and b/assets/wechat_group.png differ diff --git a/comfyui/README.md b/comfyui/README.md new file mode 100644 index 0000000..bf999e7 --- /dev/null +++ b/comfyui/README.md @@ -0,0 +1,109 @@ +## ComfyUI Nodes + +After installation, you can find the three nodes provided by this repository in the *Add Node - Ruyi* menu, as shown in the image below: + +
+ +
+ +The following sections will introduce the functions and parameters of each node. + +> Note: The new version of ComfyUI nodes is displayed in the NODE LIBRARY on the left side of the interface. + +### Load Model + +The Load Model node is used to load the model from disk. It also provides the functionality for **automatic model downloading** (auto_download parameter). + +
+ +
+ +- **model**: Select which model to use. Currently, Ruyi-Mini-7B is the only option. +- **auto_download**: Whether to automatically download. Defaults to yes. If the model is detected as missing or incomplete, it will automatically download the model to the *ComfyUI/models/Ruyi* path. + +### Load LoRA + +The Load LoRA node is used to load LoRA models, which need to be placed in the *ComfyUI/models/loras* path. + +
+ +
+ +- **lora_name**: The LoRA to be loaded; it will automatically search and display all model files in the *ComfyUI/models/loras* path. +- **strength_model**: The degree of influence of the LoRA, typically set between 1.0 and 1.4 for optimal results based on experience. + +### Sampler for Image to Video + +The Sampler for Image to Video node is used to generate videos based on input images. The starting frame image (start_img) is a required input, while the ending frame image (end_img) is optional. This node also supports **camera control** (camera_direction parameter) and **motion amplitude control** of the video subject (motion parameter). + +
+ +
+ +- **start_img**: The starting frame image. +- **end_img**: The ending frame image, optional input. +- **video_length**: The length of the video, which must be divisible by 8, with a maximum of 120 frames. +- **base_resolution**: The video resolution, such as 512, indicates that the generated video will have pixel dimensions close to 512 x 512. The model will automatically select the closest output video aspect ratio based on the input image. +- **seed**: A random number; different random numbers usually generate different videos. If the generated video does not meet requirements, this value can be adjusted to try other generation possibilities. +- **control_after_generate**: The method of changing the random number after each generation. + - **Fixed** indicates the seed is fixed. + - **Increment** indicates the seed is increased by one each time. + - **Decrement** indicates the seed is decreased by one each time. + - **Randomize** indicates the seed is randomly set each time. +- **steps**: The number of iterations for video generation. **More iterations require more time**. Typically, 25 iterations yield good results. +- **cfg**: The guidance of instructions (such as input images). A higher value indicates better guidance, with values between 7 and 10 usually achieving better generation results. +- **motion**: Controls the motion amplitude of the video subject. + - **1** indicates minimal motion, which means the video subject is nearly static. + - **2** indicates normal motion, which could be used in most case. + - **3** indicates significant motion. The video subject is trying to move as much as possible. + - **4** indicates a very large motion. Sometimes the video subject may move out of the camera frame. + - **Auto** indicates the motion is automatically determined by the model. +- **camera_direction**: Camera movement. + - **Static** indicates a stationary camera. + - **Left** indicates the camera moves left. + - **Right** indicates the camera moves right. + - **Up** indicates the camera moves up. + - **Down** indicates the camera moves down. + - **Auto** indicates automatic determination. +- **GPU_memory_mode**: Determines how GPU memory is utilized. + - **normal_mode** is the default mode, using more GPU memory and generating faster. + - **low_memory_mode** is the low memory mode, which significantly reduces GPU memory usage but severely impacts generation speed. +- **GPU_offload_steps**: Used to **optimize GPU memory usage** by moving some temporary variables from GPU memory to RAM, which **increases memory usage and decreases generation speed**. + - **0** indicates no optimization. + - **1 - 10**, where 1 has the least GPU memory usage and the slowest generation speed; 10 has the most GPU memory usage (less than the non-optimized case) and the fastest generation speed. + - Generally, with 24GB of GPU memory, you can use 7 to generate a 512 resolution, 120 frame video. For more detailed data, please refer to the following. + +## Workflow Example + +This section presents an example workflow for generating videos from images. You can import the workflow using the *Load* button in the (bottom right) menu. Note that the new version of ComfyUI move the menu to the top left, and allows you to load workflows through the *Workflow - Open* option. + +The workflows are located in the *[comfyui/workflows/](workflows/)* directory, while the assets can be found in the *[assets/](../assets/)* directory. + +After importing the workflow, you **need to manually re-specify the input image for the LoadImage input node**. Since the workflow file can only record the names of input files, it does require manual configuration. + +### Image to Video (Starting Frame) + +The workflow corresponds to the *[workflow-ruyi-i2v-start-frame.json](workflows/workflow-ruyi-i2v-start-frame.json)* file. For users with larger GPU memory, you can also use *[workflow-ruyi-i2v-start-frame-80g.json](workflows/workflow-ruyi-i2v-start-frame-80g.json)* to enhance the generation speed. + +
+ +
+ +### Image to Video (Starting and Ending Frames) + +The workflow corresponds to the *[workflow-ruyi-i2v-start-end-frames.json](workflows/workflow-ruyi-i2v-start-end-frames.json)* file. For users with larger GPU memory, you can also use *[workflow-ruyi-i2v-start-end-frames-80g.json](workflows/workflow-ruyi-i2v-start-end-frames-80g.json)* to enhance the generation speed. + +
+ +
+ +## Frequently Asked Questions + +### Model loading error: LoadModel: ConnectionError (MaxRetryError) + +This is usually caused by network issues leading to a failure in downloading from huggingface_hub. If the network is functioning properly, simply rerunning the LoadModel node should resolve the issue. + +### Video generation speed is slow and far below expectations + +- First, check if the low_memory_mode in the Load Model node is enabled. This mode significantly reduces video generation speed. +- Second, verify the version of PyTorch. PyTorch version 2.2 supports FlashAttention-2 ([link](https://pytorch.org/blog/pytorch2-2/)), which can greatly enhance computational efficiency. Installing the latest version of PyTorch can effectively improve generation speed. diff --git a/comfyui/README_CN.md b/comfyui/README_CN.md new file mode 100644 index 0000000..c645bf4 --- /dev/null +++ b/comfyui/README_CN.md @@ -0,0 +1,111 @@ +## 节点功能 + +安装完成后，可在 ComfyUI 的 *Add Node - Ruyi* 菜单找到本仓库提供的 3 个节点，如下图所示： + +
+ +
+ +下文依次介绍各个节点的功能与参数。 + +> 注：新版 ComfyUI 节点在界面左侧 NODE LIBRARY 中显示。 + +### Load Model + +用于加载模型，并提供**自动下载模型**（通过 auto_download 选项设置）的功能。 + +
+ +
+ +- **model**: 选择使用哪个模型，目前只有 Ruyi-Mini-7B 一个选项。 +- **auto_download**: 是否自动下载，默认为 yes，检测到模型不存在（或不完整）时，将自动下载模型到 *ComfyUI/models/Ruyi* 路径。 + +### Load LoRA + +用于加载 LoRA 模型，LoRA 模型需要放在 *ComfyUI/models/loras* 路径下。 + +
+ +
+ +- **lora_name**: 需要加载的 LoRA，将自动搜索并显示 *ComfyUI/models/loras* 路径下所有模型文件。 +- **strength_model**: LoRA 的影响程度，根据经验通常设置在 1.0 ~ 1.4 效果较好。 + +### Sampler for Image to Video + +用于根据输入图片生成视频，首帧图片（start_img）为必须输入，尾帧图片（end_img）为可选输入。同时，该节点支持**镜头控制**（camera_direction）与**视频主体的运动幅度控制**（motion）。 + +
+ +
+ +- **start_img**: 首帧图片。 +- **end_img**: 尾帧图片，可选输入。 +- **video_length**: 视频长度，必须能**被 8 整除**，**最大支持 120 帧**。 +- **base_resolution**: 视频分辨率，**模型将根据输入图片的长宽比自动选择最接近的输出视频长宽**。 + - **512** 表示生成的视频像素数接近 512 * 512。 + - **768** 表示生成的视频像素数接近 768 * 768。 +- **seed**: 随机种子，用于控制随机数生成器产生随机数的序列。**不同的随机种子通常能生成不同的视频**，当生成的视频不符合需求时候，可调整此值以尝试其他的生成可能。 +- **control_after_generate**: 每次生成后随机种子的变化方式。 + - **Fixed** 表示随机种子固定不变。 + - **Increment** 表示随机种子每次增加一。 + - **Decrement** 表示随机种子每次减少一。 + - **Randomize** 表示随机种子运行后随机设置。 +- **steps**: 视频生成的迭代次数，**迭代次数越多，需要的时间越久**，通常 25 次能够得到较好的结果。 +- **cfg**: 指令（如输入图片）的遵循程度，数值越大遵循程度越好，取值 7 ～ 10 通常能取得较好的生成效果。 +- **motion**: 控制视频主体的运动幅度。 + - **1** 基本不运动，适用于静态场景。 + - **2** 正常运动幅度，适用于谈话、转头等常见场合。 + - **3** 运动幅度较大，可能出现转身、走动等情况。 + - **4** 运动幅度非常大，可能出现视频主体离开画面的情况。 + - **Auto** 表示模型自动判断运动幅度大小。 +- **camera_direction**: 镜头运动。 + - **Static** 表示静止镜头。 + - **Left** 表示镜头向左移动。 + - **Right** 表示镜头向右移动。 + - **Up** 表示镜头向上移动。 + - **Down** 表示镜头向下移动。 + - **Auto** 表示模型自动判断镜头运动方向。 +- **GPU_memory_mode**: + - **normal_mode** 是默认模式，**使用显存较多，生成速度较快**。 + - **low_memory_mode** 是低显存模式，**能大幅降低显存用量，但严重影响生成速度**。 +- **GPU_offload_steps**: 用于**优化显存占用**，通过将部分临时变量从显存移动到内存而实现，会**增加内存的占用并降低生成速度**。 + - **0** 表示不优化。 + - **1 - 10**，1 显存占用最小，生成速度最慢；10 显存占用最多（少于不优化情况），生成速度最快。 + - 通常情况下，24G 显存可以使用 7 生成 512 分辨率，120 帧视频。更详细数据请参考下文。 + +## 工作流样例 + +本节展示了图生视频的工作流样例，可通过菜单中的 *Load* 按钮导入工作流。新版 ComfyUI 可通过左上方菜单的 *Workflow - Open* 加载工作流。 + +工作流位于 *[comfyui/workflows/](workflows/)* 目录中，素材位于 *[assets/](../assets/)* 目录中。 + +导入工作流后，**需要手动重新指定输入节点 LoadImage 的输入图片**，由于工作流文件只能记录输入文件的名字，所以目前只能手动设置。 + +### 图生视频（首帧） + +工作流对应 *[workflow-ruyi-i2v-start-frame.json](workflows/workflow-ruyi-i2v-start-frame.json)* 文件。对于显存较大的用户，也可以使用 *[workflow-ruyi-i2v-start-frame-80g.json](workflows/workflow-ruyi-i2v-start-frame-80g.json)* 以提高生成速度。 + +
+ +
+ +### 图生视频（首尾帧） + +工作流对应 *[workflow-ruyi-i2v-start-end-frames.json](workflows/workflow-ruyi-i2v-start-end-frames.json)* 文件。对于显存较大的用户，也可以使用 *[workflow-ruyi-i2v-start-end-frames-80g.json](workflows/workflow-ruyi-i2v-start-end-frames-80g.json)* 以提高生成速度。 + +
+ +
+ +## 常见问题 + +### 模型加载错误，LoadModel: ConnectionError (MaxRetryError) + +通常是网络问题导致 huggingface_hub 下载失败。网络正常的情况下，再次运行 LoadModel 节点即可解决。 + +### 视频生成速度慢、远低于预期 + +- 首先，请检查是否开启 Load Model 节点中 GPU_memory_mode 的 low_memory_mode，此模式会大幅降低视频生成速度。 +- 其次，请检查 PyTorch 版本。PyTorch 在 2.2 版本支持了 FlashAttention-2（[链接](https://pytorch.org/blog/pytorch2-2/)），能大幅提升计算效率。安装新版本的 PyTorch 能够有效提升生成速度。 diff --git a/comfyui/comfyui_nodes.py b/comfyui/comfyui_nodes.py new file mode 100644 index 0000000..b6923c5 --- /dev/null +++ b/comfyui/comfyui_nodes.py @@ -0,0 +1,426 @@ +"""Modified from https://github.com/kijai/ComfyUI-EasyAnimateWrapper/blob/main/nodes.py +""" +import gc +import os + +import comfy.model_management as mm +import cv2 +import folder_paths +import numpy as np +import torch +from comfy.utils import ProgressBar +from diffusers import (DDIMScheduler, DPMSolverMultistepScheduler, PNDMScheduler, + EulerAncestralDiscreteScheduler, EulerDiscreteScheduler) +from einops import rearrange +from omegaconf import OmegaConf +from PIL import Image +from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection +from safetensors.torch import load_file as load_safetensors +from huggingface_hub import snapshot_download + +from ..ruyi.data.bucket_sampler import ASPECT_RATIO_512, get_closest_ratio +from ..ruyi.models.autoencoder_magvit import AutoencoderKLMagvit +from ..ruyi.models.transformer3d import HunyuanTransformer3DModel +from ..ruyi.pipeline.pipeline_ruyi_inpaint import RuyiInpaintPipeline +from ..ruyi.utils.lora_utils import merge_lora, unmerge_lora +from ..ruyi.utils.utils import get_image_to_video_latent + +# The directory of scripts +script_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +class Ruyi_LoadModel: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ( + [ + "Ruyi-Mini-7B", + ], + { "default": "Ruyi-Mini-7B", } + ), + "auto_download": ( + ["yes", "no"], + { + "default": "yes" + } + ), + }, + } + + RETURN_TYPES = ("RUYI_MODEL",) + RETURN_NAMES = ("ruyi_model",) + FUNCTION = "load_model" + CATEGORY = "Ruyi" + + RUYI_MODEL_TYPE_DICT = { + "Ruyi-Mini-7B": "Inpaint" + } + + def get_model_type(self, model_name): + return self.RUYI_MODEL_TYPE_DICT.get(model_name, "Inpaint") + + def try_setup_pipeline(self, model_path, weight_dtype, config): + try: + # Init processbar + pbar = ProgressBar(5) + + # Get Vae + vae = AutoencoderKLMagvit.from_pretrained( + model_path, + subfolder="vae" + ).to(weight_dtype) + # Update pbar + pbar.update(1) + + # Get Transformer + transformer_additional_kwargs = OmegaConf.to_container(config['transformer_additional_kwargs']) + transformer = HunyuanTransformer3DModel.from_pretrained_2d( + model_path, + subfolder="transformer", + transformer_additional_kwargs=transformer_additional_kwargs + ).to(weight_dtype) + # Update pbar + pbar.update(1) + + # Load Clip + clip_image_encoder = CLIPVisionModelWithProjection.from_pretrained( + model_path, subfolder="image_encoder" + ).to(weight_dtype) + clip_image_processor = CLIPImageProcessor.from_pretrained( + model_path, subfolder="image_encoder" + ) + # Update pbar + pbar.update(1) + + # Load sampler and create pipeline + Choosen_Scheduler = DDIMScheduler + scheduler = Choosen_Scheduler.from_pretrained( + model_path, + subfolder="scheduler" + ) + pipeline = RuyiInpaintPipeline.from_pretrained( + model_path, + vae=vae, + transformer=transformer, + scheduler=scheduler, + torch_dtype=weight_dtype, + clip_image_encoder=clip_image_encoder, + clip_image_processor=clip_image_processor, + ) + + # Load embeddings + embeddings = load_safetensors(os.path.join(model_path, "embeddings.safetensors")) + pipeline.embeddings = embeddings + # Update pbar + pbar.update(1) + + return pipeline + except Exception as e: + print("[Ruyi] Setup pipeline failed:", e) + return None + + def load_model(self, model, auto_download): + # Init weight_dtype and device + device = mm.get_torch_device() + offload_device = mm.unet_offload_device() + + # Init model name and type + model_name = model + model_type = self.get_model_type(model_name) + weight_dtype = torch.bfloat16 + + # Load config + config_path = os.path.join(script_directory, "config", "default.yaml") + config = OmegaConf.load(config_path) + + # Init model + repo_id = f"IamCreateAI/{model_name}" + model_path = os.path.join(folder_paths.models_dir, "Ruyi", model_name) + pipeline = self.try_setup_pipeline(model_path, weight_dtype, config) + if pipeline is None and auto_download == "yes": + mm.soft_empty_cache() + gc.collect() + + # Download the model + snapshot_download(repo_id=repo_id, local_dir=model_path) + + pipeline = self.try_setup_pipeline(model_path, weight_dtype, config) + + if pipeline is None: + message = (f"[ Load Model {model_name} Failed ]\n" + f"Please download Ruyi model from huggingface repo '{repo_id}',\n" + f"And put it into '{model_path}'.") + if auto_download == "no": + message += "\n\nOr just set auto_download to 'yes'." + raise FileNotFoundError(message) + + pipeline.enable_model_cpu_offload() + ruyi_model = { + 'pipeline': pipeline, + 'dtype': weight_dtype, + 'model_path': model_path, + 'model_type': model_type, + 'loras': [], + 'strength_model': [], + } + return (ruyi_model,) + + +class Ruyi_LoadLora: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "ruyi_model": ("RUYI_MODEL",), + "lora_name": (folder_paths.get_filename_list("loras"), {"default": None,}), + "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), + } + } + RETURN_TYPES = ("RUYI_MODEL",) + RETURN_NAMES = ("ruyi_model",) + FUNCTION = "load_lora" + CATEGORY = "Ruyi" + + def load_lora(self, ruyi_model, lora_name, strength_model): + if lora_name is not None: + return ( + { + 'pipeline': ruyi_model["pipeline"], + 'dtype': ruyi_model["dtype"], + 'model_path': ruyi_model["model_path"], + 'model_type': ruyi_model["model_type"], + 'loras': ruyi_model.get("loras", []) + [folder_paths.get_full_path("loras", lora_name)], + 'strength_model': ruyi_model.get("strength_model", []) + [strength_model], + }, + ) + else: + return (ruyi_model,) + + +class Ruyi_I2VSampler: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "ruyi_model": ( + "RUYI_MODEL", + ), + "video_length": ( + "INT", {"default": 72, "min": 8, "max": 120, "step": 8} + ), + "base_resolution": ( + "INT", {"default": 512, "min": 384, "max": 1024, "step": 16} + ), + "seed": ( + "INT", {"default": 42, "min": 0, "max": 0xffffffffffffffff} + ), + "steps": ( + "INT", {"default": 25, "min": 1, "max": 200, "step": 1} + ), + "cfg": ( + "FLOAT", {"default": 7.0, "min": 1.0, "max": 20.0, "step": 0.01} + ), + "scheduler": ( + [ + "Euler", + "Euler A", + "DPM++", + "PNDM", + "DDIM", + ], + { + "default": 'DDIM' + } + ), + "motion": ( + [ "1", "2", "3", "4", "auto" ], + { "default": "2" } + ), + "camera_direction": ( + [ "static", "left", "right", "up", "down", "auto" ], + { "default": "static" } + ), + "GPU_memory_mode": ( + ["normal_mode", "low_memory_mode"], + { + "default": "normal_mode", + } + ), + "GPU_offload_steps": ( + [ "0", "1", "5", "7", "10" ], + { "default": "0" } + ), + "start_img": ( + "IMAGE", + ), + }, + "optional":{ + "end_img": ("IMAGE",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES =("images",) + FUNCTION = "process" + CATEGORY = "Ruyi" + + def tensor2pil(self, image): + return Image.fromarray(np.clip(255. * image.cpu().numpy(), 0, 255).astype(np.uint8)) + + def numpy2pil(self, image): + return Image.fromarray(np.clip(255. * image, 0, 255).astype(np.uint8)) + + def to_pil(self, image): + if isinstance(image, Image.Image): + return image + if isinstance(image, torch.Tensor): + return self.tensor2pil(image) + if isinstance(image, np.ndarray): + return self.numpy2pil(image) + raise ValueError(f"Cannot convert {type(image)} to PIL.Image") + + def get_control_embeddings(self, pipeline, aspect_ratio, motion, camera_direction): + # Default keys + p_default_key = "p.default" + n_default_key = "n.default" + + # Load embeddings + if motion == "auto": + motion = "0" + p_key = f"p.{aspect_ratio.replace(':', 'x')}movie{motion}{camera_direction}" + embeddings = pipeline.embeddings + + # Get embeddings + positive_embeds = embeddings.get(f"{p_key}.emb1", embeddings[f"{p_default_key}.emb1"]) + positive_attention_mask = embeddings.get(f"{p_key}.mask1", embeddings[f"{p_default_key}.mask1"]) + positive_embeds_2 = embeddings.get(f"{p_key}.emb2", embeddings[f"{p_default_key}.emb2"]) + positive_attention_mask_2 = embeddings.get(f"{p_key}.mask2", embeddings[f"{p_default_key}.mask2"]) + + negative_embeds = embeddings[f"{n_default_key}.emb1"] + negative_attention_mask = embeddings[f"{n_default_key}.mask1"] + negative_embeds_2 = embeddings[f"{n_default_key}.emb2"] + negative_attention_mask_2 = embeddings[f"{n_default_key}.mask2"] + + return { + "positive_embeds": positive_embeds, + "positive_attention_mask": positive_attention_mask, + "positive_embeds_2": positive_embeds_2, + "positive_attention_mask_2": positive_attention_mask_2, + + "negative_embeds": negative_embeds, + "negative_attention_mask": negative_attention_mask, + "negative_embeds_2": negative_embeds_2, + "negative_attention_mask_2": negative_attention_mask_2, + } + + def get_scheduler(self, model_path, scheduler_name): + if scheduler_name == "DPM++": + noise_scheduler = DPMSolverMultistepScheduler.from_pretrained(model_path, subfolder='scheduler') + elif scheduler_name == "Euler": + noise_scheduler = EulerDiscreteScheduler.from_pretrained(model_path, subfolder='scheduler') + elif scheduler_name == "Euler A": + noise_scheduler = EulerAncestralDiscreteScheduler.from_pretrained(model_path, subfolder='scheduler') + elif scheduler_name == "PNDM": + noise_scheduler = PNDMScheduler.from_pretrained(model_path, subfolder='scheduler') + elif scheduler_name == "DDIM": + noise_scheduler = DDIMScheduler.from_pretrained(model_path, subfolder='scheduler') + + return noise_scheduler + + def process( + self, ruyi_model, video_length, base_resolution, seed, steps, cfg, scheduler, + motion, camera_direction, GPU_memory_mode, GPU_offload_steps, + start_img, end_img=None + ): + device = mm.get_torch_device() + offload_device = mm.unet_offload_device() + + mm.soft_empty_cache() + gc.collect() + + start_img = [self.to_pil(_start_img) for _start_img in start_img] if start_img is not None else None + end_img = [self.to_pil(_end_img) for _end_img in end_img] if end_img is not None else None + + # Count most suitable height and width + aspect_ratio_sample_size = {key : [x / 512 * base_resolution for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + original_width, original_height = start_img[0].size if type(start_img) is list else Image.open(start_img).size + closest_size, closest_ratio = get_closest_ratio(original_height, original_width, ratios=aspect_ratio_sample_size) + height, width = [int(x / 16) * 16 for x in closest_size] + aspect_ratio = "16:9" if width > height else "9:16" + + # Get Pipeline + pipeline = ruyi_model['pipeline'] + model_path = ruyi_model['model_path'] + + # Set GPU memory mode + if GPU_memory_mode == "low_memory_mode" and pipeline.model_cpu_offload_flag: + # Switch to low_memory_mode + pipeline.enable_sequential_cpu_offload() + elif GPU_memory_mode == "normal_mode" and not pipeline.model_cpu_offload_flag: + # Switch to normal_mode + pipeline.enable_model_cpu_offload() + + # Set GPU offload steps + pipeline.transformer.hidden_cache_size = int(GPU_offload_steps) + + # Load Sampler + pipeline.scheduler = self.get_scheduler(model_path, scheduler) + + # Set random seed + generator= torch.Generator(device).manual_seed(seed) + + # Load control embeddings + embeddings = self.get_control_embeddings(pipeline, aspect_ratio, motion, camera_direction) + + # Inference + with torch.no_grad(): + video_length = int(video_length // pipeline.vae.mini_batch_encoder * pipeline.vae.mini_batch_encoder) if video_length != 1 else 1 + input_video, input_video_mask, clip_image = get_image_to_video_latent(start_img, end_img, video_length=video_length, sample_size=(height, width)) + + for _lora_path, _lora_weight in zip(ruyi_model.get("loras", []), ruyi_model.get("strength_model", [])): + pipeline = merge_lora(pipeline, _lora_path, _lora_weight) + + sample = pipeline( + prompt_embeds = embeddings["positive_embeds"], + prompt_attention_mask = embeddings["positive_attention_mask"], + prompt_embeds_2 = embeddings["positive_embeds_2"], + prompt_attention_mask_2 = embeddings["positive_attention_mask_2"], + + negative_prompt_embeds = embeddings["negative_embeds"], + negative_prompt_attention_mask = embeddings["negative_attention_mask"], + negative_prompt_embeds_2 = embeddings["negative_embeds_2"], + negative_prompt_attention_mask_2 = embeddings["negative_attention_mask_2"], + + video_length = video_length, + height = height, + width = width, + generator = generator, + guidance_scale = cfg, + num_inference_steps = steps, + + video = input_video, + mask_video = input_video_mask, + clip_image = clip_image, + comfyui_progressbar = True, + ).videos + videos = rearrange(sample, "b c t h w -> (b t) h w c") + + for _lora_path, _lora_weight in zip(ruyi_model.get("loras", []), ruyi_model.get("strength_model", [])): + pipeline = unmerge_lora(pipeline, _lora_path, _lora_weight) + + return (videos,) + + +NODE_CLASS_MAPPINGS = { + "Ruyi_LoadModel": Ruyi_LoadModel, + "Ruyi_LoadLora": Ruyi_LoadLora, + "Ruyi_I2VSampler": Ruyi_I2VSampler, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "Ruyi_LoadModel": "Load Model", + "Ruyi_LoadLora": "Load LoRA", + "Ruyi_I2VSampler": "Sampler for Image to Video", +} diff --git a/comfyui/workflows/workflow-ruyi-i2v-start-end-frames-80g.json b/comfyui/workflows/workflow-ruyi-i2v-start-end-frames-80g.json new file mode 100644 index 0000000..4b16047 --- /dev/null +++ b/comfyui/workflows/workflow-ruyi-i2v-start-end-frames-80g.json @@ -0,0 +1,293 @@ +{ + "last_node_id": 5, + "last_link_id": 4, + "nodes": [ + { + "id": 1, + "type": "Ruyi_LoadModel", + "pos": { + "0": 295, + "1": 238 + }, + "size": { + "0": 315, + "1": 82 + }, + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "ruyi_model", + "type": "RUYI_MODEL", + "links": [ + 2 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "Ruyi_LoadModel" + }, + "widgets_values": [ + "Ruyi-Mini-7B", + "yes" + ] + }, + { + "id": 2, + "type": "Ruyi_I2VSampler", + "pos": { + "0": 753, + "1": 326 + }, + "size": { + "0": 327.5999755859375, + "1": 338 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "ruyi_model", + "type": "RUYI_MODEL", + "link": 2 + }, + { + "name": "start_img", + "type": "IMAGE", + "link": 3 + }, + { + "name": "end_img", + "type": "IMAGE", + "link": 4, + "shape": 7 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 1 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "Ruyi_I2VSampler" + }, + "widgets_values": [ + 120, + 512, + 112876945129710, + "randomize", + 25, + 7, + "DDIM", + "auto", + "auto", + "normal_mode", + "0" + ] + }, + { + "id": 5, + "type": "VHS_VideoCombine", + "pos": { + "0": 1183, + "1": 89 + }, + "size": [ + 317, + 758.0769230769231 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 1 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null, + "shape": 7 + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null, + "shape": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 24, + "loop_count": 0, + "filename_prefix": "Ruyi-I2V-StartEndFrames", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "Ruyi-I2V-StartEndFrames_00001.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 24 + }, + "muted": false + } + } + }, + { + "id": 3, + "type": "LoadImage", + "pos": { + "0": 180, + "1": 398 + }, + "size": { + "0": 240.88999938964844, + "1": 419.6499938964844 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 3 + ], + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "girl_01.jpg", + "image" + ] + }, + { + "id": 4, + "type": "LoadImage", + "pos": { + "0": 453, + "1": 398 + }, + "size": [ + 242.0999999999999, + 421.6500000000001 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 4 + ], + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "girl_02.jpg", + "image" + ] + } + ], + "links": [ + [ + 1, + 2, + 0, + 5, + 0, + "IMAGE" + ], + [ + 2, + 1, + 0, + 2, + 0, + "RUYI_MODEL" + ], + [ + 3, + 3, + 0, + 2, + 1, + "IMAGE" + ], + [ + 4, + 4, + 0, + 2, + 2, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1, + "offset": [ + -43.00999999999996, + 33.75000000000006 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/comfyui/workflows/workflow-ruyi-i2v-start-end-frames.json b/comfyui/workflows/workflow-ruyi-i2v-start-end-frames.json new file mode 100644 index 0000000..f604b26 --- /dev/null +++ b/comfyui/workflows/workflow-ruyi-i2v-start-end-frames.json @@ -0,0 +1,293 @@ +{ + "last_node_id": 5, + "last_link_id": 4, + "nodes": [ + { + "id": 1, + "type": "Ruyi_LoadModel", + "pos": { + "0": 295, + "1": 238 + }, + "size": { + "0": 315, + "1": 82 + }, + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "ruyi_model", + "type": "RUYI_MODEL", + "links": [ + 2 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "Ruyi_LoadModel" + }, + "widgets_values": [ + "Ruyi-Mini-7B", + "yes" + ] + }, + { + "id": 5, + "type": "VHS_VideoCombine", + "pos": { + "0": 1183, + "1": 89 + }, + "size": [ + 317, + 758.0769230769231 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 1 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null, + "shape": 7 + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null, + "shape": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 24, + "loop_count": 0, + "filename_prefix": "Ruyi-I2V-StartEndFrames", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "Ruyi-I2V-StartEndFrames_00001.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 24 + }, + "muted": false + } + } + }, + { + "id": 3, + "type": "LoadImage", + "pos": { + "0": 180, + "1": 398 + }, + "size": { + "0": 240.88999938964844, + "1": 419.6499938964844 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 3 + ], + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "gir_01.jpg", + "image" + ] + }, + { + "id": 4, + "type": "LoadImage", + "pos": { + "0": 453, + "1": 398 + }, + "size": [ + 242.0999999999999, + 421.6500000000001 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 4 + ], + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "girl_02.jpg", + "image" + ] + }, + { + "id": 2, + "type": "Ruyi_I2VSampler", + "pos": { + "0": 753, + "1": 326 + }, + "size": { + "0": 327.5999755859375, + "1": 338 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "ruyi_model", + "type": "RUYI_MODEL", + "link": 2 + }, + { + "name": "start_img", + "type": "IMAGE", + "link": 3 + }, + { + "name": "end_img", + "type": "IMAGE", + "link": 4, + "shape": 7 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 1 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "Ruyi_I2VSampler" + }, + "widgets_values": [ + 120, + 512, + 112876945129710, + "randomize", + 25, + 7, + "DDIM", + "auto", + "auto", + "normal_mode", + "5" + ] + } + ], + "links": [ + [ + 1, + 2, + 0, + 5, + 0, + "IMAGE" + ], + [ + 2, + 1, + 0, + 2, + 0, + "RUYI_MODEL" + ], + [ + 3, + 3, + 0, + 2, + 1, + "IMAGE" + ], + [ + 4, + 4, + 0, + 2, + 2, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1, + "offset": [ + -43.00999999999996, + 33.75000000000006 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/comfyui/workflows/workflow-ruyi-i2v-start-frame-80g.json b/comfyui/workflows/workflow-ruyi-i2v-start-frame-80g.json new file mode 100644 index 0000000..9159b92 --- /dev/null +++ b/comfyui/workflows/workflow-ruyi-i2v-start-frame-80g.json @@ -0,0 +1,247 @@ +{ + "last_node_id": 6, + "last_link_id": 3, + "nodes": [ + { + "id": 2, + "type": "Ruyi_I2VSampler", + "pos": { + "0": 628, + "1": 284 + }, + "size": { + "0": 327.5999755859375, + "1": 338 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "ruyi_model", + "type": "RUYI_MODEL", + "link": 1 + }, + { + "name": "start_img", + "type": "IMAGE", + "link": 2 + }, + { + "name": "end_img", + "type": "IMAGE", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 3 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "Ruyi_I2VSampler" + }, + "widgets_values": [ + 120, + 512, + 925247271358454, + "randomize", + 25, + 7, + "DDIM", + "2", + "static", + "normal_mode", + "0" + ] + }, + { + "id": 1, + "type": "Ruyi_LoadModel", + "pos": { + "0": 210, + "1": 162 + }, + "size": { + "0": 315, + "1": 82 + }, + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "ruyi_model", + "type": "RUYI_MODEL", + "links": [ + 1 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "Ruyi_LoadModel" + }, + "widgets_values": [ + "Ruyi-Mini-7B", + "yes" + ] + }, + { + "id": 4, + "type": "VHS_VideoCombine", + "pos": { + "0": 1045, + "1": 133 + }, + "size": [ + 404.73553466796875, + 601.8645528157551 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 3 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null, + "shape": 7 + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null, + "shape": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 24, + "loop_count": 0, + "filename_prefix": "Ruyi-I2V-StartFrame", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "Ruyi-I2V-StartFrame_00001.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 24 + }, + "muted": false + } + } + }, + { + "id": 3, + "type": "LoadImage", + "pos": { + "0": 200, + "1": 439 + }, + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 2 + ], + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "example_03.jpg", + "image" + ] + } + ], + "links": [ + [ + 1, + 1, + 0, + 2, + 0, + "RUYI_MODEL" + ], + [ + 2, + 3, + 0, + 2, + 1, + "IMAGE" + ], + [ + 3, + 2, + 0, + 4, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1, + "offset": [ + 0, + 0 + ] + } + }, + "version": 0.4 +} diff --git a/comfyui/workflows/workflow-ruyi-i2v-start-frame.json b/comfyui/workflows/workflow-ruyi-i2v-start-frame.json new file mode 100644 index 0000000..1929a35 --- /dev/null +++ b/comfyui/workflows/workflow-ruyi-i2v-start-frame.json @@ -0,0 +1,247 @@ +{ + "last_node_id": 6, + "last_link_id": 3, + "nodes": [ + { + "id": 1, + "type": "Ruyi_LoadModel", + "pos": { + "0": 210, + "1": 162 + }, + "size": { + "0": 315, + "1": 82 + }, + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "ruyi_model", + "type": "RUYI_MODEL", + "links": [ + 1 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "Ruyi_LoadModel" + }, + "widgets_values": [ + "Ruyi-Mini-7B", + "yes" + ] + }, + { + "id": 4, + "type": "VHS_VideoCombine", + "pos": { + "0": 1045, + "1": 133 + }, + "size": [ + 404.73553466796875, + 601.8645528157551 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 3 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null, + "shape": 7 + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null, + "shape": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 24, + "loop_count": 0, + "filename_prefix": "Ruyi-I2V-StartFrame", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "Ruyi-I2V-StartFrame_00001.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 24 + }, + "muted": false + } + } + }, + { + "id": 3, + "type": "LoadImage", + "pos": { + "0": 200, + "1": 439 + }, + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 2 + ], + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "example_03.jpg", + "image" + ] + }, + { + "id": 2, + "type": "Ruyi_I2VSampler", + "pos": { + "0": 628, + "1": 284 + }, + "size": { + "0": 327.5999755859375, + "1": 338 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "ruyi_model", + "type": "RUYI_MODEL", + "link": 1 + }, + { + "name": "start_img", + "type": "IMAGE", + "link": 2 + }, + { + "name": "end_img", + "type": "IMAGE", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 3 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "Ruyi_I2VSampler" + }, + "widgets_values": [ + 120, + 512, + 925247271358454, + "randomize", + 25, + 7, + "DDIM", + "2", + "static", + "normal_mode", + "5" + ] + } + ], + "links": [ + [ + 1, + 1, + 0, + 2, + 0, + "RUYI_MODEL" + ], + [ + 2, + 3, + 0, + 2, + 1, + "IMAGE" + ], + [ + 3, + 2, + 0, + 4, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1, + "offset": [ + 0, + 0 + ] + } + }, + "version": 0.4 +} diff --git a/config/default.yaml b/config/default.yaml new file mode 100644 index 0000000..15a036e --- /dev/null +++ b/config/default.yaml @@ -0,0 +1,15 @@ +transformer_additional_kwargs: + basic_block_type: "basic" + after_norm: false + time_position_encoding: true + +noise_scheduler_kwargs: + beta_start: 0.00085 + beta_end: 0.03 + beta_schedule: "scaled_linear" + steps_offset: 1 + prediction_type: "v_prediction" + clip_sample: false + +vae_kwargs: + enable_magvit: true diff --git a/predict_i2v.py b/predict_i2v.py new file mode 100644 index 0000000..7c66838 --- /dev/null +++ b/predict_i2v.py @@ -0,0 +1,260 @@ +import os + +import torch +from PIL import Image +from diffusers import (EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, + DPMSolverMultistepScheduler, PNDMScheduler, DDIMScheduler) +from omegaconf import OmegaConf +from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection +from safetensors.torch import load_file as load_safetensors +from huggingface_hub import snapshot_download + +from ruyi.data.bucket_sampler import ASPECT_RATIO_512, get_closest_ratio +from ruyi.models.autoencoder_magvit import AutoencoderKLMagvit +from ruyi.models.transformer3d import HunyuanTransformer3DModel +from ruyi.pipeline.pipeline_ruyi_inpaint import RuyiInpaintPipeline +from ruyi.utils.lora_utils import merge_lora, unmerge_lora +from ruyi.utils.utils import get_image_to_video_latent, save_videos_grid + +# Input and output +start_image_path = "assets/girl_01.jpg" +end_image_path = "assets/girl_02.jpg" # Can be None for start-image-to-video +output_video_path = "outputs/example_01.mp4" + +# Video settings +video_length = 120 # The max video length is 120 frames (24 frames per second) +base_resolution = 512 # # The pixels in the generated video are approximately 512 x 512. Values in the range of [384, 896] typically produce good video quality. +video_size = None # Override base_resolution. Format: [height, width], e.g., [384, 672] +# Control settings +aspect_ratio = "16:9" # Choose in ["16:9", "9:16"], note that this is only the hint +motion = "auto" # Motion control, choose in ["1", "2", "3", "4", "auto"] +camera_direction = "auto" # Camera control, choose in ["static", "left", "right", "up", "down", "auto"] +# Sampler settings +steps = 25 +cfg = 7.0 +scheduler_name = "DDIM" # Choose in ["Euler", "Euler A", "DPM++", "PNDM","DDIM"] + +# GPU memory settings +low_gpu_memory_mode = False # Low gpu memory mode +gpu_offload_steps = 5 # Choose in [0, 10, 7, 5, 1], the latter number requires less GPU memory but longer time + +# Random seed +seed = 42 # The Answer to the Ultimate Question of Life, The Universe, and Everything + +# Model settings +config_path = "config/default.yaml" +model_name = "Ruyi-Mini-7B" +model_type = "Inpaint" +model_path = f"models/{model_name}" # (Down)load mode in this path +auto_download = True + +# LoRA settings +lora_path = None +lora_weight = 1.0 + +# Other settings +weight_dtype = torch.bfloat16 +device = torch.device("cuda") + + +def get_control_embeddings(pipeline, aspect_ratio, motion, camera_direction): + # Default keys + p_default_key = "p.default" + n_default_key = "n.default" + + # Load embeddings + if motion == "auto": + motion = "0" + p_key = f"p.{aspect_ratio.replace(':', 'x')}movie{motion}{camera_direction}" + embeddings = pipeline.embeddings + + # Get embeddings + positive_embeds = embeddings.get(f"{p_key}.emb1", embeddings[f"{p_default_key}.emb1"]) + positive_attention_mask = embeddings.get(f"{p_key}.mask1", embeddings[f"{p_default_key}.mask1"]) + positive_embeds_2 = embeddings.get(f"{p_key}.emb2", embeddings[f"{p_default_key}.emb2"]) + positive_attention_mask_2 = embeddings.get(f"{p_key}.mask2", embeddings[f"{p_default_key}.mask2"]) + + negative_embeds = embeddings[f"{n_default_key}.emb1"] + negative_attention_mask = embeddings[f"{n_default_key}.mask1"] + negative_embeds_2 = embeddings[f"{n_default_key}.emb2"] + negative_attention_mask_2 = embeddings[f"{n_default_key}.mask2"] + + return { + "positive_embeds": positive_embeds, + "positive_attention_mask": positive_attention_mask, + "positive_embeds_2": positive_embeds_2, + "positive_attention_mask_2": positive_attention_mask_2, + + "negative_embeds": negative_embeds, + "negative_attention_mask": negative_attention_mask, + "negative_embeds_2": negative_embeds_2, + "negative_attention_mask_2": negative_attention_mask_2, + } + + +def try_setup_pipeline(model_path, weight_dtype, config): + try: + # Get Vae + vae = AutoencoderKLMagvit.from_pretrained( + model_path, + subfolder="vae" + ).to(weight_dtype) + print("Vae loaded ...") + + # Get Transformer + transformer_additional_kwargs = OmegaConf.to_container(config['transformer_additional_kwargs']) + transformer = HunyuanTransformer3DModel.from_pretrained_2d( + model_path, + subfolder="transformer", + transformer_additional_kwargs=transformer_additional_kwargs + ).to(weight_dtype) + print("Transformer loaded ...") + + # Load Clip + clip_image_encoder = CLIPVisionModelWithProjection.from_pretrained( + model_path, subfolder="image_encoder" + ).to(weight_dtype) + clip_image_processor = CLIPImageProcessor.from_pretrained( + model_path, subfolder="image_encoder" + ) + + # Load sampler and create pipeline + Choosen_Scheduler = DDIMScheduler + scheduler = Choosen_Scheduler.from_pretrained( + model_path, + subfolder="scheduler" + ) + pipeline = RuyiInpaintPipeline.from_pretrained( + model_path, + vae=vae, + transformer=transformer, + scheduler=scheduler, + torch_dtype=weight_dtype, + clip_image_encoder=clip_image_encoder, + clip_image_processor=clip_image_processor, + ) + + # Load embeddings + embeddings = load_safetensors(os.path.join(model_path, "embeddings.safetensors")) + pipeline.embeddings = embeddings + print("Pipeline loaded ...") + + return pipeline + except Exception as e: + print("[Ruyi] Setup pipeline failed:", e) + return None + + +# Init weight_dtype and device +weight_dtype = torch.bfloat16 + +# Load config +config = OmegaConf.load(config_path) + +# Load images +start_img = [Image.open(start_image_path).convert("RGB")] +end_img = [Image.open(end_image_path).convert("RGB")] if end_image_path is not None else None + +# Init model +repo_id = f"IamCreateAI/{model_name}" +pipeline = try_setup_pipeline(model_path, weight_dtype, config) +if pipeline is None and auto_download: + print(f"Downloading {model_name} ...") + + # Download the model + snapshot_download(repo_id=repo_id, local_dir=model_path) + + pipeline = try_setup_pipeline(model_path, weight_dtype, config) + +if pipeline is None: + message = (f"[Load Model Failed] " + f"Please download Ruyi model from huggingface repo '{repo_id}', " + f"And put it into '{model_path}'.") + if not auto_download: + message += "\nOr just set auto_download to 'True'." + raise FileNotFoundError(message) + +# Setup GPU memory mode +if low_gpu_memory_mode: + pipeline.enable_sequential_cpu_offload() +else: + pipeline.enable_model_cpu_offload() + +# Prepare LoRA config +loras = { + 'models': [lora_path] if lora_path is not None else [], + 'weights': [lora_weight] if lora_path is not None else [], +} + +# Count most suitable height and width +if video_size is None: + aspect_ratio_sample_size = {key : [x / 512 * base_resolution for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + original_width, original_height = start_img[0].size if type(start_img) is list else Image.open(start_img).size + closest_size, closest_ratio = get_closest_ratio(original_height, original_width, ratios=aspect_ratio_sample_size) + height, width = [int(x / 16) * 16 for x in closest_size] +else: + height, width = video_size + +# Set hidden states offload steps +if pipeline.model_cpu_offload_flag: + pipeline.transformer.hidden_cache_size = gpu_offload_steps +else: + pipeline.transformer.hidden_cache_size = 0 + +# Load Sampler +if scheduler_name == "DPM++": + noise_scheduler = DPMSolverMultistepScheduler.from_pretrained(model_path, subfolder='scheduler') +elif scheduler_name == "Euler": + noise_scheduler = EulerDiscreteScheduler.from_pretrained(model_path, subfolder='scheduler') +elif scheduler_name == "Euler A": + noise_scheduler = EulerAncestralDiscreteScheduler.from_pretrained(model_path, subfolder='scheduler') +elif scheduler_name == "PNDM": + noise_scheduler = PNDMScheduler.from_pretrained(model_path, subfolder='scheduler') +elif scheduler_name == "DDIM": + noise_scheduler = DDIMScheduler.from_pretrained(model_path, subfolder='scheduler') +pipeline.scheduler = noise_scheduler + +# Set random seed +generator= torch.Generator(device).manual_seed(seed) + +# Load control embeddings +embeddings = get_control_embeddings(pipeline, aspect_ratio, motion, camera_direction) + +with torch.no_grad(): + video_length = int(video_length // pipeline.vae.mini_batch_encoder * pipeline.vae.mini_batch_encoder) if video_length != 1 else 1 + input_video, input_video_mask, clip_image = get_image_to_video_latent(start_img, end_img, video_length=video_length, sample_size=(height, width)) + + for _lora_path, _lora_weight in zip(loras.get("models", []), loras.get("weights", [])): + pipeline = merge_lora(pipeline, _lora_path, _lora_weight) + + sample = pipeline( + prompt_embeds = embeddings["positive_embeds"], + prompt_attention_mask = embeddings["positive_attention_mask"], + prompt_embeds_2 = embeddings["positive_embeds_2"], + prompt_attention_mask_2 = embeddings["positive_attention_mask_2"], + + negative_prompt_embeds = embeddings["negative_embeds"], + negative_prompt_attention_mask = embeddings["negative_attention_mask"], + negative_prompt_embeds_2 = embeddings["negative_embeds_2"], + negative_prompt_attention_mask_2 = embeddings["negative_attention_mask_2"], + + video_length = video_length, + height = height, + width = width, + generator = generator, + guidance_scale = cfg, + num_inference_steps = steps, + + video = input_video, + mask_video = input_video_mask, + clip_image = clip_image, + ).videos + + for _lora_path, _lora_weight in zip(loras.get("models", []), loras.get("weights", [])): + pipeline = unmerge_lora(pipeline, _lora_path, _lora_weight) + +# Save the video +output_folder = os.path.dirname(output_video_path) +if output_folder != '': + os.makedirs(output_folder, exist_ok=True) +save_videos_grid(sample, output_video_path, fps=24) diff --git a/predict_i2v_80g.py b/predict_i2v_80g.py new file mode 100644 index 0000000..67455a9 --- /dev/null +++ b/predict_i2v_80g.py @@ -0,0 +1,260 @@ +import os + +import torch +from PIL import Image +from diffusers import (EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, + DPMSolverMultistepScheduler, PNDMScheduler, DDIMScheduler) +from omegaconf import OmegaConf +from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection +from safetensors.torch import load_file as load_safetensors +from huggingface_hub import snapshot_download + +from ruyi.data.bucket_sampler import ASPECT_RATIO_512, get_closest_ratio +from ruyi.models.autoencoder_magvit import AutoencoderKLMagvit +from ruyi.models.transformer3d import HunyuanTransformer3DModel +from ruyi.pipeline.pipeline_ruyi_inpaint import RuyiInpaintPipeline +from ruyi.utils.lora_utils import merge_lora, unmerge_lora +from ruyi.utils.utils import get_image_to_video_latent, save_videos_grid + +# Input and output +start_image_path = "assets/girl_01.jpg" +end_image_path = "assets/girl_02.jpg" # Can be None for start-image-to-video +output_video_path = "outputs/example_01.mp4" + +# Video settings +video_length = 120 # The max video length is 120 frames (24 frames per second) +base_resolution = 512 # # The pixels in the generated video are approximately 512 x 512. Values in the range of [384, 896] typically produce good video quality. +video_size = None # Override base_resolution. Format: [height, width], e.g., [384, 672] +# Control settings +aspect_ratio = "9:16" # Choose in ["16:9", "9:16"], note that this is only the hint +motion = "auto" # Motion control, choose in ["1", "2", "3", "4", "auto"] +camera_direction = "auto" # Camera control, choose in ["static", "left", "right", "up", "down", "auto"] +# Sampler settings +steps = 25 +cfg = 7.0 +scheduler_name = "DDIM" # Choose in ["Euler", "Euler A", "DPM++", "PNDM","DDIM"] + +# GPU memory settings +low_gpu_memory_mode = False # Low gpu memory mode +gpu_offload_steps = 0 # Choose in [0, 10, 7, 5, 1], the latter number requires less GPU memory but longer time + +# Random seed +seed = 42 # The Answer to the Ultimate Question of Life, The Universe, and Everything + +# Model settings +config_path = "config/default.yaml" +model_name = "Ruyi-Mini-7B" +model_type = "Inpaint" +model_path = f"models/{model_name}" # (Down)load mode in this path +auto_download = True + +# LoRA settings +lora_path = None +lora_weight = 1.0 + +# Other settings +weight_dtype = torch.bfloat16 +device = torch.device("cuda") + + +def get_control_embeddings(pipeline, aspect_ratio, motion, camera_direction): + # Default keys + p_default_key = "p.default" + n_default_key = "n.default" + + # Load embeddings + if motion == "auto": + motion = "0" + p_key = f"p.{aspect_ratio.replace(':', 'x')}movie{motion}{camera_direction}" + embeddings = pipeline.embeddings + + # Get embeddings + positive_embeds = embeddings.get(f"{p_key}.emb1", embeddings[f"{p_default_key}.emb1"]) + positive_attention_mask = embeddings.get(f"{p_key}.mask1", embeddings[f"{p_default_key}.mask1"]) + positive_embeds_2 = embeddings.get(f"{p_key}.emb2", embeddings[f"{p_default_key}.emb2"]) + positive_attention_mask_2 = embeddings.get(f"{p_key}.mask2", embeddings[f"{p_default_key}.mask2"]) + + negative_embeds = embeddings[f"{n_default_key}.emb1"] + negative_attention_mask = embeddings[f"{n_default_key}.mask1"] + negative_embeds_2 = embeddings[f"{n_default_key}.emb2"] + negative_attention_mask_2 = embeddings[f"{n_default_key}.mask2"] + + return { + "positive_embeds": positive_embeds, + "positive_attention_mask": positive_attention_mask, + "positive_embeds_2": positive_embeds_2, + "positive_attention_mask_2": positive_attention_mask_2, + + "negative_embeds": negative_embeds, + "negative_attention_mask": negative_attention_mask, + "negative_embeds_2": negative_embeds_2, + "negative_attention_mask_2": negative_attention_mask_2, + } + + +def try_setup_pipeline(model_path, weight_dtype, config): + try: + # Get Vae + vae = AutoencoderKLMagvit.from_pretrained( + model_path, + subfolder="vae" + ).to(weight_dtype) + print("Vae loaded ...") + + # Get Transformer + transformer_additional_kwargs = OmegaConf.to_container(config['transformer_additional_kwargs']) + transformer = HunyuanTransformer3DModel.from_pretrained_2d( + model_path, + subfolder="transformer", + transformer_additional_kwargs=transformer_additional_kwargs + ).to(weight_dtype) + print("Transformer loaded ...") + + # Load Clip + clip_image_encoder = CLIPVisionModelWithProjection.from_pretrained( + model_path, subfolder="image_encoder" + ).to(weight_dtype) + clip_image_processor = CLIPImageProcessor.from_pretrained( + model_path, subfolder="image_encoder" + ) + + # Load sampler and create pipeline + Choosen_Scheduler = DDIMScheduler + scheduler = Choosen_Scheduler.from_pretrained( + model_path, + subfolder="scheduler" + ) + pipeline = RuyiInpaintPipeline.from_pretrained( + model_path, + vae=vae, + transformer=transformer, + scheduler=scheduler, + torch_dtype=weight_dtype, + clip_image_encoder=clip_image_encoder, + clip_image_processor=clip_image_processor, + ) + + # Load embeddings + embeddings = load_safetensors(os.path.join(model_path, "embeddings.safetensors")) + pipeline.embeddings = embeddings + print("Pipeline loaded ...") + + return pipeline + except Exception as e: + print("[Ruyi] Setup pipeline failed:", e) + return None + + +# Init weight_dtype and device +weight_dtype = torch.bfloat16 + +# Load config +config = OmegaConf.load(config_path) + +# Load images +start_img = [Image.open(start_image_path).convert("RGB")] +end_img = [Image.open(end_image_path).convert("RGB")] if end_image_path is not None else None + +# Init model +repo_id = f"IamCreateAI/{model_name}" +pipeline = try_setup_pipeline(model_path, weight_dtype, config) +if pipeline is None and auto_download: + print(f"Downloading {model_name} ...") + + # Download the model + snapshot_download(repo_id=repo_id, local_dir=model_path) + + pipeline = try_setup_pipeline(model_path, weight_dtype, config) + +if pipeline is None: + message = (f"[Load Model Failed] " + f"Please download Ruyi model from huggingface repo '{repo_id}', " + f"And put it into '{model_path}'.") + if not auto_download: + message += "\nOr just set auto_download to 'True'." + raise FileNotFoundError(message) + +# Setup GPU memory mode +if low_gpu_memory_mode: + pipeline.enable_sequential_cpu_offload() +else: + pipeline.enable_model_cpu_offload() + +# Prepare LoRA config +loras = { + 'models': [lora_path] if lora_path is not None else [], + 'weights': [lora_weight] if lora_path is not None else [], +} + +# Count most suitable height and width +if video_size is None: + aspect_ratio_sample_size = {key : [x / 512 * base_resolution for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()} + original_width, original_height = start_img[0].size if type(start_img) is list else Image.open(start_img).size + closest_size, closest_ratio = get_closest_ratio(original_height, original_width, ratios=aspect_ratio_sample_size) + height, width = [int(x / 16) * 16 for x in closest_size] +else: + height, width = video_size + +# Set hidden states offload steps +if pipeline.model_cpu_offload_flag: + pipeline.transformer.hidden_cache_size = gpu_offload_steps +else: + pipeline.transformer.hidden_cache_size = 0 + +# Load Sampler +if scheduler_name == "DPM++": + noise_scheduler = DPMSolverMultistepScheduler.from_pretrained(model_path, subfolder='scheduler') +elif scheduler_name == "Euler": + noise_scheduler = EulerDiscreteScheduler.from_pretrained(model_path, subfolder='scheduler') +elif scheduler_name == "Euler A": + noise_scheduler = EulerAncestralDiscreteScheduler.from_pretrained(model_path, subfolder='scheduler') +elif scheduler_name == "PNDM": + noise_scheduler = PNDMScheduler.from_pretrained(model_path, subfolder='scheduler') +elif scheduler_name == "DDIM": + noise_scheduler = DDIMScheduler.from_pretrained(model_path, subfolder='scheduler') +pipeline.scheduler = noise_scheduler + +# Set random seed +generator= torch.Generator(device).manual_seed(seed) + +# Load control embeddings +embeddings = get_control_embeddings(pipeline, aspect_ratio, motion, camera_direction) + +with torch.no_grad(): + video_length = int(video_length // pipeline.vae.mini_batch_encoder * pipeline.vae.mini_batch_encoder) if video_length != 1 else 1 + input_video, input_video_mask, clip_image = get_image_to_video_latent(start_img, end_img, video_length=video_length, sample_size=(height, width)) + + for _lora_path, _lora_weight in zip(loras.get("models", []), loras.get("weights", [])): + pipeline = merge_lora(pipeline, _lora_path, _lora_weight) + + sample = pipeline( + prompt_embeds = embeddings["positive_embeds"], + prompt_attention_mask = embeddings["positive_attention_mask"], + prompt_embeds_2 = embeddings["positive_embeds_2"], + prompt_attention_mask_2 = embeddings["positive_attention_mask_2"], + + negative_prompt_embeds = embeddings["negative_embeds"], + negative_prompt_attention_mask = embeddings["negative_attention_mask"], + negative_prompt_embeds_2 = embeddings["negative_embeds_2"], + negative_prompt_attention_mask_2 = embeddings["negative_attention_mask_2"], + + video_length = video_length, + height = height, + width = width, + generator = generator, + guidance_scale = cfg, + num_inference_steps = steps, + + video = input_video, + mask_video = input_video_mask, + clip_image = clip_image, + ).videos + + for _lora_path, _lora_weight in zip(loras.get("models", []), loras.get("weights", [])): + pipeline = unmerge_lora(pipeline, _lora_path, _lora_weight) + +# Save the video +output_folder = os.path.dirname(output_video_path) +if output_folder != '': + os.makedirs(output_folder, exist_ok=True) +save_videos_grid(sample, output_video_path, fps=24) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..13fd833 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,26 @@ +Pillow +einops +safetensors +timm +tomesd +torch +torchdiffeq +torchsde +decord +datasets +numpy +scikit-image +opencv-python +omegaconf +SentencePiece +albumentations +imageio[ffmpeg] +imageio[pyav] +tensorboard +beautifulsoup4 +ftfy +func_timeout +huggingface_hub +accelerate>=0.26.0 +diffusers>=0.28.2 +transformers>=4.37.2 diff --git a/ruyi/__init__.py b/ruyi/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ruyi/data/bucket_sampler.py b/ruyi/data/bucket_sampler.py new file mode 100644 index 0000000..e73b3ad --- /dev/null +++ b/ruyi/data/bucket_sampler.py @@ -0,0 +1,379 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import os +from typing import (Generic, Iterable, Iterator, List, Optional, Sequence, + Sized, TypeVar, Union) + +import cv2 +import numpy as np +import torch +from PIL import Image +from torch.utils.data import BatchSampler, Dataset, Sampler + +ASPECT_RATIO_512 = { + '0.25': [256.0, 1024.0], '0.26': [256.0, 992.0], '0.27': [256.0, 960.0], '0.28': [256.0, 928.0], + '0.32': [288.0, 896.0], '0.33': [288.0, 864.0], '0.35': [288.0, 832.0], '0.4': [320.0, 800.0], + '0.42': [320.0, 768.0], '0.48': [352.0, 736.0], '0.5': [352.0, 704.0], '0.52': [352.0, 672.0], + '0.57': [384.0, 672.0], '0.6': [384.0, 640.0], '0.68': [416.0, 608.0], '0.72': [416.0, 576.0], + '0.78': [448.0, 576.0], '0.82': [448.0, 544.0], '0.88': [480.0, 544.0], '0.94': [480.0, 512.0], + '1.0': [512.0, 512.0], '1.07': [512.0, 480.0], '1.13': [544.0, 480.0], '1.21': [544.0, 448.0], + '1.29': [576.0, 448.0], '1.38': [576.0, 416.0], '1.46': [608.0, 416.0], '1.67': [640.0, 384.0], + '1.75': [672.0, 384.0], '2.0': [704.0, 352.0], '2.09': [736.0, 352.0], '2.4': [768.0, 320.0], + '2.5': [800.0, 320.0], '2.89': [832.0, 288.0], '3.0': [864.0, 288.0], '3.11': [896.0, 288.0], + '3.62': [928.0, 256.0], '3.75': [960.0, 256.0], '3.88': [992.0, 256.0], '4.0': [1024.0, 256.0] +} +ASPECT_RATIO_RANDOM_CROP_512 = { + '0.42': [320.0, 768.0], '0.5': [352.0, 704.0], + '0.57': [384.0, 672.0], '0.68': [416.0, 608.0], '0.78': [448.0, 576.0], '0.88': [480.0, 544.0], + '0.94': [480.0, 512.0], '1.0': [512.0, 512.0], '1.07': [512.0, 480.0], + '1.13': [544.0, 480.0], '1.29': [576.0, 448.0], '1.46': [608.0, 416.0], '1.75': [672.0, 384.0], + '2.0': [704.0, 352.0], '2.4': [768.0, 320.0] +} +ASPECT_RATIO_RANDOM_CROP_PROB = [ + 1, 2, + 4, 4, 4, 4, + 8, 8, 8, + 4, 4, 4, 4, + 2, 1 +] +ASPECT_RATIO_RANDOM_CROP_PROB = np.array(ASPECT_RATIO_RANDOM_CROP_PROB) / sum(ASPECT_RATIO_RANDOM_CROP_PROB) + +def get_closest_ratio(height: float, width: float, ratios: dict = ASPECT_RATIO_512): + aspect_ratio = height / width + closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - aspect_ratio)) + return ratios[closest_ratio], float(closest_ratio) + +def get_image_size_without_loading(path): + with Image.open(path) as img: + return img.size # (width, height) + +class RandomSampler(Sampler[int]): + r"""Samples elements randomly. If without replacement, then sample from a shuffled dataset. + + If with replacement, then user can specify :attr:`num_samples` to draw. + + Args: + data_source (Dataset): dataset to sample from + replacement (bool): samples are drawn on-demand with replacement if ``True``, default=``False`` + num_samples (int): number of samples to draw, default=`len(dataset)`. + generator (Generator): Generator used in sampling. + """ + + data_source: Sized + replacement: bool + + def __init__(self, data_source: Sized, replacement: bool = False, + num_samples: Optional[int] = None, generator=None) -> None: + self.data_source = data_source + self.replacement = replacement + self._num_samples = num_samples + self.generator = generator + self._pos_start = 0 + + if not isinstance(self.replacement, bool): + raise TypeError(f"replacement should be a boolean value, but got replacement={self.replacement}") + + if not isinstance(self.num_samples, int) or self.num_samples <= 0: + raise ValueError(f"num_samples should be a positive integer value, but got num_samples={self.num_samples}") + + @property + def num_samples(self) -> int: + # dataset size might change at runtime + if self._num_samples is None: + return len(self.data_source) + return self._num_samples + + def __iter__(self) -> Iterator[int]: + n = len(self.data_source) + if self.generator is None: + seed = int(torch.empty((), dtype=torch.int64).random_().item()) + generator = torch.Generator() + generator.manual_seed(seed) + else: + generator = self.generator + + if self.replacement: + for _ in range(self.num_samples // 32): + yield from torch.randint(high=n, size=(32,), dtype=torch.int64, generator=generator).tolist() + yield from torch.randint(high=n, size=(self.num_samples % 32,), dtype=torch.int64, generator=generator).tolist() + else: + for _ in range(self.num_samples // n): + xx = torch.randperm(n, generator=generator).tolist() + if self._pos_start >= n: + self._pos_start = 0 + print("xx top 10", xx[:10], self._pos_start) + for idx in range(self._pos_start, n): + yield xx[idx] + self._pos_start = (self._pos_start + 1) % n + self._pos_start = 0 + yield from torch.randperm(n, generator=generator).tolist()[:self.num_samples % n] + + def __len__(self) -> int: + return self.num_samples + +class AspectRatioBatchImageSampler(BatchSampler): + """A sampler wrapper for grouping images with similar aspect ratio into a same batch. + + Args: + sampler (Sampler): Base sampler. + dataset (Dataset): Dataset providing data information. + batch_size (int): Size of mini-batch. + drop_last (bool): If ``True``, the sampler will drop the last batch if + its size would be less than ``batch_size``. + aspect_ratios (dict): The predefined aspect ratios. + """ + def __init__( + self, + sampler: Sampler, + dataset: Dataset, + batch_size: int, + train_folder: str = None, + aspect_ratios: dict = ASPECT_RATIO_512, + drop_last: bool = False, + config=None, + **kwargs + ) -> None: + if not isinstance(sampler, Sampler): + raise TypeError('sampler should be an instance of ``Sampler``, ' + f'but got {sampler}') + if not isinstance(batch_size, int) or batch_size <= 0: + raise ValueError('batch_size should be a positive integer value, ' + f'but got batch_size={batch_size}') + self.sampler = sampler + self.dataset = dataset + self.train_folder = train_folder + self.batch_size = batch_size + self.aspect_ratios = aspect_ratios + self.drop_last = drop_last + self.config = config + # buckets for each aspect ratio + self._aspect_ratio_buckets = {ratio: [] for ratio in aspect_ratios} + # [str(k) for k, v in aspect_ratios] + self.current_available_bucket_keys = list(aspect_ratios.keys()) + + def __iter__(self): + for idx in self.sampler: + try: + image_dict = self.dataset[idx] + + width, height = image_dict.get("weight", None), image_dict.get("height", None) + if width is None or height is None: + image_id, name = image_dict['file_path'], image_dict['text'] + if self.train_folder is None: + image_dir = image_id + else: + image_dir = os.path.join(self.train_folder, image_id) + + width, height = get_image_size_without_loading(image_dir) + + ratio = height / width # self.dataset[idx] + else: + height = int(height) + width = int(width) + ratio = height / width # self.dataset[idx] + except Exception as e: + print(e) + continue + # find the closest aspect ratio + closest_ratio = min(self.aspect_ratios.keys(), key=lambda r: abs(float(r) - ratio)) + if closest_ratio not in self.current_available_bucket_keys: + continue + bucket = self._aspect_ratio_buckets[closest_ratio] + bucket.append(idx) + # yield a batch of indices in the same aspect ratio group + if len(bucket) == self.batch_size: + yield bucket[:] + del bucket[:] + +class AspectRatioBatchSampler(BatchSampler): + """A sampler wrapper for grouping images with similar aspect ratio into a same batch. + + Args: + sampler (Sampler): Base sampler. + dataset (Dataset): Dataset providing data information. + batch_size (int): Size of mini-batch. + drop_last (bool): If ``True``, the sampler will drop the last batch if + its size would be less than ``batch_size``. + aspect_ratios (dict): The predefined aspect ratios. + """ + def __init__( + self, + sampler: Sampler, + dataset: Dataset, + batch_size: int, + video_folder: str = None, + train_data_format: str = "webvid", + aspect_ratios: dict = ASPECT_RATIO_512, + drop_last: bool = False, + config=None, + **kwargs + ) -> None: + if not isinstance(sampler, Sampler): + raise TypeError('sampler should be an instance of ``Sampler``, ' + f'but got {sampler}') + if not isinstance(batch_size, int) or batch_size <= 0: + raise ValueError('batch_size should be a positive integer value, ' + f'but got batch_size={batch_size}') + self.sampler = sampler + self.dataset = dataset + self.video_folder = video_folder + self.train_data_format = train_data_format + self.batch_size = batch_size + self.aspect_ratios = aspect_ratios + self.drop_last = drop_last + self.config = config + # buckets for each aspect ratio + self._aspect_ratio_buckets = {ratio: [] for ratio in aspect_ratios} + # [str(k) for k, v in aspect_ratios] + self.current_available_bucket_keys = list(aspect_ratios.keys()) + + def __iter__(self): + for idx in self.sampler: + try: + video_dict = self.dataset[idx] + width, more = video_dict.get("width", None), video_dict.get("height", None) + + if width is None or height is None: + if self.train_data_format == "normal": + video_id, name = video_dict['file_path'], video_dict['text'] + if self.video_folder is None: + video_dir = video_id + else: + video_dir = os.path.join(self.video_folder, video_id) + else: + videoid, name, page_dir = video_dict['videoid'], video_dict['name'], video_dict['page_dir'] + video_dir = os.path.join(self.video_folder, f"{videoid}.mp4") + cap = cv2.VideoCapture(video_dir) + + # 获取视频尺寸 + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) # 浮点数转换为整数 + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) # 浮点数转换为整数 + + ratio = height / width # self.dataset[idx] + else: + height = int(height) + width = int(width) + ratio = height / width # self.dataset[idx] + except Exception as e: + print(e) + continue + # find the closest aspect ratio + closest_ratio = min(self.aspect_ratios.keys(), key=lambda r: abs(float(r) - ratio)) + if closest_ratio not in self.current_available_bucket_keys: + continue + bucket = self._aspect_ratio_buckets[closest_ratio] + bucket.append(idx) + # yield a batch of indices in the same aspect ratio group + if len(bucket) == self.batch_size: + yield bucket[:] + del bucket[:] + +class AspectRatioBatchImageVideoSampler(BatchSampler): + """A sampler wrapper for grouping images with similar aspect ratio into a same batch. + + Args: + sampler (Sampler): Base sampler. + dataset (Dataset): Dataset providing data information. + batch_size (int): Size of mini-batch. + drop_last (bool): If ``True``, the sampler will drop the last batch if + its size would be less than ``batch_size``. + aspect_ratios (dict): The predefined aspect ratios. + """ + + def __init__(self, + sampler: Sampler, + dataset: Dataset, + batch_size: int, + train_folder: str = None, + aspect_ratios: dict = ASPECT_RATIO_512, + drop_last: bool = False + ) -> None: + if not isinstance(sampler, Sampler): + raise TypeError('sampler should be an instance of ``Sampler``, ' + f'but got {sampler}') + if not isinstance(batch_size, int) or batch_size <= 0: + raise ValueError('batch_size should be a positive integer value, ' + f'but got batch_size={batch_size}') + self.sampler = sampler + self.dataset = dataset + self.train_folder = train_folder + self.batch_size = batch_size + self.aspect_ratios = aspect_ratios + self.drop_last = drop_last + + # buckets for each aspect ratio + self.current_available_bucket_keys = list(aspect_ratios.keys()) + self.bucket = { + 'image':{ratio: [] for ratio in aspect_ratios}, + 'video':{ratio: [] for ratio in aspect_ratios} + } + + def __iter__(self): + for idx in self.sampler: + content_type = self.dataset[idx].get('type', 'image') + if content_type == 'image': + try: + image_dict = self.dataset[idx] + + width, height = image_dict.get("width", None), image_dict.get("height", None) + if width is None or height is None: + image_id, name = image_dict['file_path'], image_dict['text'] + if self.train_folder is None: + image_dir = image_id + else: + image_dir = os.path.join(self.train_folder, image_id) + + width, height = get_image_size_without_loading(image_dir) + + ratio = height / width # self.dataset[idx] + else: + height = int(height) + width = int(width) + ratio = height / width # self.dataset[idx] + except Exception as e: + print(e) + continue + # find the closest aspect ratio + closest_ratio = min(self.aspect_ratios.keys(), key=lambda r: abs(float(r) - ratio)) + if closest_ratio not in self.current_available_bucket_keys: + continue + bucket = self.bucket['image'][closest_ratio] + bucket.append(idx) + # yield a batch of indices in the same aspect ratio group + if len(bucket) == self.batch_size: + yield bucket[:] + del bucket[:] + else: + try: + video_dict = self.dataset[idx] + width, height = video_dict.get("width", None), video_dict.get("height", None) + + if width is None or height is None: + video_id, name = video_dict['file_path'], video_dict['text'] + if self.train_folder is None: + video_dir = video_id + else: + video_dir = os.path.join(self.train_folder, video_id) + cap = cv2.VideoCapture(video_dir) + + # 获取视频尺寸 + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) # 浮点数转换为整数 + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) # 浮点数转换为整数 + + ratio = height / width # self.dataset[idx] + else: + height = int(height) + width = int(width) + ratio = height / width # self.dataset[idx] + except Exception as e: + print(e) + continue + # find the closest aspect ratio + closest_ratio = min(self.aspect_ratios.keys(), key=lambda r: abs(float(r) - ratio)) + if closest_ratio not in self.current_available_bucket_keys: + continue + bucket = self.bucket['video'][closest_ratio] + bucket.append(idx) + # yield a batch of indices in the same aspect ratio group + if len(bucket) == self.batch_size: + yield bucket[:] + del bucket[:] \ No newline at end of file diff --git a/ruyi/models/__init__.py b/ruyi/models/__init__.py new file mode 100644 index 0000000..d4265f8 --- /dev/null +++ b/ruyi/models/__init__.py @@ -0,0 +1,10 @@ +from .attention import * +from .transformer2d import * +from .transformer3d import * +from .autoencoder_magvit import * +from .embeddings import * +from .motion_module import * +from .norm import * +from .patch import * +from .resampler import * + diff --git a/ruyi/models/attention.py b/ruyi/models/attention.py new file mode 100644 index 0000000..8630f0c --- /dev/null +++ b/ruyi/models/attention.py @@ -0,0 +1,1820 @@ +# Copyright 2023 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. +from typing import Any, Dict, Optional + +import diffusers +import pkg_resources +import torch +import torch.nn.functional as F +import torch.nn.init as init + +installed_version = diffusers.__version__ + +if pkg_resources.parse_version(installed_version) >= pkg_resources.parse_version("0.28.2"): + from diffusers.models.attention_processor import (Attention, + AttnProcessor2_0, + HunyuanAttnProcessor2_0) +else: + from diffusers.models.attention_processor import Attention, AttnProcessor2_0 + +from diffusers.models.attention import AdaLayerNorm, FeedForward +from diffusers.models.embeddings import SinusoidalPositionalEmbedding +from diffusers.models.normalization import AdaLayerNormZero +from diffusers.utils import USE_PEFT_BACKEND +from diffusers.utils.import_utils import is_xformers_available +from diffusers.utils.torch_utils import maybe_allow_in_graph +from einops import rearrange, repeat +from torch import nn + +from .motion_module import PositionalEncoding, get_motion_module +from .norm import FP32LayerNorm, AdaLayerNormShift + +if is_xformers_available(): + import xformers + import xformers.ops +else: + xformers = None + + + +@maybe_allow_in_graph +class GatedSelfAttentionDense(nn.Module): + r""" + A gated self-attention dense layer that combines visual features and object features. + + Parameters: + query_dim (`int`): The number of channels in the query. + context_dim (`int`): The number of channels in the context. + n_heads (`int`): The number of heads to use for attention. + d_head (`int`): The number of channels in each head. + """ + + def __init__(self, query_dim: int, context_dim: int, n_heads: int, d_head: int): + super().__init__() + + # we need a linear projection since we need cat visual feature and obj feature + self.linear = nn.Linear(context_dim, query_dim) + + self.attn = Attention(query_dim=query_dim, heads=n_heads, dim_head=d_head) + self.ff = FeedForward(query_dim, activation_fn="geglu") + + self.norm1 = FP32LayerNorm(query_dim) + self.norm2 = FP32LayerNorm(query_dim) + + self.register_parameter("alpha_attn", nn.Parameter(torch.tensor(0.0))) + self.register_parameter("alpha_dense", nn.Parameter(torch.tensor(0.0))) + + self.enabled = True + + def forward(self, x: torch.Tensor, objs: torch.Tensor) -> torch.Tensor: + if not self.enabled: + return x + + n_visual = x.shape[1] + objs = self.linear(objs) + + x = x + self.alpha_attn.tanh() * self.attn(self.norm1(torch.cat([x, objs], dim=1)))[:, :n_visual, :] + x = x + self.alpha_dense.tanh() * self.ff(self.norm2(x)) + + return x + + +def zero_module(module): + # Zero out the parameters of a module and return it. + for p in module.parameters(): + p.detach().zero_() + return module + + +class KVCompressionCrossAttention(nn.Module): + r""" + A cross attention layer. + + Parameters: + query_dim (`int`): The number of channels in the query. + cross_attention_dim (`int`, *optional*): + The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`. + heads (`int`, *optional*, defaults to 8): The number of heads to use for multi-head attention. + dim_head (`int`, *optional*, defaults to 64): The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + bias (`bool`, *optional*, defaults to False): + Set to `True` for the query, key, and value linear layers to contain a bias parameter. + """ + + def __init__( + self, + query_dim: int, + cross_attention_dim: Optional[int] = None, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + bias=False, + upcast_attention: bool = False, + upcast_softmax: bool = False, + added_kv_proj_dim: Optional[int] = None, + norm_num_groups: Optional[int] = None, + ): + super().__init__() + inner_dim = dim_head * heads + cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim + self.upcast_attention = upcast_attention + self.upcast_softmax = upcast_softmax + + self.scale = dim_head**-0.5 + + self.heads = heads + # for slice_size > 0 the attention score computation + # is split across the batch axis to save memory + # You can set slice_size with `set_attention_slice` + self.sliceable_head_dim = heads + self._slice_size = None + self._use_memory_efficient_attention_xformers = True + self.added_kv_proj_dim = added_kv_proj_dim + + if norm_num_groups is not None: + self.group_norm = nn.GroupNorm(num_channels=inner_dim, num_groups=norm_num_groups, eps=1e-5, affine=True) + else: + self.group_norm = None + + self.to_q = nn.Linear(query_dim, inner_dim, bias=bias) + self.to_k = nn.Linear(cross_attention_dim, inner_dim, bias=bias) + self.to_v = nn.Linear(cross_attention_dim, inner_dim, bias=bias) + + if self.added_kv_proj_dim is not None: + self.add_k_proj = nn.Linear(added_kv_proj_dim, cross_attention_dim) + self.add_v_proj = nn.Linear(added_kv_proj_dim, cross_attention_dim) + + self.kv_compression = nn.Conv2d( + query_dim, + query_dim, + groups=query_dim, + kernel_size=2, + stride=2, + bias=True + ) + self.kv_compression_norm = FP32LayerNorm(query_dim) + init.constant_(self.kv_compression.weight, 1 / 4) + if self.kv_compression.bias is not None: + init.constant_(self.kv_compression.bias, 0) + + self.to_out = nn.ModuleList([]) + self.to_out.append(nn.Linear(inner_dim, query_dim)) + self.to_out.append(nn.Dropout(dropout)) + + def reshape_heads_to_batch_dim(self, tensor): + batch_size, seq_len, dim = tensor.shape + head_size = self.heads + tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size) + tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size * head_size, seq_len, dim // head_size) + return tensor + + def reshape_batch_dim_to_heads(self, tensor): + batch_size, seq_len, dim = tensor.shape + head_size = self.heads + tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim) + tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size) + return tensor + + def set_attention_slice(self, slice_size): + if slice_size is not None and slice_size > self.sliceable_head_dim: + raise ValueError(f"slice_size {slice_size} has to be smaller or equal to {self.sliceable_head_dim}.") + + self._slice_size = slice_size + + def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, num_frames: int = 16, height: int = 32, width: int = 32): + batch_size, sequence_length, _ = hidden_states.shape + + encoder_hidden_states = encoder_hidden_states + + if self.group_norm is not None: + hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + + query = self.to_q(hidden_states) + dim = query.shape[-1] + query = self.reshape_heads_to_batch_dim(query) + + if self.added_kv_proj_dim is not None: + key = self.to_k(hidden_states) + value = self.to_v(hidden_states) + + encoder_hidden_states_key_proj = self.add_k_proj(encoder_hidden_states) + encoder_hidden_states_value_proj = self.add_v_proj(encoder_hidden_states) + + key = rearrange(key, "b (f h w) c -> (b f) c h w", f=num_frames, h=height, w=width) + key = self.kv_compression(key) + key = rearrange(key, "(b f) c h w -> b (f h w) c", f=num_frames) + key = self.kv_compression_norm(key) + key = key.to(query.dtype) + + value = rearrange(value, "b (f h w) c -> (b f) c h w", f=num_frames, h=height, w=width) + value = self.kv_compression(value) + value = rearrange(value, "(b f) c h w -> b (f h w) c", f=num_frames) + value = self.kv_compression_norm(value) + value = value.to(query.dtype) + + key = self.reshape_heads_to_batch_dim(key) + value = self.reshape_heads_to_batch_dim(value) + encoder_hidden_states_key_proj = self.reshape_heads_to_batch_dim(encoder_hidden_states_key_proj) + encoder_hidden_states_value_proj = self.reshape_heads_to_batch_dim(encoder_hidden_states_value_proj) + + key = torch.concat([encoder_hidden_states_key_proj, key], dim=1) + value = torch.concat([encoder_hidden_states_value_proj, value], dim=1) + else: + encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states + key = self.to_k(encoder_hidden_states) + value = self.to_v(encoder_hidden_states) + + key = rearrange(key, "b (f h w) c -> (b f) c h w", f=num_frames, h=height, w=width) + key = self.kv_compression(key) + key = rearrange(key, "(b f) c h w -> b (f h w) c", f=num_frames) + key = self.kv_compression_norm(key) + key = key.to(query.dtype) + + value = rearrange(value, "b (f h w) c -> (b f) c h w", f=num_frames, h=height, w=width) + value = self.kv_compression(value) + value = rearrange(value, "(b f) c h w -> b (f h w) c", f=num_frames) + value = self.kv_compression_norm(value) + value = value.to(query.dtype) + + key = self.reshape_heads_to_batch_dim(key) + value = self.reshape_heads_to_batch_dim(value) + + if attention_mask is not None: + if attention_mask.shape[-1] != query.shape[1]: + target_length = query.shape[1] + attention_mask = F.pad(attention_mask, (0, target_length), value=0.0) + attention_mask = attention_mask.repeat_interleave(self.heads, dim=0) + + # attention, what we cannot get enough of + if self._use_memory_efficient_attention_xformers: + hidden_states = self._memory_efficient_attention_xformers(query, key, value, attention_mask) + # Some versions of xformers return output in fp32, cast it back to the dtype of the input + hidden_states = hidden_states.to(query.dtype) + else: + if self._slice_size is None or query.shape[0] // self._slice_size == 1: + hidden_states = self._attention(query, key, value, attention_mask) + else: + hidden_states = self._sliced_attention(query, key, value, sequence_length, dim, attention_mask) + + # linear proj + hidden_states = self.to_out[0](hidden_states) + + # dropout + hidden_states = self.to_out[1](hidden_states) + return hidden_states + + def _attention(self, query, key, value, attention_mask=None): + if self.upcast_attention: + query = query.float() + key = key.float() + + attention_scores = torch.baddbmm( + torch.empty(query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device), + query, + key.transpose(-1, -2), + beta=0, + alpha=self.scale, + ) + + if attention_mask is not None: + attention_scores = attention_scores + attention_mask + + if self.upcast_softmax: + attention_scores = attention_scores.float() + + attention_probs = attention_scores.softmax(dim=-1) + + # cast back to the original dtype + attention_probs = attention_probs.to(value.dtype) + + # compute attention output + hidden_states = torch.bmm(attention_probs, value) + + # reshape hidden_states + hidden_states = self.reshape_batch_dim_to_heads(hidden_states) + return hidden_states + + def _sliced_attention(self, query, key, value, sequence_length, dim, attention_mask): + batch_size_attention = query.shape[0] + hidden_states = torch.zeros( + (batch_size_attention, sequence_length, dim // self.heads), device=query.device, dtype=query.dtype + ) + slice_size = self._slice_size if self._slice_size is not None else hidden_states.shape[0] + for i in range(hidden_states.shape[0] // slice_size): + start_idx = i * slice_size + end_idx = (i + 1) * slice_size + + query_slice = query[start_idx:end_idx] + key_slice = key[start_idx:end_idx] + + if self.upcast_attention: + query_slice = query_slice.float() + key_slice = key_slice.float() + + attn_slice = torch.baddbmm( + torch.empty(slice_size, query.shape[1], key.shape[1], dtype=query_slice.dtype, device=query.device), + query_slice, + key_slice.transpose(-1, -2), + beta=0, + alpha=self.scale, + ) + + if attention_mask is not None: + attn_slice = attn_slice + attention_mask[start_idx:end_idx] + + if self.upcast_softmax: + attn_slice = attn_slice.float() + + attn_slice = attn_slice.softmax(dim=-1) + + # cast back to the original dtype + attn_slice = attn_slice.to(value.dtype) + attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx]) + + hidden_states[start_idx:end_idx] = attn_slice + + # reshape hidden_states + hidden_states = self.reshape_batch_dim_to_heads(hidden_states) + return hidden_states + + def _memory_efficient_attention_xformers(self, query, key, value, attention_mask): + # TODO attention_mask + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + hidden_states = xformers.ops.memory_efficient_attention(query, key, value, attn_bias=attention_mask) + hidden_states = self.reshape_batch_dim_to_heads(hidden_states) + return hidden_states + + +@maybe_allow_in_graph +class TemporalTransformerBlock(nn.Module): + r""" + A Temporal Transformer block. + + Parameters: + dim (`int`): The number of channels in the input and output. + num_attention_heads (`int`): The number of heads to use for multi-head attention. + attention_head_dim (`int`): The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + num_embeds_ada_norm (: + obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. + attention_bias (: + obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. + only_cross_attention (`bool`, *optional*): + Whether to use only cross-attention layers. In this case two cross attention layers are used. + double_self_attention (`bool`, *optional*): + Whether to use two self-attention layers. In this case no cross attention layers are used. + upcast_attention (`bool`, *optional*): + Whether to upcast the attention computation to float32. This is useful for mixed precision training. + norm_elementwise_affine (`bool`, *optional*, defaults to `True`): + Whether to use learnable elementwise affine parameters for normalization. + norm_type (`str`, *optional*, defaults to `"layer_norm"`): + The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`. + final_dropout (`bool` *optional*, defaults to False): + Whether to apply a final dropout after the last feed-forward layer. + attention_type (`str`, *optional*, defaults to `"default"`): + The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`. + positional_embeddings (`str`, *optional*, defaults to `None`): + The type of positional embeddings to apply to. + num_positional_embeddings (`int`, *optional*, defaults to `None`): + The maximum number of positional embeddings to apply. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout=0.0, + cross_attention_dim: Optional[int] = None, + activation_fn: str = "geglu", + num_embeds_ada_norm: Optional[int] = None, + attention_bias: bool = False, + only_cross_attention: bool = False, + double_self_attention: bool = False, + upcast_attention: bool = False, + norm_elementwise_affine: bool = True, + norm_type: str = "layer_norm", # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single' + norm_eps: float = 1e-5, + final_dropout: bool = False, + attention_type: str = "default", + positional_embeddings: Optional[str] = None, + num_positional_embeddings: Optional[int] = None, + # kv compression + kvcompression: Optional[bool] = False, + # motion module kwargs + motion_module_type = "VanillaGrid", + motion_module_kwargs = None, + qk_norm = False, + after_norm = False, + ): + super().__init__() + self.only_cross_attention = only_cross_attention + + self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero" + self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" + self.use_ada_layer_norm_single = norm_type == "ada_norm_single" + self.use_layer_norm = norm_type == "layer_norm" + + if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: + raise ValueError( + f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" + f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." + ) + + if positional_embeddings and (num_positional_embeddings is None): + raise ValueError( + "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined." + ) + + if positional_embeddings == "sinusoidal": + self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings) + else: + self.pos_embed = None + + # Define 3 blocks. Each block has its own normalization layer. + # 1. Self-Attn + if self.use_ada_layer_norm: + self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) + elif self.use_ada_layer_norm_zero: + self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) + else: + self.norm1 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + + self.kvcompression = kvcompression + if kvcompression: + self.attn1 = KVCompressionCrossAttention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim if only_cross_attention else None, + upcast_attention=upcast_attention, + ) + else: + if pkg_resources.parse_version(installed_version) >= pkg_resources.parse_version("0.28.2"): + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim if only_cross_attention else None, + upcast_attention=upcast_attention, + qk_norm="layer_norm" if qk_norm else None, + processor=HunyuanAttnProcessor2_0() if qk_norm else AttnProcessor2_0(), + ) + else: + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim if only_cross_attention else None, + upcast_attention=upcast_attention, + ) + + self.attn_temporal = get_motion_module( + in_channels = dim, + motion_module_type = motion_module_type, + motion_module_kwargs = motion_module_kwargs, + ) + + # 2. Cross-Attn + if cross_attention_dim is not None or double_self_attention: + # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. + # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during + # the second cross attention block. + self.norm2 = ( + AdaLayerNorm(dim, num_embeds_ada_norm) + if self.use_ada_layer_norm + else FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + ) + if pkg_resources.parse_version(installed_version) >= pkg_resources.parse_version("0.28.2"): + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim if not double_self_attention else None, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + qk_norm="layer_norm" if qk_norm else None, + processor=HunyuanAttnProcessor2_0() if qk_norm else AttnProcessor2_0(), + ) # is self-attn if encoder_hidden_states is none + else: + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim if not double_self_attention else None, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + ) # is self-attn if encoder_hidden_states is none + else: + self.norm2 = None + self.attn2 = None + + # 3. Feed-forward + if not self.use_ada_layer_norm_single: + self.norm3 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + + self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout) + + if after_norm: + self.norm4 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + else: + self.norm4 = None + + # 4. Fuser + if attention_type == "gated" or attention_type == "gated-text-image": + self.fuser = GatedSelfAttentionDense(dim, cross_attention_dim, num_attention_heads, attention_head_dim) + + # 5. Scale-shift for PixArt-Alpha. + if self.use_ada_layer_norm_single: + self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) + + # let chunk size default to None + self._chunk_size = None + self._chunk_dim = 0 + + def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int): + # Sets chunk feed-forward + self._chunk_size = chunk_size + self._chunk_dim = dim + + def forward( + self, + hidden_states: torch.FloatTensor, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + timestep: Optional[torch.LongTensor] = None, + cross_attention_kwargs: Dict[str, Any] = None, + class_labels: Optional[torch.LongTensor] = None, + num_frames: int = 16, + height: int = 32, + width: int = 32, + ) -> torch.FloatTensor: + # Notice that normalization is always applied before the real computation in the following blocks. + # 0. Self-Attention + batch_size = hidden_states.shape[0] + + if self.use_ada_layer_norm: + norm_hidden_states = self.norm1(hidden_states, timestep) + elif self.use_ada_layer_norm_zero: + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( + hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + elif self.use_layer_norm: + norm_hidden_states = self.norm1(hidden_states) + elif self.use_ada_layer_norm_single: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) + ).chunk(6, dim=1) + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + norm_hidden_states = norm_hidden_states.squeeze(1) + else: + raise ValueError("Incorrect norm used") + + if self.pos_embed is not None: + norm_hidden_states = self.pos_embed(norm_hidden_states) + + # 1. Retrieve lora scale. + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + # 2. Prepare GLIGEN inputs + cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} + gligen_kwargs = cross_attention_kwargs.pop("gligen", None) + + norm_hidden_states = rearrange(norm_hidden_states, "b (f d) c -> (b f) d c", f=num_frames) + if self.kvcompression: + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, + attention_mask=attention_mask, + num_frames=1, + height=height, + width=width, + **cross_attention_kwargs, + ) + else: + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + attn_output = rearrange(attn_output, "(b f) d c -> b (f d) c", f=num_frames) + if self.use_ada_layer_norm_zero: + attn_output = gate_msa.unsqueeze(1) * attn_output + elif self.use_ada_layer_norm_single: + attn_output = gate_msa * attn_output + + hidden_states = attn_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + # 2.5 GLIGEN Control + if gligen_kwargs is not None: + hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"]) + + # 2.75. Temp-Attention + if self.attn_temporal is not None: + attn_output = rearrange(hidden_states, "b (f h w) c -> b c f h w", f=num_frames, h=height, w=width) + attn_output = self.attn_temporal(attn_output) + hidden_states = rearrange(attn_output, "b c f h w -> b (f h w) c") + + # 3. Cross-Attention + if self.attn2 is not None: + if self.use_ada_layer_norm: + norm_hidden_states = self.norm2(hidden_states, timestep) + elif self.use_ada_layer_norm_zero or self.use_layer_norm: + norm_hidden_states = self.norm2(hidden_states) + elif self.use_ada_layer_norm_single: + # For PixArt norm2 isn't applied here: + # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103 + norm_hidden_states = hidden_states + else: + raise ValueError("Incorrect norm") + + if self.pos_embed is not None and self.use_ada_layer_norm_single is None: + norm_hidden_states = self.pos_embed(norm_hidden_states) + + if norm_hidden_states.dtype != encoder_hidden_states.dtype or norm_hidden_states.dtype != encoder_attention_mask.dtype: + norm_hidden_states = norm_hidden_states.to(encoder_hidden_states.dtype) + + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + **cross_attention_kwargs, + ) + hidden_states = attn_output + hidden_states + + # 4. Feed-forward + if not self.use_ada_layer_norm_single: + norm_hidden_states = self.norm3(hidden_states) + + if self.use_ada_layer_norm_zero: + norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + + if self.use_ada_layer_norm_single: + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + + if self._chunk_size is not None: + # "feed_forward_chunk_size" can be used to save memory + if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0: + raise ValueError( + f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." + ) + + num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size + ff_output = torch.cat( + [ + self.ff(hid_slice, scale=lora_scale) + for hid_slice in norm_hidden_states.chunk(num_chunks, dim=self._chunk_dim) + ], + dim=self._chunk_dim, + ) + else: + ff_output = self.ff(norm_hidden_states, scale=lora_scale) + + if self.norm4 is not None: + ff_output = self.norm4(ff_output) + + if self.use_ada_layer_norm_zero: + ff_output = gate_mlp.unsqueeze(1) * ff_output + elif self.use_ada_layer_norm_single: + ff_output = gate_mlp * ff_output + + hidden_states = ff_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + return hidden_states + + +@maybe_allow_in_graph +class SelfAttentionTemporalTransformerBlock(nn.Module): + r""" + A Temporal Transformer block. + + Parameters: + dim (`int`): The number of channels in the input and output. + num_attention_heads (`int`): The number of heads to use for multi-head attention. + attention_head_dim (`int`): The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + num_embeds_ada_norm (: + obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. + attention_bias (: + obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. + only_cross_attention (`bool`, *optional*): + Whether to use only cross-attention layers. In this case two cross attention layers are used. + double_self_attention (`bool`, *optional*): + Whether to use two self-attention layers. In this case no cross attention layers are used. + upcast_attention (`bool`, *optional*): + Whether to upcast the attention computation to float32. This is useful for mixed precision training. + norm_elementwise_affine (`bool`, *optional*, defaults to `True`): + Whether to use learnable elementwise affine parameters for normalization. + norm_type (`str`, *optional*, defaults to `"layer_norm"`): + The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`. + final_dropout (`bool` *optional*, defaults to False): + Whether to apply a final dropout after the last feed-forward layer. + attention_type (`str`, *optional*, defaults to `"default"`): + The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`. + positional_embeddings (`str`, *optional*, defaults to `None`): + The type of positional embeddings to apply to. + num_positional_embeddings (`int`, *optional*, defaults to `None`): + The maximum number of positional embeddings to apply. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout=0.0, + cross_attention_dim: Optional[int] = None, + activation_fn: str = "geglu", + num_embeds_ada_norm: Optional[int] = None, + attention_bias: bool = False, + only_cross_attention: bool = False, + double_self_attention: bool = False, + upcast_attention: bool = False, + norm_elementwise_affine: bool = True, + norm_type: str = "layer_norm", # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single' + norm_eps: float = 1e-5, + final_dropout: bool = False, + attention_type: str = "default", + positional_embeddings: Optional[str] = None, + num_positional_embeddings: Optional[int] = None, + qk_norm = False, + after_norm = False, + ): + super().__init__() + self.only_cross_attention = only_cross_attention + + self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero" + self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" + self.use_ada_layer_norm_single = norm_type == "ada_norm_single" + self.use_layer_norm = norm_type == "layer_norm" + + if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: + raise ValueError( + f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" + f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." + ) + + if positional_embeddings and (num_positional_embeddings is None): + raise ValueError( + "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined." + ) + + if positional_embeddings == "sinusoidal": + self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings) + else: + self.pos_embed = None + + # Define 3 blocks. Each block has its own normalization layer. + # 1. Self-Attn + if self.use_ada_layer_norm: + self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) + elif self.use_ada_layer_norm_zero: + self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) + else: + self.norm1 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + + if pkg_resources.parse_version(installed_version) >= pkg_resources.parse_version("0.28.2"): + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim if only_cross_attention else None, + upcast_attention=upcast_attention, + qk_norm="layer_norm" if qk_norm else None, + processor=HunyuanAttnProcessor2_0() if qk_norm else AttnProcessor2_0(), + ) + else: + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim if only_cross_attention else None, + upcast_attention=upcast_attention, + ) + + # 2. Cross-Attn + if cross_attention_dim is not None or double_self_attention: + # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. + # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during + # the second cross attention block. + self.norm2 = ( + AdaLayerNorm(dim, num_embeds_ada_norm) + if self.use_ada_layer_norm + else FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + ) + if pkg_resources.parse_version(installed_version) >= pkg_resources.parse_version("0.28.2"): + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim if not double_self_attention else None, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + qk_norm="layer_norm" if qk_norm else None, + processor=HunyuanAttnProcessor2_0() if qk_norm else AttnProcessor2_0(), + ) # is self-attn if encoder_hidden_states is none + else: + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim if not double_self_attention else None, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + ) # is self-attn if encoder_hidden_states is none + else: + self.norm2 = None + self.attn2 = None + + # 3. Feed-forward + if not self.use_ada_layer_norm_single: + self.norm3 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + + self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout) + + if after_norm: + self.norm4 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + else: + self.norm4 = None + + # 4. Fuser + if attention_type == "gated" or attention_type == "gated-text-image": + self.fuser = GatedSelfAttentionDense(dim, cross_attention_dim, num_attention_heads, attention_head_dim) + + # 5. Scale-shift for PixArt-Alpha. + if self.use_ada_layer_norm_single: + self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) + + # let chunk size default to None + self._chunk_size = None + self._chunk_dim = 0 + + def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int): + # Sets chunk feed-forward + self._chunk_size = chunk_size + self._chunk_dim = dim + + def forward( + self, + hidden_states: torch.FloatTensor, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + timestep: Optional[torch.LongTensor] = None, + cross_attention_kwargs: Dict[str, Any] = None, + class_labels: Optional[torch.LongTensor] = None, + ) -> torch.FloatTensor: + # Notice that normalization is always applied before the real computation in the following blocks. + # 0. Self-Attention + batch_size = hidden_states.shape[0] + + if self.use_ada_layer_norm: + norm_hidden_states = self.norm1(hidden_states, timestep) + elif self.use_ada_layer_norm_zero: + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( + hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + elif self.use_layer_norm: + norm_hidden_states = self.norm1(hidden_states) + elif self.use_ada_layer_norm_single: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) + ).chunk(6, dim=1) + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + norm_hidden_states = norm_hidden_states.squeeze(1) + else: + raise ValueError("Incorrect norm used") + + if self.pos_embed is not None: + norm_hidden_states = self.pos_embed(norm_hidden_states) + + # 1. Retrieve lora scale. + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + # 2. Prepare GLIGEN inputs + cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} + gligen_kwargs = cross_attention_kwargs.pop("gligen", None) + + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + + if self.use_ada_layer_norm_zero: + attn_output = gate_msa.unsqueeze(1) * attn_output + elif self.use_ada_layer_norm_single: + attn_output = gate_msa * attn_output + + hidden_states = attn_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + # 2.5 GLIGEN Control + if gligen_kwargs is not None: + hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"]) + + # 3. Cross-Attention + if self.attn2 is not None: + if self.use_ada_layer_norm: + norm_hidden_states = self.norm2(hidden_states, timestep) + elif self.use_ada_layer_norm_zero or self.use_layer_norm: + norm_hidden_states = self.norm2(hidden_states) + elif self.use_ada_layer_norm_single: + # For PixArt norm2 isn't applied here: + # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103 + norm_hidden_states = hidden_states + else: + raise ValueError("Incorrect norm") + + if self.pos_embed is not None and self.use_ada_layer_norm_single is None: + norm_hidden_states = self.pos_embed(norm_hidden_states) + + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + **cross_attention_kwargs, + ) + hidden_states = attn_output + hidden_states + + # 4. Feed-forward + if not self.use_ada_layer_norm_single: + norm_hidden_states = self.norm3(hidden_states) + + if self.use_ada_layer_norm_zero: + norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + + if self.use_ada_layer_norm_single: + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + + if self._chunk_size is not None: + # "feed_forward_chunk_size" can be used to save memory + if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0: + raise ValueError( + f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." + ) + + num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size + ff_output = torch.cat( + [ + self.ff(hid_slice, scale=lora_scale) + for hid_slice in norm_hidden_states.chunk(num_chunks, dim=self._chunk_dim) + ], + dim=self._chunk_dim, + ) + else: + ff_output = self.ff(norm_hidden_states, scale=lora_scale) + + if self.norm4 is not None: + ff_output = self.norm4(ff_output) + + if self.use_ada_layer_norm_zero: + ff_output = gate_mlp.unsqueeze(1) * ff_output + elif self.use_ada_layer_norm_single: + ff_output = gate_mlp * ff_output + + hidden_states = ff_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + return hidden_states + + +@maybe_allow_in_graph +class KVCompressionTransformerBlock(nn.Module): + r""" + A Temporal Transformer block. + + Parameters: + dim (`int`): The number of channels in the input and output. + num_attention_heads (`int`): The number of heads to use for multi-head attention. + attention_head_dim (`int`): The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + num_embeds_ada_norm (: + obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. + attention_bias (: + obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. + only_cross_attention (`bool`, *optional*): + Whether to use only cross-attention layers. In this case two cross attention layers are used. + double_self_attention (`bool`, *optional*): + Whether to use two self-attention layers. In this case no cross attention layers are used. + upcast_attention (`bool`, *optional*): + Whether to upcast the attention computation to float32. This is useful for mixed precision training. + norm_elementwise_affine (`bool`, *optional*, defaults to `True`): + Whether to use learnable elementwise affine parameters for normalization. + norm_type (`str`, *optional*, defaults to `"layer_norm"`): + The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`. + final_dropout (`bool` *optional*, defaults to False): + Whether to apply a final dropout after the last feed-forward layer. + attention_type (`str`, *optional*, defaults to `"default"`): + The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`. + positional_embeddings (`str`, *optional*, defaults to `None`): + The type of positional embeddings to apply to. + num_positional_embeddings (`int`, *optional*, defaults to `None`): + The maximum number of positional embeddings to apply. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout=0.0, + cross_attention_dim: Optional[int] = None, + activation_fn: str = "geglu", + num_embeds_ada_norm: Optional[int] = None, + attention_bias: bool = False, + only_cross_attention: bool = False, + double_self_attention: bool = False, + upcast_attention: bool = False, + norm_elementwise_affine: bool = True, + norm_type: str = "layer_norm", # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single' + norm_eps: float = 1e-5, + final_dropout: bool = False, + attention_type: str = "default", + positional_embeddings: Optional[str] = None, + num_positional_embeddings: Optional[int] = None, + kvcompression: Optional[bool] = False, + qk_norm = False, + after_norm = False, + ): + super().__init__() + self.only_cross_attention = only_cross_attention + + self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero" + self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" + self.use_ada_layer_norm_single = norm_type == "ada_norm_single" + self.use_layer_norm = norm_type == "layer_norm" + + if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: + raise ValueError( + f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" + f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." + ) + + if positional_embeddings and (num_positional_embeddings is None): + raise ValueError( + "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined." + ) + + if positional_embeddings == "sinusoidal": + self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings) + else: + self.pos_embed = None + + # Define 3 blocks. Each block has its own normalization layer. + # 1. Self-Attn + if self.use_ada_layer_norm: + self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) + elif self.use_ada_layer_norm_zero: + self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) + else: + self.norm1 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + + self.kvcompression = kvcompression + if kvcompression: + self.attn1 = KVCompressionCrossAttention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim if only_cross_attention else None, + upcast_attention=upcast_attention, + ) + else: + if pkg_resources.parse_version(installed_version) >= pkg_resources.parse_version("0.28.2"): + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim if only_cross_attention else None, + upcast_attention=upcast_attention, + qk_norm="layer_norm" if qk_norm else None, + processor=HunyuanAttnProcessor2_0() if qk_norm else AttnProcessor2_0(), + ) + else: + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim if only_cross_attention else None, + upcast_attention=upcast_attention, + ) + + # 2. Cross-Attn + if cross_attention_dim is not None or double_self_attention: + # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. + # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during + # the second cross attention block. + self.norm2 = ( + AdaLayerNorm(dim, num_embeds_ada_norm) + if self.use_ada_layer_norm + else FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + ) + if pkg_resources.parse_version(installed_version) >= pkg_resources.parse_version("0.28.2"): + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim if not double_self_attention else None, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + qk_norm="layer_norm" if qk_norm else None, + processor=HunyuanAttnProcessor2_0() if qk_norm else AttnProcessor2_0(), + ) # is self-attn if encoder_hidden_states is none + else: + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim if not double_self_attention else None, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + ) # is self-attn if encoder_hidden_states is none + else: + self.norm2 = None + self.attn2 = None + + # 3. Feed-forward + if not self.use_ada_layer_norm_single: + self.norm3 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + + self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout) + + if after_norm: + self.norm4 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + else: + self.norm4 = None + + # 4. Fuser + if attention_type == "gated" or attention_type == "gated-text-image": + self.fuser = GatedSelfAttentionDense(dim, cross_attention_dim, num_attention_heads, attention_head_dim) + + # 5. Scale-shift for PixArt-Alpha. + if self.use_ada_layer_norm_single: + self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) + + # let chunk size default to None + self._chunk_size = None + self._chunk_dim = 0 + + def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int): + # Sets chunk feed-forward + self._chunk_size = chunk_size + self._chunk_dim = dim + + def forward( + self, + hidden_states: torch.FloatTensor, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + timestep: Optional[torch.LongTensor] = None, + cross_attention_kwargs: Dict[str, Any] = None, + class_labels: Optional[torch.LongTensor] = None, + num_frames: int = 16, + height: int = 32, + width: int = 32, + use_reentrant: bool = False, + ) -> torch.FloatTensor: + # Notice that normalization is always applied before the real computation in the following blocks. + # 0. Self-Attention + batch_size = hidden_states.shape[0] + + if self.use_ada_layer_norm: + norm_hidden_states = self.norm1(hidden_states, timestep) + elif self.use_ada_layer_norm_zero: + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( + hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + elif self.use_layer_norm: + norm_hidden_states = self.norm1(hidden_states) + elif self.use_ada_layer_norm_single: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) + ).chunk(6, dim=1) + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + norm_hidden_states = norm_hidden_states.squeeze(1) + else: + raise ValueError("Incorrect norm used") + + if self.pos_embed is not None: + norm_hidden_states = self.pos_embed(norm_hidden_states) + + # 1. Retrieve lora scale. + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + # 2. Prepare GLIGEN inputs + cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} + gligen_kwargs = cross_attention_kwargs.pop("gligen", None) + + if self.kvcompression: + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, + attention_mask=attention_mask, + num_frames=num_frames, + height=height, + width=width, + **cross_attention_kwargs, + ) + else: + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + + if self.use_ada_layer_norm_zero: + attn_output = gate_msa.unsqueeze(1) * attn_output + elif self.use_ada_layer_norm_single: + attn_output = gate_msa * attn_output + + hidden_states = attn_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + # 2.5 GLIGEN Control + if gligen_kwargs is not None: + hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"]) + + # 3. Cross-Attention + if self.attn2 is not None: + if self.use_ada_layer_norm: + norm_hidden_states = self.norm2(hidden_states, timestep) + elif self.use_ada_layer_norm_zero or self.use_layer_norm: + norm_hidden_states = self.norm2(hidden_states) + elif self.use_ada_layer_norm_single: + # For PixArt norm2 isn't applied here: + # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103 + norm_hidden_states = hidden_states + else: + raise ValueError("Incorrect norm") + + if self.pos_embed is not None and self.use_ada_layer_norm_single is None: + norm_hidden_states = self.pos_embed(norm_hidden_states) + + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + **cross_attention_kwargs, + ) + hidden_states = attn_output + hidden_states + + # 4. Feed-forward + if not self.use_ada_layer_norm_single: + norm_hidden_states = self.norm3(hidden_states) + + if self.use_ada_layer_norm_zero: + norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + + if self.use_ada_layer_norm_single: + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + + if self._chunk_size is not None: + # "feed_forward_chunk_size" can be used to save memory + if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0: + raise ValueError( + f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`." + ) + + num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size + ff_output = torch.cat( + [ + self.ff(hid_slice, scale=lora_scale) + for hid_slice in norm_hidden_states.chunk(num_chunks, dim=self._chunk_dim) + ], + dim=self._chunk_dim, + ) + else: + ff_output = self.ff(norm_hidden_states, scale=lora_scale) + + if self.norm4 is not None: + ff_output = self.norm4(ff_output) + + if self.use_ada_layer_norm_zero: + ff_output = gate_mlp.unsqueeze(1) * ff_output + elif self.use_ada_layer_norm_single: + ff_output = gate_mlp * ff_output + + hidden_states = ff_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + return hidden_states + +class GEGLU(nn.Module): + def __init__(self, dim_in, dim_out, norm_elementwise_affine): + super().__init__() + self.norm = FP32LayerNorm(dim_in, dim_in, norm_elementwise_affine) + self.proj = nn.Linear(dim_in, dim_out * 2) + + def forward(self, x): + x, gate = self.proj(self.norm(x)).chunk(2, dim=-1) + return x * F.gelu(gate) + +@maybe_allow_in_graph +class HunyuanDiTBlock(nn.Module): + r""" + Transformer block used in Hunyuan-DiT model (https://github.com/Tencent/HunyuanDiT). Allow skip connection and + QKNorm + + Parameters: + dim (`int`): + The number of channels in the input and output. + num_attention_heads (`int`): + The number of headsto use for multi-head attention. + cross_attention_dim (`int`,*optional*): + The size of the encoder_hidden_states vector for cross attention. + dropout(`float`, *optional*, defaults to 0.0): + The dropout probability to use. + activation_fn (`str`,*optional*, defaults to `"geglu"`): + Activation function to be used in feed-forward. . + norm_elementwise_affine (`bool`, *optional*, defaults to `True`): + Whether to use learnable elementwise affine parameters for normalization. + norm_eps (`float`, *optional*, defaults to 1e-6): + A small constant added to the denominator in normalization layers to prevent division by zero. + final_dropout (`bool` *optional*, defaults to False): + Whether to apply a final dropout after the last feed-forward layer. + ff_inner_dim (`int`, *optional*): + The size of the hidden layer in the feed-forward block. Defaults to `None`. + ff_bias (`bool`, *optional*, defaults to `True`): + Whether to use bias in the feed-forward block. + skip (`bool`, *optional*, defaults to `False`): + Whether to use skip connection. Defaults to `False` for down-blocks and mid-blocks. + qk_norm (`bool`, *optional*, defaults to `True`): + Whether to use normalization in QK calculation. Defaults to `True`. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + cross_attention_dim: int = 1024, + dropout=0.0, + activation_fn: str = "geglu", + norm_elementwise_affine: bool = True, + norm_eps: float = 1e-6, + final_dropout: bool = False, + ff_inner_dim: Optional[int] = None, + ff_bias: bool = True, + skip: bool = False, + qk_norm: bool = True, + time_position_encoding: bool = False, + after_norm: bool = False, + is_local_attention: bool = False, + local_attention_frames: int = 2, + enable_inpaint: bool = False, + ): + super().__init__() + + # Define 3 blocks. Each block has its own normalization layer. + # NOTE: when new version comes, check norm2 and norm 3 + # 1. Self-Attn + self.norm1 = AdaLayerNormShift(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + self.t_embed = PositionalEncoding(dim, dropout=0., max_len=512) if time_position_encoding else nn.Identity() + + self.attn1 = Attention( + query_dim=dim, + cross_attention_dim=None, + dim_head=dim // num_attention_heads, + heads=num_attention_heads, + qk_norm="layer_norm" if qk_norm else None, + eps=1e-6, + bias=True, + processor=HunyuanAttnProcessor2_0(), + ) + + # 2. Cross-Attn + self.norm2 = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) + + self.is_local_attention = is_local_attention + self.local_attention_frames = local_attention_frames + if self.is_local_attention: + from mamba_ssm import Mamba2 + self.mamba_norm_in = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) + self.in_linear = nn.Linear(dim, 1536) + self.mamba_norm_1 = FP32LayerNorm(1536, norm_eps, norm_elementwise_affine) + self.mamba_norm_2 = FP32LayerNorm(1536, norm_eps, norm_elementwise_affine) + + self.mamba_block_1 = Mamba2( + d_model=1536, + d_state=64, + d_conv=4, + expand=2, + ) + self.mamba_block_2 = Mamba2( + d_model=1536, + d_state=64, + d_conv=4, + expand=2, + ) + self.mamba_norm_after_mamba_block = FP32LayerNorm(1536, norm_eps, norm_elementwise_affine) + + self.out_linear = nn.Linear(1536, dim) + self.out_linear = zero_module(self.out_linear) + self.mamba_norm_out = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) + + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim, + dim_head=dim // num_attention_heads, + heads=num_attention_heads, + qk_norm="layer_norm" if qk_norm else None, + eps=1e-6, + bias=True, + processor=HunyuanAttnProcessor2_0(), + ) + + if enable_inpaint: + self.norm_clip = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) + self.attn_clip = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim, + dim_head=dim // num_attention_heads, + heads=num_attention_heads, + qk_norm="layer_norm" if qk_norm else None, + eps=1e-6, + bias=True, + processor=HunyuanAttnProcessor2_0(), + ) + self.gate_clip = GEGLU(dim, dim, norm_elementwise_affine) + self.norm_clip_out = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) + else: + self.attn_clip = None + self.norm_clip = None + self.gate_clip = None + self.norm_clip_out = None + + # 3. Feed-forward + self.norm3 = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) + + self.ff = FeedForward( + dim, + dropout=dropout, ### 0.0 + activation_fn=activation_fn, ### approx GeLU + final_dropout=final_dropout, ### 0.0 + inner_dim=ff_inner_dim, ### int(dim * mlp_ratio) + bias=ff_bias, + ) + + # 4. Skip Connection + if skip: + self.skip_norm = FP32LayerNorm(2 * dim, norm_eps, elementwise_affine=True) + self.skip_linear = nn.Linear(2 * dim, dim) + else: + self.skip_linear = None + + if after_norm: + print("add after norm") + self.norm4 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + else: + self.norm4 = None + + # let chunk size default to None + self._chunk_size = None + self._chunk_dim = 0 + self.is_local_attention = is_local_attention + self.local_attention_frames = local_attention_frames + + def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): + # Sets chunk feed-forward + self._chunk_size = chunk_size + self._chunk_dim = dim + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + temb: Optional[torch.Tensor] = None, + image_rotary_emb=None, + skip=None, + num_frames: int = 1, + height: int = 32, + width: int = 32, + clip_encoder_hidden_states: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + # Notice that normalization is always applied before the real computation in the following blocks. + # 0. Long Skip Connection + if self.skip_linear is not None: + cat = torch.cat([hidden_states, skip], dim=-1) + cat = self.skip_norm(cat) + hidden_states = self.skip_linear(cat) + + if num_frames != 1: + image_rotary_emb = (torch.cat([image_rotary_emb[0] for i in range(num_frames)], dim=0), torch.cat([image_rotary_emb[1] for i in range(num_frames)], dim=0)) + + # add time embedding + hidden_states = rearrange(hidden_states, "b (f d) c -> (b d) f c", f=num_frames) + if self.t_embed is not None: + hidden_states = self.t_embed(hidden_states) + hidden_states = rearrange(hidden_states, "(b d) f c -> b (f d) c", d=height * width) + + # 1. Self-Attention + norm_hidden_states = self.norm1(hidden_states, temb) ### checked: self.norm1 is correct + if num_frames > 2 and self.is_local_attention: + attn1_image_rotary_emb = (image_rotary_emb[0][:int(height * width * 2)], image_rotary_emb[1][:int(height * width * 2)]) + + norm_hidden_states_1 = rearrange(norm_hidden_states, "b (f d) c -> b f d c", d=height * width) + norm_hidden_states_1 = rearrange(norm_hidden_states_1, "b (f p) d c -> (b f) (p d) c", p = 2) + + attn_output = self.attn1( + norm_hidden_states_1, + image_rotary_emb=attn1_image_rotary_emb, + ) + attn_output = rearrange(attn_output, "(b f) (p d) c -> b (f p) d c", p = 2, f = num_frames // 2) + + norm_hidden_states_2 = rearrange(norm_hidden_states, "b (f d) c -> b f d c", d = height * width)[:, 1:-1] + local_attention_frames_num = norm_hidden_states_2.size()[1] // 2 + norm_hidden_states_2 = rearrange(norm_hidden_states_2, "b (f p) d c -> (b f) (p d) c", p = 2) + attn_output_2 = self.attn1( + norm_hidden_states_2, + image_rotary_emb=attn1_image_rotary_emb, + ) + attn_output_2 = rearrange(attn_output_2, "(b f) (p d) c -> b (f p) d c", p = 2, f = local_attention_frames_num) + attn_output[:, 1:-1] = (attn_output[:, 1:-1] + attn_output_2) / 2 + + attn_output = rearrange(attn_output, "b f d c -> b (f d) c") + else: + attn_output = self.attn1( + norm_hidden_states, + image_rotary_emb=image_rotary_emb, + ) + hidden_states = hidden_states + attn_output + + if num_frames > 2 and self.is_local_attention: + hidden_states_in = self.in_linear(self.mamba_norm_in(hidden_states)) + hidden_states = hidden_states + self.mamba_norm_out( + self.out_linear( + self.mamba_norm_after_mamba_block( + self.mamba_block_1( + self.mamba_norm_1(hidden_states_in) + ) + + self.mamba_block_2( + self.mamba_norm_2(hidden_states_in.flip(1)) + ).flip(1) + ) + ) + ) + + # 2. Cross-Attention + hidden_states = hidden_states + self.attn2( + self.norm2(hidden_states), + encoder_hidden_states=encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + ) + + if self.attn_clip is not None: + hidden_states = hidden_states + self.norm_clip_out( + self.gate_clip( + self.attn_clip( + self.norm_clip(hidden_states), + encoder_hidden_states=clip_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + ) + ) + ) + + # FFN Layer ### TODO: switch norm2 and norm3 in the state dict + mlp_inputs = self.norm3(hidden_states) + if self.norm4 is not None: + hidden_states = hidden_states + self.norm4(self.ff(mlp_inputs)) + else: + hidden_states = hidden_states + self.ff(mlp_inputs) + + return hidden_states + + +@maybe_allow_in_graph +class HunyuanTemporalTransformerBlock(nn.Module): + r""" + Transformer block used in Hunyuan-DiT model (https://github.com/Tencent/HunyuanDiT). Allow skip connection and + QKNorm + + Parameters: + dim (`int`): + The number of channels in the input and output. + num_attention_heads (`int`): + The number of headsto use for multi-head attention. + cross_attention_dim (`int`,*optional*): + The size of the encoder_hidden_states vector for cross attention. + dropout(`float`, *optional*, defaults to 0.0): + The dropout probability to use. + activation_fn (`str`,*optional*, defaults to `"geglu"`): + Activation function to be used in feed-forward. . + norm_elementwise_affine (`bool`, *optional*, defaults to `True`): + Whether to use learnable elementwise affine parameters for normalization. + norm_eps (`float`, *optional*, defaults to 1e-6): + A small constant added to the denominator in normalization layers to prevent division by zero. + final_dropout (`bool` *optional*, defaults to False): + Whether to apply a final dropout after the last feed-forward layer. + ff_inner_dim (`int`, *optional*): + The size of the hidden layer in the feed-forward block. Defaults to `None`. + ff_bias (`bool`, *optional*, defaults to `True`): + Whether to use bias in the feed-forward block. + skip (`bool`, *optional*, defaults to `False`): + Whether to use skip connection. Defaults to `False` for down-blocks and mid-blocks. + qk_norm (`bool`, *optional*, defaults to `True`): + Whether to use normalization in QK calculation. Defaults to `True`. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + cross_attention_dim: int = 1024, + dropout=0.0, + activation_fn: str = "geglu", + norm_elementwise_affine: bool = True, + norm_eps: float = 1e-6, + final_dropout: bool = False, + ff_inner_dim: Optional[int] = None, + ff_bias: bool = True, + skip: bool = False, + qk_norm: bool = True, + after_norm: bool = False, + # motion module kwargs + motion_module_type = "VanillaGrid", + motion_module_kwargs = None, + use_reentrant: bool = False, + ): + super().__init__() + + # Define 3 blocks. Each block has its own normalization layer. + # NOTE: when new version comes, check norm2 and norm 3 + # 1. Self-Attn + self.norm1 = AdaLayerNormShift(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + + self.attn1 = Attention( + query_dim=dim, + cross_attention_dim=None, + dim_head=dim // num_attention_heads, + heads=num_attention_heads, + qk_norm="layer_norm" if qk_norm else None, + eps=1e-6, + bias=True, + processor=HunyuanAttnProcessor2_0(), + ) + + self.attn_temporal = get_motion_module( + in_channels = dim, + motion_module_type = motion_module_type, + motion_module_kwargs = motion_module_kwargs, + ) + + # 2. Cross-Attn + self.norm2 = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) + + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim, + dim_head=dim // num_attention_heads, + heads=num_attention_heads, + qk_norm="layer_norm" if qk_norm else None, + eps=1e-6, + bias=True, + processor=HunyuanAttnProcessor2_0(), + ) + # 3. Feed-forward + self.norm3 = FP32LayerNorm(dim, norm_eps, norm_elementwise_affine) + + self.ff = FeedForward( + dim, + dropout=dropout, ### 0.0 + activation_fn=activation_fn, ### approx GeLU + final_dropout=final_dropout, ### 0.0 + inner_dim=ff_inner_dim, ### int(dim * mlp_ratio) + bias=ff_bias, + ) + + # 4. Skip Connection + if skip: + self.skip_norm = FP32LayerNorm(2 * dim, norm_eps, elementwise_affine=True) + self.skip_linear = nn.Linear(2 * dim, dim) + else: + self.skip_linear = None + + if after_norm: + print("add after norm") + self.norm4 = FP32LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + else: + self.norm4 = None + + # let chunk size default to None + self._chunk_size = None + self._chunk_dim = 0 + + def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): + # Sets chunk feed-forward + self._chunk_size = chunk_size + self._chunk_dim = dim + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + temb: Optional[torch.Tensor] = None, + image_rotary_emb=None, + skip=None, + num_frames: int = 16, + height: int = 32, + width: int = 32, + use_reentrant: bool = False, + ) -> torch.Tensor: + # Notice that normalization is always applied before the real computation in the following blocks. + # 0. Long Skip Connection + if self.skip_linear is not None: + cat = torch.cat([hidden_states, skip], dim=-1) + cat = self.skip_norm(cat) + hidden_states = self.skip_linear(cat) + + # 1. Self-Attention + norm_hidden_states = self.norm1(hidden_states, temb) ### checked: self.norm1 is correct + norm_hidden_states = rearrange(norm_hidden_states, "b (f d) c -> (b f) d c", f=num_frames) + attn_output = self.attn1( + norm_hidden_states, + image_rotary_emb=image_rotary_emb, + ) + attn_output = rearrange(attn_output, "(b f) d c -> b (f d) c", f=num_frames) + hidden_states = hidden_states + attn_output + + # 1.5. Temp-Attention + if self.attn_temporal is not None: + attn_output = rearrange(hidden_states, "b (f h w) c -> b c f h w", f=num_frames, h=height, w=width) + attn_output = self.attn_temporal(attn_output) + hidden_states = rearrange(attn_output, "b c f h w -> b (f h w) c") + + if num_frames != 1: + image_rotary_emb = (torch.cat([image_rotary_emb[0] for i in range(num_frames)], dim=0), torch.cat([image_rotary_emb[1] for i in range(num_frames)], dim=0)) + + # 2. Cross-Attention + hidden_states = hidden_states + self.attn2( + self.norm2(hidden_states), + encoder_hidden_states=encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + ) + + # FFN Layer ### TODO: switch norm2 and norm3 in the state dict + mlp_inputs = self.norm3(hidden_states) + if self.norm4 is not None: + hidden_states = hidden_states + self.norm4(self.ff(mlp_inputs)) + else: + hidden_states = hidden_states + self.ff(mlp_inputs) + + return hidden_states diff --git a/ruyi/models/autoencoder_magvit.py b/ruyi/models/autoencoder_magvit.py new file mode 100644 index 0000000..a96cc56 --- /dev/null +++ b/ruyi/models/autoencoder_magvit.py @@ -0,0 +1,569 @@ +# Copyright 2024 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. +from typing import Dict, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config + +try: + from diffusers.loaders import FromOriginalVAEMixin +except: + from diffusers.loaders import FromOriginalModelMixin as FromOriginalVAEMixin + +from diffusers.models.attention_processor import ( + ADDED_KV_ATTENTION_PROCESSORS, CROSS_ATTENTION_PROCESSORS, Attention, + AttentionProcessor, AttnAddedKVProcessor, AttnProcessor) +from diffusers.models.autoencoders.vae import (DecoderOutput, + DiagonalGaussianDistribution) +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.utils.accelerate_utils import apply_forward_hook +from torch import nn + +from ..vae.ldm.models.omnigen_enc_dec import Decoder as omnigen_Mag_Decoder +from ..vae.ldm.models.omnigen_enc_dec import Encoder as omnigen_Mag_Encoder + + +def str_eval(item): + if type(item) == str: + return eval(item) + else: + return item + +class AutoencoderKLMagvit(ModelMixin, ConfigMixin, FromOriginalVAEMixin): + r""" + A VAE model with KL loss for encoding images into latents and decoding latent representations into images. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + + Parameters: + in_channels (int, *optional*, defaults to 3): Number of channels in the input image. + out_channels (int, *optional*, defaults to 3): Number of channels in the output. + down_block_types (`Tuple[str]`, *optional*, defaults to `("DownEncoderBlock2D",)`): + Tuple of downsample block types. + up_block_types (`Tuple[str]`, *optional*, defaults to `("UpDecoderBlock2D",)`): + Tuple of upsample block types. + block_out_channels (`Tuple[int]`, *optional*, defaults to `(64,)`): + Tuple of block output channels. + act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. + latent_channels (`int`, *optional*, defaults to 4): Number of channels in the latent space. + sample_size (`int`, *optional*, defaults to `32`): Sample input size. + scaling_factor (`float`, *optional*, defaults to 0.18215): + The component-wise standard deviation of the trained latent space computed using the first batch of the + training set. This is used to scale the latent space to have unit variance when training the diffusion + model. The latents are scaled with the formula `z = z * scaling_factor` before being passed to the + diffusion model. When decoding, the latents are scaled back to the original scale with the formula: `z = 1 + / scaling_factor * z`. For more details, refer to sections 4.3.2 and D.1 of the [High-Resolution Image + Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752) paper. + force_upcast (`bool`, *optional*, default to `True`): + If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE + can be fine-tuned / trained to a lower range without loosing too much precision in which case + `force_upcast` can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + ch = 128, + ch_mult = [ 1,2,4,4 ], + use_gc_blocks = None, + down_block_types: tuple = None, + up_block_types: tuple = None, + mid_block_type: str = "MidBlock3D", + mid_block_use_attention: bool = True, + mid_block_attention_type: str = "3d", + mid_block_num_attention_heads: int = 1, + layers_per_block: int = 2, + act_fn: str = "silu", + num_attention_heads: int = 1, + latent_channels: int = 4, + norm_num_groups: int = 32, + scaling_factor: float = 0.1825, + slice_mag_vae=True, + slice_compression_vae=False, + cache_compression_vae=False, + use_tiling=False, # True + use_tiling_encoder=False, + use_tiling_decoder=False, + mini_batch_encoder=9, + mini_batch_decoder=3, + upcast_vae=False, + spatial_group_norm=False, + tile_sample_min_size=384, + tile_overlap_factor=0.25, + ): + super().__init__() + down_block_types = str_eval(down_block_types) + up_block_types = str_eval(up_block_types) + self.encoder = omnigen_Mag_Encoder( + in_channels=in_channels, + out_channels=latent_channels, + down_block_types=down_block_types, + ch = ch, + ch_mult = ch_mult, + use_gc_blocks=use_gc_blocks, + mid_block_type=mid_block_type, + mid_block_use_attention=mid_block_use_attention, + mid_block_attention_type=mid_block_attention_type, + mid_block_num_attention_heads=mid_block_num_attention_heads, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + num_attention_heads=num_attention_heads, + double_z=True, + slice_mag_vae=slice_mag_vae, + slice_compression_vae=slice_compression_vae, + cache_compression_vae=cache_compression_vae, + mini_batch_encoder=mini_batch_encoder, + spatial_group_norm=spatial_group_norm, + ) + + self.decoder = omnigen_Mag_Decoder( + in_channels=latent_channels, + out_channels=out_channels, + up_block_types=up_block_types, + ch = ch, + ch_mult = ch_mult, + use_gc_blocks=use_gc_blocks, + mid_block_type=mid_block_type, + mid_block_use_attention=mid_block_use_attention, + mid_block_attention_type=mid_block_attention_type, + mid_block_num_attention_heads=mid_block_num_attention_heads, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + num_attention_heads=num_attention_heads, + slice_mag_vae=slice_mag_vae, + slice_compression_vae=slice_compression_vae, + cache_compression_vae=cache_compression_vae, + mini_batch_decoder=mini_batch_decoder, + spatial_group_norm=spatial_group_norm, + ) + + self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1) + self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1) + + self.slice_mag_vae = slice_mag_vae + self.slice_compression_vae = slice_compression_vae + self.cache_compression_vae = cache_compression_vae + self.mini_batch_encoder = mini_batch_encoder + self.mini_batch_decoder = mini_batch_decoder + self.use_slicing = False + self.use_tiling = use_tiling + self.use_tiling_encoder = use_tiling_encoder + self.use_tiling_decoder = use_tiling_decoder + self.upcast_vae = upcast_vae + self.tile_sample_min_size = tile_sample_min_size + self.tile_overlap_factor = tile_overlap_factor + self.tile_latent_min_size = int(self.tile_sample_min_size / (2 ** (len(ch_mult) - 1))) + self.scaling_factor = scaling_factor + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, (omnigen_Mag_Encoder, omnigen_Mag_Decoder)): + module.gradient_checkpointing = value + + @property + # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]): + if hasattr(module, "get_processor"): + processors["%s.processor"%(name)] = module.get_processor(return_deprecated_lora=True) + + for sub_name, child in module.named_children(): + fn_recursive_add_processors("%s.%s"%(name, sub_name), child, processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor + def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes." + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor) + else: + module.set_processor(processor.pop("%s.processor"%(name))) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor("%s.%s"%(name, sub_name), child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor + def set_default_attn_processor(self): + """ + Disables custom attention processors and sets the default attention implementation. + """ + if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): + processor = AttnAddedKVProcessor() + elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()): + processor = AttnProcessor() + else: + raise ValueError( + f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" + ) + + self.set_attn_processor(processor) + + @apply_forward_hook + def encode( + self, x: torch.FloatTensor, return_dict: bool = True + ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]: + """ + Encode a batch of images into latents. + + Args: + x (`torch.FloatTensor`): Input batch of images. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + The latent representations of the encoded images. If `return_dict` is True, a + [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. + """ + if self.upcast_vae: # False + x = x.float() + self.encoder = self.encoder.float() + self.quant_conv = self.quant_conv.float() + if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size): # True, Almost False + x = self.tiled_encode(x, return_dict=return_dict) + return x + if self.use_tiling_encoder and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size): # False, Almost False + x = self.tiled_encode(x, return_dict=return_dict) + return x + + if self.use_slicing and x.shape[0] > 1: # False, False when B=1 + encoded_slices = [self.encoder(x_slice) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = self.encoder(x) + + moments = self.quant_conv(h) + posterior = DiagonalGaussianDistribution(moments) + + if not return_dict: # False + return (posterior,) + + return AutoencoderKLOutput(latent_dist=posterior) + + def _decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]: + if self.upcast_vae: + z = z.float() + self.decoder = self.decoder.float() + self.post_quant_conv = self.post_quant_conv.float() + if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size): + return self.tiled_decode(z, return_dict=return_dict) + if self.use_tiling_decoder and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size): + return self.tiled_decode(z, return_dict=return_dict) + + z = self.post_quant_conv(z) + dec = self.decoder(z) + + if not return_dict: + return (dec,) + + return DecoderOutput(sample=dec) + + @apply_forward_hook + def decode( + self, z: torch.FloatTensor, return_dict: bool = True, generator=None + ) -> Union[DecoderOutput, torch.FloatTensor]: + """ + Decode a batch of images. + + Args: + z (`torch.FloatTensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is + returned. + + """ + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = self._decode(z).sample + + if not return_dict: + return (decoded,) + + return DecoderOutput(sample=decoded) + + def blend_v( + self, a: torch.Tensor, b: torch.Tensor, blend_extent: int + ) -> torch.Tensor: + blend_extent = min(a.shape[3], b.shape[3], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * ( + 1 - y / blend_extent + ) + b[:, :, :, y, :] * (y / blend_extent) + return b + + def blend_h( + self, a: torch.Tensor, b: torch.Tensor, blend_extent: int + ) -> torch.Tensor: + blend_extent = min(a.shape[4], b.shape[4], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * ( + 1 - x / blend_extent + ) + b[:, :, :, :, x] * (x / blend_extent) + return b + + def tiled_encode(self, x: torch.FloatTensor, return_dict: bool = True) -> AutoencoderKLOutput: + overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) + row_limit = self.tile_latent_min_size - blend_extent + + # Split the image into 512x512 tiles and encode them separately. + rows = [] + for i in range(0, x.shape[3], overlap_size): + row = [] + for j in range(0, x.shape[4], overlap_size): + tile = x[ + :, + :, + :, + i : i + self.tile_sample_min_size, + j : j + self.tile_sample_min_size, + ] + tile = self.encoder(tile) + tile = self.quant_conv(tile) + row.append(tile) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=4)) + + moments = torch.cat(result_rows, dim=3) + posterior = DiagonalGaussianDistribution(moments) + + if not return_dict: + return (posterior,) + + return AutoencoderKLOutput(latent_dist=posterior) + + def tiled_decode(self, z: torch.FloatTensor, return_dict: bool = True) -> Union[DecoderOutput, torch.FloatTensor]: + overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) + blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) + row_limit = self.tile_sample_min_size - blend_extent + + # Split z into overlapping 64x64 tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + rows = [] + for i in range(0, z.shape[3], overlap_size): + row = [] + for j in range(0, z.shape[4], overlap_size): + tile = z[ + :, + :, + :, + i : i + self.tile_latent_min_size, + j : j + self.tile_latent_min_size, + ] + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + row.append(decoded) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=4)) + + dec = torch.cat(result_rows, dim=3) + + # Handle the lower right corner tile separately + lower_right_original = z[ + :, + :, + :, + -self.tile_latent_min_size:, + -self.tile_latent_min_size: + ] + quantized_lower_right = self.decoder(self.post_quant_conv(lower_right_original)) + + # Combine + H, W = quantized_lower_right.size(-2), quantized_lower_right.size(-1) + x_weights = torch.linspace(0, 1, W).unsqueeze(0).repeat(H, 1) + y_weights = torch.linspace(0, 1, H).unsqueeze(1).repeat(1, W) + weights = torch.min(x_weights, y_weights) + + if len(dec.size()) == 4: + weights = weights.unsqueeze(0).unsqueeze(0) + elif len(dec.size()) == 5: + weights = weights.unsqueeze(0).unsqueeze(0).unsqueeze(0) + + weights = weights.to(dec.device) + quantized_area = dec[:, :, :, -H:, -W:] + combined = weights * quantized_lower_right + (1 - weights) * quantized_area + + dec[:, :, :, -H:, -W:] = combined + + if not return_dict: + return (dec,) + + return DecoderOutput(sample=dec) + + def forward( + self, + sample: torch.FloatTensor, + sample_posterior: bool = False, + return_dict: bool = True, + generator: Optional[torch.Generator] = None, + ) -> Union[DecoderOutput, torch.FloatTensor]: + r""" + Args: + sample (`torch.FloatTensor`): Input sample. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the posterior. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + """ + x = sample + posterior = self.encode(x).latent_dist + if sample_posterior: + z = posterior.sample(generator=generator) + else: + z = posterior.mode() + dec = self.decode(z).sample + + if not return_dict: + return (dec,) + + return DecoderOutput(sample=dec) + + # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections + def fuse_qkv_projections(self): + """ + Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, + key, value) are fused. For cross-attention modules, key and value projection matrices are fused. + + + + This API is 🧪 experimental. + + + """ + self.original_attn_processors = None + + for _, attn_processor in self.attn_processors.items(): + if "Added" in str(attn_processor.__class__.__name__): + raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.") + + self.original_attn_processors = self.attn_processors + + for module in self.modules(): + if isinstance(module, Attention): + module.fuse_projections(fuse=True) + + # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections + def unfuse_qkv_projections(self): + """Disables the fused QKV projection if enabled. + + + + This API is 🧪 experimental. + + + + """ + if self.original_attn_processors is not None: + self.set_attn_processor(self.original_attn_processors) + + @classmethod + def from_pretrained(cls, pretrained_model_path, subfolder=None, **vae_additional_kwargs): + import json + import os + if subfolder is not None: + pretrained_model_path = os.path.join(pretrained_model_path, subfolder) + + config_file = os.path.join(pretrained_model_path, 'config.json') + if not os.path.isfile(config_file): + raise RuntimeError(f"{config_file} does not exist") + with open(config_file, "r") as f: + config = json.load(f) + + model = cls.from_config(config, **vae_additional_kwargs) + from diffusers.utils import WEIGHTS_NAME + model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME) + model_file_safetensors = model_file.replace(".bin", ".safetensors") + if os.path.exists(model_file_safetensors): + from safetensors.torch import load_file, safe_open + state_dict = load_file(model_file_safetensors) + else: + if not os.path.isfile(model_file): + raise RuntimeError(f"{model_file} does not exist") + state_dict = torch.load(model_file, map_location="cpu") + m, u = model.load_state_dict(state_dict, strict=False) + return model diff --git a/ruyi/models/embeddings.py b/ruyi/models/embeddings.py new file mode 100644 index 0000000..ca4e5be --- /dev/null +++ b/ruyi/models/embeddings.py @@ -0,0 +1,222 @@ +import math +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + +from diffusers.utils import deprecate +from diffusers.models.activations import FP32SiLU, get_activation +from diffusers.models.attention_processor import Attention + + +def get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + flip_sin_to_cos: bool = False, + downscale_freq_shift: float = 1, + scale: float = 1, + max_period: int = 10000, +): + """ + This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. + + :param timesteps: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the + embeddings. :return: an [N x dim] Tensor of positional embeddings. + """ + assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" + + half_dim = embedding_dim // 2 + exponent = -math.log(max_period) * torch.arange( + start=0, end=half_dim, dtype=torch.float32, device=timesteps.device + ) + exponent = exponent / (half_dim - downscale_freq_shift) + + emb = torch.exp(exponent) + emb = timesteps[:, None].float() * emb[None, :] + + # scale embeddings + emb = scale * emb + + # concat sine and cosine embeddings + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + + # flip sine and cosine embeddings + if flip_sin_to_cos: + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + + # zero pad + if embedding_dim % 2 == 1: + emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) + return emb + + +class Timesteps(nn.Module): + def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float): + super().__init__() + self.num_channels = num_channels + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + + def forward(self, timesteps): + t_emb = get_timestep_embedding( + timesteps, + self.num_channels, + flip_sin_to_cos=self.flip_sin_to_cos, + downscale_freq_shift=self.downscale_freq_shift, + ) + return t_emb + + +class TimestepEmbedding(nn.Module): + def __init__( + self, + in_channels: int, + time_embed_dim: int, + act_fn: str = "silu", + out_dim: int = None, + post_act_fn: Optional[str] = None, + cond_proj_dim=None, + sample_proj_bias=True, + ): + super().__init__() + + self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias) + + if cond_proj_dim is not None: + self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False) + else: + self.cond_proj = None + + self.act = get_activation(act_fn) + + if out_dim is not None: + time_embed_dim_out = out_dim + else: + time_embed_dim_out = time_embed_dim + self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias) + + if post_act_fn is None: + self.post_act = None + else: + self.post_act = get_activation(post_act_fn) + + def forward(self, sample, condition=None): + if condition is not None: + sample = sample + self.cond_proj(condition) + sample = self.linear_1(sample) + + if self.act is not None: + sample = self.act(sample) + + sample = self.linear_2(sample) + + if self.post_act is not None: + sample = self.post_act(sample) + return sample + + +class PixArtAlphaTextProjection(nn.Module): + """ + Projects caption embeddings. Also handles dropout for classifier-free guidance. + + Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py + """ + + def __init__(self, in_features, hidden_size, out_features=None, act_fn="gelu_tanh"): + super().__init__() + if out_features is None: + out_features = hidden_size + self.linear_1 = nn.Linear(in_features=in_features, out_features=hidden_size, bias=True) + if act_fn == "gelu_tanh": + self.act_1 = nn.GELU(approximate="tanh") + elif act_fn == "silu_fp32": + self.act_1 = FP32SiLU() + else: + raise ValueError(f"Unknown activation function: {act_fn}") + self.linear_2 = nn.Linear(in_features=hidden_size, out_features=out_features, bias=True) + + def forward(self, caption): + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class HunyuanDiTAttentionPool(nn.Module): + def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None): + super().__init__() + self.positional_embedding = nn.Parameter(torch.randn(spacial_dim + 1, embed_dim) / embed_dim**0.5) + self.k_proj = nn.Linear(embed_dim, embed_dim) + self.q_proj = nn.Linear(embed_dim, embed_dim) + self.v_proj = nn.Linear(embed_dim, embed_dim) + self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim) + self.num_heads = num_heads + + def forward(self, x): + x = torch.cat([x.mean(dim=1, keepdim=True), x], dim=1) + x = x + self.positional_embedding[None, :, :].to(x.dtype) + + query = self.q_proj(x[:, :1]) + key = self.k_proj(x) + value = self.v_proj(x) + batch_size, _, _ = query.size() + + query = query.reshape(batch_size, -1, self.num_heads, query.size(-1) // self.num_heads).transpose(1, 2) # (1, H, N, E/H) + key = key.reshape(batch_size, -1, self.num_heads, key.size(-1) // self.num_heads).transpose(1, 2) # (L+1, H, N, E/H) + value = value.reshape(batch_size, -1, self.num_heads, value.size(-1) // self.num_heads).transpose(1, 2) # (L+1, H, N, E/H) + + x = F.scaled_dot_product_attention(query=query, key=key, value=value, attn_mask=None, dropout_p=0.0, is_causal=False) + x = x.transpose(1, 2).reshape(batch_size, 1, -1) + x = x.to(query.dtype) + x = self.c_proj(x) + + return x.squeeze(1) + + +class HunyuanCombinedTimestepTextSizeStyleEmbedding(nn.Module): + def __init__(self, embedding_dim, pooled_projection_dim=1024, seq_len=256, cross_attention_dim=2048): + super().__init__() + + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + self.pooler = HunyuanDiTAttentionPool( + seq_len, cross_attention_dim, num_heads=8, output_dim=pooled_projection_dim + ) + # Here we use a default learned embedder layer for future extension. + self.style_embedder = nn.Embedding(1, embedding_dim) + extra_in_dim = 256 * 6 + embedding_dim + pooled_projection_dim + self.extra_embedder = PixArtAlphaTextProjection( + in_features=extra_in_dim, + hidden_size=embedding_dim * 4, + out_features=embedding_dim, + act_fn="silu_fp32", + ) + + def forward(self, timestep, encoder_hidden_states, image_meta_size, style, hidden_dtype=None): + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, 256) + + # extra condition1: text + pooled_projections = self.pooler(encoder_hidden_states) # (N, 1024) + + # extra condition2: image meta size embdding + image_meta_size = get_timestep_embedding(image_meta_size.view(-1), 256, True, 0) + image_meta_size = image_meta_size.to(dtype=hidden_dtype) + image_meta_size = image_meta_size.view(-1, 6 * 256) # (N, 1536) + + # extra condition3: style embedding + style_embedding = self.style_embedder(style) # (N, embedding_dim) + + # Concatenate all extra vectors + extra_cond = torch.cat([pooled_projections, image_meta_size, style_embedding], dim=1) + conditioning = timesteps_emb + self.extra_embedder(extra_cond) # [B, D] + + return conditioning \ No newline at end of file diff --git a/ruyi/models/motion_module.py b/ruyi/models/motion_module.py new file mode 100644 index 0000000..8a5d57f --- /dev/null +++ b/ruyi/models/motion_module.py @@ -0,0 +1,444 @@ +"""Modified from https://github.com/guoyww/AnimateDiff/blob/main/animatediff/models/motion_module.py +""" +import math + +import diffusers +import pkg_resources +import torch + +installed_version = diffusers.__version__ + +if pkg_resources.parse_version(installed_version) >= pkg_resources.parse_version("0.28.2"): + from diffusers.models.attention_processor import (Attention, + AttnProcessor2_0, + HunyuanAttnProcessor2_0) +else: + from diffusers.models.attention_processor import Attention, AttnProcessor2_0 + +from diffusers.models.attention import FeedForward +from diffusers.utils.import_utils import is_xformers_available +from einops import rearrange, repeat +from torch import nn + +from .norm import FP32LayerNorm + +if is_xformers_available(): + import xformers + import xformers.ops +else: + xformers = None + +def zero_module(module): + # Zero out the parameters of a module and return it. + for p in module.parameters(): + p.detach().zero_() + return module + +def get_motion_module( + in_channels, + motion_module_type: str, + motion_module_kwargs: dict, +): + if motion_module_type == "Vanilla": + return VanillaTemporalModule(in_channels=in_channels, **motion_module_kwargs,) + elif motion_module_type == "VanillaGrid": + return VanillaTemporalModule(in_channels=in_channels, grid=True, **motion_module_kwargs,) + else: + raise ValueError + +class VanillaTemporalModule(nn.Module): + def __init__( + self, + in_channels, + num_attention_heads = 8, + num_transformer_block = 2, + attention_block_types =( "Temporal_Self", "Temporal_Self" ), + cross_frame_attention_mode = None, + temporal_position_encoding = False, + temporal_position_encoding_max_len = 4096, + temporal_attention_dim_div = 1, + zero_initialize = True, + block_size = 1, + grid = False, + remove_time_embedding_in_photo = False, + + global_num_attention_heads = 16, + global_attention = False, + qk_norm = False, + ): + super().__init__() + + self.temporal_transformer = TemporalTransformer3DModel( + in_channels=in_channels, + num_attention_heads=num_attention_heads, + attention_head_dim=in_channels // num_attention_heads // temporal_attention_dim_div, + num_layers=num_transformer_block, + attention_block_types=attention_block_types, + cross_frame_attention_mode=cross_frame_attention_mode, + temporal_position_encoding=temporal_position_encoding, + temporal_position_encoding_max_len=temporal_position_encoding_max_len, + grid=grid, + block_size=block_size, + remove_time_embedding_in_photo=remove_time_embedding_in_photo, + qk_norm=qk_norm, + ) + self.global_transformer = GlobalTransformer3DModel( + in_channels=in_channels, + num_attention_heads=global_num_attention_heads, + attention_head_dim=in_channels // global_num_attention_heads // temporal_attention_dim_div, + qk_norm=qk_norm, + ) if global_attention else None + if zero_initialize: + self.temporal_transformer.proj_out = zero_module(self.temporal_transformer.proj_out) + if global_attention: + self.global_transformer.proj_out = zero_module(self.global_transformer.proj_out) + + def forward(self, input_tensor, encoder_hidden_states=None, attention_mask=None, anchor_frame_idx=None): + hidden_states = input_tensor + hidden_states = self.temporal_transformer(hidden_states, encoder_hidden_states, attention_mask) + if self.global_transformer is not None: + hidden_states = self.global_transformer(hidden_states) + + output = hidden_states + return output + +class GlobalTransformer3DModel(nn.Module): + def __init__( + self, + in_channels, + num_attention_heads, + attention_head_dim, + dropout = 0.0, + attention_bias = False, + upcast_attention = False, + qk_norm = False, + ): + super().__init__() + + inner_dim = num_attention_heads * attention_head_dim + + self.norm1 = FP32LayerNorm(inner_dim) + self.proj_in = nn.Linear(in_channels, inner_dim) + self.norm2 = FP32LayerNorm(inner_dim) + if pkg_resources.parse_version(installed_version) >= pkg_resources.parse_version("0.28.2"): + self.attention = Attention( + query_dim=inner_dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + qk_norm="layer_norm" if qk_norm else None, + processor=HunyuanAttnProcessor2_0() if qk_norm else AttnProcessor2_0(), + ) + else: + self.attention = Attention( + query_dim=inner_dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + ) + self.proj_out = nn.Linear(inner_dim, in_channels) + + def forward(self, hidden_states): + assert hidden_states.dim() == 5, f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}." + video_length, height, width = hidden_states.shape[2], hidden_states.shape[3], hidden_states.shape[4] + hidden_states = rearrange(hidden_states, "b c f h w -> b (f h w) c") + + residual = hidden_states + hidden_states = self.norm1(hidden_states) + hidden_states = self.proj_in(hidden_states) + + # Attention Blocks + hidden_states = self.norm2(hidden_states) + hidden_states = self.attention(hidden_states) + hidden_states = self.proj_out(hidden_states) + + output = hidden_states + residual + output = rearrange(output, "b (f h w) c -> b c f h w", f=video_length, h=height, w=width) + return output + +class TemporalTransformer3DModel(nn.Module): + def __init__( + self, + in_channels, + num_attention_heads, + attention_head_dim, + + num_layers, + attention_block_types = ( "Temporal_Self", "Temporal_Self", ), + dropout = 0.0, + norm_num_groups = 32, + cross_attention_dim = 768, + activation_fn = "geglu", + attention_bias = False, + upcast_attention = False, + + cross_frame_attention_mode = None, + temporal_position_encoding = False, + temporal_position_encoding_max_len = 4096, + grid = False, + block_size = 1, + remove_time_embedding_in_photo = False, + qk_norm = False, + ): + super().__init__() + + inner_dim = num_attention_heads * attention_head_dim + + self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True) + self.proj_in = nn.Linear(in_channels, inner_dim) + + self.block_size = block_size + self.transformer_blocks = nn.ModuleList( + [ + TemporalTransformerBlock( + dim=inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + attention_block_types=attention_block_types, + dropout=dropout, + norm_num_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim, + activation_fn=activation_fn, + attention_bias=attention_bias, + upcast_attention=upcast_attention, + cross_frame_attention_mode=cross_frame_attention_mode, + temporal_position_encoding=temporal_position_encoding, + temporal_position_encoding_max_len=temporal_position_encoding_max_len, + block_size=block_size, + grid=grid, + remove_time_embedding_in_photo=remove_time_embedding_in_photo, + qk_norm=qk_norm + ) + for d in range(num_layers) + ] + ) + self.proj_out = nn.Linear(inner_dim, in_channels) + + def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None): + assert hidden_states.dim() == 5, f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}." + video_length = hidden_states.shape[2] + hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w") + + batch, channel, height, weight = hidden_states.shape + residual = hidden_states + + hidden_states = self.norm(hidden_states) + inner_dim = hidden_states.shape[1] + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * weight, inner_dim) + hidden_states = self.proj_in(hidden_states) + + # Transformer Blocks + for block in self.transformer_blocks: + hidden_states = block(hidden_states, encoder_hidden_states=encoder_hidden_states, video_length=video_length, height=height, weight=weight) + + # output + hidden_states = self.proj_out(hidden_states) + hidden_states = hidden_states.reshape(batch, height, weight, inner_dim).permute(0, 3, 1, 2).contiguous() + + output = hidden_states + residual + output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length) + + return output + +class TemporalTransformerBlock(nn.Module): + def __init__( + self, + dim, + num_attention_heads, + attention_head_dim, + attention_block_types = ( "Temporal_Self", "Temporal_Self", ), + dropout = 0.0, + norm_num_groups = 32, + cross_attention_dim = 768, + activation_fn = "geglu", + attention_bias = False, + upcast_attention = False, + cross_frame_attention_mode = None, + temporal_position_encoding = False, + temporal_position_encoding_max_len = 4096, + block_size = 1, + grid = False, + remove_time_embedding_in_photo = False, + qk_norm = False, + ): + super().__init__() + + attention_blocks = [] + norms = [] + + for block_name in attention_block_types: + attention_blocks.append( + VersatileAttention( + attention_mode=block_name.split("_")[0], + cross_attention_dim=cross_attention_dim if block_name.endswith("_Cross") else None, + + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + + cross_frame_attention_mode=cross_frame_attention_mode, + temporal_position_encoding=temporal_position_encoding, + temporal_position_encoding_max_len=temporal_position_encoding_max_len, + block_size=block_size, + grid=grid, + remove_time_embedding_in_photo=remove_time_embedding_in_photo, + qk_norm="layer_norm" if qk_norm else None, + processor=HunyuanAttnProcessor2_0() if qk_norm else AttnProcessor2_0(), + ) if pkg_resources.parse_version(installed_version) >= pkg_resources.parse_version("0.28.2") else \ + VersatileAttention( + attention_mode=block_name.split("_")[0], + cross_attention_dim=cross_attention_dim if block_name.endswith("_Cross") else None, + + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + + cross_frame_attention_mode=cross_frame_attention_mode, + temporal_position_encoding=temporal_position_encoding, + temporal_position_encoding_max_len=temporal_position_encoding_max_len, + block_size=block_size, + grid=grid, + remove_time_embedding_in_photo=remove_time_embedding_in_photo, + ) + ) + norms.append(FP32LayerNorm(dim)) + + self.attention_blocks = nn.ModuleList(attention_blocks) + self.norms = nn.ModuleList(norms) + + self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn) + self.ff_norm = FP32LayerNorm(dim) + + def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, video_length=None, height=None, weight=None): + for attention_block, norm in zip(self.attention_blocks, self.norms): + norm_hidden_states = norm(hidden_states) + hidden_states = attention_block( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states if attention_block.is_cross_attention else None, + video_length=video_length, + height=height, + weight=weight, + ) + hidden_states + + hidden_states = self.ff(self.ff_norm(hidden_states)) + hidden_states + + output = hidden_states + return output + +class PositionalEncoding(nn.Module): + def __init__( + self, + d_model, + dropout = 0., + max_len = 4096 + ): + super().__init__() + self.dropout = nn.Dropout(p=dropout) + position = torch.arange(max_len).unsqueeze(1) + div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) + pe = torch.zeros(1, max_len, d_model) + pe[0, :, 0::2] = torch.sin(position * div_term) + pe[0, :, 1::2] = torch.cos(position * div_term) + self.register_buffer('pe', pe) + + def forward(self, x): + x = x + self.pe[:, :x.size(1)] + return self.dropout(x) + +class VersatileAttention(Attention): + def __init__( + self, + attention_mode = None, + cross_frame_attention_mode = None, + temporal_position_encoding = False, + temporal_position_encoding_max_len = 4096, + grid = False, + block_size = 1, + remove_time_embedding_in_photo = False, + *args, **kwargs + ): + super().__init__(*args, **kwargs) + assert attention_mode == "Temporal" or attention_mode == "Global" + + self.attention_mode = attention_mode + self.is_cross_attention = kwargs["cross_attention_dim"] is not None + + self.block_size = block_size + self.grid = grid + self.remove_time_embedding_in_photo = remove_time_embedding_in_photo + self.pos_encoder = PositionalEncoding( + kwargs["query_dim"], + dropout=0., + max_len=temporal_position_encoding_max_len + ) if (temporal_position_encoding and attention_mode == "Temporal") or (temporal_position_encoding and attention_mode == "Global") else None + + def extra_repr(self): + return f"(Module Info) Attention_Mode: {self.attention_mode}, Is_Cross_Attention: {self.is_cross_attention}" + + def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, video_length=None, height=None, weight=None): + batch_size, sequence_length, _ = hidden_states.shape + + if self.attention_mode == "Temporal": + # for add pos_encoder + _, before_d, _c = hidden_states.size() + hidden_states = rearrange(hidden_states, "(b f) d c -> (b d) f c", f=video_length) + + if self.remove_time_embedding_in_photo: + if self.pos_encoder is not None and video_length > 1: + hidden_states = self.pos_encoder(hidden_states) + else: + if self.pos_encoder is not None: + hidden_states = self.pos_encoder(hidden_states) + + if self.grid: + hidden_states = rearrange(hidden_states, "(b d) f c -> b f d c", f=video_length, d=before_d) + hidden_states = rearrange(hidden_states, "b f (h w) c -> b f h w c", h=height, w=weight) + + hidden_states = rearrange(hidden_states, "b f (h n) (w m) c -> (b h w) (f n m) c", n=self.block_size, m=self.block_size) + d = before_d // self.block_size // self.block_size + else: + d = before_d + encoder_hidden_states = repeat(encoder_hidden_states, "b n c -> (b d) n c", d=d) if encoder_hidden_states is not None else encoder_hidden_states + elif self.attention_mode == "Global": + # for add pos_encoder + _, d, _c = hidden_states.size() + hidden_states = rearrange(hidden_states, "(b f) d c -> (b d) f c", f=video_length) + if self.pos_encoder is not None: + hidden_states = self.pos_encoder(hidden_states) + hidden_states = rearrange(hidden_states, "(b d) f c -> b (f d) c", f=video_length, d=d) + else: + raise NotImplementedError + + encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states + + bs = 512 + new_hidden_states = [] + for i in range(0, hidden_states.shape[0], bs): + __hidden_states = super().forward( + hidden_states[i : i + bs], + encoder_hidden_states=encoder_hidden_states[i : i + bs], + attention_mask=attention_mask + ) + new_hidden_states.append(__hidden_states) + hidden_states = torch.cat(new_hidden_states, dim = 0) + + if self.attention_mode == "Temporal": + hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d) + if self.grid: + hidden_states = rearrange(hidden_states, "(b f n m) (h w) c -> (b f) h n w m c", f=video_length, n=self.block_size, m=self.block_size, h=height // self.block_size, w=weight // self.block_size) + hidden_states = rearrange(hidden_states, "b h n w m c -> b (h n) (w m) c") + hidden_states = rearrange(hidden_states, "b h w c -> b (h w) c") + elif self.attention_mode == "Global": + hidden_states = rearrange(hidden_states, "b (f d) c -> (b f) d c", f=video_length, d=d) + + return hidden_states \ No newline at end of file diff --git a/ruyi/models/norm.py b/ruyi/models/norm.py new file mode 100644 index 0000000..c6b72b6 --- /dev/null +++ b/ruyi/models/norm.py @@ -0,0 +1,117 @@ +from typing import Any, Dict, Optional, Tuple + +import torch +import torch.nn.functional as F +from diffusers.models.embeddings import TimestepEmbedding, Timesteps +from torch import nn + + +def zero_module(module): + # Zero out the parameters of a module and return it. + for p in module.parameters(): + p.detach().zero_() + return module + + +class FP32LayerNorm(nn.LayerNorm): + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + origin_dtype = inputs.dtype + if hasattr(self, 'weight') and self.weight is not None: + return F.layer_norm( + inputs.float(), self.normalized_shape, self.weight.float(), self.bias.float(), self.eps + ).to(origin_dtype) + else: + return F.layer_norm( + inputs.float(), self.normalized_shape, None, None, self.eps + ).to(origin_dtype) + +class PixArtAlphaCombinedTimestepSizeEmbeddings(nn.Module): + """ + For PixArt-Alpha. + + Reference: + https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L164C9-L168C29 + """ + + def __init__(self, embedding_dim, size_emb_dim, use_additional_conditions: bool = False): + super().__init__() + + self.outdim = size_emb_dim + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + self.use_additional_conditions = use_additional_conditions + if use_additional_conditions: + self.additional_condition_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.resolution_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=size_emb_dim) + self.aspect_ratio_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=size_emb_dim) + + self.resolution_embedder.linear_2 = zero_module(self.resolution_embedder.linear_2) + self.aspect_ratio_embedder.linear_2 = zero_module(self.aspect_ratio_embedder.linear_2) + + def forward(self, timestep, resolution, aspect_ratio, batch_size, hidden_dtype): + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) + + if self.use_additional_conditions: + resolution_emb = self.additional_condition_proj(resolution.flatten()).to(hidden_dtype) + resolution_emb = self.resolution_embedder(resolution_emb).reshape(batch_size, -1) + aspect_ratio_emb = self.additional_condition_proj(aspect_ratio.flatten()).to(hidden_dtype) + aspect_ratio_emb = self.aspect_ratio_embedder(aspect_ratio_emb).reshape(batch_size, -1) + conditioning = timesteps_emb + torch.cat([resolution_emb, aspect_ratio_emb], dim=1) + else: + conditioning = timesteps_emb + + return conditioning + +class AdaLayerNormSingle(nn.Module): + r""" + Norm layer adaptive layer norm single (adaLN-single). + + As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3). + + Parameters: + embedding_dim (`int`): The size of each embedding vector. + use_additional_conditions (`bool`): To use additional conditions for normalization or not. + """ + + def __init__(self, embedding_dim: int, use_additional_conditions: bool = False): + super().__init__() + + self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings( + embedding_dim, size_emb_dim=embedding_dim // 3, use_additional_conditions=use_additional_conditions + ) + + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True) + + def forward( + self, + timestep: torch.Tensor, + added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, + batch_size: Optional[int] = None, + hidden_dtype: Optional[torch.dtype] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + # No modulation happening here. + embedded_timestep = self.emb(timestep, **added_cond_kwargs, batch_size=batch_size, hidden_dtype=hidden_dtype) + return self.linear(self.silu(embedded_timestep)), embedded_timestep + +class AdaLayerNormShift(nn.Module): + r""" + Norm layer modified to incorporate timestep embeddings. + + Parameters: + embedding_dim (`int`): The size of each embedding vector. + num_embeddings (`int`): The size of the embeddings dictionary. + """ + + def __init__(self, embedding_dim: int, elementwise_affine=True, eps=1e-6): + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, embedding_dim) + self.norm = FP32LayerNorm(embedding_dim, elementwise_affine=elementwise_affine, eps=eps) + + def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: + shift = self.linear(self.silu(emb.to(torch.float32)).to(emb.dtype)) + x = self.norm(x) + shift.unsqueeze(dim=1) + return x \ No newline at end of file diff --git a/ruyi/models/patch.py b/ruyi/models/patch.py new file mode 100644 index 0000000..4391bed --- /dev/null +++ b/ruyi/models/patch.py @@ -0,0 +1,426 @@ +import math +from typing import Optional + +import numpy as np +import torch +import torch.nn.functional as F +import torch.nn.init as init +from einops import rearrange +from torch import nn + + +def get_2d_sincos_pos_embed( + embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16 +): + """ + grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or + [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token) + """ + if isinstance(grid_size, int): + grid_size = (grid_size, grid_size) + + grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size) / interpolation_scale + grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size) / interpolation_scale + grid = np.meshgrid(grid_w, grid_h) # here w goes first + grid = np.stack(grid, axis=0) + + grid = grid.reshape([2, 1, grid_size[1], grid_size[0]]) + pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid) + if cls_token and extra_tokens > 0: + pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0) + return pos_embed + + +def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): + if embed_dim % 2 != 0: + raise ValueError("embed_dim must be divisible by 2") + + # use half of dimensions to encode grid_h + emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2) + emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2) + + emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D) + return emb + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D) + """ + if embed_dim % 2 != 0: + raise ValueError("embed_dim must be divisible by 2") + + omega = np.arange(embed_dim // 2, dtype=np.float64) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb + +class Patch1D(nn.Module): + def __init__( + self, + channels: int, + use_conv: bool = False, + out_channels: Optional[int] = None, + stride: int = 2, + padding: int = 0, + name: str = "conv", + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.padding = padding + self.name = name + + if use_conv: + self.conv = nn.Conv1d(self.channels, self.out_channels, stride, stride=stride, padding=padding) + init.constant_(self.conv.weight, 0.0) + with torch.no_grad(): + for i in range(len(self.conv.weight)): self.conv.weight[i, i] = 1 / stride + init.constant_(self.conv.bias, 0.0) + else: + assert self.channels == self.out_channels + self.conv = nn.AvgPool1d(kernel_size=stride, stride=stride) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + assert inputs.shape[1] == self.channels + return self.conv(inputs) + +class UnPatch1D(nn.Module): + def __init__( + self, + channels: int, + use_conv: bool = False, + use_conv_transpose: bool = False, + out_channels: Optional[int] = None, + name: str = "conv", + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_conv_transpose = use_conv_transpose + self.name = name + + self.conv = None + if use_conv_transpose: + self.conv = nn.ConvTranspose1d(channels, self.out_channels, 4, 2, 1) + elif use_conv: + self.conv = nn.Conv1d(self.channels, self.out_channels, 3, padding=1) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + assert inputs.shape[1] == self.channels + if self.use_conv_transpose: + return self.conv(inputs) + + outputs = F.interpolate(inputs, scale_factor=2.0, mode="nearest") + + if self.use_conv: + outputs = self.conv(outputs) + + return outputs + +class Upsampler(nn.Module): + def __init__( + self, + spatial_upsample_factor: int = 1, + temporal_upsample_factor: int = 1, + ): + super().__init__() + + self.spatial_upsample_factor = spatial_upsample_factor + self.temporal_upsample_factor = temporal_upsample_factor + +class TemporalUpsampler3D(Upsampler): + def __init__(self): + super().__init__( + spatial_upsample_factor=1, + temporal_upsample_factor=2, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.shape[2] > 1: + first_frame, x = x[:, :, :1], x[:, :, 1:] + x = F.interpolate(x, scale_factor=(2, 1, 1), mode="trilinear") + x = torch.cat([first_frame, x], dim=2) + return x + +def cast_tuple(t, length = 1): + return t if isinstance(t, tuple) else ((t,) * length) + +def divisible_by(num, den): + return (num % den) == 0 + +def is_odd(n): + return not divisible_by(n, 2) + +class CausalConv3d(nn.Conv3d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size=3, # : int | tuple[int, int, int], + stride=1, # : int | tuple[int, int, int] = 1, + padding=1, # : int | tuple[int, int, int], # TODO: change it to 0. + dilation=1, # : int | tuple[int, int, int] = 1, + **kwargs, + ): + kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size,) * 3 + assert len(kernel_size) == 3, f"Kernel size must be a 3-tuple, got {kernel_size} instead." + + stride = stride if isinstance(stride, tuple) else (stride,) * 3 + assert len(stride) == 3, f"Stride must be a 3-tuple, got {stride} instead." + + dilation = dilation if isinstance(dilation, tuple) else (dilation,) * 3 + assert len(dilation) == 3, f"Dilation must be a 3-tuple, got {dilation} instead." + + t_ks, h_ks, w_ks = kernel_size + _, h_stride, w_stride = stride + t_dilation, h_dilation, w_dilation = dilation + + t_pad = (t_ks - 1) * t_dilation + # TODO: align with SD + if padding is None: + h_pad = math.ceil(((h_ks - 1) * h_dilation + (1 - h_stride)) / 2) + w_pad = math.ceil(((w_ks - 1) * w_dilation + (1 - w_stride)) / 2) + elif isinstance(padding, int): + h_pad = w_pad = padding + else: + assert NotImplementedError + + self.temporal_padding = t_pad + + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + padding=(0, h_pad, w_pad), + **kwargs, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x: (B, C, T, H, W) + x = F.pad( + x, + pad=(0, 0, 0, 0, self.temporal_padding, 0), + mode="replicate", # TODO: check if this is necessary + ) + return super().forward(x) + +class PatchEmbed3D(nn.Module): + """3D Image to Patch Embedding""" + + def __init__( + self, + height=224, + width=224, + patch_size=16, + time_patch_size=4, + in_channels=3, + embed_dim=768, + layer_norm=False, + flatten=True, + bias=True, + interpolation_scale=1, + ): + super().__init__() + + num_patches = (height // patch_size) * (width // patch_size) + self.flatten = flatten + self.layer_norm = layer_norm + + self.proj = nn.Conv3d( + in_channels, embed_dim, kernel_size=(time_patch_size, patch_size, patch_size), stride=(time_patch_size, patch_size, patch_size), bias=bias + ) + if layer_norm: + self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6) + else: + self.norm = None + + self.patch_size = patch_size + # See: + # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L161 + self.height, self.width = height // patch_size, width // patch_size + self.base_size = height // patch_size + self.interpolation_scale = interpolation_scale + pos_embed = get_2d_sincos_pos_embed( + embed_dim, int(num_patches**0.5), base_size=self.base_size, interpolation_scale=self.interpolation_scale + ) + self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False) + + def forward(self, latent): + height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size + + latent = self.proj(latent) + latent = rearrange(latent, "b c f h w -> (b f) c h w") + if self.flatten: + latent = latent.flatten(2).transpose(1, 2) # BCFHW -> BNC + if self.layer_norm: + latent = self.norm(latent) + # Interpolate positional embeddings if needed. + # (For PixArt-Alpha: https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L162C151-L162C160) + if self.height != height or self.width != width: + pos_embed = get_2d_sincos_pos_embed( + embed_dim=self.pos_embed.shape[-1], + grid_size=(height, width), + base_size=self.base_size, + interpolation_scale=self.interpolation_scale, + ) + pos_embed = torch.from_numpy(pos_embed) + pos_embed = pos_embed.float().unsqueeze(0).to(latent.device) + else: + pos_embed = self.pos_embed + + return (latent + pos_embed).to(latent.dtype) + +class PatchEmbedF3D(nn.Module): + """Fake 3D Image to Patch Embedding""" + + def __init__( + self, + height=224, + width=224, + patch_size=16, + in_channels=3, + embed_dim=768, + layer_norm=False, + flatten=True, + bias=True, + interpolation_scale=1, + ): + super().__init__() + + num_patches = (height // patch_size) * (width // patch_size) + self.flatten = flatten + self.layer_norm = layer_norm + + self.proj = nn.Conv2d( + in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias + ) + self.proj_t = Patch1D( + embed_dim, True, stride=patch_size + ) + if layer_norm: + self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6) + else: + self.norm = None + + self.patch_size = patch_size + # See: + # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L161 + self.height, self.width = height // patch_size, width // patch_size + self.base_size = height // patch_size + self.interpolation_scale = interpolation_scale + pos_embed = get_2d_sincos_pos_embed( + embed_dim, int(num_patches**0.5), base_size=self.base_size, interpolation_scale=self.interpolation_scale + ) + self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False) + + def forward(self, latent): + height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size + b, c, f, h, w = latent.size() + latent = rearrange(latent, "b c f h w -> (b f) c h w") + latent = self.proj(latent) + latent = rearrange(latent, "(b f) c h w -> b c f h w", f=f) + + latent = rearrange(latent, "b c f h w -> (b h w) c f") + latent = self.proj_t(latent) + latent = rearrange(latent, "(b h w) c f -> b c f h w", h=h//2, w=w//2) + + latent = rearrange(latent, "b c f h w -> (b f) c h w") + if self.flatten: + latent = latent.flatten(2).transpose(1, 2) # BCFHW -> BNC + if self.layer_norm: + latent = self.norm(latent) + + # Interpolate positional embeddings if needed. + # (For PixArt-Alpha: https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L162C151-L162C160) + if self.height != height or self.width != width: + pos_embed = get_2d_sincos_pos_embed( + embed_dim=self.pos_embed.shape[-1], + grid_size=(height, width), + base_size=self.base_size, + interpolation_scale=self.interpolation_scale, + ) + pos_embed = torch.from_numpy(pos_embed) + pos_embed = pos_embed.float().unsqueeze(0).to(latent.device) + else: + pos_embed = self.pos_embed + + return (latent + pos_embed).to(latent.dtype) + +class CasualPatchEmbed3D(nn.Module): + """3D Image to Patch Embedding""" + + def __init__( + self, + height=224, + width=224, + patch_size=16, + time_patch_size=4, + in_channels=3, + embed_dim=768, + layer_norm=False, + flatten=True, + bias=True, + interpolation_scale=1, + ): + super().__init__() + + num_patches = (height // patch_size) * (width // patch_size) + self.flatten = flatten + self.layer_norm = layer_norm + + self.proj = CausalConv3d( + in_channels, embed_dim, kernel_size=(time_patch_size, patch_size, patch_size), stride=(time_patch_size, patch_size, patch_size), bias=bias, padding=None + ) + if layer_norm: + self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6) + else: + self.norm = None + + self.patch_size = patch_size + # See: + # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L161 + self.height, self.width = height // patch_size, width // patch_size + self.base_size = height // patch_size + self.interpolation_scale = interpolation_scale + pos_embed = get_2d_sincos_pos_embed( + embed_dim, int(num_patches**0.5), base_size=self.base_size, interpolation_scale=self.interpolation_scale + ) + self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False) + + def forward(self, latent): + height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size + + latent = self.proj(latent) + latent = rearrange(latent, "b c f h w -> (b f) c h w") + if self.flatten: + latent = latent.flatten(2).transpose(1, 2) # BCFHW -> BNC + if self.layer_norm: + latent = self.norm(latent) + # Interpolate positional embeddings if needed. + # (For PixArt-Alpha: https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L162C151-L162C160) + if self.height != height or self.width != width: + pos_embed = get_2d_sincos_pos_embed( + embed_dim=self.pos_embed.shape[-1], + grid_size=(height, width), + base_size=self.base_size, + interpolation_scale=self.interpolation_scale, + ) + pos_embed = torch.from_numpy(pos_embed) + pos_embed = pos_embed.float().unsqueeze(0).to(latent.device) + else: + pos_embed = self.pos_embed + + return (latent + pos_embed).to(latent.dtype) diff --git a/ruyi/models/resampler.py b/ruyi/models/resampler.py new file mode 100644 index 0000000..43a0612 --- /dev/null +++ b/ruyi/models/resampler.py @@ -0,0 +1,146 @@ +# Copyright (c) Alibaba Cloud. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +import math + +import numpy as np +import torch +from torch import nn +from torch.nn import functional as F +from torch.nn.init import normal_ + + +def get_abs_pos(abs_pos, tgt_size): + # abs_pos: L, C + # tgt_size: M + # return: M, C + src_size = int(math.sqrt(abs_pos.size(0))) + tgt_size = int(math.sqrt(tgt_size)) + dtype = abs_pos.dtype + + if src_size != tgt_size: + return F.interpolate( + abs_pos.float().reshape(1, src_size, src_size, -1).permute(0, 3, 1, 2), + size=(tgt_size, tgt_size), + mode="bicubic", + align_corners=False, + ).permute(0, 2, 3, 1).flatten(0, 2).to(dtype=dtype) + else: + return abs_pos + +# https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/util/pos_embed.py#L20 +def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False): + """ + grid_size: int of the grid height and width + return: + pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token) + """ + grid_h = np.arange(grid_size, dtype=np.float32) + grid_w = np.arange(grid_size, dtype=np.float32) + grid = np.meshgrid(grid_w, grid_h) # here w goes first + grid = np.stack(grid, axis=0) + + grid = grid.reshape([2, 1, grid_size, grid_size]) + pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid) + if cls_token: + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) + return pos_embed + + +def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): + assert embed_dim % 2 == 0 + + # use half of dimensions to encode grid_h + emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2) + emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2) + + emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D) + return emb + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + out: (M, D) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=np.float32) + omega /= embed_dim / 2. + omega = 1. / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb + +class Resampler(nn.Module): + """ + A 2D perceiver-resampler network with one cross attention layers by + (grid_size**2) learnable queries and 2d sincos pos_emb + Outputs: + A tensor with the shape of (grid_size**2, embed_dim) + """ + def __init__( + self, + grid_size, + embed_dim, + num_heads, + kv_dim=None, + norm_layer=nn.LayerNorm + ): + super().__init__() + self.num_queries = grid_size ** 2 + self.embed_dim = embed_dim + self.num_heads = num_heads + + self.pos_embed = nn.Parameter( + torch.from_numpy(get_2d_sincos_pos_embed(embed_dim, grid_size)).float() + ).requires_grad_(False) + + self.query = nn.Parameter(torch.zeros(self.num_queries, embed_dim)) + normal_(self.query, std=.02) + + if kv_dim is not None and kv_dim != embed_dim: + self.kv_proj = nn.Linear(kv_dim, embed_dim, bias=False) + else: + self.kv_proj = nn.Identity() + + self.attn = nn.MultiheadAttention(embed_dim, num_heads) + self.ln_q = norm_layer(embed_dim) + self.ln_kv = norm_layer(embed_dim) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + def forward(self, x, key_padding_mask=None): + pos_embed = get_abs_pos(self.pos_embed, x.size(1)) + + x = self.kv_proj(x) + x = self.ln_kv(x).permute(1, 0, 2) + + N = x.shape[1] + q = self.ln_q(self.query) + out = self.attn( + self._repeat(q, N) + self.pos_embed.unsqueeze(1), + x + pos_embed.unsqueeze(1), + x, + key_padding_mask=key_padding_mask)[0] + return out.permute(1, 0, 2) + + def _repeat(self, query, N: int): + return query.unsqueeze(1).repeat(1, N, 1) \ No newline at end of file diff --git a/ruyi/models/transformer2d.py b/ruyi/models/transformer2d.py new file mode 100644 index 0000000..836d8da --- /dev/null +++ b/ruyi/models/transformer2d.py @@ -0,0 +1,555 @@ +# Copyright 2023 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 json +import os +from dataclasses import dataclass +from typing import Any, Dict, Optional + +import numpy as np +import torch +import torch.nn.functional as F +import torch.nn.init as init +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.attention import BasicTransformerBlock +from diffusers.models.embeddings import ImagePositionalEmbeddings, PatchEmbed +from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear +from diffusers.models.modeling_utils import ModelMixin +from diffusers.models.normalization import AdaLayerNormSingle +from diffusers.utils import (USE_PEFT_BACKEND, BaseOutput, deprecate, + is_torch_version) +from einops import rearrange +from torch import nn + +try: + from diffusers.models.embeddings import PixArtAlphaTextProjection +except: + from diffusers.models.embeddings import \ + CaptionProjection as PixArtAlphaTextProjection + +from .attention import (KVCompressionTransformerBlock, + SelfAttentionTemporalTransformerBlock, + TemporalTransformerBlock) + + +@dataclass +class Transformer2DModelOutput(BaseOutput): + """ + The output of [`Transformer2DModel`]. + + Args: + sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete): + The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability + distributions for the unnoised latent pixels. + """ + + sample: torch.FloatTensor + + +class Transformer2DModel(ModelMixin, ConfigMixin): + """ + A 2D Transformer model for image-like data. + + Parameters: + num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention. + attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head. + in_channels (`int`, *optional*): + The number of channels in the input and output (specify if the input is **continuous**). + num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use. + sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**). + This is fixed during training since it is used to learn a number of position embeddings. + num_vector_embeds (`int`, *optional*): + The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**). + Includes the class for the masked latent pixel. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward. + num_embeds_ada_norm ( `int`, *optional*): + The number of diffusion steps used during training. Pass if at least one of the norm_layers is + `AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are + added to the hidden states. + + During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`. + attention_bias (`bool`, *optional*): + Configure if the `TransformerBlocks` attention should contain a bias parameter. + """ + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + num_attention_heads: int = 16, + attention_head_dim: int = 88, + in_channels: Optional[int] = None, + out_channels: Optional[int] = None, + num_layers: int = 1, + dropout: float = 0.0, + norm_num_groups: int = 32, + cross_attention_dim: Optional[int] = None, + attention_bias: bool = False, + sample_size: Optional[int] = None, + num_vector_embeds: Optional[int] = None, + patch_size: Optional[int] = None, + activation_fn: str = "geglu", + num_embeds_ada_norm: Optional[int] = None, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + double_self_attention: bool = False, + upcast_attention: bool = False, + norm_type: str = "layer_norm", + norm_elementwise_affine: bool = True, + norm_eps: float = 1e-5, + attention_type: str = "default", + caption_channels: int = None, + # block type + basic_block_type: str = "basic", + ): + super().__init__() + self.use_linear_projection = use_linear_projection + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + self.basic_block_type = basic_block_type + inner_dim = num_attention_heads * attention_head_dim + + conv_cls = nn.Conv2d if USE_PEFT_BACKEND else LoRACompatibleConv + linear_cls = nn.Linear if USE_PEFT_BACKEND else LoRACompatibleLinear + + # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)` + # Define whether input is continuous or discrete depending on configuration + self.is_input_continuous = (in_channels is not None) and (patch_size is None) + self.is_input_vectorized = num_vector_embeds is not None + self.is_input_patches = in_channels is not None and patch_size is not None + + if norm_type == "layer_norm" and num_embeds_ada_norm is not None: + deprecation_message = ( + f"The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or" + " incorrectly set to `'layer_norm'`.Make sure to set `norm_type` to `'ada_norm'` in the config." + " Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect" + " results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it" + " would be very nice if you could open a Pull request for the `transformer/config.json` file" + ) + deprecate("norm_type!=num_embeds_ada_norm", "1.0.0", deprecation_message, standard_warn=False) + norm_type = "ada_norm" + + if self.is_input_continuous and self.is_input_vectorized: + raise ValueError( + f"Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make" + " sure that either `in_channels` or `num_vector_embeds` is None." + ) + elif self.is_input_vectorized and self.is_input_patches: + raise ValueError( + f"Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make" + " sure that either `num_vector_embeds` or `num_patches` is None." + ) + elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches: + raise ValueError( + f"Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:" + f" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None." + ) + + # 2. Define input layers + if self.is_input_continuous: + self.in_channels = in_channels + + self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True) + if use_linear_projection: + self.proj_in = linear_cls(in_channels, inner_dim) + else: + self.proj_in = conv_cls(in_channels, inner_dim, kernel_size=1, stride=1, padding=0) + elif self.is_input_vectorized: + assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size" + assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed" + + self.height = sample_size + self.width = sample_size + self.num_vector_embeds = num_vector_embeds + self.num_latent_pixels = self.height * self.width + + self.latent_image_embedding = ImagePositionalEmbeddings( + num_embed=num_vector_embeds, embed_dim=inner_dim, height=self.height, width=self.width + ) + elif self.is_input_patches: + assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size" + + self.height = sample_size + self.width = sample_size + + self.patch_size = patch_size + interpolation_scale = self.config.sample_size // 64 # => 64 (= 512 pixart) has interpolation scale 1 + interpolation_scale = max(interpolation_scale, 1) + self.pos_embed = PatchEmbed( + height=sample_size, + width=sample_size, + patch_size=patch_size, + in_channels=in_channels, + embed_dim=inner_dim, + interpolation_scale=interpolation_scale, + ) + + basic_block = { + "basic": BasicTransformerBlock, + "kvcompression": KVCompressionTransformerBlock, + }[self.basic_block_type] + if self.basic_block_type == "kvcompression": + self.transformer_blocks = nn.ModuleList( + [ + basic_block( + inner_dim, + num_attention_heads, + attention_head_dim, + dropout=dropout, + cross_attention_dim=cross_attention_dim, + activation_fn=activation_fn, + num_embeds_ada_norm=num_embeds_ada_norm, + attention_bias=attention_bias, + only_cross_attention=only_cross_attention, + double_self_attention=double_self_attention, + upcast_attention=upcast_attention, + norm_type=norm_type, + norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, + attention_type=attention_type, + kvcompression=False if d < 14 else True, + ) + for d in range(num_layers) + ] + ) + else: + # 3. Define transformers blocks + self.transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + inner_dim, + num_attention_heads, + attention_head_dim, + dropout=dropout, + cross_attention_dim=cross_attention_dim, + activation_fn=activation_fn, + num_embeds_ada_norm=num_embeds_ada_norm, + attention_bias=attention_bias, + only_cross_attention=only_cross_attention, + double_self_attention=double_self_attention, + upcast_attention=upcast_attention, + norm_type=norm_type, + norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, + attention_type=attention_type, + ) + for d in range(num_layers) + ] + ) + + # 4. Define output layers + self.out_channels = in_channels if out_channels is None else out_channels + if self.is_input_continuous: + # TODO: should use out_channels for continuous projections + if use_linear_projection: + self.proj_out = linear_cls(inner_dim, in_channels) + else: + self.proj_out = conv_cls(inner_dim, in_channels, kernel_size=1, stride=1, padding=0) + elif self.is_input_vectorized: + self.norm_out = nn.LayerNorm(inner_dim) + self.out = nn.Linear(inner_dim, self.num_vector_embeds - 1) + elif self.is_input_patches and norm_type != "ada_norm_single": + self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out_1 = nn.Linear(inner_dim, 2 * inner_dim) + self.proj_out_2 = nn.Linear(inner_dim, patch_size * patch_size * self.out_channels) + elif self.is_input_patches and norm_type == "ada_norm_single": + self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6) + self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5) + self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * self.out_channels) + + # 5. PixArt-Alpha blocks. + self.adaln_single = None + self.use_additional_conditions = False + if norm_type == "ada_norm_single": + self.use_additional_conditions = self.config.sample_size == 128 + # TODO(Sayak, PVP) clean this, for now we use sample size to determine whether to use + # additional conditions until we find better name + self.adaln_single = AdaLayerNormSingle(inner_dim, use_additional_conditions=self.use_additional_conditions) + + self.caption_projection = None + if caption_channels is not None: + self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) + + self.gradient_checkpointing = False + + def _set_gradient_checkpointing(self, module, value=False): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = value + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + timestep: Optional[torch.LongTensor] = None, + added_cond_kwargs: Dict[str, torch.Tensor] = None, + class_labels: Optional[torch.LongTensor] = None, + cross_attention_kwargs: Dict[str, Any] = None, + attention_mask: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + return_dict: bool = True, + ): + """ + The [`Transformer2DModel`] forward method. + + Args: + hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous): + Input `hidden_states`. + encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*): + Conditional embeddings for cross attention layer. If not given, cross-attention defaults to + self-attention. + timestep ( `torch.LongTensor`, *optional*): + Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`. + class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*): + Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in + `AdaLayerZeroNorm`. + cross_attention_kwargs ( `Dict[str, Any]`, *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). + attention_mask ( `torch.Tensor`, *optional*): + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask + is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large + negative values to the attention scores corresponding to "discard" tokens. + encoder_attention_mask ( `torch.Tensor`, *optional*): + Cross-attention mask applied to `encoder_hidden_states`. Two formats supported: + + * Mask `(batch, sequence_length)` True = keep, False = discard. + * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard. + + If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format + above. This bias will be added to the cross-attention scores. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain + tuple. + + Returns: + If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a + `tuple` where the first element is the sample tensor. + """ + # ensure attention_mask is a bias, and give it a singleton query_tokens dimension. + # we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward. + # we can tell by counting dims; if ndim == 2: it's a mask rather than a bias. + # expects mask of shape: + # [batch, key_tokens] + # adds singleton query_tokens dimension: + # [batch, 1, key_tokens] + # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: + # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) + # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) + if attention_mask is not None and attention_mask.ndim == 2: + # assume that mask is expressed as: + # (1 = keep, 0 = discard) + # convert mask into a bias that can be added to attention scores: + # (keep = +0, discard = -10000.0) + attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0 + attention_mask = attention_mask.unsqueeze(1) + + # convert encoder_attention_mask to a bias the same way we do for attention_mask + if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2: + encoder_attention_mask = (1 - encoder_attention_mask.to(encoder_hidden_states.dtype)) * -10000.0 + encoder_attention_mask = encoder_attention_mask.unsqueeze(1) + + # Retrieve lora scale. + lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0 + + # 1. Input + if self.is_input_continuous: + batch, _, height, width = hidden_states.shape + residual = hidden_states + + hidden_states = self.norm(hidden_states) + if not self.use_linear_projection: + hidden_states = ( + self.proj_in(hidden_states, scale=lora_scale) + if not USE_PEFT_BACKEND + else self.proj_in(hidden_states) + ) + inner_dim = hidden_states.shape[1] + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim) + else: + inner_dim = hidden_states.shape[1] + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim) + hidden_states = ( + self.proj_in(hidden_states, scale=lora_scale) + if not USE_PEFT_BACKEND + else self.proj_in(hidden_states) + ) + + elif self.is_input_vectorized: + hidden_states = self.latent_image_embedding(hidden_states) + elif self.is_input_patches: + height, width = hidden_states.shape[-2] // self.patch_size, hidden_states.shape[-1] // self.patch_size + hidden_states = self.pos_embed(hidden_states) + + if self.adaln_single is not None: + if self.use_additional_conditions and added_cond_kwargs is None: + raise ValueError( + "`added_cond_kwargs` cannot be None when using additional conditions for `adaln_single`." + ) + batch_size = hidden_states.shape[0] + timestep, embedded_timestep = self.adaln_single( + timestep, added_cond_kwargs, batch_size=batch_size, hidden_dtype=hidden_states.dtype + ) + + # 2. Blocks + if self.caption_projection is not None: + batch_size = hidden_states.shape[0] + encoder_hidden_states = self.caption_projection(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1]) + + for block in self.transformer_blocks: + if self.training and self.gradient_checkpointing: + args = { + "basic": [], + "kvcompression": [1, height, width], + }[self.basic_block_type] + hidden_states = torch.utils.checkpoint.checkpoint( + block, + hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + timestep, + cross_attention_kwargs, + class_labels, + *args, + use_reentrant=False, + ) + else: + kwargs = { + "basic": {}, + "kvcompression": {"num_frames":1, "height":height, "width":width}, + }[self.basic_block_type] + hidden_states = block( + hidden_states, + attention_mask=attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + timestep=timestep, + cross_attention_kwargs=cross_attention_kwargs, + class_labels=class_labels, + **kwargs + ) + + # 3. Output + if self.is_input_continuous: + if not self.use_linear_projection: + hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous() + hidden_states = ( + self.proj_out(hidden_states, scale=lora_scale) + if not USE_PEFT_BACKEND + else self.proj_out(hidden_states) + ) + else: + hidden_states = ( + self.proj_out(hidden_states, scale=lora_scale) + if not USE_PEFT_BACKEND + else self.proj_out(hidden_states) + ) + hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous() + + output = hidden_states + residual + elif self.is_input_vectorized: + hidden_states = self.norm_out(hidden_states) + logits = self.out(hidden_states) + # (batch, self.num_vector_embeds - 1, self.num_latent_pixels) + logits = logits.permute(0, 2, 1) + + # log(p(x_0)) + output = F.log_softmax(logits.double(), dim=1).float() + + if self.is_input_patches: + if self.config.norm_type != "ada_norm_single": + conditioning = self.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = self.proj_out_2(hidden_states) + elif self.config.norm_type == "ada_norm_single": + shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None]).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states) + # Modulation + hidden_states = hidden_states * (1 + scale) + shift + hidden_states = self.proj_out(hidden_states) + hidden_states = hidden_states.squeeze(1) + + # unpatchify + if self.adaln_single is None: + height = width = int(hidden_states.shape[1] ** 0.5) + hidden_states = hidden_states.reshape( + shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels) + ) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size) + ) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) + + @classmethod + def from_pretrained(cls, pretrained_model_path, subfolder=None, patch_size=2, transformer_additional_kwargs={}): + if subfolder is not None: + pretrained_model_path = os.path.join(pretrained_model_path, subfolder) + print(f"loaded 2D transformer's pretrained weights from {pretrained_model_path} ...") + + config_file = os.path.join(pretrained_model_path, 'config.json') + if not os.path.isfile(config_file): + raise RuntimeError(f"{config_file} does not exist") + with open(config_file, "r") as f: + config = json.load(f) + + from diffusers.utils import WEIGHTS_NAME + model = cls.from_config(config, **transformer_additional_kwargs) + model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME) + model_file_safetensors = model_file.replace(".bin", ".safetensors") + if os.path.exists(model_file_safetensors): + from safetensors.torch import load_file, safe_open + state_dict = load_file(model_file_safetensors) + else: + if not os.path.isfile(model_file): + raise RuntimeError(f"{model_file} does not exist") + state_dict = torch.load(model_file, map_location="cpu") + + if model.state_dict()['pos_embed.proj.weight'].size() != state_dict['pos_embed.proj.weight'].size(): + new_shape = model.state_dict()['pos_embed.proj.weight'].size() + state_dict['pos_embed.proj.weight'] = torch.tile(state_dict['proj_out.weight'], [1, 2, 1, 1]) + + if model.state_dict()['proj_out.weight'].size() != state_dict['proj_out.weight'].size(): + new_shape = model.state_dict()['proj_out.weight'].size() + state_dict['proj_out.weight'] = torch.tile(state_dict['proj_out.weight'], [patch_size, 1]) + + if model.state_dict()['proj_out.bias'].size() != state_dict['proj_out.bias'].size(): + new_shape = model.state_dict()['proj_out.bias'].size() + state_dict['proj_out.bias'] = torch.tile(state_dict['proj_out.bias'], [patch_size]) + + tmp_state_dict = {} + for key in state_dict: + if key in model.state_dict().keys() and model.state_dict()[key].size() == state_dict[key].size(): + tmp_state_dict[key] = state_dict[key] + else: + print(key, "Size don't match, skip") + state_dict = tmp_state_dict + + m, u = model.load_state_dict(state_dict, strict=False) + print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};") + + params = [p.numel() if "motion_modules." in n else 0 for n, p in model.named_parameters()] + print(f"### Postion Parameters: {sum(params) / 1e6} M") + + return model \ No newline at end of file diff --git a/ruyi/models/transformer3d.py b/ruyi/models/transformer3d.py new file mode 100644 index 0000000..1e8b67b --- /dev/null +++ b/ruyi/models/transformer3d.py @@ -0,0 +1,1239 @@ +# Copyright 2024 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 json +import math +import os +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple + +import numpy as np +import torch +import torch.nn.functional as F +import torch.nn.init as init +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.attention import BasicTransformerBlock, FeedForward +from diffusers.models.embeddings import (PatchEmbed, + PixArtAlphaTextProjection, TimestepEmbedding, Timesteps) +from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.models.normalization import AdaLayerNormContinuous +from diffusers.utils import (USE_PEFT_BACKEND, BaseOutput, is_torch_version, + logging) +from diffusers.utils.torch_utils import maybe_allow_in_graph +from einops import rearrange +from torch import nn + +from .attention import (HunyuanDiTBlock, HunyuanTemporalTransformerBlock, + SelfAttentionTemporalTransformerBlock, + TemporalTransformerBlock) +from .embeddings import HunyuanCombinedTimestepTextSizeStyleEmbedding +from .norm import AdaLayerNormSingle +from .patch import (CasualPatchEmbed3D, Patch1D, PatchEmbed3D, PatchEmbedF3D, + TemporalUpsampler3D, UnPatch1D) +from .resampler import Resampler + +try: + from diffusers.models.embeddings import PixArtAlphaTextProjection +except: + from diffusers.models.embeddings import \ + CaptionProjection as PixArtAlphaTextProjection + +def zero_module(module): + # Zero out the parameters of a module and return it. + for p in module.parameters(): + p.detach().zero_() + return module + + +class CLIPProjection(nn.Module): + """ + Projects caption embeddings. Also handles dropout for classifier-free guidance. + + Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py + """ + + def __init__(self, in_features, hidden_size, num_tokens=120): + super().__init__() + self.linear_1 = nn.Linear(in_features=in_features, out_features=hidden_size, bias=True) + self.act_1 = nn.GELU(approximate="tanh") + self.linear_2 = nn.Linear(in_features=hidden_size, out_features=hidden_size, bias=True) + self.linear_2 = zero_module(self.linear_2) + def forward(self, caption): + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + +class TimePositionalEncoding(nn.Module): + def __init__( + self, + d_model, + dropout = 0., + max_len = 24 + ): + super().__init__() + self.dropout = nn.Dropout(p=dropout) + position = torch.arange(max_len).unsqueeze(1) + div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) + pe = torch.zeros(1, max_len, d_model) + pe[0, :, 0::2] = torch.sin(position * div_term) + pe[0, :, 1::2] = torch.cos(position * div_term) + self.register_buffer('pe', pe) + + def forward(self, x): + b, c, f, h, w = x.size() + x = rearrange(x, "b c f h w -> (b h w) f c") + x = x + self.pe[:, :x.size(1)] + x = rearrange(x, "(b h w) f c -> b c f h w", b=b, h=h, w=w) + return self.dropout(x) + +@dataclass +class Transformer3DModelOutput(BaseOutput): + """ + The output of [`Transformer2DModel`]. + + Args: + sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete): + The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability + distributions for the unnoised latent pixels. + """ + + sample: torch.FloatTensor + + +class Transformer3DModel(ModelMixin, ConfigMixin): + """ + A 3D Transformer model for image-like data. + + Parameters: + num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention. + attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head. + in_channels (`int`, *optional*): + The number of channels in the input and output (specify if the input is **continuous**). + num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use. + sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**). + This is fixed during training since it is used to learn a number of position embeddings. + num_vector_embeds (`int`, *optional*): + The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**). + Includes the class for the masked latent pixel. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward. + num_embeds_ada_norm ( `int`, *optional*): + The number of diffusion steps used during training. Pass if at least one of the norm_layers is + `AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are + added to the hidden states. + + During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`. + attention_bias (`bool`, *optional*): + Configure if the `TransformerBlocks` attention should contain a bias parameter. + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + num_attention_heads: int = 16, + attention_head_dim: int = 88, + in_channels: Optional[int] = None, + out_channels: Optional[int] = None, + num_layers: int = 1, + dropout: float = 0.0, + norm_num_groups: int = 32, + cross_attention_dim: Optional[int] = None, + attention_bias: bool = False, + sample_size: Optional[int] = None, + num_vector_embeds: Optional[int] = None, + patch_size: Optional[int] = None, + activation_fn: str = "geglu", + num_embeds_ada_norm: Optional[int] = None, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + double_self_attention: bool = False, + upcast_attention: bool = False, + norm_type: str = "layer_norm", + norm_elementwise_affine: bool = True, + norm_eps: float = 1e-5, + attention_type: str = "default", + caption_channels: int = None, + # block type + basic_block_type: str = "motionmodule", + # enable_uvit + enable_uvit: bool = False, + + # 3d patch params + patch_3d: bool = False, + fake_3d: bool = False, + time_patch_size: Optional[int] = None, + + casual_3d: bool = False, + casual_3d_upsampler_index: Optional[list] = None, + + # motion module kwargs + motion_module_type = "VanillaGrid", + motion_module_kwargs = None, + motion_module_kwargs_odd = None, + motion_module_kwargs_even = None, + + # time position encoding + time_position_encoding_before_transformer = False, + + qk_norm = False, + after_norm = False, + ): + super().__init__() + self.use_linear_projection = use_linear_projection + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + self.enable_uvit = enable_uvit + inner_dim = num_attention_heads * attention_head_dim + self.basic_block_type = basic_block_type + self.patch_3d = patch_3d + self.fake_3d = fake_3d + self.casual_3d = casual_3d + self.casual_3d_upsampler_index = casual_3d_upsampler_index + + conv_cls = nn.Conv2d if USE_PEFT_BACKEND else LoRACompatibleConv + linear_cls = nn.Linear if USE_PEFT_BACKEND else LoRACompatibleLinear + + assert sample_size is not None, "Transformer3DModel over patched input must provide sample_size" + + self.height = sample_size + self.width = sample_size + + self.patch_size = patch_size + self.time_patch_size = self.patch_size if time_patch_size is None else time_patch_size + interpolation_scale = self.config.sample_size // 64 # => 64 (= 512 pixart) has interpolation scale 1 + interpolation_scale = max(interpolation_scale, 1) + + if self.casual_3d: + self.pos_embed = CasualPatchEmbed3D( + height=sample_size, + width=sample_size, + patch_size=patch_size, + time_patch_size=self.time_patch_size, + in_channels=in_channels, + embed_dim=inner_dim, + interpolation_scale=interpolation_scale, + ) + elif self.patch_3d: + if self.fake_3d: + self.pos_embed = PatchEmbedF3D( + height=sample_size, + width=sample_size, + patch_size=patch_size, + in_channels=in_channels, + embed_dim=inner_dim, + interpolation_scale=interpolation_scale, + ) + else: + self.pos_embed = PatchEmbed3D( + height=sample_size, + width=sample_size, + patch_size=patch_size, + time_patch_size=self.time_patch_size, + in_channels=in_channels, + embed_dim=inner_dim, + interpolation_scale=interpolation_scale, + ) + else: + self.pos_embed = PatchEmbed( + height=sample_size, + width=sample_size, + patch_size=patch_size, + in_channels=in_channels, + embed_dim=inner_dim, + interpolation_scale=interpolation_scale, + ) + + # 3. Define transformers blocks + if self.basic_block_type == "motionmodule": + self.transformer_blocks = nn.ModuleList( + [ + TemporalTransformerBlock( + inner_dim, + num_attention_heads, + attention_head_dim, + dropout=dropout, + cross_attention_dim=cross_attention_dim, + activation_fn=activation_fn, + num_embeds_ada_norm=num_embeds_ada_norm, + attention_bias=attention_bias, + only_cross_attention=only_cross_attention, + double_self_attention=double_self_attention, + upcast_attention=upcast_attention, + norm_type=norm_type, + norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, + attention_type=attention_type, + motion_module_type=motion_module_type, + motion_module_kwargs=motion_module_kwargs, + qk_norm=qk_norm, + after_norm=after_norm, + ) + for d in range(num_layers) + ] + ) + elif self.basic_block_type == "global_motionmodule": + self.transformer_blocks = nn.ModuleList( + [ + TemporalTransformerBlock( + inner_dim, + num_attention_heads, + attention_head_dim, + dropout=dropout, + cross_attention_dim=cross_attention_dim, + activation_fn=activation_fn, + num_embeds_ada_norm=num_embeds_ada_norm, + attention_bias=attention_bias, + only_cross_attention=only_cross_attention, + double_self_attention=double_self_attention, + upcast_attention=upcast_attention, + norm_type=norm_type, + norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, + attention_type=attention_type, + motion_module_type=motion_module_type, + motion_module_kwargs=motion_module_kwargs_even if d % 2 == 0 else motion_module_kwargs_odd, + qk_norm=qk_norm, + after_norm=after_norm, + ) + for d in range(num_layers) + ] + ) + elif self.basic_block_type == "kvcompression_motionmodule": + self.transformer_blocks = nn.ModuleList( + [ + TemporalTransformerBlock( + inner_dim, + num_attention_heads, + attention_head_dim, + dropout=dropout, + cross_attention_dim=cross_attention_dim, + activation_fn=activation_fn, + num_embeds_ada_norm=num_embeds_ada_norm, + attention_bias=attention_bias, + only_cross_attention=only_cross_attention, + double_self_attention=double_self_attention, + upcast_attention=upcast_attention, + norm_type=norm_type, + norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, + attention_type=attention_type, + kvcompression=False if d < 14 else True, + motion_module_type=motion_module_type, + motion_module_kwargs=motion_module_kwargs, + qk_norm=qk_norm, + after_norm=after_norm, + ) + for d in range(num_layers) + ] + ) + elif self.basic_block_type == "selfattentiontemporal": + self.transformer_blocks = nn.ModuleList( + [ + SelfAttentionTemporalTransformerBlock( + inner_dim, + num_attention_heads, + attention_head_dim, + dropout=dropout, + cross_attention_dim=cross_attention_dim, + activation_fn=activation_fn, + num_embeds_ada_norm=num_embeds_ada_norm, + attention_bias=attention_bias, + only_cross_attention=only_cross_attention, + double_self_attention=double_self_attention, + upcast_attention=upcast_attention, + norm_type=norm_type, + norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, + attention_type=attention_type, + qk_norm=qk_norm, + after_norm=after_norm, + ) + for d in range(num_layers) + ] + ) + else: + self.transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + inner_dim, + num_attention_heads, + attention_head_dim, + dropout=dropout, + cross_attention_dim=cross_attention_dim, + activation_fn=activation_fn, + num_embeds_ada_norm=num_embeds_ada_norm, + attention_bias=attention_bias, + only_cross_attention=only_cross_attention, + double_self_attention=double_self_attention, + upcast_attention=upcast_attention, + norm_type=norm_type, + norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, + attention_type=attention_type, + ) + for d in range(num_layers) + ] + ) + + if self.casual_3d: + self.unpatch1d = TemporalUpsampler3D() + elif self.patch_3d and self.fake_3d: + self.unpatch1d = UnPatch1D(inner_dim, True) + + if self.enable_uvit: + self.long_connect_fc = nn.ModuleList( + [ + nn.Linear(inner_dim, inner_dim, True) for d in range(13) + ] + ) + for index in range(13): + self.long_connect_fc[index] = zero_module(self.long_connect_fc[index]) + + # 4. Define output layers + self.out_channels = in_channels if out_channels is None else out_channels + if norm_type != "ada_norm_single": + self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out_1 = nn.Linear(inner_dim, 2 * inner_dim) + if self.patch_3d and not self.fake_3d: + self.proj_out_2 = nn.Linear(inner_dim, self.time_patch_size * patch_size * patch_size * self.out_channels) + else: + self.proj_out_2 = nn.Linear(inner_dim, patch_size * patch_size * self.out_channels) + elif norm_type == "ada_norm_single": + self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6) + self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5) + if self.patch_3d and not self.fake_3d: + self.proj_out = nn.Linear(inner_dim, self.time_patch_size * patch_size * patch_size * self.out_channels) + else: + self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * self.out_channels) + + # 5. PixArt-Alpha blocks. + self.adaln_single = None + self.use_additional_conditions = False + if norm_type == "ada_norm_single": + self.use_additional_conditions = self.config.sample_size == 128 + # TODO(Sayak, PVP) clean this, for now we use sample size to determine whether to use + # additional conditions until we find better name + self.adaln_single = AdaLayerNormSingle(inner_dim, use_additional_conditions=self.use_additional_conditions) + + self.caption_projection = None + self.clip_projection = None + if caption_channels is not None: + self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) + if in_channels == 12: + self.clip_projection = CLIPProjection(in_features=768, hidden_size=inner_dim * 8) + + self.gradient_checkpointing = False + + self.time_position_encoding_before_transformer = time_position_encoding_before_transformer + if self.time_position_encoding_before_transformer: + self.t_pos = TimePositionalEncoding(max_len = 4096, d_model = inner_dim) + + def _set_gradient_checkpointing(self, module, value=False): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = value + + def forward( + self, + hidden_states: torch.Tensor, + inpaint_latents: torch.Tensor = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + clip_encoder_hidden_states: Optional[torch.Tensor] = None, + timestep: Optional[torch.LongTensor] = None, + added_cond_kwargs: Dict[str, torch.Tensor] = None, + class_labels: Optional[torch.LongTensor] = None, + cross_attention_kwargs: Dict[str, Any] = None, + attention_mask: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + clip_attention_mask: Optional[torch.Tensor] = None, + return_dict: bool = True, + ): + """ + The [`Transformer2DModel`] forward method. + + Args: + hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous): + Input `hidden_states`. + encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*): + Conditional embeddings for cross attention layer. If not given, cross-attention defaults to + self-attention. + timestep ( `torch.LongTensor`, *optional*): + Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`. + class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*): + Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in + `AdaLayerZeroNorm`. + cross_attention_kwargs ( `Dict[str, Any]`, *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). + attention_mask ( `torch.Tensor`, *optional*): + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask + is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large + negative values to the attention scores corresponding to "discard" tokens. + encoder_attention_mask ( `torch.Tensor`, *optional*): + Cross-attention mask applied to `encoder_hidden_states`. Two formats supported: + + * Mask `(batch, sequence_length)` True = keep, False = discard. + * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard. + + If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format + above. This bias will be added to the cross-attention scores. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain + tuple. + + Returns: + If `return_dict` is True, an [`~models.transformer_2d.Transformer3DModelOutput`] is returned, otherwise a + `tuple` where the first element is the sample tensor. + """ + # ensure attention_mask is a bias, and give it a singleton query_tokens dimension. + # we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward. + # we can tell by counting dims; if ndim == 2: it's a mask rather than a bias. + # expects mask of shape: + # [batch, key_tokens] + # adds singleton query_tokens dimension: + # [batch, 1, key_tokens] + # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: + # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) + # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) + hidden_states = hidden_states.to(encoder_hidden_states.dtype) + + if attention_mask is not None and attention_mask.ndim == 2: + # assume that mask is expressed as: + # (1 = keep, 0 = discard) + # convert mask into a bias that can be added to attention scores: + # (keep = +0, discard = -10000.0) + attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0 + attention_mask = attention_mask.unsqueeze(1) + + if clip_attention_mask is not None: + encoder_attention_mask = torch.cat([encoder_attention_mask, clip_attention_mask], dim=1) + # convert encoder_attention_mask to a bias the same way we do for attention_mask + if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2: + encoder_attention_mask = (1 - encoder_attention_mask.to(encoder_hidden_states.dtype)) * -10000.0 + encoder_attention_mask = encoder_attention_mask.unsqueeze(1) + + if inpaint_latents is not None: + hidden_states = torch.concat([hidden_states, inpaint_latents], 1) + # 1. Input + if self.casual_3d: + video_length, height, width = (hidden_states.shape[-3] - 1) // self.time_patch_size + 1, hidden_states.shape[-2] // self.patch_size, hidden_states.shape[-1] // self.patch_size + elif self.patch_3d: + video_length, height, width = hidden_states.shape[-3] // self.time_patch_size, hidden_states.shape[-2] // self.patch_size, hidden_states.shape[-1] // self.patch_size + else: + video_length, height, width = hidden_states.shape[-3], hidden_states.shape[-2] // self.patch_size, hidden_states.shape[-1] // self.patch_size + hidden_states = rearrange(hidden_states, "b c f h w ->(b f) c h w") + + hidden_states = self.pos_embed(hidden_states) + if self.adaln_single is not None: + if self.use_additional_conditions and added_cond_kwargs is None: + raise ValueError( + "`added_cond_kwargs` cannot be None when using additional conditions for `adaln_single`." + ) + batch_size = hidden_states.shape[0] // video_length + timestep, embedded_timestep = self.adaln_single( + timestep, added_cond_kwargs, batch_size=batch_size, hidden_dtype=hidden_states.dtype + ) + hidden_states = rearrange(hidden_states, "(b f) (h w) c -> b c f h w", f=video_length, h=height, w=width) + + # hidden_states + # bs, c, f, h, w => b (f h w ) c + if self.time_position_encoding_before_transformer: + hidden_states = self.t_pos(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + # 2. Blocks + if self.caption_projection is not None: + batch_size = hidden_states.shape[0] + encoder_hidden_states = self.caption_projection(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1]) + + if clip_encoder_hidden_states is not None and encoder_hidden_states is not None: + batch_size = hidden_states.shape[0] + clip_encoder_hidden_states = self.clip_projection(clip_encoder_hidden_states) + clip_encoder_hidden_states = clip_encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1]) + + encoder_hidden_states = torch.cat([encoder_hidden_states, clip_encoder_hidden_states], dim = 1) + + skips = [] + skip_index = 0 + for index, block in enumerate(self.transformer_blocks): + if self.enable_uvit: + if index >= 15: + long_connect = self.long_connect_fc[skip_index](skips.pop()) + hidden_states = hidden_states + long_connect + skip_index += 1 + + if self.casual_3d_upsampler_index is not None and index in self.casual_3d_upsampler_index: + hidden_states = rearrange(hidden_states, "b (f h w) c -> b c f h w", f=video_length, h=height, w=width) + hidden_states = self.unpatch1d(hidden_states) + video_length = (video_length - 1) * 2 + 1 + hidden_states = rearrange(hidden_states, "b c f h w -> b (f h w) c", f=video_length, h=height, w=width) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + args = { + "basic": [], + "motionmodule": [video_length, height, width], + "global_motionmodule": [video_length, height, width], + "selfattentiontemporal": [], + "kvcompression_motionmodule": [video_length, height, width], + }[self.basic_block_type] + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + timestep, + cross_attention_kwargs, + class_labels, + *args, + **ckpt_kwargs, + ) + else: + kwargs = { + "basic": {}, + "motionmodule": {"num_frames":video_length, "height":height, "width":width}, + "global_motionmodule": {"num_frames":video_length, "height":height, "width":width}, + "selfattentiontemporal": {}, + "kvcompression_motionmodule": {"num_frames":video_length, "height":height, "width":width}, + }[self.basic_block_type] + hidden_states = block( + hidden_states, + attention_mask=attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + timestep=timestep, + cross_attention_kwargs=cross_attention_kwargs, + class_labels=class_labels, + **kwargs + ) + + if self.enable_uvit: + if index < 13: + skips.append(hidden_states) + + if self.fake_3d and self.patch_3d: + hidden_states = rearrange(hidden_states, "b (f h w) c -> (b h w) c f", f=video_length, w=width, h=height) + hidden_states = self.unpatch1d(hidden_states) + hidden_states = rearrange(hidden_states, "(b h w) c f -> b (f h w) c", w=width, h=height) + + # 3. Output + if self.config.norm_type != "ada_norm_single": + conditioning = self.transformer_blocks[0].norm1.emb( + timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + hidden_states = self.proj_out_2(hidden_states) + elif self.config.norm_type == "ada_norm_single": + shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None]).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states) + # Modulation + hidden_states = hidden_states * (1 + scale) + shift + hidden_states = self.proj_out(hidden_states) + hidden_states = hidden_states.squeeze(1) + + # unpatchify + if self.adaln_single is None: + height = width = int(hidden_states.shape[1] ** 0.5) + if self.patch_3d: + if self.fake_3d: + hidden_states = hidden_states.reshape( + shape=(-1, video_length * self.patch_size, height, width, self.patch_size, self.patch_size, self.out_channels) + ) + hidden_states = torch.einsum("nfhwpqc->ncfhpwq", hidden_states) + else: + hidden_states = hidden_states.reshape( + shape=(-1, video_length, height, width, self.time_patch_size, self.patch_size, self.patch_size, self.out_channels) + ) + hidden_states = torch.einsum("nfhwopqc->ncfohpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.out_channels, video_length * self.time_patch_size, height * self.patch_size, width * self.patch_size) + ) + else: + hidden_states = hidden_states.reshape( + shape=(-1, video_length, height, width, self.patch_size, self.patch_size, self.out_channels) + ) + hidden_states = torch.einsum("nfhwpqc->ncfhpwq", hidden_states) + output = hidden_states.reshape( + shape=(-1, self.out_channels, video_length, height * self.patch_size, width * self.patch_size) + ) + + if not return_dict: + return (output,) + + return Transformer3DModelOutput(sample=output) + + @classmethod + def from_pretrained_2d(cls, pretrained_model_path, subfolder=None, patch_size=2, transformer_additional_kwargs={}): + if subfolder is not None: + pretrained_model_path = os.path.join(pretrained_model_path, subfolder) + print(f"loaded 3D transformer's pretrained weights from {pretrained_model_path} ...") + + config_file = os.path.join(pretrained_model_path, 'config.json') + if not os.path.isfile(config_file): + raise RuntimeError(f"{config_file} does not exist") + with open(config_file, "r") as f: + config = json.load(f) + + from diffusers.utils import WEIGHTS_NAME + model = cls.from_config(config, **transformer_additional_kwargs) + model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME) + model_file_safetensors = model_file.replace(".bin", ".safetensors") + if os.path.exists(model_file_safetensors): + from safetensors.torch import load_file, safe_open + state_dict = load_file(model_file_safetensors) + else: + if not os.path.isfile(model_file): + raise RuntimeError(f"{model_file} does not exist") + state_dict = torch.load(model_file, map_location="cpu") + + if model.state_dict()['pos_embed.proj.weight'].size() != state_dict['pos_embed.proj.weight'].size(): + new_shape = model.state_dict()['pos_embed.proj.weight'].size() + if len(new_shape) == 5: + state_dict['pos_embed.proj.weight'] = state_dict['pos_embed.proj.weight'].unsqueeze(2).expand(new_shape).clone() + state_dict['pos_embed.proj.weight'][:, :, :-1] = 0 + else: + model.state_dict()['pos_embed.proj.weight'][:, :4, :, :] = state_dict['pos_embed.proj.weight'] + model.state_dict()['pos_embed.proj.weight'][:, 4:, :, :] = 0 + state_dict['pos_embed.proj.weight'] = model.state_dict()['pos_embed.proj.weight'] + + if model.state_dict()['proj_out.weight'].size() != state_dict['proj_out.weight'].size(): + new_shape = model.state_dict()['proj_out.weight'].size() + state_dict['proj_out.weight'] = torch.tile(state_dict['proj_out.weight'], [patch_size, 1]) + + if model.state_dict()['proj_out.bias'].size() != state_dict['proj_out.bias'].size(): + new_shape = model.state_dict()['proj_out.bias'].size() + state_dict['proj_out.bias'] = torch.tile(state_dict['proj_out.bias'], [patch_size]) + + tmp_state_dict = {} + for key in state_dict: + if key in model.state_dict().keys() and model.state_dict()[key].size() == state_dict[key].size(): + tmp_state_dict[key] = state_dict[key] + # else: + # print(key, "Size don't match, skip") + state_dict = tmp_state_dict + + m, u = model.load_state_dict(state_dict, strict=False) + print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};") + + params = [p.numel() if "attn_temporal." in n else 0 for n, p in model.named_parameters()] + print(f"### Attn temporal Parameters: {sum(params) / 1e6} M") + + return model + +class HunyuanTransformer3DModel(ModelMixin, ConfigMixin): + """ + HunYuanDiT: Diffusion model with a Transformer backbone. + + Inherit ModelMixin and ConfigMixin to be compatible with the sampler StableDiffusionPipeline of diffusers. + + Parameters: + num_attention_heads (`int`, *optional*, defaults to 16): + The number of heads to use for multi-head attention. + attention_head_dim (`int`, *optional*, defaults to 88): + The number of channels in each head. + in_channels (`int`, *optional*): + The number of channels in the input and output (specify if the input is **continuous**). + patch_size (`int`, *optional*): + The size of the patch to use for the input. + activation_fn (`str`, *optional*, defaults to `"geglu"`): + Activation function to use in feed-forward. + sample_size (`int`, *optional*): + The width of the latent images. This is fixed during training since it is used to learn a number of + position embeddings. + dropout (`float`, *optional*, defaults to 0.0): + The dropout probability to use. + cross_attention_dim (`int`, *optional*): + The number of dimension in the clip text embedding. + hidden_size (`int`, *optional*): + The size of hidden layer in the conditioning embedding layers. + num_layers (`int`, *optional*, defaults to 1): + The number of layers of Transformer blocks to use. + mlp_ratio (`float`, *optional*, defaults to 4.0): + The ratio of the hidden layer size to the input size. + learn_sigma (`bool`, *optional*, defaults to `True`): + Whether to predict variance. + cross_attention_dim_t5 (`int`, *optional*): + The number dimensions in t5 text embedding. + pooled_projection_dim (`int`, *optional*): + The size of the pooled projection. + text_len (`int`, *optional*): + The length of the clip text embedding. + text_len_t5 (`int`, *optional*): + The length of the T5 text embedding. + """ + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + num_attention_heads: int = 16, + attention_head_dim: int = 88, + in_channels: Optional[int] = None, + out_channels: Optional[int] = None, + patch_size: Optional[int] = None, + n_query=16, + projection_dim=768, + activation_fn: str = "gelu-approximate", + sample_size=32, + hidden_size=1152, + num_layers: int = 28, + mlp_ratio: float = 4.0, + learn_sigma: bool = True, + cross_attention_dim: int = 1024, + norm_type: str = "layer_norm", + cross_attention_dim_t5: int = 2048, + pooled_projection_dim: int = 1024, + text_len: int = 77, + text_len_t5: int = 256, + + # block type + basic_block_type: str = "basic", + # motion module kwargs + motion_module_type = "VanillaGrid", + motion_module_kwargs = None, + motion_module_kwargs_odd = None, + motion_module_kwargs_even = None, + + time_position_encoding = False, + after_norm = False, + + ): + super().__init__() + # 4. Define output layers + if learn_sigma: + self.out_channels = in_channels * 2 if out_channels is None else out_channels + else: + self.out_channels = in_channels if out_channels is None else out_channels + self.enable_inpaint = in_channels * 2 != self.out_channels if learn_sigma else in_channels != self.out_channels + self.num_heads = num_attention_heads + self.inner_dim = num_attention_heads * attention_head_dim + self.basic_block_type = basic_block_type + self.text_embedder = PixArtAlphaTextProjection( + in_features=cross_attention_dim_t5, + hidden_size=cross_attention_dim_t5 * 4, + out_features=cross_attention_dim, + act_fn="silu_fp32", + ) + + self.text_embedding_padding = nn.Parameter( + torch.randn(text_len + text_len_t5, cross_attention_dim, dtype=torch.float32) + ) + + self.pos_embed = PatchEmbed( + height=sample_size, + width=sample_size, + in_channels=in_channels, + embed_dim=hidden_size, + patch_size=patch_size, + pos_embed_type=None, + ) + + self.time_extra_emb = HunyuanCombinedTimestepTextSizeStyleEmbedding( + hidden_size, + pooled_projection_dim=pooled_projection_dim, + seq_len=text_len_t5, + cross_attention_dim=cross_attention_dim_t5, + ) + + # 3. Define transformers blocks + if self.basic_block_type == "motionmodule": + self.blocks = nn.ModuleList( + [ + HunyuanTemporalTransformerBlock( + dim=self.inner_dim, + num_attention_heads=self.config.num_attention_heads, + activation_fn=activation_fn, + ff_inner_dim=int(self.inner_dim * mlp_ratio), + cross_attention_dim=cross_attention_dim, + qk_norm=True, # See http://arxiv.org/abs/2302.05442 for details. + skip=layer > num_layers // 2, + motion_module_type=motion_module_type, + motion_module_kwargs=motion_module_kwargs, + after_norm=after_norm, + ) + for layer in range(num_layers) + ] + ) + elif self.basic_block_type == "global_motionmodule": + self.blocks = nn.ModuleList( + [ + HunyuanTemporalTransformerBlock( + dim=self.inner_dim, + num_attention_heads=self.config.num_attention_heads, + activation_fn=activation_fn, + ff_inner_dim=int(self.inner_dim * mlp_ratio), + cross_attention_dim=cross_attention_dim, + qk_norm=True, # See http://arxiv.org/abs/2302.05442 for details. + skip=layer > num_layers // 2, + motion_module_type=motion_module_type, + motion_module_kwargs=motion_module_kwargs_even if layer % 2 == 0 else motion_module_kwargs_odd, + after_norm=after_norm, + ) + for layer in range(num_layers) + ] + ) + elif self.basic_block_type == "hybrid_attention": + self.blocks = nn.ModuleList( + [ + HunyuanDiTBlock( + dim=self.inner_dim, + num_attention_heads=self.config.num_attention_heads, + activation_fn=activation_fn, + ff_inner_dim=int(self.inner_dim * mlp_ratio), + cross_attention_dim=cross_attention_dim, + qk_norm=True, # See http://arxiv.org/abs/2302.05442 for details. + skip=layer > num_layers // 2, + after_norm=after_norm, + time_position_encoding=time_position_encoding, + is_local_attention=False if layer % 2 == 0 else True, + local_attention_frames=2, + enable_inpaint=self.enable_inpaint, + ) + for layer in range(num_layers) + ] + ) + else: + # HunyuanDiT Blocks + self.blocks = nn.ModuleList( + [ + HunyuanDiTBlock( + dim=self.inner_dim, + num_attention_heads=self.config.num_attention_heads, + activation_fn=activation_fn, + ff_inner_dim=int(self.inner_dim * mlp_ratio), + cross_attention_dim=cross_attention_dim, + qk_norm=True, # See http://arxiv.org/abs/2302.05442 for details. + skip=layer > num_layers // 2, + after_norm=after_norm, + time_position_encoding=time_position_encoding, + enable_inpaint=self.enable_inpaint, + ) + for layer in range(num_layers) + ] + ) + + self.n_query = n_query + if self.enable_inpaint: + self.clip_padding = nn.Parameter( + torch.randn((self.n_query, cross_attention_dim)) * 0.02 + ) + self.clip_projection = Resampler( + int(math.sqrt(n_query)), + embed_dim=cross_attention_dim, + num_heads=self.config.num_attention_heads, + kv_dim=projection_dim, + norm_layer=nn.LayerNorm, + ) + else: + self.clip_padding = None + self.clip_projection = None + + self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True) + + self.gradient_checkpointing = False + self.hidden_cache_size = 0 + + def _set_gradient_checkpointing(self, module, value=False): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = value + + def forward( + self, + hidden_states, + timestep, + encoder_hidden_states=None, + text_embedding_mask=None, + encoder_hidden_states_t5=None, + text_embedding_mask_t5=None, + image_meta_size=None, + style=None, + image_rotary_emb=None, + inpaint_latents=None, + clip_encoder_hidden_states: Optional[torch.Tensor]=None, + clip_attention_mask: Optional[torch.Tensor]=None, + return_dict=True, + ): + """ + The [`HunyuanDiT2DModel`] forward method. + + Args: + hidden_states (`torch.Tensor` of shape `(batch size, dim, height, width)`): + The input tensor. + timestep ( `torch.LongTensor`, *optional*): + Used to indicate denoising step. + encoder_hidden_states ( `torch.Tensor` of shape `(batch size, sequence len, embed dims)`, *optional*): + Conditional embeddings for cross attention layer. This is the output of `BertModel`. + text_embedding_mask: torch.Tensor + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. This is the output + of `BertModel`. + encoder_hidden_states_t5 ( `torch.Tensor` of shape `(batch size, sequence len, embed dims)`, *optional*): + Conditional embeddings for cross attention layer. This is the output of T5 Text Encoder. + text_embedding_mask_t5: torch.Tensor + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. This is the output + of T5 Text Encoder. + image_meta_size (torch.Tensor): + Conditional embedding indicate the image sizes + style: torch.Tensor: + Conditional embedding indicate the style + image_rotary_emb (`torch.Tensor`): + The image rotary embeddings to apply on query and key tensors during attention calculation. + return_dict: bool + Whether to return a dictionary. + """ + if inpaint_latents is not None: + hidden_states = torch.concat([hidden_states, inpaint_latents], 1) + + # unpatchify: (N, out_channels, H, W) + patch_size = self.pos_embed.patch_size + video_length, height, width = hidden_states.shape[-3], hidden_states.shape[-2] // patch_size, hidden_states.shape[-1] // patch_size + hidden_states = rearrange(hidden_states, "b c f h w ->(b f) c h w") + hidden_states = self.pos_embed(hidden_states) + hidden_states = rearrange(hidden_states, "(b f) (h w) c -> b c f h w", f=video_length, h=height, w=width) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + temb = self.time_extra_emb( + timestep, encoder_hidden_states_t5, image_meta_size, style, hidden_dtype=timestep.dtype + ) # [B, D] + + # text projection + batch_size, sequence_length, _ = encoder_hidden_states_t5.shape + encoder_hidden_states_t5 = self.text_embedder( + encoder_hidden_states_t5.view(-1, encoder_hidden_states_t5.shape[-1]) + ) + encoder_hidden_states_t5 = encoder_hidden_states_t5.view(batch_size, sequence_length, -1) + + encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states_t5], dim=1) + text_embedding_mask = torch.cat([text_embedding_mask, text_embedding_mask_t5], dim=-1) + text_embedding_mask = text_embedding_mask.unsqueeze(2).bool() + + encoder_hidden_states = torch.where(text_embedding_mask, encoder_hidden_states, self.text_embedding_padding) + + if clip_encoder_hidden_states is not None: + batch_size = encoder_hidden_states.shape[0] + + clip_encoder_hidden_states = self.clip_projection(clip_encoder_hidden_states) + clip_encoder_hidden_states = clip_encoder_hidden_states.view(batch_size, -1, encoder_hidden_states.shape[-1]) + + clip_attention_mask = clip_attention_mask.unsqueeze(2).bool() + clip_encoder_hidden_states = torch.where(clip_attention_mask, clip_encoder_hidden_states, self.clip_padding) + + skips = [] + skips_cpu_cache = [] + skips_cache_size = self.hidden_cache_size # Equal or bigger than this value, skips will be moved to cpu + if skips_cache_size <= 0: + skips_cache_size = self.config.num_layers + + for layer, block in enumerate(self.blocks): + if layer > self.config.num_layers // 2: + if skips_cache_size == 1: + skip_cpu = skips.pop() + skip = skip_cpu.to(hidden_states.device) + else: + if len(skips) == 0: + skips_cache_block = skips_cpu_cache.pop() + for si in range(len(skips_cache_block)): + skips.append(skips_cache_block[si].to(hidden_states.device)) + skip = skips.pop() + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + args = { + "basic": [video_length, height, width, clip_encoder_hidden_states], + "hybrid_attention": [video_length, height, width, clip_encoder_hidden_states], + "motionmodule": [video_length, height, width], + "global_motionmodule": [video_length, height, width], + }[self.basic_block_type] + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb, + skip, + *args, + **ckpt_kwargs, + ) + else: + kwargs = { + "basic": {"num_frames":video_length, "height":height, "width":width, "clip_encoder_hidden_states":clip_encoder_hidden_states}, + "hybrid_attention": {"num_frames":video_length, "height":height, "width":width, "clip_encoder_hidden_states":clip_encoder_hidden_states}, + "motionmodule": {"num_frames":video_length, "height":height, "width":width}, + "global_motionmodule": {"num_frames":video_length, "height":height, "width":width}, + }[self.basic_block_type] + hidden_states = block( + hidden_states, + temb=temb, + encoder_hidden_states=encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + skip=skip, + **kwargs + ) # (N, L, D) + else: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + args = { + "basic": [None, video_length, height, width, clip_encoder_hidden_states], + "hybrid_attention": [None, video_length, height, width, clip_encoder_hidden_states], + "motionmodule": [None, video_length, height, width], + "global_motionmodule": [None, video_length, height, width], + }[self.basic_block_type] + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb, + *args, + **ckpt_kwargs, + ) + else: + kwargs = { + "basic": {"num_frames":video_length, "height":height, "width":width, "clip_encoder_hidden_states":clip_encoder_hidden_states}, + "hybrid_attention": {"num_frames":video_length, "height":height, "width":width, "clip_encoder_hidden_states":clip_encoder_hidden_states}, + "motionmodule": {"num_frames":video_length, "height":height, "width":width}, + "global_motionmodule": {"num_frames":video_length, "height":height, "width":width}, + }[self.basic_block_type] + hidden_states = block( + hidden_states, + temb=temb, + encoder_hidden_states=encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + **kwargs + ) # (N, L, D) + + if layer < (self.config.num_layers // 2 - 1): + if skips_cache_size == 1: + skips.append(hidden_states.to("cpu")) + else: + skips.append(hidden_states) + if len(skips) >= skips_cache_size: + skips_cache_block = [] + for si in range(len(skips)): + skips_cache_block.append(skips[si].to("cpu")) + + skips = [] + skips_cpu_cache.append(skips_cache_block) + + # final layer + hidden_states = self.norm_out(hidden_states, temb.to(torch.float32)) + hidden_states = self.proj_out(hidden_states) + # (N, L, patch_size ** 2 * out_channels) + + hidden_states = hidden_states.reshape( + shape=(hidden_states.shape[0], video_length, height, width, patch_size, patch_size, self.out_channels) + ) + hidden_states = torch.einsum("nfhwpqc->ncfhpwq", hidden_states) + output = hidden_states.reshape( + shape=(hidden_states.shape[0], self.out_channels, video_length, height * patch_size, width * patch_size) + ) + + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) + + @classmethod + def from_pretrained_2d(cls, pretrained_model_path, subfolder=None, patch_size=2, transformer_additional_kwargs={}): + if subfolder is not None: + pretrained_model_path = os.path.join(pretrained_model_path, subfolder) + # print(f"loaded 3D transformer's pretrained weights from {pretrained_model_path} ...") + + config_file = os.path.join(pretrained_model_path, 'config.json') + if not os.path.isfile(config_file): + raise RuntimeError(f"{config_file} does not exist") + with open(config_file, "r") as f: + config = json.load(f) + + from diffusers.utils import WEIGHTS_NAME + model = cls.from_config(config, **transformer_additional_kwargs) + model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME) + model_file_safetensors = model_file.replace(".bin", ".safetensors") + if os.path.exists(model_file_safetensors): + from safetensors.torch import load_file, safe_open + state_dict = load_file(model_file_safetensors) + else: + if not os.path.isfile(model_file): + raise RuntimeError(f"{model_file} does not exist") + state_dict = torch.load(model_file, map_location="cpu") + + if model.state_dict()['pos_embed.proj.weight'].size() != state_dict['pos_embed.proj.weight'].size(): + new_shape = model.state_dict()['pos_embed.proj.weight'].size() + if len(new_shape) == 5: + state_dict['pos_embed.proj.weight'] = state_dict['pos_embed.proj.weight'].unsqueeze(2).expand(new_shape).clone() + state_dict['pos_embed.proj.weight'][:, :, :-1] = 0 + + if model.state_dict()['proj_out.bias'].size() != state_dict['proj_out.bias'].size(): + if model.state_dict()['proj_out.bias'].size()[0] > state_dict['proj_out.bias'].size()[0]: + model.state_dict()['proj_out.bias'][:state_dict['proj_out.bias'].size()[0]] = state_dict['proj_out.bias'] + state_dict['proj_out.bias'] = model.state_dict()['proj_out.bias'] + + tmp_state_dict = {} + for key in state_dict: + if key in model.state_dict().keys() and model.state_dict()[key].size() == state_dict[key].size(): + tmp_state_dict[key] = state_dict[key] + state_dict = tmp_state_dict + + m, u = model.load_state_dict(state_dict, strict=False) + # print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};") + # print(m) + + params = [p.numel() if "mamba" in n else 0 for n, p in model.named_parameters()] + # print(f"### Mamba Parameters: {sum(params) / 1e6} M") + + params = [p.numel() if "attn1." in n else 0 for n, p in model.named_parameters()] + # print(f"### attn1 Parameters: {sum(params) / 1e6} M") + + params = [p.numel() for n, p in model.named_parameters()] + # print(f"### Total Parameters: {sum(params) / 1e6} M") + + return model \ No newline at end of file diff --git a/ruyi/pipeline/pipeline_ruyi_inpaint.py b/ruyi/pipeline/pipeline_ruyi_inpaint.py new file mode 100644 index 0000000..f725c2c --- /dev/null +++ b/ruyi/pipeline/pipeline_ruyi_inpaint.py @@ -0,0 +1,1072 @@ +# Copyright 2024 HunyuanDiT Authors 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 copy +import gc +import html +import inspect +import re +import urllib.parse as ul +from dataclasses import dataclass +from typing import Callable, Dict, List, Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn.functional as F +from diffusers import DiffusionPipeline, ImagePipelineOutput +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback +from diffusers.image_processor import VaeImageProcessor +from diffusers.models import AutoencoderKL, HunyuanDiT2DModel +from diffusers.models.embeddings import get_2d_rotary_pos_embed +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput +from diffusers.pipelines.stable_diffusion.safety_checker import \ + StableDiffusionSafetyChecker +from diffusers.schedulers import DDPMScheduler, DPMSolverMultistepScheduler +from diffusers.utils import (BACKENDS_MAPPING, BaseOutput, deprecate, + is_bs4_available, is_ftfy_available, + is_torch_xla_available, logging, + replace_example_docstring) +from diffusers.utils.torch_utils import randn_tensor +from einops import rearrange +from PIL import Image +from tqdm import tqdm +from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import HunyuanDiTPipeline + + >>> pipe = HunyuanDiTPipeline.from_pretrained("Tencent-Hunyuan/HunyuanDiT", torch_dtype=torch.float16) + >>> pipe.to("cuda") + + >>> # You may also use English prompt as HunyuanDiT supports both English and Chinese + >>> # prompt = "An astronaut riding a horse" + >>> prompt = "一个宇航员在骑马" + >>> image = pipe(prompt).images[0] + ``` +""" + +STANDARD_RATIO = np.array( + [ + 1.0, # 1:1 + 4.0 / 3.0, # 4:3 + 3.0 / 4.0, # 3:4 + 16.0 / 9.0, # 16:9 + 9.0 / 16.0, # 9:16 + ] +) +STANDARD_SHAPE = [ + [(1024, 1024), (1280, 1280)], # 1:1 + [(1024, 768), (1152, 864), (1280, 960)], # 4:3 + [(768, 1024), (864, 1152), (960, 1280)], # 3:4 + [(1280, 720)], # 16:9 + [(720, 1280)], # 9:16 +] +STANDARD_AREA = [np.array([w * h for w, h in shapes]) for shapes in STANDARD_SHAPE] +SUPPORTED_SHAPE = [ + (1024, 1024), + (1280, 1280), # 1:1 + (1024, 768), + (1152, 864), + (1280, 960), # 4:3 + (768, 1024), + (864, 1152), + (960, 1280), # 3:4 + (1280, 720), # 16:9 + (720, 1280), # 9:16 +] + + +def map_to_standard_shapes(target_width, target_height): + target_ratio = target_width / target_height + closest_ratio_idx = np.argmin(np.abs(STANDARD_RATIO - target_ratio)) + closest_area_idx = np.argmin(np.abs(STANDARD_AREA[closest_ratio_idx] - target_width * target_height)) + width, height = STANDARD_SHAPE[closest_ratio_idx][closest_area_idx] + return width, height + + +def get_resize_crop_region_for_grid(src, tgt_size): + th = tw = tgt_size + h, w = src + + r = h / w + + # resize + if r > 1: + 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.rescale_noise_cfg +def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): + """ + Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4 + """ + std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) + std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) + # rescale the results from guidance (fixes overexposure) + noise_pred_rescaled = noise_cfg * (std_text / std_cfg) + # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images + noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg + return noise_cfg + + +@dataclass +class RuyiPipelineOutput(BaseOutput): + videos: Union[torch.Tensor, np.ndarray] + + +class RuyiInpaintPipeline(DiffusionPipeline): + r""" + 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 ([`AutoencoderKLMagvit`]): + Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations. + transformer ([`HunyuanTransformer3DModel`]): + The HunyuanDiT model designed by Tencent Hunyuan. + scheduler ([`DDPMScheduler`]): + A scheduler to be used in combination with HunyuanDiT to denoise the encoded image latents. + """ + + model_cpu_offload_seq = "clip_image_encoder->transformer->vae" + _optional_components = [ + "safety_checker", + "feature_extractor", + "clip_image_encoder", + ] + _exclude_from_cpu_offload = ["safety_checker"] + _callback_tensor_inputs = [ + "latents", + "prompt_embeds", + "negative_prompt_embeds", + "prompt_embeds_2", + "negative_prompt_embeds_2", + ] + + def __init__( + self, + vae: AutoencoderKL, + transformer: HunyuanDiT2DModel, + scheduler: DDPMScheduler, + safety_checker: StableDiffusionSafetyChecker, + feature_extractor: CLIPImageProcessor, + requires_safety_checker: bool = True, + clip_image_processor:CLIPImageProcessor = None, + clip_image_encoder:CLIPVisionModelWithProjection = None, + ): + super().__init__() + + self.register_modules( + vae=vae, + transformer=transformer, + scheduler=scheduler, + safety_checker=safety_checker, + feature_extractor=feature_extractor, + clip_image_processor=clip_image_processor, + clip_image_encoder=clip_image_encoder, + ) + + if safety_checker is None and requires_safety_checker: + logger.warning( + f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure" + " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered" + " results in services or applications open to the public. Both the diffusers team and Hugging Face" + " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling" + " it only for use-cases that involve analyzing network behavior or auditing its results. For more" + " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ." + ) + + if safety_checker is not None and feature_extractor is None: + raise ValueError( + "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety" + " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead." + ) + + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self.register_to_config(requires_safety_checker=requires_safety_checker) + self.default_sample_size = self.transformer.config.sample_size + self.mask_processor = VaeImageProcessor( + vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True + ) + self.model_cpu_offload_flag = False + + def enable_sequential_cpu_offload(self, *args, **kwargs): + super().enable_sequential_cpu_offload(*args, **kwargs) + self.model_cpu_offload_flag = False + if hasattr(self.transformer, "clip_projection") and self.transformer.clip_projection is not None: + import accelerate + accelerate.hooks.remove_hook_from_module(self.transformer.clip_projection, recurse=True) + self.transformer.clip_projection = self.transformer.clip_projection.to("cuda") + + def enable_model_cpu_offload(self, *args, **kwargs): + super().enable_model_cpu_offload(*args, **kwargs) + self.model_cpu_offload_flag = True + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker + def run_safety_checker(self, image, device, dtype): + if self.safety_checker is None: + has_nsfw_concept = None + else: + if torch.is_tensor(image): + feature_extractor_input = self.image_processor.postprocess(image, output_type="pil") + else: + feature_extractor_input = self.image_processor.numpy_to_pil(image) + safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device) + image, has_nsfw_concept = self.safety_checker( + images=image, clip_input=safety_checker_input.pixel_values.to(dtype) + ) + return image, has_nsfw_concept + + # 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 + + def check_inputs( + self, + prompt, + height, + width, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + prompt_embeds_2=None, + negative_prompt_embeds_2=None, + prompt_attention_mask_2=None, + negative_prompt_attention_mask_2=None, + callback_on_step_end_tensor_inputs=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 None and prompt_embeds_2 is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds_2`. Cannot leave both `prompt` and `prompt_embeds_2` 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_embeds is not None and prompt_attention_mask is None: + raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") + + if prompt_embeds_2 is not None and prompt_attention_mask_2 is None: + raise ValueError("Must provide `prompt_attention_mask_2` when specifying `prompt_embeds_2`.") + + 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 negative_prompt_embeds is not None and negative_prompt_attention_mask is None: + raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") + + if negative_prompt_embeds_2 is not None and negative_prompt_attention_mask_2 is None: + raise ValueError( + "Must provide `negative_prompt_attention_mask_2` when specifying `negative_prompt_embeds_2`." + ) + 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}." + ) + if prompt_embeds_2 is not None and negative_prompt_embeds_2 is not None: + if prompt_embeds_2.shape != negative_prompt_embeds_2.shape: + raise ValueError( + "`prompt_embeds_2` and `negative_prompt_embeds_2` must have the same shape when passed directly, but" + f" got: `prompt_embeds_2` {prompt_embeds_2.shape} != `negative_prompt_embeds_2`" + f" {negative_prompt_embeds_2.shape}." + ) + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps + def get_timesteps(self, num_inference_steps, strength, device): + # get the original timestep using init_timestep + init_timestep = min(int(num_inference_steps * strength), num_inference_steps) + + t_start = max(num_inference_steps - init_timestep, 0) + timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :] + + return timesteps, num_inference_steps - t_start + + def prepare_mask_latents( + self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance + ): + # resize the mask to latents shape as we concatenate the mask to the latents + # we do that before converting to dtype to avoid breaking in case we're using cpu_offload + # and half precision + video_length = mask.shape[2] + + mask = mask.to(device=device, dtype=self.vae.dtype) + if self.vae.quant_conv.weight.ndim==5: + bs = 1 + mini_batch_encoder = self.vae.mini_batch_encoder + new_mask = [] + for i in range(0, mask.shape[0], bs): + mask_bs = mask[i : i + bs] + mask_bs = self.vae.encode(mask_bs)[0] + mask_bs = mask_bs.sample() + new_mask.append(mask_bs) + mask = torch.cat(new_mask, dim = 0) + mask = mask * self.vae.config.scaling_factor + + else: + if mask.shape[1] == 4: + mask = mask + else: + video_length = mask.shape[2] + mask = rearrange(mask, "b c f h w -> (b f) c h w") + mask = self._encode_vae_image(mask, generator=generator) + mask = rearrange(mask, "(b f) c h w -> b c f h w", f=video_length) + + masked_image = masked_image.to(device=device, dtype=self.vae.dtype) + if self.vae.quant_conv.weight.ndim==5: + bs = 1 + new_mask_pixel_values = [] + for i in range(0, masked_image.shape[0], bs): + mask_pixel_values_bs = masked_image[i : i + bs] + mask_pixel_values_bs = self.vae.encode(mask_pixel_values_bs)[0] + mask_pixel_values_bs = mask_pixel_values_bs.sample() + new_mask_pixel_values.append(mask_pixel_values_bs) + masked_image_latents = torch.cat(new_mask_pixel_values, dim = 0) + masked_image_latents = masked_image_latents * self.vae.config.scaling_factor + + else: + if masked_image.shape[1] == 4: + masked_image_latents = masked_image + else: + video_length = mask.shape[2] + masked_image = rearrange(masked_image, "b c f h w -> (b f) c h w") + masked_image_latents = self._encode_vae_image(masked_image, generator=generator) + masked_image_latents = rearrange(masked_image_latents, "(b f) c h w -> b c f h w", f=video_length) + + # aligning device to prevent device errors when concating it with the latent model input + masked_image_latents = masked_image_latents.to(device=device, dtype=dtype) + + return mask, masked_image_latents + + def prepare_latents( + self, + batch_size, + num_channels_latents, + height, + width, + video_length, + dtype, + device, + generator, + latents=None, + video=None, + timestep=None, + is_strength_max=True, + return_noise=False, + return_video_latents=False, + ): + video_latents = None + + if self.vae.quant_conv.weight.ndim==5: + mini_batch_encoder = self.vae.mini_batch_encoder + mini_batch_decoder = self.vae.mini_batch_decoder + shape = (batch_size, num_channels_latents, int(video_length // mini_batch_encoder * mini_batch_decoder) if video_length != 1 else 1, height // self.vae_scale_factor, width // self.vae_scale_factor) + else: + shape = (batch_size, num_channels_latents, video_length, height // self.vae_scale_factor, width // self.vae_scale_factor) + + 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." + ) + + if return_video_latents or (latents is None and not is_strength_max): + video = video.to(device=device, dtype=self.vae.dtype) + if self.vae.quant_conv.weight.ndim==5: + bs = 1 + new_video = [] + for i in range(0, video.shape[0], bs): + video_bs = video[i : i + bs] + video_bs = self.vae.encode(video_bs)[0] + video_bs = video_bs.sample() + new_video.append(video_bs) + video = torch.cat(new_video, dim = 0) + video = video * self.vae.config.scaling_factor + + else: + if video.shape[1] == 4: + video = video + else: + video_length = video.shape[2] + video = rearrange(video, "b c f h w -> (b f) c h w") + video = self._encode_vae_image(video, generator=generator) + video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length) + video_latents = video.repeat(batch_size // video.shape[0], 1, 1, 1, 1) + video_latents = video_latents.to(device=device, dtype=dtype) + + if latents is None: + # TODO: a fast but brute force fix, sometimes the computed shape is not equals to the video latent's shape + if video_latents is not None: + shape = video_latents.shape + noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + # if strength is 1. then initialise the latents to noise, else initial to image + noise + latents = noise if is_strength_max else self.scheduler.add_noise(video_latents, noise, timestep) + # if pure noise then scale the initial latents by the Scheduler's init sigma + latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents + else: + noise = latents.to(device) + latents = noise * self.scheduler.init_noise_sigma + + # scale the initial noise by the standard deviation required by the scheduler + outputs = (latents,) + + if return_noise: + outputs += (noise,) + + if return_video_latents: + outputs += (video_latents,) + + return outputs + + def smooth_output(self, video, mini_batch_encoder, mini_batch_decoder): + if video.size()[2] <= mini_batch_encoder: + return video + prefix_index_before = mini_batch_encoder // 2 + prefix_index_after = mini_batch_encoder - prefix_index_before + pixel_values = video[:, :, prefix_index_before:-prefix_index_after] + + # Encode middle videos + latents = self.vae.encode(pixel_values)[0] + latents = latents.mode() + middle_video = self.vae.decode(latents)[0] + + video[:, :, prefix_index_before:-prefix_index_after] = (video[:, :, prefix_index_before:-prefix_index_after] + middle_video) / 2 + return video + + + def decode_latents(self, latents): + video_length = latents.shape[2] + latents = 1 / self.vae.config.scaling_factor * latents + if self.vae.quant_conv.weight.ndim==5: + mini_batch_encoder = self.vae.mini_batch_encoder + mini_batch_decoder = self.vae.mini_batch_decoder + video = self.vae.decode(latents)[0] + video = video.clamp(-1, 1) + if not self.vae.cache_compression_vae: + video = self.smooth_output(video, mini_batch_encoder, mini_batch_decoder).cpu().clamp(-1, 1) + else: + latents = rearrange(latents, "b c f h w -> (b f) c h w") + video = [] + for frame_idx in tqdm(range(latents.shape[0])): + video.append(self.vae.decode(latents[frame_idx:frame_idx+1]).sample) + video = torch.cat(video) + video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length) + video = (video / 2 + 0.5).clamp(0, 1) + # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16 + video = video.cpu().float().numpy() + return video + + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def guidance_rescale(self): + return self._guidance_rescale + + # 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. + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + video_length: Optional[int] = None, + video: Union[torch.FloatTensor] = None, + mask_video: Union[torch.FloatTensor] = None, + masked_video_latents: Union[torch.FloatTensor] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = 50, + guidance_scale: Optional[float] = 5.0, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: Optional[float] = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds_2: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds_2: Optional[torch.Tensor] = None, + prompt_attention_mask: Optional[torch.Tensor] = None, + prompt_attention_mask_2: Optional[torch.Tensor] = None, + negative_prompt_attention_mask: Optional[torch.Tensor] = None, + negative_prompt_attention_mask_2: Optional[torch.Tensor] = None, + output_type: Optional[str] = "latent", + return_dict: bool = True, + callback_on_step_end: Optional[ + Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] + ] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + guidance_rescale: float = 0.0, + original_size: Optional[Tuple[int, int]] = (1024, 1024), + target_size: Optional[Tuple[int, int]] = None, + crops_coords_top_left: Tuple[int, int] = (0, 0), + use_resolution_binning: bool = False, + clip_image: Image = None, + clip_apply_ratio: float = 0.40, + strength: float = 1.0, + comfyui_progressbar: bool = False, + ): + r""" + The call function to the pipeline for generation with HunyuanDiT. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + 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. This parameter is modulated by `strength`. + guidance_scale (`float`, *optional*, defaults to 7.5): + A higher guidance scale value encourages the model to generate images closely linked to the text + `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide what to not include in image generation. If not defined, you need to + pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`). + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies + to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not + provided, text embeddings are generated from the `prompt` input argument. + prompt_embeds_2 (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not + provided, text embeddings are generated from the `prompt` input argument. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If + not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. + negative_prompt_embeds_2 (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If + not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument. + prompt_attention_mask (`torch.Tensor`, *optional*): + Attention mask for the prompt. Required when `prompt_embeds` is passed directly. + prompt_attention_mask_2 (`torch.Tensor`, *optional*): + Attention mask for the prompt. Required when `prompt_embeds_2` is passed directly. + negative_prompt_attention_mask (`torch.Tensor`, *optional*): + Attention mask for the negative prompt. Required when `negative_prompt_embeds` is passed directly. + negative_prompt_attention_mask_2 (`torch.Tensor`, *optional*): + Attention mask for the negative prompt. Required when `negative_prompt_embeds_2` is passed directly. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated image. Choose between `PIL.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a + plain tuple. + callback_on_step_end (`Callable[[int, int, Dict], None]`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + A callback function or a list of callback functions to be called at the end of each denoising step. + callback_on_step_end_tensor_inputs (`List[str]`, *optional*): + A list of tensor inputs that should be passed to the callback function. If not defined, all tensor + inputs will be passed. + guidance_rescale (`float`, *optional*, defaults to 0.0): + Rescale the noise_cfg according to `guidance_rescale`. Based on findings of [Common Diffusion Noise + Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4 + original_size (`Tuple[int, int]`, *optional*, defaults to `(1024, 1024)`): + The original size of the image. Used to calculate the time ids. + target_size (`Tuple[int, int]`, *optional*): + The target size of the image. Used to calculate the time ids. + crops_coords_top_left (`Tuple[int, int]`, *optional*, defaults to `(0, 0)`): + The top left coordinates of the crop. Used to calculate the time ids. + use_resolution_binning (`bool`, *optional*, defaults to `True`): + Whether to use resolution binning or not. If `True`, the input resolution will be mapped to the closest + standard resolution. Supported resolutions are 1024x1024, 1280x1280, 1024x768, 1152x864, 1280x960, + 768x1024, 864x1152, 960x1280, 1280x768, and 768x1280. It is recommended to set this to `True`. + + Examples: + + Returns: + [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned, + otherwise a `tuple` is returned where the first element is a list with the generated images and the + second element is a list of `bool`s indicating whether the corresponding generated image contains + "not-safe-for-work" (nsfw) content. + """ + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 1. default height and width + height = height or self.default_sample_size * self.vae_scale_factor + width = width or self.default_sample_size * self.vae_scale_factor + height = int(height // 16 * 16) + width = int(width // 16 * 16) + + if use_resolution_binning and (height, width) not in SUPPORTED_SHAPE: + width, height = map_to_standard_shapes(width, height) + height = int(height) + width = int(width) + logger.warning(f"Reshaped to (height, width)=({height}, {width}), Supported shapes are {SUPPORTED_SHAPE}") + + # 2. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + height, + width, + negative_prompt, + prompt_embeds, + negative_prompt_embeds, + prompt_attention_mask, + negative_prompt_attention_mask, + prompt_embeds_2, + negative_prompt_embeds_2, + prompt_attention_mask_2, + negative_prompt_attention_mask_2, + callback_on_step_end_tensor_inputs, + ) + self._guidance_scale = guidance_scale + self._guidance_rescale = guidance_rescale + self._interrupt = False + + # 3. Define 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 + + # 4. set timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps, num_inference_steps = self.get_timesteps( + num_inference_steps=num_inference_steps, strength=strength, device=device + ) + if comfyui_progressbar: + from comfy.utils import ProgressBar + pbar = ProgressBar(num_inference_steps + 3) + # at which timestep to set the initial noise (n.b. 50% if strength is 0.5) + latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) + # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise + is_strength_max = strength == 1.0 + + if video is not None: + video_length = video.shape[2] + init_video = self.image_processor.preprocess(rearrange(video, "b c f h w -> (b f) c h w"), height=height, width=width) + init_video = init_video.to(dtype=torch.float32) + init_video = rearrange(init_video, "(b f) c h w -> b c f h w", f=video_length) + else: + init_video = None + + # Prepare latent variables + num_channels_latents = self.vae.config.latent_channels + num_channels_transformer = self.transformer.config.in_channels + return_image_latents = num_channels_transformer == num_channels_latents + + # Make vae to cuda + if self.model_cpu_offload_flag: + self.vae = self.vae.to(device) + torch.cuda.empty_cache() + + # 5. Prepare latents. + latents_outputs = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + video_length, + prompt_embeds.dtype, + device, + generator, + latents, + video=init_video, + timestep=latent_timestep, + is_strength_max=is_strength_max, + return_noise=True, + return_video_latents=return_image_latents, + ) + if return_image_latents: + latents, noise, image_latents = latents_outputs + else: + latents, noise = latents_outputs + + latents_dtype = latents.dtype + if comfyui_progressbar: + pbar.update(1) + + if clip_image is not None: + inputs = self.clip_image_processor(images=clip_image, return_tensors="pt") + inputs["pixel_values"] = inputs["pixel_values"].to(latents.device, dtype=latents.dtype) + clip_encoder_hidden_states = self.clip_image_encoder(**inputs).last_hidden_state[:, 1:] + clip_encoder_hidden_states_neg = torch.zeros( + [ + batch_size, + int(self.clip_image_encoder.config.image_size / self.clip_image_encoder.config.patch_size) ** 2, + int(self.clip_image_encoder.config.hidden_size) + ] + ).to(latents.device, dtype=latents.dtype) + + clip_attention_mask = torch.ones([batch_size, self.transformer.n_query]).to(latents.device, dtype=latents.dtype) + clip_attention_mask_neg = torch.zeros([batch_size, self.transformer.n_query]).to(latents.device, dtype=latents.dtype) + + clip_encoder_hidden_states_input = torch.cat([clip_encoder_hidden_states_neg, clip_encoder_hidden_states]) if self.do_classifier_free_guidance else clip_encoder_hidden_states + clip_attention_mask_input = torch.cat([clip_attention_mask_neg, clip_attention_mask]) if self.do_classifier_free_guidance else clip_attention_mask + + elif clip_image is None and num_channels_transformer != num_channels_latents: + clip_encoder_hidden_states = torch.zeros( + [ + batch_size, + int(self.clip_image_encoder.config.image_size / self.clip_image_encoder.config.patch_size) ** 2, + int(self.clip_image_encoder.config.hidden_size) + ] + ).to(latents.device, dtype=latents.dtype) + + clip_attention_mask = torch.zeros([batch_size, self.transformer.n_query]) + clip_attention_mask = clip_attention_mask.to(latents.device, dtype=latents.dtype) + + clip_encoder_hidden_states_input = torch.cat([clip_encoder_hidden_states] * 2) if self.do_classifier_free_guidance else clip_encoder_hidden_states + clip_attention_mask_input = torch.cat([clip_attention_mask] * 2) if self.do_classifier_free_guidance else clip_attention_mask + + else: + clip_encoder_hidden_states_input = None + clip_attention_mask_input = None + if comfyui_progressbar: + pbar.update(1) + + if mask_video is not None: + if (mask_video == 255).all(): + mask_latents = torch.zeros_like(latents).to(latents.device, latents.dtype) + masked_video_latents = torch.zeros_like(latents).to(latents.device, latents.dtype) + + mask_input = torch.cat([mask_latents] * 2) if self.do_classifier_free_guidance else mask_latents + masked_video_latents_input = ( + torch.cat([masked_video_latents] * 2) if self.do_classifier_free_guidance else masked_video_latents + ) + inpaint_latents = torch.cat([mask_input, masked_video_latents_input], dim=1).to(latents.dtype) + else: + # Prepare mask latent variables + video_length = video.shape[2] + mask_condition = self.mask_processor.preprocess(rearrange(mask_video, "b c f h w -> (b f) c h w"), height=height, width=width) + mask_condition = mask_condition.to(dtype=torch.float32) + mask_condition = rearrange(mask_condition, "(b f) c h w -> b c f h w", f=video_length) + + if num_channels_transformer != num_channels_latents: + mask_condition_tile = torch.tile(mask_condition, [1, 3, 1, 1, 1]) + if masked_video_latents is None: + masked_video = init_video * (mask_condition_tile < 0.5) + torch.ones_like(init_video) * (mask_condition_tile > 0.5) * -1 + else: + masked_video = masked_video_latents + + mask_latents, masked_video_latents = self.prepare_mask_latents( + mask_condition_tile, + masked_video, + batch_size, + height, + width, + prompt_embeds.dtype, + device, + generator, + self.do_classifier_free_guidance, + ) + mask = torch.tile(mask_condition, [1, num_channels_latents, 1, 1, 1]) + mask = F.interpolate(mask, size=latents.size()[-3:], mode='trilinear', align_corners=True).to(latents.device, latents.dtype) + + mask_input = torch.cat([mask_latents] * 2) if self.do_classifier_free_guidance else mask_latents + masked_video_latents_input = ( + torch.cat([masked_video_latents] * 2) if self.do_classifier_free_guidance else masked_video_latents + ) + inpaint_latents = torch.cat([mask_input, masked_video_latents_input], dim=1).to(latents.dtype) + else: + mask = torch.tile(mask_condition, [1, num_channels_latents, 1, 1, 1]) + mask = F.interpolate(mask, size=latents.size()[-3:], mode='trilinear', align_corners=True).to(latents.device, latents.dtype) + + inpaint_latents = None + else: + if num_channels_transformer != num_channels_latents: + mask = torch.zeros_like(latents).to(latents.device, latents.dtype) + masked_video_latents = torch.zeros_like(latents).to(latents.device, latents.dtype) + + mask_input = torch.cat([mask] * 2) if self.do_classifier_free_guidance else mask + masked_video_latents_input = ( + torch.cat([masked_video_latents] * 2) if self.do_classifier_free_guidance else masked_video_latents + ) + inpaint_latents = torch.cat([mask_input, masked_video_latents_input], dim=1).to(latents.dtype) + else: + mask = torch.zeros_like(init_video[:, :1]) + mask = torch.tile(mask, [1, num_channels_latents, 1, 1, 1]) + mask = F.interpolate(mask, size=latents.size()[-3:], mode='trilinear', align_corners=True).to(latents.device, latents.dtype) + + inpaint_latents = None + if comfyui_progressbar: + pbar.update(1) + + # Check that sizes of mask, masked image and latents match + if num_channels_transformer == 48: + # default case for runwayml/stable-diffusion-inpainting + num_channels_mask = mask_latents.shape[1] + num_channels_masked_image = masked_video_latents.shape[1] + if num_channels_latents + num_channels_mask + num_channels_masked_image != self.transformer.config.in_channels: + raise ValueError( + f"Incorrect configuration settings! The config of `pipeline.transformer`: {self.transformer.config} expects" + f" {self.transformer.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" + f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" + f" = {num_channels_latents+num_channels_masked_image+num_channels_mask}. Please verify the config of" + " `pipeline.transformer` or your `mask_image` or `image` input." + ) + elif num_channels_transformer != num_channels_latents: + raise ValueError( + f"The transformer {self.transformer.__class__} should have 4 input channels, not {self.transformer.config.in_channels}." + ) + + # 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 image_rotary_emb, style embedding & time ids + grid_height = height // 8 // self.transformer.config.patch_size + grid_width = width // 8 // self.transformer.config.patch_size + base_size = 512 // 8 // self.transformer.config.patch_size + grid_crops_coords = get_resize_crop_region_for_grid((grid_height, grid_width), base_size) + image_rotary_emb = get_2d_rotary_pos_embed( + self.transformer.inner_dim // self.transformer.num_heads, grid_crops_coords, (grid_height, grid_width) + ) + + style = torch.tensor([0], device=device) + + target_size = target_size or (height, width) + add_time_ids = list(original_size + target_size + crops_coords_top_left) + add_time_ids = torch.tensor([add_time_ids], dtype=prompt_embeds.dtype) + + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask]) + prompt_embeds_2 = torch.cat([negative_prompt_embeds_2, prompt_embeds_2]) + prompt_attention_mask_2 = torch.cat([negative_prompt_attention_mask_2, prompt_attention_mask_2]) + add_time_ids = torch.cat([add_time_ids] * 2, dim=0) + style = torch.cat([style] * 2, dim=0) + + prompt_embeds = prompt_embeds.to(device=device) + prompt_attention_mask = prompt_attention_mask.to(device=device) + prompt_embeds_2 = prompt_embeds_2.to(device=device) + prompt_attention_mask_2 = prompt_attention_mask_2.to(device=device) + add_time_ids = add_time_ids.to(dtype=prompt_embeds.dtype, device=device).repeat( + batch_size * num_images_per_prompt, 1 + ) + style = style.to(device=device).repeat(batch_size * num_images_per_prompt) + + # Empty vae cache + if self.model_cpu_offload_flag: + self.vae = self.vae.to("cpu") + self.transformer = self.transformer.to(device) + torch.cuda.empty_cache() + + # 10. Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + self._num_timesteps = len(timesteps) + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + if i < len(timesteps) * (1 - clip_apply_ratio) and clip_encoder_hidden_states_input is not None: + clip_encoder_hidden_states_actual_input = torch.zeros_like(clip_encoder_hidden_states_input) + clip_attention_mask_actual_input = torch.zeros_like(clip_attention_mask_input) + else: + clip_encoder_hidden_states_actual_input = clip_encoder_hidden_states_input + clip_attention_mask_actual_input = clip_attention_mask_input + + current_timestep = t + if not torch.is_tensor(current_timestep): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = latent_model_input.device.type == "mps" + if isinstance(current_timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + current_timestep = torch.tensor([current_timestep], dtype=dtype, device=latent_model_input.device) + elif len(current_timestep.shape) == 0: + current_timestep = current_timestep[None].to(latent_model_input.device) + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + current_timestep = current_timestep.expand(latent_model_input.shape[0]).to( + dtype=latent_model_input.dtype + ) + + # predict the noise residual + noise_pred = self.transformer( + latent_model_input, + current_timestep, + encoder_hidden_states=prompt_embeds, + text_embedding_mask=prompt_attention_mask, + encoder_hidden_states_t5=prompt_embeds_2, + text_embedding_mask_t5=prompt_attention_mask_2, + image_meta_size=add_time_ids, + style=style, + image_rotary_emb=image_rotary_emb, + inpaint_latents=inpaint_latents, + clip_encoder_hidden_states=clip_encoder_hidden_states_actual_input, + clip_attention_mask=clip_attention_mask_actual_input, + return_dict=False, + )[0] + + noise_pred, _ = noise_pred.chunk(2, dim=1) + + # perform guidance + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + if self.do_classifier_free_guidance and guidance_rescale > 0.0: + # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf + noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + + if num_channels_transformer == 4: + init_latents_proper = image_latents + init_mask = mask + if i < len(timesteps) - 1: + noise_timestep = timesteps[i + 1] + init_latents_proper = self.scheduler.add_noise( + init_latents_proper, noise, torch.tensor([noise_timestep]) + ) + + latents = (1 - init_mask) * init_latents_proper + init_mask * latents + + 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) + prompt_embeds_2 = callback_outputs.pop("prompt_embeds_2", prompt_embeds_2) + negative_prompt_embeds_2 = callback_outputs.pop( + "negative_prompt_embeds_2", negative_prompt_embeds_2 + ) + + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + if comfyui_progressbar: + pbar.update(1) + + # Make vae to cuda + if self.model_cpu_offload_flag: + self.transformer = self.transformer.to("cpu") + self.vae = self.vae.to(device) + torch.cuda.empty_cache() + + # Post-processing + video = self.decode_latents(latents) + + if self.model_cpu_offload_flag: + # Make vae to cpu + self.vae = self.vae.to("cpu") + torch.cuda.empty_cache() + + # Convert to tensor + if output_type == "latent": + video = torch.from_numpy(video) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return video + + return RuyiPipelineOutput(videos=video) diff --git a/ruyi/utils/__init__.py b/ruyi/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ruyi/utils/diffusion_utils.py b/ruyi/utils/diffusion_utils.py new file mode 100644 index 0000000..a5e06e3 --- /dev/null +++ b/ruyi/utils/diffusion_utils.py @@ -0,0 +1,92 @@ +# Modified from OpenAI's diffusion repos +# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py +# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion +# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py + +import numpy as np +import torch as th + + +def normal_kl(mean1, logvar1, mean2, logvar2): + """ + Compute the KL divergence between two gaussians. + Shapes are automatically broadcasted, so batches can be compared to + scalars, among other use cases. + """ + tensor = next( + ( + obj + for obj in (mean1, logvar1, mean2, logvar2) + if isinstance(obj, th.Tensor) + ), + None, + ) + assert tensor is not None, "at least one argument must be a Tensor" + + # Force variances to be Tensors. Broadcasting helps convert scalars to + # Tensors, but it does not work for th.exp(). + logvar1, logvar2 = [ + x if isinstance(x, th.Tensor) else th.tensor(x, device=tensor.device) + for x in (logvar1, logvar2) + ] + + return 0.5 * ( + -1.0 + + logvar2 + - logvar1 + + th.exp(logvar1 - logvar2) + + ((mean1 - mean2) ** 2) * th.exp(-logvar2) + ) + + +def approx_standard_normal_cdf(x): + """ + A fast approximation of the cumulative distribution function of the + standard normal. + """ + return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) + + +def continuous_gaussian_log_likelihood(x, *, means, log_scales): + """ + Compute the log-likelihood of a continuous Gaussian distribution. + :param x: the targets + :param means: the Gaussian mean Tensor. + :param log_scales: the Gaussian log stddev Tensor. + :return: a tensor like x of log probabilities (in nats). + """ + centered_x = x - means + inv_stdv = th.exp(-log_scales) + normalized_x = centered_x * inv_stdv + return th.distributions.Normal(th.zeros_like(x), th.ones_like(x)).log_prob( + normalized_x + ) + + +def discretized_gaussian_log_likelihood(x, *, means, log_scales): + """ + Compute the log-likelihood of a Gaussian distribution discretizing to a + given image. + :param x: the target images. It is assumed that this was uint8 values, + rescaled to the range [-1, 1]. + :param means: the Gaussian mean Tensor. + :param log_scales: the Gaussian log stddev Tensor. + :return: a tensor like x of log probabilities (in nats). + """ + assert x.shape == means.shape == log_scales.shape + centered_x = x - means + inv_stdv = th.exp(-log_scales) + plus_in = inv_stdv * (centered_x + 1.0 / 255.0) + cdf_plus = approx_standard_normal_cdf(plus_in) + min_in = inv_stdv * (centered_x - 1.0 / 255.0) + cdf_min = approx_standard_normal_cdf(min_in) + log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) + log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) + cdf_delta = cdf_plus - cdf_min + log_probs = th.where( + x < -0.999, + log_cdf_plus, + th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), + ) + assert log_probs.shape == x.shape + return log_probs \ No newline at end of file diff --git a/ruyi/utils/gaussian_diffusion.py b/ruyi/utils/gaussian_diffusion.py new file mode 100644 index 0000000..669b580 --- /dev/null +++ b/ruyi/utils/gaussian_diffusion.py @@ -0,0 +1,1008 @@ +# Modified from OpenAI's diffusion repos +# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py +# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion +# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py + + +import enum +import math + +import numpy as np +import torch as th +import torch.nn.functional as F +from einops import rearrange + +from .diffusion_utils import discretized_gaussian_log_likelihood, normal_kl + + +def mean_flat(tensor): + """ + Take the mean over all non-batch dimensions. + """ + return tensor.mean(dim=list(range(1, len(tensor.shape)))) + + +class ModelMeanType(enum.Enum): + """ + Which type of output the model predicts. + """ + + PREVIOUS_X = enum.auto() # the model predicts x_{t-1} + START_X = enum.auto() # the model predicts x_0 + EPSILON = enum.auto() # the model predicts epsilon + + +class ModelVarType(enum.Enum): + """ + What is used as the model's output variance. + The LEARNED_RANGE option has been added to allow the model to predict + values between FIXED_SMALL and FIXED_LARGE, making its job easier. + """ + + LEARNED = enum.auto() + FIXED_SMALL = enum.auto() + FIXED_LARGE = enum.auto() + LEARNED_RANGE = enum.auto() + + +class LossType(enum.Enum): + MSE = enum.auto() # use raw MSE loss (and KL when learning variances) + RESCALED_MSE = ( + enum.auto() + ) # use raw MSE loss (with RESCALED_KL when learning variances) + KL = enum.auto() # use the variational lower-bound + RESCALED_KL = enum.auto() # like KL, but rescale to estimate the full VLB + + def is_vb(self): + return self in [LossType.KL, LossType.RESCALED_KL] + + +def _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, warmup_frac): + betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64) + warmup_time = int(num_diffusion_timesteps * warmup_frac) + betas[:warmup_time] = np.linspace(beta_start, beta_end, warmup_time, dtype=np.float64) + return betas + + +def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps): + """ + This is the deprecated API for creating beta schedules. + See get_named_beta_schedule() for the new library of schedules. + """ + if beta_schedule == "quad": + betas = ( + np.linspace( + beta_start ** 0.5, + beta_end ** 0.5, + num_diffusion_timesteps, + dtype=np.float64, + ) + ** 2 + ) + elif beta_schedule == "linear": + betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64) + elif beta_schedule == "warmup10": + betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.1) + elif beta_schedule == "warmup50": + betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.5) + elif beta_schedule == "const": + betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64) + elif beta_schedule == "jsd": # 1/T, 1/(T-1), 1/(T-2), ..., 1 + betas = 1.0 / np.linspace( + num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64 + ) + else: + raise NotImplementedError(beta_schedule) + assert betas.shape == (num_diffusion_timesteps,) + return betas + + +def get_named_beta_schedule(schedule_name, num_diffusion_timesteps): + """ + Get a pre-defined beta schedule for the given name. + The beta schedule library consists of beta schedules which remain similar + in the limit of num_diffusion_timesteps. + Beta schedules may be added, but should not be removed or changed once + they are committed to maintain backwards compatibility. + """ + if schedule_name == "linear": + # Linear schedule from Ho et al, extended to work for any number of + # diffusion steps. + scale = 1000 / num_diffusion_timesteps + return get_beta_schedule( + "linear", + beta_start=scale * 0.0001, + beta_end=scale * 0.02, + num_diffusion_timesteps=num_diffusion_timesteps, + ) + elif schedule_name == "squaredcos_cap_v2": + return betas_for_alpha_bar( + num_diffusion_timesteps, + lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2, + ) + else: + raise NotImplementedError(f"unknown beta schedule: {schedule_name}") + + +def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): + """ + Create a beta schedule that discretizes the given alpha_t_bar function, + which defines the cumulative product of (1-beta) over time from t = [0,1]. + :param num_diffusion_timesteps: the number of betas to produce. + :param alpha_bar: a lambda that takes an argument t from 0 to 1 and + produces the cumulative product of (1-beta) up to that + part of the diffusion process. + :param max_beta: the maximum beta to use; use values lower than 1 to + prevent singularities. + """ + betas = [] + for i in range(num_diffusion_timesteps): + t1 = i / num_diffusion_timesteps + t2 = (i + 1) / num_diffusion_timesteps + betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) + return np.array(betas) + + +class GaussianDiffusion: + """ + Utilities for training and sampling diffusion models. + Original ported from this codebase: + https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42 + :param betas: a 1-D numpy array of betas for each diffusion timestep, + starting at T and going to 1. + """ + + def __init__( + self, + *, + betas, + model_mean_type, + model_var_type, + loss_type, + snr=False, + return_startx=False, + ): + + self.model_mean_type = model_mean_type + self.model_var_type = model_var_type + self.loss_type = loss_type + self.snr = snr + self.return_startx = return_startx + + # Use float64 for accuracy. + betas = np.array(betas, dtype=np.float64) + self.betas = betas + assert len(betas.shape) == 1, "betas must be 1-D" + assert (betas > 0).all() and (betas <= 1).all() + + self.num_timesteps = int(betas.shape[0]) + + alphas = 1.0 - betas + self.alphas_cumprod = np.cumprod(alphas, axis=0) + self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) + self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0) + assert self.alphas_cumprod_prev.shape == (self.num_timesteps,) + + # calculations for diffusion q(x_t | x_{t-1}) and others + self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod) + self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod) + self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod) + self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod) + self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1) + + # calculations for posterior q(x_{t-1} | x_t, x_0) + self.posterior_variance = ( + betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) + ) + # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain + self.posterior_log_variance_clipped = np.log( + np.append(self.posterior_variance[1], self.posterior_variance[1:]) + ) if len(self.posterior_variance) > 1 else np.array([]) + + self.posterior_mean_coef1 = ( + betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) + ) + self.posterior_mean_coef2 = ( + (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod) + ) + + def q_mean_variance(self, x_start, t): + """ + Get the distribution q(x_t | x_0). + :param x_start: the [N x C x ...] tensor of noiseless inputs. + :param t: the number of diffusion steps (minus 1). Here, 0 means one step. + :return: A tuple (mean, variance, log_variance), all of x_start's shape. + """ + mean = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) + log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape) + return mean, variance, log_variance + + def q_sample(self, x_start, t, noise=None): + """ + Diffuse the data for a given number of diffusion steps. + In other words, sample from q(x_t | x_0). + :param x_start: the initial data batch. + :param t: the number of diffusion steps (minus 1). Here, 0 means one step. + :param noise: if specified, the split-out normal noise. + :return: A noisy version of x_start. + """ + if noise is None: + noise = th.randn_like(x_start) + assert noise.shape == x_start.shape + return ( + _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise + ) + + def q_posterior_mean_variance(self, x_start, x_t, t): + """ + Compute the mean and variance of the diffusion posterior: + q(x_{t-1} | x_t, x_0) + """ + assert x_start.shape == x_t.shape + posterior_mean = ( + _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start + + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t + ) + posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape) + posterior_log_variance_clipped = _extract_into_tensor( + self.posterior_log_variance_clipped, t, x_t.shape + ) + assert ( + posterior_mean.shape[0] + == posterior_variance.shape[0] + == posterior_log_variance_clipped.shape[0] + == x_start.shape[0] + ) + return posterior_mean, posterior_variance, posterior_log_variance_clipped + + def p_mean_variance(self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None): + """ + Apply the model to get p(x_{t-1} | x_t), as well as a prediction of + the initial x, x_0. + :param model: the model, which takes a signal and a batch of timesteps + as input. + :param x: the [N x C x ...] tensor at time t. + :param t: a 1-D Tensor of timesteps. + :param clip_denoised: if True, clip the denoised signal into [-1, 1]. + :param denoised_fn: if not None, a function which applies to the + x_start prediction before it is used to sample. Applies before + clip_denoised. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :return: a dict with the following keys: + - 'mean': the model mean output. + - 'variance': the model variance output. + - 'log_variance': the log of 'variance'. + - 'pred_xstart': the prediction for x_0. + """ + if model_kwargs is None: + model_kwargs = {} + + B, C = x.shape[:2] + assert t.shape == (B,) + model_output = model(x, timestep=t, **model_kwargs) + if isinstance(model_output, tuple): + model_output, extra = model_output + else: + extra = None + + if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]: + assert model_output.shape == (B, C * 2, *x.shape[2:]) + model_output, model_var_values = th.split(model_output, C, dim=1) + min_log = _extract_into_tensor(self.posterior_log_variance_clipped, t, x.shape) + max_log = _extract_into_tensor(np.log(self.betas), t, x.shape) + # The model_var_values is [-1, 1] for [min_var, max_var]. + frac = (model_var_values + 1) / 2 + model_log_variance = frac * max_log + (1 - frac) * min_log + model_variance = th.exp(model_log_variance) + elif self.model_var_type in [ModelVarType.FIXED_LARGE, ModelVarType.FIXED_SMALL]: + model_variance, model_log_variance = { + # for fixedlarge, we set the initial (log-)variance like so + # to get a better decoder log likelihood. + ModelVarType.FIXED_LARGE: ( + np.append(self.posterior_variance[1], self.betas[1:]), + np.log(np.append(self.posterior_variance[1], self.betas[1:])), + ), + ModelVarType.FIXED_SMALL: ( + self.posterior_variance, + self.posterior_log_variance_clipped, + ), + }[self.model_var_type] + model_variance = _extract_into_tensor(model_variance, t, x.shape) + model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape) + else: + model_variance = th.zeros_like(model_output) + model_log_variance = th.zeros_like(model_output) + + def process_xstart(x): + if denoised_fn is not None: + x = denoised_fn(x) + return x.clamp(-1, 1) if clip_denoised else x + + if self.model_mean_type == ModelMeanType.START_X: + pred_xstart = process_xstart(model_output) + else: + pred_xstart = process_xstart( + self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output) + ) + model_mean, _, _ = self.q_posterior_mean_variance(x_start=pred_xstart, x_t=x, t=t) + + assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape + return { + "mean": model_mean, + "variance": model_variance, + "log_variance": model_log_variance, + "pred_xstart": pred_xstart, + "extra": extra, + } + + def _predict_xstart_from_eps(self, x_t, t, eps): + assert x_t.shape == eps.shape + return ( + _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t + - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps + ) + + def _predict_eps_from_xstart(self, x_t, t, pred_xstart): + return ( + _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart + ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) + + def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None): + """ + Compute the mean for the previous step, given a function cond_fn that + computes the gradient of a conditional log probability with respect to + x. In particular, cond_fn computes grad(log(p(y|x))), and we want to + condition on y. + This uses the conditioning strategy from Sohl-Dickstein et al. (2015). + """ + gradient = cond_fn(x, t, **model_kwargs) + return p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float() + + def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None): + """ + Compute what the p_mean_variance output would have been, should the + model's score function be conditioned by cond_fn. + See condition_mean() for details on cond_fn. + Unlike condition_mean(), this instead uses the conditioning strategy + from Song et al (2020). + """ + alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) + + eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"]) + eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x, t, **model_kwargs) + + out = p_mean_var.copy() + out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps) + out["mean"], _, _ = self.q_posterior_mean_variance(x_start=out["pred_xstart"], x_t=x, t=t) + return out + + def p_sample( + self, + model, + x, + t, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + ): + """ + Sample x_{t-1} from the model at the given timestep. + :param model: the model to sample from. + :param x: the current tensor at x_{t-1}. + :param t: the value of t, starting at 0 for the first diffusion step. + :param clip_denoised: if True, clip the x_start prediction to [-1, 1]. + :param denoised_fn: if not None, a function which applies to the + x_start prediction before it is used to sample. + :param cond_fn: if not None, this is a gradient function that acts + similarly to the model. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :return: a dict containing the following keys: + - 'sample': a random sample from the model. + - 'pred_xstart': a prediction of x_0. + """ + out = self.p_mean_variance( + model, + x, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + model_kwargs=model_kwargs, + ) + noise = th.randn_like(x) + nonzero_mask = ( + (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) + ) # no noise when t == 0 + if cond_fn is not None: + out["mean"] = self.condition_mean(cond_fn, out, x, t, model_kwargs=model_kwargs) + sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise + return {"sample": sample, "pred_xstart": out["pred_xstart"]} + + def p_sample_loop( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + device=None, + progress=False, + ): + """ + Generate samples from the model. + :param model: the model module. + :param shape: the shape of the samples, (N, C, H, W). + :param noise: if specified, the noise from the encoder to sample. + Should be of the same shape as `shape`. + :param clip_denoised: if True, clip x_start predictions to [-1, 1]. + :param denoised_fn: if not None, a function which applies to the + x_start prediction before it is used to sample. + :param cond_fn: if not None, this is a gradient function that acts + similarly to the model. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :param device: if specified, the device to create the samples on. + If not specified, use a model parameter's device. + :param progress: if True, show a tqdm progress bar. + :return: a non-differentiable batch of samples. + """ + final = None + for sample in self.p_sample_loop_progressive( + model, + shape, + noise=noise, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + cond_fn=cond_fn, + model_kwargs=model_kwargs, + device=device, + progress=progress, + ): + final = sample + return final["sample"] + + def p_sample_loop_progressive( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + device=None, + progress=False, + ): + """ + Generate samples from the model and yield intermediate samples from + each timestep of diffusion. + Arguments are the same as p_sample_loop(). + Returns a generator over dicts, where each dict is the return value of + p_sample(). + """ + if device is None: + device = next(model.parameters()).device + assert isinstance(shape, (tuple, list)) + img = noise if noise is not None else th.randn(*shape, device=device) + indices = list(range(self.num_timesteps))[::-1] + + if progress: + # Lazy import so that we don't depend on tqdm. + from tqdm.auto import tqdm + + indices = tqdm(indices) + + for i in indices: + t = th.tensor([i] * shape[0], device=device) + with th.no_grad(): + out = self.p_sample( + model, + img, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + cond_fn=cond_fn, + model_kwargs=model_kwargs, + ) + yield out + img = out["sample"] + + def ddim_sample( + self, + model, + x, + t, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + eta=0.0, + ): + """ + Sample x_{t-1} from the model using DDIM. + Same usage as p_sample(). + """ + out = self.p_mean_variance( + model, + x, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + model_kwargs=model_kwargs, + ) + if cond_fn is not None: + out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) + + # Usually our model outputs epsilon, but we re-derive it + # in case we used x_start or x_prev prediction. + eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"]) + + alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) + alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape) + sigma = ( + eta + * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) + * th.sqrt(1 - alpha_bar / alpha_bar_prev) + ) + # Equation 12. + noise = th.randn_like(x) + mean_pred = ( + out["pred_xstart"] * th.sqrt(alpha_bar_prev) + + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps + ) + nonzero_mask = ( + (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) + ) # no noise when t == 0 + sample = mean_pred + nonzero_mask * sigma * noise + return {"sample": sample, "pred_xstart": out["pred_xstart"]} + + def ddim_reverse_sample( + self, + model, + x, + t, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + eta=0.0, + ): + """ + Sample x_{t+1} from the model using DDIM reverse ODE. + """ + assert eta == 0.0, "Reverse ODE only for deterministic path" + out = self.p_mean_variance( + model, + x, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + model_kwargs=model_kwargs, + ) + if cond_fn is not None: + out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) + # Usually our model outputs epsilon, but we re-derive it + # in case we used x_start or x_prev prediction. + eps = ( + _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x + - out["pred_xstart"] + ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape) + alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape) + + # Equation 12. reversed + mean_pred = out["pred_xstart"] * th.sqrt(alpha_bar_next) + th.sqrt(1 - alpha_bar_next) * eps + + return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]} + + def ddim_sample_loop( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + device=None, + progress=False, + eta=0.0, + ): + """ + Generate samples from the model using DDIM. + Same usage as p_sample_loop(). + """ + final = None + for sample in self.ddim_sample_loop_progressive( + model, + shape, + noise=noise, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + cond_fn=cond_fn, + model_kwargs=model_kwargs, + device=device, + progress=progress, + eta=eta, + ): + final = sample + return final["sample"] + + def ddim_sample_loop_progressive( + self, + model, + shape, + noise=None, + clip_denoised=True, + denoised_fn=None, + cond_fn=None, + model_kwargs=None, + device=None, + progress=False, + eta=0.0, + ): + """ + Use DDIM to sample from the model and yield intermediate samples from + each timestep of DDIM. + Same usage as p_sample_loop_progressive(). + """ + if device is None: + device = next(model.parameters()).device + assert isinstance(shape, (tuple, list)) + img = noise if noise is not None else th.randn(*shape, device=device) + indices = list(range(self.num_timesteps))[::-1] + + if progress: + # Lazy import so that we don't depend on tqdm. + from tqdm.auto import tqdm + + indices = tqdm(indices) + + for i in indices: + t = th.tensor([i] * shape[0], device=device) + with th.no_grad(): + out = self.ddim_sample( + model, + img, + t, + clip_denoised=clip_denoised, + denoised_fn=denoised_fn, + cond_fn=cond_fn, + model_kwargs=model_kwargs, + eta=eta, + ) + yield out + img = out["sample"] + + def _vb_terms_bpd( + self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None + ): + """ + Get a term for the variational lower-bound. + The resulting units are bits (rather than nats, as one might expect). + This allows for comparison to other papers. + :return: a dict with the following keys: + - 'output': a shape [N] tensor of NLLs or KLs. + - 'pred_xstart': the x_0 predictions. + """ + true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance( + x_start=x_start, x_t=x_t, t=t + ) + out = self.p_mean_variance( + model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs + ) + kl = normal_kl( + true_mean, true_log_variance_clipped, out["mean"], out["log_variance"] + ) + kl = mean_flat(kl) / np.log(2.0) + + decoder_nll = -discretized_gaussian_log_likelihood( + x_start, means=out["mean"], log_scales=0.5 * out["log_variance"] + ) + assert decoder_nll.shape == x_start.shape + decoder_nll = mean_flat(decoder_nll) / np.log(2.0) + + # At the first timestep return the decoder NLL, + # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t)) + output = th.where((t == 0), decoder_nll, kl) + return {"output": output, "pred_xstart": out["pred_xstart"]} + + def training_losses(self, model, x_start, timestep, model_kwargs=None, noise=None, skip_noise=False): + """ + Compute training losses for a single timestep. + :param model: the model to evaluate loss on. + :param x_start: the [N x C x ...] tensor of inputs. + :param t: a batch of timestep indices. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :param noise: if specified, the specific Gaussian noise to try to remove. + :return: a dict with the key "loss" containing a tensor of shape [N]. + Some mean or variance settings may also have other keys. + """ + t = timestep + if model_kwargs is None: + model_kwargs = {} + if skip_noise: + x_t = x_start + else: + if noise is None: + noise = th.randn_like(x_start) + x_t = self.q_sample(x_start, t, noise=noise) + + terms = {} + + if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL: + terms["loss"] = self._vb_terms_bpd( + model=model, + x_start=x_start, + x_t=x_t, + t=t, + clip_denoised=False, + model_kwargs=model_kwargs, + )["output"] + if self.loss_type == LossType.RESCALED_KL: + terms["loss"] *= self.num_timesteps + elif self.loss_type in [LossType.MSE, LossType.RESCALED_MSE]: + model_output = model(x_t, timestep=t, **model_kwargs)[0] + + if isinstance(model_output, dict) and model_output.get('x', None) is not None: + output = model_output['x'] + else: + output = model_output + + if self.return_startx and self.model_mean_type == ModelMeanType.EPSILON: + return self._extracted_from_training_losses_diffusers(x_t, output, t) + # self.model_var_type = ModelVarType.LEARNED_RANGE:4 + if self.model_var_type in [ + ModelVarType.LEARNED, + ModelVarType.LEARNED_RANGE, + ]: + B, C = x_t.shape[:2] + assert output.shape == (B, C * 2, *x_t.shape[2:]) + output, model_var_values = th.split(output, C, dim=1) + # Learn the variance using the variational bound, but don't let it affect our mean prediction. + frozen_out = th.cat([output.detach(), model_var_values], dim=1) + # vb variational bound + terms["vb"] = self._vb_terms_bpd( + model=lambda *args, r=frozen_out, **kwargs: r, + x_start=x_start, + x_t=x_t, + t=t, + clip_denoised=False, + )["output"] + if self.loss_type == LossType.RESCALED_MSE: + # Divide by 1000 for equivalence with initial implementation. + # Without a factor of 1/1000, the VB term hurts the MSE term. + terms["vb"] *= self.num_timesteps / 1000.0 + + target = { + ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance( + x_start=x_start, x_t=x_t, t=t + )[0], + ModelMeanType.START_X: x_start, + ModelMeanType.EPSILON: noise, + }[self.model_mean_type] + assert output.shape == target.shape == x_start.shape + if self.snr: + if self.model_mean_type == ModelMeanType.START_X: + pred_noise = self._predict_eps_from_xstart(x_t=x_t, t=t, pred_xstart=output) + pred_startx = output + elif self.model_mean_type == ModelMeanType.EPSILON: + pred_noise = output + pred_startx = self._predict_xstart_from_eps(x_t=x_t, t=t, eps=output) + # terms["mse_eps"] = mean_flat((noise - pred_noise) ** 2) + # terms["mse_x0"] = mean_flat((x_start - pred_startx) ** 2) + + t = t[:, None, None, None, None].expand(pred_startx.shape) # [128, 4, 32, 32] + # best + target = th.where(t > 249, noise, x_start) + output = th.where(t > 249, pred_noise, pred_startx) + loss = (target - output) ** 2 + if model_kwargs.get('mask_ratio', False) and model_kwargs['mask_ratio'] > 0: + assert 'mask' in model_output + loss = F.avg_pool2d(loss.mean(dim=1), model.model.module.patch_size).flatten(1) + mask = model_output['mask'] + unmask = 1 - mask + terms['mse'] = mean_flat(loss * unmask) * unmask.shape[1]/unmask.sum(1) + if model_kwargs['mask_loss_coef'] > 0: + terms['mae'] = model_kwargs['mask_loss_coef'] * mean_flat(loss * mask) * mask.shape[1]/mask.sum(1) + else: + terms["mse"] = mean_flat(loss) + terms["loss"] = terms["mse"] + terms["vb"] if "vb" in terms else terms["mse"] + if "mae" in terms: + terms["loss"] = terms["loss"] + terms["mae"] + else: + raise NotImplementedError(self.loss_type) + + return terms + + def training_losses_diffusers(self, model, x_start, timestep, model_kwargs=None, noise=None, skip_noise=False): + """ + Compute training losses for a single timestep. + :param model: the model to evaluate loss on. + :param x_start: the [N x C x ...] tensor of inputs. + :param t: a batch of timestep indices. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :param noise: if specified, the specific Gaussian noise to try to remove. + :return: a dict with the key "loss" containing a tensor of shape [N]. + Some mean or variance settings may also have other keys. + """ + t = timestep + if model_kwargs is None: + model_kwargs = {} + if skip_noise: + x_t = x_start + else: + if noise is None: + noise = th.randn_like(x_start) + x_t = self.q_sample(x_start, t, noise=noise) + + terms = {} + + if self.loss_type in [LossType.KL, LossType.RESCALED_KL]: + terms["loss"] = self._vb_terms_bpd( + model=model, + x_start=x_start, + x_t=x_t, + t=t, + clip_denoised=False, + model_kwargs=model_kwargs, + )["output"] + if self.loss_type == LossType.RESCALED_KL: + terms["loss"] *= self.num_timesteps + elif self.loss_type in [LossType.MSE, LossType.RESCALED_MSE]: + output = model(x_t, timestep=t, **model_kwargs, return_dict=False)[0] + + if self.return_startx and self.model_mean_type == ModelMeanType.EPSILON: + return self._extracted_from_training_losses_diffusers(x_t, output, t) + + if self.model_var_type in [ + ModelVarType.LEARNED, + ModelVarType.LEARNED_RANGE, + ]: + B, C = x_t.shape[:2] + assert output.shape == (B, C * 2, *x_t.shape[2:]) + output, model_var_values = th.split(output, C, dim=1) + # Learn the variance using the variational bound, but don't let it affect our mean prediction. + frozen_out = th.cat([output.detach(), model_var_values], dim=1) + terms["vb"] = self._vb_terms_bpd( + model=lambda *args, r=frozen_out, **kwargs: r, + x_start=x_start, + x_t=x_t, + t=t, + clip_denoised=False, + )["output"] + if self.loss_type == LossType.RESCALED_MSE: + # Divide by 1000 for equivalence with initial implementation. + # Without a factor of 1/1000, the VB term hurts the MSE term. + terms["vb"] *= self.num_timesteps / 1000.0 + + target = { + ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance( + x_start=x_start, x_t=x_t, t=t + )[0], + ModelMeanType.START_X: x_start, + ModelMeanType.EPSILON: noise, + }[self.model_mean_type] + assert output.shape == target.shape == x_start.shape + if self.snr: + if self.model_mean_type == ModelMeanType.START_X: + pred_noise = self._predict_eps_from_xstart(x_t=x_t, t=t, pred_xstart=output) + pred_startx = output + elif self.model_mean_type == ModelMeanType.EPSILON: + pred_noise = output + pred_startx = self._predict_xstart_from_eps(x_t=x_t, t=t, eps=output) + # terms["mse_eps"] = mean_flat((noise - pred_noise) ** 2) + # terms["mse_x0"] = mean_flat((x_start - pred_startx) ** 2) + + t = t[:, None, None, None].expand(pred_startx.shape) # [128, 4, 32, 32] + # best + target = th.where(t > 249, noise, x_start) + output = th.where(t > 249, pred_noise, pred_startx) + loss = (target - output) ** 2 + terms["mse"] = mean_flat(loss) + terms["loss"] = terms["mse"] + terms["vb"] if "vb" in terms else terms["mse"] + if "mae" in terms: + terms["loss"] = terms["loss"] + terms["mae"] + else: + raise NotImplementedError(self.loss_type) + + return terms + + def _extracted_from_training_losses_diffusers(self, x_t, output, t): + B, C = x_t.shape[:2] + assert output.shape == (B, C * 2, *x_t.shape[2:]) + output = th.split(output, C, dim=1)[0] + return output, self._predict_xstart_from_eps(x_t=x_t, t=t, eps=output), x_t + + def _prior_bpd(self, x_start): + """ + Get the prior KL term for the variational lower-bound, measured in + bits-per-dim. + This term can't be optimized, as it only depends on the encoder. + :param x_start: the [N x C x ...] tensor of inputs. + :return: a batch of [N] KL values (in bits), one per batch element. + """ + batch_size = x_start.shape[0] + t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device) + qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t) + kl_prior = normal_kl( + mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0 + ) + return mean_flat(kl_prior) / np.log(2.0) + + def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None): + """ + Compute the entire variational lower-bound, measured in bits-per-dim, + as well as other related quantities. + :param model: the model to evaluate loss on. + :param x_start: the [N x C x ...] tensor of inputs. + :param clip_denoised: if True, clip denoised samples. + :param model_kwargs: if not None, a dict of extra keyword arguments to + pass to the model. This can be used for conditioning. + :return: a dict containing the following keys: + - total_bpd: the total variational lower-bound, per batch element. + - prior_bpd: the prior term in the lower-bound. + - vb: an [N x T] tensor of terms in the lower-bound. + - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep. + - mse: an [N x T] tensor of epsilon MSEs for each timestep. + """ + device = x_start.device + batch_size = x_start.shape[0] + + vb = [] + xstart_mse = [] + mse = [] + for t in list(range(self.num_timesteps))[::-1]: + t_batch = th.tensor([t] * batch_size, device=device) + noise = th.randn_like(x_start) + x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise) + # Calculate VLB term at the current timestep + with th.no_grad(): + out = self._vb_terms_bpd( + model, + x_start=x_start, + x_t=x_t, + t=t_batch, + clip_denoised=clip_denoised, + model_kwargs=model_kwargs, + ) + vb.append(out["output"]) + xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2)) + eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"]) + mse.append(mean_flat((eps - noise) ** 2)) + + vb = th.stack(vb, dim=1) + xstart_mse = th.stack(xstart_mse, dim=1) + mse = th.stack(mse, dim=1) + + prior_bpd = self._prior_bpd(x_start) + total_bpd = vb.sum(dim=1) + prior_bpd + return { + "total_bpd": total_bpd, + "prior_bpd": prior_bpd, + "vb": vb, + "xstart_mse": xstart_mse, + "mse": mse, + } + + +def _extract_into_tensor(arr, timesteps, broadcast_shape): + """ + Extract values from a 1-D numpy array for a batch of indices. + :param arr: the 1-D numpy array. + :param timesteps: a tensor of indices into the array to extract. + :param broadcast_shape: a larger shape of K dimensions with the batch + dimension equal to the length of timesteps. + :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims. + """ + res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float() + while len(res.shape) < len(broadcast_shape): + res = res[..., None] + return res + th.zeros(broadcast_shape, device=timesteps.device) \ No newline at end of file diff --git a/ruyi/utils/lora_utils.py b/ruyi/utils/lora_utils.py new file mode 100644 index 0000000..ba9dd18 --- /dev/null +++ b/ruyi/utils/lora_utils.py @@ -0,0 +1,477 @@ +# LoRA network module +# reference: +# https://github.com/microsoft/LoRA/blob/main/loralib/layers.py +# https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py +# https://github.com/bmaltais/kohya_ss + +import hashlib +import math +import os +from collections import defaultdict +from io import BytesIO +from typing import List, Optional, Type, Union + +import safetensors.torch +import torch +import torch.utils.checkpoint +from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear +from safetensors.torch import load_file +from transformers import T5EncoderModel + + +class LoRAModule(torch.nn.Module): + """ + replaces forward method of the original Linear, instead of replacing the original Linear module. + """ + + def __init__( + self, + lora_name, + org_module: torch.nn.Module, + multiplier=1.0, + lora_dim=4, + alpha=1, + dropout=None, + rank_dropout=None, + module_dropout=None, + ): + """if alpha == 0 or None, alpha is rank (no scaling).""" + super().__init__() + self.lora_name = lora_name + + if org_module.__class__.__name__ == "Conv2d": + in_dim = org_module.in_channels + out_dim = org_module.out_channels + else: + in_dim = org_module.in_features + out_dim = org_module.out_features + + self.lora_dim = lora_dim + if org_module.__class__.__name__ == "Conv2d": + kernel_size = org_module.kernel_size + stride = org_module.stride + padding = org_module.padding + self.lora_down = torch.nn.Conv2d(in_dim, self.lora_dim, kernel_size, stride, padding, bias=False) + self.lora_up = torch.nn.Conv2d(self.lora_dim, out_dim, (1, 1), (1, 1), bias=False) + else: + self.lora_down = torch.nn.Linear(in_dim, self.lora_dim, bias=False) + self.lora_up = torch.nn.Linear(self.lora_dim, out_dim, bias=False) + + if type(alpha) == torch.Tensor: + alpha = alpha.detach().float().numpy() # without casting, bf16 causes error + alpha = self.lora_dim if alpha is None or alpha == 0 else alpha + self.scale = alpha / self.lora_dim + self.register_buffer("alpha", torch.tensor(alpha)) + + # same as microsoft's + torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5)) + torch.nn.init.zeros_(self.lora_up.weight) + + self.multiplier = multiplier + self.org_module = org_module # remove in applying + self.dropout = dropout + self.rank_dropout = rank_dropout + self.module_dropout = module_dropout + + def apply_to(self): + self.org_forward = self.org_module.forward + self.org_module.forward = self.forward + del self.org_module + + def forward(self, x, *args, **kwargs): + weight_dtype = x.dtype + org_forwarded = self.org_forward(x) + + # module dropout + if self.module_dropout is not None and self.training: + if torch.rand(1) < self.module_dropout: + return org_forwarded + + lx = self.lora_down(x.to(self.lora_down.weight.dtype)) + + # normal dropout + if self.dropout is not None and self.training: + lx = torch.nn.functional.dropout(lx, p=self.dropout) + + # rank dropout + if self.rank_dropout is not None and self.training: + mask = torch.rand((lx.size(0), self.lora_dim), device=lx.device) > self.rank_dropout + if len(lx.size()) == 3: + mask = mask.unsqueeze(1) # for Text Encoder + elif len(lx.size()) == 4: + mask = mask.unsqueeze(-1).unsqueeze(-1) # for Conv2d + lx = lx * mask + + # scaling for rank dropout: treat as if the rank is changed + scale = self.scale * (1.0 / (1.0 - self.rank_dropout)) # redundant for readability + else: + scale = self.scale + + lx = self.lora_up(lx) + + return org_forwarded.to(weight_dtype) + lx.to(weight_dtype) * self.multiplier * scale + + +def addnet_hash_legacy(b): + """Old model hash used by sd-webui-additional-networks for .safetensors format files""" + m = hashlib.sha256() + + b.seek(0x100000) + m.update(b.read(0x10000)) + return m.hexdigest()[0:8] + + +def addnet_hash_safetensors(b): + """New model hash used by sd-webui-additional-networks for .safetensors format files""" + hash_sha256 = hashlib.sha256() + blksize = 1024 * 1024 + + b.seek(0) + header = b.read(8) + n = int.from_bytes(header, "little") + + offset = n + 8 + b.seek(offset) + for chunk in iter(lambda: b.read(blksize), b""): + hash_sha256.update(chunk) + + return hash_sha256.hexdigest() + + +def precalculate_safetensors_hashes(tensors, metadata): + """Precalculate the model hashes needed by sd-webui-additional-networks to + save time on indexing the model later.""" + + # Because writing user metadata to the file can change the result of + # sd_models.model_hash(), only retain the training metadata for purposes of + # calculating the hash, as they are meant to be immutable + metadata = {k: v for k, v in metadata.items() if k.startswith("ss_")} + + bytes = safetensors.torch.save(tensors, metadata) + b = BytesIO(bytes) + + model_hash = addnet_hash_safetensors(b) + legacy_hash = addnet_hash_legacy(b) + return model_hash, legacy_hash + + +class LoRANetwork(torch.nn.Module): + TRANSFORMER_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Transformer3DModel", "HunyuanTransformer3DModel"] + TEXT_ENCODER_TARGET_REPLACE_MODULE = ["T5LayerSelfAttention", "T5LayerFF", "BertEncoder"] + LORA_PREFIX_TRANSFORMER = "lora_unet" + LORA_PREFIX_TEXT_ENCODER = "lora_te" + def __init__( + self, + text_encoder: Union[List[T5EncoderModel], T5EncoderModel], + unet, + multiplier: float = 1.0, + lora_dim: int = 4, + alpha: float = 1, + dropout: Optional[float] = None, + module_class: Type[object] = LoRAModule, + add_lora_in_attn_temporal: bool = False, + varbose: Optional[bool] = False, + ) -> None: + super().__init__() + self.multiplier = multiplier + + self.lora_dim = lora_dim + self.alpha = alpha + self.dropout = dropout + + print(f"create LoRA network. base dim (rank): {lora_dim}, alpha: {alpha}") + print(f"neuron dropout: p={self.dropout}") + + # create module instances + def create_modules( + is_unet: bool, + root_module: torch.nn.Module, + target_replace_modules: List[torch.nn.Module], + ) -> List[LoRAModule]: + prefix = ( + self.LORA_PREFIX_TRANSFORMER + if is_unet + else self.LORA_PREFIX_TEXT_ENCODER + ) + loras = [] + skipped = [] + for name, module in root_module.named_modules(): + if module.__class__.__name__ in target_replace_modules: + for child_name, child_module in module.named_modules(): + is_linear = child_module.__class__.__name__ == "Linear" or child_module.__class__.__name__ == "LoRACompatibleLinear" + is_conv2d = child_module.__class__.__name__ == "Conv2d" or child_module.__class__.__name__ == "LoRACompatibleConv" + is_conv2d_1x1 = is_conv2d and child_module.kernel_size == (1, 1) + + if not add_lora_in_attn_temporal: + if "attn_temporal" in child_name: + continue + + if is_linear or is_conv2d: + lora_name = prefix + "." + name + "." + child_name + lora_name = lora_name.replace(".", "_") + + dim = None + alpha = None + + if is_linear or is_conv2d_1x1: + dim = self.lora_dim + alpha = self.alpha + + if dim is None or dim == 0: + if is_linear or is_conv2d_1x1: + skipped.append(lora_name) + continue + + lora = module_class( + lora_name, + child_module, + self.multiplier, + dim, + alpha, + dropout=dropout, + ) + loras.append(lora) + return loras, skipped + + text_encoders = text_encoder if type(text_encoder) == list else [text_encoder] + + self.text_encoder_loras = [] + skipped_te = [] + for i, text_encoder in enumerate(text_encoders): + if text_encoder is not None: + text_encoder_loras, skipped = create_modules(False, text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE) + self.text_encoder_loras.extend(text_encoder_loras) + skipped_te += skipped + print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.") + + self.unet_loras, skipped_un = create_modules(True, unet, LoRANetwork.TRANSFORMER_TARGET_REPLACE_MODULE) + print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.") + + # assertion + names = set() + for lora in self.text_encoder_loras + self.unet_loras: + assert lora.lora_name not in names, f"duplicated lora name: {lora.lora_name}" + names.add(lora.lora_name) + + def apply_to(self, text_encoder, unet, apply_text_encoder=True, apply_unet=True): + if apply_text_encoder: + print("enable LoRA for text encoder") + else: + self.text_encoder_loras = [] + + if apply_unet: + print("enable LoRA for U-Net") + else: + self.unet_loras = [] + + for lora in self.text_encoder_loras + self.unet_loras: + lora.apply_to() + self.add_module(lora.lora_name, lora) + + def set_multiplier(self, multiplier): + self.multiplier = multiplier + for lora in self.text_encoder_loras + self.unet_loras: + lora.multiplier = self.multiplier + + def load_weights(self, file): + if os.path.splitext(file)[1] == ".safetensors": + from safetensors.torch import load_file + + weights_sd = load_file(file) + else: + weights_sd = torch.load(file, map_location="cpu") + info = self.load_state_dict(weights_sd, False) + return info + + def prepare_optimizer_params(self, text_encoder_lr, unet_lr, default_lr): + self.requires_grad_(True) + all_params = [] + + def enumerate_params(loras): + params = [] + for lora in loras: + params.extend(lora.parameters()) + return params + + if self.text_encoder_loras: + param_data = {"params": enumerate_params(self.text_encoder_loras)} + if text_encoder_lr is not None: + param_data["lr"] = text_encoder_lr + all_params.append(param_data) + + if self.unet_loras: + param_data = {"params": enumerate_params(self.unet_loras)} + if unet_lr is not None: + param_data["lr"] = unet_lr + all_params.append(param_data) + + return all_params + + def enable_gradient_checkpointing(self): + pass + + def get_trainable_params(self): + return self.parameters() + + def save_weights(self, file, dtype, metadata): + if metadata is not None and len(metadata) == 0: + metadata = None + + state_dict = self.state_dict() + + if dtype is not None: + for key in list(state_dict.keys()): + v = state_dict[key] + v = v.detach().clone().to("cpu").to(dtype) + state_dict[key] = v + + if os.path.splitext(file)[1] == ".safetensors": + from safetensors.torch import save_file + + # Precalculate model hashes to save time on indexing + if metadata is None: + metadata = {} + model_hash, legacy_hash = precalculate_safetensors_hashes(state_dict, metadata) + metadata["sshs_model_hash"] = model_hash + metadata["sshs_legacy_hash"] = legacy_hash + + save_file(state_dict, file, metadata) + else: + torch.save(state_dict, file) + +def create_network( + multiplier: float, + network_dim: Optional[int], + network_alpha: Optional[float], + text_encoder: Union[T5EncoderModel, List[T5EncoderModel]], + transformer, + neuron_dropout: Optional[float] = None, + add_lora_in_attn_temporal: bool = False, + **kwargs, +): + if network_dim is None: + network_dim = 4 # default + if network_alpha is None: + network_alpha = 1.0 + + network = LoRANetwork( + text_encoder, + transformer, + multiplier=multiplier, + lora_dim=network_dim, + alpha=network_alpha, + dropout=neuron_dropout, + add_lora_in_attn_temporal=add_lora_in_attn_temporal, + varbose=True, + ) + return network + +def merge_lora(pipeline, lora_path, multiplier, device='cpu', dtype=torch.float32, state_dict=None, transformer_only=False): + LORA_PREFIX_TRANSFORMER = "lora_unet" + LORA_PREFIX_TEXT_ENCODER = "lora_te" + if state_dict is None: + state_dict = load_file(lora_path, device=device) + else: + state_dict = state_dict + updates = defaultdict(dict) + for key, value in state_dict.items(): + layer, elem = key.split('.', 1) + updates[layer][elem] = value + + for layer, elems in updates.items(): + + if "lora_te" in layer: + if transformer_only: + continue + else: + layer_infos = layer.split(LORA_PREFIX_TEXT_ENCODER + "_")[-1].split("_") + curr_layer = pipeline.text_encoder + else: + layer_infos = layer.split(LORA_PREFIX_TRANSFORMER + "_")[-1].split("_") + curr_layer = pipeline.transformer + + temp_name = layer_infos.pop(0) + while len(layer_infos) > -1: + try: + curr_layer = curr_layer.__getattr__(temp_name) + if len(layer_infos) > 0: + temp_name = layer_infos.pop(0) + elif len(layer_infos) == 0: + break + except Exception: + if len(layer_infos) == 0: + print('Error loading layer') + if len(temp_name) > 0: + temp_name += "_" + layer_infos.pop(0) + else: + temp_name = layer_infos.pop(0) + + weight_up = elems['lora_up.weight'].to(dtype) + weight_down = elems['lora_down.weight'].to(dtype) + if 'alpha' in elems.keys(): + alpha = elems['alpha'].item() / weight_up.shape[1] + else: + alpha = 1.0 + + curr_layer.weight.data = curr_layer.weight.data.to(device) + if len(weight_up.shape) == 4: + curr_layer.weight.data += multiplier * alpha * torch.mm(weight_up.squeeze(3).squeeze(2), + weight_down.squeeze(3).squeeze(2)).unsqueeze( + 2).unsqueeze(3) + else: + curr_layer.weight.data += multiplier * alpha * torch.mm(weight_up, weight_down) + + return pipeline + +# TODO: Refactor with merge_lora. +def unmerge_lora(pipeline, lora_path, multiplier=1, device="cpu", dtype=torch.float32): + """Unmerge state_dict in LoRANetwork from the pipeline in diffusers.""" + LORA_PREFIX_UNET = "lora_unet" + LORA_PREFIX_TEXT_ENCODER = "lora_te" + state_dict = load_file(lora_path, device=device) + + updates = defaultdict(dict) + for key, value in state_dict.items(): + layer, elem = key.split('.', 1) + updates[layer][elem] = value + + for layer, elems in updates.items(): + + if "lora_te" in layer: + layer_infos = layer.split(LORA_PREFIX_TEXT_ENCODER + "_")[-1].split("_") + curr_layer = pipeline.text_encoder + else: + layer_infos = layer.split(LORA_PREFIX_UNET + "_")[-1].split("_") + curr_layer = pipeline.transformer + + temp_name = layer_infos.pop(0) + while len(layer_infos) > -1: + try: + curr_layer = curr_layer.__getattr__(temp_name) + if len(layer_infos) > 0: + temp_name = layer_infos.pop(0) + elif len(layer_infos) == 0: + break + except Exception: + if len(layer_infos) == 0: + print('Error loading layer') + if len(temp_name) > 0: + temp_name += "_" + layer_infos.pop(0) + else: + temp_name = layer_infos.pop(0) + + weight_up = elems['lora_up.weight'].to(dtype) + weight_down = elems['lora_down.weight'].to(dtype) + if 'alpha' in elems.keys(): + alpha = elems['alpha'].item() / weight_up.shape[1] + else: + alpha = 1.0 + + curr_layer.weight.data = curr_layer.weight.data.to(device) + if len(weight_up.shape) == 4: + curr_layer.weight.data -= multiplier * alpha * torch.mm(weight_up.squeeze(3).squeeze(2), + weight_down.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3) + else: + curr_layer.weight.data -= multiplier * alpha * torch.mm(weight_up, weight_down) + + return pipeline \ No newline at end of file diff --git a/ruyi/utils/respace.py b/ruyi/utils/respace.py new file mode 100644 index 0000000..bbdfe77 --- /dev/null +++ b/ruyi/utils/respace.py @@ -0,0 +1,131 @@ +# Modified from OpenAI's diffusion repos +# GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py +# ADM: https://github.com/openai/guided-diffusion/blob/main/guided_diffusion +# IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py + +import numpy as np +import torch as th + +from .gaussian_diffusion import GaussianDiffusion + + +def space_timesteps(num_timesteps, section_counts): + """ + Create a list of timesteps to use from an original diffusion process, + given the number of timesteps we want to take from equally-sized portions + of the original process. + For example, if there's 300 timesteps and the section counts are [10,15,20] + then the first 100 timesteps are strided to be 10 timesteps, the second 100 + are strided to be 15 timesteps, and the final 100 are strided to be 20. + If the stride is a string starting with "ddim", then the fixed striding + from the DDIM paper is used, and only one section is allowed. + :param num_timesteps: the number of diffusion steps in the original + process to divide up. + :param section_counts: either a list of numbers, or a string containing + comma-separated numbers, indicating the step count + per section. As a special case, use "ddimN" where N + is a number of steps to use the striding from the + DDIM paper. + :return: a set of diffusion steps from the original process to use. + """ + if isinstance(section_counts, str): + if section_counts.startswith("ddim"): + desired_count = int(section_counts[len("ddim") :]) + for i in range(1, num_timesteps): + if len(range(0, num_timesteps, i)) == desired_count: + return set(range(0, num_timesteps, i)) + raise ValueError( + f"cannot create exactly {num_timesteps} steps with an integer stride" + ) + section_counts = [int(x) for x in section_counts.split(",")] + size_per = num_timesteps // len(section_counts) + extra = num_timesteps % len(section_counts) + start_idx = 0 + all_steps = [] + for i, section_count in enumerate(section_counts): + size = size_per + (1 if i < extra else 0) + if size < section_count: + raise ValueError( + f"cannot divide section of {size} steps into {section_count}" + ) + frac_stride = 1 if section_count <= 1 else (size - 1) / (section_count - 1) + cur_idx = 0.0 + taken_steps = [] + for _ in range(section_count): + taken_steps.append(start_idx + round(cur_idx)) + cur_idx += frac_stride + all_steps += taken_steps + start_idx += size + return set(all_steps) + + +class SpacedDiffusion(GaussianDiffusion): + """ + A diffusion process which can skip steps in a base diffusion process. + :param use_timesteps: a collection (sequence or set) of timesteps from the + original diffusion process to retain. + :param kwargs: the kwargs to create the base diffusion process. + """ + + def __init__(self, use_timesteps, **kwargs): + self.use_timesteps = set(use_timesteps) + self.timestep_map = [] + self.original_num_steps = len(kwargs["betas"]) + + base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa + last_alpha_cumprod = 1.0 + new_betas = [] + for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod): + if i in self.use_timesteps: + new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) + last_alpha_cumprod = alpha_cumprod + self.timestep_map.append(i) + kwargs["betas"] = np.array(new_betas) + super().__init__(**kwargs) + + def p_mean_variance( + self, model, *args, **kwargs + ): # pylint: disable=signature-differs + return super().p_mean_variance(self._wrap_model(model), *args, **kwargs) + + def training_losses( + self, model, *args, **kwargs + ): # pylint: disable=signature-differs + return super().training_losses(self._wrap_model(model), *args, **kwargs) + + def training_losses_diffusers( + self, model, *args, **kwargs + ): # pylint: disable=signature-differs + return super().training_losses_diffusers(self._wrap_model(model), *args, **kwargs) + + def condition_mean(self, cond_fn, *args, **kwargs): + return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs) + + def condition_score(self, cond_fn, *args, **kwargs): + return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs) + + def _wrap_model(self, model): + if isinstance(model, _WrappedModel): + return model + return _WrappedModel( + model, self.timestep_map, self.original_num_steps + ) + + def _scale_timesteps(self, t): + # Scaling is done by the wrapped model. + return t + + +class _WrappedModel: + def __init__(self, model, timestep_map, original_num_steps): + self.model = model + self.timestep_map = timestep_map + # self.rescale_timesteps = rescale_timesteps + self.original_num_steps = original_num_steps + + def __call__(self, x, timestep, **kwargs): + map_tensor = th.tensor(self.timestep_map, device=timestep.device, dtype=timestep.dtype) + new_ts = map_tensor[timestep] + # if self.rescale_timesteps: + # new_ts = new_ts.float() * (1000.0 / self.original_num_steps) + return self.model(x, timestep=new_ts, **kwargs) \ No newline at end of file diff --git a/ruyi/utils/utils.py b/ruyi/utils/utils.py new file mode 100644 index 0000000..825fbd4 --- /dev/null +++ b/ruyi/utils/utils.py @@ -0,0 +1,191 @@ +import os +import gc +import imageio +import numpy as np +import torch +import torchvision +import cv2 +from einops import rearrange +from PIL import Image + +def get_width_and_height_from_image_and_base_resolution(image, base_resolution): + target_pixels = int(base_resolution) * int(base_resolution) + original_width, original_height = Image.open(image).size + ratio = (target_pixels / (original_width * original_height)) ** 0.5 + width_slider = round(original_width * ratio) + height_slider = round(original_height * ratio) + return height_slider, width_slider + +def color_transfer(sc, dc): + """ + Transfer color distribution from of sc, referred to dc. + + Args: + sc (numpy.ndarray): input image to be transfered. + dc (numpy.ndarray): reference image + + Returns: + numpy.ndarray: Transferred color distribution on the sc. + """ + + def get_mean_and_std(img): + x_mean, x_std = cv2.meanStdDev(img) + x_mean = np.hstack(np.around(x_mean, 2)) + x_std = np.hstack(np.around(x_std, 2)) + return x_mean, x_std + + sc = cv2.cvtColor(sc, cv2.COLOR_RGB2LAB) + s_mean, s_std = get_mean_and_std(sc) + dc = cv2.cvtColor(dc, cv2.COLOR_RGB2LAB) + t_mean, t_std = get_mean_and_std(dc) + img_n = ((sc - s_mean) * (t_std / s_std)) + t_mean + np.putmask(img_n, img_n > 255, 255) + np.putmask(img_n, img_n < 0, 0) + dst = cv2.cvtColor(cv2.convertScaleAbs(img_n), cv2.COLOR_LAB2RGB) + return dst + +def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=12, imageio_backend=True, color_transfer_post_process=False): + videos = rearrange(videos, "b c t h w -> t b c h w") + outputs = [] + for x in videos: + x = torchvision.utils.make_grid(x, nrow=n_rows) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + if rescale: + x = (x + 1.0) / 2.0 # -1,1 -> 0,1 + x = (x * 255).numpy().astype(np.uint8) + outputs.append(Image.fromarray(x)) + + if color_transfer_post_process: + for i in range(1, len(outputs)): + outputs[i] = Image.fromarray(color_transfer(np.uint8(outputs[i]), np.uint8(outputs[0]))) + + os.makedirs(os.path.dirname(path), exist_ok=True) + if imageio_backend: + if path.endswith("mp4"): + imageio.mimsave(path, outputs, fps=fps) + else: + imageio.mimsave(path, outputs, duration=(1000 * 1/fps)) + else: + if path.endswith("mp4"): + path = path.replace('.mp4', '.gif') + outputs[0].save(path, format='GIF', append_images=outputs, save_all=True, duration=100, loop=0) + +def get_image_to_video_latent(validation_image_start, validation_image_end, video_length, sample_size): + if validation_image_start is not None and validation_image_end is not None: + if type(validation_image_start) is str and os.path.isfile(validation_image_start): + image_start = clip_image = Image.open(validation_image_start).convert("RGB") + image_start = image_start.resize([sample_size[1], sample_size[0]]) + clip_image = clip_image.resize([sample_size[1], sample_size[0]]) + else: + image_start = clip_image = validation_image_start + image_start = [_image_start.resize([sample_size[1], sample_size[0]]) for _image_start in image_start] + clip_image = [_clip_image.resize([sample_size[1], sample_size[0]]) for _clip_image in clip_image] + + if type(validation_image_end) is str and os.path.isfile(validation_image_end): + image_end = Image.open(validation_image_end).convert("RGB") + image_end = image_end.resize([sample_size[1], sample_size[0]]) + else: + image_end = validation_image_end + image_end = [_image_end.resize([sample_size[1], sample_size[0]]) for _image_end in image_end] + + if type(image_start) is list: + clip_image = clip_image[0] + start_video = torch.cat( + [torch.from_numpy(np.array(_image_start)).permute(2, 0, 1).unsqueeze(1).unsqueeze(0) for _image_start in image_start], + dim=2 + ) + input_video = torch.tile(start_video[:, :, :1], [1, 1, video_length, 1, 1]) + input_video[:, :, :len(image_start)] = start_video + + input_video_mask = torch.zeros_like(input_video[:, :1]) + input_video_mask[:, :, len(image_start):] = 255 + else: + input_video = torch.tile( + torch.from_numpy(np.array(image_start)).permute(2, 0, 1).unsqueeze(1).unsqueeze(0), + [1, 1, video_length, 1, 1] + ) + input_video_mask = torch.zeros_like(input_video[:, :1]) + input_video_mask[:, :, 3:] = 255 + + if type(image_end) is list: + image_end = [_image_end.resize(image_start[0].size if type(image_start) is list else image_start.size) for _image_end in image_end] + end_video = torch.cat( + [torch.from_numpy(np.array(_image_end)).permute(2, 0, 1).unsqueeze(1).unsqueeze(0) for _image_end in image_end], + dim=2 + ) + input_video[:, :, -len(end_video):] = end_video + + input_video_mask[:, :, -len(image_end):] = 0 + else: + image_end = image_end.resize(image_start[0].size if type(image_start) is list else image_start.size) + input_video[:, :, -3:] = torch.from_numpy(np.array(image_end)).permute(2, 0, 1).unsqueeze(1).unsqueeze(0) + input_video_mask[:, :, -3:] = 0 + + input_video = input_video / 255 + + elif validation_image_start is not None: + if type(validation_image_start) is str and os.path.isfile(validation_image_start): + image_start = clip_image = Image.open(validation_image_start).convert("RGB") + image_start = image_start.resize([sample_size[1], sample_size[0]]) + clip_image = clip_image.resize([sample_size[1], sample_size[0]]) + else: + image_start = clip_image = validation_image_start + # image_start = clip_image = [Image.open(_image_start).convert("RGB") for _image_start in validation_image_start] + image_start = [_image_start.resize([sample_size[1], sample_size[0]]) for _image_start in image_start] + clip_image = [_clip_image.resize([sample_size[1], sample_size[0]]) for _clip_image in clip_image] + image_end = None + + if type(image_start) is list: + clip_image = clip_image[0] + start_video = torch.cat( + [torch.from_numpy(np.array(_image_start)).permute(2, 0, 1).unsqueeze(1).unsqueeze(0) for _image_start in image_start], + dim=2 + ) + input_video = torch.tile(start_video[:, :, :1], [1, 1, video_length, 1, 1]) + input_video[:, :, :len(image_start)] = start_video + input_video = input_video / 255 + + input_video_mask = torch.zeros_like(input_video[:, :1]) + input_video_mask[:, :, len(image_start):] = 255 + else: + input_video = torch.tile( + torch.from_numpy(np.array(image_start)).permute(2, 0, 1).unsqueeze(1).unsqueeze(0), + [1, 1, video_length, 1, 1] + ) / 255 + input_video_mask = torch.zeros_like(input_video[:, :1]) + # input_video_mask[:, :, 3:, ] = 255 + input_video_mask[:, :, 1:, ] = 255 + else: + image_start = None + image_end = None + input_video = torch.zeros([1, 3, video_length, sample_size[0], sample_size[1]]) + input_video_mask = torch.ones([1, 1, video_length, sample_size[0], sample_size[1]]) * 255 + clip_image = None + + del image_start + del image_end + gc.collect() + + return input_video, input_video_mask, clip_image + +def get_video_to_video_latent(input_video_path, video_length, sample_size): + if type(input_video_path) is str: + cap = cv2.VideoCapture(input_video_path) + input_video = [] + while True: + ret, frame = cap.read() + if not ret: + break + frame = cv2.resize(frame, (sample_size[1], sample_size[0])) + input_video.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + cap.release() + else: + input_video = input_video_path + + input_video = torch.from_numpy(np.array(input_video))[:video_length] + input_video = input_video.permute([3, 0, 1, 2]).unsqueeze(0) / 255 + + input_video_mask = torch.zeros_like(input_video[:, :1]) + input_video_mask[:, :, :] = 255 + + return input_video, input_video_mask, None \ No newline at end of file diff --git a/ruyi/vae/LICENSE b/ruyi/vae/LICENSE new file mode 100644 index 0000000..0e609df --- /dev/null +++ b/ruyi/vae/LICENSE @@ -0,0 +1,82 @@ +Copyright (c) 2022 Robin Rombach and Patrick Esser and contributors + +CreativeML Open RAIL-M +dated August 22, 2022 + +Section I: PREAMBLE + +Multimodal generative models are being widely adopted and used, and have the potential to transform the way artists, among other individuals, conceive and benefit from AI or ML technologies as a tool for content creation. + +Notwithstanding the current and potential benefits that these artifacts can bring to society at large, there are also concerns about potential misuses of them, either due to their technical limitations or ethical considerations. + +In short, this license strives for both the open and responsible downstream use of the accompanying model. When it comes to the open character, we took inspiration from open source permissive licenses regarding the grant of IP rights. Referring to the downstream responsible use, we added use-based restrictions not permitting the use of the Model in very specific scenarios, in order for the licensor to be able to enforce the license in case potential misuses of the Model may occur. At the same time, we strive to promote open and responsible research on generative models for art and content generation. + +Even though downstream derivative versions of the model could be released under different licensing terms, the latter will always have to include - at minimum - the same use-based restrictions as the ones in the original license (this license). We believe in the intersection between open and responsible AI development; thus, this License aims to strike a balance between both in order to enable responsible open-science in the field of AI. + +This License governs the use of the model (and its derivatives) and is informed by the model card associated with the model. + +NOW THEREFORE, You and Licensor agree as follows: + +1. Definitions + +- "License" means the terms and conditions for use, reproduction, and Distribution as defined in this document. +- "Data" means a collection of information and/or content extracted from the dataset used with the Model, including to train, pretrain, or otherwise evaluate the Model. The Data is not licensed under this License. +- "Output" means the results of operating a Model as embodied in informational content resulting therefrom. +- "Model" means any accompanying machine-learning based assemblies (including checkpoints), consisting of learnt weights, parameters (including optimizer states), corresponding to the model architecture as embodied in the Complementary Material, that have been trained or tuned, in whole or in part on the Data, using the Complementary Material. +- "Derivatives of the Model" means all modifications to the Model, works based on the Model, or any other model which is created or initialized by transfer of patterns of the weights, parameters, activations or output of the Model, to the other model, in order to cause the other model to perform similarly to the Model, including - but not limited to - distillation methods entailing the use of intermediate data representations or methods based on the generation of synthetic data by the Model for training the other model. +- "Complementary Material" means the accompanying source code and scripts used to define, run, load, benchmark or evaluate the Model, and used to prepare data for training or evaluation, if any. This includes any accompanying documentation, tutorials, examples, etc, if any. +- "Distribution" means any transmission, reproduction, publication or other sharing of the Model or Derivatives of the Model to a third party, including providing the Model as a hosted service made available by electronic or other remote means - e.g. API-based or web access. +- "Licensor" means the copyright owner or entity authorized by the copyright owner that is granting the License, including the persons or entities that may have rights in the Model and/or distributing the Model. +- "You" (or "Your") means an individual or Legal Entity exercising permissions granted by this License and/or making use of the Model for whichever purpose and in any field of use, including usage of the Model in an end-use application - e.g. chatbot, translator, image generator. +- "Third Parties" means individuals or legal entities that are not under common control with Licensor or You. +- "Contribution" means any work of authorship, including the original version of the Model and any modifications or additions to that Model or Derivatives of the Model thereof, that is intentionally submitted to Licensor for inclusion in the Model by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Model, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." +- "Contributor" means Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Model. + +Section II: INTELLECTUAL PROPERTY RIGHTS + +Both copyright and patent grants apply to the Model, Derivatives of the Model and Complementary Material. The Model and Derivatives of the Model are subject to additional terms as described in Section III. + +2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare, publicly display, publicly perform, sublicense, and distribute the Complementary Material, the Model, and Derivatives of the Model. +3. Grant of Patent License. Subject to the terms and conditions of this License and where and as applicable, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this paragraph) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Model and the Complementary Material, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Model to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Model and/or Complementary Material or a Contribution incorporated within the Model and/or Complementary Material constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for the Model and/or Work shall terminate as of the date such litigation is asserted or filed. + +Section III: CONDITIONS OF USAGE, DISTRIBUTION AND REDISTRIBUTION + +4. Distribution and Redistribution. You may host for Third Party remote access purposes (e.g. software-as-a-service), reproduce and distribute copies of the Model or Derivatives of the Model thereof in any medium, with or without modifications, provided that You meet the following conditions: +Use-based restrictions as referenced in paragraph 5 MUST be included as an enforceable provision by You in any type of legal agreement (e.g. a license) governing the use and/or distribution of the Model or Derivatives of the Model, and You shall give notice to subsequent users You Distribute to, that the Model or Derivatives of the Model are subject to paragraph 5. This provision does not apply to the use of Complementary Material. +You must give any Third Party recipients of the Model or Derivatives of the Model a copy of this License; +You must cause any modified files to carry prominent notices stating that You changed the files; +You must retain all copyright, patent, trademark, and attribution notices excluding those notices that do not pertain to any part of the Model, Derivatives of the Model. +You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions - respecting paragraph 4.a. - for use, reproduction, or Distribution of Your modifications, or for any such Derivatives of the Model as a whole, provided Your use, reproduction, and Distribution of the Model otherwise complies with the conditions stated in this License. +5. Use-based restrictions. The restrictions set forth in Attachment A are considered Use-based restrictions. Therefore You cannot use the Model and the Derivatives of the Model for the specified restricted uses. You may use the Model subject to this License, including only for lawful purposes and in accordance with the License. Use may include creating any content with, finetuning, updating, running, training, evaluating and/or reparametrizing the Model. You shall require all of Your users who use the Model or a Derivative of the Model to comply with the terms of this paragraph (paragraph 5). +6. The Output You Generate. Except as set forth herein, Licensor claims no rights in the Output You generate using the Model. You are accountable for the Output you generate and its subsequent uses. No use of the output can contravene any provision as stated in the License. + +Section IV: OTHER PROVISIONS + +7. Updates and Runtime Restrictions. To the maximum extent permitted by law, Licensor reserves the right to restrict (remotely or otherwise) usage of the Model in violation of this License, update the Model through electronic means, or modify the Output of the Model based on updates. You shall undertake reasonable efforts to use the latest version of the Model. +8. Trademarks and related. Nothing in this License permits You to make use of Licensors’ trademarks, trade names, logos or to otherwise suggest endorsement or misrepresent the relationship between the parties; and any rights not expressly granted herein are reserved by the Licensors. +9. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Model and the Complementary Material (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Model, Derivatives of the Model, and the Complementary Material and assume any risks associated with Your exercise of permissions under this License. +10. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Model and the Complementary Material (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. +11. Accepting Warranty or Additional Liability. While redistributing the Model, Derivatives of the Model and the Complementary Material thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. +12. If any provision of this License is held to be invalid, illegal or unenforceable, the remaining provisions shall be unaffected thereby and remain valid as if such provision had not been set forth herein. + +END OF TERMS AND CONDITIONS + + + + +Attachment A + +Use Restrictions + +You agree not to use the Model or Derivatives of the Model: +- In any way that violates any applicable national, federal, state, local or international law or regulation; +- For the purpose of exploiting, harming or attempting to exploit or harm minors in any way; +- To generate or disseminate verifiably false information and/or content with the purpose of harming others; +- To generate or disseminate personal identifiable information that can be used to harm an individual; +- To defame, disparage or otherwise harass others; +- For fully automated decision making that adversely impacts an individual’s legal rights or otherwise creates or modifies a binding, enforceable obligation; +- For any use intended to or which has the effect of discriminating against or harming individuals or groups based on online or offline social behavior or known or predicted personal or personality characteristics; +- To exploit any of the vulnerabilities of a specific group of persons based on their age, social, physical or mental characteristics, in order to materially distort the behavior of a person pertaining to that group in a manner that causes or is likely to cause that person or another person physical or psychological harm; +- For any use intended to or which has the effect of discriminating against individuals or groups based on legally protected characteristics or categories; +- To provide medical advice and medical results interpretation; +- To generate or disseminate information for the purpose to be used for administration of justice, law enforcement, immigration or asylum processes, such as predicting an individual will commit fraud/crime commitment (e.g. by text profiling, drawing causal relationships between assertions made in documents, indiscriminate and arbitrarily-targeted use). diff --git a/ruyi/vae/README.md b/ruyi/vae/README.md new file mode 100644 index 0000000..d632712 --- /dev/null +++ b/ruyi/vae/README.md @@ -0,0 +1,63 @@ +## VAE Training + +English | [简体中文](./README_zh-CN.md) + +After completing data preprocessing, we can obtain the following dataset: + +``` +📦 project/ +├── 📂 datasets/ +│ ├── 📂 internal_datasets/ +│ ├── 📂 videos/ +│ │ ├── 📄 00000001.mp4 +│ │ ├── 📄 00000001.jpg +│ │ └── 📄 ..... +│ └── 📄 json_of_internal_datasets.json +``` + +The json_of_internal_datasets.json is a standard JSON file. The file_path in the json can to be set as relative path, as shown in below: +```json +[ + { + "file_path": "videos/00000001.mp4", + "text": "A group of young men in suits and sunglasses are walking down a city street.", + "type": "video" + }, + { + "file_path": "train/00000001.jpg", + "text": "A group of young men in suits and sunglasses are walking down a city street.", + "type": "image" + }, + ..... +] +``` + +You can also set the path as absolute path as follow: +```json +[ + { + "file_path": "/mnt/data/videos/00000001.mp4", + "text": "A group of young men in suits and sunglasses are walking down a city street.", + "type": "video" + }, + { + "file_path": "/mnt/data/train/00000001.jpg", + "text": "A group of young men in suits and sunglasses are walking down a city street.", + "type": "image" + }, + ..... +] +``` + +## Train Video VAE +We need to set config in ```easyanimate/vae/configs/autoencoder``` at first. The default config is ```autoencoder_kl_32x32x4_slice.yaml```. We need to set the some params in yaml file. + +- ```data_json_path``` corresponds to the JSON file of the dataset. +- ```data_root``` corresponds to the root path of the dataset. If you want to use absolute path in json file, please delete this line. +- ```ckpt_path``` corresponds to the pretrained weights of the vae. +- ```gpus``` and num_nodes need to be set as the actual situation of your machine. + +The we run shell file as follow: +``` +sh scripts/train_vae.sh +``` \ No newline at end of file diff --git a/ruyi/vae/README_zh-CN.md b/ruyi/vae/README_zh-CN.md new file mode 100644 index 0000000..6bc0467 --- /dev/null +++ b/ruyi/vae/README_zh-CN.md @@ -0,0 +1,63 @@ +## VAE 训练 + +[English](./README.md) | 简体中文 + +在完成数据预处理后，你可以获得这样的数据格式: + +``` +📦 project/ +├── 📂 datasets/ +│ ├── 📂 internal_datasets/ +│ ├── 📂 videos/ +│ │ ├── 📄 00000001.mp4 +│ │ ├── 📄 00000001.jpg +│ │ └── 📄 ..... +│ └── 📄 json_of_internal_datasets.json +``` + +json_of_internal_datasets.json是一个标准的json文件。json中的file_path可以被设置为相对路径，如下所示： +```json +[ + { + "file_path": "videos/00000001.mp4", + "text": "A group of young men in suits and sunglasses are walking down a city street.", + "type": "video" + }, + { + "file_path": "train/00000001.jpg", + "text": "A group of young men in suits and sunglasses are walking down a city street.", + "type": "image" + }, + ..... +] +``` + +你也可以将路径设置为绝对路径： +```json +[ + { + "file_path": "/mnt/data/videos/00000001.mp4", + "text": "A group of young men in suits and sunglasses are walking down a city street.", + "type": "video" + }, + { + "file_path": "/mnt/data/train/00000001.jpg", + "text": "A group of young men in suits and sunglasses are walking down a city street.", + "type": "image" + }, + ..... +] +``` + +## 训练 Video VAE +我们首先需要修改 ```easyanimate/vae/configs/autoencoder``` 中的配置文件。默认的配置文件是 ```autoencoder_kl_32x32x4_slice.yaml```。你需要修改以下参数： + +- ```data_json_path``` json file 所在的目录。 +- ```data_root``` 数据的根目录。如果你在json file中使用了绝对路径，请设置为空。 +- ```ckpt_path``` 预训练的vae模型路径。 +- ```gpus``` 以及 ```num_nodes``` 需要设置为你机器的实际gpu数目。 + +运行以下的脚本来训练vae: +``` +sh scripts/train_vae.sh +``` \ No newline at end of file diff --git a/ruyi/vae/configs/autoencoder/autoencoder_kl_32x32x4_mag.yaml b/ruyi/vae/configs/autoencoder/autoencoder_kl_32x32x4_mag.yaml new file mode 100644 index 0000000..cfd0e0a --- /dev/null +++ b/ruyi/vae/configs/autoencoder/autoencoder_kl_32x32x4_mag.yaml @@ -0,0 +1,62 @@ +model: + base_learning_rate: 1.0e-04 + target: easyanimate.vae.ldm.models.omnigen_casual3dcnn.AutoencoderKLMagvit_fromOmnigen + params: + monitor: train/rec_loss + ckpt_path: models/videoVAE_omnigen_8x8x4_from_vae-ft-mse-840000-ema-pruned.ckpt + down_block_types: ("SpatialDownBlock3D", "SpatialTemporalDownBlock3D", "SpatialTemporalDownBlock3D", + "SpatialTemporalDownBlock3D",) + up_block_types: ("SpatialUpBlock3D", "SpatialTemporalUpBlock3D", "SpatialTemporalUpBlock3D", + "SpatialTemporalUpBlock3D",) + lossconfig: + target: easyanimate.vae.ldm.modules.losses.LPIPSWithDiscriminator + params: + disc_start: 50001 + kl_weight: 1.0e-06 + disc_weight: 0.5 + l2_loss_weight: 0.1 + l1_loss_weight: 1.0 + perceptual_weight: 1.0 + +data: + target: train_vae.DataModuleFromConfig + + params: + batch_size: 2 + wrap: true + num_workers: 4 + train: + target: easyanimate.vae.ldm.data.dataset_image_video.CustomSRTrain + params: + data_json_path: pretrain.json + data_root: /your_data_root # This is used in relative path + size: 128 + degradation: pil_nearest + video_size: 128 + video_len: 9 + slice_interval: 1 + validation: + target: easyanimate.vae.ldm.data.dataset_image_video.CustomSRValidation + params: + data_json_path: pretrain.json + data_root: /your_data_root # This is used in relative path + size: 128 + degradation: pil_nearest + video_size: 128 + video_len: 9 + slice_interval: 1 + +lightning: + callbacks: + image_logger: + target: train_vae.ImageLogger + params: + batch_frequency: 5000 + max_images: 8 + increase_log_steps: True + + trainer: + benchmark: True + accumulate_grad_batches: 1 + gpus: "0" + num_nodes: 1 \ No newline at end of file diff --git a/ruyi/vae/configs/autoencoder/autoencoder_kl_32x32x4_slice.yaml b/ruyi/vae/configs/autoencoder/autoencoder_kl_32x32x4_slice.yaml new file mode 100644 index 0000000..8fe607c --- /dev/null +++ b/ruyi/vae/configs/autoencoder/autoencoder_kl_32x32x4_slice.yaml @@ -0,0 +1,65 @@ +model: + base_learning_rate: 1.0e-04 + target: easyanimate.vae.ldm.models.omnigen_casual3dcnn.AutoencoderKLMagvit_fromOmnigen + params: + slice_compression_vae: true + mini_batch_encoder: 8 + mini_batch_decoder: 2 + monitor: train/rec_loss + ckpt_path: models/Diffusion_Transformer/EasyAnimateV2-XL-2-512x512/vae/diffusion_pytorch_model.safetensors + down_block_types: ("SpatialDownBlock3D", "SpatialTemporalDownBlock3D", "SpatialTemporalDownBlock3D", + "SpatialTemporalDownBlock3D",) + up_block_types: ("SpatialUpBlock3D", "SpatialTemporalUpBlock3D", "SpatialTemporalUpBlock3D", + "SpatialTemporalUpBlock3D",) + lossconfig: + target: easyanimate.vae.ldm.modules.losses.LPIPSWithDiscriminator + params: + disc_start: 50001 + kl_weight: 1.0e-06 + disc_weight: 0.5 + l2_loss_weight: 0.0 + l1_loss_weight: 1.0 + perceptual_weight: 1.0 + +data: + target: train_vae.DataModuleFromConfig + + params: + batch_size: 1 + wrap: true + num_workers: 8 + train: + target: easyanimate.vae.ldm.data.dataset_image_video.CustomSRTrain + params: + data_json_path: pretrain.json + data_root: /your_data_root # This is used in relative path + size: 256 + degradation: pil_nearest + video_size: 256 + video_len: 25 + slice_interval: 1 + validation: + target: easyanimate.vae.ldm.data.dataset_image_video.CustomSRValidation + params: + data_json_path: pretrain.json + data_root: /your_data_root # This is used in relative path + size: 256 + degradation: pil_nearest + video_size: 256 + video_len: 25 + slice_interval: 1 + +lightning: + callbacks: + image_logger: + target: train_vae.ImageLogger + params: + batch_frequency: 5000 + max_images: 8 + increase_log_steps: True + + trainer: + benchmark: True + accumulate_grad_batches: 1 + gpus: "0" + num_nodes: 1 \ No newline at end of file diff --git a/ruyi/vae/configs/autoencoder/autoencoder_kl_32x32x4_slice_decoder_only.yaml b/ruyi/vae/configs/autoencoder/autoencoder_kl_32x32x4_slice_decoder_only.yaml new file mode 100644 index 0000000..32fa32b --- /dev/null +++ b/ruyi/vae/configs/autoencoder/autoencoder_kl_32x32x4_slice_decoder_only.yaml @@ -0,0 +1,66 @@ +model: + base_learning_rate: 1.0e-04 + target: easyanimate.vae.ldm.models.omnigen_casual3dcnn.AutoencoderKLMagvit_fromOmnigen + params: + slice_compression_vae: true + train_decoder_only: true + mini_batch_encoder: 8 + mini_batch_decoder: 2 + monitor: train/rec_loss + ckpt_path: models/Diffusion_Transformer/EasyAnimateV2-XL-2-512x512/vae/diffusion_pytorch_model.safetensors + down_block_types: ("SpatialDownBlock3D", "SpatialTemporalDownBlock3D", "SpatialTemporalDownBlock3D", + "SpatialTemporalDownBlock3D",) + up_block_types: ("SpatialUpBlock3D", "SpatialTemporalUpBlock3D", "SpatialTemporalUpBlock3D", + "SpatialTemporalUpBlock3D",) + lossconfig: + target: easyanimate.vae.ldm.modules.losses.LPIPSWithDiscriminator + params: + disc_start: 50001 + kl_weight: 1.0e-06 + disc_weight: 0.5 + l2_loss_weight: 1.0 + l1_loss_weight: 0.0 + perceptual_weight: 1.0 + +data: + target: train_vae.DataModuleFromConfig + + params: + batch_size: 1 + wrap: true + num_workers: 8 + train: + target: easyanimate.vae.ldm.data.dataset_image_video.CustomSRTrain + params: + data_json_path: pretrain.json + data_root: /your_data_root # This is used in relative path + size: 256 + degradation: pil_nearest + video_size: 256 + video_len: 25 + slice_interval: 1 + validation: + target: easyanimate.vae.ldm.data.dataset_image_video.CustomSRValidation + params: + data_json_path: pretrain.json + data_root: /your_data_root # This is used in relative path + size: 256 + degradation: pil_nearest + video_size: 256 + video_len: 25 + slice_interval: 1 + +lightning: + callbacks: + image_logger: + target: train_vae.ImageLogger + params: + batch_frequency: 5000 + max_images: 8 + increase_log_steps: True + + trainer: + benchmark: True + accumulate_grad_batches: 1 + gpus: "0" + num_nodes: 1 \ No newline at end of file diff --git a/ruyi/vae/configs/autoencoder/autoencoder_kl_32x32x4_slice_t_downsample_8.yaml b/ruyi/vae/configs/autoencoder/autoencoder_kl_32x32x4_slice_t_downsample_8.yaml new file mode 100644 index 0000000..25f05f2 --- /dev/null +++ b/ruyi/vae/configs/autoencoder/autoencoder_kl_32x32x4_slice_t_downsample_8.yaml @@ -0,0 +1,66 @@ +model: + base_learning_rate: 1.0e-04 + target: easyanimate.vae.ldm.models.omnigen_casual3dcnn.AutoencoderKLMagvit_fromOmnigen + params: + slice_compression_vae: true + mini_batch_encoder: 8 + mini_batch_decoder: 1 + monitor: train/rec_loss + ckpt_path: models/Diffusion_Transformer/EasyAnimateV2-XL-2-512x512/vae/diffusion_pytorch_model.safetensors + down_block_types: ("SpatialTemporalDownBlock3D", "SpatialTemporalDownBlock3D", "SpatialTemporalDownBlock3D", + "SpatialTemporalDownBlock3D",) + up_block_types: ("SpatialTemporalUpBlock3D", "SpatialTemporalUpBlock3D", "SpatialTemporalUpBlock3D", + "SpatialTemporalUpBlock3D",) + lossconfig: + target: easyanimate.vae.ldm.modules.losses.LPIPSWithDiscriminator + params: + disc_start: 50001 + kl_weight: 1.0e-06 + disc_weight: 0.5 + l2_loss_weight: 0.0 + l1_loss_weight: 1.0 + perceptual_weight: 1.0 + + +data: + target: train_vae.DataModuleFromConfig + + params: + batch_size: 1 + wrap: true + num_workers: 8 + train: + target: easyanimate.vae.ldm.data.dataset_image_video.CustomSRTrain + params: + data_json_path: pretrain.json + data_root: /your_data_root # This is used in relative path + size: 256 + degradation: pil_nearest + video_size: 256 + video_len: 33 + slice_interval: 1 + validation: + target: easyanimate.vae.ldm.data.dataset_image_video.CustomSRValidation + params: + data_json_path: pretrain.json + data_root: /your_data_root # This is used in relative path + size: 256 + degradation: pil_nearest + video_size: 256 + video_len: 33 + slice_interval: 1 + +lightning: + callbacks: + image_logger: + target: train_vae.ImageLogger + params: + batch_frequency: 5000 + max_images: 8 + increase_log_steps: True + + trainer: + benchmark: True + accumulate_grad_batches: 1 + gpus: "0" + num_nodes: 1 \ No newline at end of file diff --git a/ruyi/vae/environment.yaml b/ruyi/vae/environment.yaml new file mode 100644 index 0000000..6279f23 --- /dev/null +++ b/ruyi/vae/environment.yaml @@ -0,0 +1,29 @@ +name: ldm +channels: + - pytorch + - defaults +dependencies: + - python=3.8.5 + - pip=20.3 + - cudatoolkit=11.3 + - pytorch=1.11.0 + - torchvision=0.12.0 + - numpy=1.19.2 + - pip: + - albumentations==0.4.3 + - diffusers + - opencv-python==4.1.2.30 + - pudb==2019.2 + - invisible-watermark + - imageio==2.9.0 + - imageio-ffmpeg==0.4.2 + - pytorch-lightning==1.4.2 + - omegaconf==2.1.1 + - test-tube>=0.7.5 + - streamlit>=0.73.1 + - einops==0.3.0 + - torch-fidelity==0.3.0 + - transformers==4.19.2 + - torchmetrics==0.6.0 + - kornia==0.6 + - -e . diff --git a/ruyi/vae/ldm/data/__init__.py b/ruyi/vae/ldm/data/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ruyi/vae/ldm/data/base.py b/ruyi/vae/ldm/data/base.py new file mode 100644 index 0000000..7ba1233 --- /dev/null +++ b/ruyi/vae/ldm/data/base.py @@ -0,0 +1,25 @@ +from abc import abstractmethod + +from torch.utils.data import (ChainDataset, ConcatDataset, Dataset, + IterableDataset) + + +class Txt2ImgIterableBaseDataset(IterableDataset): + ''' + Define an interface to make the IterableDatasets for text2img data chainable + ''' + def __init__(self, num_records=0, valid_ids=None, size=256): + super().__init__() + self.num_records = num_records + self.valid_ids = valid_ids + self.sample_ids = valid_ids + self.size = size + + print(f'{self.__class__.__name__} dataset contains {self.__len__()} examples.') + + def __len__(self): + return self.num_records + + @abstractmethod + def __iter__(self): + pass \ No newline at end of file diff --git a/ruyi/vae/ldm/data/dataset_callback.py b/ruyi/vae/ldm/data/dataset_callback.py new file mode 100644 index 0000000..d0d5943 --- /dev/null +++ b/ruyi/vae/ldm/data/dataset_callback.py @@ -0,0 +1,25 @@ +#-*- encoding:utf-8 -*- +from pytorch_lightning.callbacks import Callback + +class DatasetCallback(Callback): + def __init__(self): + self.sampler_pos_start = 0 + self.preload_used_idx_flag = False + + def on_train_start(self, trainer, pl_module): + if not self.preload_used_idx_flag: + self.preload_used_idx_flag = True + trainer.train_dataloader.batch_sampler.sampler_pos_reload = self.sampler_pos_start + + def on_save_checkpoint(self, trainer, pl_module, checkpoint): + if trainer.train_dataloader is not None: + # Save sampler_pos_start parameters in the checkpoint + checkpoint['sampler_pos_start'] = trainer.train_dataloader.batch_sampler.sampler_pos_start + + def on_load_checkpoint(self, trainer, pl_module, checkpoint): + # Restore sampler_pos_start parameters from the checkpoint + if 'sampler_pos_start' in checkpoint: + self.sampler_pos_start = checkpoint.get('sampler_pos_start', 0) + print('Load sampler_pos_start from checkpoint, sampler_pos_start = %d' % self.sampler_pos_start) + else: + print('The sampler_pos_start is not in checkpoint') \ No newline at end of file diff --git a/ruyi/vae/ldm/data/dataset_image_video.py b/ruyi/vae/ldm/data/dataset_image_video.py new file mode 100644 index 0000000..bd0c4f6 --- /dev/null +++ b/ruyi/vae/ldm/data/dataset_image_video.py @@ -0,0 +1,281 @@ +import glob +import json +import os +import pickle +import random +import shutil +import tarfile +from functools import partial + +import albumentations +import cv2 +import numpy as np +import PIL +import torchvision.transforms.functional as TF +import yaml +from decord import VideoReader +from func_timeout import FunctionTimedOut, func_set_timeout +from omegaconf import OmegaConf +from PIL import Image +from torch.utils.data import (BatchSampler, Dataset, Sampler) +from tqdm import tqdm + +from ..modules.image_degradation import (degradation_fn_bsr, + degradation_fn_bsr_light) + + +class ImageVideoSampler(BatchSampler): + """A sampler wrapper for grouping images with similar aspect ratio into a same batch. + + Args: + sampler (Sampler): Base sampler. + dataset (Dataset): Dataset providing data information. + batch_size (int): Size of mini-batch. + drop_last (bool): If ``True``, the sampler will drop the last batch if + its size would be less than ``batch_size``. + aspect_ratios (dict): The predefined aspect ratios. + """ + + def __init__(self, + sampler: Sampler, + dataset: Dataset, + batch_size: int, + drop_last: bool = False + ) -> None: + if not isinstance(sampler, Sampler): + raise TypeError('sampler should be an instance of ``Sampler``, ' + f'but got {sampler}') + if not isinstance(batch_size, int) or batch_size <= 0: + raise ValueError('batch_size should be a positive integer value, ' + f'but got batch_size={batch_size}') + self.sampler = sampler + self.dataset = dataset + self.batch_size = batch_size + self.drop_last = drop_last + + self.sampler_pos_start = 0 + self.sampler_pos_reload = 0 + + self.num_samples_random = len(self.sampler) + # buckets for each aspect ratio + self.bucket = {'image':[], 'video':[]} + + def set_epoch(self, epoch): + if hasattr(self.sampler, "set_epoch"): + self.sampler.set_epoch(epoch) + + def __iter__(self): + for index_sampler, idx in enumerate(self.sampler): + if self.sampler_pos_reload != 0 and self.sampler_pos_reload < self.num_samples_random: + if index_sampler < self.sampler_pos_reload: + self.sampler_pos_start = (self.sampler_pos_start + 1) % self.num_samples_random + continue + elif index_sampler == self.sampler_pos_reload: + self.sampler_pos_reload = 0 + + content_type = self.dataset.data.get_type(idx) + bucket = self.bucket[content_type] + bucket.append(idx) + # yield a batch of indices in the same aspect ratio group + if len(self.bucket['video']) == self.batch_size: + yield self.bucket['video'] + self.bucket['video'] = [] + elif len(self.bucket['image']) == self.batch_size: + yield self.bucket['image'] + self.bucket['image'] = [] + self.sampler_pos_start = (self.sampler_pos_start + 1) % self.num_samples_random + +class ImageVideoDataset(Dataset): + # update __getitem__() from ImageNetSR. If timeout for Pandas70M, throw exception. + # If caught exception(timeout or others), try another index until successful and return. + def __init__(self, size=None, video_size=128, video_len=25, + degradation=None, downscale_f=4, random_crop=True, min_crop_f=0.25, max_crop_f=1., + s_t=None, slice_interval=None, data_root=None + ): + """ + Imagenet Superresolution Dataloader + Performs following ops in order: + 1. crops a crop of size s from image either as random or center crop + 2. resizes crop to size with cv2.area_interpolation + 3. degrades resized crop with degradation_fn + + :param size: resizing to size after cropping + :param degradation: degradation_fn, e.g. cv_bicubic or bsrgan_light + :param downscale_f: Low Resolution Downsample factor + :param min_crop_f: determines crop size s, + where s = c * min_img_side_len with c sampled from interval (min_crop_f, max_crop_f) + :param max_crop_f: "" + :param data_root: + :param random_crop: + """ + self.base = self.get_base() + assert size + assert (size / downscale_f).is_integer() + self.size = size + self.LR_size = int(size / downscale_f) + self.min_crop_f = min_crop_f + self.max_crop_f = max_crop_f + assert(max_crop_f <= 1.) + self.center_crop = not random_crop + self.s_t = s_t + self.slice_interval = slice_interval + + self.image_rescaler = albumentations.SmallestMaxSize(max_size=size, interpolation=cv2.INTER_AREA) + self.video_rescaler = albumentations.SmallestMaxSize(max_size=video_size, interpolation=cv2.INTER_AREA) + self.video_len = video_len + self.video_size = video_size + self.data_root = data_root + + self.pil_interpolation = False # gets reset later if incase interp_op is from pillow + + if degradation == "bsrgan": + self.degradation_process = partial(degradation_fn_bsr, sf=downscale_f) + + elif degradation == "bsrgan_light": + self.degradation_process = partial(degradation_fn_bsr_light, sf=downscale_f) + else: + interpolation_fn = { + "cv_nearest": cv2.INTER_NEAREST, + "cv_bilinear": cv2.INTER_LINEAR, + "cv_bicubic": cv2.INTER_CUBIC, + "cv_area": cv2.INTER_AREA, + "cv_lanczos": cv2.INTER_LANCZOS4, + "pil_nearest": PIL.Image.NEAREST, + "pil_bilinear": PIL.Image.BILINEAR, + "pil_bicubic": PIL.Image.BICUBIC, + "pil_box": PIL.Image.BOX, + "pil_hamming": PIL.Image.HAMMING, + "pil_lanczos": PIL.Image.LANCZOS, + }[degradation] + + self.pil_interpolation = degradation.startswith("pil_") + + if self.pil_interpolation: + self.degradation_process = partial(TF.resize, size=self.LR_size, interpolation=interpolation_fn) + + else: + self.degradation_process = albumentations.SmallestMaxSize(max_size=self.LR_size, + interpolation=interpolation_fn) + + def __len__(self): + return len(self.base) + + def get_type(self, index): + return self.base[index].get('type', 'image') + + def __getitem__(self, i): + @func_set_timeout(3) # time wait 3 seconds + def get_video_item(example): + if self.data_root is not None: + video_reader = VideoReader(os.path.join(self.data_root, example['file_path'])) + else: + video_reader = VideoReader(example['file_path']) + video_length = len(video_reader) + + clip_length = min(video_length, (self.video_len - 1) * self.slice_interval + 1) + start_idx = random.randint(0, video_length - clip_length) + batch_index = np.linspace(start_idx, start_idx + clip_length - 1, self.video_len, dtype=int) + + pixel_values = video_reader.get_batch(batch_index).asnumpy() + + del video_reader + out_images = [] + LR_out_images = [] + min_side_len = min(pixel_values[0].shape[:2]) + + crop_side_len = min_side_len * np.random.uniform(self.min_crop_f, self.max_crop_f, size=None) + crop_side_len = int(crop_side_len) + if self.center_crop: + self.cropper = albumentations.CenterCrop(height=crop_side_len, width=crop_side_len) + else: + self.cropper = albumentations.RandomCrop(height=crop_side_len, width=crop_side_len) + + imgs = np.transpose(pixel_values, (1, 2, 3, 0)) + imgs = self.cropper(image=imgs)["image"] + imgs = np.transpose(imgs, (3, 0, 1, 2)) + for img in imgs: + image = self.video_rescaler(image=img)["image"] + out_images.append(image[None, :, :, :]) + if self.pil_interpolation: + image_pil = PIL.Image.fromarray(image) + LR_image = self.degradation_process(image_pil) + LR_image = np.array(LR_image).astype(np.uint8) + else: + LR_image = self.degradation_process(image=image)["image"] + LR_out_images.append(LR_image[None, :, :, :]) + + example = {} + example['image'] = (np.concatenate(out_images) / 127.5 - 1.0).astype(np.float32) + example['LR_image'] = (np.concatenate(LR_out_images) / 127.5 - 1.0).astype(np.float32) + return example + + example = self.base[i] + if example.get('type', 'image') == 'video': + while True: + try: + example = self.base[i] + return get_video_item(example) + except FunctionTimedOut: + print("stt catch: Function 'extract failed' timed out.") + i = random.randint(0, self.__len__() - 1) + except Exception as e: + print('stt catch', e) + i = random.randint(0, self.__len__() - 1) + elif example.get('type', 'image') == 'image': + while True: + try: + example = self.base[i] + if self.data_root is not None: + image = Image.open(os.path.join(self.data_root, example['file_path'])) + else: + image = Image.open(example['file_path']) + image = image.convert("RGB") + image = np.array(image).astype(np.uint8) + + min_side_len = min(image.shape[:2]) + crop_side_len = min_side_len * np.random.uniform(self.min_crop_f, self.max_crop_f, size=None) + crop_side_len = int(crop_side_len) + + if self.center_crop: + self.cropper = albumentations.CenterCrop(height=crop_side_len, width=crop_side_len) + + else: + self.cropper = albumentations.RandomCrop(height=crop_side_len, width=crop_side_len) + + image = self.cropper(image=image)["image"] + + image = self.image_rescaler(image=image)["image"] + + if self.pil_interpolation: + image_pil = PIL.Image.fromarray(image) + LR_image = self.degradation_process(image_pil) + LR_image = np.array(LR_image).astype(np.uint8) + + else: + LR_image = self.degradation_process(image=image)["image"] + + example = {} + example["image"] = (image/127.5 - 1.0).astype(np.float32) + example["LR_image"] = (LR_image/127.5 - 1.0).astype(np.float32) + return example + except Exception as e: + print("catch", e) + i = random.randint(0, self.__len__() - 1) + +class CustomSRTrain(ImageVideoDataset): + def __init__(self, data_json_path, **kwargs): + self.data_json_path = data_json_path + super().__init__(**kwargs) + + def get_base(self): + return [ann for ann in json.load(open(self.data_json_path))] + +class CustomSRValidation(ImageVideoDataset): + def __init__(self, data_json_path, **kwargs): + self.data_json_path = data_json_path + super().__init__(**kwargs) + self.data_json_path = data_json_path + + def get_base(self): + return [ann for ann in json.load(open(self.data_json_path))][:100] + \ + [ann for ann in json.load(open(self.data_json_path))][-100:] diff --git a/ruyi/vae/ldm/lr_scheduler.py b/ruyi/vae/ldm/lr_scheduler.py new file mode 100644 index 0000000..be39da9 --- /dev/null +++ b/ruyi/vae/ldm/lr_scheduler.py @@ -0,0 +1,98 @@ +import numpy as np + + +class LambdaWarmUpCosineScheduler: + """ + note: use with a base_lr of 1.0 + """ + def __init__(self, warm_up_steps, lr_min, lr_max, lr_start, max_decay_steps, verbosity_interval=0): + self.lr_warm_up_steps = warm_up_steps + self.lr_start = lr_start + self.lr_min = lr_min + self.lr_max = lr_max + self.lr_max_decay_steps = max_decay_steps + self.last_lr = 0. + self.verbosity_interval = verbosity_interval + + def schedule(self, n, **kwargs): + if self.verbosity_interval > 0: + if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_lr}") + if n < self.lr_warm_up_steps: + lr = (self.lr_max - self.lr_start) / self.lr_warm_up_steps * n + self.lr_start + self.last_lr = lr + return lr + else: + t = (n - self.lr_warm_up_steps) / (self.lr_max_decay_steps - self.lr_warm_up_steps) + t = min(t, 1.0) + lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * ( + 1 + np.cos(t * np.pi)) + self.last_lr = lr + return lr + + def __call__(self, n, **kwargs): + return self.schedule(n,**kwargs) + + +class LambdaWarmUpCosineScheduler2: + """ + supports repeated iterations, configurable via lists + note: use with a base_lr of 1.0. + """ + def __init__(self, warm_up_steps, f_min, f_max, f_start, cycle_lengths, verbosity_interval=0): + assert len(warm_up_steps) == len(f_min) == len(f_max) == len(f_start) == len(cycle_lengths) + self.lr_warm_up_steps = warm_up_steps + self.f_start = f_start + self.f_min = f_min + self.f_max = f_max + self.cycle_lengths = cycle_lengths + self.cum_cycles = np.cumsum([0] + list(self.cycle_lengths)) + self.last_f = 0. + self.verbosity_interval = verbosity_interval + + def find_in_interval(self, n): + interval = 0 + for cl in self.cum_cycles[1:]: + if n <= cl: + return interval + interval += 1 + + def schedule(self, n, **kwargs): + cycle = self.find_in_interval(n) + n = n - self.cum_cycles[cycle] + if self.verbosity_interval > 0: + if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_f}, " + f"current cycle {cycle}") + if n < self.lr_warm_up_steps[cycle]: + f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle] + self.last_f = f + return f + else: + t = (n - self.lr_warm_up_steps[cycle]) / (self.cycle_lengths[cycle] - self.lr_warm_up_steps[cycle]) + t = min(t, 1.0) + f = self.f_min[cycle] + 0.5 * (self.f_max[cycle] - self.f_min[cycle]) * ( + 1 + np.cos(t * np.pi)) + self.last_f = f + return f + + def __call__(self, n, **kwargs): + return self.schedule(n, **kwargs) + + +class LambdaLinearScheduler(LambdaWarmUpCosineScheduler2): + + def schedule(self, n, **kwargs): + cycle = self.find_in_interval(n) + n = n - self.cum_cycles[cycle] + if self.verbosity_interval > 0: + if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_f}, " + f"current cycle {cycle}") + + if n < self.lr_warm_up_steps[cycle]: + f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle] + self.last_f = f + return f + else: + f = self.f_min[cycle] + (self.f_max[cycle] - self.f_min[cycle]) * (self.cycle_lengths[cycle] - n) / (self.cycle_lengths[cycle]) + self.last_f = f + return f + diff --git a/ruyi/vae/ldm/models/__init__.py b/ruyi/vae/ldm/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ruyi/vae/ldm/models/autoencoder.py b/ruyi/vae/ldm/models/autoencoder.py new file mode 100644 index 0000000..7cc28ae --- /dev/null +++ b/ruyi/vae/ldm/models/autoencoder.py @@ -0,0 +1,337 @@ +import time +from contextlib import contextmanager + +import pytorch_lightning as pl +import torch +import torch.nn.functional as F + +from ..modules.diffusionmodules.model import Decoder, Encoder +from ..modules.distributions.distributions import DiagonalGaussianDistribution +from ..util import instantiate_from_config +from .enc_dec_pytorch import Decoder as Mag_Decoder +from .enc_dec_pytorch import Encoder as Mag_Encoder + + +class AutoencoderKLMagvit(pl.LightningModule): + def __init__(self, + ddconfig, + lossconfig, + embed_dim, + ckpt_path=None, + ignore_keys=[], + image_key="image", + colorize_nlabels=None, + monitor=None, + ): + super().__init__() + self.image_key = image_key + self.encoder = Mag_Encoder() + self.decoder = Mag_Decoder() + self.loss = instantiate_from_config(lossconfig) + self.quant_conv = torch.nn.Conv3d(16, 16, 1) + self.post_quant_conv = torch.nn.Conv3d(8, 8, 1) + self.embed_dim = embed_dim + if colorize_nlabels is not None: + assert type(colorize_nlabels)==int + self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1)) + if monitor is not None: + self.monitor = monitor + if ckpt_path is not None: + self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys) + + def init_from_ckpt(self, path, ignore_keys=list()): + sd = torch.load(path, map_location="cpu")["state_dict"] + keys = list(sd.keys()) + for k in keys: + for ik in ignore_keys: + if k.startswith(ik): + print("Deleting key {} from state_dict.".format(k)) + del sd[k] + self.load_state_dict(sd, strict=False) + print(f"Restored from {path}") + + def encode(self, x): + h = self.encoder(x) + moments = self.quant_conv(h) + posterior = DiagonalGaussianDistribution(moments) + return posterior + + def decode(self, z): + z = self.post_quant_conv(z) + dec = self.decoder(z) + return dec + + def forward(self, input, sample_posterior=True): + if input.ndim==4: + input = input.unsqueeze(2) + posterior = self.encode(input) + if sample_posterior: + z = posterior.sample() + else: + z = posterior.mode() + dec = self.decode(z) + return dec, posterior + + def get_input(self, batch, k): + x = batch[k] + if x.ndim==5: + x = x.permute(0, 4, 1, 2, 3).to(memory_format=torch.contiguous_format).float() + return x + if len(x.shape) == 3: + x = x[..., None] + x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float() + return x + + def training_step(self, batch, batch_idx, optimizer_idx): + # tic = time.time() + inputs = self.get_input(batch, self.image_key) + # print(f"get_input time {time.time() - tic}") + # tic = time.time() + reconstructions, posterior = self(inputs) + # print(f"model forward time {time.time() - tic}") + + if optimizer_idx == 0: + # train encoder+decoder+logvar + aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train") + self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True) + self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False) + # print(f"cal loss time {time.time() - tic}") + return aeloss + + if optimizer_idx == 1: + # train the discriminator + discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train") + + self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True) + self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False) + # print(f"cal loss time {time.time() - tic}") + return discloss + + def validation_step(self, batch, batch_idx): + with torch.no_grad(): + inputs = self.get_input(batch, self.image_key) + reconstructions, posterior = self(inputs) + aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step, + last_layer=self.get_last_layer(), split="val") + + discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step, + last_layer=self.get_last_layer(), split="val") + + self.log("val/rec_loss", log_dict_ae["val/rec_loss"]) + self.log_dict(log_dict_ae) + self.log_dict(log_dict_disc) + return self.log_dict + + def configure_optimizers(self): + lr = self.learning_rate + opt_ae = torch.optim.Adam(list(self.encoder.parameters())+ + list(self.decoder.parameters())+ + list(self.quant_conv.parameters())+ + list(self.post_quant_conv.parameters()), + lr=lr, betas=(0.5, 0.9)) + opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(), + lr=lr, betas=(0.5, 0.9)) + return [opt_ae, opt_disc], [] + + def get_last_layer(self): + return self.decoder.conv_out.weight + + @torch.no_grad() + def log_images(self, batch, only_inputs=False, **kwargs): + log = dict() + x = self.get_input(batch, self.image_key) + x = x.to(self.device) + if not only_inputs: + xrec, posterior = self(x) + if x.shape[1] > 3: + # colorize with random projection + assert xrec.shape[1] > 3 + x = self.to_rgb(x) + xrec = self.to_rgb(xrec) + log["samples"] = self.decode(torch.randn_like(posterior.sample())) + log["reconstructions"] = xrec + log["inputs"] = x + return log + + def to_rgb(self, x): + assert self.image_key == "segmentation" + if not hasattr(self, "colorize"): + self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x)) + x = F.conv2d(x, weight=self.colorize) + x = 2.*(x-x.min())/(x.max()-x.min()) - 1. + return x + +class AutoencoderKL(pl.LightningModule): + def __init__(self, + ddconfig, + lossconfig, + embed_dim, + ckpt_path=None, + ignore_keys=[], + image_key="image", + colorize_nlabels=None, + monitor=None, + ): + super().__init__() + self.image_key = image_key + self.encoder = Encoder(**ddconfig) + self.decoder = Decoder(**ddconfig) + self.loss = instantiate_from_config(lossconfig) + assert ddconfig["double_z"] + self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1) + self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1) + self.embed_dim = embed_dim + if colorize_nlabels is not None: + assert type(colorize_nlabels)==int + self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1)) + if monitor is not None: + self.monitor = monitor + if ckpt_path is not None: + self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys) + + def init_from_ckpt(self, path, ignore_keys=list()): + sd = torch.load(path, map_location="cpu")["state_dict"] + keys = list(sd.keys()) + for k in keys: + for ik in ignore_keys: + if k.startswith(ik): + print("Deleting key {} from state_dict.".format(k)) + del sd[k] + self.load_state_dict(sd, strict=False) + print(f"Restored from {path}") + + def encode(self, x): + h = self.encoder(x) + moments = self.quant_conv(h) + posterior = DiagonalGaussianDistribution(moments) + return posterior + + def decode(self, z): + z = self.post_quant_conv(z) + dec = self.decoder(z) + return dec + + def forward(self, input, sample_posterior=True): + posterior = self.encode(input) + if sample_posterior: + z = posterior.sample() + else: + z = posterior.mode() + dec = self.decode(z) + return dec, posterior + + def get_input(self, batch, k): + x = batch[k] + if len(x.shape) == 3: + x = x[..., None] + x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float() + return x + + def training_step(self, batch, batch_idx, optimizer_idx): + # tic = time.time() + inputs = self.get_input(batch, self.image_key) + # print(f"get_input time {time.time() - tic}") + # tic = time.time() + reconstructions, posterior = self(inputs) + # print(f"model forward time {time.time() - tic}") + tic = time.time() + + if optimizer_idx == 0: + # train encoder+decoder+logvar + aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train") + self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True) + self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False) + # print(f"cal loss time {time.time() - tic}") + return aeloss + + if optimizer_idx == 1: + # train the discriminator + discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train") + + self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True) + self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False) + # print(f"cal loss time {time.time() - tic}") + return discloss + + def validation_step(self, batch, batch_idx): + tic = time.time() + inputs = self.get_input(batch, self.image_key) + print(f"get_input time {time.time() - tic}") + tic = time.time() + reconstructions, posterior = self(inputs) + print(f"val forward time {time.time() - tic}") + tic = time.time() + aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step, + last_layer=self.get_last_layer(), split="val") + + discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step, + last_layer=self.get_last_layer(), split="val") + + self.log("val/rec_loss", log_dict_ae["val/rec_loss"]) + self.log_dict(log_dict_ae) + self.log_dict(log_dict_disc) + print(f"val end time {time.time() - tic}") + return self.log_dict + + def configure_optimizers(self): + lr = self.learning_rate + opt_ae = torch.optim.Adam(list(self.encoder.parameters())+ + list(self.decoder.parameters())+ + list(self.quant_conv.parameters())+ + list(self.post_quant_conv.parameters()), + lr=lr, betas=(0.5, 0.9)) + opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(), + lr=lr, betas=(0.5, 0.9)) + return [opt_ae, opt_disc], [] + + def get_last_layer(self): + return self.decoder.conv_out.weight + + @torch.no_grad() + def log_images(self, batch, only_inputs=False, **kwargs): + log = dict() + x = self.get_input(batch, self.image_key) + x = x.to(self.device) + if not only_inputs: + xrec, posterior = self(x) + if x.shape[1] > 3: + # colorize with random projection + assert xrec.shape[1] > 3 + x = self.to_rgb(x) + xrec = self.to_rgb(xrec) + log["samples"] = self.decode(torch.randn_like(posterior.sample())) + log["reconstructions"] = xrec + log["inputs"] = x + return log + + def to_rgb(self, x): + assert self.image_key == "segmentation" + if not hasattr(self, "colorize"): + self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x)) + x = F.conv2d(x, weight=self.colorize) + x = 2.*(x-x.min())/(x.max()-x.min()) - 1. + return x + + +class IdentityFirstStage(torch.nn.Module): + def __init__(self, *args, vq_interface=False, **kwargs): + self.vq_interface = vq_interface # TODO: Should be true by default but check to not break older stuff + super().__init__() + + def encode(self, x, *args, **kwargs): + return x + + def decode(self, x, *args, **kwargs): + return x + + def quantize(self, x, *args, **kwargs): + if self.vq_interface: + return x, None, [None, None, None] + return x + + def forward(self, x, *args, **kwargs): + return x diff --git a/ruyi/vae/ldm/models/enc_dec_pytorch.py b/ruyi/vae/ldm/models/enc_dec_pytorch.py new file mode 100644 index 0000000..af84461 --- /dev/null +++ b/ruyi/vae/ldm/models/enc_dec_pytorch.py @@ -0,0 +1,234 @@ +import torch +import torch.nn.functional as F +from torch import nn + + +def cast_tuple(t, length = 1): + return t if isinstance(t, tuple) else ((t,) * length) + +def divisible_by(num, den): + return (num % den) == 0 + +def is_odd(n): + return not divisible_by(n, 2) + +class CausalConv3d(nn.Module): + def __init__( + self, + chan_in, + chan_out, + kernel_size, + pad_mode = 'constant', + **kwargs + ): + super().__init__() + kernel_size = cast_tuple(kernel_size, 3) + + time_kernel_size, height_kernel_size, width_kernel_size = kernel_size + + assert is_odd(height_kernel_size) and is_odd(width_kernel_size) + + dilation = kwargs.pop('dilation', 1) + stride = kwargs.pop('stride', 1) + + self.pad_mode = pad_mode + time_pad = dilation * (time_kernel_size - 1) + (1 - stride) + height_pad = height_kernel_size // 2 + width_pad = width_kernel_size // 2 + + self.time_pad = time_pad + self.time_causal_padding = (width_pad, width_pad, height_pad, height_pad, time_pad, 0) + + stride = (stride, 1, 1) + dilation = (dilation, 1, 1) + self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride = stride, dilation = dilation, **kwargs) + + def forward(self, x): + x = F.pad(x, self.time_causal_padding, mode = 'replicate') + return self.conv(x) + +class Swish(nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x): + return x * F.sigmoid(x) + +class ResBlockX(nn.Module): + def __init__(self, inchannel) -> None: + super().__init__() + self.conv = nn.Sequential( + nn.GroupNorm(32, inchannel), + Swish(), + CausalConv3d(inchannel, inchannel, 3), + nn.GroupNorm(32, inchannel), + Swish(), + CausalConv3d(inchannel, inchannel, 3) + ) + + def forward(self, x): + return x + self.conv(x) + +class ResBlockXY(nn.Module): + def __init__(self, inchannel, outchannel) -> None: + super().__init__() + self.conv = nn.Sequential( + nn.GroupNorm(32, inchannel), + Swish(), + CausalConv3d(inchannel, outchannel, 3), + nn.GroupNorm(32, outchannel), + Swish(), + CausalConv3d(outchannel, outchannel, 3) + ) + self.conv_1 = nn.Conv3d(inchannel, outchannel, 1) + + def forward(self, x): + return self.conv_1(x) + self.conv(x) + +class PoolDown222(nn.Module): + def __init__(self) -> None: + super().__init__() + self.pool = nn.AvgPool3d(2, 2) + + def forward(self, x): + x = F.pad(x, (0, 0, 0, 0, 1, 0), 'replicate') + return self.pool(x) + +class PoolDown122(nn.Module): + def __init__(self) -> None: + super().__init__() + self.pool = nn.AvgPool3d((1, 2, 2), (1, 2, 2)) + + def forward(self, x): + return self.pool(x) + +class Unpool222(nn.Module): + def __init__(self) -> None: + super().__init__() + self.up = nn.Upsample(scale_factor=2, mode='nearest') + + def forward(self, x): + x = self.up(x) + return x[:, :, 1:] + +class Unpool122(nn.Module): + def __init__(self) -> None: + super().__init__() + self.up = nn.Upsample(scale_factor=(1, 2, 2), mode='nearest') + + def forward(self, x): + x = self.up(x) + return x + +class ResBlockDown(nn.Module): + def __init__(self, inchannel, outchannel) -> None: + super().__init__() + self.blcok = nn.Sequential( + CausalConv3d(inchannel, outchannel, 3), + nn.LeakyReLU(inplace=True), + PoolDown222(), + CausalConv3d(outchannel, outchannel, 3), + nn.LeakyReLU(inplace=True) + ) + self.res = nn.Sequential( + PoolDown222(), + nn.Conv3d(inchannel, outchannel, 1) + ) + + def forward(self, x): + return self.res(x) + self.blcok(x) + + +class Discriminator(nn.Module): + def __init__(self) -> None: + super().__init__() + self.block = nn.Sequential( + CausalConv3d(3, 64, 3), + nn.LeakyReLU(inplace=True), + ResBlockDown(64, 128), + ResBlockDown(128, 256), + ResBlockDown(256, 256), + ResBlockDown(256, 256), + ResBlockDown(256, 256), + CausalConv3d(256, 256, 3), + nn.LeakyReLU(inplace=True), + nn.AdaptiveAvgPool3d(1), + nn.Flatten(), + nn.Linear(256, 256), + nn.LeakyReLU(inplace=True), + nn.Linear(256, 1) + ) + + def forward(self, x): + if x.ndim==4: + x = x.unsqueeze(2) + return self.block(x) + + + +class Encoder(nn.Module): + def __init__(self) -> None: + super().__init__() + self.encoder = nn.Sequential( + CausalConv3d(3, 64, 3), + ResBlockX(64), + ResBlockX(64), + PoolDown222(), + ResBlockXY(64, 128), + ResBlockX(128), + PoolDown222(), + ResBlockX(128), + ResBlockX(128), + PoolDown122(), + ResBlockXY(128, 256), + ResBlockX(256), + ResBlockX(256), + ResBlockX(256), + nn.GroupNorm(32, 256), + Swish(), + nn.Conv3d(256, 16, 1) + ) + + def forward(self, x): + return self.encoder(x) + +class Decoder(nn.Module): + def __init__(self) -> None: + super().__init__() + self.decoder = nn.Sequential( + CausalConv3d(8, 256, 3), + ResBlockX(256), + ResBlockX(256), + ResBlockX(256), + ResBlockX(256), + Unpool122(), + CausalConv3d(256, 256, 3), + ResBlockXY(256, 128), + ResBlockX(128), + Unpool222(), + CausalConv3d(128, 128, 3), + ResBlockX(128), + ResBlockX(128), + Unpool222(), + CausalConv3d(128, 128, 3), + ResBlockXY(128, 64), + ResBlockX(64), + nn.GroupNorm(32, 64), + Swish(), + CausalConv3d(64, 64, 3) + ) + self.conv_out = nn.Conv3d(64, 3, 1) + + def forward(self, x): + return self.conv_out(self.decoder(x)) + + +if __name__=='__main__': + encoder = Encoder() + decoder = Decoder() + dis = Discriminator() + x = torch.randn((1, 3, 1, 64, 64)) + embedding = encoder(x) + y = decoder(embedding) + tmp = torch.randn((1, 4, 1, 64, 64)) + print('something mmm') \ No newline at end of file diff --git a/ruyi/vae/ldm/models/omnigen_casual3dcnn.py b/ruyi/vae/ldm/models/omnigen_casual3dcnn.py new file mode 100644 index 0000000..cd54b35 --- /dev/null +++ b/ruyi/vae/ldm/models/omnigen_casual3dcnn.py @@ -0,0 +1,340 @@ +import itertools +from dataclasses import dataclass +from typing import Optional + +import numpy as np +import pytorch_lightning as pl +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from ..util import instantiate_from_config +from .omnigen_enc_dec import Decoder as omnigen_Mag_Decoder +from .omnigen_enc_dec import Encoder as omnigen_Mag_Encoder + + +class DiagonalGaussianDistribution: + def __init__( + self, + mean: torch.Tensor, + logvar: torch.Tensor, + deterministic: bool = False, + ): + self.mean = mean + self.logvar = torch.clamp(logvar, -30.0, 20.0) + self.deterministic = deterministic + + if deterministic: + self.var = self.std = torch.zeros_like(self.mean) + else: + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + + def sample(self, generator = None) -> torch.FloatTensor: + x = torch.randn( + self.mean.shape, + generator=generator, + device=self.mean.device, + dtype=self.mean.dtype, + ) + return self.mean + self.std * x + + def mode(self): + return self.mean + + def kl(self, other: Optional["DiagonalGaussianDistribution"] = None) -> torch.Tensor: + dims = list(range(1, self.mean.ndim)) + + if self.deterministic: + return torch.Tensor([0.0]) + else: + if other is None: + return 0.5 * torch.sum( + torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, + dim=dims, + ) + else: + return 0.5 * torch.sum( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var + - 1.0 + - self.logvar + + other.logvar, + dim=dims, + ) + + def nll(self, sample: torch.Tensor) -> torch.Tensor: + dims = list(range(1, self.mean.ndim)) + + if self.deterministic: + return torch.Tensor([0.0]) + + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum( + logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, + dim=dims, + ) + +@dataclass +class EncoderOutput: + latent_dist: DiagonalGaussianDistribution + +@dataclass +class DecoderOutput: + sample: torch.Tensor + +def str_eval(item): + if type(item) == str: + return eval(item) + else: + return item + +class AutoencoderKLMagvit_fromOmnigen(pl.LightningModule): + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + ch = 128, + ch_mult = [ 1,2,4,4 ], + use_gc_blocks = None, + down_block_types: tuple = None, + up_block_types: tuple = None, + mid_block_type: str = "MidBlock3D", + mid_block_use_attention: bool = True, + mid_block_attention_type: str = "3d", + mid_block_num_attention_heads: int = 1, + layers_per_block: int = 2, + act_fn: str = "silu", + num_attention_heads: int = 1, + latent_channels: int = 4, + norm_num_groups: int = 32, + image_key="image", + monitor=None, + ckpt_path=None, + lossconfig=None, + slice_mag_vae=False, + slice_compression_vae=False, + cache_compression_vae=False, + spatial_group_norm=False, + mini_batch_encoder=9, + mini_batch_decoder=3, + train_decoder_only=False, + train_encoder_only=False, + ): + super().__init__() + self.image_key = image_key + down_block_types = str_eval(down_block_types) + up_block_types = str_eval(up_block_types) + self.encoder = omnigen_Mag_Encoder( + in_channels=in_channels, + out_channels=latent_channels, + down_block_types=down_block_types, + ch = ch, + ch_mult = ch_mult, + use_gc_blocks=use_gc_blocks, + mid_block_type=mid_block_type, + mid_block_use_attention=mid_block_use_attention, + mid_block_attention_type=mid_block_attention_type, + mid_block_num_attention_heads=mid_block_num_attention_heads, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + num_attention_heads=num_attention_heads, + double_z=True, + slice_mag_vae=slice_mag_vae, + slice_compression_vae=slice_compression_vae, + cache_compression_vae=cache_compression_vae, + spatial_group_norm=spatial_group_norm, + mini_batch_encoder=mini_batch_encoder, + ) + + self.decoder = omnigen_Mag_Decoder( + in_channels=latent_channels, + out_channels=out_channels, + up_block_types=up_block_types, + ch = ch, + ch_mult = ch_mult, + use_gc_blocks=use_gc_blocks, + mid_block_type=mid_block_type, + mid_block_use_attention=mid_block_use_attention, + mid_block_attention_type=mid_block_attention_type, + mid_block_num_attention_heads=mid_block_num_attention_heads, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + num_attention_heads=num_attention_heads, + slice_mag_vae=slice_mag_vae, + slice_compression_vae=slice_compression_vae, + cache_compression_vae=cache_compression_vae, + spatial_group_norm=spatial_group_norm, + mini_batch_decoder=mini_batch_decoder, + ) + + self.quant_conv = nn.Conv3d(2 * latent_channels, 2 * latent_channels, kernel_size=1) + self.post_quant_conv = nn.Conv3d(latent_channels, latent_channels, kernel_size=1) + + self.mini_batch_encoder = mini_batch_encoder + self.mini_batch_decoder = mini_batch_decoder + self.train_decoder_only = train_decoder_only + self.train_encoder_only = train_encoder_only + if train_decoder_only: + self.encoder.requires_grad_(False) + self.quant_conv.requires_grad_(False) + if train_encoder_only: + self.decoder.requires_grad_(False) + self.post_quant_conv.requires_grad_(False) + if monitor is not None: + self.monitor = monitor + if ckpt_path is not None: + self.init_from_ckpt(ckpt_path, ignore_keys="loss") + if lossconfig is not None: + self.loss = instantiate_from_config(lossconfig) + + def init_from_ckpt(self, path, ignore_keys=list()): + if path.endswith("safetensors"): + from safetensors.torch import load_file, safe_open + sd = load_file(path) + else: + sd = torch.load(path, map_location="cpu") + if "state_dict" in list(sd.keys()): + sd = sd["state_dict"] + keys = list(sd.keys()) + for k in keys: + for ik in ignore_keys: + if k.startswith(ik): + print("Deleting key {} from state_dict.".format(k)) + del sd[k] + self.load_state_dict(sd, strict=False) # loss.item can be ignored successfully + print(f"Restored from {path}") + + def encode(self, x: torch.Tensor) -> EncoderOutput: + h = self.encoder(x) + + moments: torch.Tensor = self.quant_conv(h) + mean, logvar = moments.chunk(2, dim=1) + posterior = DiagonalGaussianDistribution(mean, logvar) + + # return EncoderOutput(latent_dist=posterior) + return posterior + + def decode(self, z: torch.Tensor) -> DecoderOutput: + z = self.post_quant_conv(z) + + decoded = self.decoder(z) + + # return DecoderOutput(sample=decoded) + return decoded + + + def forward(self, input, sample_posterior=True): + if input.ndim==4: + input = input.unsqueeze(2) + posterior = self.encode(input) + if sample_posterior: + z = posterior.sample() + else: + z = posterior.mode() + # print("stt latent shape", z.shape) + dec = self.decode(z) + return dec, posterior + + def get_input(self, batch, k): + x = batch[k] + if x.ndim==5: + x = x.permute(0, 4, 1, 2, 3).to(memory_format=torch.contiguous_format).float() + return x + if len(x.shape) == 3: + x = x[..., None] + x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float() + return x + + def training_step(self, batch, batch_idx, optimizer_idx): + # tic = time.time() + inputs = self.get_input(batch, self.image_key) + # print(f"get_input time {time.time() - tic}") + # tic = time.time() + reconstructions, posterior = self(inputs) + # print(f"model forward time {time.time() - tic}") + + if optimizer_idx == 0: + # train encoder+decoder+logvar + aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train") + self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True) + self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False) + # print(f"cal loss time {time.time() - tic}") + return aeloss + + if optimizer_idx == 1: + # train the discriminator + discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train") + + self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True) + self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False) + # print(f"cal loss time {time.time() - tic}") + return discloss + + def validation_step(self, batch, batch_idx): + with torch.no_grad(): + inputs = self.get_input(batch, self.image_key) + reconstructions, posterior = self(inputs) + aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step, + last_layer=self.get_last_layer(), split="val") + + discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step, + last_layer=self.get_last_layer(), split="val") + + self.log("val/rec_loss", log_dict_ae["val/rec_loss"]) + self.log_dict(log_dict_ae) + self.log_dict(log_dict_disc) + return self.log_dict + + def configure_optimizers(self): + lr = self.learning_rate + if self.train_decoder_only: + opt_ae = torch.optim.Adam(list(self.decoder.parameters())+ + list(self.post_quant_conv.parameters()), + lr=lr, betas=(0.5, 0.9)) + elif self.train_encoder_only: + opt_ae = torch.optim.Adam(list(self.encoder.parameters())+ + list(self.quant_conv.parameters()), + lr=lr, betas=(0.5, 0.9)) + else: + opt_ae = torch.optim.Adam(list(self.encoder.parameters())+ + list(self.decoder.parameters())+ + list(self.quant_conv.parameters())+ + list(self.post_quant_conv.parameters()), + lr=lr, betas=(0.5, 0.9)) + opt_disc = torch.optim.Adam(list(self.loss.discriminator3d.parameters()) + list(self.loss.discriminator.parameters()), + lr=lr, betas=(0.5, 0.9)) + return [opt_ae, opt_disc], [] + + def get_last_layer(self): + return self.decoder.conv_out.weight + + @torch.no_grad() + def log_images(self, batch, only_inputs=False, **kwargs): + log = dict() + x = self.get_input(batch, self.image_key) + x = x.to(self.device) + if not only_inputs: + xrec, posterior = self(x) + if x.shape[1] > 3: + # colorize with random projection + assert xrec.shape[1] > 3 + x = self.to_rgb(x) + xrec = self.to_rgb(xrec) + log["samples"] = self.decode(torch.randn_like(posterior.sample())) + log["reconstructions"] = xrec + log["inputs"] = x + return log + + def to_rgb(self, x): + assert self.image_key == "segmentation" + if not hasattr(self, "colorize"): + self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x)) + x = F.conv2d(x, weight=self.colorize) + x = 2.*(x-x.min())/(x.max()-x.min()) - 1. + return x diff --git a/ruyi/vae/ldm/models/omnigen_enc_dec.py b/ruyi/vae/ldm/models/omnigen_enc_dec.py new file mode 100644 index 0000000..2156da1 --- /dev/null +++ b/ruyi/vae/ldm/models/omnigen_enc_dec.py @@ -0,0 +1,565 @@ +import torch +import torch.nn as nn +import numpy as np +from typing import Any, Dict, Optional, Tuple +from einops import rearrange, repeat +from diffusers.utils import (USE_PEFT_BACKEND, BaseOutput, is_torch_version, + logging) +from ..modules.vaemodules.activations import get_activation +from ..modules.vaemodules.common import CausalConv3d +from ..modules.vaemodules.down_blocks import get_down_block +from ..modules.vaemodules.mid_blocks import get_mid_block +from ..modules.vaemodules.up_blocks import get_up_block + + +def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + +class Encoder(nn.Module): + r""" + The `Encoder` layer of a variational autoencoder that encodes its input into a latent representation. + + Args: + in_channels (`int`, *optional*, defaults to 3): + The number of input channels. + out_channels (`int`, *optional*, defaults to 8): + The number of output channels. + down_block_types (`Tuple[str, ...]`, *optional*, defaults to `("SpatialDownBlock3D",)`): + The types of down blocks to use. + block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): + The number of output channels for each block. + use_gc_blocks (`Tuple[bool, ...]`, *optional*, defaults to `None`): + Whether to use global context blocks for each down block. + mid_block_type (`str`, *optional*, defaults to `"MidBlock3D"`): + The type of mid block to use. + layers_per_block (`int`, *optional*, defaults to 2): + The number of layers per block. + norm_num_groups (`int`, *optional*, defaults to 32): + The number of groups for normalization. + act_fn (`str`, *optional*, defaults to `"silu"`): + The activation function to use. See `~diffusers.models.activations.get_activation` for available options. + num_attention_heads (`int`, *optional*, defaults to 1): + The number of attention heads to use. + double_z (`bool`, *optional*, defaults to `True`): + Whether to double the number of output channels for the last block. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 8, + down_block_types = ("SpatialDownBlock3D",), + ch = 128, + ch_mult = [1,2,4,4,], + use_gc_blocks = None, + mid_block_type: str = "MidBlock3D", + mid_block_use_attention: bool = True, + mid_block_attention_type: str = "3d", + mid_block_num_attention_heads: int = 1, + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + num_attention_heads: int = 1, + double_z: bool = True, + slice_mag_vae: bool = False, + slice_compression_vae: bool = False, + cache_compression_vae: bool = False, + spatial_group_norm: bool = False, + mini_batch_encoder: int = 9, + verbose = False, + ): + super().__init__() + block_out_channels = [ch * i for i in ch_mult] + assert len(down_block_types) == len(block_out_channels), ( + "Number of down block types must match number of block output channels." + ) + if use_gc_blocks is not None: + assert len(use_gc_blocks) == len(down_block_types), ( + "Number of GC blocks must match number of down block types." + ) + else: + use_gc_blocks = [False] * len(down_block_types) + self.conv_in = CausalConv3d( + in_channels, + block_out_channels[0], + kernel_size=3, + ) + + self.down_blocks = nn.ModuleList([]) + + output_channels = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channels = output_channels + output_channels = block_out_channels[i] + is_final_block = (i == len(block_out_channels) - 1) + down_block = get_down_block( + down_block_type, + in_channels=input_channels, + out_channels=output_channels, + num_layers=layers_per_block, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=1e-6, + num_attention_heads=num_attention_heads, + add_gc_block=use_gc_blocks[i], + add_downsample=not is_final_block, + ) + self.down_blocks.append(down_block) + + self.mid_block = get_mid_block( + mid_block_type, + in_channels=block_out_channels[-1], + num_layers=layers_per_block, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=1e-6, + add_attention=mid_block_use_attention, + attention_type=mid_block_attention_type, + num_attention_heads=mid_block_num_attention_heads, + ) + + self.conv_norm_out = nn.GroupNorm( + num_channels=block_out_channels[-1], + num_groups=norm_num_groups, + eps=1e-6, + ) + self.conv_act = get_activation(act_fn) + + conv_out_channels = 2 * out_channels if double_z else out_channels + self.conv_out = CausalConv3d(block_out_channels[-1], conv_out_channels, kernel_size=3) + + self.slice_mag_vae = slice_mag_vae + self.slice_compression_vae = slice_compression_vae + self.cache_compression_vae = cache_compression_vae + self.mini_batch_encoder = mini_batch_encoder + self.spatial_group_norm = spatial_group_norm + self.verbose = verbose + + def set_padding_one_frame(self): + def _set_padding_one_frame(name, module): + if hasattr(module, 'padding_flag'): + if self.verbose: + print('Set pad mode for module[%s] type=%s' % (name, str(type(module)))) + module.padding_flag = 1 + for sub_name, sub_mod in module.named_children(): + _set_padding_one_frame(sub_name, sub_mod) + for name, module in self.named_children(): + _set_padding_one_frame(name, module) + + def set_padding_more_frame(self): + def _set_padding_more_frame(name, module): + if hasattr(module, 'padding_flag'): + if self.verbose: + print('Set pad mode for module[%s] type=%s' % (name, str(type(module)))) + module.padding_flag = 2 + for sub_name, sub_mod in module.named_children(): + _set_padding_more_frame(sub_name, sub_mod) + for name, module in self.named_children(): + _set_padding_more_frame(name, module) + + def set_cache_slice_vae_padding_one_frame(self): + def _set_cache_slice_vae_padding_one_frame(name, module): + if hasattr(module, 'padding_flag'): + if self.verbose: + print('Set pad mode for module[%s] type=%s' % (name, str(type(module)))) + module.padding_flag = 5 + for sub_name, sub_mod in module.named_children(): + _set_cache_slice_vae_padding_one_frame(sub_name, sub_mod) + for name, module in self.named_children(): + _set_cache_slice_vae_padding_one_frame(name, module) + + def set_cache_slice_vae_padding_more_frame(self): + def _set_cache_slice_vae_padding_more_frame(name, module): + if hasattr(module, 'padding_flag'): + if self.verbose: + print('Set pad mode for module[%s] type=%s' % (name, str(type(module)))) + module.padding_flag = 6 + for sub_name, sub_mod in module.named_children(): + _set_cache_slice_vae_padding_more_frame(sub_name, sub_mod) + for name, module in self.named_children(): + _set_cache_slice_vae_padding_more_frame(name, module) + + def set_3dgroupnorm_for_submodule(self): + def _set_3dgroupnorm_for_submodule(name, module): + if hasattr(module, 'set_3dgroupnorm'): + if self.verbose: + print('Set pad mode for module[%s] type=%s' % (name, str(type(module)))) + module.set_3dgroupnorm = True + for sub_name, sub_mod in module.named_children(): + _set_3dgroupnorm_for_submodule(sub_name, sub_mod) + for name, module in self.named_children(): + _set_3dgroupnorm_for_submodule(name, module) + + def single_forward(self, x: torch.Tensor, previous_features: torch.Tensor, after_features: torch.Tensor) -> torch.Tensor: + # x: (B, C, T, H, W) + if previous_features is not None and after_features is None: + x = torch.concat([previous_features, x], 2) + elif previous_features is None and after_features is not None: + x = torch.concat([x, after_features], 2) + elif previous_features is not None and after_features is not None: + x = torch.concat([previous_features, x, after_features], 2) + + x = self.conv_in(x) + + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + for down_block in self.down_blocks: + x = down_block(x) + + x = self.mid_block(x) + + if self.spatial_group_norm: + batch_size = x.shape[0] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.conv_norm_out(x) + x = rearrange(x, "(b t) c h w -> b c t h w", b=batch_size) + else: + x = self.conv_norm_out(x) + x = self.conv_act(x) + x = self.conv_out(x) + + if previous_features is not None and after_features is None: + x = x[:, :, 1:] + elif previous_features is None and after_features is not None: + x = x[:, :, :2] + elif previous_features is not None and after_features is not None: + x = x[:, :, 1:3] + return x + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.spatial_group_norm: + self.set_3dgroupnorm_for_submodule() + + if self.cache_compression_vae: + _, _, f, _, _ = x.size() + if f % 2 != 0: + self.set_padding_one_frame() + first_frames = self.single_forward(x[:, :, 0:1, :, :], None, None) + self.set_padding_more_frame() + + new_pixel_values = [first_frames] + start_index = 1 + else: + self.set_padding_more_frame() + new_pixel_values = [] + start_index = 0 + + for i in range(start_index, x.shape[2], self.mini_batch_encoder): + next_frames = self.single_forward(x[:, :, i: i + self.mini_batch_encoder, :, :], None, None) + new_pixel_values.append(next_frames) + new_pixel_values = torch.cat(new_pixel_values, dim=2) + elif self.slice_compression_vae: + _, _, f, _, _ = x.size() + if f % 2 != 0: + self.set_padding_one_frame() + first_frames = self.single_forward(x[:, :, 0:1, :, :], None, None) + self.set_padding_more_frame() + + new_pixel_values = [first_frames] + start_index = 1 + else: + self.set_padding_more_frame() + new_pixel_values = [] + start_index = 0 + + for i in range(start_index, x.shape[2], self.mini_batch_encoder): + next_frames = self.single_forward(x[:, :, i: i + self.mini_batch_encoder, :, :], None, None) + new_pixel_values.append(next_frames) + new_pixel_values = torch.cat(new_pixel_values, dim=2) + elif self.slice_mag_vae: + _, _, f, _, _ = x.size() + new_pixel_values = [] + for i in range(0, x.shape[2], self.mini_batch_encoder): + next_frames = self.single_forward(x[:, :, i: i + self.mini_batch_encoder, :, :], None, None) + new_pixel_values.append(next_frames) + new_pixel_values = torch.cat(new_pixel_values, dim=2) + else: + new_pixel_values = self.single_forward(x, None, None) + return new_pixel_values + +class Decoder(nn.Module): + r""" + The `Decoder` layer of a variational autoencoder that decodes its latent representation into an output sample. + + Args: + in_channels (`int`, *optional*, defaults to 8): + The number of input channels. + out_channels (`int`, *optional*, defaults to 3): + The number of output channels. + up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("SpatialUpBlock3D",)`): + The types of up blocks to use. + block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): + The number of output channels for each block. + use_gc_blocks (`Tuple[bool, ...]`, *optional*, defaults to `None`): + Whether to use global context blocks for each down block. + mid_block_type (`str`, *optional*, defaults to `"MidBlock3D"`): + The type of mid block to use. + layers_per_block (`int`, *optional*, defaults to 2): + The number of layers per block. + norm_num_groups (`int`, *optional*, defaults to 32): + The number of groups for normalization. + act_fn (`str`, *optional*, defaults to `"silu"`): + The activation function to use. See `~diffusers.models.activations.get_activation` for available options. + num_attention_heads (`int`, *optional*, defaults to 1): + The number of attention heads to use. + """ + + def __init__( + self, + in_channels: int = 8, + out_channels: int = 3, + up_block_types = ("SpatialUpBlock3D",), + ch = 128, + ch_mult = [1,2,4,4,], + use_gc_blocks = None, + mid_block_type: str = "MidBlock3D", + mid_block_use_attention: bool = True, + mid_block_attention_type: str = "3d", + mid_block_num_attention_heads: int = 1, + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + num_attention_heads: int = 1, + slice_mag_vae: bool = False, + slice_compression_vae: bool = False, + cache_compression_vae: bool = False, + spatial_group_norm: bool = False, + mini_batch_decoder: int = 3, + verbose = False, + ): + super().__init__() + block_out_channels = [ch * i for i in ch_mult] + assert len(up_block_types) == len(block_out_channels), ( + "Number of up block types must match number of block output channels." + ) + if use_gc_blocks is not None: + assert len(use_gc_blocks) == len(up_block_types), ( + "Number of GC blocks must match number of up block types." + ) + else: + use_gc_blocks = [False] * len(up_block_types) + + self.conv_in = CausalConv3d( + in_channels, + block_out_channels[-1], + kernel_size=3, + ) + + self.mid_block = get_mid_block( + mid_block_type, + in_channels=block_out_channels[-1], + num_layers=layers_per_block, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=1e-6, + add_attention=mid_block_use_attention, + attention_type=mid_block_attention_type, + num_attention_heads=mid_block_num_attention_heads, + ) + + self.up_blocks = nn.ModuleList([]) + + reversed_block_out_channels = list(reversed(block_out_channels)) + output_channels = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + input_channels = output_channels + output_channels = reversed_block_out_channels[i] + # is_first_block = i == 0 + is_final_block = i == len(block_out_channels) - 1 + + up_block = get_up_block( + up_block_type, + in_channels=input_channels, + out_channels=output_channels, + num_layers=layers_per_block + 1, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=1e-6, + num_attention_heads=num_attention_heads, + add_gc_block=use_gc_blocks[i], + add_upsample=not is_final_block, + ) + self.up_blocks.append(up_block) + + self.conv_norm_out = nn.GroupNorm( + num_channels=block_out_channels[0], + num_groups=norm_num_groups, + eps=1e-6, + ) + self.conv_act = get_activation(act_fn) + + self.conv_out = CausalConv3d(block_out_channels[0], out_channels, kernel_size=3) + + self.slice_mag_vae = slice_mag_vae + self.slice_compression_vae = slice_compression_vae + self.cache_compression_vae = cache_compression_vae + self.mini_batch_decoder = mini_batch_decoder + self.spatial_group_norm = spatial_group_norm + self.verbose = verbose + + def set_padding_one_frame(self): + def _set_padding_one_frame(name, module): + if hasattr(module, 'padding_flag'): + if self.verbose: + print('Set pad mode for module[%s] type=%s' % (name, str(type(module)))) + module.padding_flag = 1 + for sub_name, sub_mod in module.named_children(): + _set_padding_one_frame(sub_name, sub_mod) + for name, module in self.named_children(): + _set_padding_one_frame(name, module) + + def set_padding_more_frame(self): + def _set_padding_more_frame(name, module): + if hasattr(module, 'padding_flag'): + if self.verbose: + print('Set pad mode for module[%s] type=%s' % (name, str(type(module)))) + module.padding_flag = 2 + for sub_name, sub_mod in module.named_children(): + _set_padding_more_frame(sub_name, sub_mod) + for name, module in self.named_children(): + _set_padding_more_frame(name, module) + + def set_cache_slice_vae_padding_one_frame(self): + def _set_cache_slice_vae_padding_one_frame(name, module): + if hasattr(module, 'padding_flag'): + if self.verbose: + print('Set pad mode for module[%s] type=%s' % (name, str(type(module)))) + module.padding_flag = 5 + for sub_name, sub_mod in module.named_children(): + _set_cache_slice_vae_padding_one_frame(sub_name, sub_mod) + for name, module in self.named_children(): + _set_cache_slice_vae_padding_one_frame(name, module) + + def set_cache_slice_vae_padding_more_frame(self): + def _set_cache_slice_vae_padding_more_frame(name, module): + if hasattr(module, 'padding_flag'): + if self.verbose: + print('Set pad mode for module[%s] type=%s' % (name, str(type(module)))) + module.padding_flag = 6 + for sub_name, sub_mod in module.named_children(): + _set_cache_slice_vae_padding_more_frame(sub_name, sub_mod) + for name, module in self.named_children(): + _set_cache_slice_vae_padding_more_frame(name, module) + + def set_3dgroupnorm_for_submodule(self): + def _set_3dgroupnorm_for_submodule(name, module): + if hasattr(module, 'set_3dgroupnorm'): + if self.verbose: + print('Set pad mode for module[%s] type=%s' % (name, str(type(module)))) + module.set_3dgroupnorm = True + for sub_name, sub_mod in module.named_children(): + _set_3dgroupnorm_for_submodule(sub_name, sub_mod) + for name, module in self.named_children(): + _set_3dgroupnorm_for_submodule(name, module) + + def clear_cache(self): + def _clear_cache(name, module): + if hasattr(module, 'prev_features'): + if self.verbose: + print('Set pad mode for module[%s] type=%s' % (name, str(type(module)))) + module.prev_features = None + for sub_name, sub_mod in module.named_children(): + _clear_cache(sub_name, sub_mod) + for name, module in self.named_children(): + _clear_cache(name, module) + + def single_forward(self, x: torch.Tensor, previous_features: torch.Tensor, after_features: torch.Tensor) -> torch.Tensor: + # x: (B, C, T, H, W) + ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + if previous_features is not None and after_features is None: + b, c, t, h, w = x.size() + x = torch.concat([previous_features, x], 2) + x = self.conv_in(x) + x = self.mid_block(x) + x = x[:, :, -t:] + elif previous_features is None and after_features is not None: + b, c, t, h, w = x.size() + x = torch.concat([x, after_features], 2) + x = self.conv_in(x) + x = self.mid_block(x) + x = x[:, :, :t] + elif previous_features is not None and after_features is not None: + _, _, t_1, _, _ = previous_features.size() + _, _, t_2, _, _ = x.size() + x = torch.concat([previous_features, x, after_features], 2) + x = self.conv_in(x) + x = self.mid_block(x) + x = x[:, :, t_1:(t_1 + t_2)] + else: + x = self.conv_in(x) + x = self.mid_block(x) + + for up_block in self.up_blocks: + x = up_block(x) + + if self.spatial_group_norm: + batch_size = x.shape[0] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.conv_norm_out(x) + x = rearrange(x, "(b t) c h w -> b c t h w", b=batch_size) + else: + x = self.conv_norm_out(x) + + x = self.conv_act(x) + x = self.conv_out(x) + + return x + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.spatial_group_norm: + self.set_3dgroupnorm_for_submodule() + + if self.cache_compression_vae: + _, _, f, _, _ = x.size() + if f == 1: + self.set_padding_one_frame() + first_frames = self.single_forward(x[:, :, :1, :, :], None, None) + new_pixel_values = [first_frames] + start_index = 1 + else: + self.set_cache_slice_vae_padding_one_frame() + first_frames = self.single_forward(x[:, :, :self.mini_batch_decoder, :, :], None, None) + new_pixel_values = [first_frames] + start_index = self.mini_batch_decoder + + for i in range(start_index, x.shape[2], self.mini_batch_decoder): + self.set_cache_slice_vae_padding_more_frame() + next_frames = self.single_forward(x[:, :, i: i + self.mini_batch_decoder, :, :], None, None) + new_pixel_values.append(next_frames) + new_pixel_values = torch.cat(new_pixel_values, dim=2) + elif self.slice_compression_vae: + _, _, f, _, _ = x.size() + if f % 2 != 0: + self.set_padding_one_frame() + first_frames = self.single_forward(x[:, :, 0:1, :, :], None, None) + self.set_padding_more_frame() + new_pixel_values = [first_frames] + start_index = 1 + else: + self.set_padding_more_frame() + new_pixel_values = [] + start_index = 0 + + previous_features = None + for i in range(start_index, x.shape[2], self.mini_batch_decoder): + after_features = x[:, :, i + self.mini_batch_decoder: i + 2 * self.mini_batch_decoder, :, :] if i + self.mini_batch_decoder < x.shape[2] else None + next_frames = self.single_forward(x[:, :, i: i + self.mini_batch_decoder, :, :], previous_features, after_features) + previous_features = x[:, :, i: i + self.mini_batch_decoder, :, :] + new_pixel_values.append(next_frames) + new_pixel_values = torch.cat(new_pixel_values, dim=2) + elif self.slice_mag_vae: + _, _, f, _, _ = x.size() + new_pixel_values = [] + for i in range(0, x.shape[2], self.mini_batch_decoder): + next_frames = self.single_forward(x[:, :, i: i + self.mini_batch_decoder, :, :], None, None) + new_pixel_values.append(next_frames) + new_pixel_values = torch.cat(new_pixel_values, dim=2) + else: + new_pixel_values = self.single_forward(x, None, None) + return new_pixel_values diff --git a/ruyi/vae/ldm/modules/diffusionmodules/__init__.py b/ruyi/vae/ldm/modules/diffusionmodules/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ruyi/vae/ldm/modules/diffusionmodules/model.py b/ruyi/vae/ldm/modules/diffusionmodules/model.py new file mode 100644 index 0000000..5c7dd01 --- /dev/null +++ b/ruyi/vae/ldm/modules/diffusionmodules/model.py @@ -0,0 +1,701 @@ +# pytorch_diffusion + derived encoder decoder +import math + +import numpy as np +import torch +import torch.nn as nn +from einops import rearrange + +from ...util import instantiate_from_config +from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution + +def get_timestep_embedding(timesteps, embedding_dim): + """ + This matches the implementation in Denoising Diffusion Probabilistic Models: + From Fairseq. + Build sinusoidal embeddings. + This matches the implementation in tensor2tensor, but differs slightly + from the description in Section 3.5 of "Attention Is All You Need". + """ + assert len(timesteps.shape) == 1 + + half_dim = embedding_dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb) + emb = emb.to(device=timesteps.device) + emb = timesteps.float()[:, None] * emb[None, :] + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) + if embedding_dim % 2 == 1: # zero pad + emb = torch.nn.functional.pad(emb, (0,1,0,0)) + return emb + + +def nonlinearity(x): + # swish + return x*torch.sigmoid(x) + + +def Normalize(in_channels, num_groups=32): + return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True) + + +class Upsample(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + self.conv = torch.nn.Conv2d(in_channels, + in_channels, + kernel_size=3, + stride=1, + padding=1) + + def forward(self, x): + x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + if self.with_conv: + x = self.conv(x) + return x + + +class Downsample(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + # no asymmetric padding in torch conv, must do it ourselves + self.conv = torch.nn.Conv2d(in_channels, + in_channels, + kernel_size=3, + stride=2, + padding=0) + + def forward(self, x): + if self.with_conv: + pad = (0,1,0,1) + x = torch.nn.functional.pad(x, pad, mode="constant", value=0) + x = self.conv(x) + else: + x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2) + return x + + +class ResnetBlock(nn.Module): + def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, + dropout, temb_channels=512): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + + self.norm1 = Normalize(in_channels) + self.conv1 = torch.nn.Conv2d(in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + if temb_channels > 0: + self.temb_proj = torch.nn.Linear(temb_channels, + out_channels) + self.norm2 = Normalize(out_channels) + self.dropout = torch.nn.Dropout(dropout) + self.conv2 = torch.nn.Conv2d(out_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = torch.nn.Conv2d(in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + else: + self.nin_shortcut = torch.nn.Conv2d(in_channels, + out_channels, + kernel_size=1, + stride=1, + padding=0) + + def forward(self, x, temb): + h = x + h = self.norm1(h) + h = nonlinearity(h) + h = self.conv1(h) + + if temb is not None: + h = h + self.temb_proj(nonlinearity(temb))[:,:,None,None] + + h = self.norm2(h) + h = nonlinearity(h) + h = self.dropout(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + x = self.conv_shortcut(x) + else: + x = self.nin_shortcut(x) + + return x+h + + + + +class AttnBlock(nn.Module): + def __init__(self, in_channels): + super().__init__() + self.in_channels = in_channels + + self.norm = Normalize(in_channels) + self.q = torch.nn.Conv2d(in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0) + self.k = torch.nn.Conv2d(in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0) + self.v = torch.nn.Conv2d(in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0) + self.proj_out = torch.nn.Conv2d(in_channels, + in_channels, + kernel_size=1, + stride=1, + padding=0) + + + def forward(self, x): + h_ = x + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + # compute attention + b,c,h,w = q.shape + q = q.reshape(b,c,h*w) + q = q.permute(0,2,1) # b,hw,c + k = k.reshape(b,c,h*w) # b,c,hw + w_ = torch.bmm(q,k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j] + w_ = w_ * (int(c)**(-0.5)) + w_ = torch.nn.functional.softmax(w_, dim=2) + + # attend to values + v = v.reshape(b,c,h*w) + w_ = w_.permute(0,2,1) # b,hw,hw (first hw of k, second of q) + h_ = torch.bmm(v,w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j] + h_ = h_.reshape(b,c,h,w) + + h_ = self.proj_out(h_) + + return x+h_ + +class LinearAttention(nn.Module): + def __init__(self, dim, heads=4, dim_head=32): + super().__init__() + self.heads = heads + hidden_dim = dim_head * heads + self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False) + self.to_out = nn.Conv2d(hidden_dim, dim, 1) + + def forward(self, x): + b, c, h, w = x.shape + qkv = self.to_qkv(x) + q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3) + k = k.softmax(dim=-1) + context = torch.einsum('bhdn,bhen->bhde', k, v) + out = torch.einsum('bhde,bhdn->bhen', context, q) + out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w) + return self.to_out(out) + +class LinAttnBlock(LinearAttention): + """to match AttnBlock usage""" + def __init__(self, in_channels): + super().__init__(dim=in_channels, heads=1, dim_head=in_channels) + +def make_attn(in_channels, attn_type="vanilla"): + assert attn_type in ["vanilla", "linear", "none"], f'attn_type {attn_type} unknown' + print(f"making attention of type '{attn_type}' with {in_channels} in_channels") + if attn_type == "vanilla": + return AttnBlock(in_channels) + elif attn_type == "none": + return nn.Identity(in_channels) + else: + return LinAttnBlock(in_channels) + + +class Encoder(nn.Module): + def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, + attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, + resolution, z_channels, double_z=True, use_linear_attn=False, attn_type="vanilla", + **ignore_kwargs): + super().__init__() + if use_linear_attn: attn_type = "linear" + self.ch = ch + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + + # downsampling + self.conv_in = torch.nn.Conv2d(in_channels, + self.ch, + kernel_size=3, + stride=1, + padding=1) + + curr_res = resolution + in_ch_mult = (1,)+tuple(ch_mult) + self.in_ch_mult = in_ch_mult + self.down = nn.ModuleList() + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch*in_ch_mult[i_level] + block_out = ch*ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(make_attn(block_in, attn_type=attn_type)) + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions-1: + down.downsample = Downsample(block_in, resamp_with_conv) + curr_res = curr_res // 2 + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + self.mid.attn_1 = make_attn(block_in, attn_type=attn_type) + self.mid.block_2 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + + # end + self.norm_out = Normalize(block_in) + self.conv_out = torch.nn.Conv2d(block_in, + 2*z_channels if double_z else z_channels, + kernel_size=3, + stride=1, + padding=1) + + def forward(self, x): + # timestep embedding + temb = None + + # downsampling + hs = [self.conv_in(x)] + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1], temb) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if i_level != self.num_resolutions-1: + hs.append(self.down[i_level].downsample(hs[-1])) + + # middle + h = hs[-1] + h = self.mid.block_1(h, temb) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb) + + # end + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + +class Decoder(nn.Module): + def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, + attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, + resolution, z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False, + attn_type="vanilla", **ignorekwargs): + super().__init__() + if use_linear_attn: attn_type = "linear" + self.ch = ch + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.give_pre_end = give_pre_end + self.tanh_out = tanh_out + + # compute in_ch_mult, block_in and curr_res at lowest res + in_ch_mult = (1,)+tuple(ch_mult) + block_in = ch*ch_mult[self.num_resolutions-1] + curr_res = resolution // 2**(self.num_resolutions-1) + self.z_shape = (1,z_channels,curr_res,curr_res) + print("Working with z of shape {} = {} dimensions.".format( + self.z_shape, np.prod(self.z_shape))) + + # z to block_in + self.conv_in = torch.nn.Conv2d(z_channels, + block_in, + kernel_size=3, + stride=1, + padding=1) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + self.mid.attn_1 = make_attn(block_in, attn_type=attn_type) + self.mid.block_2 = ResnetBlock(in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch*ch_mult[i_level] + for i_block in range(self.num_res_blocks+1): + block.append(ResnetBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(make_attn(block_in, attn_type=attn_type)) + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = Upsample(block_in, resamp_with_conv) + curr_res = curr_res * 2 + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = Normalize(block_in) + self.conv_out = torch.nn.Conv2d(block_in, + out_ch, + kernel_size=3, + stride=1, + padding=1) + + def forward(self, z): + #assert z.shape[1:] == self.z_shape[1:] + self.last_z_shape = z.shape + + # timestep embedding + temb = None + + # z to block_in + h = self.conv_in(z) + + # middle + h = self.mid.block_1(h, temb) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb) + + # upsampling + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks+1): + h = self.up[i_level].block[i_block](h, temb) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + h = self.up[i_level].upsample(h) + + # end + if self.give_pre_end: + return h + + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + if self.tanh_out: + h = torch.tanh(h) + return h + + +class SimpleDecoder(nn.Module): + def __init__(self, in_channels, out_channels, *args, **kwargs): + super().__init__() + self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1), + ResnetBlock(in_channels=in_channels, + out_channels=2 * in_channels, + temb_channels=0, dropout=0.0), + ResnetBlock(in_channels=2 * in_channels, + out_channels=4 * in_channels, + temb_channels=0, dropout=0.0), + ResnetBlock(in_channels=4 * in_channels, + out_channels=2 * in_channels, + temb_channels=0, dropout=0.0), + nn.Conv2d(2*in_channels, in_channels, 1), + Upsample(in_channels, with_conv=True)]) + # end + self.norm_out = Normalize(in_channels) + self.conv_out = torch.nn.Conv2d(in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1) + + def forward(self, x): + for i, layer in enumerate(self.model): + if i in [1,2,3]: + x = layer(x, None) + else: + x = layer(x) + + h = self.norm_out(x) + h = nonlinearity(h) + x = self.conv_out(h) + return x + + +class UpsampleDecoder(nn.Module): + def __init__(self, in_channels, out_channels, ch, num_res_blocks, resolution, + ch_mult=(2,2), dropout=0.0): + super().__init__() + # upsampling + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + block_in = in_channels + curr_res = resolution // 2 ** (self.num_resolutions - 1) + self.res_blocks = nn.ModuleList() + self.upsample_blocks = nn.ModuleList() + for i_level in range(self.num_resolutions): + res_block = [] + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks + 1): + res_block.append(ResnetBlock(in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout)) + block_in = block_out + self.res_blocks.append(nn.ModuleList(res_block)) + if i_level != self.num_resolutions - 1: + self.upsample_blocks.append(Upsample(block_in, True)) + curr_res = curr_res * 2 + + # end + self.norm_out = Normalize(block_in) + self.conv_out = torch.nn.Conv2d(block_in, + out_channels, + kernel_size=3, + stride=1, + padding=1) + + def forward(self, x): + # upsampling + h = x + for k, i_level in enumerate(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = self.res_blocks[i_level][i_block](h, None) + if i_level != self.num_resolutions - 1: + h = self.upsample_blocks[k](h) + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + +class LatentRescaler(nn.Module): + def __init__(self, factor, in_channels, mid_channels, out_channels, depth=2): + super().__init__() + # residual block, interpolate, residual block + self.factor = factor + self.conv_in = nn.Conv2d(in_channels, + mid_channels, + kernel_size=3, + stride=1, + padding=1) + self.res_block1 = nn.ModuleList([ResnetBlock(in_channels=mid_channels, + out_channels=mid_channels, + temb_channels=0, + dropout=0.0) for _ in range(depth)]) + self.attn = AttnBlock(mid_channels) + self.res_block2 = nn.ModuleList([ResnetBlock(in_channels=mid_channels, + out_channels=mid_channels, + temb_channels=0, + dropout=0.0) for _ in range(depth)]) + + self.conv_out = nn.Conv2d(mid_channels, + out_channels, + kernel_size=1, + ) + + def forward(self, x): + x = self.conv_in(x) + for block in self.res_block1: + x = block(x, None) + x = torch.nn.functional.interpolate(x, size=(int(round(x.shape[2]*self.factor)), int(round(x.shape[3]*self.factor)))) + x = self.attn(x) + for block in self.res_block2: + x = block(x, None) + x = self.conv_out(x) + return x + + +class MergedRescaleEncoder(nn.Module): + def __init__(self, in_channels, ch, resolution, out_ch, num_res_blocks, + attn_resolutions, dropout=0.0, resamp_with_conv=True, + ch_mult=(1,2,4,8), rescale_factor=1.0, rescale_module_depth=1): + super().__init__() + intermediate_chn = ch * ch_mult[-1] + self.encoder = Encoder(in_channels=in_channels, num_res_blocks=num_res_blocks, ch=ch, ch_mult=ch_mult, + z_channels=intermediate_chn, double_z=False, resolution=resolution, + attn_resolutions=attn_resolutions, dropout=dropout, resamp_with_conv=resamp_with_conv, + out_ch=None) + self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=intermediate_chn, + mid_channels=intermediate_chn, out_channels=out_ch, depth=rescale_module_depth) + + def forward(self, x): + x = self.encoder(x) + x = self.rescaler(x) + return x + + +class MergedRescaleDecoder(nn.Module): + def __init__(self, z_channels, out_ch, resolution, num_res_blocks, attn_resolutions, ch, ch_mult=(1,2,4,8), + dropout=0.0, resamp_with_conv=True, rescale_factor=1.0, rescale_module_depth=1): + super().__init__() + tmp_chn = z_channels*ch_mult[-1] + self.decoder = Decoder(out_ch=out_ch, z_channels=tmp_chn, attn_resolutions=attn_resolutions, dropout=dropout, + resamp_with_conv=resamp_with_conv, in_channels=None, num_res_blocks=num_res_blocks, + ch_mult=ch_mult, resolution=resolution, ch=ch) + self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=z_channels, mid_channels=tmp_chn, + out_channels=tmp_chn, depth=rescale_module_depth) + + def forward(self, x): + x = self.rescaler(x) + x = self.decoder(x) + return x + + +class Upsampler(nn.Module): + def __init__(self, in_size, out_size, in_channels, out_channels, ch_mult=2): + super().__init__() + assert out_size >= in_size + num_blocks = int(np.log2(out_size//in_size))+1 + factor_up = 1.+ (out_size % in_size) + print(f"Building {self.__class__.__name__} with in_size: {in_size} --> out_size {out_size} and factor {factor_up}") + self.rescaler = LatentRescaler(factor=factor_up, in_channels=in_channels, mid_channels=2*in_channels, + out_channels=in_channels) + self.decoder = Decoder(out_ch=out_channels, resolution=out_size, z_channels=in_channels, num_res_blocks=2, + attn_resolutions=[], in_channels=None, ch=in_channels, + ch_mult=[ch_mult for _ in range(num_blocks)]) + + def forward(self, x): + x = self.rescaler(x) + x = self.decoder(x) + return x + + +class Resize(nn.Module): + def __init__(self, in_channels=None, learned=False, mode="bilinear"): + super().__init__() + self.with_conv = learned + self.mode = mode + if self.with_conv: + print(f"Note: {self.__class__.__name} uses learned downsampling and will ignore the fixed {mode} mode") + raise NotImplementedError() + assert in_channels is not None + # no asymmetric padding in torch conv, must do it ourselves + self.conv = torch.nn.Conv2d(in_channels, + in_channels, + kernel_size=4, + stride=2, + padding=1) + + def forward(self, x, scale_factor=1.0): + if scale_factor==1.0: + return x + else: + x = torch.nn.functional.interpolate(x, mode=self.mode, align_corners=False, scale_factor=scale_factor) + return x + +class FirstStagePostProcessor(nn.Module): + + def __init__(self, ch_mult:list, in_channels, + pretrained_model:nn.Module=None, + reshape=False, + n_channels=None, + dropout=0., + pretrained_config=None): + super().__init__() + if pretrained_config is None: + assert pretrained_model is not None, 'Either "pretrained_model" or "pretrained_config" must not be None' + self.pretrained_model = pretrained_model + else: + assert pretrained_config is not None, 'Either "pretrained_model" or "pretrained_config" must not be None' + self.instantiate_pretrained(pretrained_config) + + self.do_reshape = reshape + + if n_channels is None: + n_channels = self.pretrained_model.encoder.ch + + self.proj_norm = Normalize(in_channels,num_groups=in_channels//2) + self.proj = nn.Conv2d(in_channels,n_channels,kernel_size=3, + stride=1,padding=1) + + blocks = [] + downs = [] + ch_in = n_channels + for m in ch_mult: + blocks.append(ResnetBlock(in_channels=ch_in,out_channels=m*n_channels,dropout=dropout)) + ch_in = m * n_channels + downs.append(Downsample(ch_in, with_conv=False)) + + self.model = nn.ModuleList(blocks) + self.downsampler = nn.ModuleList(downs) + + + def instantiate_pretrained(self, config): + model = instantiate_from_config(config) + self.pretrained_model = model.eval() + # self.pretrained_model.train = False + for param in self.pretrained_model.parameters(): + param.requires_grad = False + + + @torch.no_grad() + def encode_with_pretrained(self,x): + c = self.pretrained_model.encode(x) + if isinstance(c, DiagonalGaussianDistribution): + c = c.mode() + return c + + def forward(self,x): + z_fs = self.encode_with_pretrained(x) + z = self.proj_norm(z_fs) + z = self.proj(z) + z = nonlinearity(z) + + for submodel, downmodel in zip(self.model,self.downsampler): + z = submodel(z,temb=None) + z = downmodel(z) + + if self.do_reshape: + z = rearrange(z,'b c h w -> b (h w) c') + return z + diff --git a/ruyi/vae/ldm/modules/diffusionmodules/util.py b/ruyi/vae/ldm/modules/diffusionmodules/util.py new file mode 100644 index 0000000..54b198c --- /dev/null +++ b/ruyi/vae/ldm/modules/diffusionmodules/util.py @@ -0,0 +1,268 @@ +# adopted from +# https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py +# and +# https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +# and +# https://github.com/openai/guided-diffusion/blob/0ba878e517b276c45d1195eb29f6f5f72659a05b/guided_diffusion/nn.py +# +# thanks! + + +import math +import os + +import numpy as np +import torch +import torch.nn as nn +from einops import repeat + +from ...util import instantiate_from_config + + +def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): + if schedule == "linear": + betas = ( + torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2 + ) + + elif schedule == "cosine": + timesteps = ( + torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s + ) + alphas = timesteps / (1 + cosine_s) * np.pi / 2 + alphas = torch.cos(alphas).pow(2) + alphas = alphas / alphas[0] + betas = 1 - alphas[1:] / alphas[:-1] + betas = np.clip(betas, a_min=0, a_max=0.999) + + elif schedule == "sqrt_linear": + betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) + elif schedule == "sqrt": + betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5 + else: + raise ValueError(f"schedule '{schedule}' unknown.") + return betas.numpy() + + +def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True): + if ddim_discr_method == 'uniform': + c = num_ddpm_timesteps // num_ddim_timesteps + ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c))) + elif ddim_discr_method == 'quad': + ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int) + else: + raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"') + + # assert ddim_timesteps.shape[0] == num_ddim_timesteps + # add one to get the final alpha values right (the ones from first scale to data during sampling) + steps_out = ddim_timesteps + 1 + if verbose: + print(f'Selected timesteps for ddim sampler: {steps_out}') + return steps_out + + +def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True): + # select alphas for computing the variance schedule + alphas = alphacums[ddim_timesteps] + alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist()) + + # according the the formula provided in https://arxiv.org/abs/2010.02502 + sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev)) + if verbose: + print(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}') + print(f'For the chosen value of eta, which is {eta}, ' + f'this results in the following sigma_t schedule for ddim sampler {sigmas}') + return sigmas, alphas, alphas_prev + + +def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): + """ + Create a beta schedule that discretizes the given alpha_t_bar function, + which defines the cumulative product of (1-beta) over time from t = [0,1]. + :param num_diffusion_timesteps: the number of betas to produce. + :param alpha_bar: a lambda that takes an argument t from 0 to 1 and + produces the cumulative product of (1-beta) up to that + part of the diffusion process. + :param max_beta: the maximum beta to use; use values lower than 1 to + prevent singularities. + """ + betas = [] + for i in range(num_diffusion_timesteps): + t1 = i / num_diffusion_timesteps + t2 = (i + 1) / num_diffusion_timesteps + betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) + return np.array(betas) + + +def extract_into_tensor(a, t, x_shape): + b, *_ = t.shape + out = a.gather(-1, t) + return out.reshape(b, *((1,) * (len(x_shape) - 1))) + + +def checkpoint(func, inputs, params, flag): + """ + Evaluate a function without caching intermediate activations, allowing for + reduced memory at the expense of extra compute in the backward pass. + :param func: the function to evaluate. + :param inputs: the argument sequence to pass to `func`. + :param params: a sequence of parameters `func` depends on but does not + explicitly take as arguments. + :param flag: if False, disable gradient checkpointing. + """ + if flag: + args = tuple(inputs) + tuple(params) + return CheckpointFunction.apply(func, len(inputs), *args) + else: + return func(*inputs) + + +class CheckpointFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, run_function, length, *args): + ctx.run_function = run_function + ctx.input_tensors = list(args[:length]) + ctx.input_params = list(args[length:]) + + with torch.no_grad(): + output_tensors = ctx.run_function(*ctx.input_tensors) + return output_tensors + + @staticmethod + def backward(ctx, *output_grads): + ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] + with torch.enable_grad(): + # Fixes a bug where the first op in run_function modifies the + # Tensor storage in place, which is not allowed for detach()'d + # Tensors. + shallow_copies = [x.view_as(x) for x in ctx.input_tensors] + output_tensors = ctx.run_function(*shallow_copies) + input_grads = torch.autograd.grad( + output_tensors, + ctx.input_tensors + ctx.input_params, + output_grads, + allow_unused=True, + ) + del ctx.input_tensors + del ctx.input_params + del output_tensors + return (None, None) + input_grads + + +def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False): + """ + Create sinusoidal timestep embeddings. + :param timesteps: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an [N x dim] Tensor of positional embeddings. + """ + if not repeat_only: + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half + ).to(device=timesteps.device) + args = timesteps[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + else: + embedding = repeat(timesteps, 'b -> b d', d=dim) + return embedding + + +def zero_module(module): + """ + Zero out the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().zero_() + return module + + +def scale_module(module, scale): + """ + Scale the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().mul_(scale) + return module + + +def mean_flat(tensor): + """ + Take the mean over all non-batch dimensions. + """ + return tensor.mean(dim=list(range(1, len(tensor.shape)))) + + +def normalization(channels): + """ + Make a standard normalization layer. + :param channels: number of input channels. + :return: an nn.Module for normalization. + """ + return GroupNorm32(32, channels) + + +# PyTorch 1.7 has SiLU, but we support PyTorch 1.5. +class SiLU(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + + +class GroupNorm32(nn.GroupNorm): + def forward(self, x): + return super().forward(x.float()).type(x.dtype) + +def conv_nd(dims, *args, **kwargs): + """ + Create a 1D, 2D, or 3D convolution module. + """ + if dims == 1: + return nn.Conv1d(*args, **kwargs) + elif dims == 2: + return nn.Conv2d(*args, **kwargs) + elif dims == 3: + return nn.Conv3d(*args, **kwargs) + raise ValueError(f"unsupported dimensions: {dims}") + + +def linear(*args, **kwargs): + """ + Create a linear module. + """ + return nn.Linear(*args, **kwargs) + + +def avg_pool_nd(dims, *args, **kwargs): + """ + Create a 1D, 2D, or 3D average pooling module. + """ + if dims == 1: + return nn.AvgPool1d(*args, **kwargs) + elif dims == 2: + return nn.AvgPool2d(*args, **kwargs) + elif dims == 3: + return nn.AvgPool3d(*args, **kwargs) + raise ValueError(f"unsupported dimensions: {dims}") + + +class HybridConditioner(nn.Module): + + def __init__(self, c_concat_config, c_crossattn_config): + super().__init__() + self.concat_conditioner = instantiate_from_config(c_concat_config) + self.crossattn_conditioner = instantiate_from_config(c_crossattn_config) + + def forward(self, c_concat, c_crossattn): + c_concat = self.concat_conditioner(c_concat) + c_crossattn = self.crossattn_conditioner(c_crossattn) + return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]} + + +def noise_like(shape, device, repeat=False): + repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1))) + noise = lambda: torch.randn(shape, device=device) + return repeat_noise() if repeat else noise() \ No newline at end of file diff --git a/ruyi/vae/ldm/modules/distributions/__init__.py b/ruyi/vae/ldm/modules/distributions/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ruyi/vae/ldm/modules/distributions/distributions.py b/ruyi/vae/ldm/modules/distributions/distributions.py new file mode 100644 index 0000000..aec0174 --- /dev/null +++ b/ruyi/vae/ldm/modules/distributions/distributions.py @@ -0,0 +1,92 @@ +import numpy as np +import torch + + +class AbstractDistribution: + def sample(self): + raise NotImplementedError() + + def mode(self): + raise NotImplementedError() + + +class DiracDistribution(AbstractDistribution): + def __init__(self, value): + self.value = value + + def sample(self): + return self.value + + def mode(self): + return self.value + + +class DiagonalGaussianDistribution(object): + def __init__(self, parameters, deterministic=False): + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device) + + def sample(self): + x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device) + return x + + def kl(self, other=None): + if self.deterministic: + return torch.Tensor([0.]) + else: + if other is None: + return 0.5 * torch.sum(torch.pow(self.mean, 2) + + self.var - 1.0 - self.logvar, + dim=[1, 2, 3]) + else: + return 0.5 * torch.sum( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var - 1.0 - self.logvar + other.logvar, + dim=[1, 2, 3]) + + def nll(self, sample, dims=[1,2,3]): + if self.deterministic: + return torch.Tensor([0.]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum( + logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, + dim=dims) + + def mode(self): + return self.mean + + +def normal_kl(mean1, logvar1, mean2, logvar2): + """ + source: https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/losses.py#L12 + Compute the KL divergence between two gaussians. + Shapes are automatically broadcasted, so batches can be compared to + scalars, among other use cases. + """ + tensor = None + for obj in (mean1, logvar1, mean2, logvar2): + if isinstance(obj, torch.Tensor): + tensor = obj + break + assert tensor is not None, "at least one argument must be a Tensor" + + # Force variances to be Tensors. Broadcasting helps convert scalars to + # Tensors, but it does not work for torch.exp(). + logvar1, logvar2 = [ + x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor) + for x in (logvar1, logvar2) + ] + + return 0.5 * ( + -1.0 + + logvar2 + - logvar1 + + torch.exp(logvar1 - logvar2) + + ((mean1 - mean2) ** 2) * torch.exp(-logvar2) + ) diff --git a/ruyi/vae/ldm/modules/ema.py b/ruyi/vae/ldm/modules/ema.py new file mode 100644 index 0000000..3657208 --- /dev/null +++ b/ruyi/vae/ldm/modules/ema.py @@ -0,0 +1,115 @@ +#-*- encoding:utf-8 -*- +import torch +from torch import nn +from pytorch_lightning.callbacks import Callback + +class LitEma(nn.Module): + def __init__(self, model, decay=0.9999, use_num_upates=True): + super().__init__() + if decay < 0.0 or decay > 1.0: + raise ValueError('Decay must be between 0 and 1') + + self.m_name2s_name = {} + self.register_buffer('decay', torch.tensor(decay, dtype=torch.float32)) + self.register_buffer('num_updates', torch.tensor(0,dtype=torch.int) if use_num_upates + else torch.tensor(-1,dtype=torch.int)) + + for name, p in model.named_parameters(): + if p.requires_grad: + #remove as '.'-character is not allowed in buffers + s_name = name.replace('.','') + self.m_name2s_name.update({name:s_name}) + self.register_buffer(s_name,p.clone().detach().data) + + self.collected_params = [] + + def forward(self,model): + decay = self.decay + + if self.num_updates >= 0: + self.num_updates += 1 + decay = min(self.decay,(1 + self.num_updates) / (10 + self.num_updates)) + + one_minus_decay = 1.0 - decay + + with torch.no_grad(): + m_param = dict(model.named_parameters()) + shadow_params = dict(self.named_buffers()) + + for key in m_param: + if m_param[key].requires_grad: + sname = self.m_name2s_name[key] + shadow_params[sname] = shadow_params[sname].type_as(m_param[key]) + shadow_params[sname].sub_(one_minus_decay * (shadow_params[sname] - m_param[key])) + else: + assert not key in self.m_name2s_name + + def copy_to(self, model): + m_param = dict(model.named_parameters()) + shadow_params = dict(self.named_buffers()) + for key in m_param: + if m_param[key].requires_grad: + m_param[key].data.copy_(shadow_params[self.m_name2s_name[key]].data) + else: + assert not key in self.m_name2s_name + + def store(self, parameters): + """ + Save the current parameters for restoring later. + Args: + parameters: Iterable of `torch.nn.Parameter`; the parameters to be + temporarily stored. + """ + self.collected_params = [param.clone() for param in parameters] + + def restore(self, parameters): + """ + Restore the parameters stored with the `store` method. + Useful to validate the model with EMA parameters without affecting the + original optimization process. Store the parameters before the + `copy_to` method. After validation (or model saving), use this to + restore the former parameters. + Args: + parameters: Iterable of `torch.nn.Parameter`; the parameters to be + updated with the stored parameters. + """ + for c_param, param in zip(self.collected_params, parameters): + param.data.copy_(c_param.data) + +class EMACallback(Callback): + def __init__(self, decay=0.9999): + self.decay = decay + self.shadow_params = {} + + def on_train_start(self, trainer, pl_module): + # initialize shadow parameters for original models + total_ema_cnt = 0 + for name, param in pl_module.named_parameters(): + if name not in self.shadow_params: + self.shadow_params[name] = param.data.clone() + else: # already in dict, maybe load from checkpoint + pass + print('will calc ema for param: %s' % name) + total_ema_cnt += 1 + print('total_ema_cnt=%d' % total_ema_cnt) + + def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): + # Update the shadow params at the end of each epoch + for name, param in pl_module.named_parameters(): + assert name in self.shadow_params + new_average = (1.0 - self.decay) * param.data + self.decay * self.shadow_params[name] + self.shadow_params[name] = new_average.clone() + + def on_save_checkpoint(self, trainer, pl_module, checkpoint): + # Save EMA parameters in the checkpoint + checkpoint['ema_params'] = self.shadow_params + + def on_load_checkpoint(self, trainer, pl_module, checkpoint): + # Restore EMA parameters from the checkpoint + if 'ema_params' in checkpoint: + self.shadow_params = checkpoint.get('ema_params', {}) + for k in self.shadow_params: + self.shadow_params[k] = self.shadow_params[k].cuda() + print('load shadow params from checkpoint, cnt=%d' % len(self.shadow_params)) + else: + print('ema_params is not in checkpoint') \ No newline at end of file diff --git a/ruyi/vae/ldm/modules/image_degradation/__init__.py b/ruyi/vae/ldm/modules/image_degradation/__init__.py new file mode 100644 index 0000000..e95d680 --- /dev/null +++ b/ruyi/vae/ldm/modules/image_degradation/__init__.py @@ -0,0 +1,3 @@ +from .bsrgan import degradation_bsrgan_variant as degradation_fn_bsr +from .bsrgan_light import \ + degradation_bsrgan_variant as degradation_fn_bsr_light diff --git a/ruyi/vae/ldm/modules/image_degradation/bsrgan.py b/ruyi/vae/ldm/modules/image_degradation/bsrgan.py new file mode 100644 index 0000000..b159d6f --- /dev/null +++ b/ruyi/vae/ldm/modules/image_degradation/bsrgan.py @@ -0,0 +1,732 @@ +# -*- coding: utf-8 -*- +""" +# -------------------------------------------- +# Super-Resolution +# -------------------------------------------- +# +# Kai Zhang (cskaizhang@gmail.com) +# https://github.com/cszn +# From 2019/03--2021/08 +# -------------------------------------------- +""" + +import random +from functools import partial + +import albumentations +import cv2 +import numpy as np +import scipy +import scipy.stats as ss +import torch +from scipy import ndimage +from scipy.interpolate import interp2d +from scipy.linalg import orth + +from . import utils_image as util + + +def modcrop_np(img, sf): + ''' + Args: + img: numpy image, WxH or WxHxC + sf: scale factor + Return: + cropped image + ''' + w, h = img.shape[:2] + im = np.copy(img) + return im[:w - w % sf, :h - h % sf, ...] + + +""" +# -------------------------------------------- +# anisotropic Gaussian kernels +# -------------------------------------------- +""" + + +def analytic_kernel(k): + """Calculate the X4 kernel from the X2 kernel (for proof see appendix in paper)""" + k_size = k.shape[0] + # Calculate the big kernels size + big_k = np.zeros((3 * k_size - 2, 3 * k_size - 2)) + # Loop over the small kernel to fill the big one + for r in range(k_size): + for c in range(k_size): + big_k[2 * r:2 * r + k_size, 2 * c:2 * c + k_size] += k[r, c] * k + # Crop the edges of the big kernel to ignore very small values and increase run time of SR + crop = k_size // 2 + cropped_big_k = big_k[crop:-crop, crop:-crop] + # Normalize to 1 + return cropped_big_k / cropped_big_k.sum() + + +def anisotropic_Gaussian(ksize=15, theta=np.pi, l1=6, l2=6): + """ generate an anisotropic Gaussian kernel + Args: + ksize : e.g., 15, kernel size + theta : [0, pi], rotation angle range + l1 : [0.1,50], scaling of eigenvalues + l2 : [0.1,l1], scaling of eigenvalues + If l1 = l2, will get an isotropic Gaussian kernel. + Returns: + k : kernel + """ + + v = np.dot(np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]), np.array([1., 0.])) + V = np.array([[v[0], v[1]], [v[1], -v[0]]]) + D = np.array([[l1, 0], [0, l2]]) + Sigma = np.dot(np.dot(V, D), np.linalg.inv(V)) + k = gm_blur_kernel(mean=[0, 0], cov=Sigma, size=ksize) + + return k + + +def gm_blur_kernel(mean, cov, size=15): + center = size / 2.0 + 0.5 + k = np.zeros([size, size]) + for y in range(size): + for x in range(size): + cy = y - center + 1 + cx = x - center + 1 + k[y, x] = ss.multivariate_normal.pdf([cx, cy], mean=mean, cov=cov) + + k = k / np.sum(k) + return k + + +def shift_pixel(x, sf, upper_left=True): + """shift pixel for super-resolution with different scale factors + Args: + x: WxHxC or WxH + sf: scale factor + upper_left: shift direction + """ + h, w = x.shape[:2] + shift = (sf - 1) * 0.5 + xv, yv = np.arange(0, w, 1.0), np.arange(0, h, 1.0) + if upper_left: + x1 = xv + shift + y1 = yv + shift + else: + x1 = xv - shift + y1 = yv - shift + + x1 = np.clip(x1, 0, w - 1) + y1 = np.clip(y1, 0, h - 1) + + if x.ndim == 2: + x = interp2d(xv, yv, x)(x1, y1) + if x.ndim == 3: + for i in range(x.shape[-1]): + x[:, :, i] = interp2d(xv, yv, x[:, :, i])(x1, y1) + + return x + + +def blur(x, k): + ''' + x: image, NxcxHxW + k: kernel, Nx1xhxw + ''' + n, c = x.shape[:2] + p1, p2 = (k.shape[-2] - 1) // 2, (k.shape[-1] - 1) // 2 + x = torch.nn.functional.pad(x, pad=(p1, p2, p1, p2), mode='replicate') + k = k.repeat(1, c, 1, 1) + k = k.view(-1, 1, k.shape[2], k.shape[3]) + x = x.view(1, -1, x.shape[2], x.shape[3]) + x = torch.nn.functional.conv2d(x, k, bias=None, stride=1, padding=0, groups=n * c) + x = x.view(n, c, x.shape[2], x.shape[3]) + + return x + + +def gen_kernel(k_size=np.array([15, 15]), scale_factor=np.array([4, 4]), min_var=0.6, max_var=10., noise_level=0): + """" + # modified version of https://github.com/assafshocher/BlindSR_dataset_generator + # Kai Zhang + # min_var = 0.175 * sf # variance of the gaussian kernel will be sampled between min_var and max_var + # max_var = 2.5 * sf + """ + # Set random eigen-vals (lambdas) and angle (theta) for COV matrix + lambda_1 = min_var + np.random.rand() * (max_var - min_var) + lambda_2 = min_var + np.random.rand() * (max_var - min_var) + theta = np.random.rand() * np.pi # random theta + noise = -noise_level + np.random.rand(*k_size) * noise_level * 2 + + # Set COV matrix using Lambdas and Theta + LAMBDA = np.diag([lambda_1, lambda_2]) + Q = np.array([[np.cos(theta), -np.sin(theta)], + [np.sin(theta), np.cos(theta)]]) + SIGMA = Q @ LAMBDA @ Q.T + INV_SIGMA = np.linalg.inv(SIGMA)[None, None, :, :] + + # Set expectation position (shifting kernel for aligned image) + MU = k_size // 2 - 0.5 * (scale_factor - 1) # - 0.5 * (scale_factor - k_size % 2) + MU = MU[None, None, :, None] + + # Create meshgrid for Gaussian + [X, Y] = np.meshgrid(range(k_size[0]), range(k_size[1])) + Z = np.stack([X, Y], 2)[:, :, :, None] + + # Calcualte Gaussian for every pixel of the kernel + ZZ = Z - MU + ZZ_t = ZZ.transpose(0, 1, 3, 2) + raw_kernel = np.exp(-0.5 * np.squeeze(ZZ_t @ INV_SIGMA @ ZZ)) * (1 + noise) + + # shift the kernel so it will be centered + # raw_kernel_centered = kernel_shift(raw_kernel, scale_factor) + + # Normalize the kernel and return + # kernel = raw_kernel_centered / np.sum(raw_kernel_centered) + kernel = raw_kernel / np.sum(raw_kernel) + return kernel + + +def fspecial_gaussian(hsize, sigma): + hsize = [hsize, hsize] + siz = [(hsize[0] - 1.0) / 2.0, (hsize[1] - 1.0) / 2.0] + std = sigma + [x, y] = np.meshgrid(np.arange(-siz[1], siz[1] + 1), np.arange(-siz[0], siz[0] + 1)) + arg = -(x * x + y * y) / (2 * std * std) + h = np.exp(arg) + h[h < scipy.finfo(float).eps * h.max()] = 0 + sumh = h.sum() + if sumh != 0: + h = h / sumh + return h + + +def fspecial_laplacian(alpha): + alpha = max([0, min([alpha, 1])]) + h1 = alpha / (alpha + 1) + h2 = (1 - alpha) / (alpha + 1) + h = [[h1, h2, h1], [h2, -4 / (alpha + 1), h2], [h1, h2, h1]] + h = np.array(h) + return h + + +def fspecial(filter_type, *args, **kwargs): + ''' + python code from: + https://github.com/ronaldosena/imagens-medicas-2/blob/40171a6c259edec7827a6693a93955de2bd39e76/Aulas/aula_2_-_uniform_filter/matlab_fspecial.py + ''' + if filter_type == 'gaussian': + return fspecial_gaussian(*args, **kwargs) + if filter_type == 'laplacian': + return fspecial_laplacian(*args, **kwargs) + + +""" +# -------------------------------------------- +# degradation models +# -------------------------------------------- +""" + + +def bicubic_degradation(x, sf=3): + ''' + Args: + x: HxWxC image, [0, 1] + sf: down-scale factor + Return: + bicubicly downsampled LR image + ''' + x = util.imresize_np(x, scale=1 / sf) + return x + + +def srmd_degradation(x, k, sf=3): + ''' blur + bicubic downsampling + Args: + x: HxWxC image, [0, 1] + k: hxw, double + sf: down-scale factor + Return: + downsampled LR image + Reference: + @inproceedings{zhang2018learning, + title={Learning a single convolutional super-resolution network for multiple degradations}, + author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei}, + booktitle={IEEE Conference on Computer Vision and Pattern Recognition}, + pages={3262--3271}, + year={2018} + } + ''' + x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap') # 'nearest' | 'mirror' + x = bicubic_degradation(x, sf=sf) + return x + + +def dpsr_degradation(x, k, sf=3): + ''' bicubic downsampling + blur + Args: + x: HxWxC image, [0, 1] + k: hxw, double + sf: down-scale factor + Return: + downsampled LR image + Reference: + @inproceedings{zhang2019deep, + title={Deep Plug-and-Play Super-Resolution for Arbitrary Blur Kernels}, + author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei}, + booktitle={IEEE Conference on Computer Vision and Pattern Recognition}, + pages={1671--1681}, + year={2019} + } + ''' + x = bicubic_degradation(x, sf=sf) + x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap') + return x + + +def classical_degradation(x, k, sf=3): + ''' blur + downsampling + Args: + x: HxWxC image, [0, 1]/[0, 255] + k: hxw, double + sf: down-scale factor + Return: + downsampled LR image + ''' + x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap') + # x = filters.correlate(x, np.expand_dims(np.flip(k), axis=2)) + st = 0 + return x[st::sf, st::sf, ...] + + +def add_sharpening(img, weight=0.5, radius=50, threshold=10): + """USM sharpening. borrowed from real-ESRGAN + Input image: I; Blurry image: B. + 1. K = I + weight * (I - B) + 2. Mask = 1 if abs(I - B) > threshold, else: 0 + 3. Blur mask: + 4. Out = Mask * K + (1 - Mask) * I + Args: + img (Numpy array): Input image, HWC, BGR; float32, [0, 1]. + weight (float): Sharp weight. Default: 1. + radius (float): Kernel size of Gaussian blur. Default: 50. + threshold (int): + """ + if radius % 2 == 0: + radius += 1 + blur = cv2.GaussianBlur(img, (radius, radius), 0) + residual = img - blur + mask = np.abs(residual) * 255 > threshold + mask = mask.astype('float32') + soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0) + + K = img + weight * residual + K = np.clip(K, 0, 1) + return soft_mask * K + (1 - soft_mask) * img + + +def add_blur(img, sf=4): + wd2 = 4.0 + sf + wd = 2.0 + 0.2 * sf + if random.random() < 0.5: + l1 = wd2 * random.random() + l2 = wd2 * random.random() + k = anisotropic_Gaussian(ksize=2 * random.randint(2, 11) + 3, theta=random.random() * np.pi, l1=l1, l2=l2) + else: + k = fspecial('gaussian', 2 * random.randint(2, 11) + 3, wd * random.random()) + img = ndimage.filters.convolve(img, np.expand_dims(k, axis=2), mode='mirror') + + return img + + +def add_resize(img, sf=4): + rnum = np.random.rand() + if rnum > 0.8: # up + sf1 = random.uniform(1, 2) + elif rnum < 0.7: # down + sf1 = random.uniform(0.5 / sf, 1) + else: + sf1 = 1.0 + img = cv2.resize(img, (int(sf1 * img.shape[1]), int(sf1 * img.shape[0])), interpolation=random.choice([1, 2, 3])) + img = np.clip(img, 0.0, 1.0) + + return img + + +# def add_Gaussian_noise(img, noise_level1=2, noise_level2=25): +# noise_level = random.randint(noise_level1, noise_level2) +# rnum = np.random.rand() +# if rnum > 0.6: # add color Gaussian noise +# img += np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32) +# elif rnum < 0.4: # add grayscale Gaussian noise +# img += np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32) +# else: # add noise +# L = noise_level2 / 255. +# D = np.diag(np.random.rand(3)) +# U = orth(np.random.rand(3, 3)) +# conv = np.dot(np.dot(np.transpose(U), D), U) +# img += np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32) +# img = np.clip(img, 0.0, 1.0) +# return img + +def add_Gaussian_noise(img, noise_level1=2, noise_level2=25): + noise_level = random.randint(noise_level1, noise_level2) + rnum = np.random.rand() + if rnum > 0.6: # add color Gaussian noise + img = img + np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32) + elif rnum < 0.4: # add grayscale Gaussian noise + img = img + np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32) + else: # add noise + L = noise_level2 / 255. + D = np.diag(np.random.rand(3)) + U = orth(np.random.rand(3, 3)) + conv = np.dot(np.dot(np.transpose(U), D), U) + img = img + np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32) + img = np.clip(img, 0.0, 1.0) + return img + + +def add_speckle_noise(img, noise_level1=2, noise_level2=25): + noise_level = random.randint(noise_level1, noise_level2) + img = np.clip(img, 0.0, 1.0) + rnum = random.random() + if rnum > 0.6: + img += img * np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32) + elif rnum < 0.4: + img += img * np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32) + else: + L = noise_level2 / 255. + D = np.diag(np.random.rand(3)) + U = orth(np.random.rand(3, 3)) + conv = np.dot(np.dot(np.transpose(U), D), U) + img += img * np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32) + img = np.clip(img, 0.0, 1.0) + return img + + +def add_Poisson_noise(img): + img = np.clip((img * 255.0).round(), 0, 255) / 255. + vals = 10 ** (2 * random.random() + 2.0) # [2, 4] + if random.random() < 0.5: + img = np.random.poisson(img * vals).astype(np.float32) / vals + else: + img_gray = np.dot(img[..., :3], [0.299, 0.587, 0.114]) + img_gray = np.clip((img_gray * 255.0).round(), 0, 255) / 255. + noise_gray = np.random.poisson(img_gray * vals).astype(np.float32) / vals - img_gray + img += noise_gray[:, :, np.newaxis] + img = np.clip(img, 0.0, 1.0) + return img + + +def add_JPEG_noise(img): + quality_factor = random.randint(30, 95) + img = cv2.cvtColor(util.single2uint(img), cv2.COLOR_RGB2BGR) + result, encimg = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), quality_factor]) + img = cv2.imdecode(encimg, 1) + img = cv2.cvtColor(util.uint2single(img), cv2.COLOR_BGR2RGB) + return img + + +def random_crop(lq, hq, sf=4, lq_patchsize=64): + h, w = lq.shape[:2] + rnd_h = random.randint(0, h - lq_patchsize) + rnd_w = random.randint(0, w - lq_patchsize) + lq = lq[rnd_h:rnd_h + lq_patchsize, rnd_w:rnd_w + lq_patchsize, :] + + rnd_h_H, rnd_w_H = int(rnd_h * sf), int(rnd_w * sf) + hq = hq[rnd_h_H:rnd_h_H + lq_patchsize * sf, rnd_w_H:rnd_w_H + lq_patchsize * sf, :] + return lq, hq + + +def degradation_bsrgan(img, sf=4, lq_patchsize=72, isp_model=None): + """ + This is the degradation model of BSRGAN from the paper + "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution" + ---------- + img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf) + sf: scale factor + isp_model: camera ISP model + Returns + ------- + img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1] + hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1] + """ + isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25 + sf_ori = sf + + h1, w1 = img.shape[:2] + img = img.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop + h, w = img.shape[:2] + + if h < lq_patchsize * sf or w < lq_patchsize * sf: + raise ValueError(f'img size ({h1}X{w1}) is too small!') + + hq = img.copy() + + if sf == 4 and random.random() < scale2_prob: # downsample1 + if np.random.rand() < 0.5: + img = cv2.resize(img, (int(1 / 2 * img.shape[1]), int(1 / 2 * img.shape[0])), + interpolation=random.choice([1, 2, 3])) + else: + img = util.imresize_np(img, 1 / 2, True) + img = np.clip(img, 0.0, 1.0) + sf = 2 + + shuffle_order = random.sample(range(7), 7) + idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3) + if idx1 > idx2: # keep downsample3 last + shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1] + + for i in shuffle_order: + + if i == 0: + img = add_blur(img, sf=sf) + + elif i == 1: + img = add_blur(img, sf=sf) + + elif i == 2: + a, b = img.shape[1], img.shape[0] + # downsample2 + if random.random() < 0.75: + sf1 = random.uniform(1, 2 * sf) + img = cv2.resize(img, (int(1 / sf1 * img.shape[1]), int(1 / sf1 * img.shape[0])), + interpolation=random.choice([1, 2, 3])) + else: + k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf)) + k_shifted = shift_pixel(k, sf) + k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel + img = ndimage.filters.convolve(img, np.expand_dims(k_shifted, axis=2), mode='mirror') + img = img[0::sf, 0::sf, ...] # nearest downsampling + img = np.clip(img, 0.0, 1.0) + + elif i == 3: + # downsample3 + img = cv2.resize(img, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3])) + img = np.clip(img, 0.0, 1.0) + + elif i == 4: + # add Gaussian noise + img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25) + + elif i == 5: + # add JPEG noise + if random.random() < jpeg_prob: + img = add_JPEG_noise(img) + + elif i == 6: + # add processed camera sensor noise + if random.random() < isp_prob and isp_model is not None: + with torch.no_grad(): + img, hq = isp_model.forward(img.copy(), hq) + + # add final JPEG compression noise + img = add_JPEG_noise(img) + + # random crop + img, hq = random_crop(img, hq, sf_ori, lq_patchsize) + + return img, hq + + +# todo no isp_model? +def degradation_bsrgan_variant(image, sf=4, isp_model=None): + """ + This is the degradation model of BSRGAN from the paper + "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution" + ---------- + sf: scale factor + isp_model: camera ISP model + Returns + ------- + img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1] + hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1] + """ + image = util.uint2single(image) + isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25 + sf_ori = sf + + h1, w1 = image.shape[:2] + image = image.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop + h, w = image.shape[:2] + + hq = image.copy() + + if sf == 4 and random.random() < scale2_prob: # downsample1 + if np.random.rand() < 0.5: + image = cv2.resize(image, (int(1 / 2 * image.shape[1]), int(1 / 2 * image.shape[0])), + interpolation=random.choice([1, 2, 3])) + else: + image = util.imresize_np(image, 1 / 2, True) + image = np.clip(image, 0.0, 1.0) + sf = 2 + + shuffle_order = random.sample(range(7), 7) + idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3) + if idx1 > idx2: # keep downsample3 last + shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1] + + for i in shuffle_order: + + if i == 0: + image = add_blur(image, sf=sf) + + elif i == 1: + image = add_blur(image, sf=sf) + + elif i == 2: + a, b = image.shape[1], image.shape[0] + # downsample2 + if random.random() < 0.75: + sf1 = random.uniform(1, 2 * sf) + image = cv2.resize(image, (int(1 / sf1 * image.shape[1]), int(1 / sf1 * image.shape[0])), + interpolation=random.choice([1, 2, 3])) + else: + k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf)) + k_shifted = shift_pixel(k, sf) + k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel + image = ndimage.filters.convolve(image, np.expand_dims(k_shifted, axis=2), mode='mirror') + image = image[0::sf, 0::sf, ...] # nearest downsampling + image = np.clip(image, 0.0, 1.0) + + elif i == 3: + # downsample3 + image = cv2.resize(image, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3])) + image = np.clip(image, 0.0, 1.0) + + elif i == 4: + # add Gaussian noise + image = add_Gaussian_noise(image, noise_level1=2, noise_level2=25) + + elif i == 5: + # add JPEG noise + if random.random() < jpeg_prob: + image = add_JPEG_noise(image) + + # elif i == 6: + # # add processed camera sensor noise + # if random.random() < isp_prob and isp_model is not None: + # with torch.no_grad(): + # img, hq = isp_model.forward(img.copy(), hq) + + # add final JPEG compression noise + image = add_JPEG_noise(image) + image = util.single2uint(image) + example = {"image":image} + return example + + +# TODO incase there is a pickle error one needs to replace a += x with a = a + x in add_speckle_noise etc... +def degradation_bsrgan_plus(img, sf=4, shuffle_prob=0.5, use_sharp=True, lq_patchsize=64, isp_model=None): + """ + This is an extended degradation model by combining + the degradation models of BSRGAN and Real-ESRGAN + ---------- + img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf) + sf: scale factor + use_shuffle: the degradation shuffle + use_sharp: sharpening the img + Returns + ------- + img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1] + hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1] + """ + + h1, w1 = img.shape[:2] + img = img.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop + h, w = img.shape[:2] + + if h < lq_patchsize * sf or w < lq_patchsize * sf: + raise ValueError(f'img size ({h1}X{w1}) is too small!') + + if use_sharp: + img = add_sharpening(img) + hq = img.copy() + + if random.random() < shuffle_prob: + shuffle_order = random.sample(range(13), 13) + else: + shuffle_order = list(range(13)) + # local shuffle for noise, JPEG is always the last one + shuffle_order[2:6] = random.sample(shuffle_order[2:6], len(range(2, 6))) + shuffle_order[9:13] = random.sample(shuffle_order[9:13], len(range(9, 13))) + + poisson_prob, speckle_prob, isp_prob = 0.1, 0.1, 0.1 + + for i in shuffle_order: + if i == 0: + img = add_blur(img, sf=sf) + elif i == 1: + img = add_resize(img, sf=sf) + elif i == 2: + img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25) + elif i == 3: + if random.random() < poisson_prob: + img = add_Poisson_noise(img) + elif i == 4: + if random.random() < speckle_prob: + img = add_speckle_noise(img) + elif i == 5: + if random.random() < isp_prob and isp_model is not None: + with torch.no_grad(): + img, hq = isp_model.forward(img.copy(), hq) + elif i == 6: + img = add_JPEG_noise(img) + elif i == 7: + img = add_blur(img, sf=sf) + elif i == 8: + img = add_resize(img, sf=sf) + elif i == 9: + img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25) + elif i == 10: + if random.random() < poisson_prob: + img = add_Poisson_noise(img) + elif i == 11: + if random.random() < speckle_prob: + img = add_speckle_noise(img) + elif i == 12: + if random.random() < isp_prob and isp_model is not None: + with torch.no_grad(): + img, hq = isp_model.forward(img.copy(), hq) + else: + print('check the shuffle!') + + # resize to desired size + img = cv2.resize(img, (int(1 / sf * hq.shape[1]), int(1 / sf * hq.shape[0])), + interpolation=random.choice([1, 2, 3])) + + # add final JPEG compression noise + img = add_JPEG_noise(img) + + # random crop + img, hq = random_crop(img, hq, sf, lq_patchsize) + + return img, hq + + +if __name__ == '__main__': + ''' + print("hey") + img = util.imread_uint('utils/test.png', 3) + print(img) + img = util.uint2single(img) + print(img) + img = img[:448, :448] + h = img.shape[0] // 4 + print("resizing to", h) + sf = 4 + deg_fn = partial(degradation_bsrgan_variant, sf=sf) + for i in range(20): + print(i) + img_lq = deg_fn(img) + print(img_lq) + img_lq_bicubic = albumentations.SmallestMaxSize(max_size=h, interpolation=cv2.INTER_CUBIC)(image=img)["image"] + print(img_lq.shape) + print("bicubic", img_lq_bicubic.shape) + print(img_hq.shape) + lq_nearest = cv2.resize(util.single2uint(img_lq), (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])), + interpolation=0) + lq_bicubic_nearest = cv2.resize(util.single2uint(img_lq_bicubic), (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])), + interpolation=0) + img_concat = np.concatenate([lq_bicubic_nearest, lq_nearest, util.single2uint(img_hq)], axis=1) + util.imsave(img_concat, str(i) + '.png') + ''' + + diff --git a/ruyi/vae/ldm/modules/image_degradation/bsrgan_light.py b/ruyi/vae/ldm/modules/image_degradation/bsrgan_light.py new file mode 100644 index 0000000..2b5be11 --- /dev/null +++ b/ruyi/vae/ldm/modules/image_degradation/bsrgan_light.py @@ -0,0 +1,650 @@ +# -*- coding: utf-8 -*- +import random +from functools import partial + +import albumentations +import cv2 +import numpy as np +import scipy +import scipy.stats as ss +import torch +from scipy import ndimage +from scipy.interpolate import interp2d +from scipy.linalg import orth + +from . import utils_image as util + +""" +# -------------------------------------------- +# Super-Resolution +# -------------------------------------------- +# +# Kai Zhang (cskaizhang@gmail.com) +# https://github.com/cszn +# From 2019/03--2021/08 +# -------------------------------------------- +""" + + +def modcrop_np(img, sf): + ''' + Args: + img: numpy image, WxH or WxHxC + sf: scale factor + Return: + cropped image + ''' + w, h = img.shape[:2] + im = np.copy(img) + return im[:w - w % sf, :h - h % sf, ...] + + +""" +# -------------------------------------------- +# anisotropic Gaussian kernels +# -------------------------------------------- +""" + + +def analytic_kernel(k): + """Calculate the X4 kernel from the X2 kernel (for proof see appendix in paper)""" + k_size = k.shape[0] + # Calculate the big kernels size + big_k = np.zeros((3 * k_size - 2, 3 * k_size - 2)) + # Loop over the small kernel to fill the big one + for r in range(k_size): + for c in range(k_size): + big_k[2 * r:2 * r + k_size, 2 * c:2 * c + k_size] += k[r, c] * k + # Crop the edges of the big kernel to ignore very small values and increase run time of SR + crop = k_size // 2 + cropped_big_k = big_k[crop:-crop, crop:-crop] + # Normalize to 1 + return cropped_big_k / cropped_big_k.sum() + + +def anisotropic_Gaussian(ksize=15, theta=np.pi, l1=6, l2=6): + """ generate an anisotropic Gaussian kernel + Args: + ksize : e.g., 15, kernel size + theta : [0, pi], rotation angle range + l1 : [0.1,50], scaling of eigenvalues + l2 : [0.1,l1], scaling of eigenvalues + If l1 = l2, will get an isotropic Gaussian kernel. + Returns: + k : kernel + """ + + v = np.dot(np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]), np.array([1., 0.])) + V = np.array([[v[0], v[1]], [v[1], -v[0]]]) + D = np.array([[l1, 0], [0, l2]]) + Sigma = np.dot(np.dot(V, D), np.linalg.inv(V)) + k = gm_blur_kernel(mean=[0, 0], cov=Sigma, size=ksize) + + return k + + +def gm_blur_kernel(mean, cov, size=15): + center = size / 2.0 + 0.5 + k = np.zeros([size, size]) + for y in range(size): + for x in range(size): + cy = y - center + 1 + cx = x - center + 1 + k[y, x] = ss.multivariate_normal.pdf([cx, cy], mean=mean, cov=cov) + + k = k / np.sum(k) + return k + + +def shift_pixel(x, sf, upper_left=True): + """shift pixel for super-resolution with different scale factors + Args: + x: WxHxC or WxH + sf: scale factor + upper_left: shift direction + """ + h, w = x.shape[:2] + shift = (sf - 1) * 0.5 + xv, yv = np.arange(0, w, 1.0), np.arange(0, h, 1.0) + if upper_left: + x1 = xv + shift + y1 = yv + shift + else: + x1 = xv - shift + y1 = yv - shift + + x1 = np.clip(x1, 0, w - 1) + y1 = np.clip(y1, 0, h - 1) + + if x.ndim == 2: + x = interp2d(xv, yv, x)(x1, y1) + if x.ndim == 3: + for i in range(x.shape[-1]): + x[:, :, i] = interp2d(xv, yv, x[:, :, i])(x1, y1) + + return x + + +def blur(x, k): + ''' + x: image, NxcxHxW + k: kernel, Nx1xhxw + ''' + n, c = x.shape[:2] + p1, p2 = (k.shape[-2] - 1) // 2, (k.shape[-1] - 1) // 2 + x = torch.nn.functional.pad(x, pad=(p1, p2, p1, p2), mode='replicate') + k = k.repeat(1, c, 1, 1) + k = k.view(-1, 1, k.shape[2], k.shape[3]) + x = x.view(1, -1, x.shape[2], x.shape[3]) + x = torch.nn.functional.conv2d(x, k, bias=None, stride=1, padding=0, groups=n * c) + x = x.view(n, c, x.shape[2], x.shape[3]) + + return x + + +def gen_kernel(k_size=np.array([15, 15]), scale_factor=np.array([4, 4]), min_var=0.6, max_var=10., noise_level=0): + """" + # modified version of https://github.com/assafshocher/BlindSR_dataset_generator + # Kai Zhang + # min_var = 0.175 * sf # variance of the gaussian kernel will be sampled between min_var and max_var + # max_var = 2.5 * sf + """ + # Set random eigen-vals (lambdas) and angle (theta) for COV matrix + lambda_1 = min_var + np.random.rand() * (max_var - min_var) + lambda_2 = min_var + np.random.rand() * (max_var - min_var) + theta = np.random.rand() * np.pi # random theta + noise = -noise_level + np.random.rand(*k_size) * noise_level * 2 + + # Set COV matrix using Lambdas and Theta + LAMBDA = np.diag([lambda_1, lambda_2]) + Q = np.array([[np.cos(theta), -np.sin(theta)], + [np.sin(theta), np.cos(theta)]]) + SIGMA = Q @ LAMBDA @ Q.T + INV_SIGMA = np.linalg.inv(SIGMA)[None, None, :, :] + + # Set expectation position (shifting kernel for aligned image) + MU = k_size // 2 - 0.5 * (scale_factor - 1) # - 0.5 * (scale_factor - k_size % 2) + MU = MU[None, None, :, None] + + # Create meshgrid for Gaussian + [X, Y] = np.meshgrid(range(k_size[0]), range(k_size[1])) + Z = np.stack([X, Y], 2)[:, :, :, None] + + # Calcualte Gaussian for every pixel of the kernel + ZZ = Z - MU + ZZ_t = ZZ.transpose(0, 1, 3, 2) + raw_kernel = np.exp(-0.5 * np.squeeze(ZZ_t @ INV_SIGMA @ ZZ)) * (1 + noise) + + # shift the kernel so it will be centered + # raw_kernel_centered = kernel_shift(raw_kernel, scale_factor) + + # Normalize the kernel and return + # kernel = raw_kernel_centered / np.sum(raw_kernel_centered) + kernel = raw_kernel / np.sum(raw_kernel) + return kernel + + +def fspecial_gaussian(hsize, sigma): + hsize = [hsize, hsize] + siz = [(hsize[0] - 1.0) / 2.0, (hsize[1] - 1.0) / 2.0] + std = sigma + [x, y] = np.meshgrid(np.arange(-siz[1], siz[1] + 1), np.arange(-siz[0], siz[0] + 1)) + arg = -(x * x + y * y) / (2 * std * std) + h = np.exp(arg) + h[h < scipy.finfo(float).eps * h.max()] = 0 + sumh = h.sum() + if sumh != 0: + h = h / sumh + return h + + +def fspecial_laplacian(alpha): + alpha = max([0, min([alpha, 1])]) + h1 = alpha / (alpha + 1) + h2 = (1 - alpha) / (alpha + 1) + h = [[h1, h2, h1], [h2, -4 / (alpha + 1), h2], [h1, h2, h1]] + h = np.array(h) + return h + + +def fspecial(filter_type, *args, **kwargs): + ''' + python code from: + https://github.com/ronaldosena/imagens-medicas-2/blob/40171a6c259edec7827a6693a93955de2bd39e76/Aulas/aula_2_-_uniform_filter/matlab_fspecial.py + ''' + if filter_type == 'gaussian': + return fspecial_gaussian(*args, **kwargs) + if filter_type == 'laplacian': + return fspecial_laplacian(*args, **kwargs) + + +""" +# -------------------------------------------- +# degradation models +# -------------------------------------------- +""" + + +def bicubic_degradation(x, sf=3): + ''' + Args: + x: HxWxC image, [0, 1] + sf: down-scale factor + Return: + bicubicly downsampled LR image + ''' + x = util.imresize_np(x, scale=1 / sf) + return x + + +def srmd_degradation(x, k, sf=3): + ''' blur + bicubic downsampling + Args: + x: HxWxC image, [0, 1] + k: hxw, double + sf: down-scale factor + Return: + downsampled LR image + Reference: + @inproceedings{zhang2018learning, + title={Learning a single convolutional super-resolution network for multiple degradations}, + author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei}, + booktitle={IEEE Conference on Computer Vision and Pattern Recognition}, + pages={3262--3271}, + year={2018} + } + ''' + x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap') # 'nearest' | 'mirror' + x = bicubic_degradation(x, sf=sf) + return x + + +def dpsr_degradation(x, k, sf=3): + ''' bicubic downsampling + blur + Args: + x: HxWxC image, [0, 1] + k: hxw, double + sf: down-scale factor + Return: + downsampled LR image + Reference: + @inproceedings{zhang2019deep, + title={Deep Plug-and-Play Super-Resolution for Arbitrary Blur Kernels}, + author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei}, + booktitle={IEEE Conference on Computer Vision and Pattern Recognition}, + pages={1671--1681}, + year={2019} + } + ''' + x = bicubic_degradation(x, sf=sf) + x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap') + return x + + +def classical_degradation(x, k, sf=3): + ''' blur + downsampling + Args: + x: HxWxC image, [0, 1]/[0, 255] + k: hxw, double + sf: down-scale factor + Return: + downsampled LR image + ''' + x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap') + # x = filters.correlate(x, np.expand_dims(np.flip(k), axis=2)) + st = 0 + return x[st::sf, st::sf, ...] + + +def add_sharpening(img, weight=0.5, radius=50, threshold=10): + """USM sharpening. borrowed from real-ESRGAN + Input image: I; Blurry image: B. + 1. K = I + weight * (I - B) + 2. Mask = 1 if abs(I - B) > threshold, else: 0 + 3. Blur mask: + 4. Out = Mask * K + (1 - Mask) * I + Args: + img (Numpy array): Input image, HWC, BGR; float32, [0, 1]. + weight (float): Sharp weight. Default: 1. + radius (float): Kernel size of Gaussian blur. Default: 50. + threshold (int): + """ + if radius % 2 == 0: + radius += 1 + blur = cv2.GaussianBlur(img, (radius, radius), 0) + residual = img - blur + mask = np.abs(residual) * 255 > threshold + mask = mask.astype('float32') + soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0) + + K = img + weight * residual + K = np.clip(K, 0, 1) + return soft_mask * K + (1 - soft_mask) * img + + +def add_blur(img, sf=4): + wd2 = 4.0 + sf + wd = 2.0 + 0.2 * sf + + wd2 = wd2/4 + wd = wd/4 + + if random.random() < 0.5: + l1 = wd2 * random.random() + l2 = wd2 * random.random() + k = anisotropic_Gaussian(ksize=random.randint(2, 11) + 3, theta=random.random() * np.pi, l1=l1, l2=l2) + else: + k = fspecial('gaussian', random.randint(2, 4) + 3, wd * random.random()) + img = ndimage.filters.convolve(img, np.expand_dims(k, axis=2), mode='mirror') + + return img + + +def add_resize(img, sf=4): + rnum = np.random.rand() + if rnum > 0.8: # up + sf1 = random.uniform(1, 2) + elif rnum < 0.7: # down + sf1 = random.uniform(0.5 / sf, 1) + else: + sf1 = 1.0 + img = cv2.resize(img, (int(sf1 * img.shape[1]), int(sf1 * img.shape[0])), interpolation=random.choice([1, 2, 3])) + img = np.clip(img, 0.0, 1.0) + + return img + + +# def add_Gaussian_noise(img, noise_level1=2, noise_level2=25): +# noise_level = random.randint(noise_level1, noise_level2) +# rnum = np.random.rand() +# if rnum > 0.6: # add color Gaussian noise +# img += np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32) +# elif rnum < 0.4: # add grayscale Gaussian noise +# img += np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32) +# else: # add noise +# L = noise_level2 / 255. +# D = np.diag(np.random.rand(3)) +# U = orth(np.random.rand(3, 3)) +# conv = np.dot(np.dot(np.transpose(U), D), U) +# img += np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32) +# img = np.clip(img, 0.0, 1.0) +# return img + +def add_Gaussian_noise(img, noise_level1=2, noise_level2=25): + noise_level = random.randint(noise_level1, noise_level2) + rnum = np.random.rand() + if rnum > 0.6: # add color Gaussian noise + img = img + np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32) + elif rnum < 0.4: # add grayscale Gaussian noise + img = img + np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32) + else: # add noise + L = noise_level2 / 255. + D = np.diag(np.random.rand(3)) + U = orth(np.random.rand(3, 3)) + conv = np.dot(np.dot(np.transpose(U), D), U) + img = img + np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32) + img = np.clip(img, 0.0, 1.0) + return img + + +def add_speckle_noise(img, noise_level1=2, noise_level2=25): + noise_level = random.randint(noise_level1, noise_level2) + img = np.clip(img, 0.0, 1.0) + rnum = random.random() + if rnum > 0.6: + img += img * np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32) + elif rnum < 0.4: + img += img * np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32) + else: + L = noise_level2 / 255. + D = np.diag(np.random.rand(3)) + U = orth(np.random.rand(3, 3)) + conv = np.dot(np.dot(np.transpose(U), D), U) + img += img * np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32) + img = np.clip(img, 0.0, 1.0) + return img + + +def add_Poisson_noise(img): + img = np.clip((img * 255.0).round(), 0, 255) / 255. + vals = 10 ** (2 * random.random() + 2.0) # [2, 4] + if random.random() < 0.5: + img = np.random.poisson(img * vals).astype(np.float32) / vals + else: + img_gray = np.dot(img[..., :3], [0.299, 0.587, 0.114]) + img_gray = np.clip((img_gray * 255.0).round(), 0, 255) / 255. + noise_gray = np.random.poisson(img_gray * vals).astype(np.float32) / vals - img_gray + img += noise_gray[:, :, np.newaxis] + img = np.clip(img, 0.0, 1.0) + return img + + +def add_JPEG_noise(img): + quality_factor = random.randint(80, 95) + img = cv2.cvtColor(util.single2uint(img), cv2.COLOR_RGB2BGR) + result, encimg = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), quality_factor]) + img = cv2.imdecode(encimg, 1) + img = cv2.cvtColor(util.uint2single(img), cv2.COLOR_BGR2RGB) + return img + + +def random_crop(lq, hq, sf=4, lq_patchsize=64): + h, w = lq.shape[:2] + rnd_h = random.randint(0, h - lq_patchsize) + rnd_w = random.randint(0, w - lq_patchsize) + lq = lq[rnd_h:rnd_h + lq_patchsize, rnd_w:rnd_w + lq_patchsize, :] + + rnd_h_H, rnd_w_H = int(rnd_h * sf), int(rnd_w * sf) + hq = hq[rnd_h_H:rnd_h_H + lq_patchsize * sf, rnd_w_H:rnd_w_H + lq_patchsize * sf, :] + return lq, hq + + +def degradation_bsrgan(img, sf=4, lq_patchsize=72, isp_model=None): + """ + This is the degradation model of BSRGAN from the paper + "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution" + ---------- + img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf) + sf: scale factor + isp_model: camera ISP model + Returns + ------- + img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1] + hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1] + """ + isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25 + sf_ori = sf + + h1, w1 = img.shape[:2] + img = img.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop + h, w = img.shape[:2] + + if h < lq_patchsize * sf or w < lq_patchsize * sf: + raise ValueError(f'img size ({h1}X{w1}) is too small!') + + hq = img.copy() + + if sf == 4 and random.random() < scale2_prob: # downsample1 + if np.random.rand() < 0.5: + img = cv2.resize(img, (int(1 / 2 * img.shape[1]), int(1 / 2 * img.shape[0])), + interpolation=random.choice([1, 2, 3])) + else: + img = util.imresize_np(img, 1 / 2, True) + img = np.clip(img, 0.0, 1.0) + sf = 2 + + shuffle_order = random.sample(range(7), 7) + idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3) + if idx1 > idx2: # keep downsample3 last + shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1] + + for i in shuffle_order: + + if i == 0: + img = add_blur(img, sf=sf) + + elif i == 1: + img = add_blur(img, sf=sf) + + elif i == 2: + a, b = img.shape[1], img.shape[0] + # downsample2 + if random.random() < 0.75: + sf1 = random.uniform(1, 2 * sf) + img = cv2.resize(img, (int(1 / sf1 * img.shape[1]), int(1 / sf1 * img.shape[0])), + interpolation=random.choice([1, 2, 3])) + else: + k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf)) + k_shifted = shift_pixel(k, sf) + k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel + img = ndimage.filters.convolve(img, np.expand_dims(k_shifted, axis=2), mode='mirror') + img = img[0::sf, 0::sf, ...] # nearest downsampling + img = np.clip(img, 0.0, 1.0) + + elif i == 3: + # downsample3 + img = cv2.resize(img, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3])) + img = np.clip(img, 0.0, 1.0) + + elif i == 4: + # add Gaussian noise + img = add_Gaussian_noise(img, noise_level1=2, noise_level2=8) + + elif i == 5: + # add JPEG noise + if random.random() < jpeg_prob: + img = add_JPEG_noise(img) + + elif i == 6: + # add processed camera sensor noise + if random.random() < isp_prob and isp_model is not None: + with torch.no_grad(): + img, hq = isp_model.forward(img.copy(), hq) + + # add final JPEG compression noise + img = add_JPEG_noise(img) + + # random crop + img, hq = random_crop(img, hq, sf_ori, lq_patchsize) + + return img, hq + + +# todo no isp_model? +def degradation_bsrgan_variant(image, sf=4, isp_model=None): + """ + This is the degradation model of BSRGAN from the paper + "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution" + ---------- + sf: scale factor + isp_model: camera ISP model + Returns + ------- + img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1] + hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1] + """ + image = util.uint2single(image) + isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25 + sf_ori = sf + + h1, w1 = image.shape[:2] + image = image.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop + h, w = image.shape[:2] + + hq = image.copy() + + if sf == 4 and random.random() < scale2_prob: # downsample1 + if np.random.rand() < 0.5: + image = cv2.resize(image, (int(1 / 2 * image.shape[1]), int(1 / 2 * image.shape[0])), + interpolation=random.choice([1, 2, 3])) + else: + image = util.imresize_np(image, 1 / 2, True) + image = np.clip(image, 0.0, 1.0) + sf = 2 + + shuffle_order = random.sample(range(7), 7) + idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3) + if idx1 > idx2: # keep downsample3 last + shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1] + + for i in shuffle_order: + + if i == 0: + image = add_blur(image, sf=sf) + + # elif i == 1: + # image = add_blur(image, sf=sf) + + if i == 0: + pass + + elif i == 2: + a, b = image.shape[1], image.shape[0] + # downsample2 + if random.random() < 0.8: + sf1 = random.uniform(1, 2 * sf) + image = cv2.resize(image, (int(1 / sf1 * image.shape[1]), int(1 / sf1 * image.shape[0])), + interpolation=random.choice([1, 2, 3])) + else: + k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf)) + k_shifted = shift_pixel(k, sf) + k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel + image = ndimage.filters.convolve(image, np.expand_dims(k_shifted, axis=2), mode='mirror') + image = image[0::sf, 0::sf, ...] # nearest downsampling + + image = np.clip(image, 0.0, 1.0) + + elif i == 3: + # downsample3 + image = cv2.resize(image, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3])) + image = np.clip(image, 0.0, 1.0) + + elif i == 4: + # add Gaussian noise + image = add_Gaussian_noise(image, noise_level1=1, noise_level2=2) + + elif i == 5: + # add JPEG noise + if random.random() < jpeg_prob: + image = add_JPEG_noise(image) + # + # elif i == 6: + # # add processed camera sensor noise + # if random.random() < isp_prob and isp_model is not None: + # with torch.no_grad(): + # img, hq = isp_model.forward(img.copy(), hq) + + # add final JPEG compression noise + image = add_JPEG_noise(image) + image = util.single2uint(image) + example = {"image": image} + return example + + + + +if __name__ == '__main__': + print("hey") + img = util.imread_uint('utils/test.png', 3) + img = img[:448, :448] + h = img.shape[0] // 4 + print("resizing to", h) + sf = 4 + deg_fn = partial(degradation_bsrgan_variant, sf=sf) + for i in range(20): + print(i) + img_hq = img + img_lq = deg_fn(img)["image"] + img_hq, img_lq = util.uint2single(img_hq), util.uint2single(img_lq) + print(img_lq) + img_lq_bicubic = albumentations.SmallestMaxSize(max_size=h, interpolation=cv2.INTER_CUBIC)(image=img_hq)["image"] + print(img_lq.shape) + print("bicubic", img_lq_bicubic.shape) + print(img_hq.shape) + lq_nearest = cv2.resize(util.single2uint(img_lq), (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])), + interpolation=0) + lq_bicubic_nearest = cv2.resize(util.single2uint(img_lq_bicubic), + (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])), + interpolation=0) + img_concat = np.concatenate([lq_bicubic_nearest, lq_nearest, util.single2uint(img_hq)], axis=1) + util.imsave(img_concat, str(i) + '.png') diff --git a/ruyi/vae/ldm/modules/image_degradation/utils/test.png b/ruyi/vae/ldm/modules/image_degradation/utils/test.png new file mode 100644 index 0000000..4249b43 Binary files /dev/null and b/ruyi/vae/ldm/modules/image_degradation/utils/test.png differ diff --git a/ruyi/vae/ldm/modules/image_degradation/utils_image.py b/ruyi/vae/ldm/modules/image_degradation/utils_image.py new file mode 100644 index 0000000..891e310 --- /dev/null +++ b/ruyi/vae/ldm/modules/image_degradation/utils_image.py @@ -0,0 +1,920 @@ +import math +import os +import random +from datetime import datetime + +import cv2 +import numpy as np +import torch +from torchvision.utils import make_grid + +import matplotlib.pyplot as plt + +#import matplotlib.pyplot as plt # TODO: check with Dominik, also bsrgan.py vs bsrgan_light.py + + +os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE" + + +''' +# -------------------------------------------- +# Kai Zhang (github: https://github.com/cszn) +# 03/Mar/2019 +# -------------------------------------------- +# https://github.com/twhui/SRGAN-pyTorch +# https://github.com/xinntao/BasicSR +# -------------------------------------------- +''' + + +IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tif'] + + +def is_image_file(filename): + return any(filename.endswith(extension) for extension in IMG_EXTENSIONS) + + +def get_timestamp(): + return datetime.now().strftime('%y%m%d-%H%M%S') + + +def imshow(x, title=None, cbar=False, figsize=None): + plt.figure(figsize=figsize) + plt.imshow(np.squeeze(x), interpolation='nearest', cmap='gray') + if title: + plt.title(title) + if cbar: + plt.colorbar() + plt.show() + + +def surf(Z, cmap='rainbow', figsize=None): + plt.figure(figsize=figsize) + ax3 = plt.axes(projection='3d') + + w, h = Z.shape[:2] + xx = np.arange(0,w,1) + yy = np.arange(0,h,1) + X, Y = np.meshgrid(xx, yy) + ax3.plot_surface(X,Y,Z,cmap=cmap) + #ax3.contour(X,Y,Z, zdim='z',offset=-2，cmap=cmap) + plt.show() + + +''' +# -------------------------------------------- +# get image pathes +# -------------------------------------------- +''' + + +def get_image_paths(dataroot): + paths = None # return None if dataroot is None + if dataroot is not None: + paths = sorted(_get_paths_from_images(dataroot)) + return paths + + +def _get_paths_from_images(path): + assert os.path.isdir(path), '{:s} is not a valid directory'.format(path) + images = [] + for dirpath, _, fnames in sorted(os.walk(path)): + for fname in sorted(fnames): + if is_image_file(fname): + img_path = os.path.join(dirpath, fname) + images.append(img_path) + assert images, '{:s} has no valid image file'.format(path) + return images + + +''' +# -------------------------------------------- +# split large images into small images +# -------------------------------------------- +''' + + +def patches_from_image(img, p_size=512, p_overlap=64, p_max=800): + w, h = img.shape[:2] + patches = [] + if w > p_max and h > p_max: + w1 = list(np.arange(0, w-p_size, p_size-p_overlap, dtype=np.int)) + h1 = list(np.arange(0, h-p_size, p_size-p_overlap, dtype=np.int)) + w1.append(w-p_size) + h1.append(h-p_size) +# print(w1) +# print(h1) + for i in w1: + for j in h1: + patches.append(img[i:i+p_size, j:j+p_size,:]) + else: + patches.append(img) + + return patches + + +def imssave(imgs, img_path): + """ + imgs: list, N images of size WxHxC + """ + img_name, ext = os.path.splitext(os.path.basename(img_path)) + + for i, img in enumerate(imgs): + if img.ndim == 3: + img = img[:, :, [2, 1, 0]] + new_path = os.path.join(os.path.dirname(img_path), img_name+str('_s{:04d}'.format(i))+'.png') + cv2.imwrite(new_path, img) + + +def split_imageset(original_dataroot, taget_dataroot, n_channels=3, p_size=800, p_overlap=96, p_max=1000): + """ + split the large images from original_dataroot into small overlapped images with size (p_size)x(p_size), + and save them into taget_dataroot; only the images with larger size than (p_max)x(p_max) + will be splitted. + Args: + original_dataroot: + taget_dataroot: + p_size: size of small images + p_overlap: patch size in training is a good choice + p_max: images with smaller size than (p_max)x(p_max) keep unchanged. + """ + paths = get_image_paths(original_dataroot) + for img_path in paths: + # img_name, ext = os.path.splitext(os.path.basename(img_path)) + img = imread_uint(img_path, n_channels=n_channels) + patches = patches_from_image(img, p_size, p_overlap, p_max) + imssave(patches, os.path.join(taget_dataroot,os.path.basename(img_path))) + #if original_dataroot == taget_dataroot: + #del img_path + +''' +# -------------------------------------------- +# makedir +# -------------------------------------------- +''' + + +def mkdir(path): + if not os.path.exists(path): + os.makedirs(path) + + +def mkdirs(paths): + if isinstance(paths, str): + mkdir(paths) + else: + for path in paths: + mkdir(path) + + +def mkdir_and_rename(path): + if os.path.exists(path): + new_name = path + '_archived_' + get_timestamp() + print('Path already exists. Rename it to [{:s}]'.format(new_name)) + os.rename(path, new_name) + os.makedirs(path) + + +''' +# -------------------------------------------- +# read image from path +# opencv is fast, but read BGR numpy image +# -------------------------------------------- +''' + + +# -------------------------------------------- +# get uint8 image of size HxWxn_channles (RGB) +# -------------------------------------------- +def imread_uint(path, n_channels=3): + # input: path + # output: HxWx3(RGB or GGG), or HxWx1 (G) + if n_channels == 1: + img = cv2.imread(path, 0) # cv2.IMREAD_GRAYSCALE + img = np.expand_dims(img, axis=2) # HxWx1 + elif n_channels == 3: + img = cv2.imread(path, cv2.IMREAD_UNCHANGED) # BGR or G + if img.ndim == 2: + img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) # GGG + else: + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # RGB + return img + + +# -------------------------------------------- +# matlab's imwrite +# -------------------------------------------- +def imsave(img, img_path): + img = np.squeeze(img) + if img.ndim == 3: + img = img[:, :, [2, 1, 0]] + cv2.imwrite(img_path, img) + +def imwrite(img, img_path): + img = np.squeeze(img) + if img.ndim == 3: + img = img[:, :, [2, 1, 0]] + cv2.imwrite(img_path, img) + + + +# -------------------------------------------- +# get single image of size HxWxn_channles (BGR) +# -------------------------------------------- +def read_img(path): + # read image by cv2 + # return: Numpy float32, HWC, BGR, [0,1] + img = cv2.imread(path, cv2.IMREAD_UNCHANGED) # cv2.IMREAD_GRAYSCALE + img = img.astype(np.float32) / 255. + if img.ndim == 2: + img = np.expand_dims(img, axis=2) + # some images have 4 channels + if img.shape[2] > 3: + img = img[:, :, :3] + return img + + +''' +# -------------------------------------------- +# image format conversion +# -------------------------------------------- +# numpy(single) <---> numpy(unit) +# numpy(single) <---> tensor +# numpy(unit) <---> tensor +# -------------------------------------------- +''' + + +# -------------------------------------------- +# numpy(single) [0, 1] <---> numpy(unit) +# -------------------------------------------- + + +def uint2single(img): + + return np.float32(img/255.) + + +def single2uint(img): + + return np.uint8((img.clip(0, 1)*255.).round()) + + +def uint162single(img): + + return np.float32(img/65535.) + + +def single2uint16(img): + + return np.uint16((img.clip(0, 1)*65535.).round()) + + +# -------------------------------------------- +# numpy(unit) (HxWxC or HxW) <---> tensor +# -------------------------------------------- + + +# convert uint to 4-dimensional torch tensor +def uint2tensor4(img): + if img.ndim == 2: + img = np.expand_dims(img, axis=2) + return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().div(255.).unsqueeze(0) + + +# convert uint to 3-dimensional torch tensor +def uint2tensor3(img): + if img.ndim == 2: + img = np.expand_dims(img, axis=2) + return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().div(255.) + + +# convert 2/3/4-dimensional torch tensor to uint +def tensor2uint(img): + img = img.data.squeeze().float().clamp_(0, 1).cpu().numpy() + if img.ndim == 3: + img = np.transpose(img, (1, 2, 0)) + return np.uint8((img*255.0).round()) + + +# -------------------------------------------- +# numpy(single) (HxWxC) <---> tensor +# -------------------------------------------- + + +# convert single (HxWxC) to 3-dimensional torch tensor +def single2tensor3(img): + return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float() + + +# convert single (HxWxC) to 4-dimensional torch tensor +def single2tensor4(img): + return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().unsqueeze(0) + + +# convert torch tensor to single +def tensor2single(img): + img = img.data.squeeze().float().cpu().numpy() + if img.ndim == 3: + img = np.transpose(img, (1, 2, 0)) + + return img + +# convert torch tensor to single +def tensor2single3(img): + img = img.data.squeeze().float().cpu().numpy() + if img.ndim == 3: + img = np.transpose(img, (1, 2, 0)) + elif img.ndim == 2: + img = np.expand_dims(img, axis=2) + return img + + +def single2tensor5(img): + return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1, 3).float().unsqueeze(0) + + +def single32tensor5(img): + return torch.from_numpy(np.ascontiguousarray(img)).float().unsqueeze(0).unsqueeze(0) + + +def single42tensor4(img): + return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1, 3).float() + + +# from skimage.io import imread, imsave +def tensor2img(tensor, out_type=np.uint8, min_max=(0, 1)): + ''' + Converts a torch Tensor into an image Numpy array of BGR channel order + Input: 4D(B,(3/1),H,W), 3D(C,H,W), or 2D(H,W), any range, RGB channel order + Output: 3D(H,W,C) or 2D(H,W), [0,255], np.uint8 (default) + ''' + tensor = tensor.squeeze().float().cpu().clamp_(*min_max) # squeeze first, then clamp + tensor = (tensor - min_max[0]) / (min_max[1] - min_max[0]) # to range [0,1] + n_dim = tensor.dim() + if n_dim == 4: + n_img = len(tensor) + img_np = make_grid(tensor, nrow=int(math.sqrt(n_img)), normalize=False).numpy() + img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0)) # HWC, BGR + elif n_dim == 3: + img_np = tensor.numpy() + img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0)) # HWC, BGR + elif n_dim == 2: + img_np = tensor.numpy() + else: + raise TypeError( + 'Only support 4D, 3D and 2D tensor. But received with dimension: {:d}'.format(n_dim)) + if out_type == np.uint8: + img_np = (img_np * 255.0).round() + # Important. Unlike matlab, numpy.unit8() WILL NOT round by default. + return img_np.astype(out_type) + + +''' +# -------------------------------------------- +# Augmentation, flipe and/or rotate +# -------------------------------------------- +# The following two are enough. +# (1) augmet_img: numpy image of WxHxC or WxH +# (2) augment_img_tensor4: tensor image 1xCxWxH +# -------------------------------------------- +''' + + +def augment_img(img, mode=0): + '''Kai Zhang (github: https://github.com/cszn) + ''' + if mode == 0: + return img + elif mode == 1: + return np.flipud(np.rot90(img)) + elif mode == 2: + return np.flipud(img) + elif mode == 3: + return np.rot90(img, k=3) + elif mode == 4: + return np.flipud(np.rot90(img, k=2)) + elif mode == 5: + return np.rot90(img) + elif mode == 6: + return np.rot90(img, k=2) + elif mode == 7: + return np.flipud(np.rot90(img, k=3)) + + +def augment_img_tensor4(img, mode=0): + '''Kai Zhang (github: https://github.com/cszn) + ''' + if mode == 0: + return img + elif mode == 1: + return img.rot90(1, [2, 3]).flip([2]) + elif mode == 2: + return img.flip([2]) + elif mode == 3: + return img.rot90(3, [2, 3]) + elif mode == 4: + return img.rot90(2, [2, 3]).flip([2]) + elif mode == 5: + return img.rot90(1, [2, 3]) + elif mode == 6: + return img.rot90(2, [2, 3]) + elif mode == 7: + return img.rot90(3, [2, 3]).flip([2]) + + +def augment_img_tensor(img, mode=0): + '''Kai Zhang (github: https://github.com/cszn) + ''' + img_size = img.size() + img_np = img.data.cpu().numpy() + if len(img_size) == 3: + img_np = np.transpose(img_np, (1, 2, 0)) + elif len(img_size) == 4: + img_np = np.transpose(img_np, (2, 3, 1, 0)) + img_np = augment_img(img_np, mode=mode) + img_tensor = torch.from_numpy(np.ascontiguousarray(img_np)) + if len(img_size) == 3: + img_tensor = img_tensor.permute(2, 0, 1) + elif len(img_size) == 4: + img_tensor = img_tensor.permute(3, 2, 0, 1) + + return img_tensor.type_as(img) + + +def augment_img_np3(img, mode=0): + if mode == 0: + return img + elif mode == 1: + return img.transpose(1, 0, 2) + elif mode == 2: + return img[::-1, :, :] + elif mode == 3: + img = img[::-1, :, :] + img = img.transpose(1, 0, 2) + return img + elif mode == 4: + return img[:, ::-1, :] + elif mode == 5: + img = img[:, ::-1, :] + img = img.transpose(1, 0, 2) + return img + elif mode == 6: + img = img[:, ::-1, :] + img = img[::-1, :, :] + return img + elif mode == 7: + img = img[:, ::-1, :] + img = img[::-1, :, :] + img = img.transpose(1, 0, 2) + return img + + +def augment_imgs(img_list, hflip=True, rot=True): + # horizontal flip OR rotate + hflip = hflip and random.random() < 0.5 + vflip = rot and random.random() < 0.5 + rot90 = rot and random.random() < 0.5 + + def _augment(img): + if hflip: + img = img[:, ::-1, :] + if vflip: + img = img[::-1, :, :] + if rot90: + img = img.transpose(1, 0, 2) + return img + + return [_augment(img) for img in img_list] + + +''' +# -------------------------------------------- +# modcrop and shave +# -------------------------------------------- +''' + + +def modcrop(img_in, scale): + # img_in: Numpy, HWC or HW + img = np.copy(img_in) + if img.ndim == 2: + H, W = img.shape + H_r, W_r = H % scale, W % scale + img = img[:H - H_r, :W - W_r] + elif img.ndim == 3: + H, W, C = img.shape + H_r, W_r = H % scale, W % scale + img = img[:H - H_r, :W - W_r, :] + else: + raise ValueError('Wrong img ndim: [{:d}].'.format(img.ndim)) + return img + + +def shave(img_in, border=0): + # img_in: Numpy, HWC or HW + img = np.copy(img_in) + h, w = img.shape[:2] + img = img[border:h-border, border:w-border] + return img + + +''' +# -------------------------------------------- +# image processing process on numpy image +# channel_convert(in_c, tar_type, img_list): +# rgb2ycbcr(img, only_y=True): +# bgr2ycbcr(img, only_y=True): +# ycbcr2rgb(img): +# -------------------------------------------- +''' + + +def rgb2ycbcr(img, only_y=True): + '''same as matlab rgb2ycbcr + only_y: only return Y channel + Input: + uint8, [0, 255] + float, [0, 1] + ''' + in_img_type = img.dtype + img.astype(np.float32) + if in_img_type != np.uint8: + img *= 255. + # convert + if only_y: + rlt = np.dot(img, [65.481, 128.553, 24.966]) / 255.0 + 16.0 + else: + rlt = np.matmul(img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786], + [24.966, 112.0, -18.214]]) / 255.0 + [16, 128, 128] + if in_img_type == np.uint8: + rlt = rlt.round() + else: + rlt /= 255. + return rlt.astype(in_img_type) + + +def ycbcr2rgb(img): + '''same as matlab ycbcr2rgb + Input: + uint8, [0, 255] + float, [0, 1] + ''' + in_img_type = img.dtype + img.astype(np.float32) + if in_img_type != np.uint8: + img *= 255. + # convert + rlt = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0, -0.00153632, 0.00791071], + [0.00625893, -0.00318811, 0]]) * 255.0 + [-222.921, 135.576, -276.836] + if in_img_type == np.uint8: + rlt = rlt.round() + else: + rlt /= 255. + return rlt.astype(in_img_type) + + +def bgr2ycbcr(img, only_y=True): + '''bgr version of rgb2ycbcr + only_y: only return Y channel + Input: + uint8, [0, 255] + float, [0, 1] + ''' + in_img_type = img.dtype + img.astype(np.float32) + if in_img_type != np.uint8: + img *= 255. + # convert + if only_y: + rlt = np.dot(img, [24.966, 128.553, 65.481]) / 255.0 + 16.0 + else: + rlt = np.matmul(img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786], + [65.481, -37.797, 112.0]]) / 255.0 + [16, 128, 128] + if in_img_type == np.uint8: + rlt = rlt.round() + else: + rlt /= 255. + return rlt.astype(in_img_type) + + +def channel_convert(in_c, tar_type, img_list): + # conversion among BGR, gray and y + if in_c == 3 and tar_type == 'gray': # BGR to gray + gray_list = [cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) for img in img_list] + return [np.expand_dims(img, axis=2) for img in gray_list] + elif in_c == 3 and tar_type == 'y': # BGR to y + y_list = [bgr2ycbcr(img, only_y=True) for img in img_list] + return [np.expand_dims(img, axis=2) for img in y_list] + elif in_c == 1 and tar_type == 'RGB': # gray/y to BGR + return [cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) for img in img_list] + else: + return img_list + + +''' +# -------------------------------------------- +# metric, PSNR and SSIM +# -------------------------------------------- +''' + + +# -------------------------------------------- +# PSNR +# -------------------------------------------- +def calculate_psnr(img1, img2, border=0): + # img1 and img2 have range [0, 255] + #img1 = img1.squeeze() + #img2 = img2.squeeze() + if not img1.shape == img2.shape: + raise ValueError('Input images must have the same dimensions.') + h, w = img1.shape[:2] + img1 = img1[border:h-border, border:w-border] + img2 = img2[border:h-border, border:w-border] + + img1 = img1.astype(np.float64) + img2 = img2.astype(np.float64) + mse = np.mean((img1 - img2)**2) + if mse == 0: + return float('inf') + return 20 * math.log10(255.0 / math.sqrt(mse)) + + +# -------------------------------------------- +# SSIM +# -------------------------------------------- +def calculate_ssim(img1, img2, border=0): + '''calculate SSIM + the same outputs as MATLAB's + img1, img2: [0, 255] + ''' + #img1 = img1.squeeze() + #img2 = img2.squeeze() + if not img1.shape == img2.shape: + raise ValueError('Input images must have the same dimensions.') + h, w = img1.shape[:2] + img1 = img1[border:h-border, border:w-border] + img2 = img2[border:h-border, border:w-border] + + if img1.ndim == 2: + return ssim(img1, img2) + elif img1.ndim == 3: + if img1.shape[2] == 3: + ssims = [] + for i in range(3): + ssims.append(ssim(img1[:,:,i], img2[:,:,i])) + return np.array(ssims).mean() + elif img1.shape[2] == 1: + return ssim(np.squeeze(img1), np.squeeze(img2)) + else: + raise ValueError('Wrong input image dimensions.') + + +def ssim(img1, img2): + C1 = (0.01 * 255)**2 + C2 = (0.03 * 255)**2 + + img1 = img1.astype(np.float64) + img2 = img2.astype(np.float64) + kernel = cv2.getGaussianKernel(11, 1.5) + window = np.outer(kernel, kernel.transpose()) + + mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5] # valid + mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5] + mu1_sq = mu1**2 + mu2_sq = mu2**2 + mu1_mu2 = mu1 * mu2 + sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq + sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq + sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2 + + ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * + (sigma1_sq + sigma2_sq + C2)) + return ssim_map.mean() + + +''' +# -------------------------------------------- +# matlab's bicubic imresize (numpy and torch) [0, 1] +# -------------------------------------------- +''' + + +# matlab 'imresize' function, now only support 'bicubic' +def cubic(x): + absx = torch.abs(x) + absx2 = absx**2 + absx3 = absx**3 + return (1.5*absx3 - 2.5*absx2 + 1) * ((absx <= 1).type_as(absx)) + \ + (-0.5*absx3 + 2.5*absx2 - 4*absx + 2) * (((absx > 1)*(absx <= 2)).type_as(absx)) + + +def calculate_weights_indices(in_length, out_length, scale, kernel, kernel_width, antialiasing): + if (scale < 1) and (antialiasing): + # Use a modified kernel to simultaneously interpolate and antialias- larger kernel width + kernel_width = kernel_width / scale + + # Output-space coordinates + x = torch.linspace(1, out_length, out_length) + + # Input-space coordinates. Calculate the inverse mapping such that 0.5 + # in output space maps to 0.5 in input space, and 0.5+scale in output + # space maps to 1.5 in input space. + u = x / scale + 0.5 * (1 - 1 / scale) + + # What is the left-most pixel that can be involved in the computation? + left = torch.floor(u - kernel_width / 2) + + # What is the maximum number of pixels that can be involved in the + # computation? Note: it's OK to use an extra pixel here; if the + # corresponding weights are all zero, it will be eliminated at the end + # of this function. + P = math.ceil(kernel_width) + 2 + + # The indices of the input pixels involved in computing the k-th output + # pixel are in row k of the indices matrix. + indices = left.view(out_length, 1).expand(out_length, P) + torch.linspace(0, P - 1, P).view( + 1, P).expand(out_length, P) + + # The weights used to compute the k-th output pixel are in row k of the + # weights matrix. + distance_to_center = u.view(out_length, 1).expand(out_length, P) - indices + # apply cubic kernel + if (scale < 1) and (antialiasing): + weights = scale * cubic(distance_to_center * scale) + else: + weights = cubic(distance_to_center) + # Normalize the weights matrix so that each row sums to 1. + weights_sum = torch.sum(weights, 1).view(out_length, 1) + weights = weights / weights_sum.expand(out_length, P) + + # If a column in weights is all zero, get rid of it. only consider the first and last column. + weights_zero_tmp = torch.sum((weights == 0), 0) + if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6): + indices = indices.narrow(1, 1, P - 2) + weights = weights.narrow(1, 1, P - 2) + if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6): + indices = indices.narrow(1, 0, P - 2) + weights = weights.narrow(1, 0, P - 2) + weights = weights.contiguous() + indices = indices.contiguous() + sym_len_s = -indices.min() + 1 + sym_len_e = indices.max() - in_length + indices = indices + sym_len_s - 1 + return weights, indices, int(sym_len_s), int(sym_len_e) + + +# -------------------------------------------- +# imresize for tensor image [0, 1] +# -------------------------------------------- +def imresize(img, scale, antialiasing=True): + # Now the scale should be the same for H and W + # input: img: pytorch tensor, CHW or HW [0,1] + # output: CHW or HW [0,1] w/o round + need_squeeze = True if img.dim() == 2 else False + if need_squeeze: + img.unsqueeze_(0) + in_C, in_H, in_W = img.size() + out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale) + kernel_width = 4 + kernel = 'cubic' + + # Return the desired dimension order for performing the resize. The + # strategy is to perform the resize first along the dimension with the + # smallest scale factor. + # Now we do not support this. + + # get weights and indices + weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices( + in_H, out_H, scale, kernel, kernel_width, antialiasing) + weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices( + in_W, out_W, scale, kernel, kernel_width, antialiasing) + # process H dimension + # symmetric copying + img_aug = torch.FloatTensor(in_C, in_H + sym_len_Hs + sym_len_He, in_W) + img_aug.narrow(1, sym_len_Hs, in_H).copy_(img) + + sym_patch = img[:, :sym_len_Hs, :] + inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long() + sym_patch_inv = sym_patch.index_select(1, inv_idx) + img_aug.narrow(1, 0, sym_len_Hs).copy_(sym_patch_inv) + + sym_patch = img[:, -sym_len_He:, :] + inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long() + sym_patch_inv = sym_patch.index_select(1, inv_idx) + img_aug.narrow(1, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv) + + out_1 = torch.FloatTensor(in_C, out_H, in_W) + kernel_width = weights_H.size(1) + for i in range(out_H): + idx = int(indices_H[i][0]) + for j in range(out_C): + out_1[j, i, :] = img_aug[j, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i]) + + # process W dimension + # symmetric copying + out_1_aug = torch.FloatTensor(in_C, out_H, in_W + sym_len_Ws + sym_len_We) + out_1_aug.narrow(2, sym_len_Ws, in_W).copy_(out_1) + + sym_patch = out_1[:, :, :sym_len_Ws] + inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long() + sym_patch_inv = sym_patch.index_select(2, inv_idx) + out_1_aug.narrow(2, 0, sym_len_Ws).copy_(sym_patch_inv) + + sym_patch = out_1[:, :, -sym_len_We:] + inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long() + sym_patch_inv = sym_patch.index_select(2, inv_idx) + out_1_aug.narrow(2, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv) + + out_2 = torch.FloatTensor(in_C, out_H, out_W) + kernel_width = weights_W.size(1) + for i in range(out_W): + idx = int(indices_W[i][0]) + for j in range(out_C): + out_2[j, :, i] = out_1_aug[j, :, idx:idx + kernel_width].mv(weights_W[i]) + if need_squeeze: + out_2.squeeze_() + return out_2 + + +# -------------------------------------------- +# imresize for numpy image [0, 1] +# -------------------------------------------- +def imresize_np(img, scale, antialiasing=True): + # Now the scale should be the same for H and W + # input: img: Numpy, HWC or HW [0,1] + # output: HWC or HW [0,1] w/o round + img = torch.from_numpy(img) + need_squeeze = True if img.dim() == 2 else False + if need_squeeze: + img.unsqueeze_(2) + + in_H, in_W, in_C = img.size() + out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale) + kernel_width = 4 + kernel = 'cubic' + + # Return the desired dimension order for performing the resize. The + # strategy is to perform the resize first along the dimension with the + # smallest scale factor. + # Now we do not support this. + + # get weights and indices + weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices( + in_H, out_H, scale, kernel, kernel_width, antialiasing) + weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices( + in_W, out_W, scale, kernel, kernel_width, antialiasing) + # process H dimension + # symmetric copying + img_aug = torch.FloatTensor(in_H + sym_len_Hs + sym_len_He, in_W, in_C) + img_aug.narrow(0, sym_len_Hs, in_H).copy_(img) + + sym_patch = img[:sym_len_Hs, :, :] + inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long() + sym_patch_inv = sym_patch.index_select(0, inv_idx) + img_aug.narrow(0, 0, sym_len_Hs).copy_(sym_patch_inv) + + sym_patch = img[-sym_len_He:, :, :] + inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long() + sym_patch_inv = sym_patch.index_select(0, inv_idx) + img_aug.narrow(0, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv) + + out_1 = torch.FloatTensor(out_H, in_W, in_C) + kernel_width = weights_H.size(1) + for i in range(out_H): + idx = int(indices_H[i][0]) + for j in range(out_C): + out_1[i, :, j] = img_aug[idx:idx + kernel_width, :, j].transpose(0, 1).mv(weights_H[i]) + + # process W dimension + # symmetric copying + out_1_aug = torch.FloatTensor(out_H, in_W + sym_len_Ws + sym_len_We, in_C) + out_1_aug.narrow(1, sym_len_Ws, in_W).copy_(out_1) + + sym_patch = out_1[:, :sym_len_Ws, :] + inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long() + sym_patch_inv = sym_patch.index_select(1, inv_idx) + out_1_aug.narrow(1, 0, sym_len_Ws).copy_(sym_patch_inv) + + sym_patch = out_1[:, -sym_len_We:, :] + inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long() + sym_patch_inv = sym_patch.index_select(1, inv_idx) + out_1_aug.narrow(1, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv) + + out_2 = torch.FloatTensor(out_H, out_W, in_C) + kernel_width = weights_W.size(1) + for i in range(out_W): + idx = int(indices_W[i][0]) + for j in range(out_C): + out_2[:, i, j] = out_1_aug[:, idx:idx + kernel_width, j].mv(weights_W[i]) + if need_squeeze: + out_2.squeeze_() + + return out_2.numpy() + + +if __name__ == '__main__': + print('---') +# img = imread_uint('test.bmp', 3) +# img = uint2single(img) +# img_bicubic = imresize_np(img, 1/4) \ No newline at end of file diff --git a/ruyi/vae/ldm/modules/losses/__init__.py b/ruyi/vae/ldm/modules/losses/__init__.py new file mode 100644 index 0000000..9a8b4f6 --- /dev/null +++ b/ruyi/vae/ldm/modules/losses/__init__.py @@ -0,0 +1 @@ +from .contperceptual import LPIPSWithDiscriminator \ No newline at end of file diff --git a/ruyi/vae/ldm/modules/losses/contperceptual.py b/ruyi/vae/ldm/modules/losses/contperceptual.py new file mode 100644 index 0000000..f2dfcda --- /dev/null +++ b/ruyi/vae/ldm/modules/losses/contperceptual.py @@ -0,0 +1,148 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from taming.modules.losses.vqperceptual import * # TODO: taming dependency yes/no? +from ..vaemodules.discriminator import Discriminator3D + +class LPIPSWithDiscriminator(nn.Module): + def __init__(self, disc_start, logvar_init=0.0, kl_weight=1.0, pixelloss_weight=1.0, + disc_num_layers=3, disc_in_channels=3, disc_factor=1.0, disc_weight=1.0, + perceptual_weight=1.0, use_actnorm=False, disc_conditional=False, + disc_loss="hinge", l2_loss_weight=0.0, l1_loss_weight=1.0): + + super().__init__() + assert disc_loss in ["hinge", "vanilla"] + self.kl_weight = kl_weight + self.pixel_weight = pixelloss_weight + self.perceptual_loss = LPIPS().eval() + self.perceptual_weight = perceptual_weight + # output log variance + self.logvar = nn.Parameter(torch.ones(size=()) * logvar_init) + + self.discriminator = NLayerDiscriminator(input_nc=disc_in_channels, + n_layers=disc_num_layers, + use_actnorm=use_actnorm + ).apply(weights_init) + self.discriminator3d = Discriminator3D( + in_channels=disc_in_channels, + block_out_channels=(64, 128, 256) + ).apply(weights_init) + self.discriminator_iter_start = disc_start + self.disc_loss = hinge_d_loss if disc_loss == "hinge" else vanilla_d_loss + self.disc_factor = disc_factor + self.discriminator_weight = disc_weight + self.disc_conditional = disc_conditional + self.l1_loss_weight = l1_loss_weight + self.l2_loss_weight = l2_loss_weight + + def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None): + if last_layer is not None: + nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0] + g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0] + else: + nll_grads = torch.autograd.grad(nll_loss, self.last_layer[0], retain_graph=True)[0] + g_grads = torch.autograd.grad(g_loss, self.last_layer[0], retain_graph=True)[0] + + d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4) + d_weight = torch.clamp(d_weight, 0.0, 1e4).detach() + d_weight = d_weight * self.discriminator_weight + return d_weight + + def forward(self, inputs, reconstructions, posteriors, optimizer_idx, + global_step, last_layer=None, cond=None, split="train", + weights=None): + + if inputs.ndim==4: + inputs = inputs.unsqueeze(2) + if reconstructions.ndim==4: + reconstructions = reconstructions.unsqueeze(2) + + inputs_ori = inputs + reconstructions_ori = reconstructions + + # get new loss_weight + loss_weights = 1 + # b, _ ,f, _, _ = reconstructions.size() + # loss_weights = torch.ones([b, f]).view(b, 1, f, 1, 1) + # loss_weights[:, :, 0] = 3 + # for i in range(1, f, 8): + # loss_weights[:, :, i - 1] = 3 + # loss_weights[:, :, i] = 3 + # loss_weights[:, :, -1] = 3 + # loss_weights = loss_weights.permute(0, 2, 1, 3, 4).flatten(0, 1).to(reconstructions.device) + + inputs = inputs.permute(0, 2, 1, 3, 4).flatten(0, 1) + reconstructions = reconstructions.permute(0, 2, 1, 3, 4).flatten(0, 1) + + rec_loss = 0 + if self.l1_loss_weight > 0: + rec_loss += torch.abs(inputs.contiguous() - reconstructions.contiguous()) * self.l1_loss_weight + if self.l2_loss_weight > 0: + rec_loss += F.mse_loss(inputs.contiguous(), reconstructions.contiguous(), reduction="none") * self.l2_loss_weight + if self.perceptual_weight > 0: + p_loss = self.perceptual_loss(inputs.contiguous(), reconstructions.contiguous()) + rec_loss = rec_loss + self.perceptual_weight * p_loss + rec_loss = rec_loss * loss_weights + + nll_loss = rec_loss / torch.exp(self.logvar) + self.logvar + weighted_nll_loss = nll_loss + if weights is not None: + weighted_nll_loss = weights*nll_loss + weighted_nll_loss = torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0] + nll_loss = torch.sum(nll_loss) / nll_loss.shape[0] + kl_loss = posteriors.kl() + kl_loss = torch.sum(kl_loss) / kl_loss.shape[0] + + # now the GAN part + if optimizer_idx == 0: + # generator update + if cond is None: + assert not self.disc_conditional + logits_fake = self.discriminator(reconstructions.contiguous()) + else: + assert self.disc_conditional + logits_fake = self.discriminator(torch.cat((reconstructions.contiguous(), cond), dim=1)) + logits_fake_3d = self.discriminator3d(reconstructions_ori.contiguous()) + g_loss = -torch.mean(logits_fake) - torch.mean(logits_fake_3d) + + if self.disc_factor > 0.0: + try: + d_weight = self.calculate_adaptive_weight(nll_loss, g_loss, last_layer=last_layer) + except RuntimeError: + assert not self.training + d_weight = torch.tensor(0.0) + else: + d_weight = torch.tensor(0.0) + + disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start) + loss = weighted_nll_loss + self.kl_weight * kl_loss + d_weight * disc_factor * g_loss + + log = {"{}/total_loss".format(split): loss.clone().detach().mean(), "{}/logvar".format(split): self.logvar.detach(), + "{}/kl_loss".format(split): kl_loss.detach().mean(), "{}/nll_loss".format(split): nll_loss.detach().mean(), + "{}/rec_loss".format(split): rec_loss.detach().mean(), + "{}/d_weight".format(split): d_weight.detach(), + "{}/disc_factor".format(split): torch.tensor(disc_factor), + "{}/g_loss".format(split): g_loss.detach().mean(), + } + return loss, log + + if optimizer_idx == 1: + # second pass for discriminator update + if cond is None: + logits_real = self.discriminator(inputs.contiguous().detach()) + logits_fake = self.discriminator(reconstructions.contiguous().detach()) + else: + logits_real = self.discriminator(torch.cat((inputs.contiguous().detach(), cond), dim=1)) + logits_fake = self.discriminator(torch.cat((reconstructions.contiguous().detach(), cond), dim=1)) + logits_real_3d = self.discriminator3d(inputs_ori.contiguous().detach()) + logits_fake_3d = self.discriminator3d(reconstructions_ori.contiguous().detach()) + + disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start) + d_loss = disc_factor * self.disc_loss(logits_real, logits_fake) + disc_factor * self.disc_loss(logits_real_3d, logits_fake_3d) + + log = {"{}/disc_loss".format(split): d_loss.clone().detach().mean(), + "{}/logits_real".format(split): logits_real.detach().mean(), + "{}/logits_fake".format(split): logits_fake.detach().mean() + } + return d_loss, log + diff --git a/ruyi/vae/ldm/modules/losses/vqperceptual.py b/ruyi/vae/ldm/modules/losses/vqperceptual.py new file mode 100644 index 0000000..d6d0f9e --- /dev/null +++ b/ruyi/vae/ldm/modules/losses/vqperceptual.py @@ -0,0 +1,168 @@ +import torch +import torch.nn.functional as F +from einops import repeat +from taming.modules.discriminator.model import (NLayerDiscriminator, + weights_init) +from taming.modules.losses.lpips import LPIPS +from taming.modules.losses.vqperceptual import hinge_d_loss, vanilla_d_loss +from torch import nn + + +def hinge_d_loss_with_exemplar_weights(logits_real, logits_fake, weights): + assert weights.shape[0] == logits_real.shape[0] == logits_fake.shape[0] + loss_real = torch.mean(F.relu(1. - logits_real), dim=[1,2,3]) + loss_fake = torch.mean(F.relu(1. + logits_fake), dim=[1,2,3]) + loss_real = (weights * loss_real).sum() / weights.sum() + loss_fake = (weights * loss_fake).sum() / weights.sum() + d_loss = 0.5 * (loss_real + loss_fake) + return d_loss + +def adopt_weight(weight, global_step, threshold=0, value=0.): + if global_step < threshold: + weight = value + return weight + + +def measure_perplexity(predicted_indices, n_embed): + # src: https://github.com/karpathy/deep-vector-quantization/blob/main/model.py + # eval cluster perplexity. when perplexity == num_embeddings then all clusters are used exactly equally + encodings = F.one_hot(predicted_indices, n_embed).float().reshape(-1, n_embed) + avg_probs = encodings.mean(0) + perplexity = (-(avg_probs * torch.log(avg_probs + 1e-10)).sum()).exp() + cluster_use = torch.sum(avg_probs > 0) + return perplexity, cluster_use + +def l1(x, y): + return torch.abs(x-y) + + +def l2(x, y): + return torch.pow((x-y), 2) + +''' +class VQLPIPSWithDiscriminator(nn.Module): + def __init__(self, disc_start, codebook_weight=1.0, pixelloss_weight=1.0, + disc_num_layers=3, disc_in_channels=3, disc_factor=1.0, disc_weight=1.0, + perceptual_weight=1.0, use_actnorm=False, disc_conditional=False, + disc_ndf=64, disc_loss="hinge", n_classes=None, perceptual_loss="lpips", + pixel_loss="l1"): + super().__init__() + assert disc_loss in ["hinge", "vanilla"] + assert perceptual_loss in ["lpips", "clips", "dists"] + assert pixel_loss in ["l1", "l2"] + self.codebook_weight = codebook_weight + self.pixel_weight = pixelloss_weight + if perceptual_loss == "lpips": + print(f"{self.__class__.__name__}: Running with LPIPS.") + self.perceptual_loss = LPIPS().eval() + else: + raise ValueError(f"Unknown perceptual loss: >> {perceptual_loss} <<") + self.perceptual_weight = perceptual_weight + + if pixel_loss == "l1": + self.pixel_loss = l1 + else: + self.pixel_loss = l2 + + self.discriminator = NLayerDiscriminator(input_nc=disc_in_channels, + n_layers=disc_num_layers, + use_actnorm=use_actnorm, + ndf=disc_ndf + ).apply(weights_init) + self.discriminator_iter_start = disc_start + if disc_loss == "hinge": + self.disc_loss = hinge_d_loss + elif disc_loss == "vanilla": + self.disc_loss = vanilla_d_loss + else: + raise ValueError(f"Unknown GAN loss '{disc_loss}'.") + print(f"VQLPIPSWithDiscriminator running with {disc_loss} loss.") + self.disc_factor = disc_factor + self.discriminator_weight = disc_weight + self.disc_conditional = disc_conditional + self.n_classes = n_classes + + def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None): + if last_layer is not None: + nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0] + g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0] + else: + nll_grads = torch.autograd.grad(nll_loss, self.last_layer[0], retain_graph=True)[0] + g_grads = torch.autograd.grad(g_loss, self.last_layer[0], retain_graph=True)[0] + + d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4) + d_weight = torch.clamp(d_weight, 0.0, 1e4).detach() + d_weight = d_weight * self.discriminator_weight + return d_weight + + def forward(self, codebook_loss, inputs, reconstructions, optimizer_idx, + global_step, last_layer=None, cond=None, split="train", predicted_indices=None): + if not exists(codebook_loss): + codebook_loss = torch.tensor([0.]).to(inputs.device) + #rec_loss = torch.abs(inputs.contiguous() - reconstructions.contiguous()) + rec_loss = self.pixel_loss(inputs.contiguous(), reconstructions.contiguous()) + if self.perceptual_weight > 0: + p_loss = self.perceptual_loss(inputs.contiguous(), reconstructions.contiguous()) + rec_loss = rec_loss + self.perceptual_weight * p_loss + else: + p_loss = torch.tensor([0.0]) + + nll_loss = rec_loss + #nll_loss = torch.sum(nll_loss) / nll_loss.shape[0] + nll_loss = torch.mean(nll_loss) + + # now the GAN part + if optimizer_idx == 0: + # generator update + if cond is None: + assert not self.disc_conditional + logits_fake = self.discriminator(reconstructions.contiguous()) + else: + assert self.disc_conditional + logits_fake = self.discriminator(torch.cat((reconstructions.contiguous(), cond), dim=1)) + g_loss = -torch.mean(logits_fake) + + try: + d_weight = self.calculate_adaptive_weight(nll_loss, g_loss, last_layer=last_layer) + except RuntimeError: + assert not self.training + d_weight = torch.tensor(0.0) + + disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start) + loss = nll_loss + d_weight * disc_factor * g_loss + self.codebook_weight * codebook_loss.mean() + + log = {"{}/total_loss".format(split): loss.clone().detach().mean(), + "{}/quant_loss".format(split): codebook_loss.detach().mean(), + "{}/nll_loss".format(split): nll_loss.detach().mean(), + "{}/rec_loss".format(split): rec_loss.detach().mean(), + "{}/p_loss".format(split): p_loss.detach().mean(), + "{}/d_weight".format(split): d_weight.detach(), + "{}/disc_factor".format(split): torch.tensor(disc_factor), + "{}/g_loss".format(split): g_loss.detach().mean(), + } + if predicted_indices is not None: + assert self.n_classes is not None + with torch.no_grad(): + perplexity, cluster_usage = measure_perplexity(predicted_indices, self.n_classes) + log[f"{split}/perplexity"] = perplexity + log[f"{split}/cluster_usage"] = cluster_usage + return loss, log + + if optimizer_idx == 1: + # second pass for discriminator update + if cond is None: + logits_real = self.discriminator(inputs.contiguous().detach()) + logits_fake = self.discriminator(reconstructions.contiguous().detach()) + else: + logits_real = self.discriminator(torch.cat((inputs.contiguous().detach(), cond), dim=1)) + logits_fake = self.discriminator(torch.cat((reconstructions.contiguous().detach(), cond), dim=1)) + + disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start) + d_loss = disc_factor * self.disc_loss(logits_real, logits_fake) + + log = {"{}/disc_loss".format(split): d_loss.clone().detach().mean(), + "{}/logits_real".format(split): logits_real.detach().mean(), + "{}/logits_fake".format(split): logits_fake.detach().mean() + } + return d_loss, log +''' \ No newline at end of file diff --git a/ruyi/vae/ldm/modules/vaemodules/__init__.py b/ruyi/vae/ldm/modules/vaemodules/__init__.py new file mode 100755 index 0000000..e69de29 diff --git a/ruyi/vae/ldm/modules/vaemodules/activations.py b/ruyi/vae/ldm/modules/vaemodules/activations.py new file mode 100755 index 0000000..ab2d6a9 --- /dev/null +++ b/ruyi/vae/ldm/modules/vaemodules/activations.py @@ -0,0 +1,27 @@ +import torch.nn as nn + +ACTIVATION_FUNCTIONS = { + "elu": nn.ELU(), + "swish": nn.SiLU(), + "silu": nn.SiLU(), + "mish": nn.Mish(), + "gelu": nn.GELU(), + "relu": nn.ReLU(), +} + + +def get_activation(act_fn: str) -> nn.Module: + """Helper function to get activation function from string. + + Args: + act_fn (str): Name of activation function. + + Returns: + nn.Module: Activation function. + """ + + act_fn = act_fn.lower() + if act_fn in ACTIVATION_FUNCTIONS: + return ACTIVATION_FUNCTIONS[act_fn] + else: + raise ValueError(f"Unsupported activation function: {act_fn}") diff --git a/ruyi/vae/ldm/modules/vaemodules/attention.py b/ruyi/vae/ldm/modules/vaemodules/attention.py new file mode 100644 index 0000000..5b4dacc --- /dev/null +++ b/ruyi/vae/ldm/modules/vaemodules/attention.py @@ -0,0 +1,479 @@ +import inspect + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +from .attention_processors import AttnProcessor, AttnProcessor2_0 +from .common import SpatialNorm3D + + +class Attention(nn.Module): + r""" + A cross attention layer. + + Parameters: + query_dim (`int`): + The number of channels in the query. + cross_attention_dim (`int`, *optional*): + The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`. + nheads (`int`, *optional*, defaults to 8): + The number of heads to use for multi-head attention. + head_dim (`int`, *optional*, defaults to 64): + The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): + The dropout probability to use. + bias (`bool`, *optional*, defaults to False): + Set to `True` for the query, key, and value linear layers to contain a bias parameter. + upcast_attention (`bool`, *optional*, defaults to False): + Set to `True` to upcast the attention computation to `float32`. + upcast_softmax (`bool`, *optional*, defaults to False): + Set to `True` to upcast the softmax computation to `float32`. + cross_attention_norm (`str`, *optional*, defaults to `None`): + The type of normalization to use for the cross attention. Can be `None`, `layer_norm`, or `group_norm`. + cross_attention_norm_num_groups (`int`, *optional*, defaults to 32): + The number of groups to use for the group norm in the cross attention. + added_kv_proj_dim (`int`, *optional*, defaults to `None`): + The number of channels to use for the added key and value projections. If `None`, no projection is used. + norm_num_groups (`int`, *optional*, defaults to `None`): + The number of groups to use for the group norm in the attention. + spatial_norm_dim (`int`, *optional*, defaults to `None`): + The number of channels to use for the spatial normalization. + out_bias (`bool`, *optional*, defaults to `True`): + Set to `True` to use a bias in the output linear layer. + scale_qk (`bool`, *optional*, defaults to `True`): + Set to `True` to scale the query and key by `1 / sqrt(dim_head)`. + only_cross_attention (`bool`, *optional*, defaults to `False`): + Set to `True` to only use cross attention and not added_kv_proj_dim. Can only be set to `True` if + `added_kv_proj_dim` is not `None`. + eps (`float`, *optional*, defaults to 1e-5): + An additional value added to the denominator in group normalization that is used for numerical stability. + rescale_output_factor (`float`, *optional*, defaults to 1.0): + A factor to rescale the output by dividing it with this value. + residual_connection (`bool`, *optional*, defaults to `False`): + Set to `True` to add the residual connection to the output. + _from_deprecated_attn_block (`bool`, *optional*, defaults to `False`): + Set to `True` if the attention block is loaded from a deprecated state dict. + processor (`AttnProcessor`, *optional*, defaults to `None`): + The attention processor to use. If `None`, defaults to `AttnProcessor2_0` if `torch 2.x` is used and + `AttnProcessor` otherwise. + """ + + def __init__( + self, + query_dim: int, + cross_attention_dim: int = None, + nheads: int = 8, + head_dim: int = 64, + dropout: float = 0.0, + bias: bool = False, + upcast_attention: bool = False, + upcast_softmax: bool = False, + cross_attention_norm = None, + cross_attention_norm_num_groups: int = 32, + added_kv_proj_dim = None, + norm_num_groups = None, + spatial_norm_dim = None, + out_bias: bool = True, + scale_qk: bool = True, + only_cross_attention: bool = False, + eps: float = 1e-5, + rescale_output_factor: float = 1.0, + residual_connection: bool = False, + processor = None, + out_dim: int = None, + ): + super().__init__() + + self.query_dim = query_dim + self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim + self.inner_dim = out_dim if out_dim is not None else head_dim * nheads + self.nheads = out_dim // head_dim if out_dim is not None else nheads + self.out_dim = out_dim if out_dim is not None else query_dim + + self.upcast_attention = upcast_attention + self.upcast_softmax = upcast_softmax + self.added_kv_proj_dim = added_kv_proj_dim + self.only_cross_attention = only_cross_attention + + if self.added_kv_proj_dim is None and self.only_cross_attention: + raise ValueError( + "`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None. Make sure to set either `only_cross_attention=False` or define `added_kv_proj_dim`." + ) + + self.scale_qk = scale_qk + self.scale = head_dim ** -0.5 if scale_qk else 1.0 + + self.rescale_output_factor = rescale_output_factor + self.residual_connection = residual_connection + + if norm_num_groups is not None: + self.group_norm = nn.GroupNorm(num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True) + else: + self.group_norm = None + + if spatial_norm_dim is not None: + self.spatial_norm = SpatialNorm3D(f_channels=query_dim, zq_channels=spatial_norm_dim) + else: + self.spatial_norm = None + + if cross_attention_norm is None: + self.norm_cross = None + elif cross_attention_norm == "layer_norm": + self.norm_cross = nn.LayerNorm(self.cross_attention_dim) + elif cross_attention_norm == "group_norm": + if self.added_kv_proj_dim is not None: + # The given `encoder_hidden_states` are initially of shape + # (batch_size, seq_len, added_kv_proj_dim) before being projected + # to (batch_size, seq_len, cross_attention_dim). The norm is applied + # before the projection, so we need to use `added_kv_proj_dim` as + # the number of channels for the group norm. + norm_cross_num_channels = added_kv_proj_dim + else: + norm_cross_num_channels = self.cross_attention_dim + + self.norm_cross = nn.GroupNorm( + num_channels=norm_cross_num_channels, num_groups=cross_attention_norm_num_groups, eps=1e-5, affine=True + ) + else: + raise ValueError( + f"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'" + ) + + self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias) + + if not self.only_cross_attention: + self.to_k = nn.Linear(self.cross_attention_dim, self.inner_dim, bias=bias) + self.to_v = nn.Linear(self.cross_attention_dim, self.inner_dim, bias=bias) + else: + self.to_k = None + self.to_v = None + + if self.added_kv_proj_dim is not None: + self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_dim) + self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_dim) + + self.to_out = nn.Linear(self.inner_dim, self.out_dim, bias=out_bias) + self.dropout = nn.Dropout(dropout) + + if processor is None: + processor = ( + AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnProcessor() + ) + self.set_processor(processor) + + def set_processor(self, processor: AttnProcessor) -> None: + r""" + Set the attention processor to use. + + Args: + processor (`AttnProcessor`): + The attention processor to use. + """ + # if current processor is in `self._modules` and if passed `processor` is not, we need to + # pop `processor` from `self._modules` + if ( + hasattr(self, "processor") + and isinstance(self.processor, torch.nn.Module) + and not isinstance(processor, torch.nn.Module) + ): + self._modules.pop("processor") + + self.processor = processor + self._attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) + + def prepare_attention_mask( + self, attention_mask: torch.Tensor, target_length: int, batch_size: int, out_dim: int = 3 + ) -> torch.Tensor: + r""" + Prepare the attention mask for the attention computation. + + Args: + attention_mask (`torch.Tensor`): + The attention mask to prepare. + target_length (`int`): + The target length of the attention mask. This is the length of the attention mask after padding. + batch_size (`int`): + The batch size, which is used to repeat the attention mask. + out_dim (`int`, *optional*, defaults to `3`): + The output dimension of the attention mask. Can be either `3` or `4`. + + Returns: + `torch.Tensor`: The prepared attention mask. + """ + head_size = self.nheads + if attention_mask is None: + return attention_mask + + current_length: int = attention_mask.shape[-1] + if current_length != target_length: + if attention_mask.device.type == "mps": + # HACK: MPS: Does not support padding by greater than dimension of input tensor. + # Instead, we can manually construct the padding tensor. + padding_shape = (attention_mask.shape[0], attention_mask.shape[1], target_length) + padding = torch.zeros(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device) + attention_mask = torch.cat([attention_mask, padding], dim=2) + else: + # TODO: for pipelines such as stable-diffusion, padding cross-attn mask: + # we want to instead pad by (0, remaining_length), where remaining_length is: + # remaining_length: int = target_length - current_length + attention_mask = F.pad(attention_mask, (0, target_length), value=0.0) + + if out_dim == 3: + if attention_mask.shape[0] < batch_size * head_size: + attention_mask = attention_mask.repeat_interleave(head_size, dim=0) + elif out_dim == 4: + attention_mask = attention_mask.unsqueeze(1) + attention_mask = attention_mask.repeat_interleave(head_size, dim=1) + + return attention_mask + + def get_attention_scores( + self, query: torch.Tensor, key: torch.Tensor, attention_mask: torch.Tensor = None + ) -> torch.Tensor: + r""" + Compute the attention scores. + + Args: + query (`torch.Tensor`): The query tensor. + key (`torch.Tensor`): The key tensor. + attention_mask (`torch.Tensor`, *optional*): The attention mask to use. If `None`, no mask is applied. + + Returns: + `torch.Tensor`: The attention probabilities/scores. + """ + dtype = query.dtype + if self.upcast_attention: + query = query.float() + key = key.float() + + if attention_mask is None: + baddbmm_input = torch.empty( + query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device + ) + beta = 0 + else: + baddbmm_input = attention_mask + beta = 1 + + attention_scores = torch.baddbmm( + baddbmm_input, + query, + key.transpose(-1, -2), + beta=beta, + alpha=self.scale, + ) + del baddbmm_input + + if self.upcast_softmax: + attention_scores = attention_scores.float() + + attention_probs = attention_scores.softmax(dim=-1) + del attention_scores + + attention_probs = attention_probs.to(dtype) + + return attention_probs + + def norm_encoder_hidden_states(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor: + r""" + Normalize the encoder hidden states. Requires `self.norm_cross` to be specified when constructing the + `Attention` class. + + Args: + encoder_hidden_states (`torch.Tensor`): Hidden states of the encoder. + + Returns: + `torch.Tensor`: The normalized encoder hidden states. + """ + assert self.norm_cross is not None, "self.norm_cross must be defined to call self.norm_encoder_hidden_states" + + if isinstance(self.norm_cross, nn.LayerNorm): + encoder_hidden_states = self.norm_cross(encoder_hidden_states) + elif isinstance(self.norm_cross, nn.GroupNorm): + # Group norm norms along the channels dimension and expects + # input to be in the shape of (N, C, *). In this case, we want + # to norm along the hidden dimension, so we need to move + # (batch_size, sequence_length, hidden_size) -> + # (batch_size, hidden_size, sequence_length) + encoder_hidden_states = encoder_hidden_states.transpose(1, 2) + encoder_hidden_states = self.norm_cross(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.transpose(1, 2) + else: + assert False + + return encoder_hidden_states + + def batch_to_head_dim(self, tensor: torch.Tensor) -> torch.Tensor: + r""" + Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size // nheads, seq_len, dim * nheads]`. `nheads` + is the number of heads initialized while constructing the `Attention` class. + + Args: + tensor (`torch.Tensor`): The tensor to reshape. + + Returns: + `torch.Tensor`: The reshaped tensor. + """ + head_size = self.nheads + batch_size, seq_len, dim = tensor.shape + tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim) + tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size) + return tensor + + def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Tensor: + r""" + Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size, seq_len, nheads, dim // nheads]` `nheads` is + the number of heads initialized while constructing the `Attention` class. + + Args: + tensor (`torch.Tensor`): The tensor to reshape. + out_dim (`int`, *optional*, defaults to `3`): The output dimension of the tensor. If `3`, the tensor is + reshaped to `[batch_size * nheads, seq_len, dim // nheads]`. + + Returns: + `torch.Tensor`: The reshaped tensor. + """ + head_size = self.nheads + batch_size, seq_len, dim = tensor.shape + tensor = tensor.reshape(batch_size, seq_len, head_size, dim // head_size) + tensor = tensor.permute(0, 2, 1, 3) + + if out_dim == 3: + tensor = tensor.reshape(batch_size * head_size, seq_len, dim // head_size) + + return tensor + + def forward( + self, + hidden_states: torch.FloatTensor, + encoder_hidden_states: torch.FloatTensor = None, + attention_mask: torch.FloatTensor = None, + **cross_attention_kwargs, + ) -> torch.Tensor: + r""" + The forward method of the `Attention` class. + + Args: + hidden_states (`torch.Tensor`): + The hidden states of the query. + encoder_hidden_states (`torch.Tensor`, *optional*): + The hidden states of the encoder. + attention_mask (`torch.Tensor`, *optional*): + The attention mask to use. If `None`, no mask is applied. + **cross_attention_kwargs: + Additional keyword arguments to pass along to the cross attention. + + Returns: + `torch.Tensor`: The output of the attention layer. + """ + # The `Attention` class can call different attention processors / attention functions + # here we simply pass along all tensors to the selected processor class + # For standard processors that are defined here, `**cross_attention_kwargs` is empty + + unused_kwargs = [k for k, _ in cross_attention_kwargs.items() if k not in self._attn_parameters] + # if len(unused_kwargs) > 0: + # logger.warning( + # f"cross_attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored." + # ) + cross_attention_kwargs = {k: w for k, w in cross_attention_kwargs.items() if k in self._attn_parameters} + + return self.processor( + self, + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + + +class SpatialAttention(Attention): + def forward( + self, + hidden_states: torch.FloatTensor, + encoder_hidden_states: torch.FloatTensor = None, + attention_mask: torch.FloatTensor = None, + **cross_attention_kwargs, + ) -> torch.Tensor: + is_image = hidden_states.ndim == 4 + if is_image: + hidden_states = rearrange(hidden_states, "b c h w -> b c 1 h w") + + bsz, h = hidden_states.shape[0], hidden_states.shape[3] + hidden_states = rearrange(hidden_states, "b c t h w -> (b t) (h w) c") + if encoder_hidden_states is not None: + encoder_hidden_states = rearrange(encoder_hidden_states, "b c t h w -> (b t) (h w) c") + + if attention_mask is not None: + attention_mask = rearrange(attention_mask, "b t h w -> (b t) (h w)") + + hidden_states = super().forward( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + + hidden_states = rearrange(hidden_states, "(b t) (h w) c -> b c t h w", b=bsz, h=h) + + if is_image: + hidden_states = rearrange(hidden_states, "b c 1 h w -> b c h w") + + return hidden_states + + +class TemporalAttention(Attention): + def forward( + self, + hidden_states: torch.FloatTensor, + encoder_hidden_states: torch.FloatTensor = None, + attention_mask: torch.FloatTensor = None, + **cross_attention_kwargs, + ) -> torch.Tensor: + bsz, h = hidden_states.shape[0], hidden_states.shape[3] + hidden_states = rearrange(hidden_states, "b c t h w -> (b h w) t c") + if encoder_hidden_states is not None: + encoder_hidden_states = rearrange(encoder_hidden_states, "b c t h w -> (b h w) t c") + + if attention_mask is not None: + attention_mask = rearrange(attention_mask, "b t h w -> (b h w) t") + + hidden_states = super().forward( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + + hidden_states = rearrange(hidden_states, "(b h w) t c -> b c t h w", b=bsz, h=h) + + return hidden_states + + +class Attention3D(Attention): + def forward( + self, + hidden_states: torch.FloatTensor, + encoder_hidden_states: torch.FloatTensor = None, + attention_mask: torch.FloatTensor = None, + **cross_attention_kwargs, + ) -> torch.Tensor: + t, h = hidden_states.shape[2], hidden_states.shape[3] + hidden_states = rearrange(hidden_states, "b c t h w -> b (t h w) c") + if encoder_hidden_states is not None: + encoder_hidden_states = rearrange(encoder_hidden_states, "b c t h w -> b (t h w) c") + + if attention_mask is not None: + attention_mask = rearrange(attention_mask, "b t h w -> b (t h w)") + + hidden_states = super().forward( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + + hidden_states = rearrange(hidden_states, "b (t h w) c -> b c t h w", t=t, h=h) + + return hidden_states diff --git a/ruyi/vae/ldm/modules/vaemodules/attention_processors.py b/ruyi/vae/ldm/modules/vaemodules/attention_processors.py new file mode 100644 index 0000000..0d7c05b --- /dev/null +++ b/ruyi/vae/ldm/modules/vaemodules/attention_processors.py @@ -0,0 +1,139 @@ +from typing import TYPE_CHECKING + +import torch +import torch.nn.functional as F + +if TYPE_CHECKING: + from .attention import Attention + +class AttnProcessor: + r""" + Default processor for performing attention-related computations. + """ + + def __call__( + self, + attn: "Attention", + hidden_states: torch.FloatTensor, + encoder_hidden_states, + attention_mask, + temb = None, + ) -> torch.Tensor: + residual = hidden_states + + if attn.spatial_norm is not None: + hidden_states = attn.spatial_norm(hidden_states, temb = None) + + # B, L, C + assert hidden_states.ndim == 3, f"Hidden states must be 3-dimensional, got {hidden_states.ndim}" + + batch_size, sequence_length, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + + if attn.group_norm is not None: + hidden_states = attn.group_norm(hidden_states.transpose(1, 2)) + hidden_states = hidden_states.transpose(1, 2) + + query = attn.to_q(hidden_states) + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + elif attn.norm_cross: + encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + query = attn.head_to_batch_dim(query) + key = attn.head_to_batch_dim(key) + value = attn.head_to_batch_dim(value) + + attention_probs = attn.get_attention_scores(query, key, attention_mask) + hidden_states = torch.bmm(attention_probs, value) + hidden_states = attn.batch_to_head_dim(hidden_states) + + hidden_states = attn.to_out(hidden_states) + hidden_states = attn.dropout(hidden_states) + + if attn.residual_connection: + hidden_states = hidden_states + residual + + hidden_states = hidden_states / attn.rescale_output_factor + + return hidden_states + +class AttnProcessor2_0: + r""" + Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). + """ + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + + def __call__( + self, + attn: "Attention", + hidden_states: torch.FloatTensor, + encoder_hidden_states, + attention_mask, + temb = None, + ) -> torch.FloatTensor: + residual = hidden_states + if attn.spatial_norm is not None: + hidden_states = attn.spatial_norm(hidden_states, temb = None) + + # B, L, C + assert hidden_states.ndim == 3, f"Hidden states must be 3-dimensional, got {hidden_states.ndim}" + + batch_size, sequence_length, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + + if attention_mask is not None: + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + # scaled_dot_product_attention expects attention_mask shape to be + # (batch, heads, source_length, target_length) + attention_mask = attention_mask.view(batch_size, attn.nheads, -1, attention_mask.shape[-1]) + + if attn.group_norm is not None: + hidden_states = attn.group_norm(hidden_states.transpose(1, 2)) + hidden_states = hidden_states.transpose(1, 2) + + query: torch.Tensor = attn.to_q(hidden_states) + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + elif attn.norm_cross: + encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + + key: torch.Tensor = attn.to_k(encoder_hidden_states) + value: torch.Tensor = attn.to_v(encoder_hidden_states) + + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.nheads + + query = query.view(batch_size, -1, attn.nheads, head_dim).transpose(1, 2) + + key = key.view(batch_size, -1, attn.nheads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.nheads, head_dim).transpose(1, 2) + + # the output of sdp = (batch, num_heads, seq_len, head_dim) + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False, scale=attn.scale + ) + + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.nheads * head_dim) + hidden_states = hidden_states.to(query.dtype) + + hidden_states = attn.to_out(hidden_states) + hidden_states = attn.dropout(hidden_states) + + if attn.residual_connection: + hidden_states = hidden_states + residual + + hidden_states = hidden_states / attn.rescale_output_factor + + return hidden_states diff --git a/ruyi/vae/ldm/modules/vaemodules/common.py b/ruyi/vae/ldm/modules/vaemodules/common.py new file mode 100755 index 0000000..b82cb6a --- /dev/null +++ b/ruyi/vae/ldm/modules/vaemodules/common.py @@ -0,0 +1,302 @@ +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange, repeat + +from .activations import get_activation + + +def cast_tuple(t, length = 1): + return t if isinstance(t, tuple) else ((t,) * length) + +def divisible_by(num, den): + return (num % den) == 0 + +def is_odd(n): + return not divisible_by(n, 2) + +class CausalConv3d(nn.Conv3d): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size=3, # : int | tuple[int, int, int], + stride=1, # : int | tuple[int, int, int] = 1, + padding=1, # : int | tuple[int, int, int], # TODO: change it to 0. + dilation=1, # : int | tuple[int, int, int] = 1, + **kwargs, + ): + kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size,) * 3 + assert len(kernel_size) == 3, f"Kernel size must be a 3-tuple, got {kernel_size} instead." + + stride = stride if isinstance(stride, tuple) else (stride,) * 3 + assert len(stride) == 3, f"Stride must be a 3-tuple, got {stride} instead." + + dilation = dilation if isinstance(dilation, tuple) else (dilation,) * 3 + assert len(dilation) == 3, f"Dilation must be a 3-tuple, got {dilation} instead." + + t_ks, h_ks, w_ks = kernel_size + self.t_stride, h_stride, w_stride = stride + t_dilation, h_dilation, w_dilation = dilation + + t_pad = (t_ks - 1) * t_dilation + # TODO: align with SD + if padding is None: + h_pad = math.ceil(((h_ks - 1) * h_dilation + (1 - h_stride)) / 2) + w_pad = math.ceil(((w_ks - 1) * w_dilation + (1 - w_stride)) / 2) + elif isinstance(padding, int): + h_pad = w_pad = padding + else: + assert NotImplementedError + + self.temporal_padding = t_pad + self.temporal_padding_origin = math.ceil(((t_ks - 1) * w_dilation + (1 - w_stride)) / 2) + self.padding_flag = 0 + self.prev_features = None + + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + padding=(0, h_pad, w_pad), + **kwargs, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x: (B, C, T, H, W) + dtype = x.dtype + x = x.float() + if self.padding_flag == 0: + x = F.pad( + x, + pad=(0, 0, 0, 0, self.temporal_padding, 0), + mode="replicate", # TODO: check if this is necessary + ) + x = x.to(dtype=dtype) + return super().forward(x) + elif self.padding_flag == 5: + x = F.pad( + x, + pad=(0, 0, 0, 0, self.temporal_padding, 0), + mode="replicate", # TODO: check if this is necessary + ) + x = x.to(dtype=dtype) + self.prev_features = x[:, :, -self.temporal_padding:] + return super().forward(x) + elif self.padding_flag == 6: + if self.t_stride == 2: + x = torch.concat( + [self.prev_features[:, :, -(self.temporal_padding - 1):], x], dim = 2 + ) + else: + x = torch.concat( + [self.prev_features, x], dim = 2 + ) + self.prev_features = x[:, :, -self.temporal_padding:] + x = x.to(dtype=dtype) + return super().forward(x) + else: + x = F.pad( + x, + pad=(0, 0, 0, 0, self.temporal_padding_origin, self.temporal_padding_origin), + ) + x = x.to(dtype=dtype) + return super().forward(x) + +class ResidualBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + non_linearity: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + output_scale_factor: float = 1.0, + ): + super().__init__() + + self.output_scale_factor = output_scale_factor + + self.norm1 = nn.GroupNorm( + num_groups=norm_num_groups, + num_channels=in_channels, + eps=norm_eps, + affine=True, + ) + + self.nonlinearity = get_activation(non_linearity) + + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) + + self.norm2 = nn.GroupNorm( + num_groups=norm_num_groups, + num_channels=out_channels, + eps=norm_eps, + affine=True, + ) + + self.dropout = nn.Dropout(dropout) + + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1) + + if in_channels != out_channels: + self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1) + else: + self.shortcut = nn.Identity() + + self.set_3dgroupnorm = False + + def forward(self, x: torch.Tensor) -> torch.Tensor: + shortcut = self.shortcut(x) + + if self.set_3dgroupnorm: + batch_size = x.shape[0] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.norm1(x) + x = rearrange(x, "(b t) c h w -> b c t h w", b=batch_size) + else: + x = self.norm1(x) + x = self.nonlinearity(x) + + x = self.conv1(x) + + if self.set_3dgroupnorm: + batch_size = x.shape[0] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.norm2(x) + x = rearrange(x, "(b t) c h w -> b c t h w", b=batch_size) + else: + x = self.norm2(x) + x = self.nonlinearity(x) + + x = self.dropout(x) + x = self.conv2(x) + + return (x + shortcut) / self.output_scale_factor + + +class ResidualBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + non_linearity: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + output_scale_factor: float = 1.0, + ): + super().__init__() + + self.output_scale_factor = output_scale_factor + + self.norm1 = nn.GroupNorm( + num_groups=norm_num_groups, + num_channels=in_channels, + eps=norm_eps, + affine=True, + ) + + self.nonlinearity = get_activation(non_linearity) + + self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3) + + self.norm2 = nn.GroupNorm( + num_groups=norm_num_groups, + num_channels=out_channels, + eps=norm_eps, + affine=True, + ) + + self.dropout = nn.Dropout(dropout) + + self.conv2 = CausalConv3d(out_channels, out_channels, kernel_size=3) + + if in_channels != out_channels: + self.shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1) + else: + self.shortcut = nn.Identity() + + self.set_3dgroupnorm = False + + def forward(self, x: torch.Tensor) -> torch.Tensor: + shortcut = self.shortcut(x) + + if self.set_3dgroupnorm: + batch_size = x.shape[0] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.norm1(x) + x = rearrange(x, "(b t) c h w -> b c t h w", b=batch_size) + else: + x = self.norm1(x) + x = self.nonlinearity(x) + + x = self.conv1(x) + + if self.set_3dgroupnorm: + batch_size = x.shape[0] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.norm2(x) + x = rearrange(x, "(b t) c h w -> b c t h w", b=batch_size) + else: + x = self.norm2(x) + x = self.nonlinearity(x) + + x = self.dropout(x) + x = self.conv2(x) + return (x + shortcut) / self.output_scale_factor + + +class SpatialNorm2D(nn.Module): + """ + Spatially conditioned normalization as defined in https://arxiv.org/abs/2209.09002. + + Args: + f_channels (`int`): + The number of channels for input to group normalization layer, and output of the spatial norm layer. + zq_channels (`int`): + The number of channels for the quantized vector as described in the paper. + """ + + def __init__( + self, + f_channels: int, + zq_channels: int, + ): + super().__init__() + + self.norm = nn.GroupNorm(num_channels=f_channels, num_groups=32, eps=1e-6, affine=True) + self.conv_y = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) + self.conv_b = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) + self.set_3dgroupnorm = False + + def forward(self, f: torch.FloatTensor, zq: torch.FloatTensor) -> torch.FloatTensor: + f_size = f.shape[-2:] + zq = F.interpolate(zq, size=f_size, mode="nearest") + if self.set_3dgroupnorm: + batch_size = f.shape[0] + f = rearrange(f, "b c t h w -> (b t) c h w") + norm_f = self.norm(f) + norm_f = rearrange(norm_f, "(b t) c h w -> b c t h w", b=batch_size) + else: + norm_f = self.norm(f) + new_f = norm_f * self.conv_y(zq) + self.conv_b(zq) + return new_f + + +class SpatialNorm3D(SpatialNorm2D): + def forward(self, f: torch.FloatTensor, zq: torch.FloatTensor) -> torch.FloatTensor: + batch_size = f.shape[0] + f = rearrange(f, "b c t h w -> (b t) c h w") + zq = rearrange(zq, "b c t h w -> (b t) c h w") + + x = super().forward(f, zq) + + x = rearrange(x, "(b t) c h w -> b c t h w", b=batch_size) + + return x diff --git a/ruyi/vae/ldm/modules/vaemodules/discriminator.py b/ruyi/vae/ldm/modules/vaemodules/discriminator.py new file mode 100644 index 0000000..a88e07b --- /dev/null +++ b/ruyi/vae/ldm/modules/vaemodules/discriminator.py @@ -0,0 +1,214 @@ +import math + +import torch +import torch.nn as nn + +from .downsamplers import BlurPooling2D, BlurPooling3D + + +class DiscriminatorBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + output_scale_factor: float = 1.0, + add_downsample: bool = True, + ): + super().__init__() + + self.output_scale_factor = output_scale_factor + + self.norm1 = nn.BatchNorm2d(in_channels) + + self.nonlinearity = nn.LeakyReLU(0.2) + + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) + + if add_downsample: + self.downsampler = BlurPooling2D(out_channels, out_channels) + else: + self.downsampler = nn.Identity() + + self.norm2 = nn.BatchNorm2d(out_channels) + + self.dropout = nn.Dropout(dropout) + + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1) + + if add_downsample: + self.shortcut = nn.Sequential( + BlurPooling2D(in_channels, in_channels), + nn.Conv2d(in_channels, out_channels, kernel_size=1), + ) + else: + self.shortcut = nn.Identity() + + self.spatial_downsample_factor = 2 + self.temporal_downsample_factor = 1 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + shortcut = self.shortcut(x) + + x = self.norm1(x) + x = self.nonlinearity(x) + + x = self.conv1(x) + + x = self.norm2(x) + x = self.nonlinearity(x) + + x = self.dropout(x) + x = self.downsampler(x) + x = self.conv2(x) + + return (x + shortcut) / self.output_scale_factor + + +class Discriminator2D(nn.Module): + def __init__( + self, + in_channels: int = 3, + block_out_channels = (64,), + ): + super().__init__() + + self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, padding=1) + + self.blocks = nn.ModuleList([]) + + output_channels = block_out_channels[0] + for i, out_channels in enumerate(block_out_channels): + input_channels = output_channels + output_channels = out_channels + is_final_block = i == len(block_out_channels) - 1 + + self.blocks.append( + DiscriminatorBlock2D( + in_channels=input_channels, + out_channels=output_channels, + output_scale_factor=math.sqrt(2), + add_downsample=not is_final_block, + ) + ) + + self.conv_norm_out = nn.BatchNorm2d(block_out_channels[-1]) + self.conv_act = nn.LeakyReLU(0.2) + + self.conv_out = nn.Conv2d(block_out_channels[-1], 1, kernel_size=3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x: (B, C, H, W) + x = self.conv_in(x) + + for block in self.blocks: + x = block(x) + + x = self.conv_out(x) + + return x + + +class DiscriminatorBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + output_scale_factor: float = 1.0, + add_downsample: bool = True, + ): + super().__init__() + + self.output_scale_factor = output_scale_factor + + self.norm1 = nn.GroupNorm(32, in_channels) + + self.nonlinearity = nn.LeakyReLU(0.2) + + self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1) + + if add_downsample: + self.downsampler = BlurPooling3D(out_channels, out_channels) + else: + self.downsampler = nn.Identity() + + self.norm2 = nn.GroupNorm(32, out_channels) + + self.dropout = nn.Dropout(dropout) + + self.conv2 = nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1) + + if add_downsample: + self.shortcut = nn.Sequential( + BlurPooling3D(in_channels, in_channels), + nn.Conv3d(in_channels, out_channels, kernel_size=1), + ) + else: + self.shortcut = nn.Sequential( + nn.Conv3d(in_channels, out_channels, kernel_size=1), + ) + + self.spatial_downsample_factor = 2 + self.temporal_downsample_factor = 2 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + shortcut = self.shortcut(x) + + x = self.norm1(x) + x = self.nonlinearity(x) + + x = self.conv1(x) + + x = self.norm2(x) + x = self.nonlinearity(x) + + x = self.dropout(x) + x = self.downsampler(x) + x = self.conv2(x) + + return (x + shortcut) / self.output_scale_factor + + +class Discriminator3D(nn.Module): + def __init__( + self, + in_channels: int = 3, + block_out_channels = (64,), + ): + super().__init__() + + self.conv_in = nn.Conv3d(in_channels, block_out_channels[0], kernel_size=3, padding=1, stride=2) + + self.blocks = nn.ModuleList([]) + + output_channels = block_out_channels[0] + for i, out_channels in enumerate(block_out_channels): + input_channels = output_channels + output_channels = out_channels + is_final_block = i == len(block_out_channels) - 1 + + self.blocks.append( + DiscriminatorBlock3D( + in_channels=input_channels, + out_channels=output_channels, + output_scale_factor=math.sqrt(2), + add_downsample=not is_final_block, + ) + ) + + self.conv_norm_out = nn.GroupNorm(32, block_out_channels[-1]) + self.conv_act = nn.LeakyReLU(0.2) + + self.conv_out = nn.Conv3d(block_out_channels[-1], 1, kernel_size=3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x: (B, C, T, H, W) + x = self.conv_in(x) + + for block in self.blocks: + x = block(x) + + x = self.conv_out(x) + + return x diff --git a/ruyi/vae/ldm/modules/vaemodules/down_blocks.py b/ruyi/vae/ldm/modules/vaemodules/down_blocks.py new file mode 100755 index 0000000..48b2c1f --- /dev/null +++ b/ruyi/vae/ldm/modules/vaemodules/down_blocks.py @@ -0,0 +1,533 @@ +import torch +import torch.nn as nn + +from .attention import SpatialAttention, TemporalAttention +from .common import ResidualBlock3D +from .downsamplers import (SpatialDownsampler3D, SpatialTemporalDownsampler3D, + TemporalDownsampler3D) +from .gc_block import GlobalContextBlock + + +def get_down_block( + down_block_type: str, + in_channels: int, + out_channels: int, + num_layers: int, + act_fn: str, + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + num_attention_heads: int = 1, + output_scale_factor: float = 1.0, + add_gc_block: bool = False, + add_downsample: bool = True, +) -> nn.Module: + if down_block_type == "DownBlock3D": + return DownBlock3D( + in_channels=in_channels, + out_channels=out_channels, + num_layers=num_layers, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + add_gc_block=add_gc_block, + ) + elif down_block_type == "SpatialDownBlock3D": + return SpatialDownBlock3D( + in_channels=in_channels, + out_channels=out_channels, + num_layers=num_layers, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + add_gc_block=add_gc_block, + add_downsample=add_downsample, + ) + elif down_block_type == "SpatialAttnDownBlock3D": + return SpatialAttnDownBlock3D( + in_channels=in_channels, + out_channels=out_channels, + num_layers=num_layers, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + attention_head_dim=out_channels // num_attention_heads, + output_scale_factor=output_scale_factor, + add_gc_block=add_gc_block, + add_downsample=add_downsample, + ) + elif down_block_type == "TemporalDownBlock3D": + return TemporalDownBlock3D( + in_channels=in_channels, + out_channels=out_channels, + num_layers=num_layers, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + add_gc_block=add_gc_block, + add_downsample=add_downsample, + ) + elif down_block_type == "TemporalAttnDownBlock3D": + return TemporalAttnDownBlock3D( + in_channels=in_channels, + out_channels=out_channels, + num_layers=num_layers, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + attention_head_dim=out_channels // num_attention_heads, + output_scale_factor=output_scale_factor, + add_gc_block=add_gc_block, + add_downsample=add_downsample, + ) + elif down_block_type == "SpatialTemporalDownBlock3D": + return SpatialTemporalDownBlock3D( + in_channels=in_channels, + out_channels=out_channels, + num_layers=num_layers, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + add_gc_block=add_gc_block, + add_downsample=add_downsample, + ) + else: + raise ValueError(f"Unknown down block type: {down_block_type}") + + +class DownBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + output_scale_factor: float = 1.0, + add_gc_block: bool = False, + ): + super().__init__() + + self.convs = nn.ModuleList([]) + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.convs.append( + ResidualBlock3D( + in_channels=in_channels, + out_channels=out_channels, + non_linearity=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + ) + ) + + if add_gc_block: + self.gc_block = GlobalContextBlock(out_channels, out_channels, fusion_type="mul") + else: + self.gc_block = None + + self.spatial_downsample_factor = 1 + self.temporal_downsample_factor = 1 + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + for conv in self.convs: + x = conv(x) + + if self.gc_block is not None: + x = self.gc_block(x) + + return x + + +class SpatialDownBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + output_scale_factor: float = 1.0, + add_gc_block: bool = False, + add_downsample: bool = True, + ): + super().__init__() + + self.convs = nn.ModuleList([]) + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.convs.append( + ResidualBlock3D( + in_channels=in_channels, + out_channels=out_channels, + non_linearity=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + ) + ) + + if add_gc_block: + self.gc_block = GlobalContextBlock(out_channels, out_channels, fusion_type="mul") + else: + self.gc_block = None + + if add_downsample: + self.downsampler = SpatialDownsampler3D(out_channels, out_channels) + self.spatial_downsample_factor = 2 + else: + self.downsampler = None + self.spatial_downsample_factor = 1 + + self.temporal_downsample_factor = 1 + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + for conv in self.convs: + x = conv(x) + + if self.gc_block is not None: + x = self.gc_block(x) + + if self.downsampler is not None: + x = self.downsampler(x) + + return x + + +class TemporalDownBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + output_scale_factor: float = 1.0, + add_gc_block: bool = False, + add_downsample: bool = True, + ): + super().__init__() + + self.convs = nn.ModuleList([]) + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.convs.append( + ResidualBlock3D( + in_channels=in_channels, + out_channels=out_channels, + non_linearity=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + ) + ) + + if add_gc_block: + self.gc_block = GlobalContextBlock(out_channels, out_channels, fusion_type="mul") + else: + self.gc_block = None + + if add_downsample: + self.downsampler = TemporalDownsampler3D(out_channels, out_channels) + self.temporal_downsample_factor = 2 + else: + self.downsampler = None + self.temporal_downsample_factor = 1 + + self.spatial_downsample_factor = 1 + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + for conv in self.convs: + x = conv(x) + + if self.gc_block is not None: + x = self.gc_block(x) + + if self.downsampler is not None: + x = self.downsampler(x) + + return x + + +class SpatialTemporalDownBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + output_scale_factor: float = 1.0, + add_gc_block: bool = False, + add_downsample: bool = True, + ): + super().__init__() + + self.convs = nn.ModuleList([]) + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.convs.append( + ResidualBlock3D( + in_channels=in_channels, + out_channels=out_channels, + non_linearity=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + ) + ) + + if add_gc_block: + self.gc_block = GlobalContextBlock(out_channels, out_channels, fusion_type="mul") + else: + self.gc_block = None + + if add_downsample: + self.downsampler = SpatialTemporalDownsampler3D(out_channels, out_channels) + self.spatial_downsample_factor = 2 + self.temporal_downsample_factor = 2 + else: + self.downsampler = None + self.spatial_downsample_factor = 1 + self.temporal_downsample_factor = 1 + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + for conv in self.convs: + x = conv(x) + + if self.gc_block is not None: + x = self.gc_block(x) + + if self.downsampler is not None: + x = self.downsampler(x) + + return x + + +class SpatialAttnDownBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + attention_head_dim: int = 1, + output_scale_factor: float = 1.0, + add_gc_block: bool = False, + add_downsample: bool = True, + ): + super().__init__() + + self.convs = nn.ModuleList([]) + self.attentions = nn.ModuleList([]) + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.convs.append( + ResidualBlock3D( + in_channels=in_channels, + out_channels=out_channels, + non_linearity=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + ) + ) + self.attentions.append( + SpatialAttention( + out_channels, + nheads=out_channels // attention_head_dim, + head_dim=attention_head_dim, + bias=True, + upcast_softmax=True, + norm_num_groups=norm_num_groups, + eps=norm_eps, + rescale_output_factor=output_scale_factor, + residual_connection=True, + ) + ) + + if add_gc_block: + self.gc_block = GlobalContextBlock(out_channels, out_channels, fusion_type="mul") + else: + self.gc_block = None + + if add_downsample: + self.downsampler = SpatialDownsampler3D(out_channels, out_channels) + self.spatial_downsample_factor = 2 + else: + self.downsampler = None + self.spatial_downsample_factor = 1 + + self.temporal_downsample_factor = 1 + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + for conv, attn in zip(self.convs, self.attentions): + x = conv(x) + x = attn(x) + + if self.gc_block is not None: + x = self.gc_block(x) + + if self.downsampler is not None: + x = self.downsampler(x) + + return x + + +class TemporalDownBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + output_scale_factor: float = 1.0, + add_gc_block: bool = False, + add_downsample: bool = True, + ): + super().__init__() + + self.convs = nn.ModuleList([]) + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.convs.append( + ResidualBlock3D( + in_channels=in_channels, + out_channels=out_channels, + non_linearity=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + ) + ) + + if add_gc_block: + self.gc_block = GlobalContextBlock(out_channels, out_channels, fusion_type="mul") + else: + self.gc_block = None + + if add_downsample: + self.downsampler = TemporalDownsampler3D(out_channels, out_channels) + self.temporal_downsample_factor = 2 + else: + self.downsampler = None + self.temporal_downsample_factor = 1 + + self.spatial_downsample_factor = 1 + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + for conv in self.convs: + x = conv(x) + + if self.gc_block is not None: + x = self.gc_block(x) + + if self.downsampler is not None: + x = self.downsampler(x) + + return x + + +class TemporalAttnDownBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + attention_head_dim: int = 1, + output_scale_factor: float = 1.0, + add_gc_block: bool = False, + add_downsample: bool = True, + ): + super().__init__() + + self.convs = nn.ModuleList([]) + self.attentions = nn.ModuleList([]) + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.convs.append( + ResidualBlock3D( + in_channels=in_channels, + out_channels=out_channels, + non_linearity=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + ) + ) + self.attentions.append( + TemporalAttention( + out_channels, + nheads=out_channels // attention_head_dim, + head_dim=attention_head_dim, + bias=True, + upcast_softmax=True, + norm_num_groups=norm_num_groups, + eps=norm_eps, + rescale_output_factor=output_scale_factor, + residual_connection=True, + ) + ) + + if add_gc_block: + self.gc_block = GlobalContextBlock(out_channels, out_channels, fusion_type="mul") + else: + self.gc_block = None + + if add_downsample: + self.downsampler = TemporalDownsampler3D(out_channels, out_channels) + self.temporal_downsample_factor = 2 + else: + self.downsampler = None + self.temporal_downsample_factor = 1 + + self.spatial_downsample_factor = 1 + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + for conv, attn in zip(self.convs, self.attentions): + x = conv(x) + x = attn(x) + + if self.gc_block is not None: + x = self.gc_block(x) + + if self.downsampler is not None: + x = self.downsampler(x) + + return x diff --git a/ruyi/vae/ldm/modules/vaemodules/downsamplers.py b/ruyi/vae/ldm/modules/vaemodules/downsamplers.py new file mode 100644 index 0000000..027fe9b --- /dev/null +++ b/ruyi/vae/ldm/modules/vaemodules/downsamplers.py @@ -0,0 +1,148 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .common import CausalConv3d + + +class Downsampler(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + spatial_downsample_factor: int = 1, + temporal_downsample_factor: int = 1, + ): + super().__init__() + + self.in_channels = in_channels + self.out_channels = out_channels + self.spatial_downsample_factor = spatial_downsample_factor + self.temporal_downsample_factor = temporal_downsample_factor + + +class SpatialDownsampler3D(Downsampler): + def __init__(self, in_channels: int, out_channels): + if out_channels is None: + out_channels = in_channels + + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + spatial_downsample_factor=2, + temporal_downsample_factor=1, + ) + + self.conv = CausalConv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=(1, 2, 2), + padding=0, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.pad(x, (0, 1, 0, 1)) + return self.conv(x) + + +class TemporalDownsampler3D(Downsampler): + def __init__(self, in_channels: int, out_channels): + if out_channels is None: + out_channels = in_channels + + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + spatial_downsample_factor=1, + temporal_downsample_factor=2, + ) + + self.conv = CausalConv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=(2, 1, 1), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + +class SpatialTemporalDownsampler3D(Downsampler): + def __init__(self, in_channels: int, out_channels): + if out_channels is None: + out_channels = in_channels + + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + spatial_downsample_factor=2, + temporal_downsample_factor=2, + ) + + self.conv = CausalConv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=(2, 2, 2), + padding=0, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.pad(x, (0, 1, 0, 1)) + return self.conv(x) + + +class BlurPooling2D(Downsampler): + def __init__(self, in_channels: int, out_channels): + if out_channels is None: + out_channels = in_channels + + assert in_channels == out_channels + + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + spatial_downsample_factor=2, + temporal_downsample_factor=1, + ) + + filt = torch.tensor([1, 2, 1], dtype=torch.float32) + filt = torch.einsum("i,j -> ij", filt, filt) + filt = filt / filt.sum() + filt = filt[None, None].repeat(out_channels, 1, 1, 1) + + self.register_buffer("filt", filt) + self.filt: torch.Tensor + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x: (B, C, H, W) + return F.conv2d(x, self.filt, stride=2, padding=1, groups=self.in_channels) + + +class BlurPooling3D(Downsampler): + def __init__(self, in_channels: int, out_channels): + if out_channels is None: + out_channels = in_channels + + assert in_channels == out_channels + + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + spatial_downsample_factor=2, + temporal_downsample_factor=2, + ) + + filt = torch.tensor([1, 2, 1], dtype=torch.float32) + filt = torch.einsum("i,j,k -> ijk", filt, filt, filt) + filt = filt / filt.sum() + filt = filt[None, None].repeat(out_channels, 1, 1, 1, 1) + + self.register_buffer("filt", filt) + self.filt: torch.Tensor + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x: (B, C, T, H, W) + return F.conv3d(x, self.filt, stride=2, padding=1, groups=self.in_channels) diff --git a/ruyi/vae/ldm/modules/vaemodules/gc_block.py b/ruyi/vae/ldm/modules/vaemodules/gc_block.py new file mode 100644 index 0000000..43d2eeb --- /dev/null +++ b/ruyi/vae/ldm/modules/vaemodules/gc_block.py @@ -0,0 +1,79 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + + +class GlobalContextBlock(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + min_channels: int = 16, + init_bias: float = -10., + fusion_type: str = "mul", + ): + super().__init__() + + assert fusion_type in ("mul", "add"), f"Unsupported fusion type: {fusion_type}" + self.fusion_type = fusion_type + + self.conv_ctx = nn.Conv2d(in_channels, 1, kernel_size=1) + + num_channels = max(min_channels, out_channels // 2) + + if fusion_type == "mul": + self.conv_mul = nn.Sequential( + nn.Conv2d(in_channels, num_channels, kernel_size=1), + nn.LayerNorm([num_channels, 1, 1]), # TODO: LayerNorm or GroupNorm? + nn.LeakyReLU(0.1), + nn.Conv2d(num_channels, out_channels, kernel_size=1), + nn.Sigmoid(), + ) + + nn.init.zeros_(self.conv_mul[-2].weight) + nn.init.constant_(self.conv_mul[-2].bias, init_bias) + else: + self.conv_add = nn.Sequential( + nn.Conv2d(in_channels, num_channels, kernel_size=1), + nn.LayerNorm([num_channels, 1, 1]), # TODO: LayerNorm or GroupNorm? + nn.LeakyReLU(0.1), + nn.Conv2d(num_channels, out_channels, kernel_size=1), + ) + + nn.init.zeros_(self.conv_add[-1].weight) + nn.init.constant_(self.conv_add[-1].bias, init_bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + is_image = x.ndim == 4 + if is_image: + x = rearrange(x, "b c h w -> b c 1 h w") + + # x: (B, C, T, H, W) + orig_x = x + batch_size = x.shape[0] + + x = rearrange(x, "b c t h w -> (b t) c h w") + + ctx = self.conv_ctx(x) + ctx = rearrange(ctx, "b c h w -> b c (h w)") + ctx = F.softmax(ctx, dim=-1) + + flattened_x = rearrange(x, "b c h w -> b c (h w)") + + x = torch.einsum("b c1 n, b c2 n -> b c2 c1", ctx, flattened_x) + x = rearrange(x, "... -> ... 1") + + if self.fusion_type == "mul": + mul_term = self.conv_mul(x) + mul_term = rearrange(mul_term, "(b t) c h w -> b c t h w", b=batch_size) + x = orig_x * mul_term + else: + add_term = self.conv_add(x) + add_term = rearrange(add_term, "(b t) c h w -> b c t h w", b=batch_size) + x = orig_x + add_term + + if is_image: + x = rearrange(x, "b c 1 h w -> b c h w") + + return x diff --git a/ruyi/vae/ldm/modules/vaemodules/mid_blocks.py b/ruyi/vae/ldm/modules/vaemodules/mid_blocks.py new file mode 100755 index 0000000..12776f4 --- /dev/null +++ b/ruyi/vae/ldm/modules/vaemodules/mid_blocks.py @@ -0,0 +1,196 @@ +import torch +import torch.nn as nn + +from .attention import Attention3D, SpatialAttention, TemporalAttention +from .common import ResidualBlock3D + + +def get_mid_block( + mid_block_type: str, + in_channels: int, + num_layers: int, + act_fn: str, + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + add_attention: bool = True, + attention_type: str = "3d", + num_attention_heads: int = 1, + output_scale_factor: float = 1.0, +) -> nn.Module: + if mid_block_type == "MidBlock3D": + return MidBlock3D( + in_channels=in_channels, + num_layers=num_layers, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + add_attention=add_attention, + attention_type=attention_type, + attention_head_dim=in_channels // num_attention_heads, + output_scale_factor=output_scale_factor, + ) + else: + raise ValueError(f"Unknown mid block type: {mid_block_type}") + + +class MidBlock3D(nn.Module): + """ + A 3D UNet mid-block [`MidBlock3D`] with multiple residual blocks and optional attention blocks. + + Args: + in_channels (`int`): The number of input channels. + num_layers (`int`, *optional*, defaults to 1): The number of residual blocks. + act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks. + norm_num_groups (`int`, *optional*, defaults to 32): + The number of groups to use in the group normalization layers of the resnet blocks. + norm_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks. + dropout (`float`, *optional*, defaults to 0.0): The dropout rate. + add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks. + attention_type: (`str`, *optional*, defaults to `3d`): The type of attention to use. Defaults to `3d`. + attention_head_dim (`int`, *optional*, defaults to 1): + Dimension of a single attention head. The number of attention heads is determined based on this value and + the number of input channels. + output_scale_factor (`float`, *optional*, defaults to 1.0): The output scale factor. + + Returns: + `torch.FloatTensor`: The output of the last residual block, which is a tensor of shape `(batch_size, + in_channels, temporal_length, height, width)`. + + """ + + def __init__( + self, + in_channels: int, + num_layers: int = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + add_attention: bool = True, + attention_type: str = "3d", + attention_head_dim: int = 1, + output_scale_factor: float = 1.0, + ): + super().__init__() + + self.attention_type = attention_type + + norm_num_groups = norm_num_groups if norm_num_groups is not None else min(in_channels // 4, 32) + + self.convs = nn.ModuleList([ + ResidualBlock3D( + in_channels=in_channels, + out_channels=in_channels, + non_linearity=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + ) + ]) + + self.attentions = nn.ModuleList([]) + for _ in range(num_layers - 1): + if add_attention: + if attention_type == "3d": + self.attentions.append( + Attention3D( + in_channels, + nheads=in_channels // attention_head_dim, + head_dim=attention_head_dim, + bias=True, + upcast_softmax=True, + norm_num_groups=norm_num_groups, + eps=norm_eps, + rescale_output_factor=output_scale_factor, + residual_connection=True, + ) + ) + elif attention_type == "spatial_temporal": + self.attentions.append( + nn.ModuleList([ + SpatialAttention( + in_channels, + nheads=in_channels // attention_head_dim, + head_dim=attention_head_dim, + bias=True, + upcast_softmax=True, + norm_num_groups=norm_num_groups, + eps=norm_eps, + rescale_output_factor=output_scale_factor, + residual_connection=True, + ), + TemporalAttention( + in_channels, + nheads=in_channels // attention_head_dim, + head_dim=attention_head_dim, + bias=True, + upcast_softmax=True, + norm_num_groups=norm_num_groups, + eps=norm_eps, + rescale_output_factor=output_scale_factor, + residual_connection=True, + ), + ]) + ) + elif attention_type == "spatial": + self.attentions.append( + SpatialAttention( + in_channels, + nheads=in_channels // attention_head_dim, + head_dim=attention_head_dim, + bias=True, + upcast_softmax=True, + norm_num_groups=norm_num_groups, + eps=norm_eps, + rescale_output_factor=output_scale_factor, + residual_connection=True, + ) + ) + elif attention_type == "temporal": + self.attentions.append( + TemporalAttention( + in_channels, + nheads=in_channels // attention_head_dim, + head_dim=attention_head_dim, + bias=True, + upcast_softmax=True, + norm_num_groups=norm_num_groups, + eps=norm_eps, + rescale_output_factor=output_scale_factor, + residual_connection=True, + ) + ) + else: + raise ValueError(f"Unknown attention type: {attention_type}") + else: + self.attentions.append(None) + + self.convs.append( + ResidualBlock3D( + in_channels=in_channels, + out_channels=in_channels, + non_linearity=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + ) + ) + + def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor: + hidden_states = self.convs[0](hidden_states) + + for attn, resnet in zip(self.attentions, self.convs[1:]): + if attn is not None: + if self.attention_type == "spatial_temporal": + spatial_attn, temporal_attn = attn + hidden_states = spatial_attn(hidden_states) + hidden_states = temporal_attn(hidden_states) + else: + hidden_states = attn(hidden_states) + hidden_states = resnet(hidden_states) + + return hidden_states diff --git a/ruyi/vae/ldm/modules/vaemodules/up_blocks.py b/ruyi/vae/ldm/modules/vaemodules/up_blocks.py new file mode 100755 index 0000000..ee781cb --- /dev/null +++ b/ruyi/vae/ldm/modules/vaemodules/up_blocks.py @@ -0,0 +1,395 @@ +import torch +import torch.nn as nn + +from .attention import SpatialAttention, TemporalAttention +from .common import ResidualBlock3D +from .gc_block import GlobalContextBlock +from .upsamplers import (SpatialTemporalUpsampler3D, SpatialUpsampler3D, + TemporalUpsampler3D) + + +def get_up_block( + up_block_type: str, + in_channels: int, + out_channels: int, + num_layers: int, + act_fn: str, + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + num_attention_heads: int = 1, + output_scale_factor: float = 1.0, + add_gc_block: bool = False, + add_upsample: bool = True, +) -> nn.Module: + if up_block_type == "SpatialUpBlock3D": + return SpatialUpBlock3D( + in_channels=in_channels, + out_channels=out_channels, + num_layers=num_layers, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + add_gc_block=add_gc_block, + add_upsample=add_upsample, + ) + elif up_block_type == "SpatialAttnUpBlock3D": + return SpatialAttnUpBlock3D( + in_channels=in_channels, + out_channels=out_channels, + num_layers=num_layers, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + attention_head_dim=out_channels // num_attention_heads, + output_scale_factor=output_scale_factor, + add_gc_block=add_gc_block, + add_upsample=add_upsample, + ) + elif up_block_type == "TemporalUpBlock3D": + return TemporalUpBlock3D( + in_channels=in_channels, + out_channels=out_channels, + num_layers=num_layers, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + add_gc_block=add_gc_block, + add_upsample=add_upsample, + ) + elif up_block_type == "TemporalAttnUpBlock3D": + return TemporalAttnUpBlock3D( + in_channels=in_channels, + out_channels=out_channels, + num_layers=num_layers, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + attention_head_dim=out_channels // num_attention_heads, + output_scale_factor=output_scale_factor, + add_gc_block=add_gc_block, + add_upsample=add_upsample, + ) + elif up_block_type == "SpatialTemporalUpBlock3D": + return SpatialTemporalUpBlock3D( + in_channels=in_channels, + out_channels=out_channels, + num_layers=num_layers, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + add_gc_block=add_gc_block, + add_upsample=add_upsample, + ) + else: + raise ValueError(f"Unknown up block type: {up_block_type}") + + +class SpatialUpBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + output_scale_factor: float = 1.0, + add_gc_block: bool = False, + add_upsample: bool = True, + ): + super().__init__() + + if add_upsample: + self.upsampler = SpatialUpsampler3D(in_channels, in_channels) + else: + self.upsampler = None + + if add_gc_block: + self.gc_block = GlobalContextBlock(in_channels, in_channels, fusion_type="mul") + else: + self.gc_block = None + + self.convs = nn.ModuleList([]) + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.convs.append( + ResidualBlock3D( + in_channels=in_channels, + out_channels=out_channels, + non_linearity=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + ) + ) + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + for conv in self.convs: + x = conv(x) + + if self.gc_block is not None: + x = self.gc_block(x) + + if self.upsampler is not None: + x = self.upsampler(x) + + return x + + +class SpatialAttnUpBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + attention_head_dim: int = 1, + output_scale_factor: float = 1.0, + add_gc_block: bool = False, + add_upsample: bool = True, + ): + super().__init__() + + self.convs = nn.ModuleList([]) + self.attentions = nn.ModuleList([]) + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.convs.append( + ResidualBlock3D( + in_channels=in_channels, + out_channels=out_channels, + non_linearity=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + ) + ) + self.attentions.append( + SpatialAttention( + out_channels, + nheads=out_channels // attention_head_dim, + head_dim=attention_head_dim, + bias=True, + upcast_softmax=True, + norm_num_groups=norm_num_groups, + eps=norm_eps, + rescale_output_factor=output_scale_factor, + residual_connection=True, + ) + ) + + if add_gc_block: + self.gc_block = GlobalContextBlock(out_channels, out_channels, fusion_type="mul") + else: + self.gc_block = None + + if add_upsample: + self.upsampler = SpatialUpsampler3D(out_channels, out_channels) + else: + self.upsampler = None + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + for conv, attn in zip(self.convs, self.attentions): + x = conv(x) + x = attn(x) + + if self.gc_block is not None: + x = self.gc_block(x) + + if self.upsampler is not None: + x = self.upsampler(x) + + return x + + +class TemporalUpBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + output_scale_factor: float = 1.0, + add_gc_block: bool = False, + add_upsample: bool = True, + ): + super().__init__() + + self.convs = nn.ModuleList([]) + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.convs.append( + ResidualBlock3D( + in_channels=in_channels, + out_channels=out_channels, + non_linearity=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + ) + ) + + if add_gc_block: + self.gc_block = GlobalContextBlock(out_channels, out_channels, fusion_type="mul") + else: + self.gc_block = None + + if add_upsample: + self.upsampler = TemporalUpsampler3D(out_channels, out_channels) + else: + self.upsampler = None + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + for conv in self.convs: + x = conv(x) + + if self.gc_block is not None: + x = self.gc_block(x) + + if self.upsampler is not None: + x = self.upsampler(x) + + return x + + +class TemporalAttnUpBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + attention_head_dim: int = 1, + output_scale_factor: float = 1.0, + add_gc_block: bool = False, + add_upsample: bool = True, + ): + super().__init__() + + self.convs = nn.ModuleList([]) + self.attentions = nn.ModuleList([]) + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.convs.append( + ResidualBlock3D( + in_channels=in_channels, + out_channels=out_channels, + non_linearity=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + ) + ) + self.attentions.append( + TemporalAttention( + out_channels, + nheads=out_channels // attention_head_dim, + head_dim=attention_head_dim, + bias=True, + upcast_softmax=True, + norm_num_groups=norm_num_groups, + eps=norm_eps, + rescale_output_factor=output_scale_factor, + residual_connection=True, + ) + ) + + if add_gc_block: + self.gc_block = GlobalContextBlock(out_channels, out_channels, fusion_type="mul") + else: + self.gc_block = None + + if add_upsample: + self.upsampler = TemporalUpsampler3D(out_channels, out_channels) + else: + self.upsampler = None + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + for conv, attn in zip(self.convs, self.attentions): + x = conv(x) + x = attn(x) + + if self.gc_block is not None: + x = self.gc_block(x) + + if self.upsampler is not None: + x = self.upsampler(x) + + return x + + +class SpatialTemporalUpBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-6, + dropout: float = 0.0, + output_scale_factor: float = 1.0, + add_gc_block: bool = False, + add_upsample: bool = True, + ): + super().__init__() + + self.convs = nn.ModuleList([]) + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + self.convs.append( + ResidualBlock3D( + in_channels=in_channels, + out_channels=out_channels, + non_linearity=act_fn, + norm_num_groups=norm_num_groups, + norm_eps=norm_eps, + dropout=dropout, + output_scale_factor=output_scale_factor, + ) + ) + + if add_gc_block: + self.gc_block = GlobalContextBlock(out_channels, out_channels, fusion_type="mul") + else: + self.gc_block = None + + if add_upsample: + self.upsampler = SpatialTemporalUpsampler3D(out_channels, out_channels) + else: + self.upsampler = None + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + for conv in self.convs: + x = conv(x) + + if self.gc_block is not None: + x = self.gc_block(x) + + if self.upsampler is not None: + x = self.upsampler(x) + + return x diff --git a/ruyi/vae/ldm/modules/vaemodules/upsamplers.py b/ruyi/vae/ldm/modules/vaemodules/upsamplers.py new file mode 100644 index 0000000..d16b29c --- /dev/null +++ b/ruyi/vae/ldm/modules/vaemodules/upsamplers.py @@ -0,0 +1,183 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange, repeat + +from .common import CausalConv3d + + +class Upsampler(nn.Module): + def __init__( + self, + spatial_upsample_factor: int = 1, + temporal_upsample_factor: int = 1, + ): + super().__init__() + + self.spatial_upsample_factor = spatial_upsample_factor + self.temporal_upsample_factor = temporal_upsample_factor + + +class SpatialUpsampler3D(Upsampler): + def __init__(self, in_channels: int, out_channels: int): + super().__init__(spatial_upsample_factor=2) + + if out_channels is None: + out_channels = in_channels + + self.conv = CausalConv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.interpolate(x, scale_factor=(1, 2, 2), mode="nearest") + x = self.conv(x) + return x + + +class SpatialUpsamplerD2S3D(Upsampler): + def __init__(self, in_channels: int, out_channels: int): + super().__init__(spatial_upsample_factor=2) + + if out_channels is None: + out_channels = in_channels + + self.conv = CausalConv3d( + in_channels=in_channels, + out_channels=out_channels * 4, + kernel_size=3, + ) + + o, i, t, h, w = self.conv.weight.shape + conv_weight = torch.empty(o // 4, i, t, h, w) + nn.init.kaiming_normal_(conv_weight) + conv_weight = repeat(conv_weight, "o ... -> (o 4) ...") + self.conv.weight.data.copy_(conv_weight) + + nn.init.zeros_(self.conv.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.conv(x) + x = rearrange(x, "b (c p1 p2) t h w -> b c t (h p1) (w p2)", p1=2, p2=2) + return x + + +class TemporalUpsampler3D(Upsampler): + def __init__(self, in_channels: int, out_channels: int): + super().__init__( + spatial_upsample_factor=1, + temporal_upsample_factor=2, + ) + + if out_channels is None: + out_channels = in_channels + + self.conv = CausalConv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.shape[2] > 1: + first_frame, x = x[:, :, :1], x[:, :, 1:] + x = F.interpolate(x, scale_factor=(2, 1, 1), mode="trilinear") + x = torch.cat([first_frame, x], dim=2) + x = self.conv(x) + return x + + +class TemporalUpsamplerD2S3D(Upsampler): + def __init__(self, in_channels: int, out_channels: int): + super().__init__( + spatial_upsample_factor=1, + temporal_upsample_factor=2, + ) + + if out_channels is None: + out_channels = in_channels + + self.conv = CausalConv3d( + in_channels=in_channels, + out_channels=out_channels * 2, + kernel_size=3, + ) + + o, i, t, h, w = self.conv.weight.shape + conv_weight = torch.empty(o // 2, i, t, h, w) + nn.init.kaiming_normal_(conv_weight) + conv_weight = repeat(conv_weight, "o ... -> (o 2) ...") + self.conv.weight.data.copy_(conv_weight) + + nn.init.zeros_(self.conv.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.conv(x) + x = rearrange(x, "b (c p1) t h w -> b c (t p1) h w", p1=2) + x = x[:, :, 1:] + return x + + +class SpatialTemporalUpsampler3D(Upsampler): + def __init__(self, in_channels: int, out_channels: int): + super().__init__( + spatial_upsample_factor=2, + temporal_upsample_factor=2, + ) + + if out_channels is None: + out_channels = in_channels + + self.conv = CausalConv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + ) + + self.padding_flag = 0 + self.set_3dgroupnorm = False + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.interpolate(x, scale_factor=(1, 2, 2), mode="nearest") + x = self.conv(x) + + if self.padding_flag == 0: + if x.shape[2] > 1: + first_frame, x = x[:, :, :1], x[:, :, 1:] + x = F.interpolate(x, scale_factor=(2, 1, 1), mode="trilinear" if not self.set_3dgroupnorm else "nearest") + x = torch.cat([first_frame, x], dim=2) + elif self.padding_flag == 2 or self.padding_flag == 5 or self.padding_flag == 6: + x = F.interpolate(x, scale_factor=(2, 1, 1), mode="trilinear" if not self.set_3dgroupnorm else "nearest") + return x + +class SpatialTemporalUpsamplerD2S3D(Upsampler): + def __init__(self, in_channels: int, out_channels: int): + super().__init__( + spatial_upsample_factor=2, + temporal_upsample_factor=2, + ) + + if out_channels is None: + out_channels = in_channels + + self.conv = CausalConv3d( + in_channels=in_channels, + out_channels=out_channels * 8, + kernel_size=3, + ) + + o, i, t, h, w = self.conv.weight.shape + conv_weight = torch.empty(o // 8, i, t, h, w) + nn.init.kaiming_normal_(conv_weight) + conv_weight = repeat(conv_weight, "o ... -> (o 8) ...") + self.conv.weight.data.copy_(conv_weight) + + nn.init.zeros_(self.conv.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.conv(x) + x = rearrange(x, "b (c p1 p2 p3) t h w -> b c (t p1) (h p2) (w p3)", p1=2, p2=2, p3=2) + x = x[:, :, 1:] + return x diff --git a/ruyi/vae/ldm/util.py b/ruyi/vae/ldm/util.py new file mode 100644 index 0000000..4757117 --- /dev/null +++ b/ruyi/vae/ldm/util.py @@ -0,0 +1,201 @@ +import importlib +import multiprocessing as mp +from collections import abc +from functools import partial +from inspect import isfunction +from queue import Queue +from threading import Thread + +import numpy as np +import torch +from einops import rearrange +from PIL import Image, ImageDraw, ImageFont + + +def log_txt_as_img(wh, xc, size=10): + # wh a tuple of (width, height) + # xc a list of captions to plot + b = len(xc) + txts = list() + for bi in range(b): + txt = Image.new("RGB", wh, color="white") + draw = ImageDraw.Draw(txt) + font = ImageFont.truetype('data/DejaVuSans.ttf', size=size) + nc = int(40 * (wh[0] / 256)) + lines = "\n".join(xc[bi][start:start + nc] for start in range(0, len(xc[bi]), nc)) + + try: + draw.text((0, 0), lines, fill="black", font=font) + except UnicodeEncodeError: + print("Cant encode string for logging. Skipping.") + + txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.0 + txts.append(txt) + txts = np.stack(txts) + txts = torch.tensor(txts) + return txts + + +def ismap(x): + if not isinstance(x, torch.Tensor): + return False + return (len(x.shape) == 4) and (x.shape[1] > 3) + + +def isimage(x): + if not isinstance(x, torch.Tensor): + return False + return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1) + + +def exists(x): + return x is not None + + +def default(val, d): + if exists(val): + return val + return d() if isfunction(d) else d + + +def mean_flat(tensor): + """ + https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L86 + Take the mean over all non-batch dimensions. + """ + return tensor.mean(dim=list(range(1, len(tensor.shape)))) + + +def count_params(model, verbose=False): + total_params = sum(p.numel() for p in model.parameters()) + if verbose: + print(f"{model.__class__.__name__} has {total_params * 1.e-6:.2f} M params.") + return total_params + + +def instantiate_from_config(config): + if not "target" in config: + if config == '__is_first_stage__': + return None + elif config == "__is_unconditional__": + return None + raise KeyError("Expected key `target` to instantiate.") + return get_obj_from_str(config["target"])(**config.get("params", dict())) + + +def get_obj_from_str(string, reload=False): + module, cls = string.rsplit(".", 1) + if reload: + module_imp = importlib.import_module(module) + importlib.reload(module_imp) + return getattr(importlib.import_module(module, package=None), cls) + + +def _do_parallel_data_prefetch(func, Q, data, idx, idx_to_fn=False): + # create dummy dataset instance + + # run prefetching + if idx_to_fn: + res = func(data, worker_id=idx) + else: + res = func(data) + Q.put([idx, res]) + Q.put("Done") + + +def parallel_data_prefetch( + func: callable, data, n_proc, target_data_type="ndarray", cpu_intensive=True, use_worker_id=False +): + # if target_data_type not in ["ndarray", "list"]: + # raise ValueError( + # "Data, which is passed to parallel_data_prefetch has to be either of type list or ndarray." + # ) + if isinstance(data, np.ndarray) and target_data_type == "list": + raise ValueError("list expected but function got ndarray.") + elif isinstance(data, abc.Iterable): + if isinstance(data, dict): + print( + f'WARNING:"data" argument passed to parallel_data_prefetch is a dict: Using only its values and disregarding keys.' + ) + data = list(data.values()) + if target_data_type == "ndarray": + data = np.asarray(data) + else: + data = list(data) + else: + raise TypeError( + f"The data, that shall be processed parallel has to be either an np.ndarray or an Iterable, but is actually {type(data)}." + ) + + if cpu_intensive: + Q = mp.Queue(1000) + proc = mp.Process + else: + Q = Queue(1000) + proc = Thread + # spawn processes + if target_data_type == "ndarray": + arguments = [ + [func, Q, part, i, use_worker_id] + for i, part in enumerate(np.array_split(data, n_proc)) + ] + else: + step = ( + int(len(data) / n_proc + 1) + if len(data) % n_proc != 0 + else int(len(data) / n_proc) + ) + arguments = [ + [func, Q, part, i, use_worker_id] + for i, part in enumerate( + [data[i: i + step] for i in range(0, len(data), step)] + ) + ] + processes = [] + for i in range(n_proc): + p = proc(target=_do_parallel_data_prefetch, args=arguments[i]) + processes += [p] + + # start processes + print(f"Start prefetching...") + import time + + start = time.time() + gather_res = [[] for _ in range(n_proc)] + try: + for p in processes: + p.start() + + k = 0 + while k < n_proc: + # get result + res = Q.get() + if res == "Done": + k += 1 + else: + gather_res[res[0]] = res[1] + + except Exception as e: + print("Exception: ", e) + for p in processes: + p.terminate() + + raise e + finally: + for p in processes: + p.join() + print(f"Prefetching complete. [{time.time() - start} sec.]") + + if target_data_type == 'ndarray': + if not isinstance(gather_res[0], np.ndarray): + return np.concatenate([np.asarray(r) for r in gather_res], axis=0) + + # order outputs + return np.concatenate(gather_res, axis=0) + elif target_data_type == 'list': + out = [] + for r in gather_res: + out.extend(r) + return out + else: + return gather_res diff --git a/ruyi/vae/setup.py b/ruyi/vae/setup.py new file mode 100644 index 0000000..baaf21c --- /dev/null +++ b/ruyi/vae/setup.py @@ -0,0 +1,13 @@ +from setuptools import find_packages, setup + +setup( + name='latent-diffusion', + version='0.0.1', + description='', + packages=find_packages(), + install_requires=[ + 'torch', + 'numpy', + 'tqdm', + ], +) \ No newline at end of file