190 lines
6.6 KiB
Python
190 lines
6.6 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import os
|
|
import folder_paths
|
|
from huggingface_hub import snapshot_download, hf_hub_download
|
|
import comfy.utils
|
|
|
|
# 模型目录
|
|
TELESTYLE_MODEL_DIR = os.path.join(folder_paths.models_dir, "TeleStyle")
|
|
os.makedirs(TELESTYLE_MODEL_DIR, exist_ok=True)
|
|
|
|
# Wan2.1 模型配置
|
|
WAN_MODELS = {
|
|
"Wan2.1-T2V-1.3B": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
|
"Wan2.1-T2V-14B": "Wan-AI/Wan2.1-T2V-14B-Diffusers",
|
|
}
|
|
|
|
|
|
def download_wan_model(model_name):
|
|
"""下载 Wan2.1 模型到 models/TeleStyle 目录"""
|
|
if model_name not in WAN_MODELS:
|
|
raise ValueError(f"Unknown model: {model_name}")
|
|
|
|
repo_id = WAN_MODELS[model_name]
|
|
local_dir = os.path.join(TELESTYLE_MODEL_DIR, f"{model_name}-Diffusers")
|
|
|
|
if not os.path.exists(local_dir) or not os.listdir(local_dir):
|
|
print(f"[TeleStyle] Downloading {model_name} to {local_dir}...")
|
|
snapshot_download(
|
|
repo_id=repo_id,
|
|
local_dir=local_dir,
|
|
local_dir_use_symlinks=False
|
|
)
|
|
print(f"[TeleStyle] {model_name} downloaded successfully.")
|
|
|
|
return local_dir
|
|
|
|
|
|
def download_telestyle_video_models():
|
|
"""下载 TeleStyle 视频模型到 models/TeleStyle 目录"""
|
|
dit_path = os.path.join(TELESTYLE_MODEL_DIR, "dit.ckpt")
|
|
prompt_embeds_path = os.path.join(TELESTYLE_MODEL_DIR, "prompt_embeds.pth")
|
|
|
|
if not os.path.exists(dit_path):
|
|
print(f"[TeleStyle] Downloading dit.ckpt to {TELESTYLE_MODEL_DIR}...")
|
|
hf_hub_download(
|
|
repo_id="Tele-AI/TeleStyle",
|
|
filename="weights/dit.ckpt",
|
|
local_dir=TELESTYLE_MODEL_DIR,
|
|
local_dir_use_symlinks=False
|
|
)
|
|
src = os.path.join(TELESTYLE_MODEL_DIR, "weights", "dit.ckpt")
|
|
if os.path.exists(src):
|
|
os.rename(src, dit_path)
|
|
|
|
if not os.path.exists(prompt_embeds_path):
|
|
print(f"[TeleStyle] Downloading prompt_embeds.pth to {TELESTYLE_MODEL_DIR}...")
|
|
hf_hub_download(
|
|
repo_id="Tele-AI/TeleStyle",
|
|
filename="weights/prompt_embeds.pth",
|
|
local_dir=TELESTYLE_MODEL_DIR,
|
|
local_dir_use_symlinks=False
|
|
)
|
|
src = os.path.join(TELESTYLE_MODEL_DIR, "weights", "prompt_embeds.pth")
|
|
if os.path.exists(src):
|
|
os.rename(src, prompt_embeds_path)
|
|
|
|
return dit_path, prompt_embeds_path
|
|
|
|
|
|
class TeleStyleVideoLoader:
|
|
"""TeleStyle 视频模型加载节点"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"wan_model": (list(WAN_MODELS.keys()),),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TELESTYLE_VIDEO_MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "TeleStyle"
|
|
|
|
def load_model(self, wan_model):
|
|
from .TeleStyle.telestylevideo_inference import VideoStyleInference
|
|
|
|
ckpt_t2v_path = download_wan_model(wan_model)
|
|
dit_path, prompt_embeds_path = download_telestyle_video_models()
|
|
|
|
# 返回模型路径信息，实际推理参数在 TeleStyleVideo 节点设置
|
|
model_info = {
|
|
"ckpt_t2v_path": ckpt_t2v_path,
|
|
"ckpt_dit_path": dit_path,
|
|
"prompt_embeds_path": prompt_embeds_path,
|
|
}
|
|
return (model_info,)
|
|
|
|
|
|
class TeleStyleVideo:
|
|
"""TeleStyle 视频风格迁移节点"""
|
|
|
|
_engine_cache = {} # 缓存已加载的引擎
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("TELESTYLE_VIDEO_MODEL",),
|
|
"video_frames": ("IMAGE",),
|
|
"style_first_frame": ("IMAGE",),
|
|
"height": ("INT", {"default": 1024, "min": 256, "max": 9999, "step": 16}),
|
|
"width": ("INT", {"default": 1024, "min": 256, "max": 9999, "step": 16}),
|
|
"video_length": ("INT", {"default": 81, "min": 17, "max": 999, "step": 1}),
|
|
"num_inference_steps": ("INT", {"default": 25, "min": 1, "max": 50, "step": 1}),
|
|
"guidance_scale": ("FLOAT", {"default": 3.0, "min": 1.0, "max": 10.0, "step": 0.5}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("video_frames",)
|
|
FUNCTION = "transfer_style"
|
|
CATEGORY = "TeleStyle"
|
|
|
|
def transfer_style(self, model, video_frames, style_first_frame, height, width, video_length, num_inference_steps, guidance_scale, seed):
|
|
from .TeleStyle.telestylevideo_inference import VideoStyleInference
|
|
|
|
# 构建配置
|
|
config = {
|
|
"random_seed": seed,
|
|
"video_length": video_length,
|
|
"height": height,
|
|
"width": width,
|
|
"num_inference_steps": num_inference_steps,
|
|
"guidance_scale": guidance_scale,
|
|
"ckpt_t2v_path": model["ckpt_t2v_path"],
|
|
"ckpt_dit_path": model["ckpt_dit_path"],
|
|
"prompt_embeds_path": model["prompt_embeds_path"],
|
|
"output_path": folder_paths.get_temp_directory()
|
|
}
|
|
|
|
# 创建推理引擎
|
|
engine = VideoStyleInference(config)
|
|
|
|
# 创建进度条
|
|
pbar = comfy.utils.ProgressBar(num_inference_steps)
|
|
def progress_callback(step, total):
|
|
pbar.update(1)
|
|
|
|
# ComfyUI IMAGE 格式: (frames, H, W, C), 值范围 [0, 1]
|
|
# 原始代码期望格式: (batch, frames, H, W, C), 值范围 [-1, 1]
|
|
|
|
# 转换视频帧
|
|
source_video = video_frames.cpu() * 2.0 - 1.0 # [0,1] -> [-1,1]
|
|
source_video = source_video.unsqueeze(0) # (f, h, w, c) -> (1, f, h, w, c)
|
|
|
|
# 转换风格参考帧（只取第一帧）
|
|
style_frame = style_first_frame[0:1].cpu() * 2.0 - 1.0 # 取第一帧
|
|
style_frame = style_frame.unsqueeze(0) # (1, h, w, c) -> (1, 1, h, w, c)
|
|
|
|
with torch.no_grad():
|
|
generated_video = engine.inference(source_video, style_frame, "", 0, progress_callback)
|
|
|
|
# 转换输出格式回 ComfyUI IMAGE
|
|
frames = []
|
|
for frame in generated_video:
|
|
frame_np = np.array(frame).astype(np.float32)
|
|
if frame_np.max() > 1.0:
|
|
frame_np = frame_np / 255.0
|
|
frames.append(torch.from_numpy(frame_np))
|
|
|
|
result_tensor = torch.stack(frames, dim=0)
|
|
|
|
return (result_tensor,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"TeleStyleVideoLoader": TeleStyleVideoLoader,
|
|
"TeleStyleVideo": TeleStyleVideo,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"TeleStyleVideoLoader": "TeleStyle Video Loader",
|
|
"TeleStyleVideo": "TeleStyle Video",
|
|
}
|