Spaces:
Running on Zero
Running on Zero
File size: 2,714 Bytes
c009d4f 32997dd c009d4f 32997dd c009d4f | 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 56 57 58 59 60 61 62 63 64 65 | 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)
try:
return gpu_runner(*args, **kwargs)
except BaseException as e:
err_msg = str(e)
if "uncorrectable ECC error" in err_msg or "cudaErrorECCUncorrectable" in err_msg:
print("\n" + "="*80)
print(f"🚨 [Fatal GPU Error] Captured uncorrectable ECC error during inference: {err_msg}")
print("🚨 Terminating process to trigger an automatic container restart...")
print("="*80 + "\n")
os._exit(1)
raise e
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 |