957 lines
39 KiB
Python
957 lines
39 KiB
Python
import os
|
|
import torch
|
|
from transformers import LlamaConfig, LlamaForCausalLM, LlavaForConditionalGeneration, AutoModelForCausalLM, AutoTokenizer, AutoConfig
|
|
from huggingface_hub import hf_hub_download, snapshot_download
|
|
from safetensors.torch import load_file, save_file, safe_open
|
|
import glob
|
|
import comfy.model_management as mm
|
|
import comfy.sd as sd
|
|
import folder_paths
|
|
import nodes
|
|
import tempfile
|
|
import gc # Import garbage collector
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
from PIL import Image
|
|
import folder_paths
|
|
from comfy.utils import ProgressBar
|
|
import comfy.clip_vision
|
|
import math
|
|
import json
|
|
from comfy.utils import ProgressBar
|
|
|
|
# Constants from StoryMem
|
|
IMAGE_FACTOR = 28
|
|
VIDEO_MIN_PIXELS = 48 * IMAGE_FACTOR * IMAGE_FACTOR # 37,632
|
|
MIN_FRAME_SIMILARITY = 0.9
|
|
MAX_KEYFRAME_NUM = 3
|
|
ADAPTIVE_ALPHA = 0.01
|
|
HPSV3_QUALITY_THRESHOLD = 3.0
|
|
|
|
class FunPackStoryMemJSONConverterQwenOutput:
|
|
"""
|
|
FunPack StoryMem LoRA JSON Converter - supports both manual input and Qwen-generated output.
|
|
|
|
If 'use_qwen_output' is enabled:
|
|
- Takes 'qwen_output' (string) from Kijai's WanVideoPromptExtender/Qwen node
|
|
- Parses the full story JSON (as produced by StoryMem-style system prompt)
|
|
- Splits into up to 3 separate scene JSONs (one per output)
|
|
- Each output contains story_name, story_overview + exactly ONE scene
|
|
- Extra scenes beyond 3 are ignored
|
|
- Handles 'cut' as single bool by replicating it across all prompts in the scene
|
|
- Forces first cut to True per StoryMem guidelines
|
|
|
|
If 'use_qwen_output' is false → falls back to original manual scene inputs.
|
|
|
|
Expected Qwen output structure (example: black.json):
|
|
{
|
|
"story_name": str,
|
|
"story_overview": str,
|
|
"scenes": [
|
|
{
|
|
"scene_num": int,
|
|
"video_prompts": list[str],
|
|
"first_frame_prompt": list[str],
|
|
"cut": list[bool] or bool # Now supports single bool
|
|
},
|
|
...
|
|
]
|
|
}
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"use_qwen_output": ("BOOLEAN", {
|
|
"default": False,
|
|
"label_on": "Use Qwen Output (ignores manual scenes below)",
|
|
"label_off": "Use Manual Scene Inputs",
|
|
"tooltip": "Toggle between Qwen-generated full story JSON or classic manual per-scene setup"
|
|
}),
|
|
"qwen_output": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "Connect STRING output from Kijai's WanVideoPromptExtender / Qwen enhancer here.\nDrag wire onto the left side of this box to convert to input socket."
|
|
}),
|
|
"story_name": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Story title (used in all output JSONs)"
|
|
}),
|
|
"story_overview": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "Overall story summary (shared across all scenes)"
|
|
}),
|
|
# Scene 1
|
|
"scene1_video_prompts": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "One prompt per line (each = 5-sec clip)"
|
|
}),
|
|
"scene1_first_frame_prompts": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "First frame description per prompt (must match count)"
|
|
}),
|
|
"scene1_cut_values": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Comma-separated booleans: true,false,true,... (must match prompt count; first usually true)"
|
|
}),
|
|
# Scene 2
|
|
"scene2_video_prompts": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "One prompt per line (each = 5-sec clip)"
|
|
}),
|
|
"scene2_first_frame_prompts": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "First frame description per prompt (must match count)"
|
|
}),
|
|
"scene2_cut_values": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Comma-separated booleans: true,false,true,... (must match prompt count; first usually true)"
|
|
}),
|
|
# Scene 3
|
|
"scene3_video_prompts": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "One prompt per line (each = 5-sec clip)"
|
|
}),
|
|
"scene3_first_frame_prompts": ("STRING", {
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "First frame description per prompt (must match count)"
|
|
}),
|
|
"scene3_cut_values": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Comma-separated booleans: true,false,true,... (must match prompt count; first usually true)"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"qwen_mode_notice": ("STRING", {
|
|
"multiline": True,
|
|
"default": "!!! QWEN MODE ACTIVE !!!\nAll manual inputs below are completely ignored\nConnect the Qwen JSON string above.",
|
|
"tooltip": "Reminder - only visible/meaningful when use_qwen_output = True"
|
|
})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("json_scene_1", "json_scene_2", "json_scene_3")
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "convert"
|
|
CATEGORY = "FunPack"
|
|
|
|
def convert(self,
|
|
use_qwen_output,
|
|
qwen_output,
|
|
story_name, story_overview,
|
|
scene1_video_prompts, scene1_first_frame_prompts, scene1_cut_values,
|
|
scene2_video_prompts, scene2_first_frame_prompts, scene2_cut_values,
|
|
scene3_video_prompts, scene3_first_frame_prompts, scene3_cut_values,
|
|
qwen_mode_notice=""):
|
|
|
|
if use_qwen_output:
|
|
if not qwen_output.strip():
|
|
raise ValueError("Qwen mode enabled but 'qwen_output' is empty! Connect Kijai's prompt enhancer output.")
|
|
|
|
pbar = ProgressBar(2)
|
|
|
|
try:
|
|
full_data = json.loads(qwen_output)
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"Failed to parse Qwen output as JSON:\n{str(e)}\n\nMake sure system prompt instructs Qwen to return **only** valid JSON.")
|
|
|
|
story_name = full_data.get("story_name", "Generated Story")
|
|
story_overview = full_data.get("story_overview", "Generated overview")
|
|
scenes = full_data.get("scenes", [])
|
|
|
|
if not isinstance(scenes, list) or not scenes:
|
|
raise ValueError("Qwen JSON must contain non-empty 'scenes' list.")
|
|
|
|
outputs = ["", "", ""]
|
|
|
|
for i, scene in enumerate(scenes[:3]):
|
|
if not isinstance(scene, dict):
|
|
continue
|
|
|
|
scene_num = scene.get("scene_num", i+1)
|
|
video_prompts = scene.get("video_prompts", [])
|
|
first_frame_prompt = scene.get("first_frame_prompt", []) # Note: renamed from first_frame_prompts for consistency
|
|
cut = scene.get("cut", [])
|
|
|
|
if not video_prompts:
|
|
continue
|
|
|
|
num_prompts = len(video_prompts)
|
|
|
|
if len(first_frame_prompt) != num_prompts:
|
|
raise ValueError(
|
|
f"Scene {scene_num}: video_prompts ({num_prompts}) ≠ "
|
|
f"first_frame_prompt ({len(first_frame_prompt)})"
|
|
)
|
|
|
|
# Handle cut: if single bool, replicate across all prompts
|
|
if isinstance(cut, bool):
|
|
cut = [cut] * num_prompts
|
|
elif not isinstance(cut, list):
|
|
raise ValueError(f"Scene {scene_num}: 'cut' must be bool or list of bools, got {type(cut)}")
|
|
|
|
# Check length after potential conversion
|
|
if len(cut) != num_prompts:
|
|
raise ValueError(
|
|
f"Scene {scene_num}: 'cut' length ({len(cut)}) ≠ "
|
|
f"prompt count ({num_prompts})"
|
|
)
|
|
|
|
# Force first cut to True (per StoryMem: "The first prompt in the story must always have True")
|
|
# Note: This applies per scene, but if the whole story is multi-scene, the first of scene 1 is the global first
|
|
if num_prompts > 0:
|
|
cut[0] = True
|
|
|
|
single_scene_data = {
|
|
"scene_num": scene_num,
|
|
"video_prompts": video_prompts,
|
|
"first_frame_prompt": first_frame_prompt,
|
|
"cut": cut
|
|
}
|
|
|
|
json_data = {
|
|
"story_name": story_name,
|
|
"story_overview": story_overview,
|
|
"scenes": [single_scene_data]
|
|
}
|
|
|
|
outputs[i] = json.dumps(json_data, indent=2, ensure_ascii=False)
|
|
|
|
pbar.update(2)
|
|
return tuple(outputs)
|
|
|
|
else:
|
|
# ── Original manual mode ───────────────────────────────────────
|
|
scenes_input = [
|
|
(scene1_video_prompts, scene1_first_frame_prompts, scene1_cut_values, 1),
|
|
(scene2_video_prompts, scene2_first_frame_prompts, scene2_cut_values, 2),
|
|
(scene3_video_prompts, scene3_first_frame_prompts, scene3_cut_values, 3),
|
|
]
|
|
|
|
outputs = ["", "", ""]
|
|
|
|
for i, (video_text, first_text, cuts_text, scene_num) in enumerate(scenes_input):
|
|
if not video_text.strip():
|
|
continue
|
|
|
|
video_prompts = [p.strip() for p in video_text.split("\n") if p.strip()]
|
|
if not video_prompts:
|
|
continue
|
|
|
|
first_prompts = [f.strip() for f in first_text.split("\n") if f.strip()]
|
|
|
|
if len(first_prompts) != len(video_prompts):
|
|
raise ValueError(
|
|
f"Manual scene {scene_num}: video prompts ({len(video_prompts)}) ≠ "
|
|
f"first frames ({len(first_prompts)})"
|
|
)
|
|
|
|
num_prompts = len(video_prompts)
|
|
|
|
if cuts_text.strip():
|
|
cuts_str = [c.strip().lower() for c in cuts_text.split(",") if c.strip()]
|
|
cuts = []
|
|
for c in cuts_str:
|
|
if c in ("true", "t", "1", "yes"):
|
|
cuts.append(True)
|
|
elif c in ("false", "f", "0", "no"):
|
|
cuts.append(False)
|
|
else:
|
|
raise ValueError(f"Invalid cut value in manual scene {scene_num}: '{c}'")
|
|
else:
|
|
cuts = [False] * num_prompts
|
|
|
|
if len(cuts) != num_prompts:
|
|
raise ValueError(
|
|
f"Manual scene {scene_num}: cut count ({len(cuts)}) ≠ prompt count"
|
|
)
|
|
|
|
# Force first cut to True for consistency
|
|
if num_prompts > 0:
|
|
cuts[0] = True
|
|
|
|
scene = {
|
|
"scene_num": scene_num,
|
|
"video_prompts": video_prompts,
|
|
"first_frame_prompt": first_prompts,
|
|
"cut": cuts
|
|
}
|
|
|
|
full_json = {
|
|
"story_name": story_name.strip(),
|
|
"story_overview": story_overview.strip(),
|
|
"scenes": [scene]
|
|
}
|
|
|
|
outputs[i] = json.dumps(full_json, indent=2, ensure_ascii=False)
|
|
|
|
return tuple(outputs)
|
|
class FunPackImg2LatentInterpolation:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"frame_count": ("INT", {"default": 25, "min": 1, "max": 125, "step": 4}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("img_batch_for_encode", "img_for_start_image")
|
|
FUNCTION = "process"
|
|
CATEGORY = "FunPack"
|
|
|
|
def process(self, images, frame_count):
|
|
device = images.device
|
|
total_input_frames = images.shape[0]
|
|
|
|
# 1. Take the last frame as our starting point
|
|
last_frame = images[-1:]
|
|
|
|
# 2. Generate smooth transition from this frame
|
|
interpolated = []
|
|
|
|
# First frame is exact copy of last input frame (for perfect continuity)
|
|
interpolated.append(last_frame.clone())
|
|
|
|
# Generate remaining frames with increasing denoise
|
|
for i in range(1, frame_count):
|
|
# Calculate denoise strength (0 at start, 1 at end)
|
|
denoise_strength = i / (frame_count - 1)
|
|
|
|
# Apply denoising
|
|
noise = torch.randn_like(last_frame)
|
|
blended = (1 - denoise_strength) * last_frame + denoise_strength * noise
|
|
interpolated.append(blended)
|
|
|
|
# Convert to tensor
|
|
output = torch.cat(interpolated, dim=0)
|
|
|
|
# Preview is first frame (same as input's last frame)
|
|
preview = interpolated[0].clone()
|
|
|
|
return (output, preview)
|
|
|
|
# The standalone FunPackPromptEnhancer remains completely unchanged as requested
|
|
class FunPackPromptEnhancer:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"user_prompt": ("STRING", {"multiline": True, "default": "A photo of a [subject] in a [setting]. [action]."}),
|
|
"system_prompt": ("STRING", {
|
|
"multiline": True,
|
|
"default": "<|start_header_id|>system<|end_header_id|>\n\nYou are a creative AI assistant tasked with describing videos.\n\nDescribe the video by detailing the following aspects:\n1. The main content and theme of the video.\n2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects.\n3. Actions, events, behaviors temporal relationships, physical movement changes of the objects.\n4. background environment, light, style and atmosphere.\n5. camera angles, movements, and transitions used in the video:<|eot_id|>"
|
|
}),
|
|
"model_path_type": (["Local Safetensors", "HuggingFace Pretrained"],),
|
|
"model_path": ("STRING", {"multiline": False, "default": "mlabonne/NeuralLlama-3-8B-Instruct-abliterated"}),
|
|
"llm_safetensors_file": (folder_paths.get_filename_list('clip'),),
|
|
"top_p": ("FLOAT", {"min": 0.0, "max": 1.0, "step": 0.05, "default": 0.75}),
|
|
"top_k": ("INT", {"min": 0, "max": 1000, "step": 1, "default": 40}),
|
|
"temperature": ("FLOAT", {"min": 0.0, "max": 1.0, "step": 0.01, "default": 0.6}),
|
|
"max_new_tokens": ("INT", {"min": 1, "max": 2048, "step": 64, "default": 1024}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("enhanced_prompt",)
|
|
FUNCTION = "enhance_prompt"
|
|
CATEGORY = "FunPack"
|
|
|
|
def enhance_prompt(self, user_prompt, system_prompt, model_path_type, model_path, llm_safetensors_file, top_p, top_k, temperature, max_new_tokens):
|
|
llm_model = None
|
|
llm_tokenizer = None
|
|
llm_model_device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
print(f"[FunPackPromptEnhancer] Starting prompt enhancement...")
|
|
|
|
try:
|
|
if model_path_type == "HuggingFace Pretrained":
|
|
print(f"[FunPackPromptEnhancer] Loading LLM from HuggingFace pretrained: {model_path}")
|
|
llm_tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
|
llm_model = AutoModelForCausalLM.from_pretrained(model_path, ignore_mismatched_sizes=True, trust_remote_code=True)
|
|
elif model_path_type == "Local Safetensors":
|
|
print(f"[FunPackPromptEnhancer] Loading LLM from local safetensors file: {llm_safetensors_file}")
|
|
full_safetensors_path = folder_paths.get_full_path('clip', llm_safetensors_file)
|
|
|
|
llm_tokenizer = AutoTokenizer.from_pretrained("xtuner/llava-llama-3-8b-v1_1-transformers", trust_remote_code=True)
|
|
|
|
config = AutoConfig.from_pretrained("xtuner/llava-llama-3-8b-v1_1-transformers", trust_remote_code=True)
|
|
model_base = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
|
|
|
|
state_dict = load_file(full_safetensors_path, device="cpu")
|
|
model_base.load_state_dict(state_dict, strict=False)
|
|
llm_model = model_base
|
|
|
|
llm_model = llm_model.eval().to(torch.bfloat16 if llm_model_device == "cuda" else torch.float32).to(llm_model_device).requires_grad_(False)
|
|
print(f"[FunPackPromptEnhancer] LLM model loaded successfully to {llm_model_device}!")
|
|
|
|
messages = [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_prompt}
|
|
]
|
|
|
|
llm_tokens = llm_tokenizer.apply_chat_template(
|
|
messages,
|
|
add_generation_prompt=True,
|
|
return_tensors="pt",
|
|
tokenize=True
|
|
).to(llm_model_device)
|
|
|
|
print("[FunPackPromptEnhancer] Generating enhanced prompt...")
|
|
with torch.no_grad():
|
|
generated_ids = llm_model.generate(
|
|
input_ids=llm_tokens,
|
|
do_sample=True,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
temperature=temperature,
|
|
max_new_tokens=max_new_tokens,
|
|
pad_token_id=llm_tokenizer.pad_token_id
|
|
)
|
|
|
|
output_text = llm_tokenizer.decode(generated_ids[0][llm_tokens.shape[1]:], skip_special_tokens=True)
|
|
print(f"[FunPackPromptEnhancer] Enhanced prompt generated: {output_text}")
|
|
|
|
return (output_text,)
|
|
|
|
except Exception as e:
|
|
print(f"[FunPackPromptEnhancer] Error during prompt enhancement: {e}")
|
|
raise
|
|
|
|
finally:
|
|
if llm_model is not None:
|
|
del llm_model
|
|
llm_model = None
|
|
if llm_tokenizer is not None:
|
|
del llm_tokenizer
|
|
llm_tokenizer = None
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
print("[FunPackPromptEnhancer] LLM model and tokenizer unloaded and memory cleared.")
|
|
|
|
|
|
class FunPackVideoStitch:
|
|
CATEGORY = "FunPack"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("STITCHED",)
|
|
FUNCTION = "stitch"
|
|
INPUT_TYPES = lambda: {
|
|
"required": {
|
|
"blend_frames": ("INT", {"default": 8, "min": 1, "max": 64}),
|
|
},
|
|
"optional": {
|
|
"video1": ("IMAGE",),
|
|
"video2": ("IMAGE",),
|
|
"video3": ("IMAGE",),
|
|
"video4": ("IMAGE",),
|
|
"video5": ("IMAGE",),
|
|
"video6": ("IMAGE",),
|
|
"video7": ("IMAGE",),
|
|
"video8": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
def linear_blend(self, batch_a, batch_b, blend_frames):
|
|
if blend_frames == 1:
|
|
blended_frame = 0.5 * batch_a[-1] + 0.5 * batch_b[0]
|
|
return blended_frame.unsqueeze(0)
|
|
|
|
blended = []
|
|
for i in range(blend_frames):
|
|
alpha = i / (blend_frames - 1)
|
|
blended_frame = (1 - alpha) * batch_a[-blend_frames + i] + alpha * batch_b[i]
|
|
blended.append(blended_frame.unsqueeze(0))
|
|
return torch.cat(blended, dim=0)
|
|
|
|
def stitch(self, blend_frames, video1=None, video2=None, video3=None, video4=None, video5=None, video6=None, video7=None, video8=None):
|
|
input_videos = [video1, video2, video3, video4, video5, video6, video7, video8]
|
|
video_batches = [v for v in input_videos if v is not None]
|
|
|
|
if len(video_batches) < 2:
|
|
raise ValueError("VideoStitch requires at least 2 connected video inputs.")
|
|
|
|
output_frames = []
|
|
|
|
for i in range(len(video_batches) - 1):
|
|
batch_a = video_batches[i]
|
|
batch_b = video_batches[i + 1]
|
|
|
|
if batch_a.shape[0] < blend_frames or batch_b.shape[0] < blend_frames:
|
|
raise ValueError(f"Each video batch must have at least {blend_frames} frames.")
|
|
|
|
stable_a = batch_a[:-blend_frames]
|
|
stable_b = batch_b[blend_frames:]
|
|
transition = self.linear_blend(batch_a, batch_b, blend_frames)
|
|
|
|
if i == 0:
|
|
output_frames.append(stable_a)
|
|
output_frames.append(transition)
|
|
output_frames.append(stable_b if i == len(video_batches) - 2 else batch_b[blend_frames:-blend_frames])
|
|
|
|
final_video = torch.cat(output_frames, dim=0)
|
|
return (final_video,)
|
|
|
|
class FunPackContinueVideo:
|
|
CATEGORY = "FunPack"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("CONTINUED",)
|
|
FUNCTION = "continue_video"
|
|
INPUT_TYPES = lambda: {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"frame_count": ("INT", {"default": 1, "min": 1, "max": 9999}),
|
|
}
|
|
}
|
|
|
|
def continue_video(self, images, frame_count):
|
|
total_frames = images.shape[0]
|
|
|
|
if frame_count > total_frames:
|
|
raise ValueError(f"Cannot extract {frame_count} frames from video with only {total_frames} frames.")
|
|
|
|
continued = images[-frame_count:]
|
|
return (continued,)
|
|
|
|
class FunPackStoryMemKeyframeExtractor:
|
|
"""
|
|
Extracts keyframes from video frames using:
|
|
1. HPSv3 for quality assessment (optional)
|
|
2. CLIP Vision for frame similarity
|
|
3. Adaptive threshold to limit keyframe count
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"frames": ("IMAGE",), # ComfyUI IMAGE format [B, H, W, C]
|
|
"clip_vision": (folder_paths.get_filename_list("clip_vision"),),
|
|
"max_keyframes": ("INT", {
|
|
"default": 3,
|
|
"min": 1,
|
|
"max": 20,
|
|
"step": 1,
|
|
"tooltip": "Maximum number of keyframes to extract"
|
|
}),
|
|
"similarity_threshold": ("FLOAT", {
|
|
"default": 0.9,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"tooltip": "CLIP similarity threshold (lower = more keyframes)"
|
|
}),
|
|
"use_quality_filter": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Use HPSv3 to filter low-quality frames (requires hpsv3 package)"
|
|
}),
|
|
"quality_threshold": ("FLOAT", {
|
|
"default": 3.0,
|
|
"min": 0.0,
|
|
"max": 10.0,
|
|
"step": 0.1,
|
|
"tooltip": "HPSv3 quality threshold (higher = stricter)"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"memory_frames": ("IMAGE", {
|
|
"tooltip": "Previous keyframes to compare against (avoid duplicates)"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT",)
|
|
RETURN_NAMES = ("keyframes", "keyframe_count",)
|
|
FUNCTION = "extract_keyframes"
|
|
CATEGORY = "FunPack"
|
|
DESCRIPTION = "Extract keyframes using CLIP similarity + HPSv3 quality (StoryMem algorithm)"
|
|
|
|
def __init__(self):
|
|
self.quality_model = None
|
|
|
|
def load_clip_model(self, clip_vision_name):
|
|
"""Load CLIP Vision model from ComfyUI models/clip_vision folder"""
|
|
clip_path = folder_paths.get_full_path("clip_vision", clip_vision_name)
|
|
clip_vision = comfy.clip_vision.load(clip_path)
|
|
return clip_vision
|
|
|
|
def load_quality_model(self):
|
|
"""Load HPSv3 quality assessment model"""
|
|
if self.quality_model is not None:
|
|
return
|
|
|
|
try:
|
|
from hpsv3 import HPSv3RewardInferencer
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
self.quality_model = HPSv3RewardInferencer(device=device)
|
|
except ImportError:
|
|
print("WARNING: HPSv3 not installed. Install with: pip install hpsv3")
|
|
print("Quality filtering will be disabled.")
|
|
self.quality_model = None
|
|
|
|
def smart_resize(self, height: int, width: int) -> tuple:
|
|
"""Resize frame to efficient size for processing"""
|
|
factor = IMAGE_FACTOR
|
|
min_pixels = VIDEO_MIN_PIXELS
|
|
max_pixels = 256 * IMAGE_FACTOR * IMAGE_FACTOR
|
|
|
|
h_bar = max(factor, round(height / factor) * factor)
|
|
w_bar = max(factor, round(width / factor) * factor)
|
|
|
|
if h_bar * w_bar > max_pixels:
|
|
beta = math.sqrt((height * width) / max_pixels)
|
|
h_bar = math.floor(height / beta / factor) * factor
|
|
w_bar = math.floor(width / beta / factor) * factor
|
|
elif h_bar * w_bar < min_pixels:
|
|
beta = math.sqrt(min_pixels / (height * width))
|
|
h_bar = math.ceil(height * beta / factor) * factor
|
|
w_bar = math.ceil(width * beta / factor) * factor
|
|
|
|
return max(h_bar, factor), max(w_bar, factor)
|
|
|
|
def clip_preprocess(self, frame_chw: torch.Tensor, clip_vision) -> torch.Tensor:
|
|
"""Preprocess frame for CLIP Vision model"""
|
|
# ComfyUI CLIP Vision expects [B, H, W, C] format in range [0, 1]
|
|
# Convert from [C, H, W] to [1, H, W, C]
|
|
frame = frame_chw.permute(1, 2, 0).unsqueeze(0)
|
|
|
|
# Ensure [0, 1] range
|
|
if not torch.is_floating_point(frame):
|
|
frame = frame.float()
|
|
if frame.max() > 1.5:
|
|
frame = frame / 255.0
|
|
frame = frame.clamp(0.0, 1.0)
|
|
|
|
return frame
|
|
|
|
@torch.no_grad()
|
|
def get_clip_similarity(self, frame1: torch.Tensor, frame2: torch.Tensor, clip_vision) -> float:
|
|
|
|
# Preprocess frames to [1, H, W, C] format
|
|
x1 = self.clip_preprocess(frame1, clip_vision)
|
|
x2 = self.clip_preprocess(frame2, clip_vision)
|
|
|
|
# Get CLIP Vision embeddings
|
|
z1_raw = clip_vision.encode_image(x1)
|
|
z2_raw = clip_vision.encode_image(x2)
|
|
|
|
# Extract the actual embedding tensor from various possible return formats
|
|
def extract_embedding(output):
|
|
# Case 1: Direct tensor (older/basic models)
|
|
if isinstance(output, torch.Tensor):
|
|
return output
|
|
|
|
# Case 2: ComfyUI's custom Output wrapper (common with projection models)
|
|
if isinstance(output, comfy.clip_vision.Output): # Import at top if needed: import comfy.clip_vision
|
|
if hasattr(output, 'image_embeds'):
|
|
return output.image_embeds
|
|
elif hasattr(output, 'pooled_output'):
|
|
return output.pooled_output
|
|
# Fallback: treat like dict
|
|
try:
|
|
return output['image_embeds']
|
|
except:
|
|
pass
|
|
|
|
# Case 3: Dictionary (some models)
|
|
if isinstance(output, dict):
|
|
if 'image_embeds' in output:
|
|
return output['image_embeds']
|
|
if 'pooled_output' in output:
|
|
return output['pooled_output']
|
|
if 'last_hidden_state' in output:
|
|
return output['last_hidden_state'][:, 0] # CLS token if sequence
|
|
# Fallback: first tensor value
|
|
for v in output.values():
|
|
if isinstance(v, torch.Tensor) and v.ndim >= 2:
|
|
return v
|
|
|
|
# Case 4: Tuple (rare here, but safe)
|
|
if isinstance(output, (tuple, list)) and len(output) == 1:
|
|
return extract_embedding(output[0])
|
|
|
|
raise TypeError(f"Unexpected output from encode_image: {type(output)}. "
|
|
"Supported: tensor, dict, or comfy.clip_vision.Output with 'image_embeds'.")
|
|
|
|
z1 = extract_embedding(z1_raw)
|
|
z2 = extract_embedding(z2_raw)
|
|
|
|
# Final check
|
|
if not (isinstance(z1, torch.Tensor) and isinstance(z2, torch.Tensor)):
|
|
raise RuntimeError(f"Failed to extract tensor embeddings: {type(z1)}, {type(z2)}")
|
|
|
|
# Normalize and compute cosine similarity
|
|
z1 = F.normalize(z1, dim=-1)
|
|
z2 = F.normalize(z2, dim=-1)
|
|
similarity = (z1 * z2).sum(dim=-1).item()
|
|
|
|
return similarity
|
|
|
|
def is_low_quality(self, frame: torch.Tensor, threshold: float) -> bool:
|
|
"""Check if frame quality is below threshold using HPSv3"""
|
|
if self.quality_model is None:
|
|
return False # Skip quality check if model not available
|
|
|
|
# Convert to PIL Image
|
|
frame_np = frame.permute(1, 2, 0).cpu().numpy()
|
|
frame_np = (frame_np * 255).astype(np.uint8).clip(0, 255)
|
|
pil_image = Image.fromarray(frame_np)
|
|
|
|
# Get quality score
|
|
try:
|
|
rewards = self.quality_model.reward(image_paths=[pil_image], prompts=[""])
|
|
score = rewards[0][0].item()
|
|
return score < threshold
|
|
except Exception as e:
|
|
print(f"Quality check failed: {e}")
|
|
return False
|
|
|
|
def extract_keyframe_indices(self, frames: torch.Tensor, threshold: float,
|
|
quality_threshold: float, use_quality: bool, clip_vision) -> list:
|
|
"""
|
|
Extract keyframe indices using CLIP similarity and quality filtering
|
|
|
|
Args:
|
|
frames: [N, C, H, W] tensor
|
|
threshold: CLIP similarity threshold
|
|
quality_threshold: HPSv3 quality threshold
|
|
use_quality: Whether to use quality filtering
|
|
clip_vision: ComfyUI CLIP Vision model
|
|
|
|
Returns:
|
|
List of keyframe indices
|
|
"""
|
|
num_frames, _, height, width = frames.shape
|
|
|
|
# Resize frames for efficient processing
|
|
resized_h, resized_w = self.smart_resize(height, width)
|
|
resized_frames = F.interpolate(
|
|
frames,
|
|
size=(resized_h, resized_w),
|
|
mode="bilinear",
|
|
align_corners=False
|
|
).float()
|
|
|
|
# Load quality model if needed
|
|
if use_quality:
|
|
self.load_quality_model()
|
|
|
|
# Find first high-quality frame
|
|
first_idx = 0
|
|
if use_quality and self.quality_model is not None:
|
|
while first_idx < num_frames:
|
|
if not self.is_low_quality(resized_frames[first_idx], quality_threshold):
|
|
break
|
|
first_idx += 1
|
|
|
|
if first_idx >= num_frames:
|
|
return [] # No high-quality frames found
|
|
|
|
# Initialize keyframes
|
|
keyframe_indices = [first_idx]
|
|
last_keyframe = resized_frames[first_idx]
|
|
|
|
# Iterate through remaining frames
|
|
pbar = ProgressBar(num_frames - first_idx - 1)
|
|
for i in range(first_idx + 1, num_frames):
|
|
current_frame = resized_frames[i]
|
|
|
|
# Calculate similarity with last keyframe
|
|
similarity = self.get_clip_similarity(last_keyframe, current_frame, clip_vision)
|
|
|
|
# Check if frame is different enough and high quality
|
|
is_different = similarity < threshold
|
|
is_quality = True
|
|
if use_quality and self.quality_model is not None:
|
|
is_quality = not self.is_low_quality(current_frame, quality_threshold)
|
|
|
|
if is_different and is_quality:
|
|
keyframe_indices.append(i)
|
|
last_keyframe = current_frame
|
|
|
|
pbar.update(1)
|
|
|
|
return keyframe_indices
|
|
|
|
def check_memory_duplicates(self, keyframes: torch.Tensor,
|
|
memory_frames: torch.Tensor,
|
|
clip_vision,
|
|
threshold: float = 0.9) -> list:
|
|
"""
|
|
Filter out keyframes that are too similar to memory frames
|
|
|
|
Returns:
|
|
List of boolean flags (True = keep, False = duplicate)
|
|
"""
|
|
keep_flags = []
|
|
|
|
for keyframe in keyframes:
|
|
is_duplicate = False
|
|
for memory_frame in memory_frames:
|
|
similarity = self.get_clip_similarity(keyframe, memory_frame, clip_vision)
|
|
if similarity > threshold:
|
|
is_duplicate = True
|
|
break
|
|
keep_flags.append(not is_duplicate)
|
|
|
|
return keep_flags
|
|
|
|
def extract_keyframes(self, frames, clip_vision, max_keyframes, similarity_threshold,
|
|
use_quality_filter, quality_threshold, memory_frames=None):
|
|
"""
|
|
Main extraction function
|
|
|
|
Args:
|
|
frames: ComfyUI IMAGE format [B, H, W, C] in range [0, 1]
|
|
clip_vision: CLIP Vision model name from dropdown
|
|
max_keyframes: Maximum number of keyframes
|
|
similarity_threshold: Initial CLIP similarity threshold
|
|
use_quality_filter: Whether to use HPSv3 filtering
|
|
quality_threshold: HPSv3 threshold
|
|
memory_frames: Optional previous keyframes to avoid duplicates
|
|
|
|
Returns:
|
|
(keyframes, keyframe_count)
|
|
"""
|
|
# Load CLIP Vision model from ComfyUI models folder
|
|
clip_vision_model = self.load_clip_model(clip_vision)
|
|
|
|
# Convert ComfyUI format [B, H, W, C] to PyTorch [B, C, H, W]
|
|
frames_tensor = frames.permute(0, 3, 1, 2).contiguous()
|
|
|
|
# Adaptive threshold loop
|
|
threshold = similarity_threshold
|
|
while True:
|
|
keyframe_indices = self.extract_keyframe_indices(
|
|
frames_tensor,
|
|
threshold,
|
|
quality_threshold,
|
|
use_quality_filter,
|
|
clip_vision_model
|
|
)
|
|
|
|
# Check if we have too many keyframes
|
|
if len(keyframe_indices) <= max_keyframes:
|
|
break
|
|
|
|
# Increase threshold to get fewer keyframes
|
|
threshold -= ADAPTIVE_ALPHA
|
|
|
|
# Safety check
|
|
if threshold < 0.5:
|
|
# Take first N keyframes
|
|
keyframe_indices = keyframe_indices[:max_keyframes]
|
|
break
|
|
|
|
print(f"Extracted {len(keyframe_indices)} keyframes at threshold {threshold:.3f}")
|
|
|
|
# Extract keyframes
|
|
if len(keyframe_indices) == 0:
|
|
# Return first frame as fallback
|
|
keyframes_out = frames[:1]
|
|
return (keyframes_out, 1)
|
|
|
|
keyframes_tensor = frames_tensor[keyframe_indices]
|
|
|
|
# Check against memory frames to avoid duplicates
|
|
if memory_frames is not None:
|
|
memory_tensor = memory_frames.permute(0, 3, 1, 2).contiguous()
|
|
keep_flags = self.check_memory_duplicates(
|
|
keyframes_tensor,
|
|
memory_tensor,
|
|
clip_vision_model,
|
|
threshold=MIN_FRAME_SIMILARITY
|
|
)
|
|
|
|
# Filter keyframes
|
|
kept_indices = [i for i, keep in enumerate(keep_flags) if keep]
|
|
if len(kept_indices) > 0:
|
|
keyframes_tensor = keyframes_tensor[kept_indices]
|
|
else:
|
|
# Keep at least one keyframe
|
|
keyframes_tensor = keyframes_tensor[:1]
|
|
|
|
# Convert back to ComfyUI format [B, H, W, C]
|
|
keyframes_out = keyframes_tensor.permute(0, 2, 3, 1).contiguous()
|
|
|
|
return (keyframes_out, keyframes_out.shape[0])
|
|
|
|
|
|
class FunPackStoryMemLastFrameExtractor:
|
|
"""Extract last frame and last N frames for MI2V/MM2V continuity"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"frames": ("IMAGE",),
|
|
"n_frames": ("INT", {
|
|
"default": 5,
|
|
"min": 1,
|
|
"max": 20,
|
|
"step": 1,
|
|
"tooltip": "Number of frames to extract from end (for MM2V)"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE",)
|
|
RETURN_NAMES = ("last_frame", "motion_frames",)
|
|
FUNCTION = "extract"
|
|
CATEGORY = "FunPack"
|
|
DESCRIPTION = "Extract last frame and last N frames for shot continuity (MI2V/MM2V)"
|
|
|
|
def extract(self, frames, n_frames):
|
|
"""
|
|
Extract last frame and last N frames
|
|
|
|
Returns:
|
|
(last_frame [1, H, W, C], motion_frames [N, H, W, C])
|
|
"""
|
|
last_frame = frames[-1:]
|
|
motion_frames = frames[-n_frames:]
|
|
|
|
return (last_frame, motion_frames)
|
|
|
|
|
|
# Update NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS
|
|
NODE_CLASS_MAPPINGS = {
|
|
"FunPackStoryMemJSONConverter": FunPackStoryMemJSONConverter,
|
|
"FunPackStoryMemKeyframeExtractor": FunPackStoryMemKeyframeExtractor,
|
|
"FunPackStoryMemLastFrameExtractor": FunPackStoryMemLastFrameExtractor,
|
|
"FunPackImg2LatentInterpolation": FunPackImg2LatentInterpolation,
|
|
"FunPackPromptEnhancer": FunPackPromptEnhancer,
|
|
"FunPackVideoStitch": FunPackVideoStitch,
|
|
"FunPackContinueVideo": FunPackContinueVideo,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"FunPackStoryMemJSONConverter": "FunPack StoryMem JSON Converter",
|
|
"FunPackStoryMemKeyframeExtractor": "FunPack StoryMem Keyframe Extractor",
|
|
"FunPackStoryMemLastFrameExtractor": "FunPack StoryMem Last Frame Extractor",
|
|
"FunPackImg2LatentInterpolation": "FunPack Img2Latent Interpolation",
|
|
"FunPackPromptEnhancer": "FunPack Prompt Enhancer (Standalone)",
|
|
"FunPackVideoStitch": "FunPack Video Stitch",
|
|
"FunPackContinueVideo": "FunPack Continue Video"
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|