Spaces:
Running
Running
File size: 2,761 Bytes
18869bb |
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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
from typing import Dict, List, Optional
import gradio as gr
from .base import BaseModel
class GradioModelInfo:
"""Information about a Gradio model."""
def __init__(self,
model_id: str,
name: str,
description: str,
input_type: str,
output_type: List[str],
examples: List[List[str]],
api_path: str):
self.model_id = model_id
self.name = name
self.description = description
self.input_type = input_type
self.output_type = output_type
self.examples = examples
self.api_path = api_path
self.status = "unloaded"
class GradioRegistry:
"""Registry for Gradio models."""
def __init__(self):
self._models: Dict[str, BaseModel] = {}
self._model_info: Dict[str, GradioModelInfo] = {}
def register_model(self,
model_id: str,
model: BaseModel,
api_path: str) -> None:
"""Register a new Gradio model."""
self._models[model_id] = model
# Create interface to extract information
interface = model.create_interface()
# Store model information
self._model_info[model_id] = GradioModelInfo(
model_id=model_id,
name=model.name,
description=model.description,
input_type=interface.input_components[0].__class__.__name__,
output_type=[comp.__class__.__name__ for comp in interface.output_components],
examples=interface.examples or [],
api_path=api_path
)
def get_model(self, model_id: str) -> Optional[BaseModel]:
"""Get a model by ID."""
return self._models.get(model_id)
def get_model_info(self, model_id: str) -> Optional[GradioModelInfo]:
"""Get model information by ID."""
return self._model_info.get(model_id)
def list_models(self) -> List[Dict]:
"""List all registered models."""
return [
{
"id": info.model_id,
"name": info.name,
"description": info.description,
"input_type": info.input_type,
"output_type": info.output_type,
"examples": info.examples,
"api_path": info.api_path,
"status": info.status
}
for info in self._model_info.values()
]
def update_model_status(self, model_id: str, status: str) -> None:
"""Update the status of a model."""
if model_id in self._model_info:
self._model_info[model_id].status = status |