Add English support

app.py

@@ -5,15 +5,16 @@ import torch
 from transformers import pipeline
 from diffusers import StableDiffusionPipeline
 from TTS.api import TTS
+import whisper
 
 import utils
 from youtubeaudioextractor import PytubeAudioExtractor
-from transcriber import Transcriber
+from transcriber import SpanishTranscriber, WhisperTranscriber
 from textprocessor import TextProcessor
 from videocreator import VideoCreator
 
-TRANSCRIBER_MODEL_NAME = "juancopi81/whisper-medium-es"
-lang = "es"
+spanish_transcribe_model = "juancopi81/whisper-medium-es"
+languages = {"Spanish": "es", "English": "en"}
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 device_dict = {"cuda": 0, "cpu": -1}
@@ -29,14 +30,20 @@ device_print = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
 
 # Initialize components
 audio_extractor = PytubeAudioExtractor()
-transcription_pipe = pipeline(
+es_transcription_pipe = pipeline(
     task="automatic-speech-recognition",
-    model=TRANSCRIBER_MODEL_NAME,
+    model=spanish_transcribe_model,
     chunk_length_s=30,
     device=device_dict[device],
 )
-transcription_pipe.model.config.forced_decoder_ids = transcription_pipe.tokenizer.get_decoder_prompt_ids(language=lang, task="transcribe")
-audio_transcriber = Transcriber(transcription_pipe)
+es_transcription_pipe.model.config.forced_decoder_ids = es_transcription_pipe.tokenizer.get_decoder_prompt_ids(language="es",
+                                                                                                               task="transcribe")
+es_audio_transcriber = SpanishTranscriber(es_transcription_pipe)
+
+en_transcription_pipe = whisper.load_model("base")
+
+en_audio_transcriber = WhisperTranscriber(en_transcription_pipe)
+
 openai_model = "text-davinci-003"
 text_processor = TextProcessor(openai_model)
 
@@ -47,19 +54,33 @@ image_pipeline = StableDiffusionPipeline.from_pretrained(image_model_id,
 
 image_pipeline = image_pipeline.to(device)
 
-vo_model_name = TTS.list_models()[22]
+es_vo_model_name = TTS.list_models()[22]
+en_vo_model_name = TTS.list_models()[8]
 # Init TTS
-tts = TTS(vo_model_name)
-video_creator = VideoCreator(tts, image_pipeline)
+es_tts = TTS(es_vo_model_name)
+en_tts = TTS(en_vo_model_name)
 
-def datapipeline(url: str) -> Any:
+def datapipeline(url: str, 
+                 video_language: str,
+                 summary_language: str,
+                 video_styles: str) -> Any:
     audio_path_file = audio_extractor.extract(url)
     print(f"Audio file created at: {audio_path_file}")
+    # Select transcriber
+    if video_language == "Spanish":
+        audio_transcriber = es_audio_transcriber
+        video_creator = VideoCreator(es_tts, image_pipeline)
+    elif video_language == "English":
+        audio_transcriber = en_audio_transcriber
+        video_creator = VideoCreator(en_tts, image_pipeline)
+    else:
+        return "Language not supported"
     transcribed_text = audio_transcriber.transcribe(audio_path_file)
     print("Audio transcription ready!")
-    json_scenes = text_processor.get_json_scenes(transcribed_text)
+    json_scenes = text_processor.get_json_scenes(transcribed_text,
+                                                 summary_language)
     print("Scenes ready")
-    video = video_creator.create_video(json_scenes)
+    video = video_creator.create_video(json_scenes, video_styles)
     print("Video at", video)
     return video, video
 
@@ -184,6 +205,11 @@ css = """
         #generated_id{
             min-height: 700px
         }
+        #setting_id{
+          margin-bottom: 12px;
+          text-align: center;
+          font-weight: 900;
+        }
 """
 block = gr.Blocks(css=css)
 
@@ -214,6 +240,7 @@ with block as demo:
                 Some samples videos you can try:
                 <ul>
                     <li>https://www.youtube.com/watch?v=Hk5evm1NgzA (Little Red Riding Hood. Infer time: c.a. 196 seconds)</li>
+                    <li>https://www.youtube.com/watch?v=nJxWS9jZ9-c (Elon Musk's Biography. Infer time: c.a. 176 seconds)</li>
                     <li>https://www.youtube.com/watch?v=sRmmQBBln9Q (Cook recipe. Infer time: c.a. 200 seconds)</li>
                     <li>https://www.youtube.com/watch?v=qz4Wc48KITA (Poem by Edgar Allan Poe. Infer time: c.a. 200 seconds)</li>
                     <li>https://www.youtube.com/watch?v=2D8CaoIY7Lk (The history of Christmas trees. Infer time: c.a. 130 seconds)</li>
@@ -224,12 +251,32 @@ with block as demo:
             </div>
         """
     )
+    with gr.Group():
+        with gr.Box():
+            with gr.Row(elem_id="setting_id").style(mobile_collapse=False, equal_height=True):
+              gr.HTML("<h1>Setting</h1>")
+            with gr.Row():
+              with gr.Column():
+                video_language = gr.Radio(choices=["Spanish", "English"],
+                                          label="Language of your input video:",
+                                          value="Spanish")
+              with gr.Column():
+                summary_language = gr.Radio(choices=["Spanish", "English"],
+                                            label="Language of your output video:",
+                                            value="Spanish")
+            with gr.Row():
+              video_styles = gr.Textbox(label="(OPTIONAL) Enter the styles for your ouput video",
+                                        value="",
+                                        placeholder="illustration, highly detailed, digital painting, concept art, matte, art by wlop and artgerm and greg rutkowski and alphonse mucha, masterpiece")
     with gr.Group():
         with gr.Box():
             with gr.Row().style(mobile_collapse=False, equal_height=True):
 
                 url = gr.Textbox(
-                    label="Enter the URL of the YouTubeVideo", show_label=False, max_lines=1
+                    label="Enter the URL of the YouTubeVideo",
+                    show_label=False,
+                    max_lines=1,
+                    placeholder="YouTube URL"
                 ).style(
                     border=(True, False, True, True),
                     rounded=(True, False, False, True),
@@ -243,9 +290,15 @@ with block as demo:
         file_output = gr.File()
         
         btn.click(datapipeline,
-                  inputs=[url],
+                  inputs=[url,
+                          video_language,
+                          summary_language,
+                          video_styles],
                   outputs=[video_output, file_output])
 
+    #gr.Examples(
+    #    examples=[[], []]
+    #)
     gr.HTML(
             """
                 <div class="footer">

requirements.txt

@@ -8,4 +8,5 @@ accelerate
 TTS
 moviepy
 imageio==2.4.1
-tensorboard
+tensorboard
+git+https://github.com/openai/whisper.git

textprocessor.py

@@ -9,9 +9,9 @@ You are a creator of illustrated books building a series of scenes for your book
 Your boss asked you to write a summary and illustrations of this text:
 $TRANSCRIPTION
 You have to write the summary using a maximum of 7 scenes in a JSON object following these instructions:
-[{Scene": int, "Summary": Spanish str, "Illustration": English str}, ...] where:
+[{Scene": int, "Summary": $SUMMARY_LANGUAGE str, "Illustration": English str}, ...] where:
 "Scene": The number of the scene.
-"Summary": Spanish string with a summary of the scene. It should be in Spanish, and it should be less than 30 words. Readers should understand it without looking at the illustration.
+"Summary": $SUMMARY_LANGUAGE string with a summary of the scene. It should be in $SUMMARY_LANGUAGE, and it should be less than 30 words. Readers should understand it without looking at the illustration.
 "Illustration": English string with a detailed English description of an illustration for this scene. It must be written in English and in less than 20 words. It should include many details and an artistic style for the image that matches the text.
 Just answer with the JSON object:
 """
@@ -34,8 +34,10 @@ class TextProcessor:
         self.presence_penalty = presence_penalty
     
     def get_json_scenes(self,
-                        prompt: str) -> Dict:
+                        prompt: str,
+                        summary_language: str) -> Dict:
         gpt_prompt = context_prompt.replace("$TRANSCRIPTION", prompt)
+        gpt_prompt = gpt_prompt.replace("$SUMMARY_LANGUAGE ", summary_language)
         print("gpt_prompt", gpt_prompt)
         response = openai.Completion.create(
             model=self.model,

transcriber.py

@@ -1,6 +1,14 @@
+from abc import ABC, abstractmethod
+
 from transformers import pipeline
 
-class Transcriber:
+class Transcriber(ABC):
+
+    @abstractmethod
+    def transcribe(self, file_path: str) -> str:
+        pass
+        
+class SpanishTranscriber(Transcriber):
     def __init__(self, pipe: pipeline) -> None:
         self.pipe = pipe
 
@@ -8,4 +16,16 @@ class Transcriber:
         print("Pipe:", self.pipe)
         print("Audo file at:", file_path)
         transcription = self.pipe(file_path)["text"]
+        return transcription
+    
+class WhisperTranscriber(Transcriber):
+    def __init__(self, model, without_timestamps: bool=True) -> None:
+        self.model = model
+        self.without_timestamps = without_timestamps
+
+    def transcribe(self, file_path: str = "yt_audio.mp3") -> str:
+        print("Model:", self.model)
+        print("Audo file at:", file_path)
+        transcription = self.model.transcribe(file_path,
+                                              without_timestamps=self.without_timestamps)["text"]
         return transcription

videocreator.py

@@ -11,17 +11,20 @@ class VideoCreator:
         self.tts_pipeline = tts_pipeline
         self.image_pipeline = image_pipeline
         
-    def create_video(self, scenes: Dict) -> str:
+    def create_video(self,
+                     scenes: Dict,
+                     video_styles: str) -> str:
         videos_dict = {}
         for index, scene in enumerate(scenes):
-            video_scene = self._create_video_from_scene(scenes[scene])
+            video_scene = self._create_video_from_scene(scenes[scene],
+                                                        video_styles)
             videos_dict[index] = video_scene
         merged_video = self._merge_videos(videos_dict)
         return merged_video
     
-    def _create_video_from_scene(self, scene: Dict) -> str:
+    def _create_video_from_scene(self, scene: Dict, video_styles: str) -> str:
         audio_file = self._get_audio_from_text(scene["Summary"])
-        bg_image = self._get_bg_image_from_description(scene["Illustration"])
+        bg_image = self._get_bg_image_from_description(scene["Illustration"], video_styles)
         video = gr.make_waveform(audio=audio_file,
                                  bg_image=bg_image)
         return video
@@ -31,8 +34,8 @@ class VideoCreator:
                                       file_path="output.wav")
         return "output.wav"
     
-    def _get_bg_image_from_description(self, img_desc: str):
-        images = self.image_pipeline(img_desc)
+    def _get_bg_image_from_description(self, img_desc: str, video_styles: str):
+        images = self.image_pipeline(img_desc + video_styles)
         print("Image generated!")
         image_output = images.images[0]
         image_output.save("img.png")