Added movie title and optional subtitles

app.py

@@ -1,5 +1,3 @@
-
-
 import gradio as gr
 import torch
 import torchaudio
@@ -11,7 +9,6 @@ import os
 from PIL import Image
 
 def estimate_chunk_durations(text, words_per_second=2.5, min_sec=5, max_sec=7):
-    print("Estimating chunk durations...")
     words = text.split()
     chunks = []
     current_chunk = []
@@ -26,89 +23,80 @@ def estimate_chunk_durations(text, words_per_second=2.5, min_sec=5, max_sec=7):
                 current_duration = 0
     if current_chunk:
         chunks.append(" ".join(current_chunk))
-    
-    total_time = sum([min(max(len(chunk.split()) / words_per_second, min_sec), max_sec) for chunk in chunks])
-    print(f"Total estimated time for video: {total_time:.2f} seconds")
     return chunks
 
 def generate_speech(text):
-    print("Generating speech...")
     tts = TTS("tts_models/en/ljspeech/tacotron2-DDC")
-    print("TTS model loaded")
-    tts.tts_to_file(text="Hello world!", file_path="test.wav")
     wav_path = "speech.wav"
     tts.tts_to_file(text=text, file_path=wav_path)
-    print("Speech generated")
     return wav_path
 
 def generate_images(chunks, image_size=(640, 480), use_diffusion=True, num_steps=40):
-    print("Generating images...")
     image_paths = []
-    
     if use_diffusion:
         pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
-        def custom_checker(images, **kwargs):
-            return images, [False] * len(images)  # Force all images to be "safe"
-
-        pipe.safety_checker = custom_checker
         pipe.to("cuda" if torch.cuda.is_available() else "cpu")
     
     for i, chunk in enumerate(chunks):
-        print(f"Generating image for chunk {i + 1} of {len(chunks)}: {chunk[:50]}...")
         if use_diffusion:
             image = pipe(chunk, num_inference_steps=num_steps).images[0]
             image = image.resize(image_size)
         else:
             image = Image.new("RGB", image_size, (0, 0, 0))
-        
         img_path = f"image_{i}.png"
         image.save(img_path)
         image_paths.append(img_path)
-    
     return image_paths
 
-def create_video(images, durations, speech_path, image_size=(640, 480)):
+def create_video(images, durations, speech_path, movie_title, add_subtitles, chunks, image_size=(640, 480)):
     clips = []
-    for img, dur in zip(images, durations):
-        pil_image = Image.open(img)
-        pil_image = pil_image.resize(image_size, Image.Resampling.LANCZOS)
-        frame = np.array(pil_image)  # Convert to NumPy array
-        print(f"Duration: {dur}")
+    title_clip = mp.TextClip(movie_title, fontsize=50, color='white', size=image_size)
+    title_clip = title_clip.set_duration(1).set_position('center')
+    black_start = mp.ColorClip(image_size, color=(0,0,0), duration=1).set_opacity(0.8)
+    clips.append(mp.CompositeVideoClip([black_start, title_clip]))
+    
+    for img, dur, chunk in zip(images, durations, chunks):
+        frame = np.array(Image.open(img).resize(image_size, Image.Resampling.LANCZOS))
         clip = mp.ImageClip(frame).set_duration(dur)
+        if add_subtitles:
+            txt_clip = mp.TextClip(chunk, fontsize=30, color='white', size=(image_size[0] - 20, None), method='caption')
+            txt_clip = txt_clip.set_duration(dur).set_position(('center', 'bottom'))
+            clip = mp.CompositeVideoClip([clip, txt_clip])
         clips.append(clip)
-        
-    black_start = mp.ColorClip(image_size, color=(0,0,0), duration=1)
+    
     black_end = mp.ColorClip(image_size, color=(0,0,0), duration=2)
-    video = mp.concatenate_videoclips([black_start] + clips + [black_end])
+    video = mp.concatenate_videoclips(clips + [black_end])
     audio = mp.AudioFileClip(speech_path)
     final_video = video.set_audio(audio)
     final_video.write_videofile("output.mp4", fps=24)
     return "output.mp4"
 
-def process_text(text, image_size, use_diffusion, num_steps):
+def process_text(text, movie_title, image_size, use_diffusion, num_steps, add_subtitles):
     chunks = estimate_chunk_durations(text)
     speech_path = generate_speech(text)
     image_paths = generate_images(chunks, image_size, use_diffusion, num_steps)
     durations = [min(10, max(5, len(chunk.split()) / 2.5)) for chunk in chunks]
-    video_path = create_video(image_paths, durations, speech_path, image_size)
+    video_path = create_video(image_paths, durations, speech_path, movie_title, add_subtitles, chunks, image_size)
     return video_path
 
 with gr.Blocks() as demo:
     gr.Markdown("# Text-to-Video Generator using AI 🎥")
     text_input = gr.Textbox(label="Enter your text")
+    movie_title_input = gr.Textbox(label="Movie Title", value="")
     file_input = gr.File(label="Or upload a .txt file")
     image_size_input = gr.Radio(choices=["640x480", "800x600", "1024x768"], label="Select Image Size", value="640x480")
     use_diffusion_input = gr.Checkbox(label="Use Diffusion Images", value=True)
     num_steps_input = gr.Slider(minimum=1, maximum=50, step=1, value=5, label="Diffusion Model Steps")
+    add_subtitles_input = gr.Checkbox(label="Add Subtitles", value=False)
     process_btn = gr.Button("Generate Video")
     output_video = gr.Video()
     
-    def handle_request(text, file, image_size, use_diffusion, num_steps):
+    def handle_request(text, movie_title, file, image_size, use_diffusion, num_steps, add_subtitles):
         if file is not None:
             text = open(file.name, "r").read()
         image_size_dict = {"640x480": (640, 480), "800x600": (800, 600), "1024x768": (1024, 768)}
-        return process_text(text, image_size_dict[image_size], use_diffusion, num_steps)
+        return process_text(text, movie_title, image_size_dict[image_size], use_diffusion, num_steps, add_subtitles)
     
-    process_btn.click(handle_request, inputs=[text_input, file_input, image_size_input, use_diffusion_input, num_steps_input], outputs=output_video)
+    process_btn.click(handle_request, inputs=[text_input, movie_title_input, file_input, image_size_input, use_diffusion_input, num_steps_input, add_subtitles_input], outputs=output_video)
 
 demo.launch()