File size: 1,183 Bytes
4e2a1b3 | 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 | import torch
from ..core import ModelConfig
from ..models.model_loader import ModelPool
class Metric(torch.nn.Module):
@staticmethod
def tensor_to_list(value):
if torch.is_tensor(value):
value = value.detach().cpu().tolist()
return value if isinstance(value, list) else [value]
@staticmethod
def download_and_load_models(model_configs: list[ModelConfig], torch_dtype: torch.dtype = torch.float32, device="cuda", vram_limit: float = None):
model_pool = ModelPool()
for model_config in model_configs:
model_config.download_if_necessary()
vram_config = model_config.vram_config()
vram_config["computation_dtype"] = vram_config["computation_dtype"] or torch_dtype or torch.float32
vram_config["computation_device"] = vram_config["computation_device"] or device
model_pool.auto_load_model(
model_config.path,
vram_config=vram_config,
vram_limit=vram_limit,
clear_parameters=model_config.clear_parameters,
state_dict=model_config.state_dict,
)
return model_pool
|