diff --git a/novelai_api/ImagePreset.py b/novelai_api/ImagePreset.py index ea2743b..8a31a34 100644 --- a/novelai_api/ImagePreset.py +++ b/novelai_api/ImagePreset.py @@ -265,6 +265,8 @@ class ImagePreset: #: Use the old behavior of prompt separation at the 75 tokens mark (can cut words in half) legacy_v3_extend: bool + #: ??? + params_version: int _settings: Dict[str, Any] diff --git a/novelai_api/NovelAI_API.py b/novelai_api/NovelAI_API.py index 2a2b807..2ea015e 100644 --- a/novelai_api/NovelAI_API.py +++ b/novelai_api/NovelAI_API.py @@ -14,6 +14,7 @@ class NovelAIAPI: # Constants + # TODO: might want to make the base endpoint configurable #: The base address for the API BASE_ADDRESS: str = "https://api.novelai.net" LIB_ROOT: str = dirname(abspath(__file__)) diff --git a/novelai_api/_low_level.py b/novelai_api/_low_level.py index caf1655..fc2845d 100644 --- a/novelai_api/_low_level.py +++ b/novelai_api/_low_level.py @@ -21,6 +21,8 @@ PRINT_WITH_PARAMETERS = os.environ.get("NAI_PRINT", False) +IMAGE_API_ADDRESS = "https://image.novelai.net" + # === INTERNALS === # SSE_FIELDS = ["event", "data", "id", "retry"] @@ -38,15 +40,15 @@ def print_with_parameters(args: Dict[str, Any]): if "parameters" in a: a["parameters"] = {k: str(v) for k, v in a["parameters"].items()} - for k in ["image", "mask", "controlnet_condition"]: - if k in a["parameters"]: - a["parameters"][k] = ( - f"{a['parameters'][k][:10]}...{a['parameters'][k][-10:]}" - if 30 < len(a["parameters"][k]) - else a["parameters"][k] - ) + for k in ["image", "mask", "controlnet_condition"]: + if k in a["parameters"]: + a["parameters"][k] = ( + f"{a['parameters'][k][:10]}...{a['parameters'][k][-10:]}" + if 30 < len(a["parameters"][k]) + else a["parameters"][k] + ) - print(json.dumps(a, indent=4)) + print(json.dumps(a, indent=4, sort_keys=True)) # === API === # @@ -181,7 +183,7 @@ async def _parse_response(cls, rsp: ClientResponse): elif content_type in ("audio/mpeg", "audio/webm"): yield await rsp.read() - elif content_type == "application/x-zip-compressed": + elif content_type in ("application/x-zip-compressed", "binary/octet-stream"): z = zipfile.ZipFile(io.BytesIO(await rsp.read())) for name in z.namelist(): yield name, z.read(name) @@ -195,16 +197,29 @@ async def _parse_response(cls, rsp: ClientResponse): url: str = rsp.url if isinstance(rsp.url, str) else rsp.url.human_repr() raise NovelAIError(url, -1, f"Unsupported type: {rsp.content_type}") - async def request(self, method: str, endpoint: str, data: Optional[Union[Dict[str, Any], str]] = None): + async def request( + self, + method: str, + endpoint: str, + data: Optional[Union[Dict[str, Any], str]] = None, + custom_base_address: Union[str, None] = None, + ): """ Send request with support for data streaming :param method: Method of the request (get, post, delete) :param endpoint: Endpoint of the request :param data: Data to pass to the method if needed + :param custom_base_address: Custom address to use for the request """ - url = f"{self._parent.BASE_ADDRESS}{endpoint}" + if PRINT_WITH_PARAMETERS: + print_with_parameters(data) + + if custom_base_address is None: + custom_base_address = self._parent.BASE_ADDRESS + + url = f"{custom_base_address}{endpoint}" is_sync = self._parent.session is None session = ClientSession() if is_sync else self._parent.session @@ -647,9 +662,6 @@ async def generate(self, prompt: Union[List[int], str], model: Model, params: Di endpoint = "/ai/generate-stream" if stream else "/ai/generate" - if PRINT_WITH_PARAMETERS: - print_with_parameters(data) - async for rsp, content in self.request("post", endpoint, data): self._treat_response_object(rsp, content, 201) @@ -680,7 +692,7 @@ async def generate_image( "parameters": parameters, } - async for rsp, content in self.request("post", "/ai/generate-image", data): + async for rsp, content in self.request("post", "/ai/generate-image", data, IMAGE_API_ADDRESS): self._treat_response_object(rsp, content, 200) yield content @@ -784,7 +796,7 @@ async def suggest_tags(self, tag: str, model: ImageModel) -> Dict[str, Any]: quote_via=quote, ) - async for rsp, content in self.request("get", f"/ai/generate-image/suggest-tags?{query}"): + async for rsp, content in self.request("get", f"/ai/generate-image/suggest-tags?{query}", IMAGE_API_ADDRESS): self._treat_response_object(rsp, content, 200) return content diff --git a/novelai_api/image_presets/presets_v1/default.preset b/novelai_api/image_presets/presets_v1/default.preset index 32fe0a8..f9a0b09 100644 --- a/novelai_api/image_presets/presets_v1/default.preset +++ b/novelai_api/image_presets/presets_v1/default.preset @@ -8,17 +8,14 @@ "noise": 0, "uc_preset": "Preset_Low_Quality_Bad_Anatomy", "quality_toggle": true, - "auto_smea": true, "smea": false, "smea_dyn": false, "decrisper": false, "controlnet_strength": 1, "legacy": false, "add_original_image": true, - "uncond_scale": 1, - "cfg_rescale": 0, - "noise_schedule": "native", "legacy_v3_extend": false, + "params_version": 1, "seed": 0, "uc": "" diff --git a/novelai_api/image_presets/presets_v2/default.preset b/novelai_api/image_presets/presets_v2/default.preset index a42b01d..4e03798 100644 --- a/novelai_api/image_presets/presets_v2/default.preset +++ b/novelai_api/image_presets/presets_v2/default.preset @@ -19,6 +19,7 @@ "cfg_rescale": 0, "noise_schedule": "native", "legacy_v3_extend": false, + "params_version": 1, "seed": 0, "uc": "" diff --git a/novelai_api/image_presets/presets_v3/default.preset b/novelai_api/image_presets/presets_v3/default.preset index ffcf0f5..bc8ea92 100644 --- a/novelai_api/image_presets/presets_v3/default.preset +++ b/novelai_api/image_presets/presets_v3/default.preset @@ -21,6 +21,7 @@ "legacy_v3_extend": false, "reference_information_extracted": 1, "reference_strength": 0.6, + "params_version": 1, "seed": 0, "uc": ""