Update app.py

app.py

@@ -1,27 +1,26 @@
-import moviepy.editor as mpy
-from PIL import Image
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import gradio as gr
+import os
+import multiprocessing
+import subprocess
+import nltk
 import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from min_dalle import MinDalle
-from huggingface_hub import snapshot_download
+from moviepy.editor import VideoFileClip
 from PIL import Image, ImageDraw, ImageFont
-import textwrap
 from mutagen.mp3 import MP3
 from gtts import gTTS
 from pydub import AudioSegment
-import os
-import glob
-import nltk
-import subprocess
-import shutil
+import textwrap
+import gradio as gr
 import matplotlib.pyplot as plt
-import gc  # Import the garbage collector
+import gc  # Garbage collector
+from huggingface_hub import snapshot_download
 from audio import *
-import os
-import multiprocessing
-multiprocessing.set_start_method("spawn")
-# Define a fallback for environments without GPU
+
+# Ensure proper multiprocessing start method
+multiprocessing.set_start_method("spawn", force=True)
+
+# GPU Fallback Setup
 if os.environ.get("SPACES_ZERO_GPU") is not None:
     import spaces
 else:
@@ -31,45 +30,63 @@ else:
             def wrapper(*args, **kwargs):
                 return func(*args, **kwargs)
             return wrapper
-# Download necessary NLTK data
-try:
-    nltk.data.find('tokenizers/punkt')
-except LookupError:
-    nltk.download('punkt')
-
-description = "Video Story Generator with Audio \n PS:  Generation of video by using Artifical Intellingence by dalle-mini and distilbart and gtss "
-title = "Video Story Generator with Audio by using dalle-mini and distilbart and gtss   "
 
-# Load tokenizer and model for text summarization
-tokenizer = AutoTokenizer.from_pretrained("sshleifer/distilbart-cnn-12-6")
-model = AutoModelForSeq2SeqLM.from_pretrained("sshleifer/distilbart-cnn-12-6")
-
-# Check for CUDA availability and set device
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
-model.to(device)
-
-# Function to log GPU memory (optional, for debugging)
+# Download necessary NLTK data
+def setup_nltk():
+    """Ensure required NLTK data is available."""
+    try:
+        nltk.data.find('tokenizers/punkt')
+    except LookupError:
+        nltk.download('punkt')
+
+setup_nltk()
+
+# Constants
+DESCRIPTION = (
+    "Video Story Generator with Audio\n"
+    "PS: Generation of video by using Artificial Intelligence via dalle-mini, distilbart, and GTTS."
+)
+TITLE = "Video Story Generator with Audio by using dalle-mini, distilbart, and GTTS."
+
+# Load Tokenizer and Model for Text Summarization
+def load_text_summarization_model():
+    """Load the tokenizer and model for text summarization."""
+    print("Loading text summarization model...")
+    tokenizer = AutoTokenizer.from_pretrained("sshleifer/distilbart-cnn-12-6")
+    model = AutoModelForSeq2SeqLM.from_pretrained("sshleifer/distilbart-cnn-12-6")
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    model.to(device)
+    return tokenizer, model, device
+
+tokenizer, model, device = load_text_summarization_model()
+
+# Log GPU Memory (optional, for debugging)
 def log_gpu_memory():
+    """Log GPU memory usage."""
     if torch.cuda.is_available():
         print(subprocess.check_output('nvidia-smi').decode('utf-8'))
     else:
         print("CUDA is not available. Cannot log GPU memory.")
 
-# --------- MinDalle Image Generation Functions ---------
-# Check for GPU availability
+# Check GPU Availability
 def check_gpu_availability():
+    """Print GPU availability and device details."""
     if torch.cuda.is_available():
         print(f"CUDA devices: {torch.cuda.device_count()}")
         print(f"Current device: {torch.cuda.current_device()}")
         print(torch.cuda.get_device_properties(torch.cuda.current_device()))
     else:
         print("CUDA is not available. Running on CPU.")
+
 check_gpu_availability()
-# GPU-safe model loading
+
+# GPU-Safe MinDalle Model Loading
 def initialize_min_dalle_with_gpu():
+    """Load the MinDalle model with GPU support."""
     @spaces.GPU(duration=60 * 3)
     def load_model():
+        print("Loading MinDalle model...")
         return MinDalle(
             is_mega=True,
             models_root='pretrained',
@@ -80,9 +97,10 @@ def initialize_min_dalle_with_gpu():
         )
     return load_model()
 
-# Initialize MinDalle model
+# Initialize MinDalle Model
 min_dalle_model = initialize_min_dalle_with_gpu()
 
+
 def generate_image_with_min_dalle(
     model: MinDalle,
     text: str,