Video-Generator-Tools / src /caption_gen.py
userIdc2024's picture
Upload 5 files
d856b59 verified
import os
import replicate
from dotenv import load_dotenv
load_dotenv()
replicate_client = replicate.Client(api_token= os.getenv("REPLICATE_API_KEY"))
def caption_generation(
video: bytes,
caption_size: int,
highlight_color: str):
output = replicate_client.run(
"shreejalmaharjan-27/tiktok-short-captions:46bf1c12c77ad1782d6f87828d4d8ba4d48646b8e1271b490cb9e95ccdbc4504",
input={
"model": "large-v3",
"video": f"data:video/mp4;base64,{video}",
"language": "auto",
"temperature": 0,
"caption_size": caption_size,
"highlight_color": highlight_color,
"suppress_tokens": "-1",
"logprob_threshold": -1,
"no_speech_threshold": 0.6,
"condition_on_previous_text": True,
"compression_ratio_threshold": 2.4,
"temperature_increment_on_fallback": 0.2
}
)
urls = []
if isinstance(output, list) and output:
first = output[0]
url = getattr(first, "url", str(first))
urls = [url]
elif isinstance(output, str):
urls = [output]
elif hasattr(output, "url"):
urls = [getattr(output, "url")]
if urls:
return urls[0]