Add User-Agent (#12)

* Feat: add User-Agent header --------- Signed-off-by: AnyISalIn <[email protected]>
omniinfer · Jul 11, 2023 · 93f575e · 93f575e
1 parent 85938ac
commit 93f575e
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/extension/api.py b/extension/api.py
@@ -12,11 +12,18 @@
 import json
 from types import SimpleNamespace
 from .utils import image_to_base64, read_image_files
+from .version import __version__
 
 OMNIINFER_CONFIG = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                 '.omniinfer.json')
 
 
+def _user_agent(model_name=None):
+    if model_name:
+        return 'sd-webui-cloud-inference/{} (model_name: {})'.format(__version__, model_name)
+    return 'sd-webui-cloud-inference/{}'.format(__version__)
+
+
 class BaseAPI(object):
 
     def txt2img(self, p) -> list[Image.Image]:
@@ -112,6 +119,7 @@ def __init__(self, token=None):
             self._token = token
         self._models = None
         self._session = requests.Session()
+        self._session.headers.update({'User-Agent': _user_agent()})
 
     @classmethod
     def load_from_config(cls):
@@ -231,7 +239,8 @@ def _txt2img(self, model_name, prompts, neg_prompts, sampler_name,
             "accept": "application/json",
             "content-type": "application/json",
             'Accept-Encoding': 'gzip, deflate',
-            "X-OmniInfer-Source": "source"
+            "X-OmniInfer-Source": _user_agent(model_name),
+            "User-Agent": _user_agent(model_name)
         }
 
         try:
@@ -284,7 +293,8 @@ def _img2img(self, model_name, prompts, neg_prompts, sampler_name,
             "accept": "application/json",
             "content-type": "application/json",
             'Accept-Encoding': 'gzip, deflate',
-            "X-OmniInfer-Source": "sd-webui"
+            "X-OmniInfer-Source": _user_agent(model_name),
+            "User-Agent": _user_agent(model_name)
         }
 
         res = requests.post("http://api.omniinfer.io/v2/img2img",
@@ -527,7 +537,8 @@ def refresh_models(self):
         headers = {
             "accept": "application/json",
             'Accept-Encoding': 'gzip, deflate',
-            "X-OmniInfer-Source": "sd-webui"
+            "X-OmniInfer-Source": _user_agent(),
+            "User-Agent": _user_agent()
         }
 
         print("[cloud-inference] refreshing models...")

diff --git a/extension/version.py b/extension/version.py
@@ -0,0 +1 @@
+__version__ = "0.1.2"