File size: 8,389 Bytes
b7cc217
 
 
 
 
 
 
6fff6df
b7cc217
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6fff6df
 
 
 
 
 
 
 
 
 
b7cc217
 
 
 
2ba9257
b7cc217
 
 
 
 
 
 
 
 
 
 
2ba9257
b7cc217
 
6fff6df
b7cc217
6fff6df
 
 
 
 
 
 
 
 
 
 
 
 
 
b7cc217
 
 
6fff6df
 
 
 
b7cc217
 
 
 
2ba9257
b7cc217
 
 
 
 
2ba9257
 
b7cc217
 
 
 
 
 
2ba9257
b7cc217
 
 
 
2ba9257
 
 
b7cc217
 
 
 
 
 
 
 
 
 
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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
"""
Trained tab for Models view in Video Model Studio UI
"""

import gradio as gr
import logging
from typing import Dict, Any, List, Optional, Tuple
from gradio_modal import Modal

from vms.utils.base_tab import BaseTab

logger = logging.getLogger(__name__)

class TrainedTab(BaseTab):
    """Tab for managing trained models"""
    
    def __init__(self, app_state):
        super().__init__(app_state)
        self.id = "trained_tab"
        self.title = "Trained"
    
    def create(self, parent=None) -> gr.TabItem:
        """Create the Trained tab UI components"""
        with gr.TabItem(self.title, id=self.id) as tab:
            with gr.Row():
                gr.Markdown("## Completed Models")
            
            # List for displaying models
            with gr.Column() as models_container:
                self.components["models_container"] = models_container
                self.components["no_models_message"] = gr.Markdown(
                    "No trained models found yet. Train a model to see it here.",
                    visible=False
                )
                
                # Placeholder for model rows - will be filled dynamically
                self.components["model_rows"] = []
                
                # Initial load of models
                self.refresh_models()
                
        return tab
    
    def connect_events(self) -> None:
        """Connect event handlers to UI components"""
        # Add auto-refresh timer
        refresh_timer = gr.Timer(interval=5)  # Refresh every 5 seconds
        refresh_timer.tick(
            fn=self.refresh_models,
            inputs=[],
            outputs=[self.components["models_container"]]
        )
    
    def refresh_models(self) -> gr.Column:
        """Refresh the list of trained models"""
        # Get models from service
        trained_models = self.app.models_tab.models_service.get_trained_models()
        
        # Create a new Column to replace the existing one
        with gr.Column() as new_container:
            if not trained_models:
                gr.Markdown("No trained models found yet. Train a model to see it here.")
            else:
                gr.Markdown(f"Found {len(trained_models)} trained models:")
                
                # Create headers
                with gr.Row(equal_height=True):
                    with gr.Column(scale=2, min_width=20):
                        gr.Markdown("### Model ID")
                    with gr.Column(scale=2, min_width=20):
                        gr.Markdown("### Model Type")
                    with gr.Column(scale=1, min_width=20):
                        gr.Markdown("### Completed")
                    with gr.Column(scale=2, min_width=20):
                        gr.Markdown("### Actions")
                
                # Create a row for each model
                for model in trained_models:
                    with gr.Row(equal_height=True):
                        with gr.Column(scale=2, min_width=20):
                            gr.Markdown(model.id[:8] + "...")
                        with gr.Column(scale=2, min_width=20):
                            gr.Markdown(model.model_display_name or "Unknown")
                        with gr.Column(scale=1, min_width=20):
                            gr.Markdown(model.updated_at.strftime("%Y-%m-%d"))
                            
                        with gr.Column(scale=2, min_width=20):
                            with gr.Row():
                                preview_btn = gr.Button("πŸ‘οΈ Preview", size="sm")
                                download_btn = gr.Button("πŸ’Ύ Download", size="sm")
                                publish_btn = gr.Button("🌐 Publish", size="sm")
                                delete_btn = gr.Button("πŸ—‘οΈ Delete", size="sm", variant="stop")
                                
                                # Create a modal for this specific model deletion
                                with Modal(visible=False) as delete_modal:
                                    gr.Markdown("## ⚠️ Confirm Deletion")
                                    gr.Markdown(f"Are you sure you want to delete model {model.id[:8]}...?")
                                    gr.Markdown("This action cannot be undone!")
                                    
                                    with gr.Row():
                                        cancel_btn = gr.Button("🫒 No, cancel", variant="secondary")
                                        confirm_btn = gr.Button("🚨 Yes, delete", variant="primary")

                                # Connect event handlers for this specific model
                                preview_btn.click(
                                    fn=lambda model_id=model.id: self.preview_model(model_id),
                                    inputs=[],
                                    outputs=[self.app.main_tabs]
                                )
                                
                                download_btn.click(
                                    fn=lambda model_id=model.id: self.download_model(model_id),
                                    inputs=[],
                                    outputs=[]
                                )
                                
                                publish_btn.click(
                                    fn=lambda model_id=model.id: self.publish_model(model_id),
                                    inputs=[],
                                    outputs=[self.app.main_tabs]
                                )
                                
                                # Connect delete button to show modal
                                delete_btn.click(
                                    fn=lambda: Modal(visible=True),
                                    inputs=[],
                                    outputs=[delete_modal]
                                )
                                
                                # Connect cancel button to hide modal
                                cancel_btn.click(
                                    fn=lambda: Modal(visible=False),
                                    inputs=[],
                                    outputs=[delete_modal]
                                )
                                
                                # Connect confirm button to delete and hide modal
                                confirm_btn.click(
                                    fn=lambda model_id=model.id: self.delete_model(model_id),
                                    inputs=[],
                                    outputs=[new_container]
                                ).then(
                                    fn=lambda: Modal(visible=False),
                                    inputs=[],
                                    outputs=[delete_modal]
                                )
        
        return new_container
    
    def preview_model(self, model_id: str) -> gr.Tabs:
        """Open model preview"""
        if self.app:
            # Switch to project view with this model
            self.app.switch_project(model_id)
            # Set main tab to Project (index 0)
            return self.app.main_tabs.update(selected=0)
            # TODO: Navigate to preview tab
            
    def download_model(self, model_id: str) -> None:
        """Download model weights"""
        # TODO: Implement file download
        gr.Info(f"Download for model {model_id[:8]}... is not yet implemented")
        
    def publish_model(self, model_id: str) -> gr.Tabs:
        """Publish model to Hugging Face Hub"""
        if self.app:
            # Switch to the selected model project
            self.app.switch_project(model_id)
            # Navigate to the main project tab
            return self.app.main_tabs.update(selected=0)
            # TODO: Navigate to publish tab
        
    def delete_model(self, model_id: str) -> gr.Column:
        """Delete a model and refresh the list"""
        if self.app and self.app.models_tab.models_service.delete_model(model_id):
            gr.Info(f"Model {model_id[:8]}... deleted successfully")
        else:
            gr.Warning(f"Failed to delete model {model_id[:8]}...")
            
        # Refresh the models list
        return self.refresh_models()