jbilcke-hf HF Staff commited on
Commit
82a6631
Β·
1 Parent(s): f407698

working on a fix for manual captions

Browse files
app.py CHANGED
@@ -30,18 +30,19 @@ import tempfile
30
  import zipfile
31
  from typing import Any, Optional, Dict, List, Union, Tuple
32
  from typing import AsyncGenerator
33
- from training_service import TrainingService
34
- from captioning_service import CaptioningService
35
- from splitting_service import SplittingService
36
- from import_service import ImportService
37
- from config import (
 
38
  STORAGE_PATH, VIDEOS_TO_SPLIT_PATH, STAGING_PATH,
39
  TRAINING_PATH, LOG_FILE_PATH, TRAINING_PRESETS, TRAINING_VIDEOS_PATH, MODEL_PATH, OUTPUT_PATH, DEFAULT_CAPTIONING_BOT_INSTRUCTIONS,
40
  DEFAULT_PROMPT_PREFIX, HF_API_TOKEN, ASK_USER_TO_DUPLICATE_SPACE, MODEL_TYPES, SMALL_TRAINING_BUCKETS
41
  )
42
- from utils import make_archive, count_media_files, format_media_title, is_image_file, is_video_file, validate_model_repo, format_time
43
- from finetrainers_utils import copy_files_to_training_dir, prepare_finetrainers_dataset
44
- from training_log_parser import TrainingLogParser
45
 
46
  logger = logging.getLogger(__name__)
47
  logger.setLevel(logging.INFO)
@@ -462,10 +463,8 @@ class VideoTrainerUI:
462
  full_caption = preview_caption
463
 
464
  path = Path(preview_video if preview_video else preview_image)
465
- if path.suffix == '.txt':
466
- self.trainer.update_file_caption(path.with_suffix(''), full_caption)
467
- else:
468
- self.trainer.update_file_caption(path, full_caption)
469
  return gr.update(value="Caption saved successfully!")
470
  except Exception as e:
471
  return gr.update(value=f"Error saving caption: {str(e)}")
 
30
  import zipfile
31
  from typing import Any, Optional, Dict, List, Union, Tuple
32
  from typing import AsyncGenerator
33
+
34
+ from vms.training_service import TrainingService
35
+ from vms.captioning_service import CaptioningService
36
+ from vms.splitting_service import SplittingService
37
+ from vms.import_service import ImportService
38
+ from vms.config import (
39
  STORAGE_PATH, VIDEOS_TO_SPLIT_PATH, STAGING_PATH,
40
  TRAINING_PATH, LOG_FILE_PATH, TRAINING_PRESETS, TRAINING_VIDEOS_PATH, MODEL_PATH, OUTPUT_PATH, DEFAULT_CAPTIONING_BOT_INSTRUCTIONS,
41
  DEFAULT_PROMPT_PREFIX, HF_API_TOKEN, ASK_USER_TO_DUPLICATE_SPACE, MODEL_TYPES, SMALL_TRAINING_BUCKETS
42
  )
43
+ from vms.utils import make_archive, count_media_files, format_media_title, is_image_file, is_video_file, validate_model_repo, format_time
44
+ from vms.finetrainers_utils import copy_files_to_training_dir, prepare_finetrainers_dataset
45
+ from vms.training_log_parser import TrainingLogParser
46
 
47
  logger = logging.getLogger(__name__)
48
  logger.setLevel(logging.INFO)
 
463
  full_caption = preview_caption
464
 
465
  path = Path(preview_video if preview_video else preview_image)
466
+ file_path = path.with_suffix('') if path.suffix == '.txt' else path
467
+ self.captioner.update_file_caption(file_path, full_caption)
 
 
468
  return gr.update(value="Caption saved successfully!")
469
  except Exception as e:
470
  return gr.update(value=f"Error saving caption: {str(e)}")
vms/__init__.py ADDED
File without changes
captioning_service.py β†’ vms/captioning_service.py RENAMED
@@ -100,6 +100,40 @@ class CaptioningService:
100
  cls._model_loading = None
101
  raise
102
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
103
  async def ensure_model_loaded(self):
104
  """Ensure model is loaded before processing"""
105
  if USE_MOCK_CAPTIONING_MODEL:
 
100
  cls._model_loading = None
101
  raise
102
 
103
+
104
+ def update_file_caption(self, file_path: Path, caption: str) -> None:
105
+ """Update caption for a training file
106
+
107
+ Args:
108
+ file_path: Path to the media file
109
+ caption: New caption text
110
+ """
111
+ try:
112
+ # Ensure we're working with Path objects
113
+ file_path = Path(file_path)
114
+
115
+ # Create the caption file path
116
+ caption_path = file_path.with_suffix('.txt')
117
+
118
+ # Write the new caption
119
+ caption_path.write_text(caption)
120
+
121
+ logger.info(f"Updated caption for {file_path.name}")
122
+
123
+ # the following code is disabled, because we want to make the copy to prompts.txt manual
124
+ # If the file is in TRAINING_VIDEOS_PATH, update prompts.txt as well
125
+ # if TRAINING_VIDEOS_PATH in file_path.parents:
126
+ # # Try to update the training dataset
127
+ # try:
128
+ # prepare_finetrainers_dataset()
129
+ # logger.info("Updated training dataset with new caption")
130
+ # except Exception as e:
131
+ # logger.warning(f"Could not update training dataset: {str(e)}")
132
+
133
+ except Exception as e:
134
+ logger.error(f"Error updating caption: {str(e)}")
135
+ raise
136
+
137
  async def ensure_model_loaded(self):
138
  """Ensure model is loaded before processing"""
139
  if USE_MOCK_CAPTIONING_MODEL:
config.py β†’ vms/config.py RENAMED
File without changes
finetrainers_utils.py β†’ vms/finetrainers_utils.py RENAMED
File without changes
image_preprocessing.py β†’ vms/image_preprocessing.py RENAMED
File without changes
import_service.py β†’ vms/import_service.py RENAMED
File without changes
splitting_service.py β†’ vms/splitting_service.py RENAMED
File without changes
training_log_parser.py β†’ vms/training_log_parser.py RENAMED
File without changes
training_service.py β†’ vms/training_service.py RENAMED
@@ -212,8 +212,7 @@ class TrainingService:
212
  except Exception as e:
213
  logger.error(f"Error during config validation: {str(e)}")
214
  return f"Configuration validation failed: {str(e)}"
215
-
216
-
217
  def start_training(
218
  self,
219
  model_type: str,
 
212
  except Exception as e:
213
  logger.error(f"Error during config validation: {str(e)}")
214
  return f"Configuration validation failed: {str(e)}"
215
+
 
216
  def start_training(
217
  self,
218
  model_type: str,
utils.py β†’ vms/utils.py RENAMED
File without changes
video_preprocessing.py β†’ vms/video_preprocessing.py RENAMED
File without changes