core/pipelines/base_pipeline.py

from abc import ABC, abstractmethod
from typing import List, Any, Dict
import gradio as gr
import spaces
import tempfile
import imageio
import numpy as np

class BasePipeline(ABC):
    def __init__(self):
        from core.model_manager import model_manager
        self.model_manager = model_manager

    @abstractmethod
    def get_required_models(self, **kwargs) -> List[str]:
        pass

    @abstractmethod
    def run(self, *args, progress: gr.Progress, **kwargs) -> Any:
        pass

    def _ensure_models_downloaded(self, progress: gr.Progress, **kwargs):
        """Ensures model files are downloaded before requesting GPU."""
        required_models = self.get_required_models(**kwargs)
        self.model_manager.ensure_models_downloaded(required_models, progress=progress)

    def _execute_gpu_logic(self, gpu_function: callable, duration: int, default_duration: int, task_name: str, *args, **kwargs):
        final_duration = default_duration
        try:
            if duration is not None and int(duration) > 0:
                final_duration = int(duration)
        except (ValueError, TypeError):
            print(f"Invalid ZeroGPU duration input for {task_name}. Using default {default_duration}s.")
            pass
        
        print(f"Requesting ZeroGPU for {task_name} with duration: {final_duration} seconds.")
        gpu_runner = spaces.GPU(duration=final_duration)(gpu_function)
        
        return gpu_runner(*args, **kwargs)

    def _encode_video_from_frames(self, frames_tensor_cpu: 'torch.Tensor', fps: int, progress: gr.Progress) -> str:
        progress(0.9, desc="Encoding video on CPU...")
        frames_np = (frames_tensor_cpu.numpy() * 255.0).astype(np.uint8)
        
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_video_file:
            video_path = temp_video_file.name
            writer = imageio.get_writer(video_path, fps=fps, codec='libx264', quality=8)
            for frame in frames_np:
                writer.append_data(frame)
            writer.close()
            
        progress(1.0, desc="Done!")
        return video_path