app.py

import os
import shutil
import multiprocessing
import subprocess
import nltk
import gradio as gr
import matplotlib.pyplot as plt
import gc
from huggingface_hub import snapshot_download, hf_hub_download
from typing import List
import shutil
import numpy as np
import random
import spaces
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, CLIPFeatureExtractor
from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
from diffusers.utils import export_to_video
from moviepy.editor import VideoFileClip, CompositeVideoClip, TextClip
import moviepy.editor as mpy
from PIL import Image, ImageDraw, ImageFont
from mutagen.mp3 import MP3
from gtts import gTTS
from pydub import AudioSegment
import uuid
from safetensors.torch import load_file
import textwrap

# -------------------------------------------------------------------
# No more ImageMagick dependency!
# -------------------------------------------------------------------
print("ImageMagick dependency removed. Using Pillow for text rendering.")

# Ensure NLTK’s 'punkt_tab' (and other data) is present
nltk.download('punkt_tab', quiet=True)
nltk.download('punkt', quiet=True)

# -------------------------------------------------------------------
# GPU / Environment Setup
# -------------------------------------------------------------------
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.")

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()

# Ensure proper multiprocessing start method
multiprocessing.set_start_method("spawn", force=True)

# -------------------------------------------------------------------
# Constants & Model Setup
# -------------------------------------------------------------------
dtype = torch.float16
device = "cuda" if torch.cuda.is_available() else "cpu"

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE_720 = 720 # Changed maximum image size to 720, now max resolution is 720p
MAX_IMAGE_SIZE = MAX_IMAGE_SIZE_720

RESOLUTIONS = {
  "16:9": [
    {"resolution": "360p", "width": 640, "height": 360},
    {"resolution": "480p", "width": 854, "height": 480},
    {"resolution": "720p", "width": 1280, "height": 720},
    #{"resolution": "1080p", "width": 1920, "height": 1080} # Commented out resolutions higher than 720p
  ],
  "4:3": [
    {"resolution": "360p", "width": 480, "height": 360},
    {"resolution": "480p", "width": 640, "height": 480},
    {"resolution": "720p", "width": 960, "height": 720},
    #{"resolution": "1080p", "width": 1440, "height": 1080} # Commented out resolutions higher than 720p
  ],
  "1:1": [
    {"resolution": "360p", "width": 360, "height": 360},
    {"resolution": "480p", "width": 480, "height": 480},
    {"resolution": "720p", "width": 720, "height": 720},
    #{"resolution": "1080p", "width": 1080, "height": 1080}, # Commented out resolutions higher than 720p
    #{"resolution": "1920p", "width": 1920, "height": 1920}  # Commented out resolutions higher than 720p
  ],
  "9:16": [
    {"resolution": "360p", "width": 360, "height": 640},
    {"resolution": "480p", "width": 480, "height": 854},
    {"resolution": "720p", "width": 720, "height": 1280},
    #{"resolution": "1080p", "width": 1080, "height": 1920} # Commented out resolutions higher than 720p
  ]}


DESCRIPTION = (
    "Video Story Generator with Audio\n"
    "PS: Generation of video by using Artificial Intelligence via AnimateDiff, DistilBART, and GTTS."
)
TITLE = "Video Story Generator with Audio (AnimateDiff, DistilBART, and GTTS)"

@spaces.GPU()
def load_text_summarization_model():
    """Load the tokenizer and model for text summarization on GPU/CPU."""
    print("Loading text summarization model...")
    tokenizer = AutoTokenizer.from_pretrained("sshleifer/distilbart-cnn-12-6")
    model = AutoModelForSeq2SeqLM.from_pretrained("sshleifer/distilbart-cnn-12-6")
    return tokenizer, model

tokenizer, model = load_text_summarization_model()

# Base models for AnimateDiffLightning
bases = {
    "Cartoon": "frankjoshua/toonyou_beta6",
    "Realistic": "emilianJR/epiCRealism",
    "3d": "Lykon/DreamShaper",
    "Anime": "Yntec/mistoonAnime2"
}

# Keep track of what's loaded to avoid reloading each time
step_loaded = None
base_loaded = "Realistic"
motion_loaded = None

# Initialize AnimateDiff pipeline
if not torch.cuda.is_available():
    raise NotImplementedError("No GPU detected!")

pipe = AnimateDiffPipeline.from_pretrained(
    bases[base_loaded],
    torch_dtype=dtype
).to(device)

pipe.scheduler = EulerDiscreteScheduler.from_config(
    pipe.scheduler.config,
    timestep_spacing="trailing",
    beta_schedule="linear"
)

feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")


# -------------------------------------------------------------------
# Function: Generate Short Animation
# -------------------------------------------------------------------
def generate_short_animation(
    prompt_text: str,
    base: str = "Realistic",
    motion: str = "",
    step: int = 4,
    seed: int = 42,
    width: int = 512,
    height: int = 512,
) -> str:
    """
    Generates a short animated video (MP4) from a given prompt using AnimateDiffLightning.
    Returns the local path to the resulting MP4.
    """
    global step_loaded
    global base_loaded
    global motion_loaded

    # 1) Possibly reload correct step weights
    if step_loaded != step:
        repo = "ByteDance/AnimateDiff-Lightning"
        ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
        pipe.unet.load_state_dict(
            load_file(hf_hub_download(repo, ckpt), device=device),
            strict=False
        )
        step_loaded = step

    # 2) Possibly reload the correct base model
    if base_loaded != base:
        pipe.unet.load_state_dict(
            torch.load(
                hf_hub_download(bases[base], "unet/diffusion_pytorch_model.bin"),
                map_location=device
            ),
            strict=False
        )
        base_loaded = base

    # 3) Possibly unload/load motion LORA
    if motion_loaded != motion:
        pipe.unload_lora_weights()
        if motion:
            pipe.load_lora_weights(motion, adapter_name="motion")
            pipe.set_adapters(["motion"], [0.7])  # weighting can be adjusted
        motion_loaded = motion

    # 4) Generate frames
    print(f"[INFO] Generating short animation for prompt: '{prompt_text}' ...")
    generator = torch.Generator(device=device).manual_seed(seed) if seed is not None else None
    output = pipe(
        prompt=prompt_text,
        guidance_scale=1.2,
        num_inference_steps=step,
        generator=generator,
        width=width,
        height=height
    )

    # 5) Export frames to a short MP4
    short_mp4_path = f"short_{uuid.uuid4().hex}.mp4"
    export_to_video(output.frames[0], short_mp4_path, fps=10)
    return short_mp4_path

# -------------------------------------------------------------------
# Function: Merge MP3 files
# -------------------------------------------------------------------
def merge_audio_files(mp3_names: List[str]) -> str:
    """
    Merges a list of MP3 files into a single MP3 file.
    Returns the path to the merged MP3 file.
    """
    combined = AudioSegment.empty()
    for f_name in mp3_names:
        audio = AudioSegment.from_mp3(f_name)
        combined += audio
    export_path = f"merged_audio_{uuid.uuid4().hex}.mp3" # Dynamic output path for merged audio
    combined.export(export_path, format="mp3")
    print(f"DEBUG: Audio files merged and saved to {export_path}")
    return export_path


# -------------------------------------------------------------------
# Function: Overlay Subtitles on a Video
# -------------------------------------------------------------------

def add_subtitles_to_video(input_video_path: str, text: str, duration: float) -> str:
    """
    Overlays `text` as subtitles over the entire `input_video_path` for `duration` seconds using Pillow.
    Returns the path to the newly generated MP4 with subtitles.
    """
    base_clip = VideoFileClip(input_video_path)
    final_dur = max(duration, base_clip.duration)

    def make_frame(t):
        frame_pil = Image.fromarray(base_clip.get_frame(t))
        draw = ImageDraw.Draw(frame_pil)
        try:
            font = ImageFont.truetype("arial.ttf", 40)  # Change the font size if needed
        except IOError:
            font = ImageFont.load_default()  # Use default font if Arial is not found

        # Correctly compute text size using `textbbox()`
        bbox = draw.textbbox((0, 0), text, font=font)
        textwidth, textheight = bbox[2] - bbox[0], bbox[3] - bbox[1]

        x = (frame_pil.width - textwidth) / 2
        y = frame_pil.height - 70 - textheight  # Position at the bottom

        draw.text((x, y), text, font=font, fill=(255, 255, 0))  # Yellow color
        return np.array(frame_pil)

    # Create the video clip without `size` argument
    subtitled_clip = mpy.VideoClip(make_frame, duration=final_dur)

    # Composite the subtitled clip over the original video
    final_clip = CompositeVideoClip([base_clip, subtitled_clip.set_position((0, 0))])
    final_clip = final_clip.set_duration(final_dur)

    out_path = f"sub_{uuid.uuid4().hex}.mp4"
    final_clip.write_videofile(out_path, fps=24, logger=None)

    # Cleanup
    base_clip.close()
    final_clip.close()
    subtitled_clip.close()

    return out_path



# -------------------------------------------------------------------
# Main Function: Generate Output Video
# -------------------------------------------------------------------
@spaces.GPU()
def get_output_video(text, base_model_name, motion_name, num_inference_steps_backend, randomize_seed, seed, width, height):
    """
    Summarize the user prompt, generate a short animated video for each sentence,
    overlay subtitles, merge all into a final video with a single audio track.
    """
    print("DEBUG: Starting get_output_video function...")

    # Summarize the input text
    print("DEBUG: Summarizing text...")
    device_local = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device_local)  # Move summarization model to GPU/CPU as needed

    inputs = tokenizer(
        text,
        max_length=1024,
        truncation=True,
        return_tensors="pt"
    ).to(device_local)

    summary_ids = model.generate(inputs["input_ids"])
    summary = tokenizer.batch_decode(
        summary_ids,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False
    )
    plot = list(summary[0].split('.'))  # Split summary into sentences
    print(f"DEBUG: Summary generated: {plot}")

    # Prepare seed based on randomize_seed checkbox
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else seed

    # We'll generate a short video for each sentence
    # We'll also create an audio track for each sentence
    short_videos = []
    mp3_names = []
    mp3_lengths = []
    result_no_audio = f"result_no_audio_{uuid.uuid4().hex}.mp4" # Dynamic filename for no audio video
    movie_final = f'result_final_{uuid.uuid4().hex}.mp4' # Dynamic filename for final video
    merged_audio_path = "" # To store merged audio path for cleanup

    try: # Try-finally block to ensure cleanup
        for i, sentence in enumerate(plot[:-1]):
            # 1) Generate short video for this sentence
            prompt_for_animation = f"Generate a realistic video about this: {sentence}"
            print(f"DEBUG: Generating short video {i+1} of {len(plot)-1} ...")
            short_mp4_path = generate_short_animation(
                prompt_text=prompt_for_animation,
                base=base_model_name,
                motion=motion_name,
                step=int(num_inference_steps_backend),
                seed=current_seed + i, # Increment seed for each sentence for variation
                width=width,
                height=height
            )

            # 2) Generate audio for the sentence
            audio_filename = f'audio_{uuid.uuid4().hex}_{i}.mp3' # Dynamic audio filename
            tts_obj = gTTS(text=sentence, lang='en', slow=False)
            tts_obj.save(audio_filename)
            audio_info = MP3(audio_filename)
            audio_duration = audio_info.info.length
            mp3_names.append(audio_filename)
            mp3_lengths.append(audio_duration)

            # 3) Overlay subtitles on top of the short video (using Pillow now)
            final_clip_duration = audio_duration + 0.5  # half-second pad
            short_subtitled_path = add_subtitles_to_video(
                input_video_path=short_mp4_path,
                text=sentence.strip(),
                duration=final_clip_duration
            )
            short_videos.append(short_subtitled_path)

            # Clean up the original short clip (no subtitles)
            os.remove(short_mp4_path)

        # ----------------------------------------------------------------
        # Merge all MP3 files into one
        # ----------------------------------------------------------------
        merged_audio_path = merge_audio_files(mp3_names)

        # ----------------------------------------------------------------
        # Concatenate all short subtitled videos
        # ----------------------------------------------------------------
        print("DEBUG: Concatenating all short videos into a single clip...")
        clip_objects = []
        for vid_path in short_videos:
            clip = mpy.VideoFileClip(vid_path)
            clip_objects.append(clip)

        final_concat = mpy.concatenate_videoclips(clip_objects, method="compose")
        final_concat.write_videofile(result_no_audio, fps=24, logger=None)

        # ----------------------------------------------------------------
        # Combine big video with merged audio
        # ----------------------------------------------------------------
        def combine_audio(vidname, audname, outname, fps=24):
            print(f"DEBUG: Combining audio for video: '{vidname}'")
            my_clip = mpy.VideoFileClip(vidname)
            audio_background = mpy.AudioFileClip(audname)
            final_clip = my_clip.set_audio(audio_background)
            final_clip.write_videofile(outname, fps=fps, logger=None)
            my_clip.close()
            final_clip.close()

        combine_audio(result_no_audio, merged_audio_path, movie_final)

    finally: # Cleanup always executes
        print("DEBUG: Cleaning up temporary files...")
        # Remove short subtitled videos
        for path_ in short_videos:
            os.remove(path_)
        # Remove mp3 segments
        for f_mp3 in mp3_names:
            os.remove(f_mp3)
        # Remove merged audio
        if os.path.exists(merged_audio_path):
            os.remove(merged_audio_path)
        # Remove partial no-audio mp4
        if os.path.exists(result_no_audio):
            os.remove(result_no_audio)

    print("DEBUG: get_output_video function completed successfully.")
    return movie_final

# -------------------------------------------------------------------
# Example text (user can override)
# -------------------------------------------------------------------
text = (
    "Once, there was a girl called Laura who went to the supermarket to buy the ingredients to make a cake. "
    "Because today is her birthday and her friends come to her house and help her to prepare the cake."
)

# -------------------------------------------------------------------
# Gradio Interface
# -------------------------------------------------------------------
with gr.Blocks(css="style.css") as demo:
    gr.Markdown(
        """
        # Video Generator ⚡ from stories with Artificial Intelligence

        A story can be input by user. The story is summarized using DistilBART model.
        Then, the images are generated by using AnimateDiff and AnimateDiff-Lightning,
        and the subtitles and audio are created using gTTS. These are combined to generate a video.

        **Credits**: Developed by [ruslanmv.com](https://ruslanmv.com).
        """
    )

    with gr.Group():
        with gr.Row():
            input_start_text = gr.Textbox(value=text, label='Prompt')
        with gr.Row():
            select_base = gr.Dropdown(
                label='Base model',
                choices=["Cartoon", "Realistic", "3d", "Anime"],
                value=base_loaded,
                interactive=True
            )
            select_motion = gr.Dropdown(
                label='Motion',
                choices=[
                    ("Default", ""),
                    ("Zoom in", "guoyww/animatediff-motion-lora-zoom-in"),
                    ("Zoom out", "guoyww/animatediff-motion-lora-zoom-out"),
                    ("Tilt up", "guoyww/animatediff-motion-lora-tilt-up"),
                    ("Tilt down", "guoyww/animatediff-motion-lora-tilt-down"),
                    ("Pan left", "guoyww/animatediff-motion-lora-pan-left"),
                    ("Pan right", "guoyww/animatediff-motion-lora-pan-right"),
                    ("Roll left", "guoyww/animatediff-motion-lora-rolling-anticlockwise"),
                    ("Roll right", "guoyww/animatediff-motion-lora-rolling-clockwise"),
                ],
                value="",  # default: no motion lora
                interactive=True
            )
            select_step = gr.Dropdown(
                label='Inference steps',
                choices=[('1-Step', 1), ('2-Step', 2), ('4-Step', 4), ('8-Step', 8)],
                value=4,
                interactive=True
            )
            button_gen_video = gr.Button(
                scale=1,
                variant='primary',
                value="Generate Video"
            )

        with gr.Accordion("Advanced Settings", open=False):
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=42,
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE_720, # 제한 720 pixels maximum 사이즈, updated max size to 720p
                    step=1,
                    value=640, # Default width for 480p 4:3
                )
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE_720, # 제한 720 pixels maximum 사이즈, updated max size to 720p
                    step=1,
                    value=480, # Default height for 480p 4:3
                )


        with gr.Column():
            #output_interpolation = gr.Video(label="Generated Video")
            output_interpolation = gr.Video(value="video.mp4", label="Generated Video")  # Set default video


    button_gen_video.click(
        fn=get_output_video,
        inputs=[input_start_text, select_base, select_motion, select_step, randomize_seed, seed, width, height],
        outputs=output_interpolation
    )



demo.queue().launch(debug=True, share=False)