3DGen-Arena / model /model_manager.py
ZhangYuhan's picture
update serve
3e3ca46
raw
history blame
5.33 kB
import concurrent.futures
import random
import gradio as gr
import requests
import io, base64, json
# import spaces
from PIL import Image
from .model_config import model_config
from .model_worker import BaseModelWorker
class ModelManager:
def __init__(self):
self.models_config = model_config
self.models_worker: list[BaseModelWorker] = {}
self.build_model_workers()
def build_model_workers(self):
for cfg in self.models_config.values():
worker = BaseModelWorker(cfg.model_name, cfg.i2s_model, cfg.online_model, cfg.model_path)
self.models_worker[cfg.model_name] = worker
def get_all_models(self):
models = []
for model_name in self.models_config.keys():
models.append(model_name)
return models
def get_t2s_models(self):
models = []
for cfg in self.models_config.values():
if not cfg.i2s_model:
models.append(cfg.model_name)
return models
def get_i2s_models(self):
models = []
for cfg in self.models_config.values():
if cfg.i2s_model:
models.append(cfg.model_name)
return models
def get_online_models(self):
models = []
for cfg in self.models_config.values():
if cfg.online_model:
models.append(cfg.model_name)
return models
def get_models(self, i2s_model:bool, online_model:bool):
models = []
for cfg in self.models_config.values():
if cfg.i2s_model==i2s_model and cfg.online_model==online_model:
models.append(cfg.model_name)
return models
def check_online(self, name):
worker = self.models_worker[name]
if not worker.online_model:
return
# @spaces.GPU(duration=120)
def inference(self,
prompt, model_name,
offline=False, offline_idx=None):
result = None
worker = self.models_worker[model_name]
if offline:
result = worker.load_offline(offline_idx)
if not offline or result == None:
if worker.check_online():
result = worker.inference(prompt)
return result
def render(self, shape, model_name):
worker = self.models_worker[model_name]
result = worker.render(shape)
return result
def inference_parallel(self,
prompt, model_A, model_B,
offline=False, offline_idx=None):
results = []
model_names = [model_A, model_B]
with concurrent.futures.ThreadPoolExecutor() as executor:
future_to_result = {executor.submit(self.inference, prompt, model, offline, offline_idx): model
for model in model_names}
for future in concurrent.futures.as_completed(future_to_result):
result = future.result()
results.append(result)
return results[0], results[1]
def inference_parallel_anony(self,
prompt, model_A, model_B,
i2s_model: bool, offline: bool =False, offline_idx: int =None):
if model_A == model_B == "":
if offline and i2s_model:
model_A, model_B = random.sample(self.get_i2s_models(), 2)
elif offline and not i2s_model:
model_A, model_B = random.sample(self.get_t2s_models(), 2)
else:
model_A, model_B = random.sample(self.get_models(i2s_model=i2s_model, online_model=True), 2)
model_names = [model_A, model_B]
results = []
with concurrent.futures.ThreadPoolExecutor() as executor:
future_to_result = {executor.submit(self.inference, prompt, model, offline, offline_idx): model
for model in model_names}
for future in concurrent.futures.as_completed(future_to_result):
result = future.result()
results.append(result)
return results[0], results[1], model_A, model_B
def render_parallel(self, shape_A, model_A, shape_B, model_B):
results = []
model_names = [model_A, model_B]
shapes = [shape_A, shape_B]
with concurrent.futures.ThreadPoolExecutor() as executor:
future_to_result = {executor.submit(self.render, shape, model): model
for model, shape in zip(model_names, shapes)}
for future in concurrent.futures.as_completed(future_to_result):
result = future.result()
results.append(result)
return results[0], results[1]
# def i2s_inference_parallel(self, image, model_A, model_B):
# results = []
# model_names = [model_A, model_B]
# with concurrent.futures.ThreadPoolExecutor() as executor:
# future_to_result = {executor.submit(self.inference, image, model): model
# for model in model_names}
# for future in concurrent.futures.as_completed(future_to_result):
# result = future.result()
# results.append(result)
# return results[0], results[1]