feat: init
This commit is contained in:
GACLove
+308
@@ -0,0 +1,308 @@
|
||||
# Created by https://www.toptal.com/developers/gitignore/api/python,visualstudiocode,pycharm
|
||||
# Edit at https://www.toptal.com/developers/gitignore?templates=python,visualstudiocode,pycharm
|
||||
|
||||
### PyCharm ###
|
||||
# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider
|
||||
# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
|
||||
|
||||
# User-specific stuff
|
||||
.idea/**/workspace.xml
|
||||
.idea/**/tasks.xml
|
||||
.idea/**/usage.statistics.xml
|
||||
.idea/**/dictionaries
|
||||
.idea/**/shelf
|
||||
|
||||
# AWS User-specific
|
||||
.idea/**/aws.xml
|
||||
|
||||
# Generated files
|
||||
.idea/**/contentModel.xml
|
||||
|
||||
# Sensitive or high-churn files
|
||||
.idea/**/dataSources/
|
||||
.idea/**/dataSources.ids
|
||||
.idea/**/dataSources.local.xml
|
||||
.idea/**/sqlDataSources.xml
|
||||
.idea/**/dynamic.xml
|
||||
.idea/**/uiDesigner.xml
|
||||
.idea/**/dbnavigator.xml
|
||||
|
||||
# Gradle
|
||||
.idea/**/gradle.xml
|
||||
.idea/**/libraries
|
||||
|
||||
# Gradle and Maven with auto-import
|
||||
# When using Gradle or Maven with auto-import, you should exclude module files,
|
||||
# since they will be recreated, and may cause churn. Uncomment if using
|
||||
# auto-import.
|
||||
# .idea/artifacts
|
||||
# .idea/compiler.xml
|
||||
# .idea/jarRepositories.xml
|
||||
# .idea/modules.xml
|
||||
# .idea/*.iml
|
||||
# .idea/modules
|
||||
# *.iml
|
||||
# *.ipr
|
||||
|
||||
# CMake
|
||||
cmake-build-*/
|
||||
|
||||
# Mongo Explorer plugin
|
||||
.idea/**/mongoSettings.xml
|
||||
|
||||
# File-based project format
|
||||
*.iws
|
||||
|
||||
# IntelliJ
|
||||
out/
|
||||
|
||||
# mpeltonen/sbt-idea plugin
|
||||
.idea_modules/
|
||||
|
||||
# JIRA plugin
|
||||
atlassian-ide-plugin.xml
|
||||
|
||||
# Cursive Clojure plugin
|
||||
.idea/replstate.xml
|
||||
|
||||
# SonarLint plugin
|
||||
.idea/sonarlint/
|
||||
|
||||
# Crashlytics plugin (for Android Studio and IntelliJ)
|
||||
com_crashlytics_export_strings.xml
|
||||
crashlytics.properties
|
||||
crashlytics-build.properties
|
||||
fabric.properties
|
||||
|
||||
# Editor-based Rest Client
|
||||
.idea/httpRequests
|
||||
|
||||
# Android studio 3.1+ serialized cache file
|
||||
.idea/caches/build_file_checksums.ser
|
||||
|
||||
### PyCharm Patch ###
|
||||
# Comment Reason: https://github.com/joeblau/gitignore.io/issues/186#issuecomment-215987721
|
||||
|
||||
# *.iml
|
||||
# modules.xml
|
||||
# .idea/misc.xml
|
||||
# *.ipr
|
||||
|
||||
# Sonarlint plugin
|
||||
# https://plugins.jetbrains.com/plugin/7973-sonarlint
|
||||
.idea/**/sonarlint/
|
||||
|
||||
# SonarQube Plugin
|
||||
# https://plugins.jetbrains.com/plugin/7238-sonarqube-community-plugin
|
||||
.idea/**/sonarIssues.xml
|
||||
|
||||
# Markdown Navigator plugin
|
||||
# https://plugins.jetbrains.com/plugin/7896-markdown-navigator-enhanced
|
||||
.idea/**/markdown-navigator.xml
|
||||
.idea/**/markdown-navigator-enh.xml
|
||||
.idea/**/markdown-navigator/
|
||||
|
||||
# Cache file creation bug
|
||||
# See https://youtrack.jetbrains.com/issue/JBR-2257
|
||||
.idea/$CACHE_FILE$
|
||||
|
||||
# CodeStream plugin
|
||||
# https://plugins.jetbrains.com/plugin/12206-codestream
|
||||
.idea/codestream.xml
|
||||
|
||||
# Azure Toolkit for IntelliJ plugin
|
||||
# https://plugins.jetbrains.com/plugin/8053-azure-toolkit-for-intellij
|
||||
.idea/**/azureSettings.xml
|
||||
|
||||
### Python ###
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
cover/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
#poetry.lock
|
||||
|
||||
# pdm
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||
#pdm.lock
|
||||
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
||||
# in version control.
|
||||
# https://pdm.fming.dev/#use-with-ide
|
||||
.pdm.toml
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
#.idea/
|
||||
|
||||
### Python Patch ###
|
||||
# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
|
||||
poetry.toml
|
||||
|
||||
# ruff
|
||||
.ruff_cache/
|
||||
|
||||
# LSP config files
|
||||
pyrightconfig.json
|
||||
|
||||
### VisualStudioCode ###
|
||||
.vscode/*
|
||||
# !.vscode/settings.json
|
||||
!.vscode/tasks.json
|
||||
!.vscode/launch.json
|
||||
!.vscode/extensions.json
|
||||
!.vscode/*.code-snippets
|
||||
|
||||
# Local History for Visual Studio Code
|
||||
.history/
|
||||
|
||||
# Built Visual Studio Code Extensions
|
||||
*.vsix
|
||||
|
||||
### VisualStudioCode Patch ###
|
||||
# Ignore all local history of files
|
||||
.history
|
||||
.ionide
|
||||
|
||||
# End of https://www.toptal.com/developers/gitignore/api/python,visualstudiocode,pycharm
|
||||
@@ -0,0 +1,3 @@
|
||||
[submodule "lightx2v"]
|
||||
path = lightx2v
|
||||
url = https://github.com/ModelTC/lightx2v.git
|
||||
@@ -0,0 +1,3 @@
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
Submodule
+1
Submodule lightx2v added at 863b22c658
@@ -0,0 +1,641 @@
|
||||
import os
|
||||
import torch
|
||||
import gc
|
||||
from typing import cast, no_type_check, Any
|
||||
import logging
|
||||
import json
|
||||
import numpy as np
|
||||
import comfy.model_management as mm
|
||||
from comfy.utils import ProgressBar
|
||||
import folder_paths
|
||||
from tqdm import tqdm
|
||||
from easydict import EasyDict
|
||||
|
||||
|
||||
# Import LightX2V modules
|
||||
from .lightx2v.lightx2v.utils.profiler import ProfilingContext
|
||||
from .lightx2v.lightx2v.models.input_encoders.hf.t5.model import T5EncoderModel
|
||||
from .lightx2v.lightx2v.models.input_encoders.hf.xlm_roberta.model import (
|
||||
CLIPModel as ClipVisionModel,
|
||||
)
|
||||
from .lightx2v.lightx2v.utils.utils import seed_all
|
||||
from .lightx2v.lightx2v.models.video_encoders.hf.wan.vae import WanVAE
|
||||
from .lightx2v.lightx2v.models.networks.wan.model import WanModel
|
||||
from .lightx2v.lightx2v.models.networks.wan.lora_adapter import WanLoraWrapper
|
||||
from .lightx2v.lightx2v.models.schedulers.wan.scheduler import WanScheduler
|
||||
from .lightx2v.lightx2v.models.schedulers.wan.feature_caching.scheduler import (
|
||||
WanSchedulerTeaCaching,
|
||||
)
|
||||
|
||||
|
||||
class Lightx2vWanVideoT5EncoderLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"t5_model_path": (
|
||||
"STRING",
|
||||
{"default": "models/t5/models_t5_umt5-xxl-enc-bf16.pth"},
|
||||
),
|
||||
"tokenizer_path": ("STRING", {"default": "models/t5/google/umt5-xxl"}),
|
||||
"text_len": (
|
||||
"INT",
|
||||
{"default": 512, "min": 64, "max": 2048, "step": 1},
|
||||
),
|
||||
"precision": (["bf16", "fp16", "fp32"], {"default": "bf16"}),
|
||||
"device": (["cuda", "cpu"], {"default": "cuda"}),
|
||||
"additional_param": ("STRING", {"default": "default_value"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LIGHT_T5_ENCODER",)
|
||||
RETURN_NAMES = ("t5_encoder",)
|
||||
FUNCTION = "load_t5_encoder"
|
||||
CATEGORY = "LightX2V"
|
||||
|
||||
def load_t5_encoder(
|
||||
self,
|
||||
t5_model_path,
|
||||
tokenizer_path,
|
||||
text_len,
|
||||
precision,
|
||||
device,
|
||||
additional_param,
|
||||
):
|
||||
# Map precision to torch dtype
|
||||
dtype_map = {
|
||||
"bf16": torch.bfloat16,
|
||||
"fp16": torch.float16,
|
||||
"fp32": torch.float32,
|
||||
}
|
||||
dtype = dtype_map[precision]
|
||||
|
||||
# Resolve device
|
||||
if device == "cuda":
|
||||
device = mm.get_torch_device()
|
||||
else:
|
||||
device = torch.device("cpu")
|
||||
|
||||
# Load the T5 encoder
|
||||
t5_encoder = T5EncoderModel(
|
||||
text_len=text_len,
|
||||
dtype=dtype,
|
||||
device=device, # type:ignore
|
||||
checkpoint_path=t5_model_path,
|
||||
tokenizer_path=tokenizer_path,
|
||||
shard_fn=None,
|
||||
)
|
||||
|
||||
return (t5_encoder,)
|
||||
|
||||
|
||||
class Lightx2vWanVideoT5Encoder:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"t5_encoder": ("LIGHT_T5_ENCODER",),
|
||||
"prompt": (
|
||||
"STRING",
|
||||
{
|
||||
"multiline": True,
|
||||
"default": "A beautiful landscape with mountains and a lake",
|
||||
},
|
||||
),
|
||||
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LIGHT_TEXT_EMBEDDINGS",)
|
||||
RETURN_NAMES = ("text_embeddings",)
|
||||
FUNCTION = "encode_text"
|
||||
CATEGORY = "LightX2V"
|
||||
|
||||
def encode_text(self, t5_encoder, prompt, negative_prompt):
|
||||
# Create a config object with required attributes
|
||||
class Config:
|
||||
def __init__(self, cpu_offload=False):
|
||||
self.cpu_offload = cpu_offload
|
||||
|
||||
# NOTE(xxx): adapt the config if cpu_offload is set, t5 model must be on cuda device
|
||||
config = Config(cpu_offload=False)
|
||||
|
||||
# Encode the text
|
||||
context = t5_encoder.infer([prompt], config)
|
||||
context_null = t5_encoder.infer(
|
||||
[negative_prompt if negative_prompt else ""], config
|
||||
)
|
||||
|
||||
# Create text embeddings dictionary
|
||||
text_embeddings = {"context": context, "context_null": context_null}
|
||||
|
||||
return (text_embeddings,)
|
||||
|
||||
|
||||
class Lightx2vWanVideoClipVisionEncoderLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"clip_model_path": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "models/clip/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth"
|
||||
},
|
||||
),
|
||||
"tokenizer_path": (
|
||||
"STRING",
|
||||
{"default": "models/clip/xlm-roberta-large"},
|
||||
),
|
||||
"precision": (["fp16", "fp32"], {"default": "fp16"}),
|
||||
"device": (["cuda", "cpu"], {"default": "cuda"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LIGHT_CLIP_VISION_ENCODER",)
|
||||
RETURN_NAMES = ("clip_vision_encoder",)
|
||||
FUNCTION = "load_clip_vision_encoder"
|
||||
CATEGORY = "LightX2V"
|
||||
|
||||
def load_clip_vision_encoder(
|
||||
self, clip_model_path, tokenizer_path, precision, device
|
||||
):
|
||||
# Map precision to torch dtype
|
||||
dtype_map = {"fp16": torch.float16, "fp32": torch.float32}
|
||||
dtype = dtype_map[precision]
|
||||
|
||||
# Resolve device
|
||||
if device == "cuda":
|
||||
device = mm.get_torch_device()
|
||||
else:
|
||||
device = torch.device("cpu")
|
||||
|
||||
# Load the CLIP vision encoder
|
||||
clip_vision_encoder = ClipVisionModel(
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
checkpoint_path=clip_model_path,
|
||||
tokenizer_path=tokenizer_path,
|
||||
)
|
||||
|
||||
return (clip_vision_encoder,)
|
||||
|
||||
|
||||
class Lightx2vWanVideoVaeLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"vae_model_path": ("STRING", {"default": "models/vae/Wan2.1_VAE.pth"}),
|
||||
"precision": (["bf16", "fp16", "fp32"], {"default": "fp16"}),
|
||||
"device": (["cuda", "cpu"], {"default": "cuda"}),
|
||||
"parallel": ("BOOLEAN", {"default": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LIGHT_WAN_VAE",)
|
||||
RETURN_NAMES = ("wan_vae",)
|
||||
FUNCTION = "load_vae"
|
||||
CATEGORY = "LightX2V"
|
||||
|
||||
def load_vae(self, vae_model_path, precision, device, parallel):
|
||||
# Map precision to torch dtype
|
||||
dtype_map = {
|
||||
"bf16": torch.bfloat16,
|
||||
"fp16": torch.float16,
|
||||
"fp32": torch.float32,
|
||||
}
|
||||
dtype = dtype_map[precision]
|
||||
|
||||
# Resolve device
|
||||
if device == "cuda":
|
||||
device = mm.get_torch_device()
|
||||
else:
|
||||
device = torch.device("cpu")
|
||||
|
||||
# Load the VAE
|
||||
vae = WanVAE(
|
||||
z_dim=16,
|
||||
vae_pth=vae_model_path,
|
||||
dtype=dtype,
|
||||
device=device, # type:ignore
|
||||
parallel=parallel,
|
||||
)
|
||||
|
||||
return (vae,)
|
||||
|
||||
|
||||
class Lightx2vWanVideoVaeDecoder:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("LIGHT_WAN_VAE",),
|
||||
"latent": ("LIGHT_LATENT",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "decode_latent"
|
||||
CATEGORY = "LightX2V"
|
||||
|
||||
def decode_latent(self, vae, latent):
|
||||
# 创建一个带有必要属性的配置对象
|
||||
class Config:
|
||||
def __init__(self, cpu_offload=False):
|
||||
self.cpu_offload = cpu_offload
|
||||
|
||||
config = Config(cpu_offload=False)
|
||||
|
||||
# 获取潜在表示和生成器
|
||||
latents = latent["samples"]
|
||||
generator = latent["generator"]
|
||||
|
||||
# 使用VAE解码潜在表示
|
||||
with torch.no_grad():
|
||||
# 解码得到视频帧
|
||||
decoded_images = vae.decode(latents, generator=generator, config=config)
|
||||
|
||||
# 将像素值从 [-1, 1] 归一化到 [0, 1]
|
||||
images = (decoded_images + 1) / 2
|
||||
|
||||
# 重新排列维度为ComfyUI标准的图像格式 [T, H, W, C]
|
||||
# 从 [1, C, T, H, W] 转换为 [T, H, W, C]
|
||||
images = images.squeeze(0).permute(1, 2, 3, 0).cpu()
|
||||
|
||||
# 确保像素值在有效范围内
|
||||
images = torch.clamp(images, 0, 1)
|
||||
|
||||
# 清理缓存以释放GPU内存
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
return (images,)
|
||||
|
||||
|
||||
class Lightx2vWanVideoImageEncoder:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"vae": ("LIGHT_WAN_VAE",),
|
||||
"image": ("IMAGE",),
|
||||
"clip_vision_encoder": ("LIGHT_CLIP_VISION_ENCODER",),
|
||||
"target_height": (
|
||||
"INT",
|
||||
{"default": 576, "min": 256, "max": 1024, "step": 8},
|
||||
),
|
||||
"target_width": (
|
||||
"INT",
|
||||
{"default": 1024, "min": 256, "max": 1024, "step": 8},
|
||||
),
|
||||
"vae_stride": ("INT", {"default": 8, "min": 1, "max": 32, "step": 1}),
|
||||
"patch_size": ("INT", {"default": 2, "min": 1, "max": 16, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LIGHT_IMAGE_EMBEDDINGS",)
|
||||
RETURN_NAMES = ("image_embeddings",)
|
||||
FUNCTION = "encode_image"
|
||||
CATEGORY = "LightX2V"
|
||||
|
||||
def encode_image(
|
||||
self,
|
||||
vae,
|
||||
image,
|
||||
clip_vision_encoder,
|
||||
target_height,
|
||||
target_width,
|
||||
vae_stride,
|
||||
patch_size,
|
||||
):
|
||||
# 创建配置对象
|
||||
class Config:
|
||||
def __init__(
|
||||
self,
|
||||
target_height,
|
||||
target_width,
|
||||
vae_stride,
|
||||
patch_size,
|
||||
cpu_offload=False,
|
||||
):
|
||||
self.target_height = target_height
|
||||
self.target_width = target_width
|
||||
self.vae_stride = [1, vae_stride, vae_stride]
|
||||
self.patch_size = [1, patch_size, patch_size]
|
||||
self.cpu_offload = cpu_offload
|
||||
|
||||
config = Config(target_height, target_width, vae_stride, patch_size)
|
||||
|
||||
# 将图像转换为期望的张量格式
|
||||
device = mm.get_torch_device()
|
||||
img = image[0].permute(2, 0, 1).to(device) # [C, H, W]
|
||||
img = img.sub_(0.5).div_(0.5) # 归一化到 [-1, 1]
|
||||
|
||||
# 使用CLIP视觉编码器编码图像
|
||||
clip_encoder_out = clip_vision_encoder.visual(
|
||||
[img[:, None, :, :]], config
|
||||
).squeeze(0)
|
||||
|
||||
# 计算宽高比和尺寸
|
||||
h, w = img.shape[1:]
|
||||
aspect_ratio = h / w
|
||||
max_area = config.target_height * config.target_width
|
||||
lat_h = round(
|
||||
np.sqrt(max_area * aspect_ratio)
|
||||
// config.vae_stride[1]
|
||||
// config.patch_size[1]
|
||||
* config.patch_size[1]
|
||||
)
|
||||
lat_w = round(
|
||||
np.sqrt(max_area / aspect_ratio)
|
||||
// config.vae_stride[2]
|
||||
// config.patch_size[2]
|
||||
* config.patch_size[2]
|
||||
)
|
||||
h = lat_h * config.vae_stride[1]
|
||||
w = lat_w * config.vae_stride[2]
|
||||
|
||||
# 创建掩码
|
||||
msk = torch.ones(1, 81, lat_h, lat_w, device=device)
|
||||
msk[:, 1:] = 0
|
||||
msk = torch.concat(
|
||||
[torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1
|
||||
)
|
||||
msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w)
|
||||
msk = msk.transpose(1, 2)[0]
|
||||
|
||||
# 使用VAE编码图像
|
||||
resized_img = torch.nn.functional.interpolate(
|
||||
img[None].cpu(), size=(h, w), mode="bicubic"
|
||||
).transpose(0, 1)
|
||||
padding = torch.zeros(3, 80, h, w, device=device)
|
||||
concat_img = torch.concat([resized_img.to(device), padding], dim=1)
|
||||
|
||||
vae_encode_out = vae.encode([concat_img], config)[0]
|
||||
# TODO(xxx): hard code
|
||||
vae_encode_out = torch.concat([msk, vae_encode_out]).to(torch.bfloat16)
|
||||
|
||||
# 创建图像嵌入字典
|
||||
image_embeddings = {
|
||||
"clip_encoder_out": clip_encoder_out,
|
||||
"vae_encode_out": vae_encode_out,
|
||||
"lat_h": lat_h,
|
||||
"lat_w": lat_w,
|
||||
}
|
||||
|
||||
return (image_embeddings,)
|
||||
|
||||
|
||||
class Lightx2vWanVideoModelLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model_path": ("STRING", {"default": "models/wan"}),
|
||||
"model_type": (["t2v", "i2v"], {"default": "t2v"}),
|
||||
"precision": (["bf16", "fp16", "fp32"], {"default": "bf16"}),
|
||||
"device": (["cuda", "cpu"], {"default": "cuda"}),
|
||||
"feature_caching": (["NoCaching", "Tea"], {"default": "Tea"}),
|
||||
"teacache_thresh": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.26,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"tooltip": "Only used with Tea feature caching",
|
||||
},
|
||||
),
|
||||
"attention_mode": (
|
||||
["sdpa", "flash_attn_2", "flash_attn_3"],
|
||||
{"default": "flash_attn_3"},
|
||||
),
|
||||
"cpu_offload": ("BOOLEAN", {"default": False}),
|
||||
"mm_type": ("STRING", {"default": "Default"}),
|
||||
},
|
||||
"optional": {
|
||||
"lora_path": ("STRING", {"default": None}),
|
||||
"lora_strength": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LIGHT_WAN_MODEL",)
|
||||
RETURN_NAMES = ("wan_model",)
|
||||
FUNCTION = "load_model"
|
||||
CATEGORY = "LightX2V"
|
||||
|
||||
def load_model(
|
||||
self,
|
||||
model_path,
|
||||
model_type,
|
||||
precision,
|
||||
device,
|
||||
feature_caching,
|
||||
teacache_thresh,
|
||||
attention_mode,
|
||||
mm_type,
|
||||
lora_path=None,
|
||||
lora_strength=1.0,
|
||||
):
|
||||
# 映射精度到torch dtype
|
||||
dtype_map = {
|
||||
"bf16": torch.bfloat16,
|
||||
"fp16": torch.float16,
|
||||
"fp32": torch.float32,
|
||||
}
|
||||
dtype = dtype_map[precision]
|
||||
|
||||
# 解析设备
|
||||
if device == "cuda":
|
||||
device = mm.get_torch_device()
|
||||
else:
|
||||
device = torch.device("cpu")
|
||||
|
||||
# 首先在模型路径中查找config.json
|
||||
model_path_dir = os.path.dirname(model_path)
|
||||
config_json_path = os.path.join(model_path_dir, "config.json")
|
||||
config_json = {}
|
||||
if os.path.exists(config_json_path):
|
||||
with open(config_json_path, "r") as f:
|
||||
config_json = json.load(f)
|
||||
else:
|
||||
logging.error(f"Config file not found at {config_json_path}")
|
||||
raise FileNotFoundError(f"Config file not found at {config_json_path}")
|
||||
|
||||
# 创建配置字典
|
||||
config = {
|
||||
"do_mm_calib": False,
|
||||
"cpu_offload": False,
|
||||
"parallel_attn_type": None, # [None, "ulysses", "ring"]
|
||||
"parallel_vae": False,
|
||||
"max_area": False,
|
||||
"vae_stride": (4, 8, 8),
|
||||
"patch_size": (1, 2, 2),
|
||||
"feature_caching": feature_caching, # ["NoCaching", "TaylorSeer", "Tea"]
|
||||
"teacache_thresh": teacache_thresh,
|
||||
"use_ret_steps": False,
|
||||
"use_bfloat16": dtype == torch.bfloat16,
|
||||
"mm_config": {
|
||||
"mm_type": mm_type,
|
||||
"weight_auto_quant": False if mm_type == "Default" else True,
|
||||
},
|
||||
"model_path": model_path,
|
||||
"task": model_type,
|
||||
"device": device,
|
||||
"attention_type": attention_mode,
|
||||
"lora_path": lora_path if lora_path and lora_path.strip() else None,
|
||||
"strength_model": lora_strength,
|
||||
}
|
||||
# merge model dir config.json and config dict
|
||||
config.update(**config_json)
|
||||
# NOTE(xxx): adapt to Lightx2v
|
||||
config = EasyDict(config)
|
||||
|
||||
logging.info(
|
||||
"Loaded config:\n" + json.dumps(config, indent=4, ensure_ascii=False)
|
||||
)
|
||||
logging.info(f"Loading WanModel from {model_path} with type {model_type}")
|
||||
model = WanModel(model_path, config, device)
|
||||
|
||||
# 如果指定了LoRA路径，应用LoRA
|
||||
if lora_path and os.path.exists(lora_path):
|
||||
logging.info(
|
||||
f"Applying LoRA from {lora_path} with strength {lora_strength}"
|
||||
)
|
||||
lora_wrapper = WanLoraWrapper(model)
|
||||
lora_name = lora_wrapper.load_lora(lora_path)
|
||||
lora_wrapper.apply_lora(lora_name, lora_strength)
|
||||
logging.info(f"LoRA {lora_name} applied successfully")
|
||||
|
||||
wan_model = {"wan_model": model, "config": config}
|
||||
|
||||
return (wan_model,)
|
||||
|
||||
|
||||
class Lightx2vWanVideoSampler:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("LIGHT_WAN_MODEL",),
|
||||
"text_embeddings": ("LIGHT_TEXT_EMBEDDINGS",),
|
||||
"steps": ("INT", {"default": 20, "min": 1, "max": 100, "step": 1}),
|
||||
"cfg_scale": (
|
||||
"FLOAT",
|
||||
{"default": 7.5, "min": 0.0, "max": 20.0, "step": 0.1},
|
||||
),
|
||||
"seed": ("INT", {"default": 42}),
|
||||
},
|
||||
"optional": {
|
||||
"image_embeddings": ("LIGHT_IMAGE_EMBEDDINGS",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LIGHT_LATENT",)
|
||||
RETURN_NAMES = ("latent",)
|
||||
FUNCTION = "sample"
|
||||
CATEGORY = "LightX2V"
|
||||
|
||||
def sample(
|
||||
self,
|
||||
model,
|
||||
text_embeddings,
|
||||
steps,
|
||||
cfg_scale,
|
||||
seed,
|
||||
image_embeddings=None,
|
||||
):
|
||||
# Update model config
|
||||
|
||||
config = cast(Any, model.get("config"))
|
||||
model = cast(WanModel, model.get("wan_model"))
|
||||
|
||||
seed_all(seed)
|
||||
|
||||
# wan_runner.init_scheduler
|
||||
if config.feature_caching == "NoCaching": # type:ignore
|
||||
scheduler = WanScheduler(config)
|
||||
elif config.feature_caching == "Tea": # type:ignore
|
||||
scheduler = WanSchedulerTeaCaching(config)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported feature_caching type: {config.feature_caching}" # type:ignore
|
||||
)
|
||||
|
||||
# setup scheduler
|
||||
model.set_scheduler(scheduler)
|
||||
|
||||
# wan_runner.set_target_shape
|
||||
if config.task == "i2v":
|
||||
if image_embeddings is None:
|
||||
raise ValueError("image_embeddings must be provided for i2v task")
|
||||
config.lat_h = image_embeddings["lat_h"]
|
||||
config.lat_w = image_embeddings["lat_w"]
|
||||
config.target_shape = (16, 21, config.lat_h, config.lat_w)
|
||||
elif config.task == "t2v":
|
||||
config.target_shape = (
|
||||
16,
|
||||
(config.target_video_length - 1) // 4 + 1,
|
||||
int(config.target_height) // config.vae_stride[1],
|
||||
int(config.target_width) // config.vae_stride[2],
|
||||
)
|
||||
|
||||
# Set up inputs
|
||||
inputs = {
|
||||
"text_encoder_output": text_embeddings,
|
||||
"image_encoder_output": image_embeddings,
|
||||
}
|
||||
|
||||
# Prepare for sampling
|
||||
scheduler.prepare(inputs.get("image_encoder_output"))
|
||||
|
||||
# Run sampling
|
||||
progress = ProgressBar(steps)
|
||||
for step_index in tqdm(
|
||||
range(scheduler.infer_steps), desc="inference", unit="step"
|
||||
):
|
||||
with ProfilingContext("scheduler.step_pre"):
|
||||
scheduler.step_pre(step_index=step_index)
|
||||
with ProfilingContext("model.infer"):
|
||||
model.infer(inputs)
|
||||
with ProfilingContext("scheduler.step_post"):
|
||||
scheduler.step_post()
|
||||
|
||||
progress.update(1)
|
||||
|
||||
scheduler.clear()
|
||||
|
||||
del inputs, model, text_embeddings, image_embeddings
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return ({"samples": scheduler.latents, "generator": scheduler.generator},)
|
||||
|
||||
|
||||
# Register the nodes
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Lightx2vWanVideoT5EncoderLoader": Lightx2vWanVideoT5EncoderLoader,
|
||||
"Lightx2vWanVideoT5Encoder": Lightx2vWanVideoT5Encoder,
|
||||
"Lightx2vWanVideoClipVisionEncoderLoader": Lightx2vWanVideoClipVisionEncoderLoader,
|
||||
"Lightx2vWanVideoVaeLoader": Lightx2vWanVideoVaeLoader,
|
||||
"Lightx2vWanVideoImageEncoder": Lightx2vWanVideoImageEncoder,
|
||||
"Lightx2vWanVideoVaeDecoder": Lightx2vWanVideoVaeDecoder,
|
||||
"Lightx2vWanVideoModelLoader": Lightx2vWanVideoModelLoader,
|
||||
"Lightx2vWanVideoSampler": Lightx2vWanVideoSampler,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Lightx2vWanVideoT5EncoderLoader": "LightX2V WAN T5 Encoder Loader",
|
||||
"Lightx2vWanVideoT5Encoder": "LightX2V WAN T5 Encoder",
|
||||
"Lightx2vWanVideoClipVisionEncoderLoader": "LightX2V WAN CLIP Vision Encoder Loader",
|
||||
"Lightx2vWanVideoClipVisionEncoder": "LightX2V WAN CLIP Vision Encoder",
|
||||
"Lightx2vWanVideoVaeLoader": "LightX2V WAN VAE Loader",
|
||||
"Lightx2vWanVideoImageEncoder": "LightX2V WAN Image Encoder",
|
||||
"Lightx2vWanVideoVaeDecoder": "LightX2V WAN VAE Decoder",
|
||||
"Lightx2vWanVideoModelLoader": "LightX2V WAN Model Loader",
|
||||
"Lightx2vWanVideoSampler": "LightX2V WAN Video Sampler",
|
||||
}
|
||||
Reference in New Issue
Block a user