Heartsync/TRELLIS2

1from typing import *
2import torch
3import torch.nn as nn
4from .. import models
5
6
7class Pipeline:
8    """
9    A base class for pipelines.
10    """
11    def __init__(
12        self,
13        models: dict[str, nn.Module] = None,
14    ):
15        if models is None:
16            return
17        self.models = models
18        for model in self.models.values():
19            model.eval()
20
21    @staticmethod
22    def from_pretrained(path: str) -> "Pipeline":
23        """
24        Load a pretrained model.
25        """
26        import os
27        import json
28        is_local = os.path.exists(f"{path}/pipeline.json")
29
30        if is_local:
31            config_file = f"{path}/pipeline.json"
32        else:
33            from huggingface_hub import hf_hub_download
34            config_file = hf_hub_download(path, "pipeline.json")
35
36        with open(config_file, 'r') as f:
37            args = json.load(f)['args']
38
39        _models = {}
40        for k, v in args['models'].items():
41            try:
42                _models[k] = models.from_pretrained(f"{path}/{v}")
43            except Exception as e:
44                _models[k] = models.from_pretrained(v)
45
46        new_pipeline = Pipeline(_models)
47        new_pipeline._pretrained_args = args
48        return new_pipeline
49
50    @property
51    def device(self) -> torch.device:
52        if hasattr(self, '_device'):
53            return self._device
54        for model in self.models.values():
55            if hasattr(model, 'device'):
56                return model.device
57        for model in self.models.values():
58            if hasattr(model, 'parameters'):
59                return next(model.parameters()).device
60        raise RuntimeError("No device found.")
61
62    def to(self, device: torch.device) -> None:
63        for model in self.models.values():
64            model.to(device)
65
66    def cuda(self) -> None:
67        self.to(torch.device("cuda"))
68
69    def cpu(self) -> None:
70        self.to(torch.device("cpu"))