447 lines
14 KiB
Python
447 lines
14 KiB
Python
"""
|
|
ComfyUI custom nodes for WAN Video-as-Prompt (VAP).
|
|
|
|
Provides two nodes:
|
|
- `WANVAP_Load`: loads/cache the WAN pipeline for a given `model_id`.
|
|
- `WANVAP_Sample`: runs sampling on a target image and reference video frames,
|
|
producing output frames.
|
|
|
|
This file is intentionally self-contained and re-implements the minimal
|
|
helpers used by the original `infer/wan_vap.py` script.
|
|
"""
|
|
from __future__ import annotations
|
|
import os
|
|
import random
|
|
import shutil
|
|
from typing import List, Optional
|
|
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
import sys
|
|
|
|
_repo_dir = os.path.dirname(__file__)
|
|
if _repo_dir and _repo_dir not in sys.path:
|
|
sys.path.insert(0, _repo_dir)
|
|
|
|
from diffusers_vap import (
|
|
AutoencoderKLWan,
|
|
WanImageToVideoMOTPipeline,
|
|
WanTransformer3DMOTModel,
|
|
)
|
|
|
|
from transformers import CLIPVisionModel
|
|
|
|
import folder_paths
|
|
from comfy.utils import ProgressBar
|
|
from huggingface_hub import snapshot_download
|
|
|
|
# Global cache for loaded pipelines
|
|
_GLOBAL = {"pipelines": {}}
|
|
|
|
def _normalize_device(device_hint: Optional[str] = None) -> str:
|
|
if device_hint is None:
|
|
return "cuda" if torch.cuda.is_available() else "cpu"
|
|
device = str(device_hint).strip()
|
|
if not device or device.lower() in ("auto", "none"):
|
|
return "cuda" if torch.cuda.is_available() else "cpu"
|
|
return device.lower()
|
|
|
|
|
|
def _resolve_model_source(model_id: str) -> str:
|
|
if not model_id:
|
|
raise ValueError("model_id must be provided.")
|
|
|
|
candidates: List[str] = []
|
|
candidates.append(model_id)
|
|
candidates.append(os.path.join(_repo_dir, model_id))
|
|
candidates.append(os.path.join(_repo_dir, "ckpts", model_id))
|
|
|
|
base_models_dir = getattr(folder_paths, "models_dir", None)
|
|
if base_models_dir:
|
|
for subdir in ("", "video_as_prompt", "video-as-prompt", "Video-As-Prompt"):
|
|
base_dir = os.path.join(base_models_dir, subdir) if subdir else base_models_dir
|
|
candidates.append(os.path.join(base_dir, model_id))
|
|
candidates.append(os.path.join(base_models_dir, "ckpts", model_id))
|
|
|
|
seen = set()
|
|
for candidate in candidates:
|
|
if not candidate:
|
|
continue
|
|
abs_path = os.path.abspath(os.path.expanduser(candidate))
|
|
if abs_path in seen:
|
|
continue
|
|
seen.add(abs_path)
|
|
if os.path.isdir(abs_path):
|
|
pw = os.path.join(abs_path, "pretrained_weights")
|
|
if os.path.isdir(pw):
|
|
return pw
|
|
return abs_path
|
|
|
|
return model_id
|
|
|
|
|
|
def _maybe_download_model(model_id: str, base_models_dir: Optional[str]) -> str:
|
|
if not base_models_dir:
|
|
return model_id
|
|
try:
|
|
repo_name = model_id.replace("/", "_")
|
|
target_dir = os.path.join(base_models_dir, repo_name)
|
|
if not os.path.exists(target_dir):
|
|
hf_path = snapshot_download(repo_id=model_id)
|
|
try:
|
|
shutil.copytree(hf_path, target_dir)
|
|
except Exception:
|
|
target_dir = hf_path
|
|
|
|
pw = os.path.join(target_dir, "pretrained_weights")
|
|
if os.path.exists(pw):
|
|
return pw
|
|
return target_dir
|
|
except Exception:
|
|
return model_id
|
|
|
|
def _tensor_to_pil_list(frames: torch.Tensor) -> List[Image.Image]:
|
|
if frames.dim() == 3:
|
|
frames = frames.unsqueeze(0)
|
|
if frames.dim() != 4:
|
|
raise ValueError("Expected a 3D or 4D tensor for image frames.")
|
|
|
|
if frames.shape[-1] in (1, 3, 4):
|
|
frames_bhwc = frames
|
|
elif frames.shape[1] in (1, 3, 4):
|
|
frames_bhwc = frames.permute(0, 2, 3, 1)
|
|
else:
|
|
raise ValueError("Expected BHWC or BCHW tensor with 1, 3, or 4 channels.")
|
|
|
|
frames_bhwc = frames_bhwc.detach().cpu()
|
|
if frames_bhwc.is_floating_point():
|
|
frames_bhwc = (frames_bhwc.clamp(0, 1) * 255).round().to(torch.uint8)
|
|
else:
|
|
frames_bhwc = frames_bhwc.clamp(0, 255).to(torch.uint8)
|
|
|
|
pil_frames = []
|
|
for frame in frames_bhwc:
|
|
arr = frame.numpy()
|
|
if arr.shape[2] == 1:
|
|
arr = np.repeat(arr, 3, axis=2)
|
|
elif arr.shape[2] == 4:
|
|
arr = arr[:, :, :3]
|
|
pil_frames.append(Image.fromarray(arr))
|
|
return pil_frames
|
|
|
|
|
|
def _array_to_pil(frame: np.ndarray) -> Image.Image:
|
|
if frame.ndim == 2:
|
|
frame = np.repeat(frame[:, :, None], 3, axis=2)
|
|
if frame.ndim != 3:
|
|
raise ValueError("Expected an HWC array for image data.")
|
|
if frame.shape[2] == 1:
|
|
frame = np.repeat(frame, 3, axis=2)
|
|
elif frame.shape[2] > 3:
|
|
frame = frame[:, :, :3]
|
|
if frame.dtype != np.uint8:
|
|
if frame.dtype.kind == "f":
|
|
frame = np.clip(frame, 0.0, 1.0)
|
|
frame = (frame * 255.0).round().astype(np.uint8)
|
|
else:
|
|
frame = np.clip(frame, 0, 255).astype(np.uint8)
|
|
return Image.fromarray(frame)
|
|
|
|
|
|
def _coerce_pil_image(image) -> Image.Image:
|
|
if isinstance(image, Image.Image):
|
|
return image
|
|
if torch.is_tensor(image):
|
|
frames = _tensor_to_pil_list(image)
|
|
return frames[0] if frames else None
|
|
if isinstance(image, np.ndarray):
|
|
return _array_to_pil(image)
|
|
if isinstance(image, list) and image:
|
|
return _coerce_pil_image(image[0])
|
|
return None
|
|
|
|
|
|
def _coerce_video_frames(ref_video) -> List[Image.Image]:
|
|
if ref_video is None:
|
|
return []
|
|
if torch.is_tensor(ref_video):
|
|
return _tensor_to_pil_list(ref_video)
|
|
if isinstance(ref_video, list):
|
|
frames: List[Image.Image] = []
|
|
for item in ref_video:
|
|
if torch.is_tensor(item):
|
|
frames.extend(_tensor_to_pil_list(item))
|
|
continue
|
|
if isinstance(item, Image.Image):
|
|
frames.append(item)
|
|
continue
|
|
if isinstance(item, np.ndarray):
|
|
frames.append(_array_to_pil(item))
|
|
continue
|
|
coerced = _coerce_pil_image(item)
|
|
if coerced is not None:
|
|
frames.append(coerced)
|
|
return frames
|
|
if isinstance(ref_video, Image.Image):
|
|
return [ref_video]
|
|
if isinstance(ref_video, np.ndarray):
|
|
return [_array_to_pil(ref_video)]
|
|
coerced = _coerce_pil_image(ref_video)
|
|
return [coerced] if coerced is not None else []
|
|
|
|
|
|
def _pil_list_to_tensor(frames: List[Image.Image]) -> torch.Tensor:
|
|
if not frames:
|
|
raise ValueError("No frames provided for conversion.")
|
|
tensors = []
|
|
for frame in frames:
|
|
if frame.mode != "RGB":
|
|
frame = frame.convert("RGB")
|
|
arr = np.asarray(frame, dtype=np.float32) / 255.0
|
|
tensors.append(torch.from_numpy(arr))
|
|
return torch.stack(tensors, dim=0)
|
|
|
|
|
|
def select_frames(video_frames: List[Image.Image], num: int, mode: str) -> List[Image.Image]:
|
|
if len(video_frames) == 0:
|
|
return []
|
|
if mode == "first":
|
|
return video_frames[:num]
|
|
if mode == "evenly":
|
|
import torch as _torch
|
|
|
|
idx = _torch.linspace(0, len(video_frames) - 1, num).long().tolist()
|
|
return [video_frames[i] for i in idx]
|
|
if mode == "random":
|
|
if len(video_frames) <= num:
|
|
return video_frames
|
|
import random as _random
|
|
|
|
start = _random.randint(0, len(video_frames) - num)
|
|
return video_frames[start : start + num]
|
|
return video_frames
|
|
|
|
|
|
def set_global_seed(seed: int):
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
|
|
def _device_and_dtype_for(device_hint: Optional[str] = None):
|
|
device = _normalize_device(device_hint)
|
|
dtype = torch.bfloat16 if device.startswith("cuda") else torch.float32
|
|
return device, dtype
|
|
|
|
|
|
def get_folder_list():
|
|
base_dir = folder_paths.models_dir
|
|
if not os.path.exists(base_dir):
|
|
return ["video_as_prompt"]
|
|
|
|
candidates = []
|
|
for name in os.listdir(base_dir):
|
|
full_path = os.path.join(base_dir, name)
|
|
if os.path.isdir(full_path):
|
|
if os.path.exists(os.path.join(full_path, "pretrained_weights")):
|
|
candidates.append(name)
|
|
|
|
if "video_as_prompt" not in candidates:
|
|
candidates.append("video_as_prompt")
|
|
|
|
return sorted(candidates)
|
|
|
|
|
|
def load_model(model_id: str = "ByteDance/Video-As-Prompt-Wan2.1-14B", device: Optional[str] = None):
|
|
model_source = _resolve_model_source(model_id)
|
|
base_models_dir = getattr(folder_paths, "models_dir", None) or os.path.join(os.getcwd(), "models")
|
|
|
|
if not os.path.exists(model_source):
|
|
model_source = _maybe_download_model(model_source, base_models_dir)
|
|
|
|
device, dtype = _device_and_dtype_for(device)
|
|
cache_key = (model_source, device)
|
|
cached = _GLOBAL["pipelines"].get(cache_key)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
# load components similarly to the original script using resolved model_source
|
|
image_encoder = CLIPVisionModel.from_pretrained(model_source, subfolder="image_encoder", torch_dtype=torch.float32)
|
|
vae = AutoencoderKLWan.from_pretrained(model_source, subfolder="vae", torch_dtype=torch.float32)
|
|
transformer = WanTransformer3DMOTModel.from_pretrained(model_source, subfolder="transformer", torch_dtype=dtype)
|
|
|
|
pipe = WanImageToVideoMOTPipeline.from_pretrained(
|
|
model_source,
|
|
vae=vae,
|
|
image_encoder=image_encoder,
|
|
transformer=transformer,
|
|
torch_dtype=dtype,
|
|
).to(device)
|
|
|
|
# try to enable sequential CPU offload when running on CUDA to reduce memory pressure
|
|
if device.startswith("cuda"):
|
|
try:
|
|
pipe.enable_sequential_cpu_offload()
|
|
except Exception:
|
|
pass
|
|
|
|
_GLOBAL["pipelines"][cache_key] = pipe
|
|
return pipe
|
|
|
|
|
|
def sample_from_pipe(
|
|
pipe,
|
|
image: Image.Image,
|
|
ref_video,
|
|
prompt: str,
|
|
prompt_mot_ref: Optional[List[str]] = None,
|
|
height: int = 480,
|
|
width: int = 832,
|
|
num_frames: int = 49,
|
|
num_inference_steps: int = 50,
|
|
frames_selection: str = "evenly",
|
|
guidance_scale: float = 5.0,
|
|
seed: int = 42,
|
|
):
|
|
set_global_seed(int(seed))
|
|
|
|
if pipe is None:
|
|
raise ValueError("pipeline is required. Use WANVAP_Load to create it.")
|
|
|
|
image = _coerce_pil_image(image) or image
|
|
if image is None:
|
|
raise ValueError("image must be a valid ComfyUI IMAGE or PIL image.")
|
|
|
|
ref_frames = _coerce_video_frames(ref_video)
|
|
if not ref_frames:
|
|
raise ValueError("ref_video must be provided as ComfyUI IMAGE frames.")
|
|
|
|
if isinstance(prompt_mot_ref, str):
|
|
prompt_mot_ref = [prompt_mot_ref]
|
|
if not prompt_mot_ref:
|
|
raise ValueError("prompt_mot_ref is required. Provide one or more prompts separated by '|'.")
|
|
|
|
ref_frames = select_frames(ref_frames, num_frames, frames_selection)
|
|
|
|
num_inference_steps = max(int(num_inference_steps), 1)
|
|
progress = ProgressBar(num_inference_steps)
|
|
|
|
def _on_step_end(_pipe, _step, _timestep, callback_kwargs):
|
|
progress.update(1)
|
|
return callback_kwargs
|
|
|
|
out_frames = pipe(
|
|
image=image,
|
|
ref_videos=[ref_frames],
|
|
prompt=prompt,
|
|
prompt_mot_ref=prompt_mot_ref,
|
|
height=height,
|
|
width=width,
|
|
num_frames=num_frames,
|
|
num_inference_steps=num_inference_steps,
|
|
frames_selection=frames_selection,
|
|
guidance_scale=guidance_scale,
|
|
callback_on_step_end=_on_step_end,
|
|
callback_on_step_end_tensor_inputs=[],
|
|
output_type="pil",
|
|
).frames[0]
|
|
|
|
frames_tensor = _pil_list_to_tensor(out_frames)
|
|
return frames_tensor
|
|
|
|
|
|
# ------------------ ComfyUI node classes ------------------
|
|
class WANVAP_Load:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {"model_id": ("STRING", {"default": "ByteDance/Video-As-Prompt-Wan2.1-14B"})},
|
|
"optional": {"device": (["auto", "cuda", "cpu"],)},
|
|
}
|
|
|
|
RETURN_TYPES = ("WANVAP_PIPE",)
|
|
RETURN_NAMES = ("pipe",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "WAN/VAP"
|
|
|
|
def load(self, model_id: str = "ByteDance/Video-As-Prompt-Wan2.1-14B", device: Optional[str] = None):
|
|
"""Load and cache a WAN VAP pipeline for `model_id`."""
|
|
pipe = load_model(model_id, device=device)
|
|
return (pipe,)
|
|
|
|
|
|
class WANVAP_Sample:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pipe": ("WANVAP_PIPE",),
|
|
"image": ("IMAGE",),
|
|
"ref_video": ("IMAGE",),
|
|
"prompt": ("STRING",),
|
|
"prompt_mot_ref": ("STRING",),
|
|
},
|
|
"optional": {
|
|
"num_frames": ("INT", {"default": 49}),
|
|
"height": ("INT", {"default": 480}),
|
|
"width": ("INT", {"default": 832}),
|
|
"num_inference_steps": ("INT", {"default": 50}),
|
|
"frames_selection": (["evenly", "first", "random"],),
|
|
"guidance_scale": ("FLOAT", {"default": 5.0}),
|
|
"seed": ("INT", {"default": 42}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("frames",)
|
|
FUNCTION = "sample"
|
|
CATEGORY = "WAN/VAP"
|
|
|
|
def sample(
|
|
self,
|
|
pipe,
|
|
image,
|
|
ref_video,
|
|
prompt: str,
|
|
prompt_mot_ref: str,
|
|
num_frames: int = 49,
|
|
num_inference_steps: int = 50,
|
|
height: int = 480,
|
|
width: int = 832,
|
|
frames_selection: str = "evenly",
|
|
guidance_scale: float = 5.0,
|
|
seed: int = 42,
|
|
):
|
|
"""Run VAP sampling. `prompt_mot_ref` is required and may be a single string or
|
|
multiple prompts separated by `|`, which will be converted to a list.
|
|
"""
|
|
prompt_mot_ref_list = [s.strip() for s in prompt_mot_ref.split("|")] if prompt_mot_ref else []
|
|
|
|
frames = sample_from_pipe(
|
|
pipe=pipe,
|
|
image=image,
|
|
ref_video=ref_video,
|
|
prompt=prompt,
|
|
prompt_mot_ref=prompt_mot_ref_list,
|
|
height=height,
|
|
width=width,
|
|
num_frames=num_frames,
|
|
num_inference_steps=num_inference_steps,
|
|
frames_selection=frames_selection,
|
|
guidance_scale=guidance_scale,
|
|
seed=seed,
|
|
)
|
|
|
|
return (frames,)
|
|
|
|
|
|
# Expose node classes in the module-level variable expected by many ComfyUI loaders
|
|
NODE_CLASS_MAPPINGS = {
|
|
"WANVAP_Load": WANVAP_Load,
|
|
"WANVAP_Sample": WANVAP_Sample,
|
|
}
|