autobyteus 1.2.0__py3-none-any.whl → 1.2.3__py3-none-any.whl

autobyteus/multimedia/audio/api/openai_audio_client.py

@@ -0,0 +1,112 @@
+import asyncio
+import logging
+import os
+import uuid
+from pathlib import Path
+from typing import Optional, Dict, Any, TYPE_CHECKING
+
+from openai import OpenAI
+
+from autobyteus.multimedia.audio.base_audio_client import BaseAudioClient
+from autobyteus.multimedia.utils.response_types import SpeechGenerationResponse
+
+if TYPE_CHECKING:
+    from autobyteus.multimedia.audio.audio_model import AudioModel
+    from autobyteus.multimedia.utils.multimedia_config import MultimediaConfig
+
+logger = logging.getLogger(__name__)
+
+_AUDIO_TEMP_DIR = Path("/tmp/autobyteus_audio")
+
+
+def _save_audio_bytes(audio_bytes: bytes, file_extension: Optional[str]) -> str:
+    """Saves audio bytes to disk with a random file name."""
+    _AUDIO_TEMP_DIR.mkdir(parents=True, exist_ok=True)
+    suffix = (file_extension or "mp3").lstrip(".")
+    file_path = _AUDIO_TEMP_DIR / f"{uuid.uuid4()}.{suffix}"
+    file_path.write_bytes(audio_bytes)
+    logger.info(f"Successfully saved generated audio to {file_path}")
+    return str(file_path)
+
+
+class OpenAIAudioClient(BaseAudioClient):
+    """
+    An audio client that uses OpenAI's Text-to-Speech (Speech) API.
+
+    **Setup Requirements:**
+    1. Set the `OPENAI_API_KEY` environment variable with your OpenAI API key.
+    """
+
+    def __init__(self, model: "AudioModel", config: "MultimediaConfig"):
+        super().__init__(model, config)
+        api_key = os.getenv("OPENAI_API_KEY")
+        if not api_key:
+            logger.error("OPENAI_API_KEY environment variable is not set.")
+            raise ValueError("OPENAI_API_KEY environment variable is not set.")
+
+        try:
+            self.client = OpenAI(api_key=api_key, base_url="https://api.openai.com/v1")
+            logger.info(f"OpenAIAudioClient initialized for model '{self.model.name}'.")
+        except Exception as exc:
+            logger.error(f"Failed to configure OpenAI client: {exc}")
+            raise RuntimeError(f"Failed to configure OpenAI client: {exc}") from exc
+
+    async def generate_speech(
+        self,
+        prompt: str,
+        generation_config: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> SpeechGenerationResponse:
+        """
+        Generates speech using OpenAI's Speech endpoint and returns a local file path.
+        """
+        try:
+            final_config = self.config.to_dict().copy()
+            if generation_config:
+                final_config.update(generation_config)
+
+            voice = final_config.get("voice", "alloy")
+            response_format = (
+                final_config.get("response_format")
+                or final_config.get("format")
+                or "mp3"
+            )
+            instructions = final_config.get("instructions")
+
+            logger.info(
+                "Generating speech with OpenAI TTS model '%s' using voice '%s' and format '%s'.",
+                self.model.value,
+                voice,
+                response_format,
+            )
+
+            request_kwargs = {
+                "model": self.model.value,
+                "voice": voice,
+                "input": prompt,
+            }
+
+            if instructions:
+                request_kwargs["instructions"] = instructions
+
+            if response_format:
+                request_kwargs["response_format"] = response_format
+
+            response = await asyncio.to_thread(
+                self.client.audio.speech.create,
+                **request_kwargs,
+            )
+
+            audio_bytes = getattr(response, "content", None)
+            if not audio_bytes:
+                raise ValueError("OpenAI Speech API returned an empty response.")
+
+            audio_path = _save_audio_bytes(audio_bytes, response_format)
+            return SpeechGenerationResponse(audio_urls=[audio_path])
+
+        except Exception as exc:
+            logger.error("Error during OpenAI speech generation: %s", exc)
+            raise ValueError(f"OpenAI speech generation failed: {exc}") from exc
+
+    async def cleanup(self):
+        logger.debug("OpenAIAudioClient cleanup called.")

autobyteus/multimedia/audio/audio_client_factory.py

@@ -5,6 +5,7 @@ from autobyteus.multimedia.audio.base_audio_client import BaseAudioClient
 from autobyteus.multimedia.audio.audio_model import AudioModel
 from autobyteus.multimedia.providers import MultimediaProvider
 from autobyteus.multimedia.audio.api.gemini_audio_client import GeminiAudioClient
+from autobyteus.multimedia.audio.api.openai_audio_client import OpenAIAudioClient
 from autobyteus.multimedia.audio.autobyteus_audio_provider import AutobyteusAudioModelProvider
 from autobyteus.multimedia.utils.multimedia_config import MultimediaConfig
 from autobyteus.utils.singleton import SingletonMeta
@@ -12,12 +13,55 @@ from autobyteus.utils.parameter_schema import ParameterSchema, ParameterDefiniti
 
 logger = logging.getLogger(__name__)
 
-GEMINI_TTS_VOICES = [
-    "Zephyr", "Puck", "Charon", "Kore", "Fenrir", "Leda", "Orus", "Aoede", 
-    "Callirrhoe", "Autonoe", "Enceladus", "Iapetus", "Umbriel", "Algieba", 
-    "Despina", "Erinome", "Algenib", "Rasalgethi", "Laomedeia", "Achernar", 
-    "Alnilam", "Schedar", "Gacrux", "Pulcherrima", "Achird", "Zubenelgenubi", 
-    "Vindemiatrix", "Sadachbia", "Sadaltager", "Sulafat"
+# Enhanced metadata for Google Gemini TTS voices, including gender and description.
+GEMINI_VOICE_DETAILS = {
+    "Zephyr": {"gender": "female", "description": "Bright, Higher pitch"},
+    "Puck": {"gender": "male", "description": "Upbeat, Middle pitch"},
+    "Charon": {"gender": "male", "description": "Informative, Lower pitch"},
+    "Kore": {"gender": "female", "description": "Firm, Middle pitch"},
+    "Fenrir": {"gender": "male", "description": "Excitable, Lower middle pitch"},
+    "Leda": {"gender": "female", "description": "Youthful, Higher pitch"},
+    "Orus": {"gender": "male", "description": "Firm, Lower middle pitch"},
+    "Aoede": {"gender": "female", "description": "Breezy, Middle pitch"},
+    "Callirrhoe": {"gender": "female", "description": "Easy-going, Middle pitch"},
+    "Autonoe": {"gender": "female", "description": "Bright, Middle pitch"},
+    "Enceladus": {"gender": "male", "description": "Breathy, Lower pitch"},
+    "Iapetus": {"gender": "male", "description": "Clear, Lower middle pitch"},
+    "Umbriel": {"gender": "male", "description": "Easy-going, Lower middle pitch"},
+    "Algieba": {"gender": "male", "description": "Smooth, Lower pitch"},
+    "Despina": {"gender": "female", "description": "Smooth, Middle pitch"},
+    "Erinome": {"gender": "female", "description": "Clear, Middle pitch"},
+    "Algenib": {"gender": "male", "description": "Gravelly, Lower pitch"},
+    "Rasalgethi": {"gender": "male", "description": "Informative, Middle pitch"},
+    "Laomedeia": {"gender": "female", "description": "Upbeat, Higher pitch"},
+    "Achernar": {"gender": "female", "description": "Soft, Higher pitch"},
+    "Alnilam": {"gender": "male", "description": "Firm, Lower middle pitch"},
+    "Schedar": {"gender": "male", "description": "Even, Lower middle pitch"},
+    "Gacrux": {"gender": "female", "description": "Mature, Middle pitch"},
+    "Pulcherrima": {"gender": "female", "description": "Forward, Middle pitch"},
+    "Achird": {"gender": "male", "description": "Friendly, Lower middle pitch"},
+    "Zubenelgenubi": {"gender": "male", "description": "Casual, Lower middle pitch"},
+    "Vindemiatrix": {"gender": "female", "description": "Gentle, Middle pitch"},
+    "Sadachbia": {"gender": "male", "description": "Lively, Lower pitch"},
+    "Sadaltager": {"gender": "male", "description": "Knowledgeable, Middle pitch"},
+    "Sulafat": {"gender": "female", "description": "Warm, Middle pitch"},
+}
+
+# The list of voice names, derived from the keys of the details dictionary.
+# This is used for the `enum_values` to maintain compatibility.
+GEMINI_TTS_VOICES = list(GEMINI_VOICE_DETAILS.keys())
+
+# Generate a formatted string of voice metadata to be appended to parameter descriptions.
+_voice_descriptions_list = [
+    f"- {name} ({details['gender']}): {details['description']}"
+    for name, details in GEMINI_VOICE_DETAILS.items()
+]
+GEMINI_VOICE_METADATA_DESC = "\n\nDetailed Voice Options:\n" + "\n".join(_voice_descriptions_list)
+
+
+OPENAI_TTS_VOICES = [
+    "alloy", "ash", "ballad", "coral", "echo", "fable", "onyx",
+    "nova", "sage", "shimmer", "verse"
 ]
 
 class AudioClientFactory(metaclass=SingletonMeta):
@@ -58,7 +102,7 @@ class AudioClientFactory(metaclass=SingletonMeta):
             ParameterDefinition(
                 name="voice",
                 param_type=ParameterType.ENUM,
-                description="The voice to assign to this speaker.",
+                description="The voice to assign to this speaker." + GEMINI_VOICE_METADATA_DESC,
                 enum_values=GEMINI_TTS_VOICES,
                 required=True
             )
@@ -78,7 +122,7 @@ class AudioClientFactory(metaclass=SingletonMeta):
                 param_type=ParameterType.ENUM,
                 default_value="Kore",
                 enum_values=GEMINI_TTS_VOICES,
-                description="The voice to use for single-speaker generation."
+                description="The voice to use for single-speaker generation." + GEMINI_VOICE_METADATA_DESC
             ),
             ParameterDefinition(
                 name="style_instructions",
@@ -96,12 +140,43 @@ class AudioClientFactory(metaclass=SingletonMeta):
         gemini_tts_model = AudioModel(
             name="gemini-2.5-flash-tts",
             value="gemini-2.5-flash-preview-tts",
-            provider=MultimediaProvider.GOOGLE,
+            provider=MultimediaProvider.GEMINI,
             client_class=GeminiAudioClient,
             parameter_schema=gemini_tts_schema
         )
 
+        openai_tts_schema = ParameterSchema(parameters=[
+            ParameterDefinition(
+                name="voice",
+                param_type=ParameterType.ENUM,
+                default_value="alloy",
+                enum_values=OPENAI_TTS_VOICES,
+                description="The OpenAI TTS voice to use for generation."
+            ),
+            ParameterDefinition(
+                name="format",
+                param_type=ParameterType.ENUM,
+                default_value="mp3",
+                enum_values=["mp3", "wav"],
+                description="The audio format to generate."
+            ),
+            ParameterDefinition(
+                name="instructions",
+                param_type=ParameterType.STRING,
+                description="Optional delivery instructions (tone, pacing, accent, etc.)."
+            )
+        ])
+
+        openai_tts_model = AudioModel(
+            name="gpt-4o-mini-tts",
+            value="gpt-4o-mini-tts",
+            provider=MultimediaProvider.OPENAI,
+            client_class=OpenAIAudioClient,
+            parameter_schema=openai_tts_schema
+        )
+
         models_to_register = [
+            openai_tts_model,
             gemini_tts_model,
         ]
         

autobyteus/multimedia/audio/audio_model.py

@@ -79,7 +79,8 @@ class AudioModel(metaclass=AudioModelMeta):
         """Returns the unique identifier for the model."""
         if self.runtime == MultimediaRuntime.AUTOBYTEUS and self.host_url:
             try:
-                host = urlparse(self.host_url).hostname
+                parsed = urlparse(self.host_url)
+                host = parsed.netloc or parsed.hostname or self.host_url
                 return f"{self.name}@{host}"
             except Exception:
                 return f"{self.name}@{self.host_url}" # Fallback

autobyteus/multimedia/image/api/autobyteus_image_client.py

@@ -1,4 +1,5 @@
 import logging
+import uuid
 from typing import Optional, List, Dict, Any, TYPE_CHECKING
 from autobyteus.clients import AutobyteusClient
 from autobyteus.multimedia.image.base_image_client import BaseImageClient
@@ -13,6 +14,7 @@ logger = logging.getLogger(__name__)
 class AutobyteusImageClient(BaseImageClient):
     """
     An image client that connects to an Autobyteus LLM server instance for image tasks.
+    Maintains a persistent session ID for stateful interactions (e.g. conversational editing).
     """
 
     def __init__(self, model: "ImageModel", config: "MultimediaConfig"):
@@ -21,7 +23,9 @@ class AutobyteusImageClient(BaseImageClient):
             raise ValueError("AutobyteusImageClient requires a host_url in its ImageModel.")
         
         self.autobyteus_client = AutobyteusClient(server_url=model.host_url)
-        logger.info(f"AutobyteusImageClient initialized for model '{self.model.name}' on host '{model.host_url}'.")
+        self.session_id = str(uuid.uuid4())
+        logger.info(f"AutobyteusImageClient initialized for model '{self.model.name}' "
+                    f"on host '{model.host_url}' with session_id '{self.session_id}'.")
 
     async def generate_image(
         self,
@@ -72,7 +76,7 @@ class AutobyteusImageClient(BaseImageClient):
     ) -> ImageGenerationResponse:
         """Internal helper to call the remote server."""
         try:
-            logger.info(f"Sending image generation request for model '{self.model.name}' to {self.model.host_url}")
+            logger.info(f"Sending image generation request for model '{self.model.name}' to {self.model.host_url} (Session: {self.session_id})")
             
             # The model name for the remote server is the `value`, not the unique `model_identifier`
             model_name_for_server = self.model.name
@@ -84,7 +88,8 @@ class AutobyteusImageClient(BaseImageClient):
                 prompt=prompt,
                 input_image_urls=input_image_urls,
                 mask_url=mask_url,
-                generation_config=generation_config
+                generation_config=generation_config,
+                session_id=self.session_id
             )
             
             image_urls = response_data.get("image_urls", [])
@@ -98,7 +103,16 @@ class AutobyteusImageClient(BaseImageClient):
             raise
 
     async def cleanup(self):
-        """Closes the underlying AutobyteusClient."""
+        """
+        Notifies the server to cleanup the session, then closes the underlying HTTP client.
+        """
         if self.autobyteus_client:
-            await self.autobyteus_client.close()
+            try:
+                logger.info(f"Notifying server to cleanup image session '{self.session_id}'...")
+                await self.autobyteus_client.cleanup_image_session(self.session_id)
+            except Exception as e:
+                logger.error(f"Failed to cleanup remote image session '{self.session_id}': {e}")
+            finally:
+                await self.autobyteus_client.close()
+        
         logger.debug("AutobyteusImageClient cleaned up.")

autobyteus/multimedia/image/api/gemini_image_client.py

@@ -1,14 +1,13 @@
 import logging
 import base64
-import os
 from typing import Optional, List, Dict, Any, TYPE_CHECKING
-from google import genai
-from PIL import Image
-import requests
+from google.genai import types as genai_types
 
 from autobyteus.multimedia.image.base_image_client import BaseImageClient
 from autobyteus.multimedia.utils.response_types import ImageGenerationResponse
 from autobyteus.multimedia.utils.api_utils import load_image_from_url
+from autobyteus.utils.gemini_helper import initialize_gemini_client_with_runtime
+from autobyteus.utils.gemini_model_mapping import resolve_model_for_runtime
 
 if TYPE_CHECKING:
     from autobyteus.multimedia.image.image_model import ImageModel
@@ -21,17 +20,15 @@ class GeminiImageClient(BaseImageClient):
     An image client that uses Google's Gemini models for image generation tasks.
 
     **Setup Requirements:**
-    1.  **Authentication:** Set the `GEMINI_API_KEY` environment variable with your API key.
+    1.  **AI Studio Mode:** Set `GEMINI_API_KEY`.
+    2.  **Vertex AI Mode:** Set `VERTEX_AI_PROJECT` and `VERTEX_AI_LOCATION`.
     """
 
     def __init__(self, model: "ImageModel", config: "MultimediaConfig"):
         super().__init__(model, config)
-        api_key = os.getenv("GEMINI_API_KEY")
-        if not api_key:
-            raise ValueError("Please set the GEMINI_API_KEY environment variable.")
         
         try:
-            self.client = genai.Client()
+            self.client, self.runtime_info = initialize_gemini_client_with_runtime()
             self.async_client = self.client.aio
             logger.info(f"GeminiImageClient initialized for model '{self.model.name}'.")
         except Exception as e:
@@ -60,16 +57,40 @@ class GeminiImageClient(BaseImageClient):
                     except Exception as e:
                         logger.error(f"Skipping image at '{url}' due to loading error: {e}")
 
-            # Note: The google-genai library uses the synchronous client for the `.generate_content` method on a model
-            # even in an async context, as there isn't a direct async equivalent exposed for this specific call on the model object.
-            # We use the top-level async client for other potential future calls if the library API changes.
-            model_instance = self.client.get_generative_model(model_name=f"models/{self.model.value}")
-            response = await model_instance.generate_content_async(contents=content)
+            config_dict: Dict[str, Any] = {}
+            if self.config and self.config.params:
+                config_dict.update(self.config.params)
+            if generation_config:
+                config_dict.update(generation_config)
+            if "response_modalities" not in config_dict:
+                if getattr(self, "runtime_info", None) and self.runtime_info.runtime == "vertex":
+                    config_dict["response_modalities"] = ["TEXT", "IMAGE"]
+                else:
+                    config_dict["response_modalities"] = ["IMAGE"]
+            config = genai_types.GenerateContentConfig(**config_dict)
+            
+            # FIX: Removed 'models/' prefix from model_name to support Vertex AI
+            runtime_adjusted_model = resolve_model_for_runtime(
+                self.model.value,
+                modality="image",
+                runtime=getattr(self, "runtime_info", None) and self.runtime_info.runtime,
+            )
+            if runtime_adjusted_model != self.model.value:
+                logger.info(
+                    "Using runtime-adjusted Gemini image model '%s' (requested '%s').",
+                    runtime_adjusted_model,
+                    self.model.value,
+                )
+            response = await self.async_client.models.generate_content(
+                model=runtime_adjusted_model,
+                contents=content,
+                config=config,
+            )
 
 
             image_urls = []
-            for part in response.parts:
-                if part.inline_data and "image" in part.inline_data.mime_type:
+            for part in response.parts or []:
+                if part.inline_data and part.inline_data.mime_type and "image" in part.inline_data.mime_type:
                     image_bytes = part.inline_data.data
                     base64_image = base64.b64encode(image_bytes).decode("utf-8")
                     data_uri = f"data:{part.inline_data.mime_type};base64,{base64_image}"
@@ -77,7 +98,7 @@ class GeminiImageClient(BaseImageClient):
             
             if not image_urls:
                 # Check for a safety-related refusal to generate content
-                if response.prompt_feedback.block_reason:
+                if response.prompt_feedback and response.prompt_feedback.block_reason:
                     reason = response.prompt_feedback.block_reason.name
                     logger.error(f"Image generation blocked due to safety settings. Reason: {reason}")
                     raise ValueError(f"Image generation failed due to safety settings: {reason}")

autobyteus/multimedia/image/api/openai_image_client.py

@@ -1,9 +1,14 @@
 import logging
 import os
+import tempfile
+from pathlib import Path
 from typing import Optional, List, Dict, Any, TYPE_CHECKING
+
 from openai import OpenAI
+
 from autobyteus.multimedia.image.base_image_client import BaseImageClient
 from autobyteus.multimedia.utils.response_types import ImageGenerationResponse
+from autobyteus.utils.download_utils import download_file_from_url
 
 if TYPE_CHECKING:
     from autobyteus.multimedia.image.image_model import ImageModel
@@ -11,9 +16,19 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+
+def _mime_type_from_format(output_format: str) -> str:
+    fmt = (output_format or "png").lower()
+    if fmt in {"jpg", "jpeg"}:
+        return "image/jpeg"
+    if fmt == "webp":
+        return "image/webp"
+    return "image/png"
+
+
 class OpenAIImageClient(BaseImageClient):
     """
-    An image client that uses OpenAI's DALL-E models.
+    An image client that uses OpenAI's gpt-image series via the images API.
     """
 
     def __init__(self, model: "ImageModel", config: "MultimediaConfig"):
@@ -34,49 +49,68 @@ class OpenAIImageClient(BaseImageClient):
         **kwargs
     ) -> ImageGenerationResponse:
         """
-        Generates an image using an OpenAI DALL-E model via the v1/images/generations endpoint.
-        Note: This endpoint does not support image inputs, even for multimodal models like gpt-image-1.
+        Generates an image using OpenAI's images generation endpoint.
+        Note: This endpoint does not support image inputs.
         """
         if input_image_urls:
             logger.warning(
-                f"The OpenAI `images.generate` API used by this client does not support input images. "
-                f"The images provided for model '{self.model.value}' will be ignored. "
-                f"To use image inputs, a client based on the Chat Completions API is required."
+                "The OpenAI `images.generate` API used by this client does not support input images. "
+                "The images provided for model '%s' will be ignored. "
+                "To use image inputs, a client based on the Chat Completions API is required.",
+                self.model.value,
             )
 
         try:
             image_model = self.model.value
-            logger.info(f"Generating image with OpenAI model '{image_model}' and prompt: '{prompt[:50]}...'")
+            logger.info("Generating image with OpenAI model '%s' and prompt: '%s...'", image_model, prompt[:50])
 
             # Combine default config with any overrides
             final_config = self.config.to_dict().copy()
             if generation_config:
                 final_config.update(generation_config)
-
-            response = self.client.images.generate(
-                model=image_model,
-                prompt=prompt,
-                n=final_config.get("n", 1),
-                size=final_config.get("size", "1024x1024"),
-                quality=final_config.get("quality", "standard"),
-                style=final_config.get("style", "vivid"),
-                response_format="url"
+            # Always request a single image for simplicity
+            final_config["n"] = 1
+
+            request_kwargs = {
+                "model": image_model,
+                "prompt": prompt,
+                "n": 1,
+                "size": final_config.get("size", "1024x1024"),
+                "quality": final_config.get("quality", "standard"),
+            }
+            if "output_format" in final_config:
+                request_kwargs["output_format"] = final_config["output_format"]
+            if "output_compression" in final_config:
+                request_kwargs["output_compression"] = final_config["output_compression"]
+
+            response = self.client.images.generate(**request_kwargs)
+
+            output_format = final_config.get("output_format", "png")
+            mime_type = _mime_type_from_format(output_format)
+            image_urls_list: List[str] = []
+            for img in response.data:
+                if getattr(img, "url", None):
+                    image_urls_list.append(img.url)
+                elif getattr(img, "b64_json", None):
+                    image_urls_list.append(f"data:{mime_type};base64,{img.b64_json}")
+
+            revised_prompt: Optional[str] = (
+                response.data[0].revised_prompt
+                if response.data and hasattr(response.data[0], "revised_prompt")
+                else None
             )
 
-            image_urls_list: List[str] = [img.url for img in response.data if img.url]
-            revised_prompt: Optional[str] = response.data[0].revised_prompt if response.data and hasattr(response.data[0], 'revised_prompt') else None
-
             if not image_urls_list:
-                raise ValueError("OpenAI API did not return any image URLs.")
+                raise ValueError("OpenAI API did not return any image data.")
 
-            logger.info(f"Successfully generated {len(image_urls_list)} image(s).")
+            logger.info("Successfully generated %s image(s).", len(image_urls_list))
 
             return ImageGenerationResponse(
                 image_urls=image_urls_list,
                 revised_prompt=revised_prompt
             )
         except Exception as e:
-            logger.error(f"Error during OpenAI image generation: {str(e)}")
+            logger.error("Error during OpenAI image generation: %s", str(e))
             raise ValueError(f"OpenAI image generation failed: {str(e)}")
 
     async def edit_image(
@@ -95,49 +129,97 @@ class OpenAIImageClient(BaseImageClient):
 
         source_image_url = input_image_urls[0]
         if len(input_image_urls) > 1:
-            logger.warning(f"OpenAI edit endpoint only supports one input image. Using '{source_image_url}' and ignoring the rest.")
+            logger.warning(
+                "OpenAI edit endpoint only supports one input image. Using '%s' and ignoring the rest.",
+                source_image_url,
+            )
 
+        temp_image_path: Optional[Path] = None
+        temp_mask_path: Optional[Path] = None
         try:
-            logger.info(f"Editing image '{source_image_url}' with prompt: '{prompt[:50]}...'")
+            logger.info("Editing image '%s' with prompt: '%s...'", source_image_url, prompt[:50])
 
             # Combine default config with any overrides
             final_config = self.config.to_dict().copy()
             if generation_config:
                 final_config.update(generation_config)
-
-            with open(source_image_url, "rb") as image_file:
-                mask_file = open(mask_url, "rb") if mask_url else None
+            # Always request a single edited image
+            final_config["n"] = 1
+
+            source_path = Path(source_image_url)
+            if not source_path.exists():
+                temp_image_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+                temp_image_file.close()
+                temp_image_path = Path(temp_image_file.name)
+                await download_file_from_url(source_image_url, temp_image_path)
+                source_path = temp_image_path
+
+            if mask_url:
+                mask_path = Path(mask_url)
+                if not mask_path.exists():
+                    temp_mask_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+                    temp_mask_file.close()
+                    temp_mask_path = Path(temp_mask_file.name)
+                    await download_file_from_url(mask_url, temp_mask_path)
+                    mask_path = temp_mask_path
+            else:
+                mask_path = None
+
+            with open(source_path, "rb") as image_file:
+                mask_file = open(mask_path, "rb") if mask_path else None
                 try:
-                    response = self.client.images.edit(
-                        image=image_file,
-                        mask=mask_file,
-                        prompt=prompt,
-                        model=self.model.value,
-                        n=final_config.get("n", 1),
-                        size=final_config.get("size", "1024x1024"),
-                        response_format="url"
-                    )
+                    request_kwargs = {
+                        "image": image_file,
+                        "prompt": prompt,
+                        "model": self.model.value,
+                        "n": final_config.get("n", 1),
+                        "size": final_config.get("size", "1024x1024"),
+                    }
+                    if mask_file:
+                        request_kwargs["mask"] = mask_file
+                    if "output_format" in final_config:
+                        request_kwargs["output_format"] = final_config["output_format"]
+                    if "output_compression" in final_config:
+                        request_kwargs["output_compression"] = final_config["output_compression"]
+                    response = self.client.images.edit(**request_kwargs)
                 finally:
                     if mask_file:
                         mask_file.close()
 
-            image_urls_list: List[str] = [img.url for img in response.data if img.url]
+            output_format = final_config.get("output_format", "png")
+            mime_type = _mime_type_from_format(output_format)
+            image_urls_list: List[str] = []
+            for img in response.data:
+                if getattr(img, "url", None):
+                    image_urls_list.append(img.url)
+                elif getattr(img, "b64_json", None):
+                    image_urls_list.append(f"data:{mime_type};base64,{img.b64_json}")
+
             if not image_urls_list:
-                raise ValueError("OpenAI API did not return any edited image URLs.")
+                raise ValueError("OpenAI API did not return any edited image data.")
 
-            logger.info(f"Successfully edited image, generated {len(image_urls_list)} version(s).")
+            logger.info("Successfully edited image, generated %s version(s).", len(image_urls_list))
             return ImageGenerationResponse(image_urls=image_urls_list)
 
         except FileNotFoundError as e:
-            logger.error(f"Image file not found for editing: {e.filename}")
+            logger.error("Image file not found for editing: %s", e.filename)
             raise
         except Exception as e:
-            logger.error(f"Error during OpenAI image editing: {str(e)}")
-            # The API might return a 400 Bad Request if the model doesn't support edits
+            logger.error("Error during OpenAI image editing: %s", str(e))
             if "does not support image editing" in str(e):
                 raise ValueError(f"The model '{self.model.value}' does not support the image editing endpoint.")
             raise ValueError(f"OpenAI image editing failed: {str(e)}")
-
+        finally:
+            if temp_image_path and temp_image_path.exists():
+                try:
+                    temp_image_path.unlink()
+                except OSError:
+                    logger.warning("Failed to clean up temp image file: %s", temp_image_path)
+            if temp_mask_path and temp_mask_path.exists():
+                try:
+                    temp_mask_path.unlink()
+                except OSError:
+                    logger.warning("Failed to clean up temp mask file: %s", temp_mask_path)
 
     async def cleanup(self):
         # The OpenAI client does not require explicit cleanup of a session.