Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,178 Bytes
d6c92ff |
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 |
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 |