jschwab21 commited on
Commit
9846923
·
verified ·
1 Parent(s): 72a3e3b

Update video_processing.py

Browse files
Files changed (1) hide show
  1. video_processing.py +15 -2
video_processing.py CHANGED
@@ -7,6 +7,7 @@ from transformers import CLIPProcessor, CLIPModel
7
  import torch
8
  import yt_dlp
9
  from PIL import Image
 
10
 
11
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12
  model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
@@ -15,7 +16,7 @@ processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
15
  def download_video(url):
16
  ydl_opts = {
17
  'format': 'bestvideo[height<=1440]+bestaudio/best[height<=1440]',
18
- 'outtmpl': 'downloaded_video.%(ext)s',
19
  'merge_output_format': 'mp4',
20
  }
21
  with yt_dlp.YoutubeDL(ydl_opts) as ydl:
@@ -111,8 +112,20 @@ def process_video(video_url, description):
111
  if final_clip:
112
  output_dir = "output"
113
  os.makedirs(output_dir, exist_ok=True)
114
- final_clip_path = os.path.join(output_dir, "final_clip.mp4")
115
  final_clip.write_videofile(final_clip_path, codec='libx264', audio_codec='aac')
 
116
  return final_clip_path
117
  return None
118
 
 
 
 
 
 
 
 
 
 
 
 
 
7
  import torch
8
  import yt_dlp
9
  from PIL import Image
10
+ import uuid
11
 
12
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
13
  model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
 
16
  def download_video(url):
17
  ydl_opts = {
18
  'format': 'bestvideo[height<=1440]+bestaudio/best[height<=1440]',
19
+ 'outtmpl': f'temp_videos/{uuid.uuid4()}_video.%(ext)s',
20
  'merge_output_format': 'mp4',
21
  }
22
  with yt_dlp.YoutubeDL(ydl_opts) as ydl:
 
112
  if final_clip:
113
  output_dir = "output"
114
  os.makedirs(output_dir, exist_ok=True)
115
+ final_clip_path = os.path.join(output_dir, f"{uuid.uuid4()}_final_clip.mp4")
116
  final_clip.write_videofile(final_clip_path, codec='libx264', audio_codec='aac')
117
+ cleanup_temp_files()
118
  return final_clip_path
119
  return None
120
 
121
+ def cleanup_temp_files():
122
+ temp_dir = 'temp_videos'
123
+ if os.path.exists(temp_dir):
124
+ for file in os.listdir(temp_dir):
125
+ file_path = os.path.join(temp_dir, file)
126
+ try:
127
+ if os.path.isfile(file_path):
128
+ os.unlink(file_path)
129
+ except Exception as e:
130
+ print(f"Error: {e}")
131
+