vms/ui/project/tabs/preview_tab.py

"""
Preview tab for Video Model Studio UI 
"""

import gradio as gr
import logging
import json
from pathlib import Path
from typing import Dict, Any, List, Optional, Tuple
import time

from vms.utils import BaseTab
from vms.config import (
    OUTPUT_PATH, MODEL_TYPES, DEFAULT_PROMPT_PREFIX, MODEL_VERSIONS
)

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

class PreviewTab(BaseTab):
    """Preview tab for testing trained models"""
    
    def __init__(self, app_state):
        super().__init__(app_state)
        self.id = "preview_tab"
        self.title = "4️⃣ Preview"
         
    def create(self, parent=None) -> gr.TabItem:
        """Create the Preview tab UI components"""
        with gr.TabItem(self.title, id=self.id) as tab:
            with gr.Row():
                gr.Markdown("## 🔬 Preview your model")
            
            with gr.Row():
                with gr.Column(scale=2):
                                        
                    # Add dropdown to choose between LoRA and original model
                    has_lora = self.check_lora_model_exists()
                    lora_choices = []
                    default_lora_choice = ""

                    if has_lora:
                        lora_choices = ["Use LoRA model", "Use original model"]
                        default_lora_choice = "Use LoRA model"
                    else:
                        lora_choices = ["Cannot find LoRA model", "Use original model"]
                        default_lora_choice = "Use original model"

                    self.components["use_lora"] = gr.Dropdown(
                        choices=lora_choices,
                        label="Model Selection",
                        value=default_lora_choice
                    )
                    
                    self.components["prompt"] = gr.Textbox(
                        label="Prompt",
                        placeholder="Enter your prompt here...",
                        lines=3
                    )
                    
                    self.components["negative_prompt"] = gr.Textbox(
                        label="Negative Prompt",
                        placeholder="Enter negative prompt here...",
                        lines=3,
                        value="worst quality, low quality, blurry, jittery, distorted, ugly, deformed, disfigured, messy background"
                    )
                    
                    self.components["prompt_prefix"] = gr.Textbox(
                        label="Global Prompt Prefix",
                        placeholder="Prefix to add to all prompts",
                        value=DEFAULT_PROMPT_PREFIX
                    )

                    # Ensure seed is interactive with a slider
                    self.components["seed"] = gr.Slider(
                        label="Generation Seed (-1 for random)",
                        minimum=-1,
                        maximum=2147483647,  # 2^31 - 1
                        step=1,
                        value=-1,
                        info="Set to -1 for random seed or specific value for reproducible results",
                        interactive=True
                    )
                    
                    with gr.Row():
                        # Get the currently selected model type from training tab if possible
                        default_model = self.get_default_model_type()
                        
                        with gr.Column():
                            # Make model_type read-only (disabled), as it must match what was trained
                            self.components["model_type"] = gr.Dropdown(
                                choices=list(MODEL_TYPES.keys()),
                                label="Model Type (from training)",
                                value=default_model,
                                interactive=False
                            )
                            
                            # Add model version selection based on model type
                            self.components["model_version"] = gr.Dropdown(
                                label="Model Version",
                                choices=self.get_model_version_choices(default_model),
                                value=self.get_default_model_version(default_model)
                            )

                    # Add image input for image-to-video models
                    self.components["conditioning_image"] = gr.Image(
                        label="Conditioning Image (for Image-to-Video models)",
                        type="filepath",
                        visible=False
                    )
                    
                    with gr.Row():
                        self.components["resolution_preset"] = gr.Dropdown(
                            choices=["480p", "720p"],
                            label="Resolution Preset",
                            value="480p"
                        )
                    
                    with gr.Row():
                        self.components["width"] = gr.Number(
                            label="Width",
                            value=832,
                            precision=0
                        )
                        
                        self.components["height"] = gr.Number(
                            label="Height",
                            value=480,
                            precision=0
                        )
                    
                    with gr.Row():
                        self.components["num_frames"] = gr.Slider(
                            label="Number of Frames",
                            minimum=1,
                            maximum=257,
                            step=8,
                            value=49
                        )
                        
                        self.components["fps"] = gr.Slider(
                            label="FPS",
                            minimum=1,
                            maximum=60,
                            step=1,
                            value=16
                        )
                    
                    with gr.Row():
                        self.components["guidance_scale"] = gr.Slider(
                            label="Guidance Scale",
                            minimum=1.0,
                            maximum=10.0,
                            step=0.1,
                            value=5.0
                        )
                        
                        self.components["flow_shift"] = gr.Slider(
                            label="Flow Shift",
                            minimum=0.0,
                            maximum=10.0,
                            step=0.1,
                            value=3.0
                        )
                    
                    with gr.Row():
                        self.components["lora_scale"] = gr.Slider(
                            label="LoRA Scale",
                            minimum=0.0,
                            maximum=1.0,
                            step=0.01,
                            value=0.7,
                            visible=has_lora  # Only visible if using LoRA
                        )
                        
                        self.components["inference_steps"] = gr.Slider(
                            label="Inference Steps",
                            minimum=1,
                            maximum=100,
                            step=1,
                            value=20
                        )
                    
                    self.components["enable_cpu_offload"] = gr.Checkbox(
                        label="Enable Model CPU Offload (for low-VRAM GPUs)",
                        value=False # let's assume user is using a video model training rig with a good GPU
                    )
                    
                    self.components["generate_btn"] = gr.Button(
                        "Generate Video",
                        variant="primary"
                    )
                
                with gr.Column(scale=3):
                    self.components["preview_video"] = gr.Video(
                        label="Generated Video",
                        interactive=False
                    )
                    
                    self.components["status"] = gr.Textbox(
                        label="Status",
                        interactive=False
                    )
                    
                    with gr.Accordion("Log", open=False):
                        self.components["log"] = gr.TextArea(
                            label="Generation Log",
                            interactive=False,
                            lines=60
                        )
        
        return tab

    def check_lora_model_exists(self) -> bool:
        """Check if any LoRA model files exist in the output directory"""
        # Look for the standard LoRA weights file
        lora_path = OUTPUT_PATH / "pytorch_lora_weights.safetensors"
        if lora_path.exists():
            return True
        
        # If not found in the expected location, try to find in checkpoints
        checkpoints = list(OUTPUT_PATH.glob("finetrainers_step_*"))
        has_checkpoints = len(checkpoints) > 0
        if not checkpoints:
            return False
        
        for checkpoint in checkpoints:
            lora_path = checkpoint / "pytorch_lora_weights.safetensors"
            if lora_path.exists():
                return True
        
        return False
    
    def update_lora_ui(self, use_lora_value: str) -> Dict[str, Any]:
        """Update UI based on LoRA selection"""
        is_using_lora = "Use LoRA model" in use_lora_value
        
        return {
            self.components["lora_scale"]: gr.Slider(visible=is_using_lora)
        }
    
    def get_model_version_choices(self, model_type: str) -> List[str]:
        """Get model version choices based on model type"""
        # Convert UI display name to internal name
        internal_type = MODEL_TYPES.get(model_type)
        if not internal_type or internal_type not in MODEL_VERSIONS:
            logger.warning(f"No model versions found for {model_type} (internal type: {internal_type})")
            return []
            
        # Return just the model IDs as a list of simple strings
        version_ids = list(MODEL_VERSIONS.get(internal_type, {}).keys())
        logger.info(f"Found {len(version_ids)} versions for {model_type}: {version_ids}")
        
        # Ensure they're all strings
        return [str(version) for version in version_ids]
    
    def get_default_model_version(self, model_type: str) -> str:
        """Get default model version for the given model type"""
        # Convert UI display name to internal name
        internal_type = MODEL_TYPES.get(model_type)
        logger.debug(f"get_default_model_version({model_type}) = {internal_type}")
        
        if not internal_type or internal_type not in MODEL_VERSIONS:
            logger.warning(f"No valid model versions found for {model_type}")
            return ""
            
        # Get the first version available for this model type
        versions = list(MODEL_VERSIONS.get(internal_type, {}).keys())
        if versions:
            default_version = versions[0]
            logger.debug(f"Default version for {model_type}: {default_version}")
            return default_version
        return ""

    def get_default_model_type(self) -> str:
        """Get the model type from the latest training session"""
        try:
            # First check the session.json which contains the actual training data
            session_file = OUTPUT_PATH / "session.json"
            if session_file.exists():
                with open(session_file, 'r') as f:
                    session_data = json.load(f)
                    
                # Get the internal model type from the session parameters
                if "params" in session_data and "model_type" in session_data["params"]:
                    internal_model_type = session_data["params"]["model_type"]
                    
                    # Convert internal model type to display name
                    for display_name, internal_name in MODEL_TYPES.items():
                        if internal_name == internal_model_type:
                            logger.info(f"Using model type '{display_name}' from session file")
                            return display_name
                        
                    # If we couldn't map it, log a warning
                    logger.warning(f"Could not map internal model type '{internal_model_type}' to a display name")
                    
            # If we couldn't get it from session.json, try to get it from UI state
            ui_state = self.app.training.load_ui_state()
            model_type = ui_state.get("model_type")
            
            # Make sure it's a valid model type
            if model_type in MODEL_TYPES:
                return model_type
            
            # If we still couldn't get a valid model type, try to get it from the training tab
            if hasattr(self.app, 'tabs') and 'train_tab' in self.app.tabs:
                train_tab = self.app.tabs['train_tab']
                if hasattr(train_tab, 'components') and 'model_type' in train_tab.components:
                    train_model_type = train_tab.components['model_type'].value
                    if train_model_type in MODEL_TYPES:
                        return train_model_type
            
            # Fallback to first model type
            return list(MODEL_TYPES.keys())[0]
        except Exception as e:
            logger.warning(f"Failed to get default model type from session: {e}")
            return list(MODEL_TYPES.keys())[0]
    
    def extract_model_id(self, model_version_choice: str) -> str:
        """Extract model ID from model version choice string"""
        if " - " in model_version_choice:
            return model_version_choice.split(" - ")[0].strip()
        return model_version_choice
    
    def get_model_version_type(self, model_type: str, model_version: str) -> str:
        """Get the model version type (text-to-video or image-to-video)"""
        # Convert UI display name to internal name
        internal_type = MODEL_TYPES.get(model_type)
        if not internal_type:
            return "text-to-video"
            
        # Extract model_id from model version choice
        model_id = self.extract_model_id(model_version)
            
        # Get versions from preview service
        versions = self.app.previewing.get_model_versions(internal_type)
        model_version_info = versions.get(model_id, {})
        
        # Return the model version type or default to text-to-video
        return model_version_info.get("type", "text-to-video")
    
    def connect_events(self) -> None:
        """Connect event handlers to UI components"""
        # Update resolution when preset changes
        self.components["resolution_preset"].change(
            fn=self.update_resolution,
            inputs=[self.components["resolution_preset"]],
            outputs=[
                self.components["width"],
                self.components["height"],
                self.components["flow_shift"]
            ]
        )
        
        # Update model_version choices when model_type changes or tab is selected
        if hasattr(self.app, 'tabs_component') and self.app.tabs_component is not None:
            self.app.tabs_component.select(
                fn=self.sync_model_type_and_versions,
                inputs=[],
                outputs=[
                    self.components["model_type"],
                    self.components["model_version"]
                ]
            )
        
        # Update model version-specific UI elements when version changes
        self.components["model_version"].change(
            fn=self.update_model_version_ui,
            inputs=[
                self.components["model_type"],
                self.components["model_version"]
            ],
            outputs=[
                self.components["conditioning_image"]
            ]
        )
        
        # Connect LoRA selection dropdown to update LoRA weight visibility
        self.components["use_lora"].change(
            fn=self.update_lora_ui,
            inputs=[self.components["use_lora"]],
            outputs=[self.components["lora_scale"]]
        )
        
        # Load preview UI state when the tab is selected
        if hasattr(self.app, 'tabs_component') and self.app.tabs_component is not None:
            self.app.tabs_component.select(
                fn=self.load_preview_state,
                inputs=[],
                outputs=[
                    self.components["prompt"],
                    self.components["negative_prompt"],
                    self.components["prompt_prefix"],
                    self.components["width"],
                    self.components["height"],
                    self.components["num_frames"],
                    self.components["fps"],
                    self.components["guidance_scale"],
                    self.components["flow_shift"],
                    self.components["lora_scale"],
                    self.components["inference_steps"],
                    self.components["enable_cpu_offload"],
                    self.components["model_version"],
                    self.components["seed"],
                    self.components["use_lora"]
                ]
            )
        
        # Save preview UI state when values change
        for component_name in [
            "prompt", "negative_prompt", "prompt_prefix", "model_version", "resolution_preset",
            "width", "height", "num_frames", "fps", "guidance_scale", "flow_shift",
            "lora_scale", "inference_steps", "enable_cpu_offload", "seed", "use_lora"
        ]:
            if component_name in self.components:
                self.components[component_name].change(
                    fn=self.save_preview_state_value,
                    inputs=[self.components[component_name]],
                    outputs=[]
                )
        
        # Generate button click
        self.components["generate_btn"].click(
            fn=self.generate_video,
            inputs=[
                self.components["model_type"],
                self.components["model_version"],
                self.components["prompt"],
                self.components["negative_prompt"],
                self.components["prompt_prefix"],
                self.components["width"],
                self.components["height"],
                self.components["num_frames"],
                self.components["guidance_scale"],
                self.components["flow_shift"],
                self.components["lora_scale"],
                self.components["inference_steps"],
                self.components["enable_cpu_offload"],
                self.components["fps"],
                self.components["conditioning_image"],
                self.components["seed"],
                self.components["use_lora"]
            ],
            outputs=[
                self.components["preview_video"],
                self.components["status"],
                self.components["log"]
            ]
        )
    
    def update_model_version_ui(self, model_type: str, model_version: str) -> Dict[str, Any]:
        """Update UI based on the selected model version"""
        model_version_type = self.get_model_version_type(model_type, model_version)
        
        # Show conditioning image input only for image-to-video models
        show_conditioning_image = model_version_type == "image-to-video"
        
        return {
            self.components["conditioning_image"]: gr.Image(visible=show_conditioning_image)
        }
    
    def sync_model_type_and_versions(self) -> Tuple[str, str]:
        """Sync model type with training tab when preview tab is selected and update model version choices"""
        model_type = self.get_default_model_type()
        model_version = ""
        
        # Try to get model_version from session or UI state
        ui_state = self.app.training.load_ui_state()
        preview_state = ui_state.get("preview", {})
        model_version = preview_state.get("model_version", "")
        
        # If no model version specified or invalid, use default
        if not model_version:
            # Get the internal model type
            internal_type = MODEL_TYPES.get(model_type)
            if internal_type and internal_type in MODEL_VERSIONS:
                versions = list(MODEL_VERSIONS[internal_type].keys())
                if versions:
                    model_version = versions[0]
        
        return model_type, model_version
    
    def update_resolution(self, preset: str) -> Tuple[int, int, float]:
        """Update resolution and flow shift based on preset"""
        if preset == "480p":
            return 832, 480, 3.0
        elif preset == "720p":
            return 1280, 720, 5.0
        else:
            return 832, 480, 3.0
    
    def load_preview_state(self) -> Tuple:
        """Load saved preview UI state"""
        # Try to get the saved state
        try:
            state = self.app.training.load_ui_state()
            preview_state = state.get("preview", {})
            
            # Get model type (can't be changed in UI)
            model_type = self.get_default_model_type()
            
            # If model_version not in choices for current model_type, use default
            model_version = preview_state.get("model_version", "")
            model_version_choices = self.get_model_version_choices(model_type)
            if model_version not in model_version_choices and model_version_choices:
                model_version = model_version_choices[0]
            
            # Check if LoRA exists and set appropriate dropdown options
            has_lora = self.check_lora_model_exists()
            use_lora = preview_state.get("use_lora", "")
            
            # Validate use_lora value against current state
            if has_lora:
                valid_choices = ["Use LoRA model", "Use original model"]
                if use_lora not in valid_choices:
                    use_lora = "Use LoRA model"  # Default when LoRA exists
            else:
                valid_choices = ["Cannot find LoRA model", "Use original model"]
                if use_lora not in valid_choices:
                    use_lora = "Use original model"  # Default when no LoRA
                    
            # Update the dropdown choices in the UI
            try:
                self.components["use_lora"].choices = valid_choices
            except Exception as e:
                logger.error(f"Failed to update use_lora choices: {e}")
            
            return (
                preview_state.get("prompt", ""),
                preview_state.get("negative_prompt", "worst quality, low quality, blurry, jittery, distorted, ugly, deformed, disfigured, messy background"),
                preview_state.get("prompt_prefix", DEFAULT_PROMPT_PREFIX),
                preview_state.get("width", 832),
                preview_state.get("height", 480),
                preview_state.get("num_frames", 49),
                preview_state.get("fps", 16),
                preview_state.get("guidance_scale", 5.0),
                preview_state.get("flow_shift", 3.0),
                preview_state.get("lora_scale", 0.7),
                preview_state.get("inference_steps", 30),
                preview_state.get("enable_cpu_offload", True),
                model_version,
                preview_state.get("seed", -1),
                use_lora
            )
        except Exception as e:
            logger.error(f"Error loading preview state: {e}")
            # Return defaults if loading fails
            return (
                "", 
                "worst quality, low quality, blurry, jittery, distorted, ugly, deformed, disfigured, messy background", 
                DEFAULT_PROMPT_PREFIX,
                832, 480, 49, 16, 5.0, 3.0, 0.7, 30, True,
                self.get_default_model_version(self.get_default_model_type()),
                -1,
                "Use original model" if not self.check_lora_model_exists() else "Use LoRA model"
            )
    
    def save_preview_state_value(self, value: Any) -> None:
        """Save an individual preview state value"""
        try:
            # Get the component name from the event context
            import inspect
            frame = inspect.currentframe()
            frame = inspect.getouterframes(frame)[1]
            event_context = frame.frame.f_locals
            component = event_context.get('component')
            
            if component is None:
                return
            
            # Find the component name
            component_name = None
            for name, comp in self.components.items():
                if comp == component:
                    component_name = name
                    break
            
            if component_name is None:
                return
            
            # Load current state
            state = self.app.training.load_ui_state()
            if "preview" not in state:
                state["preview"] = {}
            
            # Update the value
            state["preview"][component_name] = value
            
            # Save state
            self.app.training.save_ui_state(state)
        except Exception as e:
            logger.error(f"Error saving preview state: {e}")
    
    def generate_video(
        self,
        model_type: str,
        model_version: str,
        prompt: str,
        negative_prompt: str,
        prompt_prefix: str,
        width: int,
        height: int,
        num_frames: int,
        guidance_scale: float,
        flow_shift: float,
        lora_scale: float,
        inference_steps: int,
        enable_cpu_offload: bool,
        fps: int,
        conditioning_image: Optional[str] = None,
        seed: int = -1,
        use_lora: str = "Use LoRA model"
    ) -> Tuple[Optional[str], str, str]:
        """Handler for generate button click, delegates to preview service"""
        # Save all the parameters to preview state before generating
        print("preview_tab: generate_video() has been called")
        try:
            state = self.app.training.load_ui_state()
            if "preview" not in state:
                state["preview"] = {}
                
            # Extract model ID from model version choice
            model_version_id = self.extract_model_id(model_version)
                
            # Update all values
            preview_state = {
                "prompt": prompt,
                "negative_prompt": negative_prompt,
                "prompt_prefix": prompt_prefix,
                "model_type": model_type,
                "model_version": model_version,
                "width": width,
                "height": height,
                "num_frames": num_frames,
                "fps": fps,
                "guidance_scale": guidance_scale,
                "flow_shift": flow_shift,
                "lora_scale": lora_scale,
                "inference_steps": inference_steps,
                "enable_cpu_offload": enable_cpu_offload,
                "seed": seed,
                "use_lora": use_lora
            }
            
            state["preview"] = preview_state
            self.app.training.save_ui_state(state)
        except Exception as e:
            logger.error(f"Error saving preview state before generation: {e}")
        
        # Extract model ID from model version choice string
        model_version_id = self.extract_model_id(model_version)
        
        # Initial UI update
        video_path, status, log = None, "Initializing generation...", "Starting video generation process..."
        
        # Set lora_path to None if not using LoRA
        use_lora_model = use_lora == "Use LoRA model"
        
        # Start actual generation
        # If not using LoRA, set lora_scale to 0 to disable it
        effective_lora_scale = lora_scale if use_lora_model else 0.0
        
        result = self.app.previewing.generate_video(
            model_type=model_type,
            model_version=model_version_id,
            prompt=prompt,
            negative_prompt=negative_prompt,
            prompt_prefix=prompt_prefix,
            width=width,
            height=height,
            num_frames=num_frames,
            guidance_scale=guidance_scale,
            flow_shift=flow_shift,
            lora_scale=effective_lora_scale,  # Use 0.0 if not using LoRA
            inference_steps=inference_steps,
            enable_cpu_offload=enable_cpu_offload,
            fps=fps,
            conditioning_image=conditioning_image,
            seed=seed
        )
        
        # Return final result
        return result