File size: 5,287 Bytes
89bbef2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
246c64e
89bbef2
246c64e
 
89bbef2
246c64e
 
 
 
 
 
 
 
 
 
 
89bbef2
 
 
ab45a2c
89bbef2
 
 
 
ab45a2c
 
 
 
 
 
89bbef2
 
 
 
 
 
 
 
 
 
 
 
 
 
ab45a2c
89bbef2
 
 
 
 
 
246c64e
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
"""
Upload tab for Video Model Studio UI.
Handles manual file uploads for videos, images, and archives.
"""

import gradio as gr
import logging
from pathlib import Path
from typing import Dict, Any, Optional

from vms.utils import BaseTab

logger = logging.getLogger(__name__)

class UploadTab(BaseTab):
    """Upload tab for manual file uploads"""
    
    def __init__(self, app_state):
        super().__init__(app_state)
        self.id = "upload_tab"
        self.title = "Manual Upload"
        # Initialize the components dictionary with None values for expected shared components
        if "components" not in self.__dict__:
            self.components = {}
        self.components["import_status"] = None
        self.components["enable_automatic_video_split"] = None
        self.components["enable_automatic_content_captioning"] = None
    
    def create(self, parent=None) -> gr.Tab:
        """Create the Upload tab UI components"""
        with gr.Tab(self.title, id=self.id) as tab:
            with gr.Column():
                with gr.Row():
                    gr.Markdown("## Manual upload of video files")
                
                with gr.Row():
                    with gr.Column():
                        with gr.Row():
                            gr.Markdown("You can upload either:")
                        with gr.Row():
                            gr.Markdown("- A single MP4 video file")
                        with gr.Row():
                            gr.Markdown("- A ZIP archive containing multiple videos/images and optional caption files")
                        with gr.Row():
                            gr.Markdown("- A WebDataset shard (.tar file)")
                        with gr.Row():
                            gr.Markdown("- A ZIP archive containing WebDataset shards (.tar files)")
                    with gr.Column():
                        with gr.Row():
                            self.components["files"] = gr.Files(
                                label="Upload Images, Videos, ZIP or WebDataset",
                                file_types=[".jpg", ".jpeg", ".png", ".webp", ".webp", ".avif", ".heic", ".mp4", ".zip", ".tar"],
                                type="filepath"
                            )
            
            return tab
    
    def connect_events(self) -> None:
        """Connect event handlers to UI components"""
        # Check if required shared components exist before connecting events
        if not self.components.get("import_status"):
            logger.warning("import_status component is not set in UploadTab")
            return
        
        # File upload event with enable_splitting parameter
        upload_event = self.components["files"].upload(
            fn=self.app.importing.process_uploaded_files,
            inputs=[self.components["files"], self.components["enable_automatic_video_split"]],
            outputs=[self.components["import_status"]]
        ).success(
            fn=self.app.tabs["import_tab"].on_import_success,
            inputs=[
                self.components["enable_automatic_video_split"],
                self.components["enable_automatic_content_captioning"],
                self.app.tabs["caption_tab"].components["custom_prompt_prefix"]
            ],
            outputs=[
                self.app.project_tabs_component,
                self.components["import_status"]
            ]
        )
        
        # Only add success handler if all required components exist
        if hasattr(self.app.tabs, "import_tab") and hasattr(self.app.tabs, "caption_tab") and hasattr(self.app.tabs, "train_tab"):
            
            # Get required components for success handler
            try:
                # If the components are missing, this will raise an AttributeError
                if hasattr(self.app, "project_tabs_component"):
                    tabs_component = self.app.project_tabs_component
                else:
                    logger.warning("project_tabs_component not found in app, using None for tab switching")
                    tabs_component = None
                    
                caption_title = self.app.tabs["caption_tab"].components["caption_title"]
                train_title = self.app.tabs["train_tab"].components["train_title"]
                custom_prompt_prefix = self.app.tabs["caption_tab"].components["custom_prompt_prefix"]
                
                # Add success handler
                upload_event.success(
                    fn=self.app.tabs["import_tab"].update_titles_after_import,
                    inputs=[
                        self.components["enable_automatic_video_split"], 
                        self.components["enable_automatic_content_captioning"], 
                        custom_prompt_prefix
                    ],
                    outputs=[
                        tabs_component,
                        self.components["import_status"],
                        caption_title,
                        train_title
                    ]
                )
            except (AttributeError, KeyError) as e:
                logger.error(f"Error connecting event handlers in UploadTab: {str(e)}")
                # Continue without the success handler