RobinHood19/vlog_generation

1# --- 1. Imports ---2import os3import shutil4import tempfile5import uuid6import random7import traceback8from typing import List9from contextlib import asynccontextmanager10 11# Core libraries12import cv213import requests14import torch15import scipy.io.wavfile16import uvicorn17from PIL import Image18from moviepy import VideoFileClip, AudioFileClip19 20# FastAPI and Pydantic21from fastapi import FastAPI, HTTPException, BackgroundTasks22from pydantic import BaseModel, HttpUrl23from starlette.responses import FileResponse24 25# Transformers for music generation26from transformers import pipeline27 28# --- 2. Global State & Configuration for MusicGen ---29os.environ["TOKENIZERS_PARALLELISM"] = "false"30model_pipeline = None # This will hold the loaded MusicGen model31 32# --- 3. Lifespan Management (Model Loading/Unloading) ---33@asynccontextmanager34async def lifespan(app: FastAPI):35    """36    Manages the application's startup and shutdown events.37    The MusicGen model is loaded on startup to avoid reloading it on every request.38    """39    global model_pipeline40    print("Server starting up...")41    try:42        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")43        print(f"Using device for MusicGen: {'CUDA' if device.type == 'cuda' else 'CPU'}")44 45        print("Loading facebook/musicgen-small model...")46        model_pipeline = pipeline(47            "text-to-audio",48            model="facebook/musicgen-small",49            device=0 if device.type == 'cuda' else -150        )51        print("MusicGen model loaded successfully.")52    except Exception as e:53        print(f"FATAL: Could not load MusicGen model. Error: {e}")54        traceback.print_exc()55 56    yield  # The application runs after this point57 58    print("Server shutting down...")59    if model_pipeline:60        del model_pipeline61    if torch.cuda.is_available():62        torch.cuda.empty_cache()63    print("Cleaned up resources.")64 65# --- 4. FastAPI App Initialization ---66app = FastAPI(lifespan=lifespan)67 68# --- 5. Pydantic Models for API Requests ---69class VideoRequest(BaseModel):70    """Pydantic model to validate the incoming request."""71    image_urls: List[HttpUrl]72    prompt: str73 74# --- 6. Helper Functions ---75def generate_music(prompt: str, duration: int, file_path: str):76    """Generates music based on a prompt and saves it to a file."""77    global model_pipeline78    if model_pipeline is None:79        raise RuntimeError("Music generation model is not available.")80    81    try:82        print(f"Generating music for prompt: '{prompt}' with duration: {duration}s")83        max_length = int(duration * 50) # Heuristic: 50 tokens per second84        85        random_seed = random.randint(0, 2**32 - 1)86        torch.manual_seed(random_seed)87        if torch.cuda.is_available():88            torch.cuda.manual_seed_all(random_seed)89 90        music = model_pipeline(prompt, forward_params={"do_sample": True, "max_new_tokens": max_length})91 92        audio_data = music["audio"][0].T93        scipy.io.wavfile.write(file_path, rate=music["sampling_rate"], data=audio_data)94        print(f"Music saved to {file_path}")95        96    except Exception as e:97        print(f"Error during music generation: {e}")98        traceback.print_exc()99        raise100    finally:101        if torch.cuda.is_available():102            torch.cuda.empty_cache()103 104def process_and_create_video(image_folder: str, music_prompt: str) -> str:105    """106    Core logic: resizes images, creates a silent video, generates music,107    and merges them into a final video file.108    """109    # --- 1. Collect and validate images ---110    image_files = sorted([111        f for f in os.listdir(image_folder)112        if f.lower().endswith((".jpg", ".jpeg", ".png"))113    ])114    if not image_files:115        raise ValueError("No valid images found in the directory.")116    print(f"Found {len(image_files)} images.")117 118    # --- 2. Calculate dimensions and video duration ---119    total_width, total_height, readable_images_count = 0, 0, 0120    for file_name in image_files:121        try:122            with Image.open(os.path.join(image_folder, file_name)) as im:123                width, height = im.size124                total_width += width125                total_height += height126                readable_images_count += 1127        except IOError:128            print(f"Warning: Could not read {file_name}, skipping.")129    130    if readable_images_count == 0:131        raise ValueError("None of the image files could be read.")132 133    mean_width = int(total_width / readable_images_count)134    mean_height = int(total_height / readable_images_count)135    136    # --- Video settings (adjusted for short videos) ---137    FPS = 30138    STILL_DURATION_SEC = 1.5139    TRANSITION_DURATION_SEC = 0.5140    141    video_duration = (len(image_files) * STILL_DURATION_SEC) + \142                     ((len(image_files) - 1) * TRANSITION_DURATION_SEC)143    144    # Enforce 10-second limit145    if video_duration > 10.0:146        raise ValueError("The combination of images and durations exceeds the 10-second limit.")147    print(f"Calculated video duration: {video_duration:.2f}s")148 149 150    # --- 3. Resize images ---151    for file_name in image_files:152        file_path = os.path.join(image_folder, file_name)153        with Image.open(file_path) as im:154            im_resized = im.resize((mean_width, mean_height), Image.LANCZOS)155            im_resized.save(file_path, 'JPEG', quality=95)156 157    # --- 4. Generate silent video with transitions ---158    print("Generating silent video with transitions...")159    silent_video_path = os.path.join(image_folder, 'silent_video.mp4')160    video = cv2.VideoWriter(161        silent_video_path, cv2.VideoWriter_fourcc(*'mp4v'), FPS, (mean_width, mean_height)162    )163    164    still_frames = int(STILL_DURATION_SEC * FPS)165    transition_frames = int(TRANSITION_DURATION_SEC * FPS)166 167    for i, file_name in enumerate(image_files):168        current_frame = cv2.imread(os.path.join(image_folder, file_name))169        for _ in range(still_frames):170            video.write(current_frame)171        172        if i < len(image_files) - 1:173            next_frame = cv2.imread(os.path.join(image_folder, image_files[i + 1]))174            for j in range(transition_frames):175                alpha = j / (transition_frames - 1) if transition_frames > 1 else 1.0176                blended = cv2.addWeighted(current_frame, 1 - alpha, next_frame, alpha, 0)177                video.write(blended)178    video.release()179    print("Silent video generated.")180 181    # --- 5. Generate music ---182    music_path = os.path.join(image_folder, "background_music.wav")183    # Generate slightly longer audio to avoid abrupt cuts184    generate_music(music_prompt, int(video_duration) + 1, music_path)185 186    # --- 6. Merge video and audio using moviepy ---187    print("Merging video and audio...")188    final_video_path = os.path.join(image_folder, "final_video.mp4")189    190    video_clip = VideoFileClip(silent_video_path)191    audio_clip = AudioFileClip(music_path)192    193    # *** THIS IS THE FIX ***194    # Trim audio to match video duration exactly195    final_audio = audio_clip.subclipped(0, video_clip.duration)196    197    final_clip = video_clip.with_audio(final_audio)198    final_clip.write_videofile(final_video_path, codec='libx264', audio_codec='aac')199    200    # Close clips to release file handles201    audio_clip.close()202    video_clip.close()203    final_clip.close()204    205    print(f"Final video with audio saved to: {final_video_path}")206    return final_video_path207 208# --- 7. API Endpoint ---209@app.post("/create-video-with-music")210def create_video_from_urls(payload: VideoRequest, background_tasks: BackgroundTasks):211    """212    API endpoint to create a video with generated music from image URLs and a prompt.213    """214    if model_pipeline is None:215        raise HTTPException(status_code=503, detail="Model is not ready. Please try again later.")216        217    temp_dir = tempfile.mkdtemp()218    background_tasks.add_task(shutil.rmtree, temp_dir) # Cleanup after response219 220    print(f"Downloading {len(payload.image_urls)} images to {temp_dir}...")221    for i, url in enumerate(payload.image_urls):222        try:223            response = requests.get(str(url), stream=True, timeout=30)224            response.raise_for_status()225            file_extension = os.path.splitext(str(url.path))[-1].lower() or '.jpg'226            if file_extension not in ['.jpg', '.jpeg', '.png']:227                file_extension = '.jpg'228            file_path = os.path.join(temp_dir, f"image_{i:03d}{file_extension}")229            with open(file_path, 'wb') as f:230                shutil.copyfileobj(response.raw, f)231        except requests.exceptions.RequestException as e:232            raise HTTPException(status_code=400, detail=f"Failed to download image: {url}. Error: {e}")233 234    try:235        video_path = process_and_create_video(temp_dir, payload.prompt)236        return FileResponse(237            path=video_path,238            media_type='video/mp4',239            filename='generated_video_with_music.mp4'240        )241    except (ValueError, RuntimeError) as e:242        raise HTTPException(status_code=400, detail=str(e))243    except Exception as e:244        traceback.print_exc()245        raise HTTPException(status_code=500, detail=f"An internal error occurred: {e}")246 247# --- 8. Main execution block ---248if __name__ == '__main__':249    uvicorn.run(app, host="0.0.0.0", port=8000)250