diff --git a/funpack.py b/funpack.py index 148f408e..0d1ec028 100644 --- a/funpack.py +++ b/funpack.py @@ -608,7 +608,7 @@ class FunPackStoryWriter: "max_new_tokens": ("INT", {"min": 64, "max": 4096, "step": 64, "default": 512}), "repetition_penalty": ("FLOAT", {"min": 0.0, "max": 3.0, "step": 0.01, "default": 1.0}), "mode": (["Sequences from story", "Sequences from user prompt"],), - "recommend_loras": ("BOOLEAN", {"default": True, "label": "Enable LoRA suggestion"}), + "vision_input": ("STRING", {"multiline": True, "default": "Put outputs of your VL model here to make the Story Writer aware of the starting image."}), "sanity_check": ("BOOLEAN", {"default": True, "label": "Enable Sanity Check"}), "sanity_check_system_prompt": ("STRING", { "multiline": True, @@ -622,7 +622,7 @@ class FunPackStoryWriter: FUNCTION = "write_story" CATEGORY = "FunPack" - def write_story(self, user_prompt, story_system_prompt, sequence_system_prompt, model_path_type, model_path, llm_safetensors_file, prompt_count, top_p, top_k, temperature, max_new_tokens, repetition_penalty, mode, recommend_loras, sanity_check, sanity_check_system_prompt): + def write_story(self, user_prompt, story_system_prompt, sequence_system_prompt, model_path_type, model_path, llm_safetensors_file, prompt_count, top_p, top_k, temperature, max_new_tokens, repetition_penalty, mode, vision_input, sanity_check, sanity_check_system_prompt): llm_model = None llm_tokenizer = None llm_model_device = "cuda" if torch.cuda.is_available() else "cpu" @@ -703,6 +703,9 @@ class FunPackStoryWriter: {"role": "assistant", "content": f"""Previous sequences for continuity:{chr(10).join([f"Sequence {i+1}: {text}" for i, text in enumerate(outputs[:seq_idx])]) if seq_idx > 0 else "This is the first sequence."}"""} ] + if vision_input is not None: + messages.append({"role": "user", "content": f"""Reference image description (the video starts from this image): {vision_input}"""}) + llm_tokens = llm_tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt", tokenize=True ).to(llm_model_device) @@ -770,37 +773,7 @@ class FunPackStoryWriter: # Append generated sequence — this is what chains everything messages.append({"role": "assistant", "content": seq_text}) - - if recommend_loras == True: - lora_list = folder_paths.get_filename_list('loras') - lora_messages = [ - {"role": "system", "content": "Analyze the given array of sequences and the list of LoRAs (low-ranking adaptations - additions to video generation models that enhance certain action, concept or character) and decide which LoRAs out of user's list are recommended to use for each sequence, based on the content of those sequences and the names of LoRAs. Output only list of recommended LoRAs in the format:\nSequence 1: lora name, lora name\nSequence 2: lora name, lora name\nSequence N: lora name, lora name"}, - {"role": "user", "content": lora_list} - ] - for seq_idx, seq_output in enumerate(outputs): - seq_output = outputs[seq_idx] - lora_messages.append({"role": "user", "content": f"This is sequence ID {seq_idx + 1}: {seq_output}"}) - - # Generate recommended LoRAs - llm_tokens = llm_tokenizer.apply_chat_template( - lora_messages, add_generation_prompt=True, return_tensors="pt", tokenize=True - ).to(llm_model_device) - - print("[FunPackStoryWriter] Generating LoRA recommendations...") - with torch.no_grad(): - generated_ids = llm_model.generate( - **llm_tokens, - do_sample=True, - top_p=top_p, top_k=top_k, temperature=temperature, - max_new_tokens=max_new_tokens, - repetition_penalty=repetition_penalty, - pad_token_id=llm_tokenizer.pad_token_id, - eos_token_id=llm_tokenizer.eos_token_id, - ) - - recommended_loras = llm_tokenizer.decode(generated_ids[0][llm_tokens['input_ids'].shape[1]:], skip_special_tokens=True).strip() - print(f"[FunPackStoryWriter] LLM recommended to use this list of LoRAs: {recommended_loras}") - + return tuple(outputs) except Exception as e: @@ -1460,5 +1433,6 @@ NODE_DISPLAY_NAME_MAPPINGS = { +