diff --git a/src/leonardo_api/leonardo_async.py b/src/leonardo_api/leonardo_async.py index 5ae4ae3..e190b3f 100644 --- a/src/leonardo_api/leonardo_async.py +++ b/src/leonardo_api/leonardo_async.py @@ -83,10 +83,10 @@ async def get_user_info(self) -> dict: try: async with session.get(url) as response: response.raise_for_status() - response = await response.json() - self.___logger.debug(f"User info: {response}") + response_dict = await response.json() + self.___logger.debug(f"User info: {response_dict}") await session.close() - return response + return response_dict except Exception as error: self.___logger.error(f"Error occurred while getting user info: {str(error)}") if not session.closed: @@ -114,7 +114,7 @@ async def post_generations( prompt_magic: bool = True, control_net: bool = False, control_net_type: Optional[str] = None, - ) -> str: + ) -> dict: """ This endpoint will generate images. @@ -158,7 +158,7 @@ async def post_generations( :param control_net_type: The type of ControlNet to use. :type control_net_type: str, optional :return: generation response - :rtype: str + :rtype: dict Raises: Exception: if error occurred while post generations @@ -191,10 +191,10 @@ async def post_generations( try: async with session.post(url, json=payload) as response: response.raise_for_status() - response = await response.json() - self.___logger.debug(f"Post generations: {response}") + response_dict = await response.json() + self.___logger.debug(f"Post generations: {response_dict}") await session.close() - return response + return response_dict except Exception as error: self.___logger.error(f"Error occurred while post generations: {str(error)}") if not session.closed: @@ -219,10 +219,10 @@ async def get_single_generation(self, generation_id: str) -> dict: try: async with session.get(url) as response: response.raise_for_status() - response = await response.json() + response_dict = await response.json() self.___logger.debug(f"Single generations: {response}") await session.close() - return response + return response_dict except Exception as error: self.___logger.error(f"Error occurred while get single generations: {str(error)}") if not session.closed: @@ -279,10 +279,10 @@ async def get_generations_by_user(self, user_id: str, offset: int = 0, limit: in try: async with session.get(url, params=params) as response: response.raise_for_status() - response = await response.json() - self.___logger.debug(f"Generations for user {user_id} are: {response}") + response_dict = await response.json() + self.___logger.debug(f"Generations for user {user_id} are: {response_dict}") await session.close() - return response + return response_dict except Exception as error: self.___logger.error(f"Error occurred while obtaining user's generations: {str(error)}") if not session.closed: @@ -358,10 +358,10 @@ async def get_single_init_image(self, image_id: str) -> dict: try: async with session.get(url) as response: response.raise_for_status() - response = await response.json() - self.___logger.debug(f"Single image provided: {response}") + response_dict = await response.json() + self.___logger.debug(f"Single image provided: {response_dict}") await session.close() - return response + return response_dict except Exception as error: self.___logger.error(f"Error occurred while obtain single init image: {str(error)}") if not session.closed: @@ -395,7 +395,7 @@ async def delete_init_image(self, image_id: str) -> aiohttp.ClientResponse: await session.close() raise error - async def create_upscale(self, image_id: str) -> dict: + async def create_upscale(self, image_id: str) -> aiohttp.ClientResponse: """ This endpoint will create an upscale for the provided image ID. @@ -414,10 +414,10 @@ async def create_upscale(self, image_id: str) -> dict: try: async with session.post(url, json=payload) as response: response.raise_for_status() - response = await response.json() - self.___logger.debug(f"Upscale created: {response}") + response_dict = await response.json() + self.___logger.debug(f"Upscale created: {response_dict}") await session.close() - return response + return response_dict except Exception as error: self.___logger.error(f"Error occurred while up-scaling image: {str(error)}") if not session.closed: @@ -442,10 +442,10 @@ async def get_variation_by_id(self, generation_id: str) -> dict: try: async with session.get(url) as response: response.raise_for_status() - response = await response.json() - self.___logger.debug(f"Get variation by ID: {response}") + response_dict = await response.json() + self.___logger.debug(f"Get variation by ID: {response_dict}") await session.close() - return response + return response_dict except Exception as error: self.___logger.error(f"Error occurred while get variation by id: {str(error)}") if not session.closed: @@ -500,10 +500,10 @@ async def get_dataset_by_id(self, dataset_id: str) -> dict: try: async with session.get(url) as response: response.raise_for_status() - response = await response.json() - self.___logger.debug(f"Dataset with dataset_id={dataset_id} provided: {response}") + response_dict = await response.json() + self.___logger.debug(f"Dataset with dataset_id={dataset_id} provided: {response_dict}") await session.close() - return response + return response_dict except Exception as error: self.___logger.error(f"Error occurred while get dataset: {str(error)}") if not session.closed: @@ -714,10 +714,10 @@ async def get_custom_model_by_id(self, model_id: str) -> dict: try: async with session.get(url) as response: response.raise_for_status() - response = await response.json() - self.___logger.debug(f"Custom modal has been trained: {response}") + response_dict = await response.json() + self.___logger.debug(f"Custom modal has been trained: {response_dict}") await session.close() - return response + return response_dict except Exception as error: self.___logger.error(f"Error obtaining custom model: {str(error)}") if not session.closed: @@ -753,7 +753,7 @@ async def delete_custom_model_by_id(self, model_id: str) -> aiohttp.ClientRespon async def wait_for_image_generation( self, generation_id: str, image_index: int = 0, poll_interval: int = 5, timeout: int = 120 - ) -> dict: + ) -> aiohttp.ClientResponse: """ This method waits for the completion of image generation. @@ -766,7 +766,7 @@ async def wait_for_image_generation( :param timeout: (Optional) Waiting timeout. Default is 120 seconds. :type timeout: int, optional :return: The completed image(s) once generation is complete. - :rtype: dict + :rtype: aiohttp.ClientResponse :raises IndexError: If an invalid image_index is provided. :raises TimeoutError: If the image has not been generated in timeout seconds. diff --git a/src/leonardo_api/leonardo_sync.py b/src/leonardo_api/leonardo_sync.py index c928962..b5f9b6b 100644 --- a/src/leonardo_api/leonardo_sync.py +++ b/src/leonardo_api/leonardo_sync.py @@ -82,10 +82,10 @@ def get_user_info(self) -> dict: try: response = session.get(url) response.raise_for_status() - response = response.json() - self.___logger.debug(f"User info: {response}") + response_dict = response.json() + self.___logger.debug(f"User info: {response_dict}") session.close() - return response + return response_dict except Exception as error: self.___logger.error(f"Error occurred while getting user info: {str(error)}") raise error @@ -111,7 +111,7 @@ def post_generations( prompt_magic: bool = True, control_net: bool = False, control_net_type: Optional[str] = None, - ) -> dict: + ) -> requests.Response: """ This endpoint will generate images. @@ -155,7 +155,7 @@ def post_generations( :param control_net_type: The type of ControlNet to use. :type control_net_type: str, optional :return: The generation response. - :rtype: dict + :rtype: requests.Response """ # pylint: disable=too-many-locals url = "https://cloud.leonardo.ai/api/rest/v1/generations" @@ -185,9 +185,9 @@ def post_generations( try: response = session.post(url, json=payload) response.raise_for_status() - response = response.json() - self.___logger.debug(f"Post generations: {response}") - return response + response_dict = response.json() + self.___logger.debug(f"Post generations: {response_dict}") + return response_dict except Exception as error: self.___logger.error(f"Error occurred while post generations: {str(error)}") raise error @@ -210,9 +210,9 @@ def get_single_generation(self, generation_id: str) -> dict: try: response = session.get(url) response.raise_for_status() - response = response.json() - self.___logger.debug(f"Single generations: {response}") - return response + response_dict = response.json() + self.___logger.debug(f"Single generations: {response_dict}") + return response_dict except Exception as error: self.___logger.error(f"Error occurred while get single generations: {str(error)}") raise error @@ -264,9 +264,9 @@ def get_generations_by_user(self, user_id: str, offset: int = 0, limit: int = 10 try: response = session.get(url, params=params) response.raise_for_status() - response = response.json() + response_dict = response.json() self.___logger.debug(f"Generations for user {user_id} are: {response}") - return response + return response_dict except Exception as error: self.___logger.error(f"Error occurred while obtaining user's generations: {str(error)}") raise error @@ -332,9 +332,9 @@ def get_single_init_image(self, image_id: str) -> dict: try: response = session.get(url) response.raise_for_status() - response = response.json() - self.___logger.debug(f"Single image provided: {response}") - return response + response_dict = response.json() + self.___logger.debug(f"Single image provided: {response_dict}") + return response_dict except Exception as error: self.___logger.error(f"Error occurred while obtain single init image: {str(error)}") raise error @@ -390,14 +390,14 @@ def create_upscale(self, image_id: str) -> requests.Response: self.___logger.error(f"Error occurred while up-scaling image: {str(error)}") raise error - def get_variation_by_id(self, generation_id: str) -> dict: + def get_variation_by_id(self, generation_id: str) -> requests.Response: """ This endpoint will get the variation by ID. :param generation_id: The ID of the variation to get. :type generation_id: str :return: The variation. - :rtype: dict + :rtype: requests.Response Raises: Exception: If an error occurs while getting variation. @@ -408,9 +408,9 @@ def get_variation_by_id(self, generation_id: str) -> dict: try: response = session.get(url) response.raise_for_status() - response = response.json() - self.___logger.debug(f"Get variation by ID: {response}") - return response + response_dict = response.json() + self.___logger.debug(f"Get variation by ID: {response_dict}") + return response_dict except Exception as error: self.___logger.error(f"Error occurred while get variation by id: {str(error)}") raise error @@ -461,9 +461,9 @@ def get_dataset_by_id(self, dataset_id: str) -> dict: try: response = session.get(url) response.raise_for_status() - response = response.json() - self.___logger.debug(f"Dataset with dataset_id={dataset_id} provided: {response}") - return response + response_dict = response.json() + self.___logger.debug(f"Dataset with dataset_id={dataset_id} provided: {response_dict}") + return response_dict except Exception as error: self.___logger.error(f"Error occurred while get dataset: {str(error)}") raise error @@ -655,10 +655,10 @@ def get_custom_model_by_id(self, model_id: str) -> dict: try: response = session.get(url) response.raise_for_status() - response = response.json() - self.___logger.debug(f"Custom modal has been trained: {response}") + response_dict = response.json() + self.___logger.debug(f"Custom modal has been trained: {response_dict}") session.close() - return response + return response_dict except Exception as error: self.___logger.error(f"Error obtaining custom model: {str(error)}") raise error @@ -689,7 +689,7 @@ def delete_custom_model_by_id(self, model_id: str) -> requests.Response: def wait_for_image_generation( self, generation_id: str, image_index: int = 0, poll_interval: int = 5, timeout: int = 120 - ) -> dict: + ) -> requests.Response: """ This method waits for the completion of image generation. @@ -705,7 +705,7 @@ def wait_for_image_generation( :raises TimeoutError: If the image(s) have not been generated within the timeout. :raises IndexError: If an invalid image_index is provided. :return: The completed image(s) once generation is complete. - :rtype: dict + :rtype: requests.Response Raises: TimeoutError: If the image(s) have not been generated within the timeout.