rairo commited on
Commit
1952ce6
·
verified ·
1 Parent(s): 40eafc7

Update sozo_gen.py

Browse files
Files changed (1) hide show
  1. sozo_gen.py +34 -5
sozo_gen.py CHANGED
@@ -567,6 +567,7 @@ def generate_single_chart(df: pd.DataFrame, description: str, uid: str, project_
567
  os.unlink(img_path)
568
  return None
569
 
 
570
  def generate_video_from_project(df: pd.DataFrame, raw_md: str, data_context: Dict, uid: str, project_id: str, voice_model: str, bucket):
571
  logging.info(f"Generating video for project {project_id} with voice {voice_model}")
572
  llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", google_api_key=API_KEY, temperature=0.2)
@@ -655,12 +656,14 @@ def generate_video_from_project(df: pd.DataFrame, raw_md: str, data_context: Dic
655
 
656
  with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_vid, \
657
  tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as temp_aud, \
658
- tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as final_vid:
 
659
 
660
  silent_vid_path = Path(temp_vid.name)
661
  audio_mix_path = Path(temp_aud.name)
662
  final_vid_path = Path(final_vid.name)
663
-
 
664
  concat_media(final_video_parts, silent_vid_path)
665
  concat_media(audio_parts, audio_mix_path)
666
 
@@ -673,18 +676,44 @@ def generate_video_from_project(df: pd.DataFrame, raw_md: str, data_context: Dic
673
  ]
674
  subprocess.run(cmd, check=True, capture_output=True)
675
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
676
  blob_name = f"sozo_projects/{uid}/{project_id}/video.mp4"
677
  blob = bucket.blob(blob_name)
678
- blob.upload_from_filename(str(final_vid_path))
679
  logging.info(f"Uploaded video to {blob.public_url}")
680
 
681
- for p in temps + [silent_vid_path, audio_mix_path, final_vid_path]:
682
  if os.path.exists(p): os.unlink(p)
683
 
684
  return blob.public_url
685
  return None
686
 
687
-
688
  # In sozo_gen.py, add these new functions at the end of the file
689
 
690
  def generate_image_with_gemini(prompt: str) -> Image.Image:
 
567
  os.unlink(img_path)
568
  return None
569
 
570
+
571
  def generate_video_from_project(df: pd.DataFrame, raw_md: str, data_context: Dict, uid: str, project_id: str, voice_model: str, bucket):
572
  logging.info(f"Generating video for project {project_id} with voice {voice_model}")
573
  llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", google_api_key=API_KEY, temperature=0.2)
 
656
 
657
  with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_vid, \
658
  tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as temp_aud, \
659
+ tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as final_vid, \
660
+ tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as branded_vid:
661
 
662
  silent_vid_path = Path(temp_vid.name)
663
  audio_mix_path = Path(temp_aud.name)
664
  final_vid_path = Path(final_vid.name)
665
+ branded_vid_path = Path(branded_vid.name)
666
+
667
  concat_media(final_video_parts, silent_vid_path)
668
  concat_media(audio_parts, audio_mix_path)
669
 
 
676
  ]
677
  subprocess.run(cmd, check=True, capture_output=True)
678
 
679
+ upload_path = final_vid_path
680
+ logo_path = Path("sozob.png")
681
+
682
+ if logo_path.exists():
683
+ logging.info("Logo 'sozob.png' found. Adding full-screen end-card.")
684
+ duration_for_filter = total_audio_duration
685
+
686
+ filter_complex = f"[1:v]scale={WIDTH}:{HEIGHT}[logo];[0:v][logo]overlay=0:0:enable='gte(t,{duration_for_filter - 2})'"
687
+
688
+ logo_cmd = [
689
+ "ffmpeg", "-y",
690
+ "-i", str(final_vid_path),
691
+ "-i", str(logo_path),
692
+ "-filter_complex", filter_complex,
693
+ "-map", "0:a",
694
+ "-c:a", "copy",
695
+ "-c:v", "libx264", "-pix_fmt", "yuv420p",
696
+ str(branded_vid_path)
697
+ ]
698
+ try:
699
+ subprocess.run(logo_cmd, check=True, capture_output=True)
700
+ upload_path = branded_vid_path
701
+ except subprocess.CalledProcessError as e:
702
+ logging.error(f"Failed to add logo end-card. Uploading unbranded video. Error: {e.stderr.decode()}")
703
+ else:
704
+ logging.warning("Logo 'sozob.png' not found in root directory. Skipping end-card.")
705
+
706
  blob_name = f"sozo_projects/{uid}/{project_id}/video.mp4"
707
  blob = bucket.blob(blob_name)
708
+ blob.upload_from_filename(str(upload_path))
709
  logging.info(f"Uploaded video to {blob.public_url}")
710
 
711
+ for p in temps + [silent_vid_path, audio_mix_path, final_vid_path, branded_vid_path]:
712
  if os.path.exists(p): os.unlink(p)
713
 
714
  return blob.public_url
715
  return None
716
 
 
717
  # In sozo_gen.py, add these new functions at the end of the file
718
 
719
  def generate_image_with_gemini(prompt: str) -> Image.Image: