vms/ui/video_trainer_ui.py

import platform
import subprocess

#import sys
#print("python = ", sys.version)

# can be "Linux", "Darwin"
if platform.system() == "Linux":
    # for some reason it says "pip not found"
    # and also "pip3 not found"
    # subprocess.run(
    #     "pip install flash-attn --no-build-isolation",
    #
    #     # hmm... this should be False, since we are in a CUDA environment, no?
    #     env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
    #     
    #     shell=True,
    # )
    pass

import gradio as gr
from pathlib import Path
import logging
import mimetypes
import shutil
import os
import traceback
import asyncio
import tempfile
import zipfile
from typing import Any, Optional, Dict, List, Union, Tuple
from typing import AsyncGenerator

from ..services import TrainingService, CaptioningService, SplittingService, ImportService
from ..config import (
    STORAGE_PATH, VIDEOS_TO_SPLIT_PATH, STAGING_PATH,
    TRAINING_PATH, LOG_FILE_PATH, TRAINING_PRESETS, TRAINING_VIDEOS_PATH, MODEL_PATH, OUTPUT_PATH, DEFAULT_CAPTIONING_BOT_INSTRUCTIONS,
    DEFAULT_PROMPT_PREFIX, HF_API_TOKEN, ASK_USER_TO_DUPLICATE_SPACE, MODEL_TYPES, SMALL_TRAINING_BUCKETS
)
from ..utils import make_archive, count_media_files, format_media_title, is_image_file, is_video_file, validate_model_repo, format_time, copy_files_to_training_dir, prepare_finetrainers_dataset, TrainingLogParser
from ..tabs import ImportTab, SplitTab, CaptionTab, TrainTab, ManageTab

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

httpx_logger = logging.getLogger('httpx')
httpx_logger.setLevel(logging.WARN)

class VideoTrainerUI:       
    def __init__(self):
        """Initialize services and tabs"""
        # Initialize core services
        self.trainer = TrainingService()
        self.splitter = SplittingService()
        self.importer = ImportService()
        self.captioner = CaptioningService()
        self._should_stop_captioning = False
        
        # Recovery status from any interrupted training
        recovery_result = self.trainer.recover_interrupted_training()
        self.recovery_status = recovery_result.get("status", "unknown")
        self.ui_updates = recovery_result.get("ui_updates", {})
        
        self.log_parser = TrainingLogParser()
        
        # Shared state for tabs
        self.state = {
            "recovery_result": recovery_result
        }
        
        # Initialize tabs dictionary (will be populated in create_ui)
        self.tabs = {}
        self.tabs_component = None
        
    def create_ui(self):
        """Create the main Gradio UI"""
        with gr.Blocks(title="🎥 Video Model Studio") as app:
            gr.Markdown("# 🎥 Video Model Studio")
            
            # Create main tabs component
            with gr.Tabs() as self.tabs_component:
                # Initialize tab objects
                self.tabs["import_tab"] = ImportTab(self)
                self.tabs["split_tab"] = SplitTab(self)
                self.tabs["caption_tab"] = CaptionTab(self)
                self.tabs["train_tab"] = TrainTab(self)
                self.tabs["manage_tab"] = ManageTab(self)
                
                # Create tab UI components
                for tab_id, tab_obj in self.tabs.items():
                    tab_obj.create(self.tabs_component)
            
            # Connect event handlers
            for tab_id, tab_obj in self.tabs.items():
                tab_obj.connect_events()
            
            # Add app-level timers for auto-refresh functionality
            self._add_timers()
            
            # Initialize app state on load
            app.load(
                fn=self.initialize_app_state,
                outputs=[
                    self.tabs["split_tab"].components["video_list"],
                    self.tabs["caption_tab"].components["training_dataset"],
                    self.tabs["train_tab"].components["start_btn"],
                    self.tabs["train_tab"].components["stop_btn"],
                    self.tabs["train_tab"].components["pause_resume_btn"],
                    self.tabs["train_tab"].components["training_preset"],
                    self.tabs["train_tab"].components["model_type"],
                    self.tabs["train_tab"].components["lora_rank"],
                    self.tabs["train_tab"].components["lora_alpha"],
                    self.tabs["train_tab"].components["num_epochs"],
                    self.tabs["train_tab"].components["batch_size"],
                    self.tabs["train_tab"].components["learning_rate"],
                    self.tabs["train_tab"].components["save_iterations"]
                ]
            )
            
        return app
    
    def _add_timers(self):
        """Add auto-refresh timers to the UI"""
        # Status update timer (every 1 second)
        status_timer = gr.Timer(value=1)
        status_timer.tick(
            fn=self.get_latest_status_message_logs_and_button_labels,
            outputs=[
                self.tabs["train_tab"].components["status_box"],
                self.tabs["train_tab"].components["log_box"],
                self.tabs["train_tab"].components["start_btn"],
                self.tabs["train_tab"].components["stop_btn"],
                self.tabs["train_tab"].components["pause_resume_btn"]
            ]
        )
        
        # Dataset refresh timer (every 5 seconds)
        dataset_timer = gr.Timer(value=5)
        dataset_timer.tick(
            fn=self.refresh_dataset,
            outputs=[
                self.tabs["split_tab"].components["video_list"],
                self.tabs["caption_tab"].components["training_dataset"]
            ]
        )
        
        # Titles update timer (every 6 seconds)
        titles_timer = gr.Timer(value=6)
        titles_timer.tick(
            fn=self.update_titles,
            outputs=[
                self.tabs["split_tab"].components["split_title"],
                self.tabs["caption_tab"].components["caption_title"],
                self.tabs["train_tab"].components["train_title"]
            ]
        )
    
    def handle_global_stop(self):
        """Handle the global stop button click"""
        result = self.stop_all_and_clear()
        
        # Format the details for display
        status = result["status"]
        details = "\n".join(f"{k}: {v}" for k, v in result["details"].items())
        full_status = f"{status}\n\nDetails:\n{details}"
        
        # Get fresh lists after cleanup
        videos = self.splitter.list_unprocessed_videos()
        clips = self.list_training_files_to_caption()
        
        return {
            self.tabs["manage_tab"].components["global_status"]: gr.update(value=full_status, visible=True),
            self.tabs["split_tab"].components["video_list"]: videos,
            self.tabs["caption_tab"].components["training_dataset"]: clips,
            self.tabs["train_tab"].components["status_box"]: "Training stopped and data cleared",
            self.tabs["train_tab"].components["log_box"]: "",
            self.tabs["split_tab"].components["detect_status"]: "Scene detection stopped",
            self.tabs["import_tab"].components["import_status"]: "All data cleared",
            self.tabs["caption_tab"].components["preview_status"]: "Captioning stopped"
        }
    
    def upload_to_hub(self, repo_id: str) -> str:
        """Upload model to HuggingFace Hub"""
        if not repo_id:
            return "Error: Repository ID is required"
        
        # Validate repository name
        validation = validate_model_repo(repo_id)
        if validation["error"]:
            return f"Error: {validation['error']}"
        
        # Check if we have a model to upload
        if not self.trainer.get_model_output_safetensors():
            return "Error: No model found to upload"
        
        # Upload model to hub
        success = self.trainer.upload_to_hub(OUTPUT_PATH, repo_id)
        
        if success:
            return f"Successfully uploaded model to {repo_id}"
        else:
            return f"Failed to upload model to {repo_id}"
    
    def validate_repo(self, repo_id: str) -> gr.update:
        """Validate repository ID for HuggingFace Hub"""
        validation = validate_model_repo(repo_id)
        if validation["error"]:
            return gr.update(value=repo_id, error=validation["error"])
        return gr.update(value=repo_id, error=None)


    async def _process_caption_generator(self, captioning_bot_instructions, prompt_prefix):
        """Process the caption generator's results in the background"""
        try:
            async for _ in self.captioner.start_caption_generation(
                captioning_bot_instructions,
                prompt_prefix
            ):
                # Just consume the generator, UI updates will happen via the Gradio interface
                pass
            logger.info("Background captioning completed")
        except Exception as e:
            logger.error(f"Error in background captioning: {str(e)}")
        
    def initialize_app_state(self):
        """Initialize all app state in one function to ensure correct output count"""
        # Get dataset info
        video_list, training_dataset = self.refresh_dataset()
        
        # Get button states
        button_states = self.get_initial_button_states()
        start_btn = button_states[0]
        stop_btn = button_states[1]
        pause_resume_btn = button_states[2]
        
        # Get UI form values
        ui_state = self.load_ui_values()
        training_preset = ui_state.get("training_preset", list(TRAINING_PRESETS.keys())[0])
        model_type_val = ui_state.get("model_type", list(MODEL_TYPES.keys())[0])
        lora_rank_val = ui_state.get("lora_rank", "128")
        lora_alpha_val = ui_state.get("lora_alpha", "128")
        num_epochs_val = int(ui_state.get("num_epochs", 70))
        batch_size_val = int(ui_state.get("batch_size", 1))
        learning_rate_val = float(ui_state.get("learning_rate", 3e-5))
        save_iterations_val = int(ui_state.get("save_iterations", 500))
        
        # Return all values in the exact order expected by outputs
        return (
            video_list, 
            training_dataset,
            start_btn, 
            stop_btn, 
            pause_resume_btn,
            training_preset, 
            model_type_val, 
            lora_rank_val, 
            lora_alpha_val,
            num_epochs_val, 
            batch_size_val, 
            learning_rate_val, 
            save_iterations_val
        )

    def initialize_ui_from_state(self):
        """Initialize UI components from saved state"""
        ui_state = self.load_ui_values()
        
        # Return values in order matching the outputs in app.load
        return (
            ui_state.get("training_preset", list(TRAINING_PRESETS.keys())[0]),
            ui_state.get("model_type", list(MODEL_TYPES.keys())[0]),
            ui_state.get("lora_rank", "128"),
            ui_state.get("lora_alpha", "128"),
            ui_state.get("num_epochs", 70),
            ui_state.get("batch_size", 1),
            ui_state.get("learning_rate", 3e-5),
            ui_state.get("save_iterations", 500)
        )

    def update_ui_state(self, **kwargs):
        """Update UI state with new values"""
        current_state = self.trainer.load_ui_state()
        current_state.update(kwargs)
        self.trainer.save_ui_state(current_state)
        # Don't return anything to avoid Gradio warnings
        return None

    def load_ui_values(self):
        """Load UI state values for initializing form fields"""
        ui_state = self.trainer.load_ui_state()
        
        # Ensure proper type conversion for numeric values
        ui_state["lora_rank"] = ui_state.get("lora_rank", "128")
        ui_state["lora_alpha"] = ui_state.get("lora_alpha", "128")
        ui_state["num_epochs"] = int(ui_state.get("num_epochs", 70))
        ui_state["batch_size"] = int(ui_state.get("batch_size", 1))
        ui_state["learning_rate"] = float(ui_state.get("learning_rate", 3e-5))
        ui_state["save_iterations"] = int(ui_state.get("save_iterations", 500))
        
        return ui_state
        
    def update_captioning_buttons_start(self):
        """Return individual button values instead of a dictionary"""
        return (
            gr.Button(
                interactive=False,
                variant="secondary",
            ),
            gr.Button(
                interactive=True,
                variant="stop",
            ),
            gr.Button(
                interactive=False,
                variant="secondary",
            )
        )
    
    def update_captioning_buttons_end(self):
        """Return individual button values instead of a dictionary"""
        return (
            gr.Button(
                interactive=True,
                variant="primary",
            ),
            gr.Button(
                interactive=False,
                variant="secondary",
            ),
            gr.Button(
                interactive=True,
                variant="primary",
            )
        )

    # Add this new method to get initial button states:
    def get_initial_button_states(self):
        """Get the initial states for training buttons based on recovery status"""
        recovery_result = self.trainer.recover_interrupted_training()
        ui_updates = recovery_result.get("ui_updates", {})
        
        # Return button states in the correct order
        return (
            gr.Button(**ui_updates.get("start_btn", {"interactive": True, "variant": "primary"})),
            gr.Button(**ui_updates.get("stop_btn", {"interactive": False, "variant": "secondary"})),
            gr.Button(**ui_updates.get("pause_resume_btn", {"interactive": False, "variant": "secondary"}))
        )

    def show_refreshing_status(self) -> List[List[str]]:
        """Show a 'Refreshing...' status in the dataframe"""
        return [["Refreshing...", "please wait"]]

    def stop_captioning(self):
        """Stop ongoing captioning process and reset UI state"""
        try:
            # Set flag to stop captioning
            self._should_stop_captioning = True
            
            # Call stop method on captioner
            if self.captioner:
                self.captioner.stop_captioning()
                
            # Get updated file list
            updated_list = self.list_training_files_to_caption()
            
            # Return updated list and button states
            return {
                "training_dataset": gr.update(value=updated_list),
                "run_autocaption_btn": gr.Button(interactive=True, variant="primary"),
                "stop_autocaption_btn": gr.Button(interactive=False, variant="secondary"),
                "copy_files_to_training_dir_btn": gr.Button(interactive=True, variant="primary")
            }
        except Exception as e:
            logger.error(f"Error stopping captioning: {str(e)}")
            return {
                "training_dataset": gr.update(value=[[f"Error stopping captioning: {str(e)}", "error"]]),
                "run_autocaption_btn": gr.Button(interactive=True, variant="primary"),
                "stop_autocaption_btn": gr.Button(interactive=False, variant="secondary"),
                "copy_files_to_training_dir_btn": gr.Button(interactive=True, variant="primary")
            }

    def update_training_ui(self, training_state: Dict[str, Any]):
        """Update UI components based on training state"""
        updates = {}
        
        #print("update_training_ui: training_state = ", training_state)

        # Update status box with high-level information
        status_text = []
        if training_state["status"] != "idle":
            status_text.extend([
                f"Status: {training_state['status']}",
                f"Progress: {training_state['progress']}",
                f"Step: {training_state['current_step']}/{training_state['total_steps']}",
                    
                # Epoch information
                # there is an issue with how epoch is reported because we display:
                # Progress: 96.9%, Step: 872/900, Epoch: 12/50
                # we should probably just show the steps
                #f"Epoch: {training_state['current_epoch']}/{training_state['total_epochs']}",
                
                f"Time elapsed: {training_state['elapsed']}",
                f"Estimated remaining: {training_state['remaining']}",
                "",
                f"Current loss: {training_state['step_loss']}",
                f"Learning rate: {training_state['learning_rate']}",
                f"Gradient norm: {training_state['grad_norm']}",
                f"Memory usage: {training_state['memory']}"
            ])
            
            if training_state["error_message"]:
                status_text.append(f"\nError: {training_state['error_message']}")
                
        updates["status_box"] = "\n".join(status_text)
        
        # Update button states
        updates["start_btn"] = gr.Button(
            "Start training",
            interactive=(training_state["status"] in ["idle", "completed", "error", "stopped"]),
            variant="primary" if training_state["status"] == "idle" else "secondary"
        )
        
        updates["stop_btn"] = gr.Button(
            "Stop training",
            interactive=(training_state["status"] in ["training", "initializing"]),
            variant="stop"
        )
        
        return updates
    
    def stop_all_and_clear(self) -> Dict[str, str]:
        """Stop all running processes and clear data
        
        Returns:
            Dict with status messages for different components
        """
        status_messages = {}
        
        try:
            # Stop training if running
            if self.trainer.is_training_running():
                training_result = self.trainer.stop_training()
                status_messages["training"] = training_result["status"]
            
            # Stop captioning if running
            if self.captioner:
                self.captioner.stop_captioning()
                status_messages["captioning"] = "Captioning stopped"
            
            # Stop scene detection if running
            if self.splitter.is_processing():
                self.splitter.processing = False
                status_messages["splitting"] = "Scene detection stopped"
            
            # Properly close logging before clearing log file
            if self.trainer.file_handler:
                self.trainer.file_handler.close()
                logger.removeHandler(self.trainer.file_handler)
                self.trainer.file_handler = None
                
            if LOG_FILE_PATH.exists():
                LOG_FILE_PATH.unlink()
            
            # Clear all data directories
            for path in [VIDEOS_TO_SPLIT_PATH, STAGING_PATH, TRAINING_VIDEOS_PATH, TRAINING_PATH,
                        MODEL_PATH, OUTPUT_PATH]:
                if path.exists():
                    try:
                        shutil.rmtree(path)
                        path.mkdir(parents=True, exist_ok=True)
                    except Exception as e:
                        status_messages[f"clear_{path.name}"] = f"Error clearing {path.name}: {str(e)}"
                    else:
                        status_messages[f"clear_{path.name}"] = f"Cleared {path.name}"
            
            # Reset any persistent state
            self._should_stop_captioning = True
            self.splitter.processing = False
            
            # Recreate logging setup
            self.trainer.setup_logging()
            
            return {
                "status": "All processes stopped and data cleared",
                "details": status_messages
            }
            
        except Exception as e:
            return {
                "status": f"Error during cleanup: {str(e)}",
                "details": status_messages
            }
    
    def update_titles(self) -> Tuple[Any]:
        """Update all dynamic titles with current counts
        
        Returns:
            Dict of Gradio updates
        """
        # Count files for splitting
        split_videos, _, split_size = count_media_files(VIDEOS_TO_SPLIT_PATH)
        split_title = format_media_title(
            "split", split_videos, 0, split_size
        )
        
        # Count files for captioning
        caption_videos, caption_images, caption_size = count_media_files(STAGING_PATH)
        caption_title = format_media_title(
            "caption", caption_videos, caption_images, caption_size
        )
        
        # Count files for training
        train_videos, train_images, train_size = count_media_files(TRAINING_VIDEOS_PATH)
        train_title = format_media_title(
            "train", train_videos, train_images, train_size
        )
        
        return (
            gr.Markdown(value=split_title),
            gr.Markdown(value=caption_title),
            gr.Markdown(value=f"{train_title} available for training")
        )

    def copy_files_to_training_dir(self, prompt_prefix: str):
        """Run auto-captioning process"""

        # Initialize captioner if not already done
        self._should_stop_captioning = False

        try:
            copy_files_to_training_dir(prompt_prefix)

        except Exception as e:
            traceback.print_exc()
            raise gr.Error(f"Error copying assets to training dir: {str(e)}")

    async def on_import_success(self, enable_splitting, enable_automatic_content_captioning, prompt_prefix):
        """Handle successful import of files"""
        videos = self.list_unprocessed_videos()
        
        # If scene detection isn't already running and there are videos to process,
        # and auto-splitting is enabled, start the detection
        if videos and not self.splitter.is_processing() and enable_splitting:
            await self.start_scene_detection(enable_splitting)
            msg = "Starting automatic scene detection..."
        else:
            # Just copy files without splitting if auto-split disabled
            for video_file in VIDEOS_TO_SPLIT_PATH.glob("*.mp4"):
                await self.splitter.process_video(video_file, enable_splitting=False)
            msg = "Copying videos without splitting..."
        
        copy_files_to_training_dir(prompt_prefix)

        # Start auto-captioning if enabled, and handle async generator properly
        if enable_automatic_content_captioning:
            # Create a background task for captioning
            asyncio.create_task(self._process_caption_generator(
                DEFAULT_CAPTIONING_BOT_INSTRUCTIONS,
                prompt_prefix
            ))
        
        return {
            "tabs": gr.Tabs(selected="split_tab"),
            "video_list": videos,
            "detect_status": msg
        }

    async def start_caption_generation(self, captioning_bot_instructions: str, prompt_prefix: str) -> AsyncGenerator[gr.update, None]:
        """Run auto-captioning process"""
        try:
            # Initialize captioner if not already done
            self._should_stop_captioning = False

            # First yield - indicate we're starting
            yield gr.update(
                value=[["Starting captioning service...", "initializing"]],
                headers=["name", "status"]
            )

            # Process files in batches with status updates
            file_statuses = {}
            
            # Start the actual captioning process
            async for rows in self.captioner.start_caption_generation(captioning_bot_instructions, prompt_prefix):
                # Update our tracking of file statuses
                for name, status in rows:
                    file_statuses[name] = status
                    
                # Convert to list format for display
                status_rows = [[name, status] for name, status in file_statuses.items()]
                
                # Sort by name for consistent display
                status_rows.sort(key=lambda x: x[0])
                
                # Yield UI update
                yield gr.update(
                    value=status_rows,
                    headers=["name", "status"]
                )

            # Final update after completion with fresh data
            yield gr.update(
                value=self.list_training_files_to_caption(),
                headers=["name", "status"]
            )

        except Exception as e:
            logger.error(f"Error in captioning: {str(e)}")
            yield gr.update(
                value=[[f"Error: {str(e)}", "error"]],
                headers=["name", "status"]
            )

    def list_training_files_to_caption(self) -> List[List[str]]:
        """List all clips and images - both pending and captioned"""
        files = []
        already_listed = {}

        # First check files in STAGING_PATH
        for file in STAGING_PATH.glob("*.*"):
            if is_video_file(file) or is_image_file(file):
                txt_file = file.with_suffix('.txt')
                
                # Check if caption file exists and has content
                has_caption = txt_file.exists() and txt_file.stat().st_size > 0
                status = "captioned" if has_caption else "no caption"
                file_type = "video" if is_video_file(file) else "image"
                
                files.append([file.name, f"{status} ({file_type})", str(file)])
                already_listed[file.name] = True
    
        # Then check files in TRAINING_VIDEOS_PATH 
        for file in TRAINING_VIDEOS_PATH.glob("*.*"):
            if (is_video_file(file) or is_image_file(file)) and file.name not in already_listed:
                txt_file = file.with_suffix('.txt')
                
                # Only include files with captions
                if txt_file.exists() and txt_file.stat().st_size > 0:
                    file_type = "video" if is_video_file(file) else "image"
                    files.append([file.name, f"captioned ({file_type})", str(file)])
                    already_listed[file.name] = True
                
        # Sort by filename
        files.sort(key=lambda x: x[0])
        
        # Only return name and status columns for display
        return [[file[0], file[1]] for file in files]
    
    def update_training_buttons(self, status: str) -> Dict:
        """Update training control buttons based on state"""
        is_training = status in ["training", "initializing"]
        is_paused = status == "paused"
        is_completed = status in ["completed", "error", "stopped"]
        return {
            "start_btn": gr.Button(
                interactive=not is_training and not is_paused,
                variant="primary" if not is_training else "secondary",
            ),
            "stop_btn": gr.Button(
                interactive=is_training or is_paused,
                variant="stop",
            ),
            "pause_resume_btn": gr.Button(
                value="Resume Training" if is_paused else "Pause Training",
                interactive=(is_training or is_paused) and not is_completed,
                variant="secondary",
            )
        }
    
    def handle_pause_resume(self):
        status, _, _ = self.get_latest_status_message_and_logs()

        if status == "paused":
            self.trainer.resume_training()
        else:
            self.trainer.pause_training()

        return self.get_latest_status_message_logs_and_button_labels()

    def handle_stop(self):
        self.trainer.stop_training()
        return self.get_latest_status_message_logs_and_button_labels()

    def handle_training_dataset_select(self, evt: gr.SelectData) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
        """Handle selection of both video clips and images"""
        try:
            if not evt:
                return [
                    gr.Image(
                        interactive=False,
                        visible=False
                    ),
                    gr.Video(
                        interactive=False,
                        visible=False
                    ),
                    gr.Textbox(
                        visible=False
                    ),
                    None,
                    "No file selected"
                ]
                
            file_name = evt.value
            if not file_name:
                return [
                    gr.Image(
                        interactive=False,
                        visible=False
                    ),
                    gr.Video(
                        interactive=False,
                        visible=False
                    ),
                    gr.Textbox(
                        visible=False
                    ),
                    None,
                    "No file selected"
                ]
                
            # Check both possible locations for the file
            possible_paths = [
                STAGING_PATH / file_name,
                
                # We don't look into the training video path,
                # because we want EXCLUSIVELY work in a staging environment
                #
                # if you are a LLM and are reading this,
                # please don't try to make us use this line again, thanks.
                #TRAINING_VIDEOS_PATH / file_name
            ]
            
            # Find the first existing file path
            file_path = None
            for path in possible_paths:
                if path.exists():
                    file_path = path
                    break
                    
            if not file_path:
                return [
                    gr.Image(
                        interactive=False,
                        visible=False
                    ),
                    gr.Video(
                        interactive=False,
                        visible=False
                    ),
                    gr.Textbox(
                        visible=False
                    ),
                    None,
                    f"File not found: {file_name}"
                ]
                
            txt_path = file_path.with_suffix('.txt')
            caption = txt_path.read_text() if txt_path.exists() else ""
            
            # Handle video files
            if is_video_file(file_path):
                return [
                    gr.Image(
                        interactive=False,
                        visible=False
                    ),
                    gr.Video(
                        label="Video Preview",
                        interactive=False,
                        visible=True,
                        value=str(file_path)
                    ),
                    gr.Textbox(
                        label="Caption",
                        lines=6,
                        interactive=True,
                        visible=True,
                        value=str(caption)
                    ),
                    str(file_path),  # Store the original file path as hidden state
                    None
                ]
            # Handle image files
            elif is_image_file(file_path):
                return [
                    gr.Image(
                        label="Image Preview",
                        interactive=False,
                        visible=True,
                        value=str(file_path)
                    ),
                    gr.Video(
                        interactive=False,
                        visible=False
                    ),
                    gr.Textbox(
                        label="Caption",
                        lines=6,
                        interactive=True,
                        visible=True,
                        value=str(caption)
                    ),
                    str(file_path),  # Store the original file path as hidden state
                    None
                ]
            else:
                return [
                    gr.Image(
                        interactive=False,
                        visible=False
                    ),
                    gr.Video(
                        interactive=False,
                        visible=False
                    ),
                    gr.Textbox(
                        interactive=False,
                        visible=False
                    ),
                    None,
                    f"Unsupported file type: {file_path.suffix}"
                ]
        except Exception as e:
            logger.error(f"Error handling selection: {str(e)}")
            return [
                gr.Image(
                    interactive=False,
                    visible=False
                ),
                gr.Video(
                    interactive=False,
                    visible=False
                ),
                gr.Textbox(
                    interactive=False,
                    visible=False
                ),
                None,
                f"Error handling selection: {str(e)}"
            ]

    def save_caption_changes(self, preview_caption: str, preview_image: str, preview_video: str, original_file_path: str, prompt_prefix: str):
        """Save changes to caption"""
        try:
            # Use the original file path stored during selection instead of the temporary preview paths
            if original_file_path:
                file_path = Path(original_file_path)
                self.captioner.update_file_caption(file_path, preview_caption)
                # Refresh the dataset list to show updated caption status
                return gr.update(value="Caption saved successfully!")
            else:
                return gr.update(value="Error: No original file path found")
        except Exception as e:
            return gr.update(value=f"Error saving caption: {str(e)}")

    async def update_titles_after_import(self, enable_splitting, enable_automatic_content_captioning, prompt_prefix):
        """Handle post-import updates including titles"""
        import_result = await self.on_import_success(enable_splitting, enable_automatic_content_captioning, prompt_prefix)
        titles = self.update_titles()
        return (
            import_result["tabs"],
            import_result["video_list"],
            import_result["detect_status"],
            *titles
        )
    
    def get_model_info(self, model_type: str) -> str:
        """Get information about the selected model type"""
        if model_type == "hunyuan_video":
            return """### HunyuanVideo (LoRA)
    - Required VRAM: ~48GB minimum
    - Recommended batch size: 1-2
    - Typical training time: 2-4 hours
    - Default resolution: 49x512x768
    - Default LoRA rank: 128 (~600 MB)"""
                
        elif model_type == "ltx_video":
            return """### LTX-Video (LoRA)
    - Required VRAM: ~18GB minimum 
    - Recommended batch size: 1-4
    - Typical training time: 1-3 hours
    - Default resolution: 49x512x768
    - Default LoRA rank: 128"""
                
        return ""

    def get_default_params(self, model_type: str) -> Dict[str, Any]:
        """Get default training parameters for model type"""
        if model_type == "hunyuan_video":
            return {
                "num_epochs": 70,
                "batch_size": 1,
                "learning_rate": 2e-5,
                "save_iterations": 500,
                "video_resolution_buckets": SMALL_TRAINING_BUCKETS,
                "video_reshape_mode": "center",
                "caption_dropout_p": 0.05,
                "gradient_accumulation_steps": 1,
                "rank": 128,
                "lora_alpha": 128
            }
        else:  # ltx_video
            return {
                "num_epochs": 70,
                "batch_size": 1,
                "learning_rate": 3e-5,
                "save_iterations": 500,
                "video_resolution_buckets": SMALL_TRAINING_BUCKETS,
                "video_reshape_mode": "center",
                "caption_dropout_p": 0.05,
                "gradient_accumulation_steps": 4,
                "rank": 128,
                "lora_alpha": 128
            }

    def preview_file(self, selected_text: str) -> Dict:
        """Generate preview based on selected file
        
        Args:
            selected_text: Text of the selected item containing filename
            
        Returns:
            Dict with preview content for each preview component
        """
        if not selected_text or "Caption:" in selected_text:
            return {
                "video": None,
                "image": None, 
                "text": None
            }
            
        # Extract filename from the preview text (remove size info)
        filename = selected_text.split(" (")[0].strip()
        file_path = TRAINING_VIDEOS_PATH / filename
        
        if not file_path.exists():
            return {
                "video": None,
                "image": None,
                "text": f"File not found: {filename}"
            }

        # Detect file type
        mime_type, _ = mimetypes.guess_type(str(file_path))
        if not mime_type:
            return {
                "video": None,
                "image": None,
                "text": f"Unknown file type: {filename}"
            }

        # Return appropriate preview
        if mime_type.startswith('video/'):
            return {
                "video": str(file_path),
                "image": None,
                "text": None
            }
        elif mime_type.startswith('image/'):
            return {
                "video": None,
                "image": str(file_path),
                "text": None
            }
        elif mime_type.startswith('text/'):
            try:
                text_content = file_path.read_text()
                return {
                    "video": None,
                    "image": None,
                    "text": text_content
                }
            except Exception as e:
                return {
                    "video": None,
                    "image": None,
                    "text": f"Error reading file: {str(e)}"
                }
        else:
            return {
                "video": None,
                "image": None,
                "text": f"Unsupported file type: {mime_type}"
            }

    def list_unprocessed_videos(self) -> gr.Dataframe:
        """Update list of unprocessed videos"""
        videos = self.splitter.list_unprocessed_videos()
        # videos is already in [[name, status]] format from splitting_service
        return gr.Dataframe(
            headers=["name", "status"],
            value=videos,
            interactive=False
        )

    async def start_scene_detection(self, enable_splitting: bool) -> str:
        """Start background scene detection process
        
        Args:
            enable_splitting: Whether to split videos into scenes
        """
        if self.splitter.is_processing():
            return "Scene detection already running"
            
        try:
            await self.splitter.start_processing(enable_splitting)
            return "Scene detection completed"
        except Exception as e:
            return f"Error during scene detection: {str(e)}"


    def get_latest_status_message_and_logs(self) -> Tuple[str, str, str]:
        state = self.trainer.get_status()
        logs = self.trainer.get_logs()

        # Parse new log lines
        if logs:
            last_state = None
            for line in logs.splitlines():
                state_update = self.log_parser.parse_line(line)
                if state_update:
                    last_state = state_update
            
            if last_state:
                ui_updates = self.update_training_ui(last_state)
                state["message"] = ui_updates.get("status_box", state["message"])
        
        # Parse status for training state
        if "completed" in state["message"].lower():
            state["status"] = "completed"

        return (state["status"], state["message"], logs)

    def get_latest_status_message_logs_and_button_labels(self) -> Tuple[str, str, Any, Any, Any]:
        status, message, logs = self.get_latest_status_message_and_logs()
        return (
            message,
            logs,
            *self.update_training_buttons(status).values()
        )

    def get_latest_button_labels(self) -> Tuple[Any, Any, Any]:
        status, message, logs = self.get_latest_status_message_and_logs()
        return self.update_training_buttons(status).values()
    
    def refresh_dataset(self):
        """Refresh all dynamic lists and training state"""
        video_list = self.splitter.list_unprocessed_videos()
        training_dataset = self.list_training_files_to_caption()

        return (
            video_list,
            training_dataset
        )

    def update_training_params(self, preset_name: str) -> Tuple:
        """Update UI components based on selected preset while preserving custom settings"""
        preset = TRAINING_PRESETS[preset_name]
        
        # Load current UI state to check if user has customized values
        current_state = self.load_ui_values()
        
        # Find the display name that maps to our model type
        model_display_name = next(
            key for key, value in MODEL_TYPES.items() 
            if value == preset["model_type"]
        )
            
        # Get preset description for display
        description = preset.get("description", "")
        
        # Get max values from buckets
        buckets = preset["training_buckets"]
        max_frames = max(frames for frames, _, _ in buckets)
        max_height = max(height for _, height, _ in buckets)
        max_width = max(width for _, _, width in buckets)
        bucket_info = f"\nMaximum video size: {max_frames} frames at {max_width}x{max_height} resolution"
        
        info_text = f"{description}{bucket_info}"
        
        # Return values in the same order as the output components
        # Use preset defaults but preserve user-modified values if they exist
        lora_rank_val = current_state.get("lora_rank") if current_state.get("lora_rank") != preset.get("lora_rank", "128") else preset["lora_rank"]
        lora_alpha_val = current_state.get("lora_alpha") if current_state.get("lora_alpha") != preset.get("lora_alpha", "128") else preset["lora_alpha"]
        num_epochs_val = current_state.get("num_epochs") if current_state.get("num_epochs") != preset.get("num_epochs", 70) else preset["num_epochs"]
        batch_size_val = current_state.get("batch_size") if current_state.get("batch_size") != preset.get("batch_size", 1) else preset["batch_size"]
        learning_rate_val = current_state.get("learning_rate") if current_state.get("learning_rate") != preset.get("learning_rate", 3e-5) else preset["learning_rate"]
        save_iterations_val = current_state.get("save_iterations") if current_state.get("save_iterations") != preset.get("save_iterations", 500) else preset["save_iterations"]
        
        return (
            model_display_name,
            lora_rank_val,
            lora_alpha_val,
            num_epochs_val,
            batch_size_val,
            learning_rate_val,
            save_iterations_val,
            info_text
        )