olivv-cs-ComfyUI-FunPack/funpack.py

"""
olivv's FunPack v1.1.1

Changelog:

v1.1.2 - Added "instruct_from_pretrained"
v1.1.1 - Added DualCLIP Instruct loader node for experimenting with Instruct models in FramePack.
v1.1.0 - Changed interpolation logic.

Custom ComfyUI node created specifically to imitate FramePack's approach, but using WAN2.1 models.
It takes the last frame of a video, then multiplies it many times,
gradually increasing denoising to the end, somewhat similar to what FramePack does,
but without frame prediction.

"""
import os
import torch
from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM, AutoModelForCausalLM, AutoTokenizer, AutoConfig
from huggingface_hub import hf_hub_download, snapshot_download
from safetensors.torch import load_file, save_file
import glob
import comfy.model_management as mm
import comfy.sd as sd
import folder_paths
import nodes

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 = "WANPack/Video"

    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)

# Helper to list files in models/clip
def list_clip_files():
    model_dir = os.path.join(mm.model_path, "clip")
    return sorted(f for f in os.listdir(model_dir) if f.endswith(".safetensors"))

class FunPackCLIPLoader:
    @classmethod
    def INPUT_TYPES(s):
        base = nodes.DualCLIPLoader.INPUT_TYPES()
        return {
            'required': 
                {
                    'clip_model_name': (s.get_filename_list(),),
                    'text_encoder_model_name': (s.get_filename_list(),),
                    'llm_vision_model_name':(s.get_filename_list(),),
                    'type': base['required']['type'],
                    'encoder_pretrained_path': ("STRING", {"multiline": False, "default": "mlabonne/NeuralLlama-3-8B-Instruct-abliterated"}),
                    'vision_pretrained_path': ("STRING", {"multiline": False, "default": "qresearch/llama-3.1-8B-vision-378"}),
                    'encoder_from_pretrained': ("BOOLEAN", {"default": False, "tooltip": "Load Instruct model from pretrained_path"}),
                    'vision_from_pretrained': ("BOOLEAN", {"default": False, "tooltip": "Load LLM+vision model from pretrained_path"}),
                    'load_te': ("BOOLEAN", {"default": True, "tooltip": "If off, does not load separate model as text encoder, using only llm_vision_model_name"}),
                    'system_prompt': ("STRING", {
                        "multiline": True,
                        "default": "You are a creative assistant optimized for generating vivid, detailed descriptions for video generation."
                    })
                }
            }
    RETURN_TYPES = 'CLIP',
    RETURN_NAMES = 'clip',
    FUNCTION = "load"
    CATEGORY = "conditioning"
    
    @classmethod
    def get_filename_list(s):
        files = []
        files += folder_paths.get_filename_list('clip')
        return sorted(files)
    
    def load(self, clip_model_name, type, text_encoder_model_name, llm_vision_model_name, encoder_pretrained_path, vision_pretrained_path, system_prompt, encoder_from_pretrained=None, vision_from_pretrained=None, load_te=None):
        # Load CLIP model using ComfyUI
        clip_path = folder_paths.get_full_path('clip', clip_model_name)
        def get_clip_type(type):
            clip_type = getattr(sd.CLIPType, type.upper(), sd.CLIPType.HUNYUAN_VIDEO)
            print("Detected clip type:", clip_type)
            return clip_type
        
        # Load TE model from weights
        encoder_path = folder_paths.get_full_path('clip', text_encoder_model_name)
        config_source = encoder_pretrained_path
        
        pretrained_vision_local_path = snapshot_download(repo_id=vision_pretrained_path)
        pvlp_model = pretrained_vision_local_path + "/model.safetensors"
        
        print("Loading TE from pretrained is set to", encoder_from_pretrained)
        print("Loading LLM+vision from pretrained is set to", vision_from_pretrained)
        print("Loading custom TE is set to", load_te)

        # Load LLM with vision capabilities (expected llava-llama-3-8b_v1_1)
        if vision_from_pretrained == True:
            print("Loading LLM+vision from", vision_pretrained_path)
            if os.path.exists(pvlp_model):
                print("Model already saved in a single file. Loading from local path...")
                model_dir = snapshot_download(repo_id=vision_pretrained_path)
                vision_path = pretrained_vision_local_path + "/model.safetensors"
                print ("Loading from", vision_path)
            else:
                print("Local model does not exist. Loading, merging and saving it locally..")
                model = AutoModelForCausalLM.from_pretrained(vision_pretrained_path, trust_remote_code=True)
                model_dir = snapshot_download(repo_id=vision_pretrained_path)
                shard_paths = sorted(glob.glob(os.path.join(model_dir, "*.safetensors")))
                shard_paths = [p for p in shard_paths if "index" not in p]
                print("Shard files:", shard_paths)
                # Step 3: Load and combine shards
                combined_state_dict = {}
                for shard_path in shard_paths:
                    print(shard_path)
                    # Load shard to CPU to minimize memory usage
                    shard_state_dict = load_file(shard_path, device="cpu")
                    combined_state_dict.update(shard_state_dict)
                    # Free memory
                    del shard_state_dict
                
                output_safetensors = os.path.join(pretrained_vision_local_path, "model.safetensors")
                print("Model output path:", output_safetensors)
                save_file(combined_state_dict, output_safetensors)
                
                print("Shards successfully transformed into single model.")
                
                vision_path = pretrained_vision_local_path + "/model.safetensors"
            
        else:
            print("Loading LLM+vision from existing local safetensors file...")
            vision_path = folder_paths.get_full_path('clip', llm_vision_model_name)
        
        if load_te == True:
            try:
                if encoder_from_pretrained == False:
                    print("Loading custom text encoder from the path:", encoder_path)
                    model = LlamaForCausalLM.from_pretrained(config_source, trust_remote_code=True)
                    state_dict = load_file(encoder_path, device="cuda")
                    model.load_state_dict(state_dict, strict=False)
                    model.eval().to(torch.float16).requires_grad_(False)
                    print("Custom text encoder from safetensors file loaded successfully!")
                else:
                    print("Loading custom text encoder from the path:", encoder_pretrained_path)
                    model = AutoModelForCausalLM.from_pretrained(encoder_pretrained_path, trust_remote_code=True)
                    tokenizer = AutoTokenizer.from_pretrained(encoder_pretrained_path, trust_remote_code=True)
                    model.eval().to(torch.float16).requires_grad_(False)
                    print("Custom text encoder from transformers loaded successfully!")
            except Exception as e:
                print(f"Error loading custom text encoder: {e}")
                raise

        # Wrap it like a CLIP-compatible text encoder
        class InstructWrapper:
            def __init__(self):
                print("TEWrapper initialized!")
                self.system_prompt = system_prompt
                print("System prompt is set to:", self.system_prompt)
            def tokenize(self, text):
                messages = [
                    {"role": "system", "content": self.system_prompt},
                    {"role": "user", "content": text}
                ]
                return tokenizer.apply_chat_template(messages, add_generation_prompt=False, return_tensors="pt").to("cuda")

            def encode_from_tokens(self, tokens, return_pooled=True):
                with torch.no_grad():
                    output = model(input_ids=tokens, output_hidden_states=True)
                    hidden = output.hidden_states[-1]  # final layer
                    pooled = hidden[:, -1, :]  # use last token for pooled
                return hidden, pooled

        # Replace text encoder in CLIP model
        clip_model = sd.load_clip(ckpt_paths=[clip_path, vision_path], embedding_directory=None, clip_type=get_clip_type(type), model_options={})
        if load_te == True:
            clip_model.text = InstructWrapper()
        print("Current TE:", clip_model.text)  # Check if encoder is replaced 
        return (clip_model,)


NODE_CLASS_MAPPINGS = {
    "FunPackImg2LatentInterpolation": FunPackImg2LatentInterpolation,
    "FunPackCLIPLoader": FunPackCLIPLoader,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "FunPackImg2LatentInterpolation": "FunPack img2latent Interpolation",
    "FunPackCLIPLoader": "FunPack CLIP Loader"
}