From 93f575eabfcbd95f5f036b599415c0d520c2e14f Mon Sep 17 00:00:00 2001
From: AnyISalIn <anyisalin@gmail.com>
Date: Tue, 11 Jul 2023 21:07:27 +0800
Subject: [PATCH] Add User-Agent (#12)

* Feat: add User-Agent header
---------

Signed-off-by: AnyISalIn <anyisalin@gmail.com>
---
 extension/api.py     | 17 ++++++++++++++---
 extension/version.py |  1 +
 2 files changed, 15 insertions(+), 3 deletions(-)
 create mode 100644 extension/version.py

diff --git a/extension/api.py b/extension/api.py
index c25ea00..493b268 100644
--- 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
new file mode 100644
index 0000000..48fef32
--- /dev/null
+++ b/extension/version.py
@@ -0,0 +1 @@
+__version__ = "0.1.2"
\ No newline at end of file