File size: 2,201 Bytes
d0dbbdd 305b06e d0dbbdd 815439b d0dbbdd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | 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
import sys
import os
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 |