mirror of
https://github.com/Paillat-dev/viralfactory.git
synced 2026-01-02 01:06:19 +00:00
Paillat
e3229518d4
✨ feat(GenerationContext.py): add support for captioning engine in the GenerationContext class
The import statement for the `moviepy.editor` module is changed to `moviepy.editor as mp` to improve code readability. Additionally, the `gradio` module is imported as `gr` to improve code readability. The `GenerationContext` class now includes a `captioningengine` parameter and initializes a `captioningengine` attribute. The `setup_dir` method is modified to include a call to create a directory for the output files. The `get_file_path` method is modified to return the file path based on the output directory. The `process` method is modified to include additional steps for captioning. The `timed_script` attribute is added to store the result of the `ttsengine.synthesize` method. The `captioningengine` is used to generate captions and store them in the `captions` attribute. The final video is rendered using the `moviepy` library and saved as "final.mp4" in the output directory.
73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
import moviepy.editor as mp
|
|
import time
|
|
import os
|
|
import gradio as gr
|
|
|
|
from .. import engines
|
|
from ..utils.prompting import get_prompt
|
|
|
|
|
|
class GenerationContext:
|
|
def __init__(
|
|
self,
|
|
powerfulllmengine,
|
|
simplellmengine,
|
|
scriptengine,
|
|
ttsengine,
|
|
captioningengine,
|
|
) -> None:
|
|
self.powerfulllmengine: engines.LLMEngine.BaseLLMEngine = powerfulllmengine[0]
|
|
self.powerfulllmengine.ctx = self
|
|
|
|
self.simplellmengine: engines.LLMEngine.BaseLLMEngine = simplellmengine[0]
|
|
self.simplellmengine.ctx = self
|
|
|
|
self.scriptengine: engines.ScriptEngine.BaseScriptEngine = scriptengine[0]
|
|
self.scriptengine.ctx = self
|
|
|
|
self.ttsengine: engines.TTSEngine.BaseTTSEngine = ttsengine[0]
|
|
self.ttsengine.ctx = self
|
|
|
|
self.captioningengine: engines.CaptioningEngine.BaseCaptioningEngine = (
|
|
captioningengine[0]
|
|
)
|
|
self.captioningengine.ctx = self
|
|
|
|
def setup_dir(self):
|
|
self.dir = f"output/{time.time()}"
|
|
os.makedirs(self.dir)
|
|
|
|
def get_file_path(self, name: str) -> str:
|
|
return os.path.join(self.dir, name)
|
|
|
|
def process(self):
|
|
# ⚠️ IMPORTANT NOTE: All methods called here are expected to be defined as abstract methods in the base classes, if not there is an issue with the engine implementation.
|
|
|
|
progress = gr.Progress()
|
|
self.width, self.height = (
|
|
1080,
|
|
1920,
|
|
) # TODO: Add support for custom resolution, for now it's tiktok's resolution
|
|
self.setup_dir()
|
|
|
|
self.script = self.scriptengine.generate()
|
|
|
|
self.timed_script = self.ttsengine.synthesize(
|
|
self.script, self.get_file_path("tts.wav")
|
|
)
|
|
|
|
if not isinstance(self.captioningengine, engines.NoneEngine):
|
|
self.captions = self.captioningengine.get_captions()
|
|
else:
|
|
self.captions = []
|
|
|
|
# add any other processing steps here
|
|
|
|
# we render to a file called final.mp4
|
|
# using moviepy CompositeVideoClip
|
|
|
|
clip = mp.CompositeVideoClip(self.captions, size=(self.width, self.height))
|
|
audio = mp.AudioFileClip(self.get_file_path("tts.wav"))
|
|
clip = clip.set_audio(audio)
|
|
clip.write_videofile(self.get_file_path("final.mp4"), fps=60)
|