-
Notifications
You must be signed in to change notification settings - Fork 2.3k
/
data_classes.py
228 lines (175 loc) 路 6.63 KB
/
data_classes.py
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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
"""Pydantic data models and other dataclasses. This is the only file that uses Optional[]
typing syntax instead of | None syntax to work with pydantic"""
from __future__ import annotations
import pathlib
import secrets
import shutil
from abc import ABC, abstractmethod
from enum import Enum, auto
from typing import TYPE_CHECKING, Any, List, Optional, Union
from fastapi import Request
from gradio_client.utils import traverse
from . import wasm_utils
if not wasm_utils.IS_WASM or TYPE_CHECKING:
from pydantic import BaseModel, RootModel, ValidationError
else:
# XXX: Currently Pyodide V2 is not available on Pyodide,
# so we install V1 for the Wasm version.
from typing import Generic, TypeVar
from pydantic import BaseModel as BaseModelV1
from pydantic import ValidationError, schema_of
# Map V2 method calls to V1 implementations.
# Ref: https://docs.pydantic.dev/latest/migration/#changes-to-pydanticbasemodel
class BaseModelMeta(type(BaseModelV1)):
def __new__(cls, name, bases, dct):
# Override `dct` to dynamically create a `Config` class based on `model_config`.
if "model_config" in dct:
config_class = type("Config", (), {})
for key, value in dct["model_config"].items():
setattr(config_class, key, value)
dct["Config"] = config_class
del dct["model_config"]
model_class = super().__new__(cls, name, bases, dct)
return model_class
class BaseModel(BaseModelV1, metaclass=BaseModelMeta):
pass
BaseModel.model_dump = BaseModel.dict # type: ignore
BaseModel.model_json_schema = BaseModel.schema # type: ignore
# RootModel is not available in V1, so we create a dummy class.
PydanticUndefined = object()
RootModelRootType = TypeVar("RootModelRootType")
class RootModel(BaseModel, Generic[RootModelRootType]):
root: RootModelRootType
def __init__(self, root: RootModelRootType = PydanticUndefined, **data):
if data:
if root is not PydanticUndefined:
raise ValueError(
'"RootModel.__init__" accepts either a single positional argument or arbitrary keyword arguments'
)
root = data # type: ignore
# XXX: No runtime validation is executed.
super().__init__(root=root) # type: ignore
def dict(self, **kwargs):
return super().dict(**kwargs)["root"]
@classmethod
def schema(cls, **_kwargs):
# XXX: kwargs are ignored.
return schema_of(cls.__fields__["root"].type_) # type: ignore
RootModel.model_dump = RootModel.dict # type: ignore
RootModel.model_json_schema = RootModel.schema # type: ignore
class SimplePredictBody(BaseModel):
data: List[Any]
session_hash: Optional[str] = None
class PredictBody(BaseModel):
model_config = {"arbitrary_types_allowed": True}
session_hash: Optional[str] = None
event_id: Optional[str] = None
data: List[Any]
event_data: Optional[Any] = None
fn_index: Optional[int] = None
trigger_id: Optional[int] = None
simple_format: bool = False
batched: Optional[
bool
] = False # Whether the data is a batch of samples (i.e. called from the queue if batch=True) or a single sample (i.e. called from the UI)
request: Optional[
Request
] = None # dictionary of request headers, query parameters, url, etc. (used to to pass in request for queuing)
class ResetBody(BaseModel):
event_id: str
class ComponentServerBody(BaseModel):
session_hash: str
component_id: int
fn_name: str
data: Any
class InterfaceTypes(Enum):
STANDARD = auto()
INPUT_ONLY = auto()
OUTPUT_ONLY = auto()
UNIFIED = auto()
class GradioBaseModel(ABC):
def copy_to_dir(self, dir: str | pathlib.Path) -> GradioDataModel:
if not isinstance(self, (BaseModel, RootModel)):
raise TypeError("must be used in a Pydantic model")
dir = pathlib.Path(dir)
# TODO: Making sure path is unique should be done in caller
def unique_copy(obj: dict):
data = FileData(**obj)
return data._copy_to_dir(
str(pathlib.Path(dir / secrets.token_hex(10)))
).model_dump()
return self.__class__.from_json(
x=traverse(
self.model_dump(),
unique_copy,
FileData.is_file_data,
)
)
@classmethod
@abstractmethod
def from_json(cls, x) -> GradioDataModel:
pass
class GradioModel(GradioBaseModel, BaseModel):
@classmethod
def from_json(cls, x) -> GradioModel:
return cls(**x)
class GradioRootModel(GradioBaseModel, RootModel):
@classmethod
def from_json(cls, x) -> GradioRootModel:
return cls(root=x)
GradioDataModel = Union[GradioModel, GradioRootModel]
class FileData(GradioModel):
path: str # server filepath
url: Optional[str] = None # normalised server url
size: Optional[int] = None # size in bytes
orig_name: Optional[str] = None # original filename
mime_type: Optional[str] = None
is_stream: bool = False
@property
def is_none(self):
return all(
f is None
for f in [
self.path,
self.url,
self.size,
self.orig_name,
self.mime_type,
]
)
@classmethod
def from_path(cls, path: str) -> FileData:
return cls(path=path)
def _copy_to_dir(self, dir: str) -> FileData:
pathlib.Path(dir).mkdir(exist_ok=True)
new_obj = dict(self)
if not self.path:
raise ValueError("Source file path is not set")
new_name = shutil.copy(self.path, dir)
new_obj["path"] = new_name
return self.__class__(**new_obj)
@classmethod
def is_file_data(cls, obj: Any):
if isinstance(obj, dict):
try:
return not FileData(**obj).is_none
except (TypeError, ValidationError):
return False
return False
class ListFiles(GradioRootModel):
root: List[FileData]
def __getitem__(self, index):
return self.root[index]
def __iter__(self):
return iter(self.root)
class _StaticFiles:
"""
Class to hold all static files for an app
"""
all_paths = []
def __init__(self, paths: list[str | pathlib.Path]) -> None:
self.paths = paths
self.all_paths.extend([pathlib.Path(p).resolve() for p in paths])
@classmethod
def clear(cls):
cls.all_paths = []