EricRollei-Local_LLM_Prompt_Enhancer/image_to_video_node.py

"""
Image-to-Video Prompt Expander Node
Uses vision model to describe image, then expands motion prompt
"""

import torch
import numpy as np
from PIL import Image
import io
import base64
from typing import Tuple, Optional
from .llm_backend import LLMBackend
from .expansion_engine import PromptExpander
from .utils import (
    save_prompts_to_file,
    parse_keywords,
    validate_positive_keywords
)
from .qwen3_vl_backend import caption_with_qwen3_vl
from .wildcard_support import (
    enforce_verbatim,
    is_changed_token,
    normalise_seed,
    wildcard_optional_inputs,
)
from .profile_support import WAN_TARGET, add_required_terms, get_profile, target_report_lines, video_target_inputs


class ImageToVideoPromptExpander:
    """
    Advanced Image-to-Video prompt expansion with vision model
    1. Analyzes the image with vision model
    2. Takes user's motion/action input
    3. Combines both into detailed video prompt
    """
    
    def __init__(self):
        self.expander = PromptExpander()
        self.type = "image_to_video_expansion"
        self.output_dir = "output/video_prompts"
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # Image input
                "image": ("IMAGE",),
                
                # Motion/Action input
                "motion_description": ("STRING", {
                    "multiline": True,
                    "default": "slowly turns head and smiles at camera",
                    "placeholder": "Describe the motion/action you want"
                }),
                
                # Core settings
                "preset": ([
                    "custom",
                    "cinematic", 
                    "surreal",
                    "action",
                    "stylized",
                    "noir",
                    "random"
                ], {
                    "default": "cinematic"
                }),
                
                "expansion_tier": ([
                    "auto",
                    "basic",
                    "enhanced", 
                    "advanced",
                    "cinematic"
                ], {
                    "default": "enhanced"
                }),
                
                # Reference image mode (Wan 2.2 optimization)
                "reference_mode": ([
                    "recreate_exact",
                    "style_transfer",
                    "character_remix",
                    "face_only"
                ], {
                    "default": "recreate_exact",
                    "tooltip": (
                        "recreate_exact: Keep exact character, outfit, lighting - animate this\n"
                        "style_transfer: Match lighting/mood only, create new pose/scene\n"
                        "character_remix: Use character but change outfit/setting\n"
                        "face_only: Preserve face identity only, ignore background/lighting"
                    )
                }),
                
                # Vision model settings
                "use_vision_model": ("BOOLEAN", {
                    "default": True
                }),
                
                "vision_backend": ([
                    "lm_studio",
                    "ollama",
                    "qwen3_vl"
                ], {
                    "default": "lm_studio",
                    "tooltip": (
                        "lm_studio: Uses currently loaded model in LM Studio\n"
                        "ollama: Uses currently loaded model in Ollama\n"
                        "qwen3_vl: Auto-detects local Qwen3-VL model (no API server needed)"
                    )
                }),
                
                "vision_endpoint": ("STRING", {
                    "default": "http://localhost:1234/v1",
                    "multiline": False,
                    "tooltip": (
                        "lm_studio/ollama: API endpoint URL\n"
                        "qwen3_vl: Leave default, or specify custom model path like 'A:\\path\\to\\model'"
                    )
                }),
                
                # Text expansion model (can be same or different)
                "expansion_backend": ([
                    "lm_studio",
                    "ollama",
                    "qwen3_vl"
                ], {
                    "default": "lm_studio",
                    "tooltip": (
                        "lm_studio: Uses currently loaded model in LM Studio\n"
                        "ollama: Uses currently loaded model in Ollama\n"
                        "qwen3_vl: Auto-detects local Qwen3-VL model (no API server needed)"
                    )
                }),
                
                "expansion_endpoint": ("STRING", {
                    "default": "http://localhost:1234/v1",
                    "multiline": False,
                    "tooltip": (
                        "lm_studio/ollama: API endpoint URL\n"
                        "qwen3_vl: Leave default, or specify custom model path like 'A:\\path\\to\\model'"
                    )
                }),
                
                "temperature": ("FLOAT", {
                    "default": 0.7,
                    "min": 0.1,
                    "max": 2.0,
                    "step": 0.1
                }),
                
                # Camera/Motion controls
                "camera_movement": ([
                    "auto",
                    "none",
                    "static shot",
                    "camera pushes in",
                    "dolly in",
                    "camera pulls back",
                    "dolly out",
                    "camera pans right",
                    "camera pans left",
                    "camera tilts up",
                    "camera tilts down",
                    "tracking shot",
                    "arc shot",
                    "handheld camera",
                    "steadicam"
                ], {
                    "default": "auto"
                }),
                
                "shot_size": ([
                    "auto",
                    "none",
                    "extreme close-up shot",
                    "close-up shot",
                    "medium close-up shot",
                    "medium shot",
                    "medium wide shot",
                    "wide shot"
                ], {
                    "default": "auto"
                }),
                
                "motion_speed": ([
                    "auto",
                    "very slow",
                    "slow",
                    "normal",
                    "fast",
                    "very fast"
                ], {
                    "default": "auto"
                }),
                
                # Lighting/Time
                "lighting_type": ([
                    "auto",
                    "none",
                    "soft lighting",
                    "hard lighting",
                    "edge lighting",
                    "rim lighting",
                    "natural lighting"
                ], {
                    "default": "auto"
                }),
                
                "time_of_day": ([
                    "auto",
                    "none",
                    "maintain current",
                    "sunrise",
                    "daytime",
                    "sunset",
                    "night"
                ], {
                    "default": "maintain current"
                }),
                
                # Keywords
                "positive_keywords": ("STRING", {
                    "default": "",
                    "multiline": True,
                    "placeholder": "LoRA triggers, style keywords"
                }),
                
                "negative_keywords": ("STRING", {
                    "default": "",
                    "multiline": True
                }),
                
                # Output
                "save_to_file": ("BOOLEAN", {
                    "default": False
                }),
                
                "filename_base": ("STRING", {
                    "default": "img2vid_prompt",
                    "multiline": False
                })
            },
            # New inputs are only ever APPENDED here, so saved workflows keep their widget order.
            "optional": {
                **wildcard_optional_inputs(),
                **video_target_inputs(),      # target_model, duration_seconds
            }
        }

    @classmethod
    def IS_CHANGED(cls, wildcard_seed: int = -1, **kwargs):
        """wildcard_seed = -1 means a fresh roll each run, which must defeat ComfyUI's cache."""
        return is_changed_token(wildcard_seed)
    
    RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING")
    RETURN_NAMES = ("positive_prompt", "negative_prompt", "image_description", "status")
    
    FUNCTION = "expand_img2vid_prompt"
    CATEGORY = "Eric Prompt Enhancers"
    OUTPUT_NODE = True
    
    def expand_img2vid_prompt(
        self,
        image: torch.Tensor,
        motion_description: str,
        preset: str,
        expansion_tier: str,
        reference_mode: str,
        use_vision_model: bool,
        vision_backend: str,
        vision_endpoint: str,
        expansion_backend: str,
        expansion_endpoint: str,
        temperature: float,
        camera_movement: str,
        shot_size: str,
        motion_speed: str,
        lighting_type: str,
        time_of_day: str,
        positive_keywords: str,
        negative_keywords: str,
        save_to_file: bool,
        filename_base: str,
        wildcard_seed: int = -1,
        wildcard_llm_block: bool = True,
        wildcard_instructions: str = "",
        # Absent only in API-format workflows saved before these inputs existed: those keep the Wan 2.2 path.
        target_model: str = WAN_TARGET,
        duration_seconds: float = 0.0
    ) -> Tuple[str, str, str, str]:
        """
        Main processing: Analyze image → Combine with motion → Expand
        """
        
        try:
            # STEP 1: Get image description
            if use_vision_model:
                image_desc_result = self._analyze_image(
                    image,
                    vision_backend,
                    vision_endpoint,
                    temperature
                )
                
                if not image_desc_result["success"]:
                    return (
                        motion_description,
                        "",
                        f"ERROR: {image_desc_result['error']}",
                        f"❌ Vision model failed"
                    )
                
                image_description = image_desc_result["description"]
            else:
                image_description = "[Image description skipped - using motion only]"
            
            # STEP 2: Build combined prompt with reference mode instruction
            combined_input = self._build_combined_prompt_with_mode(
                image_description,
                motion_description,
                motion_speed,
                reference_mode
            )
            
            # Another target model: the picture description travels as a fact about the start frame, and the
            # idea is the motion alone, which is how every image-to-video model wants to be prompted.
            profile_target = target_model != WAN_TARGET and get_profile(target_model) is not None
            start_frame_caption = ""
            if profile_target:
                if image_description and not image_description.startswith("[Image"):
                    start_frame_caption = image_description
                combined_input = self._build_motion_idea(motion_description, motion_speed, reference_mode)
            
            # STEP 3: Gather aesthetic controls
            aesthetic_controls = {}
            if camera_movement not in ["auto", "none"]:
                aesthetic_controls["camera_movement"] = camera_movement
            if shot_size not in ["auto", "none"]:
                aesthetic_controls["shot_size"] = shot_size
            if lighting_type not in ["auto", "none"]:
                aesthetic_controls["lighting_type"] = lighting_type
            if time_of_day not in ["auto", "none", "maintain current"]:
                aesthetic_controls["time_of_day"] = time_of_day
            
            # STEP 4: Expand with expansion engine
            pos_kw_list = parse_keywords(positive_keywords)
            neg_kw_list = parse_keywords(negative_keywords)
            
            system_prompt, user_prompt, breakdown_dict = self.expander.expand_prompt(
                basic_prompt=combined_input,
                preset=preset,
                tier=expansion_tier,
                mode="image-to-video",
                positive_keywords=pos_kw_list,
                variation_seed=None,
                aesthetic_controls=aesthetic_controls,
                wildcard_seed=normalise_seed(wildcard_seed),
                wildcard_llm_block=wildcard_llm_block,
                wildcard_instructions=wildcard_instructions,
                vision_caption=start_frame_caption,
                target_model=target_model,
                duration_seconds=duration_seconds
            )
            on_profile_path = breakdown_dict.get("target_model", WAN_TARGET) != WAN_TARGET
            
            # STEP 5: Call expansion LLM (model_name auto-detected)
            expansion_llm = LLMBackend(
                backend_type=expansion_backend,
                endpoint=expansion_endpoint,
                model_name=None,  # Auto-detect for all backends
                temperature=temperature
            )
            
            response = expansion_llm.send_prompt(
                system_prompt=system_prompt,
                user_prompt=user_prompt,
                max_tokens=3000 if on_profile_path else 2000
            )
            
            if not response["success"]:
                return (
                    combined_input,
                    "",
                    image_description,
                    f"❌ Expansion failed: {response['error']}"
                )
            
            # STEP 6: Parse and validate
            parsed = self.expander.parse_llm_response(
                response["response"],
                target_model=breakdown_dict.get("target_model", WAN_TARGET),
                task=breakdown_dict.get("target_task", ""),
                word_range=breakdown_dict.get("target_words")
            )
            enhanced_prompt = parsed["prompt"]
            
            if not enhanced_prompt or len(enhanced_prompt) < 20:
                enhanced_prompt = combined_input
            
            if on_profile_path:
                # Structured formats (MiniMax H3 fields) cannot take text appended at the end
                enhanced_prompt, term_warnings = add_required_terms(
                    enhanced_prompt, pos_kw_list, self.expander.last_resolution, breakdown_dict["target_model"])
                breakdown_dict["target_warnings"] = list(parsed.get("warnings") or []) + term_warnings
            else:
                # Add required keywords
                if pos_kw_list:
                    keywords_present, missing = validate_positive_keywords(pos_kw_list, enhanced_prompt)
                    if missing:
                        enhanced_prompt += f" {', '.join(missing)}"
                
                # Verbatim wildcard picks (__=name__) must survive the rewrite word for word
                enhanced_prompt, _restored = enforce_verbatim(enhanced_prompt, self.expander.last_resolution)
            
            # STEP 7: Generate negative prompt
            negative_prompt = self.expander.generate_negative_prompt(
                preset=preset,
                custom_negatives=neg_kw_list,
                mode="image-to-video",
                target_model=breakdown_dict.get("target_model", WAN_TARGET)
            )
            
            # STEP 8: Save if requested
            if save_to_file:
                metadata = {
                    "type": "image-to-video",
                    "preset": preset,
                    "tier": expansion_tier,
                    "vision_model": vision_backend if use_vision_model else "none",
                    "expansion_model": expansion_llm.model_name or "auto-detected",
                    "target_model": breakdown_dict.get("target_model", WAN_TARGET),
                    "image_description": image_description,
                    "motion_input": motion_description
                }
                
                breakdown_text = self._format_breakdown(
                    image_description,
                    motion_description,
                    combined_input,
                    aesthetic_controls,
                    breakdown_dict
                )
                
                save_result = save_prompts_to_file(
                    positive_prompt=enhanced_prompt,
                    negative_prompt=negative_prompt,
                    breakdown=breakdown_text,
                    metadata=metadata,
                    filename_base=filename_base,
                    output_dir=self.output_dir
                )
                
                file_status = f"💾 Saved" if save_result["success"] else "⚠️ Save failed"
            else:
                file_status = "Not saved"
            
            target_display = breakdown_dict.get("target_name") or breakdown_dict.get("target_model", WAN_TARGET)
            status = f"✅ Image-to-Video prompt | {target_display} | Vision: {use_vision_model} | {file_status}"
            for warning in breakdown_dict.get("target_warnings") or []:
                status += f"\n⚠️ {warning}"
            
            return (
                enhanced_prompt,
                negative_prompt,
                image_description,
                status
            )
        
        except Exception as e:
            import traceback
            error_detail = traceback.format_exc()
            print(f"ERROR in image-to-video expansion: {error_detail}")
            return (
                motion_description,
                "",
                f"ERROR: {str(e)}",
                f"❌ {str(e)}"
            )
    
    def _analyze_image(
        self,
        image: torch.Tensor,
        backend: str,
        endpoint: str,
        temperature: float
    ) -> dict:
        """Use vision model to analyze the image"""
        
        try:
            # Build vision prompts up front so all backends share them
            vision_system_prompt = """You are an expert at describing images in detail for video generation.

Describe this image comprehensively:
- Main subject(s): appearance, clothing, pose, expression
- Setting/environment: location, objects, background
- Lighting: quality, direction, mood
- Color palette: dominant colors, saturation
- Composition: arrangement of elements
- Atmosphere/mood: overall feeling

Be specific and detailed. Focus on visual elements that matter for video generation.
Output ONLY the description, no labels or meta-commentary."""
            
            vision_user_prompt = "Describe this image in detail for video generation purposes."

            pil_image = self._tensor_to_pil(image)

            if backend == "qwen3_vl":
                # Use endpoint for custom model path
                model_spec = endpoint if endpoint != "http://localhost:1234/v1" else None
                qwen_result = caption_with_qwen3_vl(
                    image=pil_image,
                    prompt=vision_user_prompt,
                    system_prompt=vision_system_prompt,
                    model_spec=model_spec,
                    backend_hint=None,
                    max_new_tokens=768,
                    temperature=temperature,
                )

                if not qwen_result.get("success"):
                    return {
                        "success": False,
                        "error": qwen_result.get("error", "Qwen3-VL caption failed")
                    }

                description = qwen_result.get("caption", "").strip()

                return {
                    "success": True,
                    "description": description
                }

            img_base64 = self._pil_to_base64(pil_image)
            
            # Call vision model (model_name auto-detected)
            llm = LLMBackend(
                backend_type=backend,
                endpoint=endpoint,
                model_name=None,  # Auto-detect for all backends
                temperature=temperature
            )
            
            # For vision models, we need to send the image
            if backend == "lm_studio":
                response = self._call_vision_lm_studio(
                    llm,
                    vision_system_prompt,
                    vision_user_prompt,
                    img_base64
                )
            elif backend == "ollama":
                response = self._call_vision_ollama(
                    llm,
                    vision_system_prompt,
                    vision_user_prompt,
                    img_base64
                )
            else:
                return {"success": False, "error": "Unknown backend"}
            
            if not response["success"]:
                return {"success": False, "error": response["error"]}
            
            description = response["response"].strip()
            
            return {
                "success": True,
                "description": description
            }
        
        except Exception as e:
            return {
                "success": False,
                "error": f"Image analysis failed: {str(e)}"
            }

    @staticmethod
    def _tensor_to_pil(image: torch.Tensor) -> Image.Image:
        """Convert a ComfyUI image tensor to a PIL image."""

        img_np = image[0].cpu().numpy()
        img_np = np.clip(img_np * 255.0, 0, 255).astype(np.uint8)
        return Image.fromarray(img_np)

    @staticmethod
    def _pil_to_base64(pil_image: Image.Image) -> str:
        """Encode a PIL image to a base64 PNG string."""

        buffered = io.BytesIO()
        pil_image.save(buffered, format="PNG")
        return base64.b64encode(buffered.getvalue()).decode('utf-8')
    
    def _call_vision_lm_studio(self, llm, system_prompt: str, user_prompt: str, img_base64: str) -> dict:
        """Call LM Studio with vision (OpenAI-compatible format)"""
        import requests
        
        try:
            url = f"{llm.endpoint}/chat/completions"
            
            payload = {
                "model": llm.model_name,
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": user_prompt},
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/png;base64,{img_base64}"
                                }
                            }
                        ]
                    }
                ],
                "temperature": llm.temperature,
                "max_tokens": 1000
            }
            
            response = requests.post(url, json=payload, timeout=120)
            response.raise_for_status()
            
            data = response.json()
            content = data['choices'][0]['message']['content']
            
            return {"success": True, "response": content, "error": None}
        
        except Exception as e:
            return {"success": False, "response": "", "error": str(e)}
    
    def _call_vision_ollama(self, llm, system_prompt: str, user_prompt: str, img_base64: str) -> dict:
        """Call Ollama with vision"""
        import requests
        
        try:
            url = f"{llm.endpoint}/api/generate"
            
            full_prompt = f"{system_prompt}\n\n{user_prompt}"
            
            payload = {
                "model": llm.model_name,
                "prompt": full_prompt,
                "images": [img_base64],
                "stream": False,
                "options": {
                    "temperature": llm.temperature,
                    "num_predict": 1000
                }
            }
            
            response = requests.post(url, json=payload, timeout=120)
            response.raise_for_status()
            
            data = response.json()
            content = data.get('response', '')
            
            return {"success": True, "response": content, "error": None}
        
        except Exception as e:
            return {"success": False, "response": "", "error": str(e)}
    
    def _build_combined_prompt(
        self,
        image_description: str,
        motion_description: str,
        motion_speed: str
    ) -> str:
        """Combine image description and motion into expansion-ready prompt"""
        
        parts = []
        
        # Add image description
        if image_description and not image_description.startswith("ERROR") and not image_description.startswith("[Image"):
            parts.append(f"Scene: {image_description}")
        
        # Add motion with speed modifier
        motion = motion_description
        if motion_speed not in ["auto", "normal"]:
            motion = f"{motion_speed}, {motion}"
        
        parts.append(f"Motion: {motion}")
        
        return " | ".join(parts)
    
    # Reference modes of this node, said in a way that does not assume Wan 2.2
    MOTION_IDEA_NOTES = {
        "style_transfer": "Keep the lighting, colour palette and mood of the start image; the pose and action may change freely.",
        "character_remix": "The character in the start image is the same person throughout; the clothing and setting may change as the motion requires.",
        "face_only": "Only the subject's face and body identity must be preserved from the start image.",
    }

    def _build_motion_idea(self, motion_description: str, motion_speed: str, reference_mode: str) -> str:
        """The idea handed to a profile-built rewriter: the motion, its speed, and the reference mode as one plain sentence."""
        motion = motion_description
        if motion_speed not in ["auto", "normal"]:
            motion = f"{motion_speed}, {motion}"
        note = self.MOTION_IDEA_NOTES.get(reference_mode, "")
        return f"{motion}\n\n{note}" if note else motion

    def _build_combined_prompt_with_mode(
        self,
        image_description: str,
        motion_description: str,
        motion_speed: str,
        reference_mode: str
    ) -> str:
        """Combine image description and motion with Wan 2.2 reference mode instruction"""
        
        # Get reference mode instruction
        mode_instructions = {
            "recreate_exact": "Use the provided image as the exact character and costume reference. Keep the same face, hair, outfit, and lighting. Animate this character without changing identity.",
            "style_transfer": "Match the lighting, color palette, and cinematic mood of the provided image, but create a new pose and new scene.",
            "character_remix": "Loosely base the main character on the provided image, but change the outfit and relocate them to a new setting.",
            "face_only": "Preserve ONLY the subject's face and body identity from the provided image. Ignore the original background and lighting. Place this character in the new described scene."
        }
        
        mode_instruction = mode_instructions.get(reference_mode, mode_instructions["recreate_exact"])
        
        parts = []
        
        # Add reference mode instruction first
        parts.append(f"REFERENCE MODE: {mode_instruction}")
        
        # Add image description
        if image_description and not image_description.startswith("ERROR") and not image_description.startswith("[Image"):
            parts.append(f"Scene Analysis: {image_description}")
        
        # Add motion with speed modifier
        motion = motion_description
        if motion_speed not in ["auto", "normal"]:
            motion = f"{motion_speed}, {motion}"
        
        parts.append(f"Motion: {motion}")
        
        return " | ".join(parts)
    
    def _format_breakdown(
        self,
        image_description: str,
        motion_input: str,
        combined_prompt: str,
        aesthetic_controls: dict,
        breakdown_dict: dict
    ) -> str:
        """Format breakdown for file output"""
        
        lines = [
            "=" * 60,
            "IMAGE-TO-VIDEO PROMPT EXPANSION",
            "=" * 60,
            "\nIMAGE ANALYSIS:",
            image_description,
            "\nMOTION INPUT:",
            motion_input,
            "\nCOMBINED PROMPT:",
            combined_prompt,
            f"\nExpansion Tier: {breakdown_dict.get('detected_tier', 'N/A')}",
            f"Preset: {breakdown_dict.get('applied_preset', 'N/A')}"
        ]
        lines.extend(target_report_lines(breakdown_dict))
        
        if aesthetic_controls:
            lines.append("\nAesthetic Controls:")
            for key, value in aesthetic_controls.items():
                lines.append(f"  - {key.replace('_', ' ').title()}: {value}")
        
        lines.append("\n" + "=" * 60)
        
        return "\n".join(lines)