.generate_image and .generate_text now returning request status class…

… for more info about request (like kudos cost)

README.md

@@ -25,7 +25,7 @@ async def image():
                 'AAM XL', 'AlbedoBase XL (SDXL)', 'Animagine XL', 'Anime Illust Diffusion', 'DreamShaper XL',
                 'ICBINP XL', 'Juggernaut XL', 'Quiet Goodnight XL', 'Unstable Diffusers XL']
     )
-    generations: list[models.GenerationStable] = await client.generate_image(generation_input)
+    generations: list[models.GenerationStable] = (await client.generate_image(generation_input)).generations
     for generation in generations:
         print(f'{generation.model}: {generation.img}')
 
@@ -52,7 +52,7 @@ async def text():
         params=params,
         models=['koboldcpp/Kunoichi-DPO-v2-7B-Q8_0-imatrix'] # i dont know which models are good for text :D
     )
-    results: list[models.GenerationKobold] = await client.generate_text(generation_input)
+    results: list[models.GenerationKobold] = (await client.generate_text(generation_input)).generations
     for result in results:
         print(result.text.strip())
 

aihorde/client.py

@@ -51,13 +51,13 @@ async def image_generation_status(
 
     async def generate_images(
         self, request: models.GenerationInputStable
-    ) -> list[models.GenerationStable]:
+    ) -> models.RequestStatusStable:
         resp: models.RequestAsync = await self.generate_image_request(request)
         request_id = resp.id
         while True:
             status = await self.image_generation_status(request_id)
             if status.done:
-                return status.generations
+                return status
             else:
                 await asyncio.sleep(int(status.wait_time / 1.5))
 
@@ -77,13 +77,13 @@ async def text_generation_status(
 
     async def generate_text(
         self, request: models.GenerationInputKobold
-    ) -> list[models.GenerationKobold]:
+    ) -> models.RequestStatusKobold:
         resp: models.RequestAsync = await self.generate_text_request(request)
         request_id = resp.id
         while True:
             status = await self.text_generation_status(request_id)
             if status.done:
-                return status.generations
+                return status
             else:
                 await asyncio.sleep(int(status.wait_time / 1.5))
 

test.py

@@ -16,7 +16,7 @@ async def text():
         params=params,
         models=['koboldcpp/Kunoichi-DPO-v2-7B-Q8_0-imatrix'] # i dont know which models are good for text
     )
-    results: list[models.GenerationKobold] = await client.generate_text(generation_input)
+    results: list[models.GenerationKobold] = (await client.generate_text(generation_input)).generations
     for result in results:
         print(result.text.strip())
 
@@ -29,7 +29,7 @@ async def image():
                 'AAM XL', 'AlbedoBase XL (SDXL)', 'Animagine XL', 'Anime Illust Diffusion', 'DreamShaper XL',
                 'ICBINP XL', 'Juggernaut XL', 'Quiet Goodnight XL', 'Unstable Diffusers XL']
     )
-    generations: list[models.GenerationStable] = await client.generate_image(generation_input)
+    generations: list[models.GenerationStable] = (await client.generate_image(generation_input)).generations
     for generation in generations:
         print(f'{generation.model}: {generation.img}')