Files
olivv-cs-ComfyUI-FunPack/image_processing.py
T

608 lines
22 KiB
Python

import copy
import math
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
import comfy.clip_vision
import folder_paths
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 FunPackClipVisionOutputCombine:
CATEGORY = "FunPack"
RETURN_TYPES = ("CLIP_VISION_OUTPUT",)
RETURN_NAMES = ("clip_vision_output",)
FUNCTION = "combine"
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"clip_vision_output1": ("CLIP_VISION_OUTPUT",),
"method": (["mean", "median", "maximum", "minimum"], {"default": "mean"}),
},
"optional": {
"clip_vision_output2": ("CLIP_VISION_OUTPUT",),
"clip_vision_output3": ("CLIP_VISION_OUTPUT",),
"clip_vision_output4": ("CLIP_VISION_OUTPUT",),
},
}
def _keys(self, clip_output):
if isinstance(clip_output, dict):
return clip_output.keys()
return vars(clip_output).keys()
def _get(self, clip_output, key):
if isinstance(clip_output, dict):
return clip_output[key]
return getattr(clip_output, key)
def _set(self, clip_output, key, value):
if isinstance(clip_output, dict):
clip_output[key] = value
else:
setattr(clip_output, key, value)
def _copy_value(self, value):
if isinstance(value, torch.Tensor):
return value.clone()
return copy.deepcopy(value)
def _new_output(self, source):
if isinstance(source, dict):
return {}
return comfy.clip_vision.Output()
def _combine_tensors(self, key, tensors, method):
first = tensors[0]
first_shape = first.shape
for tensor in tensors[1:]:
if tensor.shape != first_shape:
raise ValueError(
f"Cannot combine CLIP Vision output field '{key}' with shapes "
f"{tuple(first_shape)} and {tuple(tensor.shape)}."
)
stacked = torch.stack(
[tensor.to(device=first.device, dtype=first.dtype) for tensor in tensors],
dim=0,
)
if method == "median":
return stacked.median(dim=0).values
if method == "maximum":
return stacked.max(dim=0).values
if method == "minimum":
return stacked.min(dim=0).values
return stacked.mean(dim=0)
@torch.no_grad()
def combine(self, clip_vision_output1, method, clip_vision_output2=None, clip_vision_output3=None, clip_vision_output4=None):
outputs = [
output
for output in (clip_vision_output1, clip_vision_output2, clip_vision_output3, clip_vision_output4)
if output is not None
]
combined = self._new_output(outputs[0])
for key in self._keys(outputs[0]):
values = [self._get(output, key) for output in outputs if key in self._keys(output)]
first_value = self._get(outputs[0], key)
if len(values) == len(outputs) and all(isinstance(value, torch.Tensor) for value in values):
self._set(combined, key, self._combine_tensors(key, values, method))
else:
self._set(combined, key, self._copy_value(first_value))
return (combined,)
class FunPackVideoStitch:
CATEGORY = "FunPack"
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("STITCHED",)
FUNCTION = "stitch"
INPUT_TYPES = lambda: {
"required": {
"blend_frames": ("INT", {"default": 8, "min": 0, "max": 64}),
"transition_type": (["linear", "ease_in", "ease_out", "ease_in_out", "cosine"], {"default": "linear"}),
},
"optional": {
"video1": ("IMAGE",),
"video2": ("IMAGE",),
"video3": ("IMAGE",),
"video4": ("IMAGE",),
"video5": ("IMAGE",),
"video6": ("IMAGE",),
"video7": ("IMAGE",),
"video8": ("IMAGE",),
}
}
def get_alpha(self, step_index, blend_frames, transition_type):
if blend_frames == 1:
return 0.5
t = step_index / (blend_frames - 1)
if transition_type == "ease_in":
return t * t
if transition_type == "ease_out":
return 1 - ((1 - t) * (1 - t))
if transition_type == "ease_in_out":
return t * t * (3 - (2 * t))
if transition_type == "cosine":
return 0.5 - (0.5 * math.cos(math.pi * t))
return t
def blend_batches(self, batch_a, batch_b, blend_frames, transition_type):
blended = []
for i in range(blend_frames):
alpha = self.get_alpha(i, blend_frames, transition_type)
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, transition_type, 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.")
if blend_frames == 0:
return (torch.cat(video_batches, dim=0),)
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.blend_batches(batch_a, batch_b, blend_frames, transition_type)
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)