Hugging Face's logo

Spaces:
yerang
/
LivePortrait
Running on Zero

App Files Community
LivePortrait
File size: 6,490 Bytes 
7931de6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
 
# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import typing

import numpy as np
import tritonclient.grpc
import tritonclient.http
import tritonclient.utils

from pytriton.model_config.generator import ModelConfigGenerator
from pytriton.model_config.triton_model_config import TritonModelConfig


def verify_equalness_of_dicts_with_ndarray(a_dict, b_dict):
    assert a_dict.keys() == b_dict.keys(), f"{a_dict} != {b_dict}"
    for output_name in a_dict:
        assert isinstance(
            a_dict[output_name], type(b_dict[output_name])
        ), f"type(a[{output_name}])={type(a_dict[output_name])} != type(b[{output_name}])={type(b_dict[output_name])}"
        if isinstance(a_dict[output_name], np.ndarray):
            assert a_dict[output_name].dtype == b_dict[output_name].dtype
            assert a_dict[output_name].shape == b_dict[output_name].shape
            if np.issubdtype(a_dict[output_name].dtype, np.number):
                assert np.allclose(b_dict[output_name], a_dict[output_name])
            else:
                assert np.array_equal(b_dict[output_name], a_dict[output_name])
        else:
            assert a_dict[output_name] == b_dict[output_name]


def wrap_to_grpc_infer_result(
    model_config: TritonModelConfig, request_id: str, outputs_dict: typing.Dict[str, np.ndarray]
):
    raw_output_contents = [output_data.tobytes() for output_data in outputs_dict.values()]
    return tritonclient.grpc.InferResult(
        tritonclient.grpc.service_pb2.ModelInferResponse(
            model_name=model_config.model_name,
            model_version=str(model_config.model_version),
            id=request_id,
            outputs=[
                tritonclient.grpc.service_pb2.ModelInferResponse.InferOutputTensor(
                    name=output_name,
                    datatype=tritonclient.utils.np_to_triton_dtype(output_data.dtype),
                    shape=output_data.shape,
                )
                for output_name, output_data in outputs_dict.items()
            ],
            raw_output_contents=raw_output_contents,
        )
    )


def wrap_to_http_infer_result(
    model_config: TritonModelConfig, request_id: str, outputs_dict: typing.Dict[str, np.ndarray]
):
    raw_output_contents = [output_data.tobytes() for output_data in outputs_dict.values()]
    buffer = b"".join(raw_output_contents)

    content = {
        "outputs": [
            {
                "name": name,
                "datatype": tritonclient.utils.np_to_triton_dtype(output_data.dtype),
                "shape": list(output_data.shape),
                "parameters": {"binary_data_size": len(output_data.tobytes())},
            }
            for name, output_data in outputs_dict.items()
        ]
    }
    header = json.dumps(content).encode("utf-8")
    response_body = header + buffer

    return tritonclient.http.InferResult.from_response_body(response_body, False, len(header))


def extract_array_from_grpc_infer_input(input_: tritonclient.grpc.InferInput):
    np_array = np.frombuffer(input_._raw_content, dtype=tritonclient.utils.triton_to_np_dtype(input_.datatype()))
    np_array = np_array.reshape(input_.shape())
    return np_array


def extract_array_from_http_infer_input(input_: tritonclient.http.InferInput):
    np_array = np.frombuffer(input_._raw_data, dtype=tritonclient.utils.triton_to_np_dtype(input_.datatype()))
    np_array = np_array.reshape(input_.shape())
    return np_array


def patch_grpc_client__server_up_and_ready(mocker):
    mocker.patch.object(tritonclient.grpc.InferenceServerClient, "is_server_ready").return_value = True
    mocker.patch.object(tritonclient.grpc.InferenceServerClient, "is_server_live").return_value = True


def patch_http_client__server_up_and_ready(mocker):
    mocker.patch.object(tritonclient.http.InferenceServerClient, "is_server_ready").return_value = True
    mocker.patch.object(tritonclient.http.InferenceServerClient, "is_server_live").return_value = True


def patch_grpc_client__model_up_and_ready(mocker, model_config: TritonModelConfig):
    from google.protobuf import json_format  # pytype: disable=pyi-error
    from tritonclient.grpc import model_config_pb2, service_pb2  # pytype: disable=pyi-error

    mock_get_repo_index = mocker.patch.object(tritonclient.grpc.InferenceServerClient, "get_model_repository_index")
    mock_get_repo_index.return_value = service_pb2.RepositoryIndexResponse(
        models=[
            service_pb2.RepositoryIndexResponse.ModelIndex(
                name=model_config.model_name, version="1", state="READY", reason=""
            ),
        ]
    )

    mocker.patch.object(tritonclient.grpc.InferenceServerClient, "is_model_ready").return_value = True

    model_config_dict = ModelConfigGenerator(model_config).get_config()
    model_config_protobuf = json_format.ParseDict(model_config_dict, model_config_pb2.ModelConfig())
    response = service_pb2.ModelConfigResponse(config=model_config_protobuf)
    response_dict = json.loads(json_format.MessageToJson(response, preserving_proto_field_name=True))
    mock_get_model_config = mocker.patch.object(tritonclient.grpc.InferenceServerClient, "get_model_config")
    mock_get_model_config.return_value = response_dict


def patch_http_client__model_up_and_ready(mocker, model_config: TritonModelConfig):
    mock_get_repo_index = mocker.patch.object(tritonclient.http.InferenceServerClient, "get_model_repository_index")
    mock_get_repo_index.return_value = [
        {"name": model_config.model_name, "version": "1", "state": "READY", "reason": ""}
    ]

    mocker.patch.object(tritonclient.http.InferenceServerClient, "is_model_ready").return_value = True

    model_config_dict = ModelConfigGenerator(model_config).get_config()
    mock_get_model_config = mocker.patch.object(tritonclient.http.InferenceServerClient, "get_model_config")
    mock_get_model_config.return_value = model_config_dict