130 lines
5.2 KiB
Python
130 lines
5.2 KiB
Python
import os
|
|
import torch
|
|
|
|
# Use a simple string for imageOrLatent to avoid JSON serialization issues
|
|
imageOrLatent = "IMAGE"
|
|
|
|
class WanVideoCacheSamples:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"cache_path": ("STRING", {"default": "cache_data", "multiline": False}),
|
|
"cache_name": ("STRING", {"default": "wanvideo_cached_samples.pt", "multiline": False}),
|
|
},
|
|
"optional": {
|
|
"samples": ("LATENT", {"forceInput": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
RETURN_NAMES = ("samples",)
|
|
FUNCTION = "cache_or_load_samples"
|
|
CATEGORY = "WanVideoWrapper_QQ/utils"
|
|
|
|
def cache_or_load_samples(self, cache_path, cache_name, samples=None):
|
|
# Get the path to the ComfyUI-WanVideoWrapper_QQ directory
|
|
qq_wrapper_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
# Sanitize cache_path to prevent it from being an absolute path
|
|
# and ensure it's just a directory name.
|
|
sanitized_cache_path = os.path.basename(cache_path)
|
|
|
|
# The cache directory will be inside the node's directory.
|
|
full_cache_dir = os.path.join(qq_wrapper_path, sanitized_cache_path)
|
|
os.makedirs(full_cache_dir, exist_ok=True)
|
|
|
|
full_cache_file_path = os.path.join(full_cache_dir, cache_name)
|
|
|
|
if samples is not None:
|
|
# Samples are provided, cache them and pass them through
|
|
print(f"WanVideoCacheSamples: Caching samples to {full_cache_file_path}")
|
|
torch.save(samples, full_cache_file_path)
|
|
return (samples,)
|
|
else:
|
|
# Samples are not provided, try to load from cache
|
|
if os.path.exists(full_cache_file_path):
|
|
print(f"WanVideoCacheSamples: Loading samples from {full_cache_file_path}")
|
|
loaded_samples = torch.load(full_cache_file_path)
|
|
return (loaded_samples,)
|
|
else:
|
|
raise FileNotFoundError(f"WanVideoCacheSamples: No samples provided and no cached samples found at {full_cache_file_path}")
|
|
|
|
|
|
class WanVideoCacheVideos:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"cache_path": ("STRING", {"default": "cache_data", "multiline": False}),
|
|
"cache_name": ("STRING", {"default": "wanvideo_cached_videos.pt", "multiline": False}),
|
|
},
|
|
"optional": {
|
|
"images": (imageOrLatent, {"forceInput": True}),
|
|
"frame_count": ("INT", {"forceInput": True}),
|
|
"audio": ("AUDIO", {"forceInput": True}),
|
|
"video_info": ("VHS_VIDEOINFO", {"forceInput": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (imageOrLatent, "INT", "AUDIO", "VHS_VIDEOINFO")
|
|
RETURN_NAMES = ("images", "frame_count", "audio", "video_info")
|
|
FUNCTION = "cache_or_load_videos"
|
|
CATEGORY = "WanVideoWrapper_QQ/utils"
|
|
|
|
def cache_or_load_videos(self, cache_path, cache_name, images=None, frame_count=None, audio=None, video_info=None):
|
|
# Get the path to the ComfyUI-WanVideoWrapper_QQ directory
|
|
qq_wrapper_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
# Sanitize cache_path to prevent it from being an absolute path
|
|
# and ensure it's just a directory name.
|
|
sanitized_cache_path = os.path.basename(cache_path)
|
|
|
|
# The cache directory will be inside the node's directory.
|
|
full_cache_dir = os.path.join(qq_wrapper_path, sanitized_cache_path)
|
|
os.makedirs(full_cache_dir, exist_ok=True)
|
|
|
|
full_cache_file_path = os.path.join(full_cache_dir, cache_name)
|
|
|
|
if images is not None:
|
|
# Video data is provided, cache them and pass them through
|
|
print(f"WanVideoCacheVideos: Caching video data to {full_cache_file_path}")
|
|
video_data = {
|
|
"images": images,
|
|
"frame_count": frame_count,
|
|
"audio": audio,
|
|
"video_info": video_info
|
|
}
|
|
torch.save(video_data, full_cache_file_path)
|
|
return (images, frame_count, audio, video_info)
|
|
else:
|
|
# Video data is not provided, try to load from cache
|
|
if os.path.exists(full_cache_file_path):
|
|
print(f"WanVideoCacheVideos: Loading video data from {full_cache_file_path}")
|
|
loaded_video_data = torch.load(full_cache_file_path)
|
|
return (
|
|
loaded_video_data["images"],
|
|
loaded_video_data["frame_count"],
|
|
loaded_video_data["audio"],
|
|
loaded_video_data["video_info"]
|
|
)
|
|
else:
|
|
raise FileNotFoundError(f"WanVideoCacheVideos: No video data provided and no cached video found at {full_cache_file_path}")
|
|
|
|
# Update the node mappings to include the new class
|
|
NODE_CLASS_MAPPINGS = {
|
|
"WanVideoCacheSamples": WanVideoCacheSamples,
|
|
"WanVideoCacheVideos": WanVideoCacheVideos
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"WanVideoCacheSamples": "WanVideo Cache Samples",
|
|
"WanVideoCacheVideos": "WanVideo Cache Videos"
|
|
}
|