feat:ai画画使用异步网络通信

zmdyy0318 · Oct 26, 2022 · b527907 · b527907
1 parent 7d39466
commit b527907
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/src/plugins/novelai/image.py b/src/plugins/novelai/image.py
@@ -278,13 +278,13 @@ async def __generate_image(self, token: str, model: str, keyword_en: str, explic
                         "seed": int(time.time()),
                     }
                 }
-
-            response = httpx.post(self.__url + "/ai/generate-image", headers=header, json=body, timeout=40)
-            if response.status_code != 201:
-                logger.error(f"Image::__generate_image failed, status_code:{response.status_code}")
-                return False, str(response.status_code), None
-            base64_str = response.text[27:]
-            return True, str(response.status_code), base64_str
+            async with httpx.AsyncClient(headers=header, timeout=40) as client:
+                response = await client.post(self.__url + "/ai/generate-image", headers=header, json=body, timeout=40)
+                if response.status_code != 201:
+                    logger.error(f"Image::__generate_image failed, status_code:{response.status_code}")
+                    return False, str(response.status_code), None
+                base64_str = response.text[27:]
+                return True, str(response.status_code), base64_str
         except (Exception,) as e:
             logger.error(f"Image::__generate_image failed, e:{e}")
             return False, str(e), None