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 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)