diff options
Diffstat (limited to 'g4f')
| -rw-r--r-- | g4f/Provider/nexra/NexraSDTurbo.py | 81 | ||||
| -rw-r--r-- | g4f/models.py | 11 |
2 files changed, 51 insertions, 41 deletions
diff --git a/g4f/Provider/nexra/NexraSDTurbo.py b/g4f/Provider/nexra/NexraSDTurbo.py index da1428b8..865b4522 100644 --- a/g4f/Provider/nexra/NexraSDTurbo.py +++ b/g4f/Provider/nexra/NexraSDTurbo.py @@ -1,28 +1,26 @@ from __future__ import annotations -from aiohttp import ClientSession import json - -from ...typing import AsyncResult, Messages -from ..base_provider import AsyncGeneratorProvider, ProviderModelMixin +import requests +from ...typing import CreateResult, Messages +from ..base_provider import ProviderModelMixin, AbstractProvider from ...image import ImageResponse - -class NexraSDTurbo(AsyncGeneratorProvider, ProviderModelMixin): +class NexraSDTurbo(AbstractProvider, ProviderModelMixin): label = "Nexra Stable Diffusion Turbo" url = "https://nexra.aryahcr.cc/documentation/stable-diffusion/en" api_endpoint = "https://nexra.aryahcr.cc/api/image/complements" - working = False + working = True - default_model = 'sdxl-turbo' + default_model = "sdxl-turbo" models = [default_model] @classmethod def get_model(cls, model: str) -> str: return cls.default_model - + @classmethod - async def create_async_generator( + def create_completion( cls, model: str, messages: Messages, @@ -31,38 +29,41 @@ class NexraSDTurbo(AsyncGeneratorProvider, ProviderModelMixin): strength: str = 0.7, # Min: 0, Max: 1 steps: str = 2, # Min: 1, Max: 10 **kwargs - ) -> AsyncResult: + ) -> CreateResult: model = cls.get_model(model) - + headers = { - "Content-Type": "application/json" + 'Content-Type': 'application/json' } - async with ClientSession(headers=headers) as session: - prompt = messages[0]['content'] - data = { - "prompt": prompt, - "model": model, - "response": response, - "data": { - "strength": strength, - "steps": steps - } + + data = { + "prompt": messages[-1]["content"], + "model": model, + "response": response, + "data": { + "strength": strength, + "steps": steps } - async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response: - text_data = await response.text() - - if response.status == 200: - try: - json_start = text_data.find('{') - json_data = text_data[json_start:] - - data = json.loads(json_data) - if 'images' in data and len(data['images']) > 0: - image_url = data['images'][-1] - yield ImageResponse(image_url, prompt) - else: - yield ImageResponse("No images found in the response.", prompt) - except json.JSONDecodeError: - yield ImageResponse("Failed to parse JSON. Response might not be in JSON format.", prompt) + } + + response = requests.post(cls.api_endpoint, headers=headers, json=data) + + result = cls.process_response(response) + yield result + + @classmethod + def process_response(cls, response): + if response.status_code == 200: + try: + content = response.text.strip() + content = content.lstrip('_') # Remove the leading underscore + data = json.loads(content) + if data.get('status') and data.get('images'): + image_url = data['images'][0] + return ImageResponse(images=[image_url], alt="Generated Image") else: - yield ImageResponse(f"Request failed with status: {response.status}", prompt) + return "Error: No image URL found in the response" + except json.JSONDecodeError as e: + return f"Error: Unable to decode JSON response. Details: {str(e)}" + else: + return f"Error: {response.status_code}, Response: {response.text}" diff --git a/g4f/models.py b/g4f/models.py index 6f36892c..542967f2 100644 --- a/g4f/models.py +++ b/g4f/models.py @@ -53,6 +53,7 @@ from .Provider import ( NexraMidjourney, NexraQwen, NexraSD15, + NexraSDTurbo, OpenaiChat, PerplexityLabs, Pi, @@ -734,10 +735,17 @@ nemotron_70b = Model( ############# ### Stability AI ### +sdxl_turbo = Model( + name = 'sdxl-turbo', + base_provider = 'Stability AI', + best_provider = NexraSDTurbo + +) + sdxl = Model( name = 'sdxl', base_provider = 'Stability AI', - best_provider = IterListProvider([ReplicateHome, DeepInfraImage]) + best_provider = IterListProvider([ReplicateHome, DeepInfraImage, sdxl_turbo.best_provider]) ) @@ -1103,6 +1111,7 @@ class ModelUtils: ### Stability AI ### 'sdxl': sdxl, +'sdxl-turbo': sdxl_turbo, 'sd-1.5': sd_1_5, 'sd-3': sd_3, |