Upgrade API dependencies, add Senty.io support, remove FastAPI analytics

.pre-commit-config.yaml

@@ -1,9 +1,6 @@
 repos:
-  - repo: https://github.com/psf/black
-    rev: 23.1.0
-    hooks:
-      - id: black
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.1.4
+    rev: v0.2.2
     hooks:
       - id: ruff
+      - id: ruff-format

api/app.py

@@ -4,8 +4,7 @@
 import os
 from pathlib import Path
 
-from api_analytics.fastapi import Analytics
-from fastapi import Depends, FastAPI, Request, status
+from fastapi import FastAPI, Request, status
 from fastapi.exceptions import RequestValidationError
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, JSONResponse
@@ -25,25 +24,13 @@
 
 logger = logging.getLogger(__name__)
 
-
-async def log_request(request: Request):
-    "Log all requests"
-
-    logger.debug(
-        f"url: {request.url}"
-        # f"url: {request.url}, params: {request.query_params}, body: {await request.body()}"
-    )
-
-
 app = FastAPI(
     # Docs
     docs_url="/api/docs",
     redoc_url="/api/redoc",
     swagger_ui_parameters={
         "filter": True,
     },
-    # Middleware
-    dependencies=[Depends(log_request)],
     # Metadata
     title="VoltaML",
     license_info={
@@ -62,27 +49,35 @@ async def validation_exception_handler(_request: Request, exc: RequestValidation
 
     logger.debug(exc)
 
-    if exc._error_cache is not None and exc._error_cache[0]["loc"][0] == "body":
+    errors = exc.errors()
+
+    if errors is not None and errors[0]["loc"][0] == "body":
         from core.config._config import Configuration
 
         default_value = Configuration()
-        keys = [str(i) for i in exc._error_cache[0]["loc"][1:]]  # type: ignore
-        current_value = exc._error_cache[0]["ctx"]["given"]  # type: ignore
-
-        # Traverse the config object to find the correct value
-        for key in keys:
-            default_value = getattr(default_value, key)
-
-        websocket_manager.broadcast_sync(
-            data=Data(
-                data={
-                    "default_value": default_value,
-                    "key": keys,
-                    "current_value": current_value,
-                },
-                data_type="incorrect_settings_value",
+        keys = [str(i) for i in errors[0]["loc"][1:]]  # type: ignore
+
+        try:
+            current_value = errors[0]["ctx"]["given"]  # type: ignore
+
+            # Traverse the config object to find the correct value
+            for key in keys:
+                default_value = getattr(default_value, key)
+
+            websocket_manager.broadcast_sync(
+                data=Data(
+                    data={
+                        "default_value": default_value,
+                        "key": keys,
+                        "current_value": current_value,
+                    },
+                    data_type="incorrect_settings_value",
+                )
+            )
+        except KeyError:
+            logger.info(
+                "Unable to parse incorrect settings value, skipping validation correction"
             )
-        )
 
     try:
         websocket_manager.broadcast_sync(
@@ -165,7 +160,9 @@ async def startup_event():
 
         for model in config.api.autoloaded_models:
             if model in [i.path for i in all_models]:
-                backend: InferenceBackend = [i.backend for i in all_models if i.path == model][0]  # type: ignore
+                backend: InferenceBackend = [
+                    i.backend for i in all_models if i.path == model
+                ][0]  # type: ignore
                 model_type = determine_model_type(get_full_model_path(model))[1]
 
                 gpu.load_model(model, backend, type=model_type)
@@ -185,12 +182,17 @@ async def shutdown_event():
 
 
 # Enable FastAPI Analytics if key is provided
-key = os.getenv("FASTAPI_ANALYTICS_KEY")
-if key:
-    app.add_middleware(Analytics, api_key=key)
-    logger.info("Enabled FastAPI Analytics")
+sentry_dsn = os.getenv("SENTRY_DSN")
+if sentry_dsn:
+    import sentry_sdk
+
+    sentry_sdk.init(
+        dsn=sentry_dsn,
+        enable_tracing=True,
+    )
+    logger.info("Enabled Sentry")
 else:
-    logger.debug("No FastAPI Analytics key provided, skipping")
+    logger.debug("No Sentry DSN provided, skipping")
 
 # Mount routers
 ## HTTP
@@ -226,19 +228,21 @@ async def shutdown_event():
 app.mount("/themes", StaticFiles(directory="data/themes"), name="themes")
 
 origins = ["*"]
+methods = ["*"]
+headers = ["*"]
 
 # Allow CORS for specified origins
 app.add_middleware(
     CORSMiddleware,
     allow_origins=origins,
     allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
+    allow_methods=methods,
+    allow_headers=headers,
 )
 static_app.add_middleware(
     CORSMiddleware,
     allow_origins=origins,
     allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
+    allow_methods=methods,
+    allow_headers=headers,
 )

api/routes/models.py

@@ -46,7 +46,8 @@ class UploadFileTarget(FileTarget):
     def __init__(self, dir_: Path, *args, **kwargs):
         super().__init__(None, *args, **kwargs)  # type: ignore
         self.file = UploadFile(
-            filename=None, file=NamedTemporaryFile(delete=False, dir=dir_)  # type: ignore
+            filename=None,
+            file=NamedTemporaryFile(delete=False, dir=dir_),  # type: ignore
         )
         self._fd = self.file.file
         self.dir = dir_

core/aitemplate/src/compile_lib/clip.py

@@ -60,10 +60,16 @@ def compile_clip(
         batch_size = IntVar(values=list(batch_size), name="batch_size")
 
     input_ids_ait = Tensor(
-        [batch_size, seqlen], name="input0", dtype="int64", is_input=True  # type: ignore
+        [batch_size, seqlen],
+        name="input0",
+        dtype="int64",
+        is_input=True,  # type: ignore
     )
     position_ids_ait = Tensor(
-        [batch_size, seqlen], name="input1", dtype="int64", is_input=True  # type: ignore
+        [batch_size, seqlen],
+        name="input1",
+        dtype="int64",
+        is_input=True,  # type: ignore
     )
     Y = ait_mod(input_ids=input_ids_ait, position_ids=position_ids_ait)
     mark_output(Y)

core/aitemplate/src/compile_lib/unet.py

@@ -167,11 +167,17 @@ def compile_unet(
         embedding_size = IntVar(values=list(clip_chunks), name="embedding_size")
 
     latent_model_input_ait = Tensor(
-        [batch_size, height_d, width_d, in_channels], name="input0", is_input=True, dtype=dtype  # type: ignore
+        [batch_size, height_d, width_d, in_channels],
+        name="input0",
+        is_input=True,
+        dtype=dtype,  # type: ignore
     )
     timesteps_ait = Tensor([batch_size], name="input1", is_input=True, dtype=dtype)
     text_embeddings_pt_ait = Tensor(
-        [batch_size, embedding_size, hidden_dim], name="input2", is_input=True, dtype=dtype  # type: ignore
+        [batch_size, embedding_size, hidden_dim],
+        name="input2",
+        is_input=True,
+        dtype=dtype,  # type: ignore
     )
 
     class_labels = None
@@ -181,7 +187,10 @@ def compile_unet(
     add_embeds = None
     if xl:
         add_embeds = Tensor(
-            [batch_size, projection_class_embeddings_input_dim], name="add_embeds", is_input=True, dtype=dtype  # type: ignore
+            [batch_size, projection_class_embeddings_input_dim],
+            name="add_embeds",
+            is_input=True,
+            dtype=dtype,  # type: ignore
         )
 
     down_block_residual_0 = None

core/aitemplate/src/modeling/clip.py

@@ -52,7 +52,8 @@ def __init__(
         self.to_k = nn.Linear(context_dim, inner_dim, bias=False, dtype=dtype)
         self.to_v = nn.Linear(context_dim, inner_dim, bias=False, dtype=dtype)
         self.to_out = nn.Sequential(
-            nn.Linear(inner_dim, query_dim, dtype=dtype), nn.Dropout(dropout, dtype=dtype)  # type: ignore
+            nn.Linear(inner_dim, query_dim, dtype=dtype),
+            nn.Dropout(dropout, dtype=dtype),  # type: ignore
         )
 
     def forward(

core/aitemplate/src/modeling/embeddings.py

@@ -33,7 +33,9 @@ def get_timestep_embedding(
     half_dim = embedding_dim // 2
 
     exponent = (-math.log(max_period)) * Tensor(
-        shape=[half_dim], dtype=dtype, name=arange_name  # type: ignore
+        shape=[half_dim],
+        dtype=dtype,
+        name=arange_name,  # type: ignore
     )
 
     exponent = exponent * (1.0 / (half_dim - downscale_freq_shift))

core/aitemplate/src/modeling/unet_2d_condition.py

@@ -110,7 +110,9 @@ def __init__(
 
         if addition_embed_type == "text_time":
             self.add_embedding = TimestepEmbedding(
-                projection_class_embeddings_input_dim, time_embed_dim, dtype=dtype  # type: ignore
+                projection_class_embeddings_input_dim,
+                time_embed_dim,
+                dtype=dtype,  # type: ignore
             )
 
         self.down_blocks = nn.ModuleList([])
@@ -257,7 +259,8 @@ def forward(
                 class_labels = self.time_proj(class_labels)
 
             class_emb = ops.batch_gather()(
-                self.class_embedding.weight.tensor(), class_labels  # type: ignore
+                self.class_embedding.weight.tensor(),
+                class_labels,  # type: ignore
             )
             emb = emb + class_emb
 
@@ -314,7 +317,9 @@ def forward(
         for upsample_block in self.up_blocks:
             res_samples = down_block_res_samples[-len(upsample_block.resnets) :]  # type: ignore
             down_block_res_samples = down_block_res_samples[
-                : -len(upsample_block.resnets)  # type: ignore
+                : -len(
+                    upsample_block.resnets
+                )  # type: ignore
             ]
 
             if (

core/aitemplate/src/modeling/vae.py

@@ -295,7 +295,13 @@ def decode(self, z: Tensor, return_dict: bool = True) -> Tensor:
         dec = self.decoder(z)
         return dec
 
-    def encode(self, x: Tensor, sample: Tensor = None, return_dict: bool = True, deterministic: bool = False) -> Tensor:  # type: ignore
+    def encode(
+        self,
+        x: Tensor,
+        sample: Tensor = None,
+        return_dict: bool = True,
+        deterministic: bool = False,
+    ) -> Tensor:  # type: ignore
         h = self.encoder(x)
         moments = self.quant_conv(h)
         if sample is None:

core/inference/ait/pipeline.py

@@ -355,7 +355,8 @@ def do_denoise(
                     # Infer ControlNet only for the conditional batch.
                     control_model_input = x
                     control_model_input = self.scheduler.scale_model_input(
-                        control_model_input, t  # type: ignore
+                        control_model_input,
+                        t,  # type: ignore
                     ).half()
                     controlnet_prompt_embeds = text_embeddings.chunk(2)[1]
                 else:
@@ -414,7 +415,11 @@ def do_denoise(
 
             if not isinstance(self.scheduler, KdiffusionSchedulerAdapter):
                 x = self.scheduler.step(
-                    noise_pred, t, x, **extra_step_kwargs, return_dict=False  # type: ignore
+                    noise_pred,
+                    t,
+                    x,
+                    **extra_step_kwargs,
+                    return_dict=False,  # type: ignore
                 )[0]
             else:
                 x = noise_pred
@@ -494,5 +499,6 @@ def _call(*args, **kwargs):
             return (image, has_nsfw_concept)
 
         return StableDiffusionPipelineOutput(
-            images=image, nsfw_content_detected=has_nsfw_concept  # type: ignore
+            images=image,
+            nsfw_content_detected=has_nsfw_concept,  # type: ignore
         )

core/inference/esrgan/utils/architecture/SPSR.py

@@ -66,7 +66,11 @@ def __init__(
             n_upscale = 1
 
         fea_conv = B.conv_block(
-            self.in_nc, self.num_filters, kernel_size=3, norm_type=None, act_type=None  # type: ignore
+            self.in_nc,
+            self.num_filters,
+            kernel_size=3,
+            norm_type=None,
+            act_type=None,  # type: ignore
         )
         rb_blocks = [
             B.RRDB(
@@ -131,7 +135,11 @@ def __init__(
         self.get_g_nopadding = Get_gradient_nopadding()
 
         self.b_fea_conv = B.conv_block(
-            self.in_nc, self.num_filters, kernel_size=3, norm_type=None, act_type=None  # type: ignore
+            self.in_nc,
+            self.num_filters,
+            kernel_size=3,
+            norm_type=None,
+            act_type=None,  # type: ignore
         )
 
         self.b_concat_1 = B.conv_block(
@@ -253,7 +261,11 @@ def __init__(
         self.b_module = B.sequential(*b_upsampler, b_HR_conv0, b_HR_conv1)
 
         self.conv_w = B.conv_block(
-            self.num_filters, self.out_nc, kernel_size=1, norm_type=None, act_type=None  # type: ignore
+            self.num_filters,
+            self.out_nc,
+            kernel_size=1,
+            norm_type=None,
+            act_type=None,  # type: ignore
         )
 
         self.f_concat = B.conv_block(
@@ -284,7 +296,11 @@ def __init__(
             act_type=act,
         )
         self.f_HR_conv1 = B.conv_block(
-            self.num_filters, self.out_nc, kernel_size=3, norm_type=None, act_type=None  # type: ignore
+            self.num_filters,
+            self.out_nc,
+            kernel_size=3,
+            norm_type=None,
+            act_type=None,  # type: ignore
         )
 
         self.load_state_dict(self.state, strict=False)

core/inference/functions.py

@@ -458,10 +458,14 @@ def new_forward(
                 setattr(pipe, name, nt)
                 del state_dict, dont_convert
 
-            text_encoder.text_model.encoder.old_forward = text_encoder.text_model.encoder.forward  # type: ignore
+            text_encoder.text_model.encoder.old_forward = (
+                text_encoder.text_model.encoder.forward
+            )  # type: ignore
             # fuck you python
             # enjoy bober
-            text_encoder.text_model.encoder.forward = partial(new_forward, bober=text_encoder.text_model.encoder)  # type: ignore
+            text_encoder.text_model.encoder.forward = partial(
+                new_forward, bober=text_encoder.text_model.encoder
+            )  # type: ignore
             logger.debug(f"Overwritten {name}s final_layer_norm.")
 
     if optimize: