feat: Add LightX2V VAE custom nodes for ComfyUI
This commit is contained in:
gushiqiao
@@ -1 +1,239 @@
|
||||
# ComfyUI-LightVAE
|
||||
# ComfyUI-LightVAE
|
||||
|
||||
<div align="center">
|
||||
|
||||
|
||||
|
||||
**High-Performance VAE Custom Nodes**
|
||||
|
||||
[](https://huggingface.co/lightx2v/Autoencoders)
|
||||
[](https://github.com/ModelTC/LightX2V)
|
||||
[](LICENSE)
|
||||
|
||||
English | [简体中文](README_CN.md)
|
||||
|
||||
</div>
|
||||
|
||||
|
||||
## 📖 Introduction
|
||||
|
||||
ComfyUI-LightVAE is a collection of LightX2V VAE custom nodes designed for ComfyUI, supporting high-performance video VAE models including **LightVAE** and **LightTAE**.
|
||||
|
||||
The LightX2V team has deeply optimized VAE, creating two major series: **LightVAE** and **LightTAE**, which significantly reduce memory usage and improve inference speed while maintaining high quality.
|
||||
|
||||
|
||||
## ✨ Key Features
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td width="50%">
|
||||
|
||||
### 🎯 LightVAE Series
|
||||
**Feature: Best Balance ⚖️**
|
||||
|
||||
- ✅ Uses **Causal 3D Conv** (same as official)
|
||||
- ✅ **Near-official quality** ⭐⭐⭐⭐
|
||||
- ✅ **~50% less memory** (~4-5 GB)
|
||||
- ✅ **2-3x faster**
|
||||
- ✅ Balances quality, speed, and memory 🏆
|
||||
|
||||
</td>
|
||||
<td width="50%">
|
||||
|
||||
### ⚡ LightTAE Series
|
||||
**Feature: Ultra-fast + High Quality 🏆**
|
||||
|
||||
- ✅ Minimal memory usage (~0.4 GB)
|
||||
- ✅ Lightning-fast inference
|
||||
- ✅ **Near-official quality** ⭐⭐⭐⭐
|
||||
- ✅ **Surpasses open-source TAE**
|
||||
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
|
||||
## 🚀 Performance Comparison
|
||||
|
||||
**Test Environment**: H100 GPU, BF16, 81-frame video (480P)
|
||||
|
||||
| Model | Encode Time | Decode Time | Encode Memory | Decode Memory | Quality |
|
||||
|:------|:------------|:------------|:--------------|:--------------|:--------|
|
||||
| **lightvaew2_1** | 1.5s | 2.1s | 4.8GB | 5.6GB | ⭐⭐⭐⭐⭐ |
|
||||
| **lighttaew2_1** | 0.4s | 0.25s | 0.009GB | 0.4GB | ⭐⭐⭐⭐ |
|
||||
| [Wan2.1_VAE](https://github.com/Wan-Video/Wan2.1) | 4.2s | 5.5s | 8.5GB | 10.1GB | ⭐⭐⭐⭐ |
|
||||
| [taew2_1](https://github.com/madebyollin/taesd) | 0.4s | 0.25s | 0.009GB | 0.4GB | ⭐⭐⭐ |
|
||||
|
||||
**Performance Improvements:**
|
||||
- 🚀 LightVAE is **2-3x faster** than official VAE, **50% less memory**
|
||||
- ⚡ LightTAE is **10+ times faster** than official VAE, **95%+ less memory**
|
||||
- 🎨 Near-official [VAE](https://github.com/Wan-Video/Wan2.1) quality, surpasses open-source [TAE](https://github.com/madebyollin/taesd)
|
||||
|
||||
|
||||
## 📦 Installation
|
||||
|
||||
### 1. Install LightX2V Dependencies
|
||||
|
||||
```bash
|
||||
# Clone LightX2V repository
|
||||
git clone https://github.com/ModelTC/LightX2V
|
||||
cd LightX2V
|
||||
|
||||
python setup_vae.py install
|
||||
```
|
||||
|
||||
### 2. Install ComfyUI-WanVideoWrapper
|
||||
|
||||
LightVAE nodes depend on WanVideoWrapper for main model support:
|
||||
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes
|
||||
git clone https://github.com/kijai/ComfyUI-WanVideoWrapper
|
||||
```
|
||||
|
||||
### 3. Install ComfyUI-LightVAE
|
||||
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes
|
||||
git clone https://github.com/YOUR_USERNAME/ComfyUI-LightVAE
|
||||
```
|
||||
|
||||
### 4. Restart ComfyUI
|
||||
|
||||
|
||||
## 📥 Download Models
|
||||
|
||||
### Main Models (Diffusion Models)
|
||||
|
||||
**Option 1: Distilled Models (Recommended, 4-step)**
|
||||
- 🔗 [Wan2.1 Distilled Models](https://huggingface.co/lightx2v/Wan2.1-Distill-Models) and [Wan2.2 Distilled Models](https://huggingface.co/lightx2v/Wan2.2-Distill-Models)
|
||||
- ✅ Supports BF16 format
|
||||
- ✅ Supports FP8 format (requires models with `_comfyui.safetensors` suffix)
|
||||
|
||||
**Option 2: Original Models (20-step)**
|
||||
- 🔗 [Wan2.1 Official Models](https://huggingface.co/lightx2v/Wan2.1-Official-Models) and [Wan2.2 Official Models](https://huggingface.co/lightx2v/Wan2.2-Official-Models)
|
||||
- ✅ Supports BF16 format
|
||||
- ✅ Supports FP8 format (requires models with `_comfyui.safetensors` suffix)
|
||||
|
||||
```bash
|
||||
# Download to ComfyUI/models/diffusion_models/
|
||||
huggingface-cli download lightx2v/Wan2.1/2-Distill-Models \
|
||||
--local-dir ./ComfyUI/models/diffusion_models/
|
||||
```
|
||||
|
||||
### VAE Models
|
||||
|
||||
**All VAE Models** (Required):
|
||||
- 🔗 [LightX2V Autoencoders](https://huggingface.co/lightx2v/Autoencoders)
|
||||
|
||||
```bash
|
||||
# Download all VAE models
|
||||
huggingface-cli download lightx2v/Autoencoders \
|
||||
--local-dir ./ComfyUI/models/vae/
|
||||
|
||||
# Or download only what you need (Recommended)
|
||||
huggingface-cli download lightx2v/Autoencoders lightvaew2_1.pth \
|
||||
--local-dir ./ComfyUI/models/vae/
|
||||
```
|
||||
|
||||
**Supported VAE Models**:
|
||||
- `Wan2.1_VAE.pth` / `.safetensors` - Official VAE 2.1
|
||||
- `Wan2.2_VAE.pth` / `.safetensors` - Official VAE 2.2
|
||||
- `lightvaew2_1.pth` / `.safetensors` - Optimized VAE 2.1 ⭐ **Recommended**
|
||||
- `taew2_1.pth` / `.safetensors` - Open-source TAE 2.1
|
||||
- `taew2_2.pth` / `.safetensors` - Open-source TAE 2.2
|
||||
- `lighttaew2_1.pth` / `.safetensors` - Optimized TAE 2.1 ⚡ **Fastest**
|
||||
- `lighttaew2_2.pth` / `.safetensors` - Optimized TAE 2.2
|
||||
|
||||
|
||||
## 🎯 Node Documentation
|
||||
|
||||
### 1. LightX2V VAE Decoder Loader
|
||||
|
||||
Loads VAE decoder models.
|
||||
|
||||
**Input Parameters:**
|
||||
- `vae_filename` - VAE model filename (automatically lists from `./models/vae/`)
|
||||
- `dtype` - Data type (bfloat16 / float16 / float32)
|
||||
- `device` - Compute device (cuda / cpu)
|
||||
|
||||
**Output:**
|
||||
- `vae` - VAE model object
|
||||
|
||||
**Features:**
|
||||
- ✅ Automatically identifies VAE type from filename
|
||||
- ✅ Supports all LightX2V VAE models
|
||||
|
||||
|
||||
|
||||
### 2. LightX2V VAE Decode
|
||||
|
||||
Decodes latent to video frames.
|
||||
|
||||
**Input Parameters:**
|
||||
- `vae` - VAE object from Loader
|
||||
- `latent` - Latent representation
|
||||
|
||||
**Output:**
|
||||
- `IMAGE` - Decoded video frames
|
||||
|
||||
**Supports:**
|
||||
- ✅ All VAE series (WanVAE, LightVAE)
|
||||
- ✅ All TAE series (TAE, LightTAE)
|
||||
|
||||
|
||||
|
||||
## 🖼️ Example Workflows
|
||||
|
||||
### Wan2.1 I2V 4-step FP8 + LightVAE
|
||||
|
||||
High-performance configuration using 4-step distilled model + LightVAE optimized decoder.
|
||||
|
||||
|
||||
|
||||
**Workflow File**: [`example/workflows/wan2.1_I2V_4step_fp8_lightvae.json`](example/workflows/wan2.1_I2V_4step_fp8_lightvae.json)
|
||||
|
||||
### Wan2.2 TI2V + LightVAE
|
||||
|
||||
Wan2.2 Text-Image-to-Video + LightVAE decoding.
|
||||
|
||||
**Workflow File**: [`example/workflows/wan2.2_TI2V_lightvae.json`](example/workflows/wan2.2_TI2V_lightvae.json)
|
||||
|
||||
|
||||
|
||||
## ⚠️ Important Notes
|
||||
|
||||
### Model Compatibility
|
||||
- ⚠️ **Wan2.1 VAE** can only be used with **Wan2.1/Wan2.2-A1B** backbone models
|
||||
- ⚠️ **Wan2.2 VAE** can only be used with **Wan2.2 TI2V** backbone models
|
||||
- ❌ Do not mix different versions of VAE and backbone models
|
||||
|
||||
|
||||
## 📚 Related Resources
|
||||
|
||||
- **Project Homepage**: https://github.com/ModelTC/LightX2V
|
||||
- **VAE Models**: https://huggingface.co/lightx2v/Autoencoders
|
||||
- **Video Generation Models**: https://huggingface.co/lightx2v/
|
||||
- **ComfyUI-WanVideoWrapper**: https://github.com/kijai/ComfyUI-WanVideoWrapper
|
||||
- **TAE Series Models**: https://github.com/madebyollin/taesd
|
||||
- **Wan-AI**: https://huggingface.co/Wan-AI
|
||||
|
||||
|
||||
## 🙏 Acknowledgements
|
||||
|
||||
If this project helps you, please give a ⭐ to [LightX2V](https://github.com/ModelTC/LightX2V) and this repository!
|
||||
|
||||
|
||||
|
||||
## 📞 Support
|
||||
|
||||
- **GitHub Issues**: Issues page of this repository
|
||||
- **LightX2V Issues**: https://github.com/ModelTC/LightX2V/issues
|
||||
- **HuggingFace**: https://huggingface.co/lightx2v
|
||||
|
||||
<div align="center">
|
||||
|
||||
**Enjoy using LightX2V VAE!** 🚀
|
||||
|
||||
</div>
|
||||
|
||||
|
||||
+239
@@ -0,0 +1,239 @@
|
||||
# ComfyUI-LightVAE
|
||||
|
||||
<div align="center">
|
||||
|
||||
|
||||
|
||||
**高性能 VAE 自定义节点**
|
||||
|
||||
[](https://huggingface.co/lightx2v/Autoencoders)
|
||||
[](https://github.com/ModelTC/LightX2V)
|
||||
[](LICENSE)
|
||||
|
||||
[English](README.md) | 简体中文
|
||||
|
||||
</div>
|
||||
|
||||
|
||||
## 📖 简介
|
||||
|
||||
ComfyUI-LightVAE 是为 ComfyUI 设计的 LightX2V VAE 自定义节点集合，支持包括 **LightVAE** 和 **LightTAE** 在内的高性能视频 VAE 模型。
|
||||
|
||||
LightX2V 团队对 VAE 进行了深度优化，衍生出两大系列：**LightVAE** 和 **LightTAE**，在保持高质量的同时显著降低内存占用并提升推理速度。
|
||||
|
||||
|
||||
## ✨ 核心特性
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td width="50%">
|
||||
|
||||
### 🎯 LightVAE 系列
|
||||
**特点：最佳平衡 ⚖️**
|
||||
|
||||
- ✅ 使用 **Causal 3D Conv**（与官方相同）
|
||||
- ✅ **质量接近官方** ⭐⭐⭐⭐
|
||||
- ✅ 内存减少 **~50%**（~4-5 GB）
|
||||
- ✅ 速度提升 **2-3x**
|
||||
- ✅ 平衡质量、速度和内存 🏆
|
||||
|
||||
</td>
|
||||
<td width="50%">
|
||||
|
||||
### ⚡ LightTAE 系列
|
||||
**特点：极速 + 高质量 🏆**
|
||||
|
||||
- ✅ 极低内存占用（~0.4 GB）
|
||||
- ✅ 极速推理
|
||||
- ✅ **质量接近官方** ⭐⭐⭐⭐
|
||||
- ✅ **超越开源 TAE**
|
||||
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
|
||||
## 🚀 性能对比
|
||||
|
||||
**测试环境**: H100 GPU, BF16, 81 帧视频 (480P)
|
||||
|
||||
| 模型 | 编码时间 | 解码时间 | 编码内存 | 解码内存 | 质量 |
|
||||
|:-----|:---------|:---------|:---------|:---------|:-----|
|
||||
| **lightvaew2_1** | 1.5s | 2.1s | 4.8GB | 5.6GB | ⭐⭐⭐⭐⭐ |
|
||||
| **lighttaew2_1** | 0.4s | 0.25s | 0.009GB | 0.4GB | ⭐⭐⭐⭐ |
|
||||
| [Wan2.1_VAE](https://github.com/Wan-Video/Wan2.1) | 4.2s | 5.5s | 8.5GB | 10.1GB | ⭐⭐⭐⭐ |
|
||||
| [taew2_1](https://github.com/madebyollin/taesd) | 0.4s | 0.25s | 0.009GB | 0.4GB | ⭐⭐⭐ |
|
||||
|
||||
**性能提升：**
|
||||
- 🚀 LightVAE 比官方 VAE **快 2-3 倍**，内存**节省 50%**
|
||||
- ⚡ LightTAE 比官方 VAE **快 10+ 倍**，内存**节省 95%+**
|
||||
- 🎨 质量接近官方 [VAE](https://github.com/Wan-Video/Wan2.1)，超过开源 [TAE](https://github.com/madebyollin/taesd)
|
||||
|
||||
|
||||
## 📦 安装
|
||||
|
||||
### 1. 安装 LightX2V 依赖
|
||||
|
||||
```bash
|
||||
# 克隆 LightX2V 仓库
|
||||
git clone https://github.com/ModelTC/LightX2V
|
||||
cd LightX2V
|
||||
|
||||
python setup_vae.py install
|
||||
```
|
||||
|
||||
### 2. 安装 ComfyUI-WanVideoWrapper
|
||||
|
||||
LightVAE 节点依赖 WanVideoWrapper 提供的主模型支持：
|
||||
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes
|
||||
git clone https://github.com/kijai/ComfyUI-WanVideoWrapper
|
||||
```
|
||||
|
||||
### 3. 安装 ComfyUI-LightVAE
|
||||
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes
|
||||
git clone https://github.com/YOUR_USERNAME/ComfyUI-LightVAE
|
||||
```
|
||||
|
||||
### 4. 重启 ComfyUI
|
||||
|
||||
|
||||
## 📥 下载模型
|
||||
|
||||
### 主模型（Diffusion Models）
|
||||
|
||||
**选项 1: 蒸馏模型（推荐，4-step）**
|
||||
- 🔗 [Wan2.1蒸馏模型](https://huggingface.co/lightx2v/Wan2.1-Distill-Models) 和 [Wan2.2蒸馏模型](https://huggingface.co/lightx2v/Wan2.2-Distill-Models)
|
||||
- ✅ 支持 BF16 格式
|
||||
- ✅ 支持 FP8 格式（需使用 `_comfyui.safetensors` 后缀的模型）
|
||||
|
||||
**选项 2: 原始模型（20-step）**
|
||||
- 🔗 [Wan2.1官方模型](https://huggingface.co/lightx2v/Wan2.1-Official-Models) 和 [Wan2.2官方模型](https://huggingface.co/lightx2v/Wan2.2-Official-Models)
|
||||
- ✅ 支持 BF16 格式
|
||||
- ✅ 支持 FP8 格式（需使用 `_comfyui.safetensors` 后缀的模型）
|
||||
|
||||
```bash
|
||||
# 下载到 ComfyUI/models/diffusion_models/
|
||||
huggingface-cli download lightx2v/Wan2.1/2-Distill-Models \
|
||||
--local-dir ./ComfyUI/models/diffusion_models/
|
||||
```
|
||||
|
||||
### VAE 模型
|
||||
|
||||
**所有 VAE 模型**（必需）:
|
||||
- 🔗 [LightX2V Autoencoders](https://huggingface.co/lightx2v/Autoencoders)
|
||||
|
||||
```bash
|
||||
# 下载所有 VAE 模型
|
||||
huggingface-cli download lightx2v/Autoencoders \
|
||||
--local-dir ./ComfyUI/models/vae/
|
||||
|
||||
# 或只下载需要的（推荐）
|
||||
huggingface-cli download lightx2v/Autoencoders lightvaew2_1.pth \
|
||||
--local-dir ./ComfyUI/models/vae/
|
||||
```
|
||||
|
||||
**支持的 VAE 模型**:
|
||||
- `Wan2.1_VAE.pth` / `.safetensors` - 官方 VAE 2.1
|
||||
- `Wan2.2_VAE.pth` / `.safetensors` - 官方 VAE 2.2
|
||||
- `lightvaew2_1.pth` / `.safetensors` - 优化 VAE 2.1 ⭐ **推荐**
|
||||
- `taew2_1.pth` / `.safetensors` - 开源 TAE 2.1
|
||||
- `taew2_2.pth` / `.safetensors` - 开源 TAE 2.2
|
||||
- `lighttaew2_1.pth` / `.safetensors` - 优化 TAE 2.1 ⚡ **最快**
|
||||
- `lighttaew2_2.pth` / `.safetensors` - 优化 TAE 2.2
|
||||
|
||||
|
||||
## 🎯 节点说明
|
||||
|
||||
### 1. LightX2V VAE Decoder Loader
|
||||
|
||||
加载 VAE 解码器模型。
|
||||
|
||||
**输入参数:**
|
||||
- `vae_filename` - VAE 模型文件名（从 `./models/vae/` 目录自动列出）
|
||||
- `dtype` - 数据类型（bfloat16 / float16 / float32）
|
||||
- `device` - 计算设备（cuda / cpu）
|
||||
|
||||
**输出:**
|
||||
- `vae` - VAE 模型对象
|
||||
|
||||
**特点:**
|
||||
- ✅ 自动从文件名识别 VAE 类型
|
||||
- ✅ 支持所有 LightX2V VAE 模型
|
||||
|
||||
|
||||
|
||||
### 2. LightX2V VAE Decode
|
||||
|
||||
解码 latent 到视频帧。
|
||||
|
||||
**输入参数:**
|
||||
- `vae` - 来自 Loader 的 VAE 对象
|
||||
- `latent` - 潜在表示
|
||||
|
||||
**输出:**
|
||||
- `IMAGE` - 解码后的视频帧
|
||||
|
||||
**支持:**
|
||||
- ✅ 所有 VAE 系列（WanVAE, LightVAE）
|
||||
- ✅ 所有 TAE 系列（TAE, LightTAE）
|
||||
|
||||
|
||||
|
||||
## 🖼️ 示例工作流
|
||||
|
||||
### Wan2.1 I2V 4-step FP8 + LightVAE
|
||||
|
||||
使用 4-step 蒸馏模型 + LightVAE 优化解码器的高性能配置。
|
||||
|
||||
|
||||
|
||||
**工作流文件**: [`example/workflows/wan2.1_I2V_4step_fp8_lightvae.json`](example/workflows/wan2.1_I2V_4step_fp8_lightvae.json)
|
||||
|
||||
### Wan2.2 TI2V + LightVAE
|
||||
|
||||
Wan2.2 文本图像到视频 + LightVAE 解码。
|
||||
|
||||
**工作流文件**: [`example/workflows/wan2.2_TI2V_lightvae.json`](example/workflows/wan2.2_TI2V_lightvae.json)
|
||||
|
||||
|
||||
|
||||
## ⚠️ 注意事项
|
||||
|
||||
### 模型兼容性
|
||||
- ⚠️ **Wan2.1 VAE** 只能用于 **Wan2.1/Wan2.2-A1B** 主干模型
|
||||
- ⚠️ **Wan2.2 VAE** 只能用于 **Wan2.2 TI2V** 主干模型
|
||||
- ❌ 不要混用不同版本的 VAE 和主干模型
|
||||
|
||||
|
||||
## 📚 相关资源
|
||||
|
||||
- **项目主页**: https://github.com/ModelTC/LightX2V
|
||||
- **VAE 模型**: https://huggingface.co/lightx2v/Autoencoders
|
||||
- **视频生成模型**: https://huggingface.co/lightx2v/
|
||||
- **ComfyUI-WanVideoWrapper**: https://github.com/kijai/ComfyUI-WanVideoWrapper
|
||||
- **tae系列模型**： https://github.com/madebyollin/taesd
|
||||
- **Wan-AI**： https://huggingface.co/Wan-AI
|
||||
|
||||
|
||||
## 🙏 致谢
|
||||
|
||||
如果这个项目对你有帮助，请给 ⭐ [LightX2V](https://github.com/ModelTC/LightX2V) 和本仓库！
|
||||
|
||||
|
||||
|
||||
## 📞 支持
|
||||
|
||||
- **GitHub Issues**: 本仓库的 Issues 页面
|
||||
- **LightX2V Issues**: https://github.com/ModelTC/LightX2V/issues
|
||||
- **HuggingFace**: https://huggingface.co/lightx2v
|
||||
|
||||
<div align="center">
|
||||
|
||||
**享受使用 LightX2V VAE！** 🚀
|
||||
|
||||
</div>
|
||||
|
||||
+16
@@ -0,0 +1,16 @@
|
||||
from .nodes import (
|
||||
VAEDecoderLoader,
|
||||
LightVAEDecode
|
||||
)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Lightx2v-VAEDecoderLoader": VAEDecoderLoader,
|
||||
"Lightx2v-VAEDecode": LightVAEDecode,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Lightx2v-VAEDecoderLoader": "LightX2V VAE Decoder Loader",
|
||||
"Lightx2v-VAEDecode": "LightX2V VAE Decode",
|
||||
}
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 712 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 17 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 37 KiB |
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,143 @@
|
||||
"""
|
||||
LightX2V VAE Nodes for ComfyUI
|
||||
VAE Loader + 各类 VAE 的 Decode 节点
|
||||
"""
|
||||
|
||||
import os
|
||||
import torch
|
||||
import folder_paths
|
||||
from lightx2v.models.video_encoders.hf.wan.vae import WanVAE
|
||||
from lightx2v.models.video_encoders.hf.wan.vae_2_2 import Wan2_2_VAE
|
||||
from lightx2v.models.video_encoders.hf.wan.vae_tiny import WanVAE_tiny, Wan2_2_VAE_tiny
|
||||
|
||||
class VAEDecoderLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
vae_default_dir = "./models/vae/"
|
||||
return {
|
||||
"required": {
|
||||
"vae_filename": (os.listdir(vae_default_dir) ,{"default": "Wan2.1_VAE.safetensors"}),
|
||||
"dtype": (["bfloat16", "float16", "float32"], {"default": "bfloat16"}),
|
||||
"device": (["cuda", "cpu"], {"default": "cuda"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VAE",)
|
||||
RETURN_NAMES = ("vae",)
|
||||
FUNCTION = "load_vae"
|
||||
CATEGORY = "LightX2V"
|
||||
|
||||
def load_vae(self, vae_filename: str, dtype: str, device: str):
|
||||
dtype_map = {
|
||||
"bfloat16": torch.bfloat16,
|
||||
"float16": torch.float16,
|
||||
"float32": torch.float32
|
||||
}
|
||||
vae_name = os.path.splitext(vae_filename)[0]
|
||||
vae_default_dir = "./models/vae/"
|
||||
torch_dtype = dtype_map[dtype]
|
||||
vae_path = os.path.join(vae_default_dir, vae_filename)
|
||||
|
||||
if not os.path.exists(vae_path):
|
||||
raise FileNotFoundError(f"VAE model not found: {vae_path}")
|
||||
|
||||
if vae_name == "Wan2.1_VAE":
|
||||
model = WanVAE(
|
||||
vae_path=vae_path,
|
||||
dtype=torch_dtype,
|
||||
device=device,
|
||||
use_lightvae=False
|
||||
)
|
||||
|
||||
elif vae_name == "Wan2.2_VAE":
|
||||
model = Wan2_2_VAE(
|
||||
vae_path=vae_path,
|
||||
dtype=torch_dtype,
|
||||
device=device
|
||||
)
|
||||
|
||||
elif vae_name == "lightvaew2_1":
|
||||
model = WanVAE(
|
||||
vae_path=vae_path,
|
||||
dtype=torch_dtype,
|
||||
device=device,
|
||||
use_lightvae=True
|
||||
)
|
||||
|
||||
elif vae_name == "taew2_1":
|
||||
model = WanVAE_tiny(
|
||||
vae_path=vae_path,
|
||||
dtype=torch_dtype,
|
||||
device=device,
|
||||
need_scaled=False
|
||||
).to(device)
|
||||
|
||||
elif vae_name == "taew2_2":
|
||||
model = Wan2_2_VAE_tiny(
|
||||
vae_path=vae_path,
|
||||
dtype=torch_dtype,
|
||||
device=device,
|
||||
need_scaled=False
|
||||
).to(device)
|
||||
|
||||
elif vae_name == "lighttaew2_1":
|
||||
model = WanVAE_tiny(
|
||||
vae_path=vae_path,
|
||||
dtype=torch_dtype,
|
||||
device=device,
|
||||
need_scaled=True
|
||||
).to(device)
|
||||
|
||||
elif vae_name == "lighttaew2_2":
|
||||
model = Wan2_2_VAE_tiny(
|
||||
vae_path=vae_path,
|
||||
dtype=torch_dtype,
|
||||
device=device,
|
||||
need_scaled=True
|
||||
).to(device)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown vae_name: {vae_name}")
|
||||
|
||||
return ({
|
||||
"model": model,
|
||||
"vae_name": vae_name,
|
||||
"dtype": torch_dtype,
|
||||
"device": device,
|
||||
},)
|
||||
|
||||
|
||||
|
||||
class LightVAEDecode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("VAE",),
|
||||
"latent": ("LATENT",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "decode"
|
||||
CATEGORY = "LightX2V"
|
||||
|
||||
@torch.no_grad()
|
||||
def decode(self, vae: dict, latent: dict):
|
||||
model = vae["model"]
|
||||
dtype = vae["dtype"]
|
||||
device = vae["device"]
|
||||
|
||||
latent_samples = latent["samples"].to(device, dtype)
|
||||
with torch.no_grad():
|
||||
video = model.decode(latent_samples.squeeze(0)).half().permute(0, 2, 3, 4, 1)
|
||||
|
||||
if len(video.shape) == 5: #Combine batches
|
||||
video = video.reshape(-1, video.shape[-3], video.shape[-2], video.shape[-1])
|
||||
return (video.add_(1.0).div_(2.0),)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"VAEDecoderLoader",
|
||||
"LightVAEDecode"
|
||||
]
|
||||
Reference in New Issue
Block a user