app.py

"""
Hugging Face Wan2.2 Text-to-Video Application
A comprehensive text-to-video generation app using Wan2.2-S2V-14B model
"""

import os
import torch
import numpy as np
import cv2
from PIL import Image
from moviepy.editor import VideoFileClip, ImageSequenceClip
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from transformers import AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import hf_hub_download
import gradio as gr
from tqdm import tqdm
import yaml
import json
import logging
from datetime import datetime
import threading
import queue
import multiprocessing
from typing import List, Dict, Any, Optional
from dataclasses import dataclass
from enum import Enum
import hashlib
import base64
from io import BytesIO
import tempfile
import shutil
import subprocess
import sys
import signal
import psutil
import gc

# ==================== Configuration ====================
class Config:
    """Application configuration class"""
    
    MODEL_NAME = "Wan-AI/Wan2.2-S2V-14B"
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    MAX_VIDEO_LENGTH = 60  # seconds
    DEFAULT_FRAME_RATE = 24
    MAX_INPUT_LENGTH = 512
    BATCH_SIZE = 4
    NUM_INFERENCE_STEPS = 50
    GUIDANCE_SCALE = 7.5
    NEGATIVE_PROMPT = "blurry, low quality, distorted, ugly"
    
    # Cache settings
    CACHE_DIR = "./cache"
    MAX_CACHE_SIZE = 100
    
    # Logging
    LOG_LEVEL = logging.INFO
    LOG_FILE = "app.log"
    
    # UI settings
    THEME = "huggingface"
    DEFAULT_WIDTH = 800
    DEFAULT_HEIGHT = 600

# ==================== Data Classes ====================
@dataclass
class VideoFrame:
    """Represents a single video frame"""
    image: np.ndarray
    timestamp: float
    metadata: Dict[str, Any]

@dataclass
class VideoGenerationParams:
    """Parameters for video generation"""
    prompt: str
    negative_prompt: str
    duration: int
    frame_rate: int
    width: int
    height: int
    guidance_scale: float
    num_inference_steps: int
    seed: int

class VideoQuality(Enum):
    """Video quality presets"""
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    ULTRA = "ultra"

# ==================== Logging Setup ====================
class Logger:
    """Custom logging class"""
    
    def __init__(self):
        self.setup_logging()
    
    def setup_logging(self):
        """Setup logging configuration"""
        logging.basicConfig(
            level=Config.LOG_LEVEL,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(Config.LOG_FILE),
                logging.StreamHandler()
            ]
        )
        self.logger = logging.getLogger("Wan2App")
    
    def info(self, message: str):
        self.logger.info(message)
    
    def warning(self, message: str):
        self.logger.warning(message)
    
    def error(self, message: str):
        self.logger.error(message)
    
    def debug(self, message: str):
        self.logger.debug(message)

# ==================== Cache System ====================
class CacheManager:
    """Cache management for generated videos and frames"""
    
    def __init__(self):
        self.cache_dir = Config.CACHE_DIR
        self.max_cache_size = Config.MAX_CACHE_SIZE
        self.cache = {}
        self.setup_cache()
    
    def setup_cache(self):
        """Setup cache directory"""
        if not os.path.exists(self.cache_dir):
            os.makedirs(self.cache_dir)
        # Clean old cache
        self.cleanup_old_cache()
    
    def generate_cache_key(self, params: VideoGenerationParams) -> str:
        """Generate unique cache key from parameters"""
        param_str = f"{params.prompt}_{params.duration}_{params.frame_rate}_{params.width}_{params.height}_{params.seed}"
        return hashlib.md5(param_str.encode()).hexdigest()
    
    def save_to_cache(self, key: str, video_path: str):
        """Save video to cache"""
        if len(self.cache) >= self.max_cache_size:
            self.cleanup_oldest()
        cache_path = os.path.join(self.cache_dir, key)
        shutil.copy(video_path, cache_path)
        self.cache[key] = {"path": cache_path, "timestamp": datetime.now()}
    
    def get_from_cache(self, key: str) -> Optional[str]:
        """Get video from cache"""
        if key in self.cache:
            return self.cache[key]["path"]
        return None
    
    def cleanup_oldest(self):
        """Remove oldest cache entry"""
        if self.cache:
            oldest_key = min(self.cache.keys(), key=lambda k: self.cache[k]["timestamp"])
            cache_path = self.cache[oldest_key]["path"]
            if os.path.exists(cache_path):
                os.remove(cache_path)
            del self.cache[oldest_key]
    
    def cleanup_old_cache(self):
        """Cleanup old cache files"""
        for filename in os.listdir(self.cache_dir):
            file_path = os.path.join(self.cache_dir, filename)
            if os.path.isfile(file_path) and datetime.fromtimestamp(os.path.getmtime(file_path)) < datetime.now() - timedelta(days=7):
                os.remove(file_path)

# ==================== Model Manager ====================
class ModelManager:
    """Manages loading and using the Wan2.2 model"""
    
    def __init__(self):
        self.device = Config.DEVICE
        self.model = None
        self.tokenizer = None
        self.pipeline = None
        self.logger = Logger()
    
    def load_model(self):
        """Load the Wan2.2 model"""
        try:
            self.logger.info(f"Loading model on {self.device}...")
            
            # Load tokenizer
            self.tokenizer = AutoTokenizer.from_pretrained(Config.MODEL_NAME)
            
            # Load model
            self.model = AutoModelForCausalLM.from_pretrained(
                Config.MODEL_NAME,
                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
                device_map="auto"
            )
            
            # Setup pipeline
            self.pipeline = StableDiffusionPipeline.from_pretrained(
                Config.MODEL_NAME,
                scheduler=DPMSolverMultistepScheduler.from_config(
                    DPMSolverMultistepScheduler.from_pretrained(
                        Config.MODEL_NAME, subfolder="scheduler"
                    )
                ),
                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
            ).to(self.device)
            
            self.logger.info("Model loaded successfully!")
            return True
            
        except Exception as e:
            self.logger.error(f"Error loading model: {str(e)}")
            return False
    
    def generate_frame(self, prompt: str, negative_prompt: str = None, 
                      width: int = 512, height: int = 512, 
                      num_inference_steps: int = 50, guidance_scale: float = 7.5) -> np.ndarray:
        """Generate a single frame from text prompt"""
        try:
            # Tokenize prompt
            inputs = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
            
            # Generate image
            with torch.no_grad():
                image = self.pipeline(
                    prompt=prompt,
                    negative_prompt=negative_prompt or Config.NEGATIVE_PROMPT,
                    height=height,
                    width=width,
                    num_inference_steps=num_inference_steps,
                    guidance_scale=guidance_scale,
                    generator=torch.Generator(device=self.device).manual_seed(int(datetime.now().timestamp()))
                ).images[0]
            
            # Convert to numpy array
            frame = np.array(image)
            return frame
            
        except Exception as e:
            self.logger.error(f"Error generating frame: {str(e)}")
            return None
    
    def generate_video_frames(self, params: VideoGenerationParams) -> List[VideoFrame]:
        """Generate multiple frames for video"""
        frames = []
        total_frames = int(params.duration * params.frame_rate)
        
        self.logger.info(f"Generating {total_frames} frames...")
        
        for i in range(total_frames):
            progress = (i / total_frames) * 100
            if i % max(1, total_frames // 10) == 0:
                self.logger.info(f"Progress: {progress:.1f}%")
            
            # Create timestamped prompt
            timestamp = i / params.frame_rate
            prompt = f"{params.prompt} at time {timestamp:.2f}s"
            
            frame = self.generate_frame(
                prompt=prompt,
                negative_prompt=params.negative_prompt,
                width=params.width,
                height=params.height,
                num_inference_steps=params.num_inference_steps,
                guidance_scale=params.guidance_scale
            )
            
            if frame is not None:
                frames.append(VideoFrame(
                    image=frame,
                    timestamp=timestamp,
                    metadata={"frame_index": i, "prompt": prompt}
                ))
            
            # Memory management
            if self.device == "cuda":
                torch.cuda.empty_cache()
        
        return frames
    
    def cleanup(self):
        """Cleanup model resources"""
        if self.model:
            del self.model
        if self.pipeline:
            del self.pipeline
        if self.device == "cuda":
            torch.cuda.empty_cache()
        gc.collect()

# ==================== Video Processor ====================
class VideoProcessor:
    """Handles video creation and processing"""
    
    def __init__(self):
        self.logger = Logger()
    
    def create_video_from_frames(self, frames: List[VideoFrame], output_path: str, 
                                frame_rate: int = 24) -> bool:
        """Create video from list of frames"""
        try:
            # Convert PIL images to numpy arrays if needed
            image_arrays = []
            for frame in frames:
                if isinstance(frame.image, Image.Image):
                    image_arrays.append(np.array(frame.image))
                else:
                    image_arrays.append(frame.image)
            
            # Create video clip
            clip = ImageSequenceClip(image_arrays, fps=frame_rate)
            
            # Write video
            clip.write_videofile(output_path, codec='libx264', audio=False)
            
            self.logger.info(f"Video saved to {output_path}")
            return True
            
        except Exception as e:
            self.logger.error(f"Error creating video: {str(e)}")
            return False
    
    def add_audio_to_video(self, video_path: str, audio_path: str = None) -> str:
        """Add audio to video"""
        try:
            # If no audio provided, generate simple audio
            if not audio_path:
                audio_path = self.generate_simple_audio(video_path)
            
            if not audio_path:
                return video_path
            
            # Combine video and audio
            output_path = video_path.replace(".mp4", "_with_audio.mp4")
            cmd = [
                'ffmpeg', '-y', '-i', video_path, '-i', audio_path,
                '-c:v', 'copy', '-c:a', 'aac', '-strict', 'experimental',
                output_path
            ]
            
            subprocess.run(cmd, check=True)
            return output_path
            
        except Exception as e:
            self.logger.error(f"Error adding audio: {str(e)}")
            return video_path
    
    def generate_simple_audio(self, video_path: str) -> str:
        """Generate simple audio for video"""
        try:
            duration = self.get_video_duration(video_path)
            audio_path = video_path.replace(".mp4", ".wav")
            
            # Create simple audio
            cmd = [
                'ffmpeg', '-y', '-f', 'lavfi', '-i',
                f'sin(frequency=440:duration={duration})', '-ar', '44100',
                audio_path
            ]
            
            subprocess.run(cmd, check=True)
            return audio_path
            
        except Exception as e:
            self.logger.error(f"Error generating audio: {str(e)}")
            return None
    
    def get_video_duration(self, video_path: str) -> float:
        """Get video duration"""
        try:
            clip = VideoFileClip(video_path)
            duration = clip.duration
            clip.close()
            return duration
        except:
            return 5.0  # default duration

# ==================== Queue Manager ====================
class QueueManager:
    """Manages task queue for video generation"""
    
    def __init__(self):
        self.queue = queue.Queue()
        self.processed_tasks = 0
        self.max_workers = multiprocessing.cpu_count()
        self.lock = threading.Lock()
    
    def add_task(self, task_id: str, params: VideoGenerationParams):
        """Add task to queue"""
        self.queue.put((task_id, params))
    
    def process_queue(self, model_manager: ModelManager, video_processor: VideoProcessor):
        """Process tasks in queue"""
        while True:
            try:
                task_id, params = self.queue.get(timeout=1)
                self.process_task(task_id, params, model_manager, video_processor)
                self.queue.task_done()
            except queue.Empty:
                continue
    
    def process_task(self, task_id: str, params: VideoGenerationParams,
                    model_manager: ModelManager, video_processor: VideoProcessor):
        """Process a single task"""
        try:
            # Generate frames
            frames = model_manager.generate_video_frames(params)
            
            if not frames:
                raise Exception("No frames generated")
            
            # Create temporary video
            temp_dir = tempfile.mkdtemp()
            temp_video_path = os.path.join(temp_dir, "temp_video.mp4")
            
            if video_processor.create_video_from_frames(frames, temp_video_path, params.frame_rate):
                # Add audio
                final_video_path = video_processor.add_audio_to_video(temp_video_path)
                
                # Save result
                self.save_result(task_id, final_video_path)
                
                # Cleanup
                shutil.rmtree(temp_dir)
                
        except Exception as e:
            self.logger.error(f"Error processing task {task_id}: {str(e)}")

# ==================== UI Manager ====================
class UIManager:
    """Manages the user interface"""
    
    def __init__(self):
        self.model_manager = ModelManager()
        self.video_processor = VideoProcessor()
        self.cache_manager = CacheManager()
        self.queue_manager = QueueManager()
        self.logger = Logger()
        
        # Start queue processing thread
        self.queue_thread = threading.Thread(
            target=self.queue_manager.process_queue,
            args=(self.model_manager, self.video_processor),
            daemon=True
        )
        self.queue_thread.start()
    
    def setup_ui(self):
        """Setup Gradio UI"""
        with gr.Blocks(theme=gr.themes.HuggingFace()) as demo:
            gr.Markdown("# Hugging Face Wan2.2 Text-to-Video Generator")
            gr.Markdown("Generate videos from text prompts using Wan2.2-S2V-14B model")
            
            with gr.Row():
                with gr.Column():
                    prompt_input = gr.Textbox(
                        label="Text Prompt",
                        placeholder="Enter your video description...",
                        lines=3
                    )
                    
                    negative_prompt_input = gr.Textbox(
                        label="Negative Prompt",
                        placeholder="Enter what you don't want in the video...",
                        lines=2
                    )
                    
                    duration_slider = gr.Slider(
                        1, Config.MAX_VIDEO_LENGTH, value=10,
                        label="Video Duration (seconds)"
                    )
                    
                    frame_rate_slider = gr.Slider(
                        12, 60, value=Config.DEFAULT_FRAME_RATE,
                        label="Frame Rate (fps)"
                    )
                    
                    quality_dropdown = gr.Dropdown(
                        [q.value for q in VideoQuality],
                        value=VideoQuality.MEDIUM.value,
                        label="Quality Preset"
                    )
                    
                    seed_input = gr.Number(
                        value=int(datetime.now().timestamp()),
                        label="Random Seed"
                    )
                    
                    generate_btn = gr.Button("Generate Video")
                
                with gr.Column():
                    output_video = gr.Video(label="Generated Video")
                    progress_bar = gr.ProgressBar()
                    status_text = gr.Textbox(label="Status")
            
            # Connect buttons
            generate_btn.click(
                self.generate_video,
                inputs=[prompt_input, negative_prompt_input, duration_slider,
                       frame_rate_slider, quality_dropdown, seed_input],
                outputs=[output_video, progress_bar, status_text]
            )
        
        return demo
    
    def generate_video(self, prompt: str, negative_prompt: str,
                      duration: int, frame_rate: int, quality: str,
                      seed: int):
        """Generate video from user input"""
        if not prompt.strip():
            return None, None, "Please enter a prompt"
        
        try:
            # Get quality parameters
            params = self.get_quality_params(quality)
            
            # Create generation parameters
            gen_params = VideoGenerationParams(
                prompt=prompt,
                negative_prompt=negative_prompt or Config.NEGATIVE_PROMPT,
                duration=duration,
                frame_rate=frame_rate,
                width=params["width"],
                height=params["height"],
                guidance_scale=params["guidance_scale"],
                num_inference_steps=params["num_inference_steps"],
                seed=seed
            )
            
            # Check cache
            cache_key = self.cache_manager.generate_cache_key(gen_params)
            cached_video = self.cache_manager.get_from_cache(cache_key)
            
            if cached_video and os.path.exists(cached_video):
                self.logger.info(f"Video found in cache: {cache_key}")
                return cached_video, None, "Video loaded from cache"
            
            # Add to queue
            task_id = str(uuid.uuid4())
            self.queue_manager.add_task(task_id, gen_params)
            
            # Wait for completion (simplified)
            status = f"Processing: {prompt[:50]}..."
            return None, None, status
            
        except Exception as e:
            self.logger.error(f"Error generating video: {str(e)}")
            return None, None, f"Error: {str(e)}"
    
    def get_quality_params(self, quality: str) -> Dict[str, Any]:
        """Get parameters based on quality preset"""
        presets = {
            VideoQuality.LOW.value: {
                "width": 256,
                "height": 256,
                "guidance_scale": 5.0,
                "num_inference_steps": 20
            },
            VideoQuality.MEDIUM.value: {
                "width": 512,
                "height": 512,
                "guidance_scale": 7.5,
                "num_inference_steps": 50
            },
            VideoQuality.HIGH.value: {
                "width": 768,
                "height": 768,
                "guidance_scale": 8.5,
                "num_inference_steps": 75
            },
            VideoQuality.ULTRA.value: {
                "width": 1024,
                "height": 1024,
                "guidance_scale": 9.0,
                "num_inference_steps": 100
            }
        }
        return presets.get(quality, presets[VideoQuality.MEDIUM.value])
    
    def run(self):
        """Run the application"""
        if not self.model_manager.load_model():
            self.logger.error("Failed to load model. Exiting...")
            return
        
        demo = self.setup_ui()
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=True,
            debug=True
        )

# ==================== Main Application ====================
def main():
    """Main application entry point"""
    try:
        # Setup signal handlers
        signal.signal(signal.SIGINT, signal_handler)
        signal.signal(signal.SIGTERM, signal_handler)
        
        # Initialize UI
        ui = UIManager()
        ui.run()
        
    except KeyboardInterrupt:
        print("\nShutting down...")
    except Exception as e:
        print(f"Error: {str(e)}")

def signal_handler(signum, frame):
    """Signal handler for graceful shutdown"""
    print(f"\nReceived signal {signum}. Shutting down...")
    sys.exit(0)

if __name__ == "__main__":
    main()