fix for the custom prompt prefix

vms/ui/project/services/importing/file_upload.py

@@ -22,7 +22,7 @@ logger = logging.getLogger(__name__)
 class FileUploadHandler:
     """Handles processing of uploaded files"""
     
-    def process_uploaded_files(self, file_paths: List[str], enable_splitting: bool) -> str:
+    def process_uploaded_files(self, file_paths: List[str], enable_splitting: bool, custom_prompt_prefix: str = None) -> str:
         """Process uploaded file (ZIP, TAR, MP4, or image)
         
         Args:
@@ -48,7 +48,7 @@ class FileUploadHandler:
                 file_ext = file_path.suffix.lower()
 
                 if file_ext == '.zip':
-                    return self.process_zip_file(file_path, enable_splitting)
+                    return self.process_zip_file(file_path, enable_splitting, custom_prompt_prefix)
                 elif file_ext == '.tar':
                     return self.process_tar_file(file_path, enable_splitting)
                 elif file_ext == '.mp4' or file_ext == '.webm':
@@ -63,7 +63,7 @@ class FileUploadHandler:
                 logger.error(f"Error processing file {file_path}: {str(e)}", exc_info=True)
                 raise gr.Error(f"Error processing file: {str(e)}")
 
-    def process_zip_file(self, file_path: Path, enable_splitting: bool) -> str:
+    def process_zip_file(self, file_path: Path, enable_splitting: bool, custom_prompt_prefix: str = None) -> str:
         """Process uploaded ZIP file containing media files or WebDataset tar files
         
         Args:
@@ -138,7 +138,7 @@ class FileUploadHandler:
                                     logger.info(f"Copied caption file for {file}")
                                 elif is_image_file(file_path):
                                     caption = txt_path.read_text()
-                                    caption = add_prefix_to_caption(caption, DEFAULT_PROMPT_PREFIX)
+                                    caption = add_prefix_to_caption(caption, custom_prompt_prefix or DEFAULT_PROMPT_PREFIX)
                                     target_path.with_suffix('.txt').write_text(caption)
                                     logger.info(f"Processed caption for {file}")
                                     

vms/ui/project/services/importing/hub_dataset.py

@@ -169,7 +169,8 @@ class HubDatasetBrowser:
         dataset_id: str, 
         file_type: str, 
         enable_splitting: bool,
-        progress_callback: Optional[Callable] = None
+        progress_callback: Optional[Callable] = None,
+        custom_prompt_prefix: str = None
     ) -> str:
         """Download all files of a specific type from the dataset
         
@@ -329,7 +330,8 @@ class HubDatasetBrowser:
         self, 
         dataset_id: str, 
         enable_splitting: bool,
-        progress_callback: Optional[Callable] = None
+        progress_callback: Optional[Callable] = None,
+        custom_prompt_prefix: str = None
     ) -> Tuple[str, str]:
         """Download a dataset and process its video/image content
         
@@ -555,7 +557,7 @@ class HubDatasetBrowser:
                                     txt_path = file_path.with_suffix('.txt')
                                     if txt_path.exists():
                                         caption = txt_path.read_text()
-                                        caption = add_prefix_to_caption(caption, DEFAULT_PROMPT_PREFIX)
+                                        caption = add_prefix_to_caption(caption, custom_prompt_prefix or DEFAULT_PROMPT_PREFIX)
                                         target_path.with_suffix('.txt').write_text(caption)
                                         logger.info(f"Processed caption for {file_path.name}")
                                     

vms/ui/project/services/importing/import_service.py

@@ -28,7 +28,7 @@ class ImportingService:
         self.youtube_handler = YouTubeDownloader()
         self.hub_browser = HubDatasetBrowser(self.hf_api)
     
-    def process_uploaded_files(self, file_paths: List[str], enable_splitting: bool) -> str:
+    def process_uploaded_files(self, file_paths: List[str], enable_splitting: bool, custom_prompt_prefix: str = None) -> str:
         """Process uploaded file (ZIP, TAR, MP4, or image)
         
         Args:
@@ -45,7 +45,7 @@ class ImportingService:
         
         print(f"process_uploaded_files(..., enable_splitting = {enable_splitting:})")
         print(f"process_uploaded_files: calling self.file_handler.process_uploaded_files")
-        return self.file_handler.process_uploaded_files(file_paths, enable_splitting)
+        return self.file_handler.process_uploaded_files(file_paths, enable_splitting, custom_prompt_prefix)
     
     def download_youtube_video(self, url: str, enable_splitting: bool, progress=None) -> str:
         """Download a video from YouTube
@@ -86,7 +86,8 @@ class ImportingService:
         self, 
         dataset_id: str, 
         enable_splitting: bool,
-        progress_callback: Optional[Callable] = None
+        progress_callback: Optional[Callable] = None,
+        custom_prompt_prefix: str = None
     ) -> Tuple[str, str]:
         """Download a dataset and process its video/image content
         
@@ -98,14 +99,15 @@ class ImportingService:
         Returns:
             Tuple of (loading_msg, status_msg)
         """
-        return await self.hub_browser.download_dataset(dataset_id, enable_splitting, progress_callback)
+        return await self.hub_browser.download_dataset(dataset_id, enable_splitting, progress_callback, custom_prompt_prefix)
     
     async def download_file_group(
         self, 
         dataset_id: str, 
         file_type: str, 
         enable_splitting: bool,
-        progress_callback: Optional[Callable] = None
+        progress_callback: Optional[Callable] = None,
+        custom_prompt_prefix: str = None
     ) -> str:
         """Download a group of files (videos or WebDatasets)
         
@@ -118,4 +120,4 @@ class ImportingService:
         Returns:
             Status message
         """
-        return await self.hub_browser.download_file_group(dataset_id, file_type, enable_splitting, progress_callback)
+        return await self.hub_browser.download_file_group(dataset_id, file_type, enable_splitting, progress_callback, custom_prompt_prefix)

vms/ui/project/services/training.py

@@ -579,6 +579,7 @@ class TrainingService:
         precomputation_items: int = DEFAULT_PRECOMPUTATION_ITEMS,
         lr_warmup_steps: int = DEFAULT_NB_LR_WARMUP_STEPS,
         progress: Optional[gr.Progress] = None,
+        custom_prompt_prefix: Optional[str] = None,
     ) -> Tuple[str, str]:
         """Start training with finetrainers"""
         
@@ -669,16 +670,10 @@ class TrainingService:
             else:
                 flow_weighting_scheme = "logit_normal"
 
-            # Get the custom prompt prefix from the tabs
-            custom_prompt_prefix = None
-            if hasattr(self, 'app') and self.app is not None:
-                if hasattr(self.app, 'tabs') and 'caption_tab' in self.app.tabs:
-                    if hasattr(self.app.tabs['caption_tab'], 'components') and 'custom_prompt_prefix' in self.app.tabs['caption_tab'].components:
-                        # Get the value and clean it
-                        prefix = self.app.tabs['caption_tab'].components['custom_prompt_prefix'].value
-                        if prefix:
-                            # Clean the prefix - remove trailing comma, space or comma+space
-                            custom_prompt_prefix = prefix.rstrip(', ')
+            # Use the custom prompt prefix passed as parameter
+            # Clean the prefix - remove trailing comma, space or comma+space
+            if custom_prompt_prefix:
+                custom_prompt_prefix = custom_prompt_prefix.rstrip(', ')
 
             # Create a proper dataset configuration JSON file
             dataset_config_file = self.app.output_path / "dataset_config.json"

vms/ui/project/tabs/import_tab/hub_tab.py

@@ -267,7 +267,7 @@ class HubTab(BaseTab):
                 ""                                             # status_output
             )
 
-    async def _download_with_progress(self, dataset_id, file_type, enable_splitting, progress_callback):
+    async def _download_with_progress(self, dataset_id, file_type, enable_splitting, progress_callback, custom_prompt_prefix=None):
         """Wrapper for download_file_group that integrates with progress tracking"""
         try:
             # Set up the progress callback adapter
@@ -289,7 +289,8 @@ class HubTab(BaseTab):
                 dataset_id, 
                 file_type, 
                 enable_splitting,
-                progress_callback=progress_adapter
+                progress_callback=progress_adapter,
+                custom_prompt_prefix=custom_prompt_prefix
             )
             
             return result
@@ -298,7 +299,7 @@ class HubTab(BaseTab):
             logger.error(f"Error in download with progress: {str(e)}", exc_info=True)
             return f"Error: {str(e)}"
 
-    def download_file_group(self, dataset_id: str, enable_splitting: bool, file_type: str, progress=gr.Progress()) -> Tuple:
+    def download_file_group(self, dataset_id: str, enable_splitting: bool, file_type: str, custom_prompt_prefix: str = None, progress=gr.Progress()) -> Tuple:
         """Handle download of a group of files (videos or WebDatasets) with progress tracking"""
         try:
             if not dataset_id:
@@ -323,7 +324,8 @@ class HubTab(BaseTab):
                 dataset_id, 
                 file_type, 
                 enable_splitting,
-                progress
+                progress,
+                custom_prompt_prefix
             ))
             
             # When download is complete, update the UI

vms/ui/project/tabs/import_tab/upload_tab.py

@@ -65,7 +65,7 @@ class UploadTab(BaseTab):
         # 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"]],
+            inputs=[self.components["files"], self.components["enable_automatic_video_split"], self.app.tabs["caption_tab"].components["custom_prompt_prefix"]],
             outputs=[self.components["import_status"]]
         ).success(
             fn=self.app.tabs["import_tab"].on_import_success,

vms/ui/project/tabs/train_tab.py

@@ -906,6 +906,13 @@ For image-to-video tasks, 'index' (usually with index 0) is most common as it co
         precomputation_items = int(self.components["precomputation_items"].value)
         lr_warmup_steps = int(self.components["lr_warmup_steps"].value)
         
+        # Get custom prompt prefix from caption tab
+        custom_prompt_prefix = None
+        if hasattr(self.app, 'tabs') and 'caption_tab' in self.app.tabs:
+            caption_tab = self.app.tabs['caption_tab']
+            if hasattr(caption_tab, 'components') and 'custom_prompt_prefix' in caption_tab.components:
+                custom_prompt_prefix = caption_tab.components['custom_prompt_prefix'].value
+        
         # Start training (it will automatically use the checkpoint if provided)
         try:
             return self.app.training.start_training(
@@ -924,7 +931,8 @@ For image-to-video tasks, 'index' (usually with index 0) is most common as it co
                 precomputation_items=precomputation_items,
                 lr_warmup_steps=lr_warmup_steps,
                 progress=progress,
-                pretrained_lora_path=pretrained_lora_path
+                pretrained_lora_path=pretrained_lora_path,
+                custom_prompt_prefix=custom_prompt_prefix
             )
         except Exception as e:
             logger.exception("Error starting training")